From a3e2885e7ed7a188a312ad1c0ab5891355048d19 Mon Sep 17 00:00:00 2001 From: xlycae Date: Fri, 10 Jul 2026 15:07:24 +0800 Subject: [PATCH 1/2] feat: add fast-ulysses attention backend --- .../wan_i2v_dist_cfg_fast_ulysses.json | 18 + .../dist_infer/wan_i2v_dist_cfg_ulysses.json | 2 +- docs/fast_ulysses_backend.md | 61 ++ lightx2v/common/ops/attn/__init__.py | 1 + lightx2v/common/ops/attn/fast_ulysses_attn.py | 537 ++++++++++++++++++ lightx2v_fast_ulysses/.gitignore | 5 + lightx2v_fast_ulysses/CMakeLists.txt | 51 ++ lightx2v_fast_ulysses/NOTICE.md | 18 + lightx2v_fast_ulysses/README.md | 28 + lightx2v_fast_ulysses/__init__.py | 16 + lightx2v_fast_ulysses/comm.py | 154 +++++ lightx2v_fast_ulysses/csrc/a2a_config.cuh | 46 ++ lightx2v_fast_ulysses/csrc/all_to_all.cu | 301 ++++++++++ lightx2v_fast_ulysses/csrc/all_to_all_tma.cu | 305 ++++++++++ lightx2v_fast_ulysses/csrc/bindings.cpp | 103 ++++ lightx2v_fast_ulysses/csrc/symmetric_pool.cu | 82 +++ lightx2v_fast_ulysses/csrc/symmetric_pool.cuh | 42 ++ lightx2v_fast_ulysses/csrc/tma_ptx.cuh | 135 +++++ lightx2v_fast_ulysses/csrc/ulysses_common.cuh | 91 +++ lightx2v_fast_ulysses/csrc/ulysses_group.cu | 234 ++++++++ lightx2v_fast_ulysses/csrc/ulysses_group.cuh | 113 ++++ lightx2v_fast_ulysses/pyproject.toml | 18 + lightx2v_fast_ulysses/setup.py | 84 +++ pyproject.toml | 2 +- test_cases/test_fast_ulysses_backend.py | 352 ++++++++++++ 25 files changed, 2797 insertions(+), 2 deletions(-) create mode 100644 configs/dist_infer/wan_i2v_dist_cfg_fast_ulysses.json create mode 100644 docs/fast_ulysses_backend.md create mode 100644 lightx2v/common/ops/attn/fast_ulysses_attn.py create mode 100644 lightx2v_fast_ulysses/.gitignore create mode 100644 lightx2v_fast_ulysses/CMakeLists.txt create mode 100644 lightx2v_fast_ulysses/NOTICE.md create mode 100644 lightx2v_fast_ulysses/README.md create mode 100644 lightx2v_fast_ulysses/__init__.py create mode 100644 lightx2v_fast_ulysses/comm.py create mode 100644 lightx2v_fast_ulysses/csrc/a2a_config.cuh create mode 100644 lightx2v_fast_ulysses/csrc/all_to_all.cu create mode 100644 lightx2v_fast_ulysses/csrc/all_to_all_tma.cu create mode 100644 lightx2v_fast_ulysses/csrc/bindings.cpp create mode 100644 lightx2v_fast_ulysses/csrc/symmetric_pool.cu create mode 100644 lightx2v_fast_ulysses/csrc/symmetric_pool.cuh create mode 100644 lightx2v_fast_ulysses/csrc/tma_ptx.cuh create mode 100644 lightx2v_fast_ulysses/csrc/ulysses_common.cuh create mode 100644 lightx2v_fast_ulysses/csrc/ulysses_group.cu create mode 100644 lightx2v_fast_ulysses/csrc/ulysses_group.cuh create mode 100644 lightx2v_fast_ulysses/pyproject.toml create mode 100644 lightx2v_fast_ulysses/setup.py create mode 100644 test_cases/test_fast_ulysses_backend.py diff --git a/configs/dist_infer/wan_i2v_dist_cfg_fast_ulysses.json b/configs/dist_infer/wan_i2v_dist_cfg_fast_ulysses.json new file mode 100644 index 000000000..553e5d737 --- /dev/null +++ b/configs/dist_infer/wan_i2v_dist_cfg_fast_ulysses.json @@ -0,0 +1,18 @@ +{ + "infer_steps": 40, + "target_video_length": 81, + "target_height": 480, + "target_width": 832, + "self_attn_1_type": "flash_attn3", + "cross_attn_1_type": "flash_attn3", + "cross_attn_2_type": "flash_attn3", + "sample_guide_scale": 5, + "sample_shift": 5, + "enable_cfg": true, + "cpu_offload": false, + "parallel": { + "seq_p_size": 4, + "seq_p_attn_type": "fast_ulysses", + "cfg_p_size": 1 + } +} diff --git a/configs/dist_infer/wan_i2v_dist_cfg_ulysses.json b/configs/dist_infer/wan_i2v_dist_cfg_ulysses.json index 001f1f0c6..c0cbe83ea 100644 --- a/configs/dist_infer/wan_i2v_dist_cfg_ulysses.json +++ b/configs/dist_infer/wan_i2v_dist_cfg_ulysses.json @@ -13,6 +13,6 @@ "parallel": { "seq_p_size": 4, "seq_p_attn_type": "ulysses", - "cfg_p_size": 2 + "cfg_p_size": 1 } } diff --git a/docs/fast_ulysses_backend.md b/docs/fast_ulysses_backend.md new file mode 100644 index 000000000..9ca622335 --- /dev/null +++ b/docs/fast_ulysses_backend.md @@ -0,0 +1,61 @@ +# Fast Ulysses Backend + +LightX2V provides an optional `fast_ulysses` sequence-parallel attention backend that wraps the migrated fast-ulysses NVSHMEM all-to-all op around the existing LightX2V attention module. + +This PR only migrates the A2A path from `https://github.com/triple-mu/fast-ulysses`. Upstream fused QK/RoPE/RMSNorm ops are intentionally not included. + +## Install + +Build the optional native package only on machines that have NVSHMEM: + +```bash +NVSHMEM_HOME=/path/to/nvshmem pip install ./lightx2v_fast_ulysses +``` + +`NVSHMEM_HOME` must contain `include/nvshmem.h` and `lib/cmake/nvshmem`. + +On the local validation host: + +```bash +export NVSHMEM_HOME=/data1/lyxu18/workspace/nvshem +export LD_LIBRARY_PATH=$NVSHMEM_HOME/lib:$LD_LIBRARY_PATH +pip install -v ./lightx2v_fast_ulysses --no-build-isolation +``` + +## Configure + +Use the existing sequence-parallel attention selector: + +```json +{ + "parallel": { + "seq_p_attn_type": "fast_ulysses" + } +} +``` + +The fast path targets the validated single-node pure-image/self-attention path with fp16 or bf16 tensors. Unsupported paths such as fp8/fp4 communication, head parallelism, GQA, `q_only_img`, mixed text/image split attention, missing native extension, or cross-node sequence-parallel groups automatically fall back to the existing `ulysses` backend. + +## Validation + +Measured on one H800 node with GPUs 4/5/6/7, `seq_p_size=4`, `cfg_p_size=1`, bf16, and flash_attn3. The fast run used the native A2A path without fallback. + +| Scope | Shape / metric | `ulysses` | `fast_ulysses` | Speedup | +| --- | --- | ---: | ---: | ---: | +| Ulysses attention layer | `N=27280,H=40,D=128`, warmup=10, iters=50 | 7518.400 us | 6762.496 us | 1.112x | +| Wan I2V DiT steady state | average step time excluding first step | 1.601850 s | 1.540803 s | 1.040x | + +Full single-run e2e wall time is not used as the primary performance claim because model loading, cache state, and first-step native initialization add noise that is unrelated to the attention backend. + +## Third-Party Handling + +The fast-ulysses source used by this backend is migrated under `lightx2v_fast_ulysses/` and attributed to `https://github.com/triple-mu/fast-ulysses`. + +NVSHMEM remains an external dependency and is not a LightX2V top-level submodule. If cloning a LightX2V tree that does contain third-party submodules, use the usual recursive commands: + +```bash +git clone --recursive +git submodule update --init --recursive +``` + +This matches the existing LightX2V pattern for heavy CUDA dependencies such as CUTLASS: keep them external or fetch them in build images instead of making normal Python installs depend on them. diff --git a/lightx2v/common/ops/attn/__init__.py b/lightx2v/common/ops/attn/__init__.py index 64eff2e52..b430c1621 100755 --- a/lightx2v/common/ops/attn/__init__.py +++ b/lightx2v/common/ops/attn/__init__.py @@ -1,4 +1,5 @@ from .draft_attn import DraftAttnWeight +from .fast_ulysses_attn import FastUlyssesAttnWeight from .flash_attn import FlashAttn2Weight, FlashAttn3Weight, FlashAttn4Weight, SparseFlashAttn4Weight from .general_sparse_attn import GeneralSparseAttnWeight from .nbhd_attn import NbhdAttnWeight, NbhdAttnWeightFlashInfer diff --git a/lightx2v/common/ops/attn/fast_ulysses_attn.py b/lightx2v/common/ops/attn/fast_ulysses_attn.py new file mode 100644 index 000000000..28409e449 --- /dev/null +++ b/lightx2v/common/ops/attn/fast_ulysses_attn.py @@ -0,0 +1,537 @@ +import os +import socket + +import torch +import torch.distributed as dist +from loguru import logger + +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER + +from .template import AttnWeightTemplate +from .ulysses_attn import UlyssesAttnWeight + + +def _to_int(value): + if isinstance(value, torch.Tensor): + return int(value.item()) + return int(value) + + +def _cu_value(cu_seqlens, index): + return _to_int(cu_seqlens[index]) + + +def _env_use_tma(): + value = os.environ.get("LIGHTX2V_FAST_ULYSSES_USE_TMA", "auto").lower() + if value == "auto": + return None + if value in {"1", "true", "yes", "on"}: + return True + if value in {"0", "false", "no", "off"}: + return False + raise ValueError("LIGHTX2V_FAST_ULYSSES_USE_TMA must be auto, 1, or 0") + + +@ATTN_WEIGHT_REGISTER("fast_ulysses") +class FastUlyssesAttnWeight(AttnWeightTemplate): + _groups = {} + _single_node_cache = {} + _fallback_warnings = set() + _fast_path_logs = set() + _fallback_calls = 0 + _fast_path_calls = 0 + + def __init__(self): + self.config = {} + self._fallback_backend = UlyssesAttnWeight() + + @classmethod + def _cache_key(cls, seq_p_group, device=None): + group = seq_p_group if seq_p_group is not None else dist.group.WORLD + ranks = tuple(int(rank) for rank in dist.get_process_group_ranks(group)) + if device is None: + device_index = torch.cuda.current_device() + else: + device = torch.device(device) + device_index = device.index if device.index is not None else torch.cuda.current_device() + return ranks, device_index + + def _get_group(self, seq_p_group, device=None): + from lightx2v_fast_ulysses import UlyssesGroup + + key = self._cache_key(seq_p_group, device) + if key not in type(self)._groups: + pool_bytes = int(os.environ.get("LIGHTX2V_FAST_ULYSSES_POOL_BYTES", str(2 << 30))) + type(self)._groups[key] = UlyssesGroup(process_group=seq_p_group, device=device, initial_pool_bytes=pool_bytes) + return type(self)._groups[key] + + @classmethod + def destroy_cached_groups(cls): + for group in list(cls._groups.values()): + group.destroy() + cls._groups.clear() + cls._single_node_cache.clear() + cls.reset_runtime_stats() + + @classmethod + def reset_runtime_stats(cls): + cls._fallback_warnings.clear() + cls._fast_path_logs.clear() + cls._fallback_calls = 0 + cls._fast_path_calls = 0 + + @classmethod + def runtime_stats(cls): + return { + "fallback_calls": cls._fallback_calls, + "fast_path_calls": cls._fast_path_calls, + } + + @classmethod + def _rank_key(cls, seq_p_group): + group = seq_p_group if seq_p_group is not None else dist.group.WORLD + return tuple(int(rank) for rank in dist.get_process_group_ranks(group)) + + def _is_single_node(self, seq_p_group): + try: + key = self._rank_key(seq_p_group) + if len(key) <= 1: + return True + if key not in type(self)._single_node_cache: + host = socket.gethostname() + hosts = [None for _ in key] + dist.all_gather_object(hosts, host, group=seq_p_group) + type(self)._single_node_cache[key] = len(set(hosts)) == 1 + return type(self)._single_node_cache[key] + except Exception as exc: + self._log_fallback_once(f"single-node check failed: {exc}") + return False + + @classmethod + def _log_fallback_once(cls, reason, seq_p_group=None): + cls._fallback_calls += 1 + if reason in cls._fallback_warnings: + return + try: + rank = dist.get_rank(seq_p_group) + except Exception: + rank = 0 + if rank == 0: + logger.info("fast_ulysses fallback to ulysses: {}", reason) + cls._fallback_warnings.add(reason) + + def _log_fast_path_once(self, seq_p_group): + try: + key = self._rank_key(seq_p_group) + except Exception: + key = ("unknown",) + if key in type(self)._fast_path_logs: + return + try: + rank = dist.get_rank(seq_p_group) + except Exception: + rank = 0 + if rank == 0: + logger.info("fast_ulysses native A2A path active for ranks {}", key) + type(self)._fast_path_logs.add(key) + + def _fallback( + self, + reason, + q, + k, + v, + slice_qkv_len, + cu_seqlens_qkv, + attention_module=None, + seq_p_group=None, + use_fp8_comm=False, + use_fp4_comm=False, + use_tensor_fusion=False, + enable_head_parallel=False, + img_first=True, + q_only_img=False, + **kwargs, + ): + self._log_fallback_once(reason, seq_p_group) + return self._fallback_backend.apply( + q=q, + k=k, + v=v, + slice_qkv_len=slice_qkv_len, + cu_seqlens_qkv=cu_seqlens_qkv, + attention_module=attention_module, + seq_p_group=seq_p_group, + use_fp8_comm=use_fp8_comm, + use_fp4_comm=use_fp4_comm, + use_tensor_fusion=use_tensor_fusion, + enable_head_parallel=enable_head_parallel, + img_first=img_first, + q_only_img=q_only_img, + **kwargs, + ) + + def apply( + self, + q, + k, + v, + slice_qkv_len, + cu_seqlens_qkv, + attention_module=None, + seq_p_group=None, + use_fp8_comm=False, + use_fp4_comm=False, + use_tensor_fusion=False, + enable_head_parallel=False, + img_first=True, + q_only_img=False, + **kwargs, + ): + if attention_module is None: + raise ValueError("fast_ulysses requires attention_module") + if use_fp8_comm: + return self._fallback( + "use_fp8_comm is not supported by the fast A2A path", + q, + k, + v, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + if use_fp4_comm: + return self._fallback( + "use_fp4_comm is not supported by the fast A2A path", + q, + k, + v, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + if use_tensor_fusion: + return self._fallback( + "use_tensor_fusion is not supported by the fast A2A path", + q, + k, + v, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + if enable_head_parallel: + return self._fallback( + "enable_head_parallel is not supported by the fast A2A path", + q, + k, + v, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + if q_only_img: + return self._fallback( + "q_only_img is not supported by the fast A2A path", + q, + k, + v, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + if not img_first: + return self._fallback( + "img_first=False is not supported by the fast A2A path", + q, + k, + v, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + + q_orig, k_orig, v_orig = q, k, v + + if q.dim() == 4: + q = q.reshape(-1, q.shape[-2], q.shape[-1]) + k = k.reshape(-1, k.shape[-2], k.shape[-1]) + v = v.reshape(-1, v.shape[-2], v.shape[-1]) + + if q.shape != k.shape or q.shape != v.shape: + return self._fallback( + "GQA or mismatched q/k/v shapes are not supported by the fast A2A path", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + if q.device.type != "cuda" or k.device.type != "cuda" or v.device.type != "cuda": + return self._fallback( + "non-CUDA tensors are not supported by the fast A2A path", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + if q.device != k.device or q.device != v.device: + return self._fallback( + "q/k/v must be on the same CUDA device for the fast A2A path", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + if q.dtype != k.dtype or q.dtype != v.dtype or q.dtype not in (torch.float16, torch.bfloat16): + return self._fallback( + "CUDA tensors must be fp16/bf16 for the fast A2A path", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + + s_local, heads, head_dim = q.shape + if (head_dim * q.element_size()) % 16 != 0: + return self._fallback( + "head_dim * elem_size must be 16-byte aligned for the fast A2A path", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + world_size = dist.get_world_size(seq_p_group) + if world_size < 1 or world_size > 8: + return self._fallback( + "fast A2A path only supports single-node world_size in [1, 8]", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + if not self._is_single_node(seq_p_group): + return self._fallback( + "sequence-parallel group spans multiple nodes", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + if heads % world_size != 0: + return self._fallback( + "head count is not divisible by sequence-parallel world_size", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + + if _to_int(slice_qkv_len) != s_local or len(cu_seqlens_qkv) != 2 or _cu_value(cu_seqlens_qkv, 0) != 0 or _cu_value(cu_seqlens_qkv, 1) != s_local: + return self._fallback( + "mixed text/image split is not supported by the fast A2A path", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + + use_tma = _env_use_tma() + if use_tma is True: + device_major = torch.cuda.get_device_capability(q.device)[0] + if device_major < 9: + return self._fallback( + "LIGHTX2V_FAST_ULYSSES_USE_TMA=1 requires sm90+", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + + try: + group = self._get_group(seq_p_group, q.device) + except Exception as exc: + return self._fallback( + f"native fast-ulysses group is unavailable: {exc}", + q_orig, + k_orig, + v_orig, + slice_qkv_len, + cu_seqlens_qkv, + attention_module, + seq_p_group, + use_fp8_comm, + use_fp4_comm, + use_tensor_fusion, + enable_head_parallel, + img_first, + q_only_img, + **kwargs, + ) + global_seqlen = s_local * world_size + shard_heads = heads // world_size + self._log_fast_path_once(seq_p_group) + + qh = group.all_to_all_single_4d(q.reshape(1, s_local, heads, head_dim), mode=0, tag="lightx2v_q", use_tma=use_tma) + kh = group.all_to_all_single_4d(k.reshape(1, s_local, heads, head_dim), mode=0, tag="lightx2v_k", use_tma=use_tma) + vh = group.all_to_all_single_4d(v.reshape(1, s_local, heads, head_dim), mode=0, tag="lightx2v_v", use_tma=use_tma) + + attn = attention_module.apply( + q=qh[0], + k=kh[0], + v=vh[0], + max_seqlen_q=global_seqlen, + max_seqlen_kv=global_seqlen, + **kwargs, + ) + attn = attn.reshape(1, global_seqlen, shard_heads, head_dim) + out = group.all_to_all_single_4d(attn, mode=1, tag="lightx2v_out", use_tma=use_tma) + type(self)._fast_path_calls += 1 + return out[0].reshape(s_local, heads * head_dim) diff --git a/lightx2v_fast_ulysses/.gitignore b/lightx2v_fast_ulysses/.gitignore new file mode 100644 index 000000000..0c4157911 --- /dev/null +++ b/lightx2v_fast_ulysses/.gitignore @@ -0,0 +1,5 @@ +/_C*.so +/build/ +/dist/ +/*.egg-info/ +/__pycache__/ diff --git a/lightx2v_fast_ulysses/CMakeLists.txt b/lightx2v_fast_ulysses/CMakeLists.txt new file mode 100644 index 000000000..166aac561 --- /dev/null +++ b/lightx2v_fast_ulysses/CMakeLists.txt @@ -0,0 +1,51 @@ +# Copied and adapted from https://github.com/triple-mu/fast-ulysses +cmake_minimum_required(VERSION 3.18) +project(lightx2v_fast_ulysses LANGUAGES CXX CUDA) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CUDA_STANDARD 17) +if (NOT DEFINED CMAKE_CUDA_ARCHITECTURES OR "${CMAKE_CUDA_ARCHITECTURES}" STREQUAL "") + set(CMAKE_CUDA_ARCHITECTURES "80;90;100") +endif () + +# ccache: wrap CXX/CUDA compilers to cache objects and skip recompiling unchanged TUs (skipped if not installed) +find_program(CCACHE_PROGRAM ccache) +if (CCACHE_PROGRAM) + set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}") + set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}") + message(STATUS "ccache enabled: ${CCACHE_PROGRAM}") +endif () + +find_package(Python COMPONENTS Interpreter Development REQUIRED) +execute_process( + COMMAND ${Python_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)" + OUTPUT_VARIABLE TORCH_CMAKE_PREFIX_PATH OUTPUT_STRIP_TRAILING_WHITESPACE) +set(CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX_PATH}" CACHE STRING "Torch prefix" FORCE) +find_package(Torch REQUIRED) + +if (NOT DEFINED NVSHMEM_HOME OR NOT EXISTS "${NVSHMEM_HOME}/include/nvshmem.h") + message(FATAL_ERROR "NVSHMEM_HOME missing: setup.py must pass -DNVSHMEM_HOME=.") +endif () +find_package(NVSHMEM REQUIRED HINTS "${NVSHMEM_HOME}/lib/cmake/nvshmem") + +if (NOT DEFINED EXT_SUFFIX) + set(EXT_SUFFIX ".so") +endif () + +set(SOURCES + "${CMAKE_CURRENT_SOURCE_DIR}/csrc/bindings.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/csrc/all_to_all.cu" + "${CMAKE_CURRENT_SOURCE_DIR}/csrc/all_to_all_tma.cu" + "${CMAKE_CURRENT_SOURCE_DIR}/csrc/symmetric_pool.cu" + "${CMAKE_CURRENT_SOURCE_DIR}/csrc/ulysses_group.cu") +add_library(_C SHARED ${SOURCES}) +set_target_properties(_C PROPERTIES PREFIX "" SUFFIX "${EXT_SUFFIX}" + CUDA_SEPARABLE_COMPILATION OFF CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}") +target_include_directories(_C PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}/csrc" "${NVSHMEM_HOME}/include") +find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib" NO_DEFAULT_PATH REQUIRED) +target_link_libraries(_C PRIVATE ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY} Python::Module + nvshmem::nvshmem_host cuda) # cuda: CUDA driver, for cuTensorMapEncodeTiled (TMA) +target_compile_options(_C PRIVATE + $<$:-O3 --expt-relaxed-constexpr + -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ + -U__CUDA_NO_BFLOAT16_OPERATORS__ -U__CUDA_NO_BFLOAT16_CONVERSIONS__>) diff --git a/lightx2v_fast_ulysses/NOTICE.md b/lightx2v_fast_ulysses/NOTICE.md new file mode 100644 index 000000000..7ca26a16c --- /dev/null +++ b/lightx2v_fast_ulysses/NOTICE.md @@ -0,0 +1,18 @@ +# lightx2v-fast-ulysses Notice + +This optional package contains code migrated and adapted from: + +- Source: https://github.com/triple-mu/fast-ulysses +- Source commit: 8b0f7c65f57518ebdf593b957f1651cbde18f830 + +Copied/adapted files include: + +- `__init__.py` +- `comm.py` +- `setup.py` +- `CMakeLists.txt` +- A2A-related files under `csrc/` + +Upstream fused QK/RoPE/RMSNorm ops are not included in this LightX2V migration. + +Keep this notice with the migrated fast-ulysses code. diff --git a/lightx2v_fast_ulysses/README.md b/lightx2v_fast_ulysses/README.md new file mode 100644 index 000000000..fecb63952 --- /dev/null +++ b/lightx2v_fast_ulysses/README.md @@ -0,0 +1,28 @@ +# lightx2v-fast-ulysses + +Optional NVSHMEM-backed A2A extension for LightX2V's `fast_ulysses` attention backend. + +Only the fast-ulysses all-to-all path is migrated here. Upstream fused QK/RoPE/RMSNorm ops are out of scope for this package. + +Build: + +```bash +NVSHMEM_HOME=/path/to/nvshmem pip install ./lightx2v_fast_ulysses +``` + +Local validation path: + +```bash +NVSHMEM_HOME=/data1/workspace/nvshem \ +LD_LIBRARY_PATH=/data1/workspace/nvshem/lib:$LD_LIBRARY_PATH \ +pip install -v ./lightx2v_fast_ulysses --no-build-isolation +``` + +If cloning LightX2V with third-party submodules, use: + +```bash +git clone --recursive +git submodule update --init --recursive +``` + +`fast_ulysses` source is migrated here from https://github.com/triple-mu/fast-ulysses; NVSHMEM is still an external dependency. diff --git a/lightx2v_fast_ulysses/__init__.py b/lightx2v_fast_ulysses/__init__.py new file mode 100644 index 000000000..f019d930e --- /dev/null +++ b/lightx2v_fast_ulysses/__init__.py @@ -0,0 +1,16 @@ +# Copied and adapted from https://github.com/triple-mu/fast-ulysses +"""LightX2V-vendored fast-ulysses custom ops over the NVSHMEM symmetric heap.""" + +import torch # noqa: F401 load libtorch before dlopen of _C + +try: + from . import _C # noqa: F401,E402 trigger TORCH_LIBRARY registration +except ImportError as exc: + raise ImportError( + "lightx2v_fast_ulysses native extension is not built. " + "Build it with: NVSHMEM_HOME=/path/to/nvshmem pip install ./lightx2v_fast_ulysses" + ) from exc + +from .comm import AsyncA2AHandle, UlyssesGroup # noqa: E402 + +__all__ = ["UlyssesGroup", "AsyncA2AHandle", "_C"] diff --git a/lightx2v_fast_ulysses/comm.py b/lightx2v_fast_ulysses/comm.py new file mode 100644 index 000000000..fdf0b063c --- /dev/null +++ b/lightx2v_fast_ulysses/comm.py @@ -0,0 +1,154 @@ +# Copied and adapted from https://github.com/triple-mu/fast-ulysses +"""Python wrapper: build a C++ UlyssesGroup from a torch ProcessGroup (bootstrap is pure C++).""" + +from __future__ import annotations + +import os +from typing import Callable, Optional + +import torch +import torch.distributed as dist + + +class AsyncA2AHandle: + """Result of an async a2a: the collective runs on the group's comm stream; wait() makes the + CALLER's current stream wait for it (GPU-side event wait, host does not block) and returns the + output view. The output lives in the tag-scoped symmetric buffer -- do not issue another call + with the same tag until this result has been consumed.""" + + def __init__(self, out, ev_done: torch.cuda.Event): + self._out = out + self._ev_done = ev_done + + def wait(self): + torch.cuda.current_stream().wait_event(self._ev_done) + return self._out + + +class UlyssesGroup: + def __init__( + self, + process_group: Optional[dist.ProcessGroup] = None, + device: Optional[torch.device] = None, + initial_pool_bytes: int = 2 << 30, + ) -> None: + pg = process_group if process_group is not None else dist.group.WORLD + self.pg = pg + self.rank = dist.get_rank(pg) + self.world_size = dist.get_world_size(pg) + self.peer_global_ranks = list(dist.get_process_group_ranks(pg)) + if device is None: + device = torch.device("cuda", torch.cuda.current_device()) + self.device = device + torch.cuda.set_device(device) + + # Reservation must be set via env before NVSHMEM init. + os.environ["NVSHMEM_SYMMETRIC_SIZE"] = str(int(initial_pool_bytes)) + # P2P direct writes do not need NVLS (NVLink SHARP multicast); on some nodes its + # multicast heap mapping fails and segfaults, so disable by default for cross-node + # robustness (overridable via env). + os.environ.setdefault("NVSHMEM_DISABLE_NVLS", "1") + # This op is single-node NVLink P2P only; on nodes with IB NICs, NVSHMEM tries to init + # the IB remote transport and segfaults, so disable remote transport by default + # (verified on H200+IB nodes: init SIGSEGVs otherwise). + os.environ.setdefault("NVSHMEM_REMOTE_TRANSPORT", "none") + + cls = torch.classes.fast_ulysses.UlyssesGroup + if dist.get_rank() == 0: + uid = cls.get_uniqueid() + else: + uid = [0] * cls.uniqueid_nints() + uid_t = torch.tensor(uid, dtype=torch.int64, device=device) + dist.broadcast(uid_t, src=0, group=dist.group.WORLD) + cls.init_world(uid_t.tolist(), dist.get_rank(), dist.get_world_size()) + + dist.barrier(group=pg) + self._group = cls( + [int(r) for r in self.peer_global_ranks], + int(self.rank), + int(device.index), + int(initial_pool_bytes), + ) + dist.barrier(group=pg) + + # Dedicated high-priority stream for the ASYNC collectives (sync calls run directly on the + # caller's stream -- routing them through here costs two event hops per call, ~0.27 ms + # measured, comparable to the a2a itself). The fast_barrier epoch is one per-group monotonic + # counter, so barrier kernels must execute in submission order across streams: wait() every + # async handle before issuing the next sync collective (see all_to_all_single_4d_async). + # High priority lets the comm kernels get SM slots under concurrent compute. + _, greatest = torch.cuda.Stream.priority_range() + self._comm_stream = torch.cuda.Stream(device=device, priority=greatest) + + def _launch_on_comm_stream(self, inputs: list[torch.Tensor], fn: Callable): + """Run a collective on the group's comm stream: comm stream waits for the caller's current + stream (inputs ready -- and, since the ready-event trails everything already submitted, any + earlier consumer of the same tag's buffer), runs fn, and returns (result, done_event).""" + cur = torch.cuda.current_stream() + ev_ready = torch.cuda.Event() + ev_ready.record(cur) + self._comm_stream.wait_event(ev_ready) + with torch.cuda.stream(self._comm_stream): + out = fn() + for t in inputs: + t.record_stream(self._comm_stream) # keep the allocator from reusing x too early + ev_done = torch.cuda.Event() + ev_done.record(self._comm_stream) + return out, ev_done + + def all_to_all_single_4d( + self, + x: torch.Tensor, + *, + mode: int = 0, + tag: str = "", + use_tma: bool | None = None, + ) -> torch.Tensor: + # COLLECTIVE SEMANTICS: s/n must divide world_size (uniform). The first (shape, mode, use_tma) + # seen runs a local micro-benchmark and caches the launch config; every rank MUST issue the SAME + # (shape, mode, use_tma) call sequence (the nvshmem symmetric alloc + cross-rank barrier are + # collective; all ranks miss the same entry on the first call together). Sync AND async calls + # count in that sequence (both run on the same comm stream). + # + # use_tma (None=auto / True / False): None=auto -> sm<9 uses non-TMA; sm90+ micro-benchmarks BOTH + # paths on the first call for this shape and caches the faster (runtime path selection, replacing the + # old static table). True forces TMA (requires sm90+, else TORCH_CHECK fails); False forces non-TMA. + # Every rank MUST pass the SAME use_tma (a mismatch diverges kernel/barrier + cache key -> hang). + # + # tag scopes the symmetric-heap output buffer (reused on same tag+shape+dtype). Results that + # must stay live together (e.g. q/k/v) MUST use distinct tags, else they alias one buffer. + return torch.ops.fast_ulysses.all_to_all_single_4d( + self._group, x.contiguous(), mode, tag, use_tma + ) + + def all_to_all_single_4d_async( + self, + x: torch.Tensor, + *, + mode: int = 0, + tag: str = "", + use_tma: bool | None = None, + ) -> AsyncA2AHandle: + # Async variant: launches on the group's comm stream and returns immediately; kernels submitted + # to the caller's stream afterwards overlap with the a2a until handle.wait(). Collective + # constraints are identical to the sync call (same rank-uniform call sequence, sync and async + # counted together). + # + # ORDERING CONSTRAINT when mixing with sync calls: the fast_barrier epoch is one per-group + # monotonic counter, so barrier kernels must EXECUTE in submission order. wait() every async + # handle of this group before issuing the next sync collective on the main stream -- that data + # dependency forces the comm-stream barriers to complete first. (Sync calls run directly on the + # caller's stream: routing them through the comm stream costs two event hops per call, measured + # ~0.27 ms/call on H200 -- comparable to the a2a itself.) + x = x.contiguous() + out, ev_done = self._launch_on_comm_stream( + [x], + lambda: torch.ops.fast_ulysses.all_to_all_single_4d( + self._group, x, mode, tag, use_tma + ), + ) + return AsyncA2AHandle(out, ev_done) + + def destroy(self) -> None: + dist.barrier(group=self.pg) + self._group.destroy() diff --git a/lightx2v_fast_ulysses/csrc/a2a_config.cuh b/lightx2v_fast_ulysses/csrc/a2a_config.cuh new file mode 100644 index 000000000..80486c50a --- /dev/null +++ b/lightx2v_fast_ulysses/csrc/a2a_config.cuh @@ -0,0 +1,46 @@ +// Copied and adapted from https://github.com/triple-mu/fast-ulysses +#pragma once +#include "ulysses_common.cuh" +#include +#include + +namespace ulysses { + +// A2A launch config: TMA path uses tile_n/tile_s (stages/bdiv are fixed at 4 in all_to_all_tma.cu); +// non-TMA path uses threads/unroll/blocks. +struct A2AConfig { + int tile_n = 0, tile_s = 1; // TMA path + int threads = 512, unroll = 4, blocks = 0; // non-TMA path +}; + +// Defaults per mode (near-optimal for DiT): +// mode0 -> tile_n=max(1,n_local-1), tile_s=1 (non-divisor small tile_n yields more n-tiles, higher TMA +// concurrency; safe in mode0 -- the dst n-dim is n_local, so the trailing tile clips). mode1 -> tile_n=n_local, +// tile_s=2 (mode1 needs tile_n | n_local, so a whole-block tile is the safe default). +inline A2AConfig default_config(int mode, int n_local) +{ + A2AConfig c; + if (mode == 0) { + c.tile_n = std::max(1, n_local - 1); + c.tile_s = 1; + } + else { + c.tile_n = std::max(1, n_local); + c.tile_s = 2; + } + c.threads = 512; + c.unroll = 4; + c.blocks = 0; + return c; +} + +// Autotune/cache key (excludes b and elem -- 2B path only): ws, mode, tma(0/1), +// n_local, s_local, d. The tma bit separates TMA / non-TMA configs. +using ConfigKey = std::tuple; + +inline ConfigKey config_key(int ws, int mode, bool tma, const Ulysses4DDims& dims) +{ + return ConfigKey{ws, mode, tma ? 1 : 0, dims.n_local, dims.s_local, dims.d}; +} + +} // namespace ulysses diff --git a/lightx2v_fast_ulysses/csrc/all_to_all.cu b/lightx2v_fast_ulysses/csrc/all_to_all.cu new file mode 100644 index 000000000..e3a23b34a --- /dev/null +++ b/lightx2v_fast_ulysses/csrc/all_to_all.cu @@ -0,0 +1,301 @@ +// Copied and adapted from https://github.com/triple-mu/fast-ulysses +#include "a2a_config.cuh" +#include "ulysses_common.cuh" +#include +#include +#include +#include +#include + +namespace ulysses { + +// Each (peer, b, s) unit copies one contiguous n_local*d block (contiguous in both src and dst under +// mode0/mode1): inner (nl,v) is contiguous, so consecutive threads fill a whole block, giving coalesced +// large (typically 4KB) remote bursts. Far better NVLink efficiency than the old scattered writes +// ("write a 256B d-row then jump 4KB"). UNROLL register prefetch pipelines read/write (local reads hidden behind remote +// writes). +template +__global__ void a2a_copy_generic( + const uint8_t* __restrict__ src, PeerPtrs peers, Ulysses4DDims dims, int elem_size, Epilogue) +{ + const int row_bytes = dims.d * elem_size; // 16B aligned (guaranteed by Global Constraints) + const int vecs = row_bytes >> 4; // uint4 count per d-row + const int blk_vecs = dims.n_local * vecs; // uint4 count per contiguous n_local*d block + const int64_t units = static_cast(WORLD_SIZE) * dims.b * dims.s_local; + const int64_t total = units * blk_vecs; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + const int64_t tid = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + + for (int64_t base = tid; base < total; base += stride * UNROLL) { + const uint4* sp[UNROLL]; + uint4* dp[UNROLL]; + uint4 reg[UNROLL]; + // Read phase (local HBM): prefetch UNROLL uint4s into registers +#pragma unroll + for (int k = 0; k < UNROLL; ++k) { + int64_t idx = base + static_cast(k) * stride; + sp[k] = nullptr; + if (idx >= total) + continue; + int inner = static_cast(idx % blk_vecs); // uint4 offset within block (over nl and d) + int64_t u = idx / blk_vecs; // unit index over [WS, b, s_local] + int s = static_cast(u % dims.s_local); + u /= dims.s_local; + int b_idx = static_cast(u % dims.b); + u /= dims.b; + int peer = static_cast(u); // 0..WS-1 + + int64_t src_base_row, dst_base_row; // block base address in units of d-rows + if (MODE == 0) { + src_base_row = (static_cast(b_idx) * dims.s_local + s) * dims.n_global + peer * dims.n_local; + dst_base_row = + (static_cast(b_idx) * dims.s_global + (dims.rank * dims.s_local + s)) * dims.n_local; + } + else { + src_base_row = (static_cast(b_idx) * dims.s_global + (peer * dims.s_local + s)) * dims.n_local; + dst_base_row = + (static_cast(b_idx) * dims.s_local + s) * dims.n_global + dims.rank * dims.n_local; + } + sp[k] = reinterpret_cast(src + src_base_row * row_bytes) + inner; + dp[k] = reinterpret_cast(static_cast(peers.p[peer]) + dst_base_row * row_bytes) + inner; + reg[k] = *sp[k]; + } + // Write phase (remote NVLink): issued in bulk, local read latency hidden behind it +#pragma unroll + for (int k = 0; k < UNROLL; ++k) + if (sp[k]) + *dp[k] = reg[k]; + } + __threadfence_system(); // system-scope visibility of P2P direct writes to other GPUs +} + +template +static void launch_ws_u(const PeerPtrs& pp, + const uint8_t* src, + const Ulysses4DDims& dims, + int mode, + int elem_size, + int blocks, + int threads, + cudaStream_t stream) +{ + if (mode == 0) + a2a_copy_generic + <<>>(src, pp, dims, elem_size, EpilogueIdentity{}); + else + a2a_copy_generic + <<>>(src, pp, dims, elem_size, EpilogueIdentity{}); +} + +template +static void launch_ws(const uint8_t* src, + const std::vector& peer_ptrs, + const Ulysses4DDims& dims, + int mode, + int elem_size, + int blocks, + int threads, + int unroll, + cudaStream_t stream) +{ + PeerPtrs pp; + for (int i = 0; i < WS; ++i) + pp.p[i] = reinterpret_cast(peer_ptrs[i]); + // unroll candidates are {4, 8} (resolve_config_nontma); only these two are instantiated. + if (unroll == 8) + launch_ws_u(pp, src, dims, mode, elem_size, blocks, threads, stream); + else + launch_ws_u(pp, src, dims, mode, elem_size, blocks, threads, stream); +} + +// Dispatch by ws to the matching launch_ws (folds the duplicated switch(ws) in launch_a2a / +// resolve_config_nontma). Caller already TORCH_CHECKs ws in [1, 8], so default is a no-op. +static void nontma_dispatch(int ws, + const uint8_t* s, + const std::vector& peers, + const Ulysses4DDims& dims, + int mode, + int elem, + int blocks, + int threads, + int unroll, + cudaStream_t stream) +{ + switch (ws) { + case 1: + launch_ws<1>(s, peers, dims, mode, elem, blocks, threads, unroll, stream); + break; + case 2: + launch_ws<2>(s, peers, dims, mode, elem, blocks, threads, unroll, stream); + break; + case 3: + launch_ws<3>(s, peers, dims, mode, elem, blocks, threads, unroll, stream); + break; + case 4: + launch_ws<4>(s, peers, dims, mode, elem, blocks, threads, unroll, stream); + break; + case 5: + launch_ws<5>(s, peers, dims, mode, elem, blocks, threads, unroll, stream); + break; + case 6: + launch_ws<6>(s, peers, dims, mode, elem, blocks, threads, unroll, stream); + break; + case 7: + launch_ws<7>(s, peers, dims, mode, elem, blocks, threads, unroll, stream); + break; + case 8: + launch_ws<8>(s, peers, dims, mode, elem, blocks, threads, unroll, stream); + break; + default: + break; + } +} + +// SM count of the current device: queried once per process and cached (one device bound per process). +// Declared in ulysses_common.cuh for the A2A config resolver. +int sm_count_cached() +{ + static const int sm = [] { + int d = 0, s = 0; + ULYSSES_CUDA_CHECK(cudaGetDevice(&d)); + ULYSSES_CUDA_CHECK(cudaDeviceGetAttribute(&s, cudaDevAttrMultiProcessorCount, d)); + return s > 0 ? s : 132; // 132 = H100 SM count, defensive fallback for an unexpected 0 + }(); + return sm; +} + +static int clamp_blocks(int64_t needed, int64_t want) +{ + return static_cast(std::max(1, std::min(needed, want))); +} + +// Launch the non-TMA direct-write kernel with the given config (blocks/threads/unroll all set by +// resolve_config_nontma). threads/unroll set the number of in-flight remote write transactions (Little's +// law); 512 threads measured to raise NVLink BW from ~210-260 (256 threads) to a steady ~310 across all N +// (close to TMA), hence the default 512. +void launch_a2a(const void* src, + const std::vector& peer_ptrs, + const Ulysses4DDims& dims, + int mode, + int elem_size, + const A2AConfig& cfg, + cudaStream_t stream) +{ + const int ws = static_cast(peer_ptrs.size()); + nontma_dispatch(ws, + static_cast(src), + peer_ptrs, + dims, + mode, + elem_size, + cfg.blocks, + cfg.threads, + cfg.unroll, + stream); +} + +// Local microbench shared by both resolve_config paths: warmup then time 10 iters, return us/call. +// No collective primitive -- under SPMD all ranks miss the same (shape,mode,tma) on the first call +// and run this concurrently, so the remote writes contend just as in steady state. Correctness is +// guaranteed by the post-launch fast_barrier, not by any timing-side lockstep. +float microbench_us(const std::function& run_once, cudaStream_t stream) +{ + cudaEvent_t s, e; + ULYSSES_CUDA_CHECK(cudaEventCreate(&s)); + ULYSSES_CUDA_CHECK(cudaEventCreate(&e)); + for (int i = 0; i < 3; ++i) + run_once(); // warm up + ULYSSES_CUDA_CHECK(cudaEventRecord(s, stream)); + for (int i = 0; i < 10; ++i) + run_once(); + ULYSSES_CUDA_CHECK(cudaEventRecord(e, stream)); + ULYSSES_CUDA_CHECK(cudaEventSynchronize(e)); + float ms = 0.f; + ULYSSES_CUDA_CHECK(cudaEventElapsedTime(&ms, s, e)); + ULYSSES_CUDA_CHECK(cudaEventDestroy(s)); + ULYSSES_CUDA_CHECK(cudaEventDestroy(e)); + return ms * 100.f; // ms(10 iters) -> us/call: /10 iters * 1000 (ms->us) +} + +// non-TMA config resolution: autotune over the three launch knobs -- threads (in-flight remote writes, +// Little's law), unroll (register-prefetch pipeline depth), and grid size (blocks). Sweeps a converged +// threads x unroll x grid grid by micro-benchmark and keeps the fastest. No cache of its own: the returned +// cfg is held by UlyssesGroup::cfg_cache_ (so the sweep runs once per shape). +A2AConfig resolve_config_nontma(const void* src, + const std::vector& peer_ptrs, + const Ulysses4DDims& dims, + int mode, + int elem_size, + cudaStream_t stream, + bool verbose, + const std::function& finish) +{ + const int ws = static_cast(peer_ptrs.size()); + const int64_t row_bytes = static_cast(dims.d) * elem_size; + const int vecs = static_cast(row_bytes >> 4); + const int64_t total = static_cast(ws) * dims.b * dims.n_local * dims.s_local * vecs; + const uint8_t* s = static_cast(src); + const int sm = sm_count_cached(); + + // Converged sweep (30 candidates): threads {256,512,1024} x unroll {4,8} x grid factor {8,12,16,24,32}. + // unroll matches launch_ws_u's compile-time instantiations; the factor spread covers the measured DiT + // optima (n_local=5 wins at factor 12; the table's non-TMA cells span factors 8..32). + const int threads_cand[] = {256, 512, 1024}; + const int unroll_cand[] = {4, 8}; + const double factors[] = {8.0, 12.0, 16.0, 24.0, 32.0}; + + A2AConfig best; + best.threads = 512; + best.unroll = 4; + best.blocks = clamp_blocks((total + 511) / 512, static_cast(sm) * 16); + float best_us = 1e30f; + for (int threads : threads_cand) { + const int64_t needed = (total + threads - 1) / threads; + for (int unroll : unroll_cand) { + for (double f : factors) { + const int blocks = clamp_blocks(needed, static_cast(sm * f)); + auto launch = [&] { + nontma_dispatch(ws, s, peer_ptrs, dims, mode, elem_size, blocks, threads, unroll, stream); + }; + // Probe (bare launch, not timed): large threads x large unroll can exceed the per-block + // register budget. A failed launch (cudaErrorLaunchOutOfResources) returns synchronously and + // times near-zero, so it would be falsely picked as fastest; skip it (cudaGetLastError + // clears the error too). + launch(); + if (cudaGetLastError() != cudaSuccess) + continue; + // Timed run is the real per-call op (launch + finish=quiet+barrier) so the ranking matches + // steady state. + const float us = microbench_us( + [&] { + launch(); + finish(); + }, + stream); + if (us < best_us) { + best_us = us; + best.threads = threads; + best.unroll = unroll; + best.blocks = blocks; + } + } + } + } + // tile_n/tile_s are TMA-only fields; carry default_config's values so the struct stays consistent. + const A2AConfig dc = default_config(mode, dims.n_local); + best.tile_n = dc.tile_n; + best.tile_s = dc.tile_s; + if (verbose && dims.rank == 0) { + const double per_iter_us = best_us; // microbench_us already returns us/call + const double remote_bytes = static_cast(total) * 16.0 * (ws - 1) / ws; + const double gbps = remote_bytes / (per_iter_us * 1e3); + std::cout << "[ulysses non-TMA tune] ws=" << ws << " mode=" << mode << " n_local=" << dims.n_local + << " s_local=" << dims.s_local << " d=" << dims.d << " -> threads=" << best.threads + << " unroll=" << best.unroll << " blocks=" << best.blocks << " | " << std::fixed + << std::setprecision(1) << per_iter_us << " us/iter " << std::setprecision(0) << gbps << " GB/s" + << std::endl; + } + return best; +} + +} // namespace ulysses diff --git a/lightx2v_fast_ulysses/csrc/all_to_all_tma.cu b/lightx2v_fast_ulysses/csrc/all_to_all_tma.cu new file mode 100644 index 000000000..d3e1e25b9 --- /dev/null +++ b/lightx2v_fast_ulysses/csrc/all_to_all_tma.cu @@ -0,0 +1,305 @@ +// Copied and adapted from https://github.com/triple-mu/fast-ulysses +// TMA (cp.async.bulk.tensor) A2A: few SMs issue, TMA engine moves data (src gmem->smem->peer gmem). +// One launch covers all peers with many blocks (one tile per block) -> high concurrency saturates NVLink. mode0/mode1 +// unified. Mechanism/coordinates verified by tma_p2p_probe / tma_a2a_test. Raw PTX (mbar/TMA/wait_group/fence) +// extracted into named device functions in tma_ptx.cuh, behavior-equivalent. +#include "a2a_config.cuh" +#include "tma_ptx.cuh" +#include "ulysses_common.cuh" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ulysses { + +struct TmaMaps { + CUtensorMap m[8]; // per-peer dst tensormap (world_size <= 8) +}; + +// Global tile index -> (peer, n, s, bi). +__device__ __forceinline__ void +tma_decode(int g, int per_peer, int n_ntiles, int n_stiles, int tile_n, int tile_s, int& peer, int& n, int& s, int& bi) +{ + peer = g / per_peer; + int t = g % per_peer; + int nt = t % n_ntiles; + int tmp = t / n_ntiles; + int st = tmp % n_stiles; + bi = tmp / n_stiles; + n = nt * tile_n; + s = st * tile_s; +} + +// Software-pipelined: single thread issues, B smem buffers rotate, prefetch-1-ahead + wait_group(B-1) +// keeps B-1 TMA stores in flight -> remote NVLink write pipeline stays full (removes the single-stage serial bubble). +// grid covers the full (peer, b, s-tile, n-tile) set; each block grid-strides over its own run of tiles. +template +__global__ void a2a_tma_kernel(const __grid_constant__ CUtensorMap src_map, + const __grid_constant__ TmaMaps dst, + int ws, + int mode, + int rank, + int s_local, + int n_local, + int b, + int tile_s, + int tile_n, + uint32_t tile_bytes) +{ + if (threadIdx.x != 0) + return; + + extern __shared__ uint8_t smem_raw[]; + const uint32_t tb_al = (tile_bytes + 127u) & ~127u; // 128B-align each buffer + uintptr_t base = (reinterpret_cast(smem_raw) + 127) & ~static_cast(127); + uint64_t* mbar = reinterpret_cast(base + (uintptr_t)STAGES * tb_al); + + uint32_t buf_a[STAGES], mbar_a[STAGES]; + int parity[STAGES]; + for (int k = 0; k < STAGES; ++k) { + buf_a[k] = (uint32_t)__cvta_generic_to_shared(reinterpret_cast(base + (uintptr_t)k * tb_al)); + mbar_a[k] = (uint32_t)__cvta_generic_to_shared(mbar + k); + parity[k] = 0; + mbar_init(mbar_a[k]); + } + + const int n_ntiles = (n_local + tile_n - 1) / tile_n; + const int n_stiles = (s_local + tile_s - 1) / tile_s; + const int per_peer = b * n_stiles * n_ntiles; + const int total = ws * per_peer; + + // mode0: src head-dim offset peer*n_local / dst seq-dim offset rank*s_local; mode1 swaps them. + const int src_n_pp = (mode == 0) ? n_local : 0; // multiplied by peer + const int src_s_pp = (mode == 0) ? 0 : s_local; + const int dst_n_off = (mode == 0) ? 0 : rank * n_local; + const int dst_s_off = (mode == 0) ? rank * s_local : 0; + + // Tiles owned by this block (grid-stride run): g = blockIdx.x + j*gridDim.x, M total + const int M = (total - (int)blockIdx.x + (int)gridDim.x - 1) / (int)gridDim.x; + if (M <= 0) + return; + + // prologue: issue loads for the first min(STAGES, M) tiles + for (int k = 0; k < STAGES && k < M; ++k) { + int g = blockIdx.x + k * gridDim.x; + int peer, n, s, bi; + tma_decode(g, per_peer, n_ntiles, n_stiles, tile_n, tile_s, peer, n, s, bi); + mbar_arrive_expect(mbar_a[k], tile_bytes); + tma_load_4d(buf_a[k], &src_map, 0, n + peer * src_n_pp, s + peer * src_s_pp, bi, mbar_a[k]); + } + + for (int j = 0; j < M; ++j) { + int cur = j % STAGES; + int g = blockIdx.x + j * gridDim.x; + int peer, n, s, bi; + tma_decode(g, per_peer, n_ntiles, n_stiles, tile_n, tile_s, peer, n, s, bi); + + // wait for this tile's load to complete + mbar_wait(mbar_a[cur], parity[cur]); + parity[cur] ^= 1; + async_proxy_fence(); + + // store this tile -> peer dst + tma_store_4d(&dst.m[peer], 0, n + dst_n_off, s + dst_s_off, bi, buf_a[cur]); + tma_commit_group(); + + // prefetch next tile: first drain to <= STAGES-1 stores in flight so the target buffer was read by an old store + int nl = j + STAGES; + if (nl < M) { + tma_wait_group(); // keep <=STAGES-1 stores in flight + int slot = nl % STAGES; + int g2 = blockIdx.x + nl * gridDim.x; + int peer2, n2, s2, bi2; + tma_decode(g2, per_peer, n_ntiles, n_stiles, tile_n, tile_s, peer2, n2, s2, bi2); + mbar_arrive_expect(mbar_a[slot], tile_bytes); + tma_load_4d(buf_a[slot], &src_map, 0, n2 + peer2 * src_n_pp, s2 + peer2 * src_s_pp, bi2, mbar_a[slot]); + } + } + tma_wait_group<0>(); // drain all stores before exit (correctness) +} + +// 4D tensormap: dims(innermost first)=[d, ndim, sdim, b]; box=[d, tile_n, tile_s, 1]. +static CUtensorMap tma_make_map(void* base, int d, int ndim, int sdim, int b, int tile_n, int tile_s, int elem_size) +{ + CUtensorMap m; + uint64_t es = (uint64_t)elem_size; + uint64_t gdims[4] = {(uint64_t)d, (uint64_t)ndim, (uint64_t)sdim, (uint64_t)b}; + uint64_t gstr[3] = {(uint64_t)d * es, (uint64_t)ndim * d * es, (uint64_t)sdim * ndim * d * es}; + uint32_t box[4] = {(uint32_t)d, (uint32_t)tile_n, (uint32_t)tile_s, 1u}; + uint32_t estr[4] = {1, 1, 1, 1}; + CUtensorMapDataType dt = (elem_size == 2) ? CU_TENSOR_MAP_DATA_TYPE_UINT16 : CU_TENSOR_MAP_DATA_TYPE_UINT8; + CUresult r = cuTensorMapEncodeTiled(&m, + dt, + 4, + base, + gdims, + gstr, + box, + estr, + CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TORCH_CHECK(r == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed: ", (int)r); + return m; +} + +// Build src + per-peer dst tensormaps for the uniform path. mode0/mode1 src/dst layouts are dual: +// mode0 src=[b,s_local,n_global,d], dst=[b,s_global,n_local,d]; mode1 swaps n/s in both. +static void build_tma_maps(int mode, + const Ulysses4DDims& dims, + const void* src, + const std::vector& peers, + int elem, + int tile_n, + int tile_s, + CUtensorMap& src_map, + TmaMaps& dst) +{ + const int ws = (int)peers.size(); + const int d = dims.d, b = dims.b; + const int s_local = dims.s_local, n_local = dims.n_local; + const int s_global = dims.s_global, n_global = dims.n_global; + if (mode == 0) { + src_map = tma_make_map((void*)src, d, n_global, s_local, b, tile_n, tile_s, elem); + for (int p = 0; p < ws; ++p) + dst.m[p] = tma_make_map((void*)peers[p], d, n_local, s_global, b, tile_n, tile_s, elem); + } + else { + src_map = tma_make_map((void*)src, d, n_local, s_global, b, tile_n, tile_s, elem); + for (int p = 0; p < ws; ++p) + dst.m[p] = tma_make_map((void*)peers[p], d, n_global, s_local, b, tile_n, tile_s, elem); + } +} + +// Launch the TMA kernel for a given config (build maps + launch only). mode0/mode1 use different src/dst tensormap +// layouts. +void launch_a2a_tma(const void* src, + const std::vector& peer_ptrs, + const Ulysses4DDims& dims, + int mode, + int elem_size, + const A2AConfig& cfg, + cudaStream_t stream) +{ + const int ws = (int)peer_ptrs.size(); + const int d = dims.d, b = dims.b, rank = dims.rank; + const int s_local = dims.s_local, n_local = dims.n_local; + const int tile_n = std::min({cfg.tile_n, n_local, 256}); + const int tile_s = std::min({cfg.tile_s, s_local, 256}); + constexpr int stages = 4; // fixed (measured DiT optimum; keeps tma_wait_group at {3,0}) + const uint32_t tile_bytes = (uint32_t)tile_s * tile_n * d * elem_size; + const uint32_t tb_al = (tile_bytes + 127u) & ~127u; + const int smem = (int)(tb_al * (uint32_t)stages) + 128 + 8 * stages; + + static bool attr_set = false; + if (!attr_set) { + ULYSSES_CUDA_CHECK( + cudaFuncSetAttribute(a2a_tma_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 200 * 1024)); + attr_set = true; + } + + const int n_stiles = (s_local + tile_s - 1) / tile_s; + const int n_ntiles = (n_local + tile_n - 1) / tile_n; + const int total = ws * b * n_stiles * n_ntiles; + constexpr int bdiv = 4; // fixed (measured DiT optimum): more tiles per block -> enables pipelining + const int blocks = std::max(std::min(total, 65535) / bdiv, 1); + + CUtensorMap src_map; + TmaMaps dst; + build_tma_maps(mode, dims, src, peer_ptrs, elem_size, tile_n, tile_s, src_map, dst); + a2a_tma_kernel + <<>>(src_map, dst, ws, mode, rank, s_local, n_local, b, tile_s, tile_n, tile_bytes); +} + +// Candidate configs: cover small N (small tile_n, shallow pipeline) to large N (large tile_n, deep pipeline). +// In mode0 a non-divisor tile_n (e.g. n_local-1) often beats a whole block -- it forms multiple n-tiles for +// higher TMA concurrency (a single whole-block tile is actually slower). Seed includes default_config; rest +// are empirically tuned. +// +// CORRECTNESS GUARD: the dst dim offset by this rank (mode0: s at rank*s_local; mode1: n at rank*n_local) is +// tiled over the rank's chunk, but a TMA store only clips at the dst's GLOBAL dim, not at this rank's +// sub-segment. So if that tile does not divide the chunk, the trailing tile overruns into the neighbor rank's +// region (corruption; and during autotune the cross-rank overrun writes race). Require it to divide: mode0 +// needs s_local % tile_s == 0 (tile_s is 1 here, always ok), mode1 needs n_local % tile_n == 0. The other dim +// spans the rank's full chunk == the dst global dim, so its trailing tile is clipped correctly. +static std::vector tma_candidates(int mode, int n_local, int s_local) +{ + std::vector v; + auto add = [&](int tn, int ts) { + tn = std::max(1, std::min(tn, n_local)); + ts = std::max(1, std::min(ts, s_local)); + if (mode == 0 ? (s_local % ts != 0) : (n_local % tn != 0)) + return; // would overrun the neighbor rank's region (see CORRECTNESS GUARD above) + for (auto& c : v) + if (c.tile_n == tn && c.tile_s == ts) + return; + A2AConfig c{}; + c.tile_n = tn; + c.tile_s = ts; + v.push_back(c); + }; + // Converged set (stages/bdiv fixed at 4): vary tile_n {default, n_local-1, n_local/2, n_local, 8} and, + // for mode1, also tile_s {1,2}. The overrun guard in add() drops non-dividing tile_n for mode1. + const A2AConfig def = default_config(mode, n_local); // seed: default config enters candidates first + add(def.tile_n, def.tile_s); + const int nl1 = std::max(1, n_local - 1), nlh = std::max(1, n_local / 2); + const int ts_extra = (mode == 0) ? 1 : 2; // mode1 also probes tile_s=2 (whole-s clips, always safe) + add(nl1, 1); + add(nlh, 1); + add(n_local, 1); + add(8, 1); + add(nl1, ts_extra); + add(nlh, ts_extra); + return v; +} + +// autotune: microbench all candidates (local timing via the shared microbench_us), return the fastest +// config (remote writes are real but overwritten by the subsequent final launch, so correctness is +// unaffected). No collective primitive of its own -- under SPMD all ranks miss the same (shape,mode) +// on the first call and run this concurrently (contention captured); the cache lives in the caller +// (UlyssesGroup::cfg_cache_). verbose: rank-0 debug print. +A2AConfig resolve_config_tma(const void* src, + const std::vector& peer_ptrs, + const Ulysses4DDims& dims, + int mode, + int elem_size, + bool verbose, + cudaStream_t stream, + const std::function& finish) +{ + const int n_local = dims.n_local, s_local = dims.s_local, ws = (int)peer_ptrs.size(); + const auto cands = tma_candidates(mode, n_local, s_local); + + A2AConfig best = cands[0]; + float best_t = 1e30f; + for (const auto& c : cands) { + // Timed run is the real per-call op (launch + finish=quiet+barrier) so the ranking matches steady + // state (a launch-only microbench mis-ranks tile_n: a smaller tile_n that loses on raw launch time + // wins once the cross-rank barrier serializes per-iteration contention). + float us = microbench_us( + [&] { + launch_a2a_tma(src, peer_ptrs, dims, mode, elem_size, c, stream); + finish(); + }, + stream); + if (verbose) + std::cerr << "[tma-at] ws=" << ws << " mode=" << mode << " nl=" << n_local << " sl=" << s_local + << " | tn=" << c.tile_n << " ts=" << c.tile_s << " -> " << std::fixed << std::setprecision(1) + << us << " us/call" << std::endl; + if (us < best_t) { + best_t = us; + best = c; + } + } + return best; +} + +} // namespace ulysses diff --git a/lightx2v_fast_ulysses/csrc/bindings.cpp b/lightx2v_fast_ulysses/csrc/bindings.cpp new file mode 100644 index 000000000..ad91dd1fc --- /dev/null +++ b/lightx2v_fast_ulysses/csrc/bindings.cpp @@ -0,0 +1,103 @@ +// Copied and adapted from https://github.com/triple-mu/fast-ulysses +#include +#include +#include +#include +#include +#include + +#include "a2a_config.cuh" +#include "ulysses_common.cuh" +#include "ulysses_group.cuh" + +namespace ulysses { + +int64_t nvshmem_uniqueid_nbytes() +{ + return static_cast(sizeof(nvshmemx_uniqueid_t)); +} + +at::Tensor all_to_all_single_4d(const c10::intrusive_ptr& group, + at::Tensor input, + int64_t mode, + std::string tag, + c10::optional use_tma) +{ + TORCH_CHECK(input.is_cuda() && input.dim() == 4, "input must be a 4D CUDA tensor"); + TORCH_CHECK(input.scalar_type() == at::kHalf || input.scalar_type() == at::kBFloat16, + "dtype must be float16 or bfloat16"); + input = input.contiguous(); + const int ws = static_cast(group->world_size()); + TORCH_CHECK(ws >= 1 && ws <= 8, "world_size must be in [1, 8] (single-node NVLink), got ", ws); + const int me = static_cast(group->rank()); + const int b = static_cast(input.size(0)); + const int x1 = static_cast(input.size(1)); + const int x2 = static_cast(input.size(2)); + const int d = static_cast(input.size(3)); + const int elem = static_cast(input.element_size()); + TORCH_CHECK((static_cast(d) * elem) % 16 == 0, "d*elem_size must be 16B-aligned"); + TORCH_CHECK(mode == 0 || mode == 1, "mode must be 0 or 1"); + + const at::cuda::CUDAGuard guard(input.device()); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + Ulysses4DDims dims; + dims.b = b; + dims.d = d; + dims.rank = me; + std::vector out_shape; + if (mode == 0) { + TORCH_CHECK(x2 % ws == 0, "n_global must be divisible by world_size"); + dims.s_local = x1; + dims.n_global = x2; + dims.s_global = x1 * ws; + dims.n_local = x2 / ws; + out_shape = {b, dims.s_global, dims.n_local, d}; + } + else { + TORCH_CHECK(x1 % ws == 0, "s_global must be divisible by world_size"); + dims.s_global = x1; + dims.n_local = x2; + dims.s_local = x1 / ws; + dims.n_global = x2 * ws; + out_shape = {b, dims.s_local, dims.n_global, d}; + } + + const auto& buf = group->pool().acquire(out_shape, input.scalar_type(), tag); + const int use_tma_i = use_tma.has_value() ? (*use_tma ? 1 : 0) : -1; + if (use_tma_i > 0) + TORCH_CHECK(group->sm_major() >= 9, "use_tma=True requires sm90+ (TMA unavailable on this GPU)"); + const auto pc = + group->resolve_config(dims, static_cast(mode), use_tma_i, input.data_ptr(), buf.peer_ptrs, elem, stream); + if (pc.tma) + launch_a2a_tma(input.data_ptr(), buf.peer_ptrs, dims, static_cast(mode), elem, pc.cfg, stream); + else + launch_a2a(input.data_ptr(), buf.peer_ptrs, dims, static_cast(mode), elem, pc.cfg, stream); + ULYSSES_CUDA_CHECK(cudaGetLastError()); + nvshmemx_quiet_on_stream(stream); + group->fast_barrier(stream); + return buf.view; +} + +} // namespace ulysses + +TORCH_LIBRARY(fast_ulysses, m) +{ + m.def("nvshmem_uniqueid_nbytes() -> int"); + m.impl("nvshmem_uniqueid_nbytes", &ulysses::nvshmem_uniqueid_nbytes); + + m.class_("UlyssesGroup") + .def(torch::init, int64_t, int64_t, int64_t>()) + .def("rank", &ulysses::UlyssesGroup::rank) + .def("world_size", &ulysses::UlyssesGroup::world_size) + .def("destroy", &ulysses::UlyssesGroup::destroy) + .def_static("uniqueid_nints", &ulysses::UlyssesGroup::uniqueid_nints) + .def_static("get_uniqueid", &ulysses::UlyssesGroup::get_uniqueid) + .def_static("init_world", &ulysses::UlyssesGroup::init_world); + + m.def("all_to_all_single_4d(__torch__.torch.classes.fast_ulysses.UlyssesGroup group, " + "Tensor input, int mode, str tag, bool? use_tma=None) -> Tensor"); + m.impl("all_to_all_single_4d", c10::DispatchKey::CompositeExplicitAutograd, &ulysses::all_to_all_single_4d); +} + +PYBIND11_MODULE(_C, m) {} diff --git a/lightx2v_fast_ulysses/csrc/symmetric_pool.cu b/lightx2v_fast_ulysses/csrc/symmetric_pool.cu new file mode 100644 index 000000000..129b86773 --- /dev/null +++ b/lightx2v_fast_ulysses/csrc/symmetric_pool.cu @@ -0,0 +1,82 @@ +// Copied and adapted from https://github.com/triple-mu/fast-ulysses +#include "symmetric_pool.cuh" +#include +#include +#include +#include + +namespace ulysses { + +SymmetricHeapPool::SymmetricHeapPool(int64_t reserved_bytes, int world_size, std::vector peer_global_pes): + reserved_(reserved_bytes), + world_size_(world_size), + peer_global_pes_(std::move(peer_global_pes)) +{ +} + +const SymmetricHeapPool::Buffer& +SymmetricHeapPool::acquire(const std::vector& shape, c10::ScalarType dtype, const std::string& tag) +{ + TORCH_CHECK(!destroyed_, "SymmetricHeapPool::acquire called after destroy()"); + Key key{tag, shape, dtype}; + auto it = registry_.find(key); + if (it != registry_.end()) + return it->second; // reuse + + int64_t numel = 1; + for (auto s : shape) + numel *= s; + const int64_t elem = c10::elementSize(dtype); + int64_t nbytes = numel * elem; + nbytes = (nbytes + 15) / 16 * 16; // uint4 alignment + + // Uniform: every rank computes an identical nbytes (out_shape is rank-independent), so the + // collective (uniform-size) nvshmem_align below needs no max-reduce. + const int64_t alloc_bytes = nbytes; + TORCH_CHECK(used_ + alloc_bytes <= reserved_, + "SymmetricHeapPool OOM: need ", + alloc_bytes, + " B, used ", + used_, + " / reserved ", + reserved_, + " B. Increase initial_pool_bytes."); + + void* p = nvshmem_align(256, alloc_bytes); // collective alloc (uniform size); address never moves + TORCH_CHECK(p != nullptr, "nvshmem_align failed for ", alloc_bytes, " B"); + segments_.push_back(p); + used_ += alloc_bytes; + + Buffer buf; + buf.sym_base = p; + buf.nbytes = alloc_bytes; + buf.peer_ptrs.resize(world_size_); + for (int i = 0; i < world_size_; ++i) + buf.peer_ptrs[i] = reinterpret_cast(nvshmem_ptr(p, peer_global_pes_[i])); + // Single-node P2P: all peer pointers must be non-null. + for (int i = 0; i < world_size_; ++i) + TORCH_CHECK(buf.peer_ptrs[i] != 0, + "nvshmem_ptr returned NULL for peer ", + i, + " (non-P2P-reachable; phase-1 requires single-node NVLink)."); + + auto opts = at::TensorOptions().dtype(dtype).device(at::kCUDA, at::cuda::current_device()); + buf.view = at::from_blob( + p, shape, [](void*) {}, opts); // no-op deleter + + auto res = registry_.emplace(std::move(key), std::move(buf)); + return res.first->second; +} + +void SymmetricHeapPool::destroy() +{ + if (destroyed_) + return; + registry_.clear(); // drop from_blob views (does not free underlying memory) + for (void* p : segments_) + nvshmem_free(p); + segments_.clear(); + destroyed_ = true; +} + +} // namespace ulysses diff --git a/lightx2v_fast_ulysses/csrc/symmetric_pool.cuh b/lightx2v_fast_ulysses/csrc/symmetric_pool.cuh new file mode 100644 index 000000000..f8ba0541b --- /dev/null +++ b/lightx2v_fast_ulysses/csrc/symmetric_pool.cuh @@ -0,0 +1,42 @@ +// Copied and adapted from https://github.com/triple-mu/fast-ulysses +#pragma once +#include +#include +#include +#include +#include +#include +#include + +namespace ulysses { + +class SymmetricHeapPool { +public: + // reserved_bytes: per-group cap (must be <= NVSHMEM_SYMMETRIC_SIZE reserved at init). + SymmetricHeapPool(int64_t reserved_bytes, int world_size, std::vector peer_global_pes); + + struct Buffer { + void* sym_base; + int64_t nbytes; + std::vector peer_ptrs; // nvshmem_ptr(sym_base, peer_global_pe) + at::Tensor view; // from_blob with no-op deleter (pool owns lifetime) + }; + + // Reuse on (tag,shape,dtype) hit; otherwise collectively allocate a new segment and register it. + const Buffer& acquire(const std::vector& shape, c10::ScalarType dtype, const std::string& tag); + + // Terminal collective op: before calling, release all from_blob views returned by acquire() and + // ensure no A2A/collective is in flight, since this nvshmem_free's the segments those views alias. + void destroy(); // nvshmem_free all segments + clear registry + +private: + using Key = std::tuple, c10::ScalarType>; + int64_t reserved_, used_ = 0; + int world_size_; + std::vector peer_global_pes_; + std::vector segments_; + std::map registry_; + bool destroyed_ = false; +}; + +} // namespace ulysses diff --git a/lightx2v_fast_ulysses/csrc/tma_ptx.cuh b/lightx2v_fast_ulysses/csrc/tma_ptx.cuh new file mode 100644 index 000000000..df9d0d69c --- /dev/null +++ b/lightx2v_fast_ulysses/csrc/tma_ptx.cuh @@ -0,0 +1,135 @@ +// Copied and adapted from https://github.com/triple-mu/fast-ulysses +#pragma once +// Raw PTX wrappers for TMA / mbarrier / system-scope release-acquire: lifts the kernel's +// inline asm into named __device__ __forceinline__ functions. Byte-for-byte equivalent to the +// original inline asm (clobber lists preserved verbatim). +// - mbar_* take an already-converted uint32_t smem address (no __cvta_generic_to_shared inside). +// - tma_load_4d / tma_store_4d / async_proxy_fence keep the trailing : "memory" / ::: "memory" +// visibility fence. Dropping it lets the compiler reorder smem accesses across the fence, +// causing rare data corruption that small shapes won't catch. +// - tma_commit_group / tma_wait_group / mbar_init / mbar_arrive_expect keep their original +// no-clobber form. +// +// The TMA / mbarrier helpers below use sm90+ only PTX (cp.async.bulk.tensor, mbarrier.*expect_tx, +// mbarrier.try_wait.parity, cp.async.bulk.commit_group/wait_group, fence.proxy.async). They are +// reached only on the TMA path, which the host gates to sm90+ at runtime (resolve_config never picks TMA on sm<9). +// For multi-arch builds that include sm80, the sm80 device pass cannot assemble these features, so +// we emit __trap() stubs for __CUDA_ARCH__ < 900; they are never executed on pre-Hopper GPUs. The +// host pass (__CUDA_ARCH__ undefined) takes the real definitions but does not codegen __device__ +// bodies, so the asm is harmless there. +#include + +namespace ulysses { + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 900) + +// --- sm80 (pre-Hopper) trap stubs: the TMA path is never selected here at runtime, so every sm90+ helper +// compiles to an unreachable __trap(). One macro covers the non-template stubs. --- +#define ULYSSES_TMA_TRAP_STUB(name, ...) \ + __device__ __forceinline__ void name(__VA_ARGS__) \ + { \ + __trap(); \ + } +ULYSSES_TMA_TRAP_STUB(mbar_init, uint32_t) +ULYSSES_TMA_TRAP_STUB(mbar_arrive_expect, uint32_t, uint32_t) +ULYSSES_TMA_TRAP_STUB(mbar_wait, uint32_t, int) +ULYSSES_TMA_TRAP_STUB(tma_load_4d, uint32_t, const void*, int, int, int, int, uint32_t) +ULYSSES_TMA_TRAP_STUB(tma_store_4d, const void*, int, int, int, int, uint32_t) +ULYSSES_TMA_TRAP_STUB(tma_commit_group) +ULYSSES_TMA_TRAP_STUB(async_proxy_fence) +#undef ULYSSES_TMA_TRAP_STUB +template +__device__ __forceinline__ void tma_wait_group() +{ + __trap(); +} + +#else + +// Init mbarrier with expected arrival count 1. Arg is an already-__cvta'd uint32_t smem address. +__device__ __forceinline__ void mbar_init(uint32_t mbar) +{ + asm volatile("mbarrier.init.shared.b64 [%0], 1;" ::"r"(mbar)); +} + +// Arrive at mbarrier and declare this phase's expected transfer bytes (for TMA complete_tx). +__device__ __forceinline__ void mbar_arrive_expect(uint32_t mbar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" ::"r"(mbar), "r"(bytes)); +} + +// Spin-wait for the mbarrier parity flip. The 'L_%=' label is made unique so the asm can be +// inlined multiple times (a fixed 'L:' would trigger a ptxas duplicate-label error). +__device__ __forceinline__ void mbar_wait(uint32_t mbar, int phase) +{ + asm volatile( + "{\n .reg .pred p;\n L_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1;\n @!p bra L_%=;\n }\n" ::"r"(mbar), + "r"(phase)); +} + +// TMA 4D load: global -> smem, completion counted into mbar via complete_tx. smem/mbar are +// already-__cvta'd uint32_t addresses. +__device__ __forceinline__ void +tma_load_4d(uint32_t smem, const void* map, int c0, int c1, int c2, int c3, uint32_t mbar) +{ + asm volatile("cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes " + "[%0], [%1, {%2, %3, %4, %5}], [%6];" ::"r"(smem), + "l"(map), + "r"(c0), + "r"(c1), + "r"(c2), + "r"(c3), + "r"(mbar) + : "memory"); +} + +// TMA 4D store: smem -> global (bulk_group). smem is an already-__cvta'd uint32_t address. +__device__ __forceinline__ void tma_store_4d(const void* map, int c0, int c1, int c2, int c3, uint32_t smem) +{ + asm volatile("cp.async.bulk.tensor.4d.global.shared::cta.bulk_group " + "[%0, {%1, %2, %3, %4}], [%5];" ::"l"(map), + "r"(c0), + "r"(c1), + "r"(c2), + "r"(c3), + "r"(smem) + : "memory"); +} + +// Commit one bulk_group (enclosing the TMA stores issued above). +__device__ __forceinline__ void tma_commit_group() +{ + asm volatile("cp.async.bulk.commit_group;"); +} + +// Wait until at most N bulk_groups remain in flight (wait for all but the last N stores). N is a template +// non-type param because the PTX immediate must be a compile-time constant; the kernel (templated on STAGES) +// calls tma_wait_group() to keep the pipeline full and tma_wait_group<0>() to drain at the end. +template +__device__ __forceinline__ void tma_wait_group() +{ + asm volatile("cp.async.bulk.wait_group %0;" ::"n"(N)); +} + +// async-proxy visibility fence: makes subsequent generic writes visible to TMA (the async proxy). +__device__ __forceinline__ void async_proxy_fence() +{ + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + +#endif // __CUDA_ARCH__ < 900 + +// System-scope release store / acquire load (sm70+, used by fast_barrier on all archs; not gated). +__device__ __forceinline__ void st_release_sys_u64(uint64_t* addr, uint64_t v) +{ + asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(addr), "l"(v) : "memory"); +} + +__device__ __forceinline__ uint64_t ld_acquire_sys_u64(const uint64_t* addr) +{ + uint64_t v; + asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(addr) : "memory"); + return v; +} + +} // namespace ulysses diff --git a/lightx2v_fast_ulysses/csrc/ulysses_common.cuh b/lightx2v_fast_ulysses/csrc/ulysses_common.cuh new file mode 100644 index 000000000..08587bfb9 --- /dev/null +++ b/lightx2v_fast_ulysses/csrc/ulysses_common.cuh @@ -0,0 +1,91 @@ +// Copied and adapted from https://github.com/triple-mu/fast-ulysses +#pragma once +#include +#include +#include +#include +#include + +namespace ulysses { + +// Check a CUDA runtime call's status and throw (TORCH_CHECK) on failure, including the call text and the +// driver error string. Use for real (non-autotune-probe) runtime calls; for kernel launches pass +// cudaGetLastError(). The autotune OOR probe deliberately inspects cudaGetLastError() itself, so it does +// NOT use this macro. +#define ULYSSES_CUDA_CHECK(expr) \ + do { \ + cudaError_t err_ = (expr); \ + TORCH_CHECK(err_ == cudaSuccess, "CUDA error (" #expr "): ", cudaGetErrorString(err_)); \ + } while (0) + +// Shared host/device dim descriptor (int32 fields; kernels use int64 for addressing). +struct Ulysses4DDims { + int32_t b, s_local, s_global, n_local, n_global, d, rank; +}; + +struct A2AConfig; // a2a_config.cuh + +template +struct PeerPtrs { + void* p[WS]; +}; + +// A2A copies raw bytes; this no-op keeps the generic copy kernel simple. +struct EpilogueIdentity { + __device__ __forceinline__ void operator()(uint8_t* /*row_bytes*/, int /*row_off*/) const {} +}; + +// Free function: pure-CUDA vectorized direct write to peer symmetric memory (host already +// resolved peer_ptrs via nvshmem_ptr). Launch with the given config (no env, no autotune). +void launch_a2a(const void* src, + const std::vector& peer_ptrs, + const Ulysses4DDims& dims, + int mode, + int elem_size, + const A2AConfig& cfg, + cudaStream_t stream); + +// Local microbench shared by both resolve paths: warmup + time 10 iters, return us/call. Defined in +// all_to_all.cu. run_once is the REAL per-call op (launch + finish, where finish = quiet + fast_barrier), +// so its ranking matches steady-state perf. The fast_barrier inside is hang-safe: under pure-lazy SPMD all +// ranks miss the same (shape,mode,tma) on the first call together, so they call equal barriers in lockstep. +float microbench_us(const std::function& run_once, cudaStream_t stream); + +// SM count of the current device, queried once per process and cached. Defined in all_to_all.cu. +int sm_count_cached(); + +// non-TMA config resolution: micro-benchmark sweep over threads x unroll x blocks, keep the fastest (result +// held by the group's cfg_cache_). finish: the per-call quiet+barrier, appended to each timed run so the +// measurement reflects real per-call cost. verbose: print the chosen config (rank 0 only). +A2AConfig resolve_config_nontma(const void* src, + const std::vector& peer_ptrs, + const Ulysses4DDims& dims, + int mode, + int elem_size, + cudaStream_t stream, + bool verbose, + const std::function& finish); + +// TMA version (fewer SMs, better comm/compute overlap); uniform mode0/mode1. See all_to_all_tma.cu. +// Launch with the given config (build maps + launch only; no team, no autotune, no env). +void launch_a2a_tma(const void* src, + const std::vector& peer_ptrs, + const Ulysses4DDims& dims, + int mode, + int elem_size, + const A2AConfig& cfg, + cudaStream_t stream); + +// TMA config resolution: pick the best candidate by micro-benchmark (result cached in the group's +// cfg_cache_; a hit skips the micro-benchmark). finish: per-call quiet+barrier appended to each timed run. +// verbose: rank-0 debug printing. +A2AConfig resolve_config_tma(const void* src, + const std::vector& peer_ptrs, + const Ulysses4DDims& dims, + int mode, + int elem_size, + bool verbose, + cudaStream_t stream, + const std::function& finish); + +} // namespace ulysses diff --git a/lightx2v_fast_ulysses/csrc/ulysses_group.cu b/lightx2v_fast_ulysses/csrc/ulysses_group.cu new file mode 100644 index 000000000..1a15835d3 --- /dev/null +++ b/lightx2v_fast_ulysses/csrc/ulysses_group.cu @@ -0,0 +1,234 @@ +// Copied and adapted from https://github.com/triple-mu/fast-ulysses +#include "tma_ptx.cuh" +#include "ulysses_group.cuh" +#include +#include +#include +#include +#include +#include +#include + +namespace ulysses { + +static bool g_world_inited = false; +static int g_live_groups = 0; + +// ---- Custom NVLink flag barrier ---- +// Flag layout: each rank holds uint64 flags[ws]. On arrival, rank r writes epoch into every peer p's +// flags[r] (P2P, release/sys), then spins until its own flags[0..ws-1] all == epoch (acquire/sys). +// Monotonically increasing epoch means no reset and no ABA. Strict lockstep (SPMD collective) keeps +// epoch identical across ranks throughout. +struct BarPeers { + uint64_t p[8]; +}; + +__global__ void ulysses_barrier_kernel(uint64_t* local, BarPeers peers, int ws, int rank, uint64_t epoch) +{ + int t = threadIdx.x; + if (t >= ws) + return; + uint64_t* remote = reinterpret_cast(peers.p[t]) + rank; // peer t's flags[rank] + st_release_sys_u64(remote, epoch); + uint64_t v; + uint64_t* mine = local + t; // own flags[t] (written by peer t) + do { + v = ld_acquire_sys_u64(mine); + } while (v < epoch); +} + +void UlyssesGroup::fast_barrier(cudaStream_t stream) +{ + if (world_size_ == 1) + return; + if (!bar_ready_) { + const auto& buf = pool_->acquire({static_cast(world_size_)}, at::kLong, "__ulysses_sync__"); + bar_local_ = buf.sym_base; + bar_peers_ = buf.peer_ptrs; + ULYSSES_CUDA_CHECK( + cudaMemsetAsync(bar_local_, 0, world_size_ * sizeof(uint64_t), stream)); // init 0; epoch starts at 1 + // One slow sync: ensure all ranks finish clearing before anyone writes a flag (otherwise the + // clear could overwrite an already-written epoch). + nvshmemx_barrier_on_stream(team_, stream); + bar_ready_ = true; + } + ++bar_epoch_; + BarPeers peers; + for (int i = 0; i < world_size_; ++i) + peers.p[i] = bar_peers_[i]; + ulysses_barrier_kernel<<<1, 32, 0, stream>>>( + reinterpret_cast(bar_local_), peers, world_size_, my_rank_, bar_epoch_); + ULYSSES_CUDA_CHECK(cudaGetLastError()); // catch a barrier-kernel launch failure +} + +int64_t UlyssesGroup::uniqueid_nints() +{ + return static_cast((sizeof(nvshmemx_uniqueid_t) + 7) / 8); +} + +std::vector UlyssesGroup::get_uniqueid() +{ + nvshmemx_uniqueid_t uid; + std::memset(&uid, 0, sizeof(uid)); + TORCH_CHECK(nvshmemx_get_uniqueid(&uid) == 0, "nvshmemx_get_uniqueid failed"); + std::vector out(uniqueid_nints(), 0); + std::memcpy(out.data(), &uid, sizeof(uid)); + return out; +} + +void UlyssesGroup::init_world(std::vector uid_ints, int64_t global_rank, int64_t global_nranks) +{ + if (g_world_inited) + return; + TORCH_CHECK(static_cast(uid_ints.size()) >= uniqueid_nints(), "uid_ints too short"); + nvshmemx_uniqueid_t uid; + std::memcpy(&uid, uid_ints.data(), sizeof(uid)); + // Use INITIALIZER, not memset(0): it stamps the version field of attr/args/uid_args. + // nvshmemx_set_attr_uniqueid_args does not write version, and hostlib_init_attr dispatches the V2 + // path based on attr.args.version, so the version must be stamped first (inline nvshmemx_init_attr + // auto-stamps when version is invalid; here we explicitly substitute that step). + nvshmemx_init_attr_t attr = NVSHMEMX_INIT_ATTR_INITIALIZER; + TORCH_CHECK( + nvshmemx_set_attr_uniqueid_args(static_cast(global_rank), static_cast(global_nranks), &uid, &attr) + == 0, + "nvshmemx_set_attr_uniqueid_args failed"); + // DEVIATION (see task-5-report): use the host-lib direct entry nvshmemx_hostlib_init_attr instead of + // inline nvshmemx_init_attr. The inline version calls nvshmemi_init_thread, a symbol that lives only + // in static libnvshmem_device.a; linking it clashes with the NVSHMEM version node of torch's bundled + // libtorch_nvshmem.so (undefined symbol nvshmem_selected_device_transport). hostlib_init_attr is the + // equivalent entry exported directly by the host shared library (NVSHMEM's own python UID path uses it). + TORCH_CHECK(nvshmemx_hostlib_init_attr(NVSHMEMX_INIT_WITH_UNIQUEID, &attr) == 0, + "nvshmemx_hostlib_init_attr failed"); + g_world_inited = true; +} + +UlyssesGroup::UlyssesGroup(std::vector peer_global_pes, + int64_t my_rank, + int64_t device_id, + int64_t reserved_bytes): + my_rank_(static_cast(my_rank)), + world_size_(static_cast(peer_global_pes.size())), + device_id_(static_cast(device_id)) +{ + TORCH_CHECK(g_world_inited, "init_world must be called before constructing UlyssesGroup"); + // Cache compute capability major (TMA cp.async.bulk.tensor is Hopper(sm90)+ only); whether a call + // actually takes the TMA path is decided per call by resolve_config (explicit use_tma, or auto-measured). + int major = 0; + ULYSSES_CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_id_)); + sm_major_ = major; + TORCH_CHECK( + world_size_ >= 1 && world_size_ <= 8, "world_size must be in [1, 8] (single-node NVLink), got ", world_size_); + ULYSSES_CUDA_CHECK(cudaSetDevice(device_id_)); + peer_global_pes_.reserve(world_size_); + for (auto pe : peer_global_pes) + peer_global_pes_.push_back(static_cast(pe)); + + // team: if the group covers the whole world contiguously -> TEAM_WORLD; else a contiguous stride-1 + // subgroup via split_strided. + const int gpes = nvshmem_n_pes(); + bool is_full_world = (world_size_ == gpes); + for (int i = 0; i < world_size_ && is_full_world; ++i) + if (peer_global_pes_[i] != i) + is_full_world = false; + if (is_full_world) { + team_ = NVSHMEM_TEAM_WORLD; + owns_team_ = false; + } + else { + int start = peer_global_pes_[0]; + for (int i = 1; i < world_size_; ++i) + TORCH_CHECK(peer_global_pes_[i] == start + i, "phase-1 only supports a contiguous PE subgroup"); + nvshmem_team_config_t cfg; + std::memset(&cfg, 0, sizeof(cfg)); + TORCH_CHECK(nvshmem_team_split_strided(NVSHMEM_TEAM_WORLD, start, 1, world_size_, &cfg, 0, &team_) == 0, + "nvshmem_team_split_strided failed"); + owns_team_ = true; + } + + pool_ = std::make_unique(reserved_bytes, world_size_, peer_global_pes_); + ++g_live_groups; +} + +// See the resolve_config / hang-safety notes in ulysses_group.cuh. finish is the real per-call tail (quiet + +// fast_barrier); every micro-benchmark times launch+finish so its ranking matches steady state and the two +// paths' times are directly comparable. +UlyssesGroup::PathConfig UlyssesGroup::resolve_config(const Ulysses4DDims& dims, + int mode, + int use_tma_i, + const void* src, + const std::vector& peers, + int elem, + cudaStream_t stream) +{ + auto finish = [this, stream] { + nvshmemx_quiet_on_stream(stream); + fast_barrier(stream); + }; + // Autotune the given path (cached per (ws,mode,tma,dims)). + auto resolve_path = [&](bool tma) -> A2AConfig { + const ConfigKey key = config_key(world_size_, mode, tma, dims); + auto it = cfg_cache_.find(key); + if (it != cfg_cache_.end()) + return it->second; + A2AConfig cfg = tma ? resolve_config_tma(src, peers, dims, mode, elem, false, stream, finish) : + resolve_config_nontma(src, peers, dims, mode, elem, stream, false, finish); + cfg_cache_[key] = cfg; + return cfg; + }; + + // Explicit path: force it. + if (use_tma_i == 1) + return {true, resolve_path(true)}; + if (use_tma_i == 0) + return {false, resolve_path(false)}; + + // Auto: sm<9 has no TMA; sm90+ picks the faster of the two paths (runtime replacement of the DiT table). + if (sm_major_ < 9) + return {false, resolve_path(false)}; + const std::tuple akey{mode, dims.n_local, dims.s_local, dims.d}; + auto ap = auto_path_cache_.find(akey); + if (ap != auto_path_cache_.end()) + return {ap->second, resolve_path(ap->second)}; + // First auto call for this shape: tune both paths, time each real per-call, keep the faster. + const A2AConfig cfg_n = resolve_path(false); + const A2AConfig cfg_t = resolve_path(true); + const float t_n = microbench_us( + [&] { + launch_a2a(src, peers, dims, mode, elem, cfg_n, stream); + finish(); + }, + stream); + const float t_t = microbench_us( + [&] { + launch_a2a_tma(src, peers, dims, mode, elem, cfg_t, stream); + finish(); + }, + stream); + const bool tma = t_t <= t_n; + auto_path_cache_[akey] = tma; + return {tma, tma ? cfg_t : cfg_n}; +} + +void UlyssesGroup::destroy() +{ + if (destroyed_) + return; + if (pool_) + pool_->destroy(); + if (owns_team_) + nvshmem_team_destroy(team_); + destroyed_ = true; + if (--g_live_groups == 0 && g_world_inited) { + // DEVIATION: nvshmem_finalize() is inline and calls nvshmemi_finalize() (again only in static + // device.a). Use nvshmemx_hostlib_finalize() exported by the host shared library instead. + nvshmemx_hostlib_finalize(); + g_world_inited = false; + } +} + +UlyssesGroup::~UlyssesGroup() +{ + destroy(); +} + +} // namespace ulysses diff --git a/lightx2v_fast_ulysses/csrc/ulysses_group.cuh b/lightx2v_fast_ulysses/csrc/ulysses_group.cuh new file mode 100644 index 000000000..156d6ba11 --- /dev/null +++ b/lightx2v_fast_ulysses/csrc/ulysses_group.cuh @@ -0,0 +1,113 @@ +// Copied and adapted from https://github.com/triple-mu/fast-ulysses +#pragma once +#include "a2a_config.cuh" +#include "symmetric_pool.cuh" +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ulysses { + +class UlyssesGroup: public torch::CustomClassHolder { +public: + static int64_t uniqueid_nints(); // ceil(sizeof(nvshmemx_uniqueid_t)/8) + static std::vector get_uniqueid(); // rank0 only + static void init_world(std::vector uid_ints, int64_t global_rank, int64_t global_nranks); // idempotent + + UlyssesGroup(std::vector peer_global_pes, int64_t my_rank, int64_t device_id, int64_t reserved_bytes); + ~UlyssesGroup() override; + + int64_t rank() const + { + return my_rank_; + } + int64_t world_size() const + { + return world_size_; + } + void destroy(); + + SymmetricHeapPool& pool() + { + return *pool_; + } + nvshmem_team_t team() const + { + return team_; + } + + // Device compute capability major (cached via cudaDeviceGetAttribute at construction). resolve_config uses + // it for the auto path choice (sm<9 has no TMA). + int sm_major() const + { + return sm_major_; + } + + struct PathConfig { + bool tma; + A2AConfig cfg; + }; + + // Resolve the launch path + config for (dims,mode). use_tma_i is the tri-state -1 auto / 0 non-TMA / + // 1 TMA (the sm<9 error for an explicit 1 is enforced by the caller). Returns the path to launch and + // its config; a cache hit returns directly. + // + // Explicit 0/1: micro-benchmark that path's candidates, keep the fastest. Auto (-1): on sm<9 -> non-TMA; + // on sm90+ micro-benchmark BOTH paths and pick the faster (the runtime replacement of the old static DiT + // table -- "tune the best path+config and cache it"), memoised in auto_path_cache_. + // + // The microbench times the REAL per-call op (launch + quiet + fast_barrier) so its ranking matches steady + // state, and the two paths' times are directly comparable. The fast_barrier makes the miss branch + // cross-rank, but it is HANG-SAFE: pure-lazy SPMD (no tune()) means all ranks issue the same + // (shape,mode,use_tma) sequence and miss the same entry on the first call together -> equal barrier calls + // in lockstep (the candidate count is a function of rank-invariant dims, so it is identical on every rank). + // A cache-hit rank and a miss rank never coexist for the same call, so no rank blocks alone. + // + // The auto path choice is a per-rank local timing decision, so on a near-tie shape different ranks may pick + // different kernels. This is harmless: the two paths are functionally equivalent P2P writes (each rank + // writes its own region correctly either way), and per-call barrier counts stay equal regardless of path. + // + // Thread safety: resolve_config/all_to_all on the same group instance must be called serially (SPMD + // single-threaded; the caches are lock-free, so concurrent multi-stream use is not thread-safe). + PathConfig resolve_config(const Ulysses4DDims& dims, + int mode, + int use_tma_i, + const void* src, + const std::vector& peers, + int elem, + cudaStream_t stream); + + // Custom single-node NVLink flag barrier: replaces the slow nvshmem sync (~280us) that falls back on + // hardware without NVLS fabric. Call nvshmemx_quiet_on_stream first (so this rank's writes are globally + // visible). No-op when world_size==1. + void fast_barrier(cudaStream_t stream); + +private: + int my_rank_, world_size_, device_id_; + int sm_major_ = 0; // cached cudaDeviceGetAttribute(major) at construction + std::vector peer_global_pes_; + nvshmem_team_t team_; + bool owns_team_ = false; + bool destroyed_ = false; + std::unique_ptr pool_; + + // cfg_cache_ holds the best config per (ws,mode,tma,n_local,s_local,d). Lock-free std::map, + // must be accessed serially (see resolve_config comment). + std::map cfg_cache_; + // auto_path_cache_ memoises the best path for the auto (use_tma=None) case per (mode,n_local,s_local,d) + // -- true=TMA -- so a repeat auto call skips the two-path micro-benchmark. + std::map, bool> auto_path_cache_; + + // fast_barrier state: symmetric flag buffer (uint64[ws]) + monotonic epoch (incremented lockstep per rank). + bool bar_ready_ = false; + uint64_t bar_epoch_ = 0; + void* bar_local_ = nullptr; // this rank's flag base + std::vector bar_peers_; // per-peer flag base (including self) +}; + +} // namespace ulysses diff --git a/lightx2v_fast_ulysses/pyproject.toml b/lightx2v_fast_ulysses/pyproject.toml new file mode 100644 index 000000000..6f46efed2 --- /dev/null +++ b/lightx2v_fast_ulysses/pyproject.toml @@ -0,0 +1,18 @@ +[build-system] +requires = [ + "setuptools>=61.0", + "wheel", + "packaging", + "ninja", +] +build-backend = "setuptools.build_meta" + +[project] +name = "lightx2v-fast-ulysses" +version = "0.0.1" +description = "LightX2V-vendored fast-ulysses NVSHMEM all-to-all extension" +requires-python = ">=3.10" +dependencies = ["torch"] + +[tool.setuptools] +include-package-data = true diff --git a/lightx2v_fast_ulysses/setup.py b/lightx2v_fast_ulysses/setup.py new file mode 100644 index 000000000..b1c4a840a --- /dev/null +++ b/lightx2v_fast_ulysses/setup.py @@ -0,0 +1,84 @@ +# Copied and adapted from https://github.com/triple-mu/fast-ulysses +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +import sysconfig +from pathlib import Path + +from setuptools import Extension, setup +from setuptools.command.build_ext import build_ext +from setuptools.command.build_py import build_py + +_HERE = Path(__file__).resolve().parent +os.chdir(_HERE) + +NVSHMEM_HOME = os.environ.get("NVSHMEM_HOME") +if not NVSHMEM_HOME: + raise RuntimeError("NVSHMEM_HOME must point to an NVSHMEM installation before building lightx2v-fast-ulysses.") +if not (Path(NVSHMEM_HOME) / "include" / "nvshmem.h").exists(): + raise FileNotFoundError(f"NVSHMEM_HOME invalid: {NVSHMEM_HOME}") + + +class CMakeExtension(Extension): + def __init__(self, name: str) -> None: + super().__init__(name, sources=[]) + + +class CMakeBuild(build_ext): + def build_extension(self, ext: Extension) -> None: + outdir = Path(self.get_ext_fullpath(ext.name)).resolve().parent + ext_suffix = sysconfig.get_config_var("EXT_SUFFIX") or ".so" + # Persistent build dir (not pip's ephemeral build_temp) so CMake can + # build incrementally and skip recompiling unchanged TUs. CMake re-runs + # configure automatically when the -D args change, so no manual wipe. + builddir = _HERE / "build" + builddir.mkdir(parents=True, exist_ok=True) + env = os.environ.copy() + arch = env.get("CUSTOM_ULYSSES_CUDA_ARCH", "80;90;100") + # Torch expects dotted arch (e.g. 86 -> 8.6, 100 -> 10.0); insert the + # decimal point before the last digit of each ";"-separated token. + env["TORCH_CUDA_ARCH_LIST"] = " ".join( + f"{tok[:-1]}.{tok[-1]}" for tok in arch.split(";") if tok + ) + subprocess.check_call( + [ + "cmake", + "-S", + str(_HERE), + "-B", + str(builddir), + f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={outdir}", + f"-DPython_EXECUTABLE={sys.executable}", + f"-DEXT_SUFFIX={ext_suffix}", + "-DCMAKE_BUILD_TYPE=Release", + f"-DCMAKE_CUDA_ARCHITECTURES={arch}", + f"-DNVSHMEM_HOME={NVSHMEM_HOME}", + ], + env=env, + ) + subprocess.check_call(["cmake", "--build", str(builddir), "-j"], env=env) + built_ext = outdir / f"_C{ext_suffix}" + if built_ext.exists(): + shutil.copy2(built_ext, _HERE / built_ext.name) + + +class CleanBuildPy(build_py): + def run(self) -> None: + package_build_dir = Path(self.build_lib) / "lightx2v_fast_ulysses" + if package_build_dir.exists(): + shutil.rmtree(package_build_dir) + super().run() + + +setup( + name="lightx2v-fast-ulysses", + version="0.0.1", + packages=["lightx2v_fast_ulysses"], + package_dir={"lightx2v_fast_ulysses": "."}, + package_data={"lightx2v_fast_ulysses": ["NOTICE.md"]}, + ext_modules=[CMakeExtension("lightx2v_fast_ulysses._C")], + cmdclass={"build_ext": CMakeBuild, "build_py": CleanBuildPy}, +) diff --git a/pyproject.toml b/pyproject.toml index c62264b41..de38767d0 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,7 +79,7 @@ include-package-data = true [tool.setuptools.packages.find] include = ["lightx2v*"] -exclude = ["lightx2v_kernel*"] +exclude = ["lightx2v_kernel*", "lightx2v_fast_ulysses*"] [tool.ruff] exclude = [ diff --git a/test_cases/test_fast_ulysses_backend.py b/test_cases/test_fast_ulysses_backend.py new file mode 100644 index 000000000..762ab05f6 --- /dev/null +++ b/test_cases/test_fast_ulysses_backend.py @@ -0,0 +1,352 @@ +import unittest +import json +import sys +import types +from pathlib import Path +from unittest import mock + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + + +class FakeFastUlyssesGroup: + instances = [] + + def __init__(self, world_size): + self.world_size = world_size + self.calls = [] + self.destroyed = False + type(self).instances.append(self) + + def all_to_all_single_4d(self, x, *, mode=0, tag="", use_tma=None): + self.calls.append((mode, tag, tuple(x.shape), use_tma)) + b, x1, x2, d = x.shape + if mode == 0: + return x.reshape(b, x1 * self.world_size, x2 // self.world_size, d) + if mode == 1: + return x.reshape(b, x1 // self.world_size, x2 * self.world_size, d) + raise AssertionError(f"unexpected mode {mode}") + + def destroy(self): + self.destroyed = True + + +class FakeAttention: + def __init__(self): + self.kwargs = None + + def apply(self, **kwargs): + self.kwargs = kwargs + return kwargs["q"] + + +class FastUlyssesBackendTest(unittest.TestCase): + def tearDown(self): + from lightx2v.common.ops.attn.fast_ulysses_attn import FastUlyssesAttnWeight + + FastUlyssesAttnWeight._groups.clear() + FastUlyssesAttnWeight._single_node_cache.clear() + FastUlyssesAttnWeight.reset_runtime_stats() + FakeFastUlyssesGroup.instances.clear() + + def test_fast_ulysses_backend_is_registered(self): + from lightx2v.common.ops import attn # noqa: F401 + from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER + + self.assertIn("fast_ulysses", ATTN_WEIGHT_REGISTER) + + def test_unsupported_flags_fallback_before_native_import(self): + from lightx2v.common.ops.attn.fast_ulysses_attn import FastUlyssesAttnWeight + + q = torch.zeros(4, 4, 8) + fallback_out = torch.empty(1) + cases = ( + ("use_fp8_comm", {"use_fp8_comm": True}), + ("use_fp4_comm", {"use_fp4_comm": True}), + ("use_tensor_fusion", {"use_tensor_fusion": True}), + ("enable_head_parallel", {"enable_head_parallel": True}), + ("q_only_img", {"q_only_img": True}), + ("img_first", {"img_first": False}), + ) + for message, kwargs in cases: + with self.subTest(message=message): + backend = FastUlyssesAttnWeight() + with ( + mock.patch.object(backend, "_get_group", side_effect=AssertionError("native path must not initialize")), + mock.patch("lightx2v.common.ops.attn.fast_ulysses_attn.UlyssesAttnWeight.apply", return_value=fallback_out) as fallback, + ): + out = backend.apply( + q=q, + k=q, + v=q, + slice_qkv_len=4, + cu_seqlens_qkv=[0, 4], + attention_module=FakeAttention(), + seq_p_group=object(), + **kwargs, + ) + self.assertIs(out, fallback_out) + fallback.assert_called_once() + + def test_unsupported_shapes_fallback_before_native_import(self): + from lightx2v.common.ops.attn.fast_ulysses_attn import FastUlyssesAttnWeight + + backend = FastUlyssesAttnWeight() + q = torch.zeros(4, 4, 8, dtype=torch.bfloat16, device="cuda") + fallback_out = torch.empty(1) + with ( + mock.patch.object(backend, "_get_group", side_effect=AssertionError("native path must not initialize")), + mock.patch("lightx2v.common.ops.attn.fast_ulysses_attn.UlyssesAttnWeight.apply", return_value=fallback_out) as fallback, + ): + out = backend.apply( + q=q, + k=torch.zeros(4, 2, 8), + v=torch.zeros(4, 2, 8), + slice_qkv_len=4, + cu_seqlens_qkv=[0, 4], + attention_module=FakeAttention(), + seq_p_group=object(), + ) + self.assertIs(out, fallback_out) + fallback.assert_called_once() + + with ( + mock.patch("torch.distributed.get_world_size", return_value=1), + mock.patch.object(backend, "_get_group", side_effect=AssertionError("native path must not initialize")), + mock.patch("lightx2v.common.ops.attn.fast_ulysses_attn.UlyssesAttnWeight.apply", return_value=fallback_out) as fallback, + ): + out = backend.apply( + q=q, + k=q, + v=q, + slice_qkv_len=2, + cu_seqlens_qkv=[0, 4], + attention_module=FakeAttention(), + seq_p_group=object(), + ) + self.assertIs(out, fallback_out) + fallback.assert_called_once() + + def test_cross_node_detection_fallback_before_native_import(self): + from lightx2v.common.ops.attn.fast_ulysses_attn import FastUlyssesAttnWeight + + backend = FastUlyssesAttnWeight() + q = torch.zeros(4, 4, 8) + fallback_out = torch.empty(1) + with ( + mock.patch("torch.distributed.get_world_size", return_value=2), + mock.patch.object(backend, "_is_single_node", return_value=False), + mock.patch.object(backend, "_get_group", side_effect=AssertionError("native path must not initialize")), + mock.patch("lightx2v.common.ops.attn.fast_ulysses_attn.UlyssesAttnWeight.apply", return_value=fallback_out) as fallback, + ): + out = backend.apply( + q=q, + k=q, + v=q, + slice_qkv_len=4, + cu_seqlens_qkv=[0, 4], + attention_module=FakeAttention(), + seq_p_group=object(), + ) + self.assertIs(out, fallback_out) + fallback.assert_called_once() + + def test_native_init_error_fallback(self): + from lightx2v.common.ops.attn.fast_ulysses_attn import FastUlyssesAttnWeight + + backend = FastUlyssesAttnWeight() + q = torch.zeros(4, 4, 8, dtype=torch.bfloat16, device="cuda") + fallback_out = torch.empty(1) + with ( + mock.patch("torch.distributed.get_world_size", return_value=2), + mock.patch.object(backend, "_is_single_node", return_value=True), + mock.patch.object(backend, "_get_group", side_effect=ImportError("missing native extension")), + mock.patch("lightx2v.common.ops.attn.fast_ulysses_attn.UlyssesAttnWeight.apply", return_value=fallback_out) as fallback, + ): + out = backend.apply( + q=q, + k=q, + v=q, + slice_qkv_len=4, + cu_seqlens_qkv=[0, 4], + attention_module=FakeAttention(), + seq_p_group=object(), + ) + self.assertIs(out, fallback_out) + fallback.assert_called_once() + + def test_native_precondition_failures_fallback_before_native_import(self): + from lightx2v.common.ops.attn.fast_ulysses_attn import FastUlyssesAttnWeight + + fallback_out = torch.empty(1) + cases = ( + ( + torch.zeros(4, 4, 8), + "cpu tensors", + {}, + ), + ( + torch.zeros(4, 4, 8, device="cuda"), + "cuda float32 tensors", + {}, + ), + ( + torch.zeros(4, 4, 7, dtype=torch.bfloat16, device="cuda"), + "unaligned head dim", + {}, + ), + ( + torch.zeros(4, 4, 8, dtype=torch.bfloat16, device="cuda"), + "forced TMA on sm80", + {"LIGHTX2V_FAST_ULYSSES_USE_TMA": "1"}, + ), + ) + for q, name, env in cases: + with self.subTest(name=name): + backend = FastUlyssesAttnWeight() + patches = [ + mock.patch("torch.distributed.get_world_size", return_value=2), + mock.patch.object(backend, "_is_single_node", return_value=True), + mock.patch.object(backend, "_get_group", side_effect=AssertionError("native path must not initialize")), + mock.patch("lightx2v.common.ops.attn.fast_ulysses_attn.UlyssesAttnWeight.apply", return_value=fallback_out), + mock.patch.dict("os.environ", env, clear=False), + ] + if env: + patches.append(mock.patch("torch.cuda.get_device_capability", return_value=(8, 0))) + with patches[0], patches[1], patches[2], patches[3] as fallback, patches[4]: + if env: + with patches[5]: + out = backend.apply( + q=q, + k=q, + v=q, + slice_qkv_len=4, + cu_seqlens_qkv=[0, 4], + attention_module=FakeAttention(), + seq_p_group=object(), + ) + else: + out = backend.apply( + q=q, + k=q, + v=q, + slice_qkv_len=4, + cu_seqlens_qkv=[0, 4], + attention_module=FakeAttention(), + seq_p_group=object(), + ) + self.assertIs(out, fallback_out) + fallback.assert_called_once() + + def test_validated_pure_image_path_wraps_a2a_and_attention(self): + from lightx2v.common.ops.attn.fast_ulysses_attn import FastUlyssesAttnWeight + + q = torch.arange(4 * 4 * 8, device="cuda").reshape(4, 4, 8).to(torch.bfloat16) + k = q + 1 + v = q + 2 + group = FakeFastUlyssesGroup(world_size=2) + attention = FakeAttention() + backend = FastUlyssesAttnWeight() + + with ( + mock.patch("torch.distributed.get_world_size", return_value=2), + mock.patch("torch.distributed.get_rank", return_value=0), + mock.patch.object(backend, "_is_single_node", return_value=True), + mock.patch.object(backend, "_get_group", return_value=group), + ): + out = backend.apply( + q=q, + k=k, + v=v, + slice_qkv_len=4, + cu_seqlens_qkv=[0, 4], + attention_module=attention, + seq_p_group=object(), + block_idx=7, + ) + + self.assertEqual(tuple(out.shape), (4, 32)) + self.assertEqual([call[0] for call in group.calls], [0, 0, 0, 1]) + self.assertEqual([call[1] for call in group.calls], ["lightx2v_q", "lightx2v_k", "lightx2v_v", "lightx2v_out"]) + self.assertEqual(tuple(attention.kwargs["q"].shape), (8, 2, 8)) + self.assertEqual(attention.kwargs["max_seqlen_q"], 8) + self.assertEqual(attention.kwargs["max_seqlen_kv"], 8) + self.assertEqual(FastUlyssesAttnWeight.runtime_stats()["fast_path_calls"], 1) + self.assertEqual(FastUlyssesAttnWeight.runtime_stats()["fallback_calls"], 0) + + def test_group_cache_is_shared_across_backend_instances(self): + from lightx2v.common.ops.attn.fast_ulysses_attn import FastUlyssesAttnWeight + + fake_module = types.SimpleNamespace(UlyssesGroup=lambda **kwargs: FakeFastUlyssesGroup(world_size=2)) + group = object() + + with ( + mock.patch.dict(sys.modules, {"lightx2v_fast_ulysses": fake_module}), + mock.patch("torch.cuda.current_device", return_value=0), + mock.patch("torch.distributed.get_process_group_ranks", return_value=[0, 1]), + ): + first = FastUlyssesAttnWeight()._get_group(group) + second = FastUlyssesAttnWeight()._get_group(group) + + self.assertIs(first, second) + self.assertEqual(len(FakeFastUlyssesGroup.instances), 1) + + FastUlyssesAttnWeight.destroy_cached_groups() + self.assertTrue(first.destroyed) + self.assertEqual(FastUlyssesAttnWeight._groups, {}) + + def test_group_cache_uses_explicit_tensor_device(self): + from lightx2v.common.ops.attn.fast_ulysses_attn import FastUlyssesAttnWeight + + captured = [] + + def make_group(**kwargs): + captured.append(kwargs) + return FakeFastUlyssesGroup(world_size=2) + + fake_module = types.SimpleNamespace(UlyssesGroup=make_group) + group = object() + device = torch.device("cuda", 1) + + with ( + mock.patch.dict(sys.modules, {"lightx2v_fast_ulysses": fake_module}), + mock.patch("torch.cuda.current_device", return_value=0), + mock.patch("torch.distributed.get_process_group_ranks", return_value=[0, 1]), + ): + first = FastUlyssesAttnWeight()._get_group(group, device=device) + second = FastUlyssesAttnWeight()._get_group(group, device=device) + + self.assertIs(first, second) + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["device"], device) + self.assertIn(((0, 1), 1), FastUlyssesAttnWeight._groups) + + def test_native_package_surface_and_build_are_a2a_only(self): + root = Path(__file__).resolve().parents[1] / "lightx2v_fast_ulysses" + init_text = (root / "__init__.py").read_text() + comm_text = (root / "comm.py").read_text() + cmake_text = (root / "CMakeLists.txt").read_text() + bindings_text = (root / "csrc" / "bindings.cpp").read_text() + + csrc_text = "\n".join(path.read_text() for path in (root / "csrc").glob("*")) + for name in ("rms_norm", "norm_rope", "all_to_all_single_4d_qk", "config_key_qk", "qk_norm_rope", "all_to_all_qk"): + self.assertNotIn(name, init_text) + self.assertNotIn(name, comm_text) + self.assertNotIn(name, bindings_text) + self.assertNotIn(name, csrc_text) + self.assertNotIn("file(GLOB", cmake_text) + self.assertNotIn("all_to_all_qk.cu", cmake_text) + self.assertNotIn("qk_norm_rope.cu", cmake_text) + + def test_wan_i2v_fast_config_uses_fast_ulysses(self): + root = Path(__file__).resolve().parents[1] + fast_config = root / "configs" / "dist_infer" / "wan_i2v_dist_cfg_fast_ulysses.json" + cfg = json.loads(fast_config.read_text()) + self.assertEqual(cfg["parallel"]["seq_p_size"], 4) + self.assertEqual(cfg["parallel"]["cfg_p_size"], 1) + self.assertEqual(cfg["parallel"]["seq_p_attn_type"], "fast_ulysses") + + +if __name__ == "__main__": + unittest.main() From 7b31be81641feff52c816de333fc8fb12de945e1 Mon Sep 17 00:00:00 2001 From: xlycae Date: Fri, 10 Jul 2026 15:34:34 +0800 Subject: [PATCH 2/2] fix: address fast-ulysses review feedback --- lightx2v_fast_ulysses/__init__.py | 5 +- lightx2v_fast_ulysses/comm.py | 52 ++++---- lightx2v_fast_ulysses/csrc/ulysses_group.cu | 7 +- lightx2v_fast_ulysses/csrc/ulysses_group.cuh | 2 +- lightx2v_fast_ulysses/setup.py | 4 +- test_cases/test_fast_ulysses_backend.py | 122 ++++++++++++++++++- 6 files changed, 152 insertions(+), 40 deletions(-) diff --git a/lightx2v_fast_ulysses/__init__.py b/lightx2v_fast_ulysses/__init__.py index f019d930e..f9e2ba05c 100644 --- a/lightx2v_fast_ulysses/__init__.py +++ b/lightx2v_fast_ulysses/__init__.py @@ -6,10 +6,7 @@ try: from . import _C # noqa: F401,E402 trigger TORCH_LIBRARY registration except ImportError as exc: - raise ImportError( - "lightx2v_fast_ulysses native extension is not built. " - "Build it with: NVSHMEM_HOME=/path/to/nvshmem pip install ./lightx2v_fast_ulysses" - ) from exc + raise ImportError("lightx2v_fast_ulysses native extension is not built. Build it with: NVSHMEM_HOME=/path/to/nvshmem pip install ./lightx2v_fast_ulysses") from exc from .comm import AsyncA2AHandle, UlyssesGroup # noqa: E402 diff --git a/lightx2v_fast_ulysses/comm.py b/lightx2v_fast_ulysses/comm.py index fdf0b063c..4a198415c 100644 --- a/lightx2v_fast_ulysses/comm.py +++ b/lightx2v_fast_ulysses/comm.py @@ -27,10 +27,10 @@ def wait(self): class UlyssesGroup: def __init__( - self, - process_group: Optional[dist.ProcessGroup] = None, - device: Optional[torch.device] = None, - initial_pool_bytes: int = 2 << 30, + self, + process_group: Optional[dist.ProcessGroup] = None, + device: Optional[torch.device] = None, + initial_pool_bytes: int = 2 << 30, ) -> None: pg = process_group if process_group is not None else dist.group.WORLD self.pg = pg @@ -54,18 +54,20 @@ def __init__( os.environ.setdefault("NVSHMEM_REMOTE_TRANSPORT", "none") cls = torch.classes.fast_ulysses.UlyssesGroup - if dist.get_rank() == 0: + pg_rank = dist.get_rank(pg) + if pg_rank == 0: uid = cls.get_uniqueid() else: uid = [0] * cls.uniqueid_nints() uid_t = torch.tensor(uid, dtype=torch.int64, device=device) - dist.broadcast(uid_t, src=0, group=dist.group.WORLD) - cls.init_world(uid_t.tolist(), dist.get_rank(), dist.get_world_size()) + src_global_rank = dist.get_global_rank(pg, 0) + dist.broadcast(uid_t, src=src_global_rank, group=pg) + cls.init_world(uid_t.tolist(), pg_rank, self.world_size) dist.barrier(group=pg) self._group = cls( - [int(r) for r in self.peer_global_ranks], - int(self.rank), + list(range(self.world_size)), + int(pg_rank), int(device.index), int(initial_pool_bytes), ) @@ -97,12 +99,12 @@ def _launch_on_comm_stream(self, inputs: list[torch.Tensor], fn: Callable): return out, ev_done def all_to_all_single_4d( - self, - x: torch.Tensor, - *, - mode: int = 0, - tag: str = "", - use_tma: bool | None = None, + self, + x: torch.Tensor, + *, + mode: int = 0, + tag: str = "", + use_tma: bool | None = None, ) -> torch.Tensor: # COLLECTIVE SEMANTICS: s/n must divide world_size (uniform). The first (shape, mode, use_tma) # seen runs a local micro-benchmark and caches the launch config; every rank MUST issue the SAME @@ -117,17 +119,15 @@ def all_to_all_single_4d( # # tag scopes the symmetric-heap output buffer (reused on same tag+shape+dtype). Results that # must stay live together (e.g. q/k/v) MUST use distinct tags, else they alias one buffer. - return torch.ops.fast_ulysses.all_to_all_single_4d( - self._group, x.contiguous(), mode, tag, use_tma - ) + return torch.ops.fast_ulysses.all_to_all_single_4d(self._group, x.contiguous(), mode, tag, use_tma) def all_to_all_single_4d_async( - self, - x: torch.Tensor, - *, - mode: int = 0, - tag: str = "", - use_tma: bool | None = None, + self, + x: torch.Tensor, + *, + mode: int = 0, + tag: str = "", + use_tma: bool | None = None, ) -> AsyncA2AHandle: # Async variant: launches on the group's comm stream and returns immediately; kernels submitted # to the caller's stream afterwards overlap with the a2a until handle.wait(). Collective @@ -143,9 +143,7 @@ def all_to_all_single_4d_async( x = x.contiguous() out, ev_done = self._launch_on_comm_stream( [x], - lambda: torch.ops.fast_ulysses.all_to_all_single_4d( - self._group, x, mode, tag, use_tma - ), + lambda: torch.ops.fast_ulysses.all_to_all_single_4d(self._group, x, mode, tag, use_tma), ) return AsyncA2AHandle(out, ev_done) diff --git a/lightx2v_fast_ulysses/csrc/ulysses_group.cu b/lightx2v_fast_ulysses/csrc/ulysses_group.cu index 1a15835d3..5df10688c 100644 --- a/lightx2v_fast_ulysses/csrc/ulysses_group.cu +++ b/lightx2v_fast_ulysses/csrc/ulysses_group.cu @@ -76,7 +76,7 @@ std::vector UlyssesGroup::get_uniqueid() return out; } -void UlyssesGroup::init_world(std::vector uid_ints, int64_t global_rank, int64_t global_nranks) +void UlyssesGroup::init_world(std::vector uid_ints, int64_t rank, int64_t nranks) { if (g_world_inited) return; @@ -89,8 +89,7 @@ void UlyssesGroup::init_world(std::vector uid_ints, int64_t global_rank // auto-stamps when version is invalid; here we explicitly substitute that step). nvshmemx_init_attr_t attr = NVSHMEMX_INIT_ATTR_INITIALIZER; TORCH_CHECK( - nvshmemx_set_attr_uniqueid_args(static_cast(global_rank), static_cast(global_nranks), &uid, &attr) - == 0, + nvshmemx_set_attr_uniqueid_args(static_cast(rank), static_cast(nranks), &uid, &attr) == 0, "nvshmemx_set_attr_uniqueid_args failed"); // DEVIATION (see task-5-report): use the host-lib direct entry nvshmemx_hostlib_init_attr instead of // inline nvshmemx_init_attr. The inline version calls nvshmemi_init_thread, a symbol that lives only @@ -228,7 +227,7 @@ void UlyssesGroup::destroy() UlyssesGroup::~UlyssesGroup() { - destroy(); + // Collective cleanup must be explicit. Python synchronizes ranks before calling destroy(). } } // namespace ulysses diff --git a/lightx2v_fast_ulysses/csrc/ulysses_group.cuh b/lightx2v_fast_ulysses/csrc/ulysses_group.cuh index 156d6ba11..c1a173ac1 100644 --- a/lightx2v_fast_ulysses/csrc/ulysses_group.cuh +++ b/lightx2v_fast_ulysses/csrc/ulysses_group.cuh @@ -17,7 +17,7 @@ class UlyssesGroup: public torch::CustomClassHolder { public: static int64_t uniqueid_nints(); // ceil(sizeof(nvshmemx_uniqueid_t)/8) static std::vector get_uniqueid(); // rank0 only - static void init_world(std::vector uid_ints, int64_t global_rank, int64_t global_nranks); // idempotent + static void init_world(std::vector uid_ints, int64_t rank, int64_t nranks); // idempotent UlyssesGroup(std::vector peer_global_pes, int64_t my_rank, int64_t device_id, int64_t reserved_bytes); ~UlyssesGroup() override; diff --git a/lightx2v_fast_ulysses/setup.py b/lightx2v_fast_ulysses/setup.py index b1c4a840a..8e16839ca 100644 --- a/lightx2v_fast_ulysses/setup.py +++ b/lightx2v_fast_ulysses/setup.py @@ -40,9 +40,7 @@ def build_extension(self, ext: Extension) -> None: arch = env.get("CUSTOM_ULYSSES_CUDA_ARCH", "80;90;100") # Torch expects dotted arch (e.g. 86 -> 8.6, 100 -> 10.0); insert the # decimal point before the last digit of each ";"-separated token. - env["TORCH_CUDA_ARCH_LIST"] = " ".join( - f"{tok[:-1]}.{tok[-1]}" for tok in arch.split(";") if tok - ) + env["TORCH_CUDA_ARCH_LIST"] = " ".join(f"{tok[:-1]}.{tok[-1]}" for tok in arch.split(";") if tok) subprocess.check_call( [ "cmake", diff --git a/test_cases/test_fast_ulysses_backend.py b/test_cases/test_fast_ulysses_backend.py index 762ab05f6..a7efe5986 100644 --- a/test_cases/test_fast_ulysses_backend.py +++ b/test_cases/test_fast_ulysses_backend.py @@ -1,7 +1,9 @@ -import unittest +import importlib.util import json +import re import sys import types +import unittest from pathlib import Path from unittest import mock @@ -41,6 +43,55 @@ def apply(self, **kwargs): return kwargs["q"] +class FakeNativeUlyssesGroup: + uid = [17, 23, 42, 99] + init_calls = [] + constructor_calls = [] + + @classmethod + def reset(cls): + cls.init_calls.clear() + cls.constructor_calls.clear() + + @classmethod + def uniqueid_nints(cls): + return len(cls.uid) + + @classmethod + def get_uniqueid(cls): + return list(cls.uid) + + @classmethod + def init_world(cls, uid, rank, nranks): + cls.init_calls.append((list(uid), int(rank), int(nranks))) + + def __init__(self, peer_pes, rank, device_id, initial_pool_bytes): + type(self).constructor_calls.append((list(peer_pes), int(rank), int(device_id), int(initial_pool_bytes))) + self.destroyed = False + + def destroy(self): + self.destroyed = True + + +class FakeCudaStream: + def __init__(self, device=None, priority=0): + self.device = device + self.priority = priority + + @staticmethod + def priority_range(): + return (0, -1) + + +def _load_fast_ulysses_comm_module(): + path = Path(__file__).resolve().parents[1] / "lightx2v_fast_ulysses" / "comm.py" + spec = importlib.util.spec_from_file_location("_lightx2v_fast_ulysses_comm_test", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + class FastUlyssesBackendTest(unittest.TestCase): def tearDown(self): from lightx2v.common.ops.attn.fast_ulysses_attn import FastUlyssesAttnWeight @@ -49,6 +100,7 @@ def tearDown(self): FastUlyssesAttnWeight._single_node_cache.clear() FastUlyssesAttnWeight.reset_runtime_stats() FakeFastUlyssesGroup.instances.clear() + FakeNativeUlyssesGroup.reset() def test_fast_ulysses_backend_is_registered(self): from lightx2v.common.ops import attn # noqa: F401 @@ -322,6 +374,61 @@ def make_group(**kwargs): self.assertEqual(captured[0]["device"], device) self.assertIn(((0, 1), 1), FastUlyssesAttnWeight._groups) + def test_native_group_bootstrap_uses_process_group_scope(self): + comm = _load_fast_ulysses_comm_module() + pg = object() + original_tensor = torch.tensor + broadcasts = [] + barriers = [] + + def get_rank(group=None): + self.assertIs(group, pg) + return 1 + + def get_world_size(group=None): + self.assertIs(group, pg) + return 2 + + def get_process_group_ranks(group): + self.assertIs(group, pg) + return [4, 5] + + def get_global_rank(group, group_rank): + self.assertIs(group, pg) + self.assertEqual(group_rank, 0) + return 4 + + def broadcast(tensor, src, group): + broadcasts.append((src, group, tensor.tolist())) + tensor.copy_(original_tensor(FakeNativeUlyssesGroup.uid, dtype=tensor.dtype)) + + def barrier(group=None): + self.assertIs(group, pg) + barriers.append(group) + + fake_classes = types.SimpleNamespace(fast_ulysses=types.SimpleNamespace(UlyssesGroup=FakeNativeUlyssesGroup)) + with ( + mock.patch.object(comm.dist, "get_rank", side_effect=get_rank), + mock.patch.object(comm.dist, "get_world_size", side_effect=get_world_size), + mock.patch.object(comm.dist, "get_process_group_ranks", side_effect=get_process_group_ranks), + mock.patch.object(comm.dist, "get_global_rank", side_effect=get_global_rank), + mock.patch.object(comm.dist, "broadcast", side_effect=broadcast), + mock.patch.object(comm.dist, "barrier", side_effect=barrier), + mock.patch.object(comm.torch, "classes", fake_classes), + mock.patch.object(comm.torch, "tensor", side_effect=lambda data, dtype=None, device=None: original_tensor(data, dtype=dtype)), + mock.patch.object(comm.torch.cuda, "set_device"), + mock.patch.object(comm.torch.cuda, "Stream", FakeCudaStream), + ): + group = comm.UlyssesGroup(process_group=pg, device=torch.device("cuda", 3), initial_pool_bytes=1024) + + self.assertEqual(group.rank, 1) + self.assertEqual(group.world_size, 2) + self.assertEqual(group.peer_global_ranks, [4, 5]) + self.assertEqual(broadcasts, [(4, pg, [0, 0, 0, 0])]) + self.assertEqual(FakeNativeUlyssesGroup.init_calls, [([17, 23, 42, 99], 1, 2)]) + self.assertEqual(FakeNativeUlyssesGroup.constructor_calls, [([0, 1], 1, 3, 1024)]) + self.assertEqual(barriers, [pg, pg]) + def test_native_package_surface_and_build_are_a2a_only(self): root = Path(__file__).resolve().parents[1] / "lightx2v_fast_ulysses" init_text = (root / "__init__.py").read_text() @@ -339,6 +446,19 @@ def test_native_package_surface_and_build_are_a2a_only(self): self.assertNotIn("all_to_all_qk.cu", cmake_text) self.assertNotIn("qk_norm_rope.cu", cmake_text) + def test_native_destructor_does_not_run_collective_cleanup(self): + root = Path(__file__).resolve().parents[1] / "lightx2v_fast_ulysses" + text = (root / "csrc" / "ulysses_group.cu").read_text() + marker = "UlyssesGroup::~UlyssesGroup()" + self.assertIn(marker, text) + body = text.split(marker, 1)[1].split("\n}", 1)[0] + body_without_comments = re.sub(r"//.*", "", body) + + self.assertNotRegex(body_without_comments, r"\bdestroy\s*\(") + self.assertNotIn("nvshmem_free", body_without_comments) + self.assertNotIn("nvshmem_team_destroy", body_without_comments) + self.assertNotIn("nvshmemx_hostlib_finalize", body_without_comments) + def test_wan_i2v_fast_config_uses_fast_ulysses(self): root = Path(__file__).resolve().parents[1] fast_config = root / "configs" / "dist_infer" / "wan_i2v_dist_cfg_fast_ulysses.json"