From 1d20814c9c624c2bd813941408160ae34c674268 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Thu, 6 Aug 2026 10:08:52 +0000 Subject: [PATCH 1/3] perf(ulysses): overlap Copy Engine transfers with attention prep Add a source-built CUDA IPC Ulysses backend with cached peer target buffers, a high-priority communication stream, and grouped GPU-memory handshakes. Route eligible attention projections through it while retaining NCCL fallbacks. Start MiniMax H3's strided V transfer before Q/K norm and RoPE, then submit tagged Q/K transfers. Add correctness, parity, benchmark, and profiler coverage; refresh parallel and MiniMax H3 documentation with four-H100 measurements. Verification: ruff check; 33 focused unit tests; four-GPU CUDA IPC correctness; MiniMax H3 Ulysses2/Ulysses4/TP2+Ulysses2 parity; four-H100 Flash, Sage, and AdaTaylorCache benchmarks; CUDA trace confirming 157.7 us V-copy/QK-preprocessing overlap. --- docs/en/parallel.md | 14 +- docs/zh/parallel.md | 14 +- examples/minimax_h3/README.md | 34 ++- telefuser/distributed/ulysses_comm.py | 100 +++++++- telefuser/models/flashvsr_dit.py | 6 +- telefuser/models/lingbot_video_dit.py | 21 +- telefuser/models/liveact_dit.py | 6 +- telefuser/models/minimax_h3_dit.py | 14 +- telefuser/models/qwen_image_dit.py | 6 +- telefuser/models/wan_video_dit.py | 12 +- tests/unit/distributed/test_ulysses_comm.py | 58 +++++ tests/unit/models/test_lingbot_video_dit.py | 8 +- tests/unit/models/test_minimax_h3_dit.py | 62 +++++ tf-kernel/CMakeLists.txt | 3 + .../csrc/distributed/ulysses_cuda_ipc.cu | 215 ++++++++++++++++++ .../tests/distributed/ulysses_correctness.py | 94 ++++++++ tf-kernel/tf_kernel/__init__.py | 1 + tf-kernel/tf_kernel/ulysses.py | 213 +++++++++++++++++ .../benchmark_minimax_h3_four_gpu.py | 104 +++++++++ .../profile_minimax_h3_ulysses_overlap.py | 110 +++++++++ 20 files changed, 1049 insertions(+), 46 deletions(-) create mode 100644 tf-kernel/csrc/distributed/ulysses_cuda_ipc.cu create mode 100644 tf-kernel/tests/distributed/ulysses_correctness.py create mode 100644 tf-kernel/tf_kernel/ulysses.py create mode 100644 tools/validation/benchmark_minimax_h3_four_gpu.py create mode 100644 tools/validation/profile_minimax_h3_ulysses_overlap.py diff --git a/docs/en/parallel.md b/docs/en/parallel.md index 9af54f3a..8c3a53a5 100644 --- a/docs/en/parallel.md +++ b/docs/en/parallel.md @@ -88,6 +88,14 @@ Input: (B, S_LOCAL, H_GLOBAL, D) - Suitable for medium-length sequences - Requires number of heads to be divisible by GPU count +When the installed `tf-kernel` wheel contains the Ulysses CUDA IPC operators and every rank in the Ulysses process +group is on the same host, TeleFuser uses the source-built Copy Engine backend for grouped Q/K/V scatter. It writes +directly into each peer final-layout target buffer, caches target allocations by tag/shape/dtype, and fans out over +one high-priority copy stream. Q, K, and V stay as separate submissions so projection compute can overlap +with communication, while the three transfers share one CUDA stream-memory handshake that does not occupy an SM. +Single collectives and output gather stay on the faster PyTorch/NCCL path. Multi-host groups, missing kernels, +and unsupported CUDA IPC configurations also use the PyTorch/NCCL fallback. + ### Ring Attention Sequence parallelism based on P2P communication: @@ -145,9 +153,9 @@ Asynchronous All-to-All implementation, overlapping computation and communicatio ```python # Initiate async All-to-All -q_wait = ulysses_scatter_heads(q, group) -k_wait = ulysses_scatter_heads(k, group) -v_wait = ulysses_scatter_heads(v, group) +q_wait = ulysses_scatter_heads(q, group, tag="q", barrier=False) +k_wait = ulysses_scatter_heads(k, group, tag="k", barrier=False) +v_wait = ulysses_scatter_heads(v, group, tag="v") # Wait for completion q = q_wait() diff --git a/docs/zh/parallel.md b/docs/zh/parallel.md index 2652f51b..6b91905d 100644 --- a/docs/zh/parallel.md +++ b/docs/zh/parallel.md @@ -87,6 +87,14 @@ telefuser/distributed/ - 适合中等长度序列 - 需要头数能被 GPU 数整除 +当已安装的 `tf-kernel` wheel 包含 Ulysses CUDA IPC 算子，且 Ulysses 进程组内所有 rank 位于同一主机时， +TeleFuser 会对成组的 Q/K/V scatter 使用源码编译的 Copy Engine 后端。该后端直接写入对端最终布局的 +target buffer，按 tag/shape/dtype 缓存 target allocation，并使用一条高优先级 copy stream。 +Q、K、V 保持独立提交，因此 projection 计算仍可与通信重叠；三次传输只共享一次不占用 SM 的 CUDA +stream-memory 握手。 +单次 collective 和输出 gather 继续使用实测更快的 PyTorch/NCCL 路径。跨主机进程组、缺少算子或 CUDA IPC +不受支持时也会回退到 PyTorch/NCCL。 + ### Ring Attention 基于 P2P 通信的序列并行： @@ -144,9 +152,9 @@ config = ParallelConfig( ```python # 发起异步 All-to-All -q_wait = ulysses_scatter_heads(q, group) -k_wait = ulysses_scatter_heads(k, group) -v_wait = ulysses_scatter_heads(v, group) +q_wait = ulysses_scatter_heads(q, group, tag="q", barrier=False) +k_wait = ulysses_scatter_heads(k, group, tag="k", barrier=False) +v_wait = ulysses_scatter_heads(v, group, tag="v") # 等待完成 q = q_wait() diff --git a/examples/minimax_h3/README.md b/examples/minimax_h3/README.md index ea87bc21..9ac028b8 100644 --- a/examples/minimax_h3/README.md +++ b/examples/minimax_h3/README.md @@ -355,13 +355,29 @@ encoding, DiT, video/audio decode, host materialization, and orchestration; it e and MP4 encoding. Wall time surrounds the same `run()` call. Full-device memory is sampled from `nvidia-smi` every 100 ms during the measured request. +The measurements use PyTorch 2.11.0, CUDA 12.8, NCCL 2.28.9, and the source-built SM90 `tf-kernel` wheel from this +checkout. Reproduce the three rows from the repository root: + +```bash +python -m tools.validation.benchmark_minimax_h3_four_gpu \ + --attention FLASH_ATTN_4 \ + --output /tmp/minimax_h3_flash.json +python -m tools.validation.benchmark_minimax_h3_four_gpu \ + --attention SAGE_ATTN_2_8_8_SM90 \ + --output /tmp/minimax_h3_sage.json +python -m tools.validation.benchmark_minimax_h3_four_gpu \ + --attention FLASH_ATTN_4 \ + --feature-cache \ + --output /tmp/minimax_h3_flash_cache.json +``` + | Attention | Feature cache | Computed / skipped DiT calls | Pipeline time | Wall time | DiT time | Pipeline speedup | Peak memory GPU 0 / 1 / 2 / 3 | |---|---|---:|---:|---:|---:|---:|---:| -| FlashAttention 4 | Disabled | 49 / 0 | 79.10 s | 79.36 s | 76.48 s | 1.00x | 52.90 / 51.65 / 51.49 / 51.66 GiB | -| SageAttention 2_8_8 SM90 | Disabled | 49 / 0 | 75.96 s | 76.30 s | 73.37 s | 1.04x | 51.49 / 50.11 / 50.10 / 50.04 GiB | -| FlashAttention 4 | AdaTaylorCache | 26 / 23 | 43.53 s | 43.94 s | 40.87 s | 1.82x | 53.00 / 51.40 / 51.22 / 51.22 GiB | +| FlashAttention 4 | Disabled | 49 / 0 | 77.32 s | 77.63 s | 74.64 s | 1.00x | 51.48 / 50.28 / 50.30 / 50.28 GiB | +| SageAttention 2_8_8 SM90 | Disabled | 49 / 0 | 72.41 s | 72.72 s | 69.85 s | 1.07x | 51.42 / 50.20 / 50.24 / 50.24 GiB | +| FlashAttention 4 | AdaTaylorCache | 26 / 23 | 42.39 s | 42.68 s | 39.82 s | 1.82x | 52.73 / 52.09 / 51.56 / 51.54 GiB | -Sage SM90 reduces pipeline latency by 3.98% and DiT latency by 4.07%. It is approximate and remains an H100 opt-in: +Sage SM90 reduces pipeline latency by 6.34% and DiT latency by 6.43%. It is approximate and remains an H100 opt-in: ```bash python -m examples.minimax_h3.minimax_h3_fl2va_h100 \ @@ -375,9 +391,15 @@ Against the FlashAttention 4 output from the same seed, the Sage run measured vi 0.7683, and audio cosine similarity 0.98505. Review generated quality for the target workload before selecting it in production; FlashAttention 4 remains the default. -AdaTaylorCache reduces steady-state pipeline latency by 45.0% and increases maximum single-GPU occupancy by -0.10 GiB (0.2%) in these measurements. Against the previously matched uncached MP4, PSNR is 26.91, SSIM is 0.8619, +AdaTaylorCache reduces steady-state pipeline latency by 45.2% and increases maximum single-GPU occupancy by +1.25 GiB (2.4%) in these measurements. Against the previously matched uncached MP4, PSNR is 26.91, SSIM is 0.8619, audio cosine similarity is 0.9562, and audio duration is unchanged. The earlier matched local SGLang SP2+TP2 parity run measured 79.37 seconds and 67.8 GiB on GPU 0 under the same request shape. +With the source-built tf-kernel available, MiniMax H3 uses the direct CUDA IPC Copy Engine Ulysses scatter. The +fused-QKV projection is passed as a strided V view first; Q/K normalization and RoPE then overlap that transfer, +followed by tagged Q/K transfers and one shared GPU-memory handshake. The three destination buffers are cached per +Ulysses group, so this avoids a QKV packing copy and repeated target allocation. If the optional backend is +unavailable, the same calls fall back to NCCL. + These numbers describe this request and environment, not a general performance or quality guarantee. diff --git a/telefuser/distributed/ulysses_comm.py b/telefuser/distributed/ulysses_comm.py index 364d27fb..8c8635a2 100644 --- a/telefuser/distributed/ulysses_comm.py +++ b/telefuser/distributed/ulysses_comm.py @@ -6,12 +6,17 @@ from __future__ import annotations +import logging from collections.abc import Callable +from typing import Any import torch import torch.distributed as dist import torch.distributed._functional_collectives as fc +logger = logging.getLogger(__name__) +_cuda_ipc_groups: dict[int, Any | None] = {} + def _get_distributed_info(process_group: dist.ProcessGroup) -> tuple[int, int]: """Return the process-group rank and world size.""" @@ -35,17 +40,67 @@ def _wait_async_tensor(tensor: torch.Tensor) -> torch.Tensor: return tensor +def _get_cuda_ipc_group(tensor: torch.Tensor, process_group: dist.ProcessGroup) -> Any | None: + """Lazily create the optional same-host tf-kernel communication backend.""" + if not tensor.is_cuda or torch.compiler.is_compiling(): + return None + key = id(process_group) + if key in _cuda_ipc_groups: + return _cuda_ipc_groups[key] + try: + from tf_kernel.ulysses import CudaIpcUlyssesGroup + + if not CudaIpcUlyssesGroup.is_available(): + _cuda_ipc_groups[key] = None + else: + _cuda_ipc_groups[key] = CudaIpcUlyssesGroup(process_group, tensor.device) + except (ImportError, NotImplementedError, RuntimeError) as error: + logger.debug("CUDA IPC Ulysses is unavailable; using NCCL: %s", error) + _cuda_ipc_groups[key] = None + return _cuda_ipc_groups[key] + + +def _disable_cuda_ipc_group(process_group: dist.ProcessGroup, error: Exception) -> None: + key = id(process_group) + group = _cuda_ipc_groups.get(key) + if group is not None: + group.close() + _cuda_ipc_groups[key] = None + logger.warning("CUDA IPC Ulysses failed; falling back to NCCL: %s", error) + + def ulysses_scatter_heads( tensor: torch.Tensor, process_group: dist.ProcessGroup, *, async_comm: bool = True, + tag: str | None = None, + barrier: bool = True, ) -> Callable[[], torch.Tensor]: - """Scatter global heads and gather sequence across Ulysses ranks.""" + """Scatter global heads and gather sequence across Ulysses ranks. + + A tagged call with ``barrier=False`` starts a same-host Copy Engine group. + Tagged calls share completion until the final ``barrier=True`` call; standalone + calls remain on NCCL. + """ _, world_size = _get_distributed_info(process_group) batch, local_seq_len, num_heads, head_dim = tensor.shape local_heads = _local_head_count(num_heads, world_size) + existing_group = _cuda_ipc_groups.get(id(process_group)) + use_cuda_ipc = ( + async_comm + and tag is not None + and (not barrier or (existing_group is not None and existing_group.has_pending_group)) + ) + cuda_ipc_group = _get_cuda_ipc_group(tensor, process_group) if use_cuda_ipc else None + if cuda_ipc_group is not None: + try: + handle = cuda_ipc_group.all_to_all_single_4d_async(tensor, mode=0, tag=tag, barrier=barrier) + return handle.wait + except RuntimeError as error: + _disable_cuda_ipc_group(process_group, error) + tensor = tensor.reshape(batch, local_seq_len, world_size, local_heads, head_dim) tensor = tensor.permute(2, 1, 0, 3, 4).contiguous() comm_buffer_shape = tensor.shape @@ -84,6 +139,16 @@ def ulysses_gather_heads( expected_local_heads = _local_head_count(num_heads, world_size) if local_heads != expected_local_heads: raise ValueError(f"Ulysses local head count must be {expected_local_heads}, got {local_heads}") + + if _can_use_destination_major_kernel(tensor): + return _gather_heads_destination_major( + tensor, + process_group, + world_size=world_size, + num_heads=num_heads, + async_comm=async_comm, + ) + local_seq_len = global_seq_len // world_size tensor = tensor.reshape(batch, world_size, local_seq_len, local_heads, head_dim) @@ -133,7 +198,7 @@ def ulysses_scatter_qkv( batch, local_seq_len, num_heads, head_dim = query.shape local_heads = _local_head_count(num_heads, world_size) if batch != 1 or not _can_use_destination_major_kernel(query, key, value): - combined_wait = ulysses_scatter_heads(torch.cat((query, key, value), dim=-1), process_group) + combined_wait = ulysses_scatter_heads(torch.cat((query, key, value), dim=-1), process_group, tag="qkv") def fallback_wait() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: return combined_wait().chunk(3, dim=-1) @@ -172,15 +237,42 @@ def ulysses_gather_heads_destination_major( if not _can_use_destination_major_kernel(tensor): return ulysses_gather_heads(tensor, process_group, num_heads=num_heads) + return _gather_heads_destination_major( + tensor, + process_group, + world_size=world_size, + num_heads=num_heads, + async_comm=True, + ) + + +def _gather_heads_destination_major( + tensor: torch.Tensor, + process_group: dist.ProcessGroup, + *, + world_size: int, + num_heads: int, + async_comm: bool, +) -> Callable[[], torch.Tensor]: + """Submit sequence-major All-to-All and fuse the received head relayout.""" + from telefuser.kernel.triton.ulysses_relayout import merge_ulysses_heads + batch, global_seq_len, local_heads, head_dim = tensor.shape local_seq_len = global_seq_len // world_size packed = tensor.permute(1, 0, 2, 3).contiguous() output = torch.empty_like(packed.flatten()) - dist.all_to_all_single(output, packed.flatten(), group=process_group) + work = dist.all_to_all_single( + output, + packed.flatten(), + group=process_group, + async_op=async_comm, + ) def wait() -> torch.Tensor: + if work is not None: + work.wait() received = output.reshape(world_size, local_seq_len, batch, local_heads, head_dim) - return merge_ulysses_heads(received).flatten(2, 3) + return merge_ulysses_heads(received).flatten(2, 3).reshape(batch, local_seq_len, num_heads, head_dim) return wait diff --git a/telefuser/models/flashvsr_dit.py b/telefuser/models/flashvsr_dit.py index d2cde3de..b018eaa2 100644 --- a/telefuser/models/flashvsr_dit.py +++ b/telefuser/models/flashvsr_dit.py @@ -112,15 +112,15 @@ def async_usp_forward( sp_group = get_ulysses_group(device_mesh=device_mesh) v = self.v(x) v_4d = rearrange(v, "b s (h d) -> b s h d", h=self.num_heads) - v_wait = ulysses_scatter_heads(v_4d, sp_group) + v_wait = ulysses_scatter_heads(v_4d, sp_group, tag="v", barrier=False) q = self.norm_q(self.q(x)) q = rope_apply(q, freqs, self.num_heads) q_4d = rearrange(q, "b s (h d) -> b s h d", h=self.num_heads) - q_wait = ulysses_scatter_heads(q_4d, sp_group) + q_wait = ulysses_scatter_heads(q_4d, sp_group, tag="q", barrier=False) k = self.norm_k(self.k(x)) k = rope_apply(k, freqs, self.num_heads) k_4d = rearrange(k, "b s (h d) -> b s h d", h=self.num_heads) - k_wait = ulysses_scatter_heads(k_4d, sp_group) + k_wait = ulysses_scatter_heads(k_4d, sp_group, tag="k") q_4d = q_wait() k_4d = k_wait() v_4d = v_wait() diff --git a/telefuser/models/lingbot_video_dit.py b/telefuser/models/lingbot_video_dit.py index 510f3684..72952181 100644 --- a/telefuser/models/lingbot_video_dit.py +++ b/telefuser/models/lingbot_video_dit.py @@ -145,19 +145,24 @@ def forward( packed_sequence_lengths: list[int] | None = None, ) -> torch.Tensor: batch, sequence, _ = hidden_states.shape - query = self.to_q(hidden_states).view(batch, sequence, self.num_heads, self.head_dim) - key = self.to_k(hidden_states).view(batch, sequence, self.num_heads, self.head_dim) - value = self.to_v(hidden_states).view(batch, sequence, self.num_heads, self.head_dim) - query = apply_lingbot_video_complex_rope(self.norm_q(query), rotary_emb) - key = apply_lingbot_video_complex_rope(self.norm_k(key), rotary_emb) group = self.ulysses_group use_ulysses = ( group is not None and dist.is_available() and dist.is_initialized() and dist.get_world_size(group) > 1 ) + + value = self.to_v(hidden_states).view(batch, sequence, self.num_heads, self.head_dim) + if use_ulysses: + value_wait = ulysses_scatter_heads(value, group, tag="v", barrier=False) + + query = self.to_q(hidden_states).view(batch, sequence, self.num_heads, self.head_dim) + query = apply_lingbot_video_complex_rope(self.norm_q(query), rotary_emb) + if use_ulysses: + query_wait = ulysses_scatter_heads(query, group, tag="q", barrier=False) + + key = self.to_k(hidden_states).view(batch, sequence, self.num_heads, self.head_dim) + key = apply_lingbot_video_complex_rope(self.norm_k(key), rotary_emb) if use_ulysses: - query_wait = ulysses_scatter_heads(query, group) - key_wait = ulysses_scatter_heads(key, group) - value_wait = ulysses_scatter_heads(value, group) + key_wait = ulysses_scatter_heads(key, group, tag="k") query = query_wait() key = key_wait() value = value_wait() diff --git a/telefuser/models/liveact_dit.py b/telefuser/models/liveact_dit.py index c5e197d9..5f7de9c3 100644 --- a/telefuser/models/liveact_dit.py +++ b/telefuser/models/liveact_dit.py @@ -372,13 +372,13 @@ def forward_sp( # QKV projection [B, S/N, H, D] q = self.norm_q(self.q(x)).view(b, s, n, d) + q_wait = ulysses_scatter_heads(q, self.ulysses_group, tag="q", barrier=False) k = self.norm_k(self.k(x)).view(b, s, n, d) + k_wait = ulysses_scatter_heads(k, self.ulysses_group, tag="k", barrier=False) v = self.v(x).view(b, s, n, d) # Ulysses scatter heads: [B, S/N, H, D] -> [B, S, H/N, D] - q_wait = ulysses_scatter_heads(q, self.ulysses_group) - k_wait = ulysses_scatter_heads(k, self.ulysses_group) - v_wait = ulysses_scatter_heads(v, self.ulysses_group) + v_wait = ulysses_scatter_heads(v, self.ulysses_group, tag="v") q = q_wait() k = k_wait() v = v_wait() diff --git a/telefuser/models/minimax_h3_dit.py b/telefuser/models/minimax_h3_dit.py index 189d07d3..f29f8ae2 100644 --- a/telefuser/models/minimax_h3_dit.py +++ b/telefuser/models/minimax_h3_dit.py @@ -26,7 +26,7 @@ get_ulysses_world_size, ) from telefuser.distributed.parallel_shard import sequence_parallel_shard, sequence_parallel_unshard -from telefuser.distributed.ulysses_comm import ulysses_gather_heads_destination_major, ulysses_scatter_qkv +from telefuser.distributed.ulysses_comm import ulysses_gather_heads_destination_major, ulysses_scatter_heads from telefuser.feature_cache import AdaTaylorCacheCalibrator, NoOpCache from telefuser.ops import RMSNorm, apply_qk_norm_rope_neox, indexed_gate, indexed_scale_shift, silu_and_mul_reuse_input from telefuser.ops.attention import attention @@ -468,6 +468,10 @@ def forward( sequence, _ = hidden.shape qkv = self.qkv_proj(hidden).reshape(sequence, 3, self.num_heads, self.head_dim) query, key, value = qkv.unbind(dim=1) + group = self.ulysses_group + use_ulysses = group is not None and dist.get_world_size(group) > 1 + if use_ulysses: + value_wait = ulysses_scatter_heads(value.unsqueeze(0), group, tag="v", barrier=False) if rope_cos_sin_cache is not None: query, key = apply_qk_norm_rope_neox( query, @@ -483,10 +487,12 @@ def forward( query = query.unsqueeze(0) key = key.unsqueeze(0) value = value.unsqueeze(0) - group = self.ulysses_group - use_ulysses = group is not None and dist.get_world_size(group) > 1 if use_ulysses: - query, key, value = ulysses_scatter_qkv(query, key, value, group)() + query_wait = ulysses_scatter_heads(query, group, tag="q", barrier=False) + key_wait = ulysses_scatter_heads(key, group, tag="k") + query = query_wait() + key = key_wait() + value = value_wait() if attention_config is not None and attention_config.attn_impl == AttnImplType.SAGE_ATTN_2_8_8_SM90: total_tokens = query.shape[1] live_tokens = self._sage_live_tokens(sequence_lengths, total_tokens) diff --git a/telefuser/models/qwen_image_dit.py b/telefuser/models/qwen_image_dit.py index 02bfbc45..19c29793 100644 --- a/telefuser/models/qwen_image_dit.py +++ b/telefuser/models/qwen_image_dit.py @@ -351,7 +351,7 @@ def async_usp_forward( img_value = img_value.unflatten(-1, (self.num_heads, -1)) txt_value = txt_value.unflatten(-1, (self.num_heads, -1)) joint_value = torch.cat([txt_value, img_value], dim=1) - joint_value_wait = ulysses_scatter_heads(joint_value, group) + joint_value_wait = ulysses_scatter_heads(joint_value, group, tag="v", barrier=False) img_query = self.to_q(image) txt_query = self.add_q_proj(text) @@ -366,7 +366,7 @@ def async_usp_forward( img_query = apply_rotary_emb_qwen(img_query, img_freqs) txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs) joint_query = torch.cat([txt_query, img_query], dim=1) - joint_query_wait = ulysses_scatter_heads(joint_query, group) + joint_query_wait = ulysses_scatter_heads(joint_query, group, tag="q", barrier=False) img_key = self.to_k(image) txt_key = self.add_k_proj(text) @@ -381,7 +381,7 @@ def async_usp_forward( img_key = apply_rotary_emb_qwen(img_key, img_freqs) txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs) joint_key = torch.cat([txt_key, img_key], dim=1) - joint_key_wait = ulysses_scatter_heads(joint_key, group) + joint_key_wait = ulysses_scatter_heads(joint_key, group, tag="k") joint_value = joint_value_wait() joint_query = joint_query_wait() diff --git a/telefuser/models/wan_video_dit.py b/telefuser/models/wan_video_dit.py index 9048055b..bfbd872c 100755 --- a/telefuser/models/wan_video_dit.py +++ b/telefuser/models/wan_video_dit.py @@ -130,20 +130,20 @@ def async_usp_forward( ) -> torch.Tensor: """Async Ulysses-style sequence parallel forward.""" group = get_ulysses_group(device_mesh) + v = self.v(x) + v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads) + v_wait = ulysses_scatter_heads(v, group, tag="v", barrier=False) q = self.norm_q(self.q(x)) q = rope_apply(q, freqs_cos, freqs_sin, self.num_heads) q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads) - q_wait = ulysses_scatter_heads(q, group) + q_wait = ulysses_scatter_heads(q, group, tag="q", barrier=False) k = self.norm_k(self.k(x)) k = rope_apply(k, freqs_cos, freqs_sin, self.num_heads) k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads) - k_wait = ulysses_scatter_heads(k, group) - v = self.v(x) - v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads) - v_wait = ulysses_scatter_heads(v, group) + k_wait = ulysses_scatter_heads(k, group, tag="k") + v = v_wait() q = q_wait() k = k_wait() - v = v_wait() if sparse_state is not None and sparse_state.config.sparse_impl == "radial": seqlen = q.shape[2] q = rearrange(q, "b n s d -> (b s) n d", s=seqlen, n=self.num_heads) diff --git a/tests/unit/distributed/test_ulysses_comm.py b/tests/unit/distributed/test_ulysses_comm.py index a3018573..1fd7e5a7 100644 --- a/tests/unit/distributed/test_ulysses_comm.py +++ b/tests/unit/distributed/test_ulysses_comm.py @@ -70,3 +70,61 @@ def test_local_head_count_requires_even_partition() -> None: assert _local_head_count(32, 4) == 8 with pytest.raises(ValueError, match="divisible"): _local_head_count(30, 4) + + +@patch("telefuser.distributed.ulysses_comm.dist.get_world_size", return_value=4) +@patch("telefuser.distributed.ulysses_comm.dist.get_rank", return_value=0) +def test_scatter_routes_tag_and_grouped_barrier_to_cuda_ipc(mock_rank, mock_world_size) -> None: + del mock_rank, mock_world_size + from telefuser.distributed.ulysses_comm import ulysses_scatter_heads + + tensor = torch.randn(2, 10, 32, 64) + output = torch.randn(2, 40, 8, 64) + handle = MagicMock() + handle.wait.return_value = output + backend = MagicMock() + backend.all_to_all_single_4d_async.return_value = handle + + with patch("telefuser.distributed.ulysses_comm._get_cuda_ipc_group", return_value=backend): + wait = ulysses_scatter_heads(tensor, MagicMock(), tag="q", barrier=False) + + backend.all_to_all_single_4d_async.assert_called_once_with(tensor, mode=0, tag="q", barrier=False) + assert wait() is output + + +@patch("telefuser.distributed.ulysses_comm.dist.get_world_size", return_value=4) +@patch("telefuser.distributed.ulysses_comm.dist.get_rank", return_value=0) +def test_final_grouped_scatter_reuses_pending_cuda_ipc_group(mock_rank, mock_world_size) -> None: + del mock_rank, mock_world_size + from telefuser.distributed import ulysses_comm + + tensor = torch.randn(2, 10, 32, 64) + process_group = MagicMock() + backend = MagicMock(has_pending_group=True) + backend.all_to_all_single_4d_async.return_value.wait.return_value = tensor + + with ( + patch.dict(ulysses_comm._cuda_ipc_groups, {id(process_group): backend}, clear=True), + patch.object(ulysses_comm, "_get_cuda_ipc_group", return_value=backend), + ): + wait = ulysses_comm.ulysses_scatter_heads(tensor, process_group, tag="v") + + backend.all_to_all_single_4d_async.assert_called_once_with(tensor, mode=0, tag="v", barrier=True) + assert wait() is tensor + + +@patch("telefuser.distributed.ulysses_comm.dist.get_world_size", return_value=4) +@patch("telefuser.distributed.ulysses_comm.dist.get_rank", return_value=0) +def test_standalone_scatter_stays_on_nccl(mock_rank, mock_world_size) -> None: + del mock_rank, mock_world_size + from telefuser.distributed import ulysses_comm + + tensor = torch.randn(2, 10, 32, 64) + with ( + patch.object(ulysses_comm, "_get_cuda_ipc_group") as get_backend, + patch.object(ulysses_comm.fc, "all_to_all_single", return_value=tensor.flatten()), + ): + wait = ulysses_comm.ulysses_scatter_heads(tensor, MagicMock(), tag="qkv") + + get_backend.assert_not_called() + assert wait().shape == (2, 40, 8, 64) diff --git a/tests/unit/models/test_lingbot_video_dit.py b/tests/unit/models/test_lingbot_video_dit.py index 75c28d90..cb88759d 100644 --- a/tests/unit/models/test_lingbot_video_dit.py +++ b/tests/unit/models/test_lingbot_video_dit.py @@ -63,12 +63,14 @@ def test_ulysses_submits_all_qkv_collectives_before_waiting() -> None: events: list[str] = [] submit_index = 0 - def submit(tensor: torch.Tensor, group: object): + def submit(tensor: torch.Tensor, group: object, *, tag: str, barrier: bool = True): nonlocal submit_index del group - name = ("q", "k", "v")[submit_index] + name = ("v", "q", "k")[submit_index] submit_index += 1 events.append(f"submit-{name}") + assert tag == name + assert barrier is (name == "k") def wait() -> torch.Tensor: events.append(f"wait-{name}") @@ -94,7 +96,7 @@ def wait() -> torch.Tensor: ): module(hidden_states, rotary) - assert events == ["submit-q", "submit-k", "submit-v", "wait-q", "wait-k", "wait-v", "submit-output", "wait-output"] + assert events == ["submit-v", "submit-q", "submit-k", "wait-q", "wait-k", "wait-v", "submit-output", "wait-output"] def test_transformer_rejects_non_patch_aligned_latent_geometry() -> None: diff --git a/tests/unit/models/test_minimax_h3_dit.py b/tests/unit/models/test_minimax_h3_dit.py index 9c7f0e08..09234369 100644 --- a/tests/unit/models/test_minimax_h3_dit.py +++ b/tests/unit/models/test_minimax_h3_dit.py @@ -159,6 +159,68 @@ def sage_output(query: torch.Tensor, *_: torch.Tensor, **__: object) -> torch.Te assert torch.count_nonzero(output[61:]) == 0 +def test_ulysses_overlaps_strided_value_scatter_with_qk_preprocessing() -> None: + module = MiniMaxH3Attention(_small_config()).eval() + module.ulysses_group = MagicMock() + hidden = torch.randn(8, 32, dtype=torch.bfloat16) + events: list[tuple[object, ...]] = [] + q_norm = module.q_norm.forward + k_norm = module.k_norm.forward + + def scatter( + tensor: torch.Tensor, + _group: object, + *, + tag: str, + barrier: bool = True, + ) -> object: + events.append(("submit", tag, barrier, tensor.is_contiguous())) + + def wait() -> torch.Tensor: + events.append(("wait", tag)) + return tensor + + return wait + + def record_q_norm(tensor: torch.Tensor) -> torch.Tensor: + events.append(("norm", "q")) + return q_norm(tensor) + + def record_k_norm(tensor: torch.Tensor) -> torch.Tensor: + events.append(("norm", "k")) + return k_norm(tensor) + + with ( + patch("telefuser.models.minimax_h3_dit.dist.get_world_size", return_value=2), + patch("telefuser.models.minimax_h3_dit.ulysses_scatter_heads", side_effect=scatter), + patch( + "telefuser.models.minimax_h3_dit.ulysses_gather_heads_destination_major", + side_effect=lambda tensor, *_args, **_kwargs: lambda: tensor, + ), + patch("telefuser.models.minimax_h3_dit.attention", side_effect=lambda query, *_args, **_kwargs: query), + patch.object(module.q_norm, "forward", side_effect=record_q_norm), + patch.object(module.k_norm, "forward", side_effect=record_k_norm), + ): + output = module( + hidden, + sequence_lengths=[8], + rope_cos_sin_cache=None, + attention_config=AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), + ) + + assert output.shape == hidden.shape + assert events == [ + ("submit", "v", False, False), + ("norm", "q"), + ("norm", "k"), + ("submit", "q", False, True), + ("submit", "k", True, True), + ("wait", "q"), + ("wait", "k"), + ("wait", "v"), + ] + + def test_enable_tp_shards_fused_projections_by_logical_section() -> None: model = MiniMaxH3DiT(_small_config()) block = model.blocks[0] diff --git a/tf-kernel/CMakeLists.txt b/tf-kernel/CMakeLists.txt index af3da267..5d31fdf5 100644 --- a/tf-kernel/CMakeLists.txt +++ b/tf-kernel/CMakeLists.txt @@ -268,6 +268,9 @@ endif() set(SOURCES "csrc/common_extension.cc" + # Distributed communication + "csrc/distributed/ulysses_cuda_ipc.cu" + # Elementwise operations "csrc/elementwise/activation.cu" "csrc/elementwise/cast.cu" diff --git a/tf-kernel/csrc/distributed/ulysses_cuda_ipc.cu b/tf-kernel/csrc/distributed/ulysses_cuda_ipc.cu new file mode 100644 index 00000000..78ec891c --- /dev/null +++ b/tf-kernel/csrc/distributed/ulysses_cuda_ipc.cu @@ -0,0 +1,215 @@ +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace tf_kernel { +namespace { + +constexpr size_t kIpcHandleBytes = sizeof(cudaIpcMemHandle_t); +constexpr size_t kSerializedIpcHandleBytes = kIpcHandleBytes + sizeof(uint64_t); + +void check_driver(CUresult result, const char* operation) { + if (result == CUDA_SUCCESS) { + return; + } + const char* detail = nullptr; + cuGetErrorString(result, &detail); + TORCH_CHECK(false, operation, " failed: ", detail == nullptr ? "unknown CUDA driver error" : detail); +} + +void check_ipc_handle(const torch::Tensor& handle) { + TORCH_CHECK(handle.device().is_cpu(), "CUDA IPC handle must be on CPU"); + TORCH_CHECK(handle.scalar_type() == torch::kUInt8, "CUDA IPC handle must have dtype uint8"); + TORCH_CHECK(handle.is_contiguous(), "CUDA IPC handle must be contiguous"); + TORCH_CHECK(handle.numel() == kSerializedIpcHandleBytes, "invalid CUDA IPC handle size"); +} + +} // namespace + +torch::Tensor cuda_ipc_get_mem_handle(torch::Tensor tensor) { + TORCH_CHECK(tensor.is_cuda(), "CUDA IPC can only export a CUDA tensor"); + c10::cuda::CUDAGuard device_guard(tensor.device()); + + CUdeviceptr allocation_base = 0; + size_t allocation_size = 0; + const auto tensor_pointer = reinterpret_cast(tensor.data_ptr()); + check_driver(cuMemGetAddressRange(&allocation_base, &allocation_size, tensor_pointer), "cuMemGetAddressRange"); + const auto offset = static_cast(tensor_pointer - allocation_base); + TORCH_CHECK(offset < allocation_size, "CUDA tensor pointer is outside its allocation"); + + cudaIpcMemHandle_t handle{}; + C10_CUDA_CHECK(cudaIpcGetMemHandle(&handle, reinterpret_cast(allocation_base))); + auto result = torch::empty( + {static_cast(kSerializedIpcHandleBytes)}, + torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU)); + std::memcpy(result.data_ptr(), &handle, kIpcHandleBytes); + std::memcpy(static_cast(result.data_ptr()) + kIpcHandleBytes, &offset, sizeof(offset)); + return result; +} + +int64_t cuda_ipc_open_mem_handle(torch::Tensor handle) { + check_ipc_handle(handle); + cudaIpcMemHandle_t ipc_handle{}; + uint64_t offset = 0; + std::memcpy(&ipc_handle, handle.data_ptr(), kIpcHandleBytes); + std::memcpy(&offset, static_cast(handle.data_ptr()) + kIpcHandleBytes, sizeof(offset)); + + void* allocation_base = nullptr; + C10_CUDA_CHECK(cudaIpcOpenMemHandle(&allocation_base, ipc_handle, cudaIpcMemLazyEnablePeerAccess)); + const auto pointer = reinterpret_cast(allocation_base) + offset; + return static_cast(pointer); +} + +void cuda_ipc_close_mem_handle(int64_t pointer) { + TORCH_CHECK(pointer != 0, "cannot close a null CUDA IPC pointer"); + CUdeviceptr allocation_base = 0; + size_t allocation_size = 0; + check_driver( + cuMemGetAddressRange( + &allocation_base, &allocation_size, static_cast(static_cast(pointer))), + "cuMemGetAddressRange"); + C10_CUDA_CHECK(cudaIpcCloseMemHandle(reinterpret_cast(allocation_base))); +} + +void ulysses_all_to_all_ce( + torch::Tensor input, + int64_t peer_output_pointer, + int64_t rank, + int64_t world_size, + int64_t mode, + int64_t peer) { + TORCH_CHECK(input.is_cuda(), "Ulysses input must be CUDA"); + TORCH_CHECK(input.dim() == 4, "Ulysses input must be 4D, got ", input.dim(), "D"); + TORCH_CHECK(input.stride(3) == 1, "Ulysses head channels must be contiguous"); + TORCH_CHECK(input.stride(2) == input.size(3), "Ulysses heads must be contiguous within each sequence row"); + TORCH_CHECK( + input.stride(1) >= input.size(2) * input.size(3), + "Ulysses sequence rows must not overlap"); + TORCH_CHECK( + input.size(0) <= 1 || input.stride(0) >= input.size(1) * input.stride(1), + "Ulysses batches must not overlap"); + TORCH_CHECK(world_size > 1, "unsupported Ulysses world size: ", world_size); + TORCH_CHECK(rank >= 0 && rank < world_size, "invalid Ulysses rank: ", rank); + TORCH_CHECK(peer >= 0 && peer < world_size, "invalid Ulysses peer: ", peer); + TORCH_CHECK(mode == 0 || mode == 1, "Ulysses mode must be 0 or 1, got ", mode); + auto* target = reinterpret_cast(static_cast(peer_output_pointer)); + TORCH_CHECK(target != nullptr, "null Ulysses output pointer for peer ", peer); + + c10::cuda::CUDAGuard device_guard(input.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device()); + const size_t element_size = input.element_size(); + const int64_t batch = input.size(0); + const int64_t sequence = input.size(1); + const int64_t heads = input.size(2); + const int64_t head_dim = input.size(3); + auto* source = static_cast(input.data_ptr()); + + if (mode == 0) { + TORCH_CHECK(heads % world_size == 0, "head count must be divisible by the Ulysses world size"); + const int64_t local_heads = heads / world_size; + const size_t width = static_cast(local_heads * head_dim) * element_size; + const size_t source_pitch = static_cast(input.stride(1)) * element_size; + const size_t target_pitch = width; + const int64_t global_sequence = sequence * world_size; + for (int64_t batch_index = 0; batch_index < batch; ++batch_index) { + const int64_t source_elements = + batch_index * input.stride(0) + peer * local_heads * head_dim; + const int64_t target_elements = + (batch_index * global_sequence + rank * sequence) * local_heads * head_dim; + C10_CUDA_CHECK(cudaMemcpy2DAsync( + target + target_elements * element_size, + target_pitch, + source + source_elements * element_size, + source_pitch, + width, + sequence, + cudaMemcpyDefault, + stream)); + } + } else { + TORCH_CHECK(sequence % world_size == 0, "sequence length must be divisible by the Ulysses world size"); + const int64_t local_sequence = sequence / world_size; + const int64_t global_heads = heads * world_size; + const size_t width = static_cast(heads * head_dim) * element_size; + const size_t source_pitch = static_cast(input.stride(1)) * element_size; + const size_t target_pitch = static_cast(global_heads * head_dim) * element_size; + for (int64_t batch_index = 0; batch_index < batch; ++batch_index) { + const int64_t source_elements = + batch_index * input.stride(0) + peer * local_sequence * input.stride(1); + const int64_t target_elements = + batch_index * local_sequence * global_heads * head_dim + rank * heads * head_dim; + C10_CUDA_CHECK(cudaMemcpy2DAsync( + target + target_elements * element_size, + target_pitch, + source + source_elements * element_size, + source_pitch, + width, + local_sequence, + cudaMemcpyDefault, + stream)); + } + } +} + + +void ulysses_stream_barrier( + std::vector peer_barrier_pointers, + torch::Tensor local_barrier, + int64_t rank, + int64_t world_size, + int64_t epoch) { + TORCH_CHECK(local_barrier.is_cuda(), "Ulysses barrier must be CUDA"); + TORCH_CHECK(local_barrier.scalar_type() == torch::kInt64, "Ulysses barrier must have dtype int64"); + TORCH_CHECK(local_barrier.is_contiguous(), "Ulysses barrier must be contiguous"); + TORCH_CHECK(local_barrier.numel() == world_size, "Ulysses barrier has the wrong size"); + TORCH_CHECK(rank >= 0 && rank < world_size, "invalid Ulysses rank: ", rank); + TORCH_CHECK(epoch > 0, "Ulysses barrier epoch must be positive"); + TORCH_CHECK( + peer_barrier_pointers.size() == static_cast(world_size), + "one Ulysses barrier pointer is required per peer"); + + c10::cuda::CUDAGuard device_guard(local_barrier.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(local_barrier.get_device()); + const auto driver_stream = reinterpret_cast(stream); + for (int64_t peer = 0; peer < world_size; ++peer) { + const auto peer_base = static_cast(static_cast(peer_barrier_pointers[peer])); + check_driver( + cuStreamWriteValue64( + driver_stream, + peer_base + rank * sizeof(uint64_t), + static_cast(epoch), + CU_STREAM_WRITE_VALUE_DEFAULT), + "cuStreamWriteValue64"); + } + const auto local_base = reinterpret_cast(local_barrier.data_ptr()); + for (int64_t peer = 0; peer < world_size; ++peer) { + check_driver( + cuStreamWaitValue64( + driver_stream, + local_base + peer * sizeof(uint64_t), + static_cast(epoch), + CU_STREAM_WAIT_VALUE_EQ), + "cuStreamWaitValue64"); + } +} +} // namespace tf_kernel + +TORCH_LIBRARY_FRAGMENT(tf_kernel, m) { + m.def("cuda_ipc_get_mem_handle(Tensor tensor) -> Tensor"); + m.impl("cuda_ipc_get_mem_handle", &tf_kernel::cuda_ipc_get_mem_handle); + m.def("cuda_ipc_open_mem_handle(Tensor handle) -> int"); + m.impl("cuda_ipc_open_mem_handle", &tf_kernel::cuda_ipc_open_mem_handle); + m.def("cuda_ipc_close_mem_handle(int pointer) -> ()"); + m.impl("cuda_ipc_close_mem_handle", &tf_kernel::cuda_ipc_close_mem_handle); + m.def( + "ulysses_all_to_all_ce(Tensor input, int peer_output_pointer, int rank, int world_size, int mode, int peer) -> ()"); + m.impl("ulysses_all_to_all_ce", torch::kCUDA, &tf_kernel::ulysses_all_to_all_ce); + m.def("ulysses_stream_barrier(int[] peer_barrier_pointers, Tensor local_barrier, int rank, int world_size, int epoch) -> ()"); + m.impl("ulysses_stream_barrier", torch::kCUDA, &tf_kernel::ulysses_stream_barrier); +} diff --git a/tf-kernel/tests/distributed/ulysses_correctness.py b/tf-kernel/tests/distributed/ulysses_correctness.py new file mode 100644 index 00000000..f0dc3b4a --- /dev/null +++ b/tf-kernel/tests/distributed/ulysses_correctness.py @@ -0,0 +1,94 @@ +"""Run with: torchrun --standalone --nproc-per-node=4 tests/distributed/ulysses_correctness.py.""" + +from __future__ import annotations + +import os + +import torch +import torch.distributed as dist + +from tf_kernel.ulysses import CudaIpcUlyssesGroup + + +def _input(rank: int, shape: tuple[int, ...], dtype: torch.dtype, offset: int = 0) -> torch.Tensor: + count = 1 + for dimension in shape: + count *= dimension + values = torch.arange(count, device="cuda", dtype=torch.float32).reshape(shape) + return (values.remainder(32) + rank * 64 + offset).to(dtype) + + +def _expected_scatter( + rank: int, + shape: tuple[int, int, int, int], + dtype: torch.dtype, + offset: int = 0, +) -> torch.Tensor: + world_size = dist.get_world_size() + local_heads = shape[2] // world_size + sources = [_input(peer, shape, dtype, offset) for peer in range(world_size)] + pieces = [source[:, :, rank * local_heads : (rank + 1) * local_heads] for source in sources] + return torch.cat(pieces, dim=1) + + +def main() -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank)) + rank = dist.get_rank() + shape = (2, 12, 8, 16) + + group = CudaIpcUlyssesGroup(dist.group.WORLD) + for dtype in (torch.float16, torch.bfloat16): + input = _input(rank, shape, dtype) + packed_qkv = torch.empty( + shape[0], + shape[1], + 3, + shape[2], + shape[3], + dtype=dtype, + device="cuda", + ) + packed_qkv[:, :, 2].copy_(input) + strided_value = packed_qkv[:, :, 2] + assert not strided_value.is_contiguous() + strided = group.all_to_all_single_4d_async(strided_value, mode=0, tag=f"strided-{dtype}").wait() + torch.testing.assert_close(strided, _expected_scatter(rank, shape, dtype), rtol=0, atol=0) + + scattered = group.all_to_all_single_4d_async(input, mode=0, tag=f"scatter-{dtype}").wait() + torch.testing.assert_close(scattered, _expected_scatter(rank, shape, dtype), rtol=0, atol=0) + + restored = group.all_to_all_single_4d_async(scattered, mode=1, tag=f"gather-{dtype}").wait() + torch.testing.assert_close(restored, input, rtol=0, atol=0) + + handles = [ + group.all_to_all_single_4d_async(input + 16, mode=0, tag=f"q-{dtype}", barrier=False), + group.all_to_all_single_4d_async(input + 32, mode=0, tag=f"k-{dtype}", barrier=False), + group.all_to_all_single_4d_async(input + 48, mode=0, tag=f"v-{dtype}"), + ] + outputs = [handle.wait() for handle in handles] + for output, offset in zip(outputs, (16, 32, 48), strict=True): + torch.testing.assert_close(output, _expected_scatter(rank, shape, dtype, offset), rtol=0, atol=0) + + pointers = [output.data_ptr() for output in outputs] + handles = [ + group.all_to_all_single_4d_async(input + 80, mode=0, tag=f"q-{dtype}", barrier=False), + group.all_to_all_single_4d_async(input + 96, mode=0, tag=f"k-{dtype}", barrier=False), + group.all_to_all_single_4d_async(input + 112, mode=0, tag=f"v-{dtype}"), + ] + outputs = [handle.wait() for handle in handles] + assert [output.data_ptr() for output in outputs] == pointers + for output, offset in zip(outputs, (80, 96, 112), strict=True): + torch.testing.assert_close(output, _expected_scatter(rank, shape, dtype, offset), rtol=0, atol=0) + + torch.cuda.synchronize() + dist.barrier() + group.close() + if rank == 0: + print("CUDA IPC Ulysses correctness: PASS") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tf-kernel/tf_kernel/__init__.py b/tf-kernel/tf_kernel/__init__.py index da859c53..ac639735 100644 --- a/tf-kernel/tf_kernel/__init__.py +++ b/tf-kernel/tf_kernel/__init__.py @@ -43,6 +43,7 @@ ) from tf_kernel.memory import set_kv_buffer_kernel, weak_ref_tensor +from tf_kernel.ulysses import AsyncUlyssesHandle, CudaIpcUlyssesGroup from tf_kernel.sageattn2 import (sageattn, sageattn_qk_int8_pv_fp8_cuda, sageattn_qk_int8_pv_fp8_cuda_sm90, sageattn_qk_int8_pv_fp16_cuda, diff --git a/tf-kernel/tf_kernel/ulysses.py b/tf-kernel/tf_kernel/ulysses.py new file mode 100644 index 00000000..063adc06 --- /dev/null +++ b/tf-kernel/tf_kernel/ulysses.py @@ -0,0 +1,213 @@ +"""Single-node Ulysses all-to-all over CUDA IPC and copy engines.""" + +from __future__ import annotations + +import socket +from dataclasses import dataclass + +import torch +import torch.distributed as dist + + +@dataclass +class _Completion: + event: torch.cuda.Event | None = None + + +class AsyncUlyssesHandle: + """GPU-side completion handle for an asynchronous Ulysses all-to-all.""" + + def __init__(self, output: torch.Tensor, completion: _Completion, source: torch.Tensor) -> None: + self._output = output + self._completion = completion + self._source: torch.Tensor | None = source + + def wait(self) -> torch.Tensor: + if self._completion.event is None: + raise RuntimeError("a deferred Ulysses handle was waited before a barrier-carrying call") + torch.cuda.current_stream().wait_event(self._completion.event) + self._source = None + return self._output + + +@dataclass +class _Target: + output: torch.Tensor + peer_outputs: list[int] + remote_pointers: list[int] + + +class CudaIpcUlyssesGroup: + """A same-host process group that writes directly into peer target buffers. + + Target buffers are registered collectively and cached by tag, shape, and dtype. + Transfers run on one high-priority copy-engine stream. Several calls may defer + their stream-memory handshake so Q/K/V share one handshake without being + packed into one tensor. + """ + + def __init__(self, process_group: dist.ProcessGroup, device: torch.device | None = None) -> None: + if not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before CUDA IPC Ulysses") + self.process_group = process_group + self.rank = dist.get_rank(process_group) + self.world_size = dist.get_world_size(process_group) + if self.world_size <= 1: + raise ValueError("CUDA IPC Ulysses requires at least two ranks") + + self.device = torch.device(device or torch.device("cuda", torch.cuda.current_device())) + if self.device.type != "cuda": + raise ValueError(f"CUDA IPC Ulysses requires a CUDA device, got {self.device}") + if self.device.index is None: + self.device = torch.device("cuda", torch.cuda.current_device()) + + hosts: list[str | None] = [None] * self.world_size + dist.all_gather_object(hosts, socket.gethostname(), group=process_group) + if len(set(hosts)) != 1: + raise NotImplementedError("CUDA IPC Ulysses only supports process groups on one host") + + _, greatest_priority = torch.cuda.Stream.priority_range() + self._comm_stream = torch.cuda.Stream(device=self.device, priority=greatest_priority) + self._targets: dict[tuple[object, ...], _Target] = {} + self._barrier = torch.zeros(self.world_size, dtype=torch.int64, device=self.device) + self._peer_barriers, self._barrier_remote_pointers = self._open_peer_handles(self._barrier) + self._barrier_epoch = 0 + self._pending_completion: _Completion | None = None + self._closed = False + + @staticmethod + def is_available() -> bool: + try: + _ = torch.ops.tf_kernel.cuda_ipc_get_mem_handle + _ = torch.ops.tf_kernel.ulysses_all_to_all_ce + _ = torch.ops.tf_kernel.ulysses_stream_barrier + except AttributeError: + return False + return torch.cuda.is_available() + + @property + def has_pending_group(self) -> bool: + """Whether deferred transfers are waiting for the final grouped handshake.""" + return self._pending_completion is not None + + def _all_gather_handle(self, handle: torch.Tensor) -> torch.Tensor: + local = handle.to(device=self.device, non_blocking=False) + gathered = torch.empty(self.world_size * handle.numel(), dtype=torch.uint8, device=self.device) + dist.all_gather_into_tensor(gathered, local, group=self.process_group) + return gathered.reshape(self.world_size, handle.numel()).cpu() + + def _open_peer_handles(self, tensor: torch.Tensor) -> tuple[list[int], list[int]]: + handles = self._all_gather_handle(torch.ops.tf_kernel.cuda_ipc_get_mem_handle(tensor)) + pointers: list[int] = [] + remote_pointers: list[int] = [] + error: Exception | None = None + for peer in range(self.world_size): + if peer == self.rank: + pointers.append(tensor.data_ptr()) + continue + try: + pointer = int(torch.ops.tf_kernel.cuda_ipc_open_mem_handle(handles[peer])) + pointers.append(pointer) + remote_pointers.append(pointer) + except Exception as caught: # all ranks must reach the status collective + error = caught + pointers.append(0) + + status = torch.tensor(0 if error else 1, dtype=torch.int32, device=self.device) + dist.all_reduce(status, op=dist.ReduceOp.MIN, group=self.process_group) + if not status.item(): + for pointer in remote_pointers: + torch.ops.tf_kernel.cuda_ipc_close_mem_handle(pointer) + detail = f": {error}" if error is not None else " on another rank" + raise RuntimeError(f"failed to open CUDA IPC peer memory{detail}") + return pointers, remote_pointers + + def _output_shape(self, input: torch.Tensor, mode: int) -> tuple[int, int, int, int]: + batch, sequence, heads, head_dim = input.shape + if mode == 0: + if heads % self.world_size: + raise ValueError(f"head count {heads} is not divisible by world size {self.world_size}") + return batch, sequence * self.world_size, heads // self.world_size, head_dim + if mode == 1: + if sequence % self.world_size: + raise ValueError(f"sequence length {sequence} is not divisible by world size {self.world_size}") + return batch, sequence // self.world_size, heads * self.world_size, head_dim + raise ValueError(f"Ulysses mode must be 0 or 1, got {mode}") + + def _target(self, input: torch.Tensor, mode: int, tag: str) -> _Target: + output_shape = self._output_shape(input, mode) + key = (tag, mode, output_shape, input.dtype, input.device) + target = self._targets.get(key) + if target is not None: + return target + + output = torch.empty(output_shape, dtype=input.dtype, device=input.device) + peer_outputs, remote_pointers = self._open_peer_handles(output) + target = _Target(output=output, peer_outputs=peer_outputs, remote_pointers=remote_pointers) + self._targets[key] = target + return target + + def all_to_all_single_4d_async( + self, + input: torch.Tensor, + *, + mode: int, + tag: str, + barrier: bool = True, + ) -> AsyncUlyssesHandle: + if self._closed: + raise RuntimeError("CUDA IPC Ulysses group is closed") + if input.device != self.device: + raise ValueError(f"input is on {input.device}, expected {self.device}") + if input.ndim != 4: + raise ValueError(f"Ulysses input must be 4D, got {input.ndim}D") + + target = self._target(input, mode, tag) + caller_stream = torch.cuda.current_stream(self.device) + ready = torch.cuda.Event() + ready.record(caller_stream) + self._comm_stream.wait_event(ready) + + completion = self._pending_completion + if completion is None: + completion = _Completion() + self._pending_completion = completion + + with torch.cuda.stream(self._comm_stream): + for peer in range(self.world_size): + torch.ops.tf_kernel.ulysses_all_to_all_ce( + input, target.peer_outputs[peer], self.rank, self.world_size, mode, peer + ) + + if barrier: + self._barrier_epoch += 1 + with torch.cuda.stream(self._comm_stream): + torch.ops.tf_kernel.ulysses_stream_barrier( + self._peer_barriers, + self._barrier, + self.rank, + self.world_size, + self._barrier_epoch, + ) + completion.event = torch.cuda.Event() + completion.event.record(self._comm_stream) + self._pending_completion = None + return AsyncUlyssesHandle(target.output, completion, input) + + def close(self) -> None: + if self._closed: + return + self._comm_stream.synchronize() + for target in self._targets.values(): + for pointer in target.remote_pointers: + torch.ops.tf_kernel.cuda_ipc_close_mem_handle(pointer) + for pointer in self._barrier_remote_pointers: + torch.ops.tf_kernel.cuda_ipc_close_mem_handle(pointer) + self._targets.clear() + self._closed = True + + def __enter__(self) -> CudaIpcUlyssesGroup: + return self + + def __exit__(self, *_: object) -> None: + self.close() diff --git a/tools/validation/benchmark_minimax_h3_four_gpu.py b/tools/validation/benchmark_minimax_h3_four_gpu.py new file mode 100644 index 00000000..a98ad24b --- /dev/null +++ b/tools/validation/benchmark_minimax_h3_four_gpu.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark the documented warm MiniMax H3 four-GPU profile.""" + +from __future__ import annotations + +import argparse +import gc +import json +import subprocess +import threading +import time +from pathlib import Path + +from examples.minimax_h3.minimax_h3_fl2va_h100 import PPL_CONFIG, get_pipeline, run +from telefuser.core.config import AttnImplType + + +def _sample_device_memory(stop: threading.Event, peaks_mib: list[int]) -> None: + while not stop.is_set(): + output = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,memory.used", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + ).stdout + for line in output.splitlines(): + index, used_mib = (int(value.strip()) for value in line.split(",")) + if index < len(peaks_mib): + peaks_mib[index] = max(peaks_mib[index], used_mib) + stop.wait(0.1) + + +def _generate(pipeline: object) -> object: + return run( + pipeline, + prompt=PPL_CONFIG["prompt"], + seed=0, + aspect_ratio="16:9", + target_video_length=5, + mode="t2va", + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", default=PPL_CONFIG["model_root"]) + parser.add_argument( + "--attention", + choices=("FLASH_ATTN_4", "SAGE_ATTN_2_8_8_SM90"), + default="FLASH_ATTN_4", + ) + parser.add_argument("--feature-cache", action="store_true") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + pipeline = get_pipeline( + 4, + args.model_root, + num_inference_steps=50, + online_adaln_cache=True, + attn_impl=AttnImplType[args.attention], + enable_feature_cache=args.feature_cache, + ) + try: + warmup = _generate(pipeline) + del warmup + gc.collect() + + peaks_mib = [0, 0, 0, 0] + stop_sampling = threading.Event() + sampler = threading.Thread( + target=_sample_device_memory, + args=(stop_sampling, peaks_mib), + daemon=True, + ) + sampler.start() + try: + started = time.perf_counter() + result = _generate(pipeline) + wall_seconds = time.perf_counter() - started + finally: + stop_sampling.set() + sampler.join() + + payload = { + "attention": args.attention, + "feature_cache": args.feature_cache, + "wall_seconds": wall_seconds, + "peak_memory_mib": peaks_mib, + "runtime_metrics": result.runtime_metrics, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(payload, sort_keys=True)) + finally: + pipeline.stop() + + +if __name__ == "__main__": + main() diff --git a/tools/validation/profile_minimax_h3_ulysses_overlap.py b/tools/validation/profile_minimax_h3_ulysses_overlap.py new file mode 100644 index 00000000..155ca99c --- /dev/null +++ b/tools/validation/profile_minimax_h3_ulysses_overlap.py @@ -0,0 +1,110 @@ +"""Export a CUDA trace for the MiniMax H3 Ulysses scatter overlap. + +Run with: + PYTHONPATH=/tmp/tf-kernel-ulysses torchrun --standalone --nproc-per-node=4 \ + tools/validation/profile_minimax_h3_ulysses_overlap.py --output /tmp/h3-ulysses-trace +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path + +import torch +import torch.distributed as dist + +from telefuser.core.config import AttentionConfig, AttnImplType +from telefuser.models import minimax_h3_dit +from telefuser.models.minimax_h3_dit import MiniMaxH3Attention, MiniMaxH3DiTConfig + + +def _config() -> MiniMaxH3DiTConfig: + return MiniMaxH3DiTConfig( + hidden_size=5376, + num_layers=1, + token_refiner_num_layers=1, + num_attention_heads=28, + attention_head_dim=128, + ffn_hidden_size=64, + latents_dim=2, + audio_latents_dim=2, + patch_size=(1, 2, 2), + text_dim=16, + timestep_input_dim=8, + time_embed_hidden_size=32, + time_embed_dim=16, + rope_inv_freq_len=16, + ) + + +def _ulysses_pair(rank: int) -> dist.ProcessGroup: + pairs = ((0, 1), (2, 3)) + groups = [dist.new_group(ranks) for ranks in pairs] + return groups[rank // 2] + + +def _run_attention( + module: MiniMaxH3Attention, + hidden: torch.Tensor, + rope_cos_sin_cache: torch.Tensor, +) -> torch.Tensor: + return module( + hidden, + sequence_lengths=[hidden.shape[0]], + rope_cos_sin_cache=rope_cos_sin_cache, + attention_config=AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--sequence-length", type=int, default=9472) + args = parser.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl", device_id=device) + rank = dist.get_rank() + if dist.get_world_size() != 4: + raise ValueError("this profiler requires four ranks for the Ulysses2 x TP2 shape") + + torch.manual_seed(17 + rank) + module = MiniMaxH3Attention(_config()).eval().to(device) + module.set_ulysses_group(_ulysses_pair(rank)) + hidden = torch.randn(args.sequence_length, module.qkv_proj.in_features, device=device, dtype=torch.bfloat16) + angles = torch.randn(args.sequence_length, 16, device=device, dtype=torch.float32) + rope_cos_sin_cache = torch.cat((angles.cos(), angles.sin()), dim=-1).to(torch.bfloat16) + + # The trace isolates pre-attention communication from the much larger attention kernel. + original_attention = minimax_h3_dit.attention + minimax_h3_dit.attention = lambda query, *_args, **_kwargs: query + try: + _run_attention(module, hidden, rope_cos_sin_cache) + torch.cuda.synchronize(device) + dist.barrier() + + if rank == 0: + args.output.mkdir(parents=True, exist_ok=True) + with torch.profiler.profile( + activities=(torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA), + record_shapes=False, + profile_memory=False, + with_stack=False, + ) as profiler: + _run_attention(module, hidden, rope_cos_sin_cache) + torch.cuda.synchronize(device) + profiler.export_chrome_trace(str(args.output / "minimax_h3_ulysses_rank0.json.gz")) + else: + _run_attention(module, hidden, rope_cos_sin_cache) + torch.cuda.synchronize(device) + dist.barrier() + finally: + minimax_h3_dit.attention = original_attention + dist.destroy_process_group() + + +if __name__ == "__main__": + main() From d8fee47590d6a74f89fd5e59f2b967efa9ed64b2 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Thu, 6 Aug 2026 13:36:00 +0000 Subject: [PATCH 2/3] perf(lingbot): deliver lossless streaming frames Overlap Ulysses Q/K/V communication with attention preparation and profile the resulting execution path. Add a lossless LingBot delivery mode with FIFO backpressure, published-frame completion accounting, and a LiveKit acknowledgement before track teardown. Validate the one-minute four-H100 AIPerf replay and update the LingBot and AIPerf benchmark tables. Verification: ruff format/check; 68 focused TeleFuser tests; 6 AIPerf adapter tests; one-minute four-H100 AIPerf replay (1/1 succeeded). --- .../configs/stream_lingbot_world_v2_1min.json | 1 + .../telefuser_aiperf/adapter.py | 54 +++++++- .../tests/test_livekit_adapter.py | 46 +++++++ docs/en/benchmark_aiperf.md | 37 ++--- examples/lingbot/README.md | 61 ++++++--- telefuser/models/lingbot_world_fast_dit.py | 15 +- .../pipelines/lingbot_world_fast/service.py | 36 ++++- .../pipelines/lingbot_world_fast/session.py | 1 + telefuser/service/api/stream_schema.py | 1 + telefuser/service/livekit/data_protocol.py | 6 +- telefuser/service/livekit/room_client.py | 4 +- telefuser/service/livekit/worker.py | 37 ++++- .../models/test_lingbot_world_fast_dit.py | 34 ++++- .../test_service_action_loop.py | 38 ++++++ .../unit/service/livekit/test_room_client.py | 2 +- tests/unit/service/livekit/test_worker.py | 1 + ...rofile_lingbot_world_v2_ulysses_overlap.py | 129 ++++++++++++++++++ 17 files changed, 450 insertions(+), 53 deletions(-) create mode 100644 tools/validation/profile_lingbot_world_v2_ulysses_overlap.py diff --git a/benchmarks/telefuser_aiperf/configs/stream_lingbot_world_v2_1min.json b/benchmarks/telefuser_aiperf/configs/stream_lingbot_world_v2_1min.json index 0548e908..cb88be23 100644 --- a/benchmarks/telefuser_aiperf/configs/stream_lingbot_world_v2_1min.json +++ b/benchmarks/telefuser_aiperf/configs/stream_lingbot_world_v2_1min.json @@ -19,6 +19,7 @@ "sample_shift": 10.0, "control_mode": "cam", "show_control_hud": false, + "delivery_mode": "lossless", "benchmark_metrics": true }, "transport": { diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/adapter.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/adapter.py index 63bdab99..3cfaf83e 100644 --- a/benchmarks/telefuser_aiperf/telefuser_aiperf/adapter.py +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/adapter.py @@ -70,12 +70,16 @@ def __init__( self.connected = False self.target_ready = False self.active_event = asyncio.Event() + self.target_done_event = asyncio.Event() + self.delivery_complete_event = asyncio.Event() self.first_frame_event = asyncio.Event() self.done_event = asyncio.Event() self.control_task: asyncio.Task[None] | None = None + self.delivery_ack_task: asyncio.Task[None] | None = None self.pending_control_acks: deque[int] = deque() self.pending_control_frames: deque[int] = deque() self.control_sent_at: dict[int, float] = {} + self.expected_published_frames: int | None = None def _active_start(self) -> float: if self.active_started_at is None: @@ -104,8 +108,32 @@ def _handle_video_frame(self) -> None: self.events.record("first_frame") self.last_frame_at = now self._mark_control_frame(now) + self._mark_delivery_complete() self._try_start_active_window() + def _mark_delivery_complete(self) -> None: + if self.expected_published_frames is None or self.result.frames_received < self.expected_published_frames: + return + if not self.result.done_received: + return + if self.delivery_ack_task is None: + self.delivery_ack_task = asyncio.create_task(self._send_delivery_ack()) + + async def _send_delivery_ack(self) -> None: + if self.expected_published_frames is None: + return + try: + await self.room.publish_data( + {"type": "delivery_ack", "published_frames": self.expected_published_frames}, + topic=_CONTROL_TOPIC, + reliable=True, + ) + self.events.record("delivery_ack_sent", published_frames=self.expected_published_frames) + except Exception as exc: # noqa: BLE001 - delivery transport errors are benchmark data + self.result.error = redact_string(f"delivery acknowledgement failed: {exc}") + finally: + self.delivery_complete_event.set() + def _mark_control_frame(self, now: float) -> None: if not self.pending_control_frames: return @@ -121,6 +149,7 @@ def _handle_room_event(self, event: str, payload: Mapping[str, Any]) -> None: self.events.record(event, **dict(payload)) if event == "disconnected": self.done_event.set() + self.target_done_event.set() def _handle_data_message( self, @@ -146,8 +175,15 @@ def _handle_data_message( return if payload.get("type") == "done": self.result.done_received = True + published_frames = payload.get("published_frames") + if isinstance(published_frames, int) and published_frames >= 0: + self.expected_published_frames = published_frames + self._mark_delivery_complete() + if self.expected_published_frames is None: + self.delivery_complete_event.set() + self.target_done_event.set() self.done_event.set() - self.events.record("done_message", topic=topic) + self.events.record("done_message", topic=topic, published_frames=self.expected_published_frames) return data = payload.get("data") if isinstance(payload.get("data"), dict) else payload self._record_transport_profile(data) @@ -155,6 +191,8 @@ def _handle_data_message( error = data.get("error") or payload.get("error") self.result.error = redact_string(str(error)) self.done_event.set() + self.target_done_event.set() + self.delivery_complete_event.set() stage = data.get("stage") if stage is not None: self._handle_status_stage(str(stage), data, now) @@ -314,9 +352,21 @@ async def _wait_for_media(self) -> None: 0.0, ) try: - await asyncio.wait_for(self.done_event.wait(), timeout=remaining) + await asyncio.wait_for(self.target_done_event.wait(), timeout=remaining) except asyncio.TimeoutError: self.events.record("session_duration_elapsed") + return + if self.expected_published_frames is not None: + try: + await asyncio.wait_for( + self.delivery_complete_event.wait(), + timeout=float(self.adapter.options.shutdown_timeout_s), + ) + except asyncio.TimeoutError as exc: + raise TimeoutError( + "Target completed publishing but the client did not receive all frames " + f"({self.result.frames_received}/{self.expected_published_frames})" + ) from exc async def run(self) -> SessionResult: self.started_at = time.perf_counter() diff --git a/benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py b/benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py index bd18913e..377e158e 100644 --- a/benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py +++ b/benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py @@ -373,3 +373,49 @@ async def test_adapter_maps_control_ack_and_next_frame(tmp_path: Path) -> None: assert result.control_events[0].next_frame_latency_ms is not None room = _ControlLiveKitRoom.instances[0] assert room.published[0][1:] == ("tf.control", True) + + +class _IncompleteDeliveryRoom(_FakeLiveKitRoom): + async def connect( + self, + url: str, + token: str, + *, + timeout_s: float, + on_data: Callable[[bytes | str, str, str], None], + on_video_frame: Callable[[], None], + on_event: Callable[[str, Mapping[str, Any]], None], + ) -> None: + self.connected = (url, token) + on_data( + orjson.dumps({"type": "chunk", "data": {"stage": "worker_running"}}), + "tf.status", + "telefuser-worker-0", + ) + on_video_frame() + on_video_frame() + on_data( + orjson.dumps({"type": "done", "session_id": "livekit-session", "published_frames": 3}), + "tf.status", + "telefuser-worker-0", + ) + + +@pytest.mark.asyncio +async def test_adapter_fails_when_client_does_not_receive_target_published_frames(tmp_path: Path) -> None: + _IncompleteDeliveryRoom.instances.clear() + adapter = TeleFuserLiveKitAdapter( + contract=_contract(), + config=_config(tmp_path), + artifacts_dir=tmp_path, + room_client_factory=_IncompleteDeliveryRoom, + http_client=_FakeHttpClient(), + ) + + result = await adapter.run_session(_plan()) + + assert result.success is False + assert result.frames_received == 2 + assert result.done_received is True + assert result.error is not None + assert "2/3" in result.error diff --git a/docs/en/benchmark_aiperf.md b/docs/en/benchmark_aiperf.md index 6d243320..94c1d143 100644 --- a/docs/en/benchmark_aiperf.md +++ b/docs/en/benchmark_aiperf.md @@ -132,29 +132,34 @@ and uses `tools/validation/benchmark_lingbot_world_v2_direct.py`. ### Current one-minute streaming replay -The one-minute workload was rerun on 2026-08-03 at TeleFuser commit -`284996dd616cfd44a55523687b7f2a63a281abb9`. It validates sustained target generation, bounded KV-cache capacity, -and the paced LiveKit delivery path on the current communication-optimized revision. +The one-minute workload was rerun on 2026-08-06 with the current source tree, four H100 80 GB GPUs, Python +3.11.13, PyTorch 2.11.0+cu128, CUDA 12.8, BF16 DiT, FP32 VAE, FlashAttention-4, disabled FSDP, and disabled +`torch.compile`. The source tree includes the tagged Q/K/V Copy Engine Ulysses path. The run used the `stream_lingbot_world_v2_1min.json` workload and AIPerf 0.11.0 at commit -`e977ffbb1648510acec431b2a3fbd1a0f7bb8a35`. The 60-second request was truncated to 60 complete latent chunks: -957 generated frames representing 59.75 seconds of media. With `local_attn_size=18` and `sink_size=6`, the -240-latent-frame session reported a fixed 28,080-token KV capacity. +`e977ffbb1648510acec431b2a3fbd1a0f7bb8a35`. Its `delivery_mode=lossless` request uses FIFO backpressure and +keeps the LiveKit video track open until the AIPerf client confirms the sender declared frame count. The 60-second +request generated 957 frames across 60 complete latent chunks, representing 59.75 seconds of media. The steady +summary excludes the first 13-frame chunk; the 240-latent-frame session reported a fixed 28,080-token KV capacity +with `local_attn_size=18` and `sink_size=6`. | Metric | Result | |---|---:| +| Successful sessions | 1 / 1 | | Generated target frames / chunks | 957 / 60 | | Steady frames / chunks after excluding chunk 0 | 944 / 59 | -| Steady target compute time / FPS | 58.2791 s / **16.1979** | -| Chunk compute mean / p50 / p90 / p99 / max | 0.9878 / 0.9593 / 1.0624 / 1.0932 / 1.1149 s | -| LiveKit stream FPS / client frames | 13.1967 / 803 | -| First client frame / session runtime | 6.0682 / 66.8948 s | -| Runtime creation | 1.4176 s | -| Artifact | `20260803_095518_62ec043c` | - -The target completed all 60 chunks and cleared the average 16 FPS compute gate. It did not keep every chunk below one -second: p99 was 1.0932 seconds and the maximum was 1.1149 seconds. The lower client frame count belongs to the paced -delivery measurement and must not be conflated with target generation completeness. +| Steady target compute time / FPS | 53.6901 s / **17.5824** | +| Chunk compute mean / p50 / p90 / p99 / max | 0.9100 / 0.9088 / 0.9177 / 0.9751 / 1.0549 s | +| LiveKit declared / decoded client frames | 958 / 961 | +| Client callback FPS after first frame | 15.2313 | +| First client frame / session runtime | 5.9822 / 79.2523 s | +| Runtime creation | 1.4279 s | +| Artifact | `20260806_132517_46419f8f` | + +The target completed all 60 chunks and clears the average 16 FPS compute gate. Lossless delivery completed only +after the client confirmed all 958 declared frames. The three additional client decoder callbacks come from the +LiveKit startup track and must not be interpreted as generated model frames. + ## Reproducibility diff --git a/examples/lingbot/README.md b/examples/lingbot/README.md index f7e515d3..62c0424e 100644 --- a/examples/lingbot/README.md +++ b/examples/lingbot/README.md @@ -33,13 +33,13 @@ these versions before attributing a difference to code changes. | GPU | 4 x NVIDIA H100 80 GB HBM3 (SM90) | | NVIDIA driver | `590.48.01` | | Python | `3.11.13` | -| PyTorch | `2.11.0+cu130` | -| PyTorch CUDA runtime | `13.0` | +| PyTorch | `2.11.0+cu128` | +| PyTorch CUDA runtime | `12.8` | | FlashAttention 4 | `flash-attn-4==4.0.0b19` | | CUTLASS DSL | `nvidia-cutlass-dsl==4.6.0` | | CUDA Python | `cuda-python==13.3.1` | -Create an isolated Python 3.11 environment and install the CUDA 13.0 PyTorch build from the wheel index used by +Create an isolated Python 3.11 environment and install the CUDA 12.8 PyTorch build from the wheel index used by your deployment. Install PyTorch before TeleFuser so optional CUDA packages resolve against the intended ABI: ```bash @@ -47,17 +47,16 @@ python3.11 -m venv .venv-lingbot source .venv-lingbot/bin/activate python -m pip install --upgrade pip setuptools wheel -# Install torch==2.11.0+cu130 from your CUDA 13.0 PyTorch wheel index first. +# Install torch==2.11.0+cu128 from your CUDA 12.8 PyTorch wheel index first. python -m pip install -e ".[dev]" python -m pip install \ - "flash-attn-4[cu13]==4.0.0b19" \ + "flash-attn-4==4.0.0b19" \ "nvidia-cutlass-dsl==4.6.0" \ "cuda-python==13.3.1" ``` -The `cu13` extra installs FA4's CUDA 13 dependency variant. For a CUDA 12.8 PyTorch environment, install -`flash-attn-4==4.0.0b19` without that extra and use matching CUDA 12.x dependencies; do not mix cu128 and cu130 -interpreters in one distributed run. +The current AIPerf validation uses the CUDA 12.8 PyTorch build above. For a CUDA 13.0 PyTorch environment, install +`flash-attn-4[cu13]==4.0.0b19` and matching CUDA 13 dependencies; do not mix cu128 and cu130 interpreters in one distributed run. Verify both the package versions and TeleFuser's runtime backend selection before benchmarking: @@ -127,23 +126,22 @@ python examples/lingbot/lingbot_world_v2_image_to_video_h100.py \ ### Validated Four-H100 Real-Time Gate -Commit `540b579` was validated on 2026-08-03 with four H100 80 GB GPUs, PyTorch 2.11.0+cu128, +The direct pipeline-service path was validated with four H100 80 GB GPUs, PyTorch 2.11.0+cu128, FlashAttention-4, BF16 DiT, FP32 VAE, disabled FSDP, and disabled `torch.compile`. The default 832x480 request generated all 77 frames in five chunks at a 16 FPS playback target. -| Metric | Result | -| --- | ---: | -| Steady compute FPS | **17.14** | -| Steady chunk mean / p50 / p90 | 0.9335 / 0.9409 / 0.9410 s | -| Slowest steady chunk | 1.0058 s | -| Generated frames / chunks | 77 / 5 | +| Revision / communication path | Steady compute FPS | Steady chunk mean / p50 / p90 | Slowest steady chunk | Generated frames / chunks | +| --- | ---: | ---: | ---: | ---: | +| `540b579` (2026-08-03) | 17.14 | 0.9335 / 0.9409 / 0.9410 s | 1.0058 s | 77 / 5 | +| Current source, tagged Q/K/V Copy Engine Ulysses (2026-08-06) | **19.08** | **0.8385 / 0.8415 / 0.8706 s** | **0.9040 s** | 77 / 5 | The steady summary excludes chunk 0 and covers four 16-frame chunks. `compute_seconds` synchronizes all target CUDA devices and includes condition handling, DiT, clean-KV update, spatial VAE decode, GPU-to-CPU transfer, and frame conversion. It excludes model loading, runtime creation, LiveKit pacing/encoding, network delivery, and client -rendering. The average therefore clears the 16 FPS target-side real-time gate, while the slowest chunk exceeds its -one-second budget by 5.8 ms; treat this as a validated configuration, not a guarantee for other hardware, resolutions, -durations, concurrent sessions, or transport conditions. +rendering. The current source measurement used the local SM90 `tf-kernel` build with the CUDA IPC Copy Engine backend; +its V transfer is issued before Q projection, Q before K projection, and K is the completion barrier. The current +result clears the 16 FPS target-side real-time gate, but remains a point measurement rather than a guarantee for other +hardware, resolutions, durations, concurrent sessions, or transport conditions. Reproduce the measured direct pipeline-service path without LiveKit or codec time: @@ -162,6 +160,33 @@ The offline CLI was also validated to produce an H.264 832x480 video containing [AIPerf benchmark guide](../../docs/en/benchmark_aiperf.md) for the one-minute workload, client delivery metrics, and comparisons that require identical environments. +### Validated One-Minute AIPerf Replay + +The standard `stream_lingbot_world_v2_1min.json` workload was run on 2026-08-06 with the current source tree, +AIPerf 0.11.0 at `e977ffbb1648510acec431b2a3fbd1a0f7bb8a35`, four H100 80 GB GPUs, and the tagged Q/K/V +Copy Engine Ulysses path. The workload requests `delivery_mode=lossless`: target output uses FIFO backpressure, +the sender declares its published-frame count, and the client confirms receipt before the LiveKit track closes. It +requests 957 generated frames across 60 chunks (59.75 seconds of media), excludes chunk 0, and retains the +28,080-token KV capacity from `local_attn_size=18` and `sink_size=6`. + +| Metric | AIPerf result | +| --- | ---: | +| Successful sessions | 1 / 1 | +| Generated target frames / chunks | 957 / 60 | +| Steady target frames / chunks after warmup | 944 / 59 | +| Steady target compute time / FPS | 53.6901 s / **17.5824** | +| Chunk compute mean / p50 / p90 / p99 / max | 0.9100 / 0.9088 / 0.9177 / 0.9751 / 1.0549 s | +| LiveKit declared / decoded client frames | 958 / 961 | +| Client callback FPS after first frame | 15.2313 | +| First client frame / session runtime | 5.9822 / 79.2523 s | +| Runtime creation | 1.4279 s | +| Artifact | `20260806_132517_46419f8f` | + +The target-side compute rate clears the average 16 FPS gate. Lossless delivery completed after the client confirmed +all 958 declared frames; the three additional decoder callbacks are transport-level callbacks and are separate +from the 957 generated target frames. + + ## Usage ### Four H100 GPUs diff --git a/telefuser/models/lingbot_world_fast_dit.py b/telefuser/models/lingbot_world_fast_dit.py index 109d4976..14325922 100644 --- a/telefuser/models/lingbot_world_fast_dit.py +++ b/telefuser/models/lingbot_world_fast_dit.py @@ -157,15 +157,24 @@ def forward( ) -> torch.Tensor | None: group = get_ulysses_group(device_mesh) ulysses_enabled = group is not None and get_ulysses_world_size(device_mesh) > 1 + + v = rearrange(self.v(x), "b s (n d) -> b s n d", n=self.num_heads) + if ulysses_enabled: + v_wait = ulysses_scatter_heads(v, group, tag="v", barrier=False) + q = rearrange(self.norm_q(self.q(x)), "b s (n d) -> b s n d", n=self.num_heads) + if ulysses_enabled: + q_wait = ulysses_scatter_heads(q, group, tag="q", barrier=False) + k = rearrange(self.norm_k(self.k(x)), "b s (n d) -> b s n d", n=self.num_heads) - v = rearrange(self.v(x), "b s (n d) -> b s n d", n=self.num_heads) - qkv_wait = ulysses_scatter_heads(torch.cat((q, k, v), dim=-1), group) if ulysses_enabled else None + if ulysses_enabled: + k_wait = ulysses_scatter_heads(k, group, tag="k") + frame_tokens = grid_size[1] * grid_size[2] start_frame = current_start // frame_tokens valid_seq_len = math.prod(grid_size) if ulysses_enabled: - q, k, v = qkv_wait().chunk(3, dim=-1) + q, k, v = q_wait(), k_wait(), v_wait() padded_seq_len = q.shape[1] q = self._apply_causal_rope(q, freqs_cos, freqs_sin, grid_size, start_frame, causal_rope)[:, :valid_seq_len] k = self._apply_causal_rope(k, freqs_cos, freqs_sin, grid_size, start_frame, causal_rope)[:, :valid_seq_len] diff --git a/telefuser/pipelines/lingbot_world_fast/service.py b/telefuser/pipelines/lingbot_world_fast/service.py index d2a98055..93668e3d 100644 --- a/telefuser/pipelines/lingbot_world_fast/service.py +++ b/telefuser/pipelines/lingbot_world_fast/service.py @@ -11,6 +11,7 @@ import time import uuid from collections.abc import AsyncGenerator, Callable, Mapping +from concurrent.futures import TimeoutError as FutureTimeoutError from dataclasses import asdict, is_dataclass from pathlib import Path @@ -65,6 +66,7 @@ MAX_GENERATION_SECONDS = 20.0 DEFAULT_OUTPUT_QUEUE_SIZE = 4 _VIDEO_OUTPUT_TYPES = frozenset({"chunk", "preview"}) +_LOSSLESS_OUTPUT_TYPES = _VIDEO_OUTPUT_TYPES | frozenset({"done", "error"}) _TERMINAL_OUTPUT_TYPES = frozenset({"done", "error"}) _MAX_INPUT_IMAGE_BYTES = 10 * 1024 * 1024 _CONTROL_PREFETCH_DEPTH = 0 @@ -317,6 +319,10 @@ def create_session(self, config: dict) -> str: f"got {duration_seconds:g} seconds" ) + delivery_mode = str(config.get("delivery_mode", defaults.get("delivery_mode", "latest"))) + if delivery_mode not in {"latest", "lossless"}: + raise ValueError(f"delivery_mode must be 'latest' or 'lossless', got {delivery_mode!r}") + max_sequence_length = int(config.get("max_sequence_length", defaults.get("max_sequence_length", 512))) if capacity_profile is not None and max_sequence_length > int(capacity_profile["max_sequence_length"]): raise ValueError( @@ -363,6 +369,7 @@ def create_session(self, config: dict) -> str: defaults.get("control_pitch_limit_degrees", 85.0), ) ), + delivery_mode=delivery_mode, show_control_hud=bool(config.get("show_control_hud", defaults.get("show_control_hud", True))), benchmark_metrics=bool(config.get("benchmark_metrics", defaults.get("benchmark_metrics", False))), control_idle_timeout=control_idle_timeout, @@ -435,13 +442,36 @@ def _publish_lease_transitions(self, transitions: tuple[ExecutionLeaseTransition self._put_output(state, payload) @staticmethod - def _put_output(state: LingBotWorldFastSessionState, payload: dict) -> None: + def _put_output(state: LingBotWorldFastSessionState, payload: dict) -> bool: if state.output_queue is None or state.loop is None: - return + return False + + payload_type = str(payload.get("type", "")) + if state.config.delivery_mode == "lossless" and payload_type in _LOSSLESS_OUTPUT_TYPES: + future = asyncio.run_coroutine_threadsafe(state.output_queue.put(payload), state.loop) + while True: + try: + future.result(timeout=0.1) + with state.metrics_lock: + state.output_queue_high_watermark = max( + state.output_queue_high_watermark, + state.output_queue.qsize(), + ) + return True + except FutureTimeoutError: + if not state.active: + future.cancel() + return False + except Exception as exc: + logger.warning(f"Failed to enqueue lossless LingBot output: {exc}") + return False + try: state.loop.call_soon_threadsafe(LingBotWorldFastService._enqueue_output, state, payload) except Exception as exc: logger.warning(f"Failed to enqueue LingBotWorld output: {exc}") + return False + return True @staticmethod def _enqueue_output(state: LingBotWorldFastSessionState, payload: dict) -> None: @@ -1357,7 +1387,6 @@ def _run_worker_loop( }, ) finally: - state.active = False self._lease_manager.deactivate(session_id) try: self._release_generation_session(state) @@ -1385,6 +1414,7 @@ def _run_worker_loop( }, ) self._put_output(state, {"type": "done"}) + state.active = False with self._sessions_lock: if self._sessions.get(session_id) is state: self._sessions.pop(session_id, None) diff --git a/telefuser/pipelines/lingbot_world_fast/session.py b/telefuser/pipelines/lingbot_world_fast/session.py index ccbe3f4a..10f618cc 100644 --- a/telefuser/pipelines/lingbot_world_fast/session.py +++ b/telefuser/pipelines/lingbot_world_fast/session.py @@ -63,6 +63,7 @@ class LingBotWorldFastSessionConfig: control_pitch_step_degrees: float = 2.0 control_pitch_limit_degrees: float = 85.0 show_control_hud: bool = True + delivery_mode: str = "latest" benchmark_metrics: bool = False control_idle_timeout: float = 10.0 diff --git a/telefuser/service/api/stream_schema.py b/telefuser/service/api/stream_schema.py index e908573c..86ba7eda 100644 --- a/telefuser/service/api/stream_schema.py +++ b/telefuser/service/api/stream_schema.py @@ -27,6 +27,7 @@ class StreamDoneMessage(BaseModel): type: str = "done" session_id: str = "" total_chunks: int = 0 + published_frames: int = 0 timestamp: float = Field(default_factory=time.time) diff --git a/telefuser/service/livekit/data_protocol.py b/telefuser/service/livekit/data_protocol.py index e4ba0323..94db2ca9 100644 --- a/telefuser/service/livekit/data_protocol.py +++ b/telefuser/service/livekit/data_protocol.py @@ -10,7 +10,7 @@ TF_METRICS_TOPIC = "tf.metrics" TF_ASSET_TOPIC = "tf.asset" -KNOWN_CONTROL_TYPES = frozenset({"control_state", "control", "prompt", "reset", "stop"}) +KNOWN_CONTROL_TYPES = frozenset({"control_state", "control", "prompt", "reset", "stop", "delivery_ack"}) KNOWN_CONTROLS = frozenset( { "ArrowUp", @@ -146,4 +146,8 @@ def _normalize_legacy_message(message: dict[str, Any]) -> dict[str, Any]: raise DataProtocolError(f"Unsupported control: {control}") if event not in KNOWN_CONTROL_EVENTS: raise DataProtocolError(f"Unsupported control event: {event}") + elif msg_type == "delivery_ack": + published_frames = message.get("published_frames") + if not isinstance(published_frames, int) or published_frames < 0: + raise DataProtocolError("delivery_ack published_frames must be a non-negative integer") return dict(message) diff --git a/telefuser/service/livekit/room_client.py b/telefuser/service/livekit/room_client.py index 06371d0b..528861a2 100644 --- a/telefuser/service/livekit/room_client.py +++ b/telefuser/service/livekit/room_client.py @@ -15,6 +15,8 @@ DataMessageHandler = Callable[[bytes | str | dict[str, Any], str, str], None] _VIDEO_MAX_BITRATE = 8_000_000 +_VIDEO_ENCODER_MIN_MAX_FRAMERATE = 30.0 + class RoomClient(Protocol): """Minimal room operations required by a TeleFuser LiveKit worker.""" @@ -97,7 +99,7 @@ async def publish_video_track(self, name: str, width: int, height: int, *, fps: source=rtc.TrackSource.SOURCE_CAMERA, simulcast=False, video_encoding=rtc.VideoEncoding( - max_framerate=fps, + max_framerate=max(fps, _VIDEO_ENCODER_MIN_MAX_FRAMERATE), max_bitrate=_VIDEO_MAX_BITRATE, ), video_codec=rtc.VideoCodec.VP8, diff --git a/telefuser/service/livekit/worker.py b/telefuser/service/livekit/worker.py index 3dd696d8..42272747 100644 --- a/telefuser/service/livekit/worker.py +++ b/telefuser/service/livekit/worker.py @@ -23,6 +23,9 @@ _ROOM_DISCONNECT_TIMEOUT_SECONDS = 5.0 _CONTROLLER_JOIN_TIMEOUT_SECONDS = 60.0 +_VIDEO_DRAIN_GRACE_SECONDS = 0.5 +_DELIVERY_ACK_TIMEOUT_SECONDS = 15.0 +_VIDEO_TRACK_SUBSCRIPTION_GRACE_SECONDS = 2.0 class WorkerEventSink(Protocol): @@ -79,6 +82,7 @@ def __init__( self.gpu_num = gpu_num self._active_session_id: str | None = None self._pipeline_session_id: str | None = None + self._delivery_ack_event = asyncio.Event() self._stop_event = asyncio.Event() async def start(self, *, skip_validation: bool = False) -> None: @@ -98,6 +102,7 @@ async def run_session(self, record: SessionRecord) -> None: if self._active_session_id is not None: raise RuntimeError(f"Worker {self.worker_id} is already running a session") + self._delivery_ack_event.clear() self._active_session_id = record.session_id self._stop_event.clear() error: str | None = None @@ -142,7 +147,11 @@ async def run_session(self, record: SessionRecord) -> None: }, ).model_dump(mode="json") ) - await self._publish_pipeline_chunks(record.session_id, chunks) + await self._publish_pipeline_chunks( + record.session_id, + chunks, + wait_for_delivery_ack=record.config.get("delivery_mode") == "lossless", + ) except asyncio.CancelledError: error = "cancelled" raise @@ -195,6 +204,10 @@ def _on_data_message( except Exception as exc: logger.warning(f"LiveKit control message rejected: session={record.session_id} error={exc}") return + if chunk.get("type") == "delivery_ack": + self._delivery_ack_event.set() + return + self.pipeline_adapter.push_chunk(self._pipeline_session_id, chunk) if chunk.get("type") == "stop": self._stop_event.set() @@ -203,8 +216,11 @@ async def _publish_pipeline_chunks( self, session_id: str, chunks: AsyncGenerator[dict, None], + *, + wait_for_delivery_ack: bool, ) -> None: chunk_count = 0 + published_frames = 0 next_frame_at: float | None = None async for chunk in chunks: if self._stop_event.is_set(): @@ -223,6 +239,10 @@ async def _publish_pipeline_chunks( if fps <= 0: fps = float(self.config.default_fps) frame_interval = 1.0 / fps + if frames and published_frames == 0 and wait_for_delivery_ack: + height, width = frames[0].shape[:2] + await self.room_client.publish_video_track("telefuser-output", width, height, fps=fps) + await asyncio.sleep(_VIDEO_TRACK_SUBSCRIPTION_GRACE_SECONDS) for frame in frames: if self._stop_event.is_set(): @@ -235,6 +255,7 @@ async def _publish_pipeline_chunks( await asyncio.sleep(delay) await self.room_client.publish_video_frame(frame, fps=fps) next_frame_at += frame_interval + published_frames += 1 if audio is not None: await self.room_client.publish_audio_frame( @@ -266,8 +287,20 @@ async def _publish_pipeline_chunks( ) await self.room_client.publish_status(message.model_dump(mode="json")) - done = StreamDoneMessage(session_id=session_id, total_chunks=chunk_count).model_dump(mode="json") + done = StreamDoneMessage( + session_id=session_id, + total_chunks=chunk_count, + published_frames=published_frames, + ).model_dump(mode="json") await self.room_client.publish_status(done) + if wait_for_delivery_ack and published_frames and not self._stop_event.is_set(): + try: + await asyncio.wait_for(self._delivery_ack_event.wait(), timeout=_DELIVERY_ACK_TIMEOUT_SECONDS) + except TimeoutError: + logger.warning(f"LiveKit delivery acknowledgement timed out: session={session_id}") + elif published_frames and not self._stop_event.is_set(): + # VideoSource accepts frames before the underlying RTP sender has drained them. + await asyncio.sleep(_VIDEO_DRAIN_GRACE_SECONDS) async def _close_active_session(self) -> None: pipeline_session_id = self._pipeline_session_id diff --git a/tests/unit/models/test_lingbot_world_fast_dit.py b/tests/unit/models/test_lingbot_world_fast_dit.py index 2fb6e073..b35c663f 100644 --- a/tests/unit/models/test_lingbot_world_fast_dit.py +++ b/tests/unit/models/test_lingbot_world_fast_dit.py @@ -51,7 +51,7 @@ def fake_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: assert captured["output_layout"] == "BSND" -def test_causal_self_attention_packs_qkv_into_one_ulysses_collective() -> None: +def test_causal_self_attention_overlaps_tagged_qkv_ulysses_collectives() -> None: attention = CausalSelfAttention(dim=32, num_heads=4) freqs = precompute_freqs_cis_3d(8) freqs_cos = torch.cat([freq.real for freq in freqs], dim=-1) @@ -63,8 +63,24 @@ def test_causal_self_attention_packs_qkv_into_one_ulysses_collective() -> None: "local_end_index": 0, } - def fake_scatter(tensor: torch.Tensor, _group: object): - return lambda: tensor + events: list[str] = [] + + def fake_scatter( + tensor: torch.Tensor, + _group: object, + *, + tag: str, + barrier: bool = True, + ): + events.append(f"submit-{tag}") + assert barrier is (tag == "k") + assert tensor.shape == (1, 4, 4, 8) + + def wait() -> torch.Tensor: + events.append(f"wait-{tag}") + return tensor + + return wait def fake_gather(tensor: torch.Tensor, _group: object, *, num_heads: int): assert num_heads == 4 @@ -90,9 +106,15 @@ def fake_gather(tensor: torch.Tensor, _group: object, *, num_heads: int): max_attention_size=4, ) - scatter.assert_called_once() - packed_qkv = scatter.call_args.args[0] - assert packed_qkv.shape == (1, 4, 4, 24) + assert scatter.call_count == 3 + assert events == [ + "submit-v", + "submit-q", + "submit-k", + "wait-q", + "wait-k", + "wait-v", + ] def test_cached_cross_attention_uses_unified_attention_and_bsnd_cache() -> None: diff --git a/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py index 89a80ca9..1e15ed41 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py @@ -906,3 +906,41 @@ def test_service_start_warms_the_pipeline_with_its_default_shape() -> None: assert warmup_config.image.size == (832, 480) assert warmup_config.chunk_size == 4 assert warmup_config.frame_num == 29 + + +def test_lossless_output_queue_blocks_without_evicting_video() -> None: + async def _run() -> None: + service = LingBotWorldFastService(MagicMock()) + state = _state() + state.config.delivery_mode = "lossless" + state.loop = asyncio.get_running_loop() + state.output_queue = asyncio.Queue(maxsize=1) + producer = threading.Thread( + target=lambda: ( + service._put_output(state, {"type": "chunk", "index": 0}), + service._put_output(state, {"type": "chunk", "index": 1}), + ) + ) + producer.start() + await asyncio.sleep(0.01) + assert producer.is_alive() + assert await state.output_queue.get() == {"type": "chunk", "index": 0} + await asyncio.to_thread(producer.join, 1.0) + assert not producer.is_alive() + assert await state.output_queue.get() == {"type": "chunk", "index": 1} + assert LingBotWorldFastService._runtime_metrics(state)["dropped_video_payloads"] == 0 + + asyncio.run(_run()) + + +def test_create_session_validates_delivery_mode() -> None: + pipeline = MagicMock() + service = LingBotWorldFastService(pipeline) + session_id = service.create_session( + {"session_id": "lossless", "image": Image.new("RGB", (8, 8)), "delivery_mode": "lossless"} + ) + assert service._sessions[session_id].config.delivery_mode == "lossless" + service.close_session(session_id) + + with pytest.raises(ValueError, match="delivery_mode"): + service.create_session({"image": Image.new("RGB", (8, 8)), "delivery_mode": "unknown"}) diff --git a/tests/unit/service/livekit/test_room_client.py b/tests/unit/service/livekit/test_room_client.py index d8d66b05..35a7a402 100644 --- a/tests/unit/service/livekit/test_room_client.py +++ b/tests/unit/service/livekit/test_room_client.py @@ -156,7 +156,7 @@ async def _run() -> None: assert captured["source"] == (3, 2) assert video_options.simulcast is False assert video_options.video_codec == "VP8" - assert video_options.video_encoding.max_framerate == 16 + assert video_options.video_encoding.max_framerate == 30 assert video_options.video_encoding.max_bitrate == 8_000_000 assert captured["frame"]["buffer_type"] == "RGB24" assert captured["audio_source"] == (48_000, 1) diff --git a/tests/unit/service/livekit/test_worker.py b/tests/unit/service/livekit/test_worker.py index 82a8ca7d..5208ceb9 100644 --- a/tests/unit/service/livekit/test_worker.py +++ b/tests/unit/service/livekit/test_worker.py @@ -242,6 +242,7 @@ async def _run() -> None: assert room.disconnected is True assert sink.pipeline_sessions == [("session-1", "pipeline-session-1")] assert room.statuses[-1]["total_chunks"] == 2 + assert room.statuses[-1]["published_frames"] == 2 assert sink.finished == [("worker-0", "session-1", None)] asyncio.run(_run()) diff --git a/tools/validation/profile_lingbot_world_v2_ulysses_overlap.py b/tools/validation/profile_lingbot_world_v2_ulysses_overlap.py new file mode 100644 index 00000000..0b5bb782 --- /dev/null +++ b/tools/validation/profile_lingbot_world_v2_ulysses_overlap.py @@ -0,0 +1,129 @@ +"""Export a CUDA trace for LingBot World V2's Ulysses Copy Engine overlap. + +Run with: + PYTHONPATH=/tmp/tf-kernel-ulysses torchrun --standalone --nproc-per-node=4 \ + tools/validation/profile_lingbot_world_v2_ulysses_overlap.py --output /tmp/lingbot-v2-trace +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path + +import torch +import torch.distributed as dist + +from telefuser.models import lingbot_world_fast_dit +from telefuser.models.lingbot_world_fast_dit import CausalSelfAttention +from telefuser.models.wan_video_dit import precompute_freqs_cis_3d + + +def _frequencies(head_dim: int, sequence_length: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + frequencies = precompute_freqs_cis_3d(head_dim, end=sequence_length) + cosine = torch.cat([frequency.real for frequency in frequencies], dim=-1).to(device) + sine = torch.cat([frequency.imag for frequency in frequencies], dim=-1).to(device) + return cosine, sine + + +def _run_attention( + module: CausalSelfAttention, + hidden: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + cache: dict[str, torch.Tensor | int], + grid_size: tuple[int, int, int], +) -> torch.Tensor: + output = module( + hidden, + freqs_cos, + freqs_sin, + grid_size, + cache, + current_start=0, + max_attention_size=grid_size[-1], + device_mesh=object(), + ) + if output is None: + raise RuntimeError("LingBot World V2 attention unexpectedly returned no output") + return output + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--local-sequence-length", type=int, default=1560) + args = parser.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl", device_id=device) + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size != 4: + raise ValueError("this profiler requires four Ulysses ranks") + + torch.manual_seed(17) + module = CausalSelfAttention(dim=2048, num_heads=16).eval().to(device, dtype=torch.bfloat16) + torch.manual_seed(31 + rank) + hidden = torch.randn(args.local_sequence_length, 2048, device=device, dtype=torch.bfloat16).unsqueeze(0) + global_sequence_length = args.local_sequence_length * world_size + grid_size = (1, 1, global_sequence_length) + freqs_cos, freqs_sin = _frequencies(module.head_dim, global_sequence_length, device) + cache = { + "k": torch.zeros( + 1, + global_sequence_length, + module.num_heads // world_size, + module.head_dim, + device=device, + dtype=torch.bfloat16, + ), + "v": torch.zeros( + 1, + global_sequence_length, + module.num_heads // world_size, + module.head_dim, + device=device, + dtype=torch.bfloat16, + ), + "global_end_index": 0, + "local_end_index": 0, + } + + original_group = lingbot_world_fast_dit.get_ulysses_group + original_world_size = lingbot_world_fast_dit.get_ulysses_world_size + original_attention = lingbot_world_fast_dit.attn_func + lingbot_world_fast_dit.get_ulysses_group = lambda _mesh: dist.group.WORLD + lingbot_world_fast_dit.get_ulysses_world_size = lambda _mesh: world_size + lingbot_world_fast_dit.attn_func = lambda query, _key, _value, **_kwargs: query + try: + _run_attention(module, hidden, freqs_cos, freqs_sin, cache, grid_size) + torch.cuda.synchronize(device) + dist.barrier() + + if rank == 0: + args.output.mkdir(parents=True, exist_ok=True) + with torch.profiler.profile( + activities=(torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA), + record_shapes=False, + profile_memory=False, + with_stack=False, + ) as profiler: + _run_attention(module, hidden, freqs_cos, freqs_sin, cache, grid_size) + torch.cuda.synchronize(device) + profiler.export_chrome_trace(str(args.output / "lingbot_world_v2_ulysses_rank0.json.gz")) + else: + _run_attention(module, hidden, freqs_cos, freqs_sin, cache, grid_size) + torch.cuda.synchronize(device) + dist.barrier() + finally: + lingbot_world_fast_dit.get_ulysses_group = original_group + lingbot_world_fast_dit.get_ulysses_world_size = original_world_size + lingbot_world_fast_dit.attn_func = original_attention + dist.destroy_process_group() + + +if __name__ == "__main__": + main() From 819c2388d7bbdc259821be9c6180879643a0c347 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Thu, 6 Aug 2026 16:34:39 +0000 Subject: [PATCH 3/3] fix(ipc): harden CUDA IPC lifecycle Give each tensor-channel consumer an independent CUDA IPC refcounter and let the channel own cleanup of its producer pools and consumer mappings. Close Ulysses IPC groups before distributed teardown, keep child workers on cooperative shutdown, and allow LiveKit sessions a grace period before cancellation. Add unit and multi-GPU coverage for IPC cleanup, refcounter isolation, worker signals, group closure, and session stop behavior. Validate with the 4-GPU LingBot World V2 streaming path. --- docs/en/parallel.md | 5 +- docs/zh/parallel.md | 5 +- telefuser/distributed/ulysses_comm.py | 14 ++++ telefuser/service/livekit/worker_pool.py | 12 +++ telefuser/worker/parallel_worker.py | 15 +++- telefuser/worker/tensor_channel.py | 56 ++++++++++--- .../integration/test_worker_tensor_channel.py | 3 + tests/unit/distributed/test_ulysses_comm.py | 65 +++++++++++++++ .../unit/service/livekit/test_worker_pool.py | 83 +++++++++++++++++++ tests/unit/worker/test_parallel_worker.py | 8 +- tests/unit/worker/test_tensor_channel.py | 16 ++++ .../csrc/distributed/ulysses_cuda_ipc.cu | 4 +- .../tests/distributed/ulysses_correctness.py | 23 +++++ tf-kernel/tf_kernel/ulysses.py | 27 +++++- 14 files changed, 315 insertions(+), 21 deletions(-) create mode 100644 tests/unit/service/livekit/test_worker_pool.py diff --git a/docs/en/parallel.md b/docs/en/parallel.md index 8c3a53a5..8e32c152 100644 --- a/docs/en/parallel.md +++ b/docs/en/parallel.md @@ -90,8 +90,9 @@ Input: (B, S_LOCAL, H_GLOBAL, D) When the installed `tf-kernel` wheel contains the Ulysses CUDA IPC operators and every rank in the Ulysses process group is on the same host, TeleFuser uses the source-built Copy Engine backend for grouped Q/K/V scatter. It writes -directly into each peer final-layout target buffer, caches target allocations by tag/shape/dtype, and fans out over -one high-priority copy stream. Q, K, and V stay as separate submissions so projection compute can overlap +directly into each peer final-layout target buffer, keeps up to 12 target allocations in a tag/shape/dtype LRU cache, +and fans out over one high-priority copy stream. Eviction synchronizes participating devices before closing peer +mappings. Q, K, and V stay as separate submissions so projection compute can overlap with communication, while the three transfers share one CUDA stream-memory handshake that does not occupy an SM. Single collectives and output gather stay on the faster PyTorch/NCCL path. Multi-host groups, missing kernels, and unsupported CUDA IPC configurations also use the PyTorch/NCCL fallback. diff --git a/docs/zh/parallel.md b/docs/zh/parallel.md index 6b91905d..1cba0d69 100644 --- a/docs/zh/parallel.md +++ b/docs/zh/parallel.md @@ -89,8 +89,9 @@ telefuser/distributed/ 当已安装的 `tf-kernel` wheel 包含 Ulysses CUDA IPC 算子，且 Ulysses 进程组内所有 rank 位于同一主机时， TeleFuser 会对成组的 Q/K/V scatter 使用源码编译的 Copy Engine 后端。该后端直接写入对端最终布局的 -target buffer，按 tag/shape/dtype 缓存 target allocation，并使用一条高优先级 copy stream。 -Q、K、V 保持独立提交，因此 projection 计算仍可与通信重叠；三次传输只共享一次不占用 SM 的 CUDA +target buffer，在按 tag/shape/dtype 组织的 LRU 中最多缓存 12 个 target allocation，并使用一条高优先级 +copy stream。淘汰前会同步参与设备，再关闭 peer mapping。Q、K、V 保持独立提交，因此 projection 计算仍 +可与通信重叠；三次传输只共享一次不占用 SM 的 CUDA stream-memory 握手。 单次 collective 和输出 gather 继续使用实测更快的 PyTorch/NCCL 路径。跨主机进程组、缺少算子或 CUDA IPC 不受支持时也会回退到 PyTorch/NCCL。 diff --git a/telefuser/distributed/ulysses_comm.py b/telefuser/distributed/ulysses_comm.py index 8c8635a2..f73445c1 100644 --- a/telefuser/distributed/ulysses_comm.py +++ b/telefuser/distributed/ulysses_comm.py @@ -69,6 +69,14 @@ def _disable_cuda_ipc_group(process_group: dist.ProcessGroup, error: Exception) logger.warning("CUDA IPC Ulysses failed; falling back to NCCL: %s", error) +def _close_cuda_ipc_groups() -> None: + """Collectively close initialized CUDA IPC groups before process-group teardown.""" + groups = [group for group in _cuda_ipc_groups.values() if group is not None] + _cuda_ipc_groups.clear() + for group in groups: + group.close() + + def ulysses_scatter_heads( tensor: torch.Tensor, process_group: dist.ProcessGroup, @@ -95,10 +103,16 @@ def ulysses_scatter_heads( ) cuda_ipc_group = _get_cuda_ipc_group(tensor, process_group) if use_cuda_ipc else None if cuda_ipc_group is not None: + group_was_pending = cuda_ipc_group.has_pending_group try: handle = cuda_ipc_group.all_to_all_single_4d_async(tensor, mode=0, tag=tag, barrier=barrier) return handle.wait except RuntimeError as error: + if group_was_pending or cuda_ipc_group.has_pending_group: + raise RuntimeError( + "CUDA IPC Ulysses failed after a grouped transfer started; " + "the request cannot safely fall back to NCCL" + ) from error _disable_cuda_ipc_group(process_group, error) tensor = tensor.reshape(batch, local_seq_len, world_size, local_heads, head_dim) diff --git a/telefuser/service/livekit/worker_pool.py b/telefuser/service/livekit/worker_pool.py index 9722f58c..71a76210 100644 --- a/telefuser/service/livekit/worker_pool.py +++ b/telefuser/service/livekit/worker_pool.py @@ -10,6 +10,8 @@ from .session_registry import SessionRecord from .worker import LiveKitWorker +_SESSION_STOP_GRACE_SECONDS = 8.0 + class WorkerPool(Protocol): """Worker-pool operations used by the API runtime.""" @@ -59,6 +61,16 @@ async def stop_session(self, session_id: str) -> None: return for worker in self._workers.values(): await worker.stop_session(session_id) + + try: + await asyncio.wait_for(asyncio.shield(task), timeout=_SESSION_STOP_GRACE_SECONDS) + return + except TimeoutError: + logger.warning( + f"LiveKit session did not stop within {_SESSION_STOP_GRACE_SECONDS:g}s; " + f"cancelling runner: session={session_id}" + ) + if not task.done(): task.cancel() try: diff --git a/telefuser/worker/parallel_worker.py b/telefuser/worker/parallel_worker.py index b814e0bd..00f8374f 100644 --- a/telefuser/worker/parallel_worker.py +++ b/telefuser/worker/parallel_worker.py @@ -8,6 +8,7 @@ import gc import os +import signal import threading import time from collections.abc import Callable, Collection @@ -72,8 +73,12 @@ def _worker_loop( from telefuser.utils.profiler import mark_as_worker_process mark_as_worker_process() + previous_sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN) args = None kwargs = None + stage_inputs = None + y = None + clean_shutdown = False try: parallel_config = stage.model_runtime_config.parallel_config # Avoid host-wide launch pools in every spawned CUDA worker, including @@ -123,6 +128,7 @@ def _worker_loop( del data if name == "exit": logger.info(f"parallel worker {stage.name} on rank {rank} exits") + clean_shutdown = True break if name == _DISCARD_TENSOR_REFS: discarded = 0 @@ -169,14 +175,19 @@ def _worker_loop( finally: args = None kwargs = None + stage_inputs = None + y = None current_platform.synchronize() - for channel in tensor_input_channels: - channel.release_local_cuda_ipc() + if clean_shutdown and world_size > 1: + from telefuser.distributed.ulysses_comm import _close_cuda_ipc_groups + + _close_cuda_ipc_groups() gc.collect() current_platform.empty_cache() current_platform.ipc_collect() if world_size > 1: dist.destroy_process_group() + signal.signal(signal.SIGINT, previous_sigint_handler) class ParallelWorker: diff --git a/telefuser/worker/tensor_channel.py b/telefuser/worker/tensor_channel.py index 22f9011f..7243e58f 100644 --- a/telefuser/worker/tensor_channel.py +++ b/telefuser/worker/tensor_channel.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from dataclasses import replace as dataclass_replace from multiprocessing.queues import SimpleQueue +from multiprocessing.util import Finalize from typing import Any import torch @@ -65,7 +66,7 @@ class _CudaIpcSlot: class _CudaIpcPool: profile_index: int storage: torch.Tensor - handle: tuple[Any, ...] + handles: tuple[tuple[Any, ...], ...] slots: list[_CudaIpcSlot] next_slot_id: int = 0 @@ -114,6 +115,7 @@ def __init__( self._cuda_consumer_completion_events: dict[tuple[int, int], torch.cuda.Event] = {} self._cuda_completion_handles: dict[tuple[int, int, int], bytes] = {} self._cuda_producer_completion_events: dict[tuple[int, int, int, int], torch.cuda.Event] = {} + self._cuda_ipc_finalizer: Finalize | None = None self._next_transfer_id = 0 self._producer_bound = False self._consumer_bound = False @@ -271,6 +273,7 @@ def _rank_shape(self, shape: tuple[int, ...], shard_dim: int | None, rank: int) return tuple(rank_shape) def _create_cuda_pool(self, ref: WorkerTensorRef, tensor: torch.Tensor) -> _CudaIpcPool: + self._ensure_local_cuda_ipc_finalizer() with torch.cuda.device(tensor.device): storage = torch.empty( (self.cuda_ipc_slots, *tensor.shape), @@ -292,7 +295,7 @@ def _create_cuda_pool(self, ref: WorkerTensorRef, tensor: torch.Tensor) -> _Cuda pool = _CudaIpcPool( profile_index=len(self._cuda_pools), storage=storage, - handle=storage.untyped_storage()._share_cuda_(), + handles=tuple(storage.untyped_storage()._share_cuda_() for _ in range(self.consumer_world_size)), slots=slots, ) key = (ref.tensor_index, ref.shape, tensor.dtype, str(tensor.device)) @@ -374,14 +377,14 @@ def _stage_cuda_tensor(self, ref: WorkerTensorRef, tensor: torch.Tensor) -> tupl ref=ref, profile_index=pool.profile_index, slot_id=slot_id, - handle=pool.handle, + handle=handle, ready_event_handle=slot.ready_event_handle, byte_offset=view.storage_offset() * view.element_size(), shape=tuple(view.shape), stride=tuple(view.stride()), dtype=view.dtype, ) - for view in rank_views + for view, handle in zip(rank_views, pool.handles, strict=True) ) def _wait_for_cuda_slot_reuse( @@ -446,6 +449,7 @@ def _copy_cuda_payload( rank: int, target: torch.device, ) -> torch.Tensor: + self._ensure_local_cuda_ipc_finalizer() source = torch.device(payload.ref.source_device) rebuild_device = target if target.type == "cuda" else source if rebuild_device.index is None: @@ -486,13 +490,45 @@ def _ack_cuda_payload(self, payload: _CudaIpcPayload, *, rank: int, copied: bool generation if copied else -generation ) + @staticmethod + def _clear_local_cuda_ipc_state(*state: dict[Any, Any]) -> None: + for cache in state: + cache.clear() + + def _ensure_local_cuda_ipc_finalizer(self) -> None: + finalizer = self._cuda_ipc_finalizer + if finalizer is not None and finalizer.still_active(): + return + state = ( + self._cuda_pools, + self._cuda_consumer_completion_events, + self._cuda_completion_handles, + self._cuda_producer_completion_events, + self._cuda_event_cache, + self._cuda_storage_cache, + ) + self._cuda_ipc_finalizer = Finalize( + self, + WorkerTensorChannel._clear_local_cuda_ipc_state, + args=state, + exitpriority=10, + ) + def release_local_cuda_ipc(self) -> None: - """Release process-local CUDA IPC mappings after pending work is synchronized.""" - self._cuda_consumer_completion_events.clear() - self._cuda_completion_handles.clear() - self._cuda_producer_completion_events.clear() - self._cuda_event_cache.clear() - self._cuda_storage_cache.clear() + """Release process-local CUDA IPC state owned by this channel.""" + finalizer = self._cuda_ipc_finalizer + if finalizer is not None and finalizer.still_active(): + finalizer() + else: + self._clear_local_cuda_ipc_state( + self._cuda_pools, + self._cuda_consumer_completion_events, + self._cuda_completion_handles, + self._cuda_producer_completion_events, + self._cuda_event_cache, + self._cuda_storage_cache, + ) + self._cuda_ipc_finalizer = None def discard(self, value: Any, *, rank: int) -> int: """Consume and release referenced tensors without materializing a device copy.""" diff --git a/tests/integration/test_worker_tensor_channel.py b/tests/integration/test_worker_tensor_channel.py index 1e5fb261..348ddf89 100644 --- a/tests/integration/test_worker_tensor_channel.py +++ b/tests/integration/test_worker_tensor_channel.py @@ -123,6 +123,9 @@ def discard() -> None: try: abandoned = producer.reduce(torch.arange(4, dtype=torch.float32), sync=True) + assert abandoned._cuda_payloads is not None + ref_counters = {(payload.handle[4], payload.handle[5]) for payload in abandoned._cuda_payloads} + assert len(ref_counters) == 2 thread = threading.Thread(target=discard) thread.start() assert channel.rank_one_started.wait(timeout=10) diff --git a/tests/unit/distributed/test_ulysses_comm.py b/tests/unit/distributed/test_ulysses_comm.py index 1fd7e5a7..96e8a2af 100644 --- a/tests/unit/distributed/test_ulysses_comm.py +++ b/tests/unit/distributed/test_ulysses_comm.py @@ -128,3 +128,68 @@ def test_standalone_scatter_stays_on_nccl(mock_rank, mock_world_size) -> None: get_backend.assert_not_called() assert wait().shape == (2, 40, 8, 64) + + +class _FailingGroupedBackend: + def __init__(self) -> None: + self.pending = False + + @property + def has_pending_group(self) -> bool: + return self.pending + + def all_to_all_single_4d_async(self, *_args: object, **_kwargs: object) -> None: + self.pending = True + raise RuntimeError("injected grouped failure") + + +@patch("telefuser.distributed.ulysses_comm.dist.get_world_size", return_value=4) +@patch("telefuser.distributed.ulysses_comm.dist.get_rank", return_value=0) +def test_grouped_ipc_failure_does_not_attempt_partial_nccl_fallback(mock_rank, mock_world_size) -> None: + del mock_rank, mock_world_size + from telefuser.distributed import ulysses_comm + + tensor = torch.randn(2, 10, 32, 64) + backend = _FailingGroupedBackend() + with ( + patch.object(ulysses_comm, "_get_cuda_ipc_group", return_value=backend), + patch.object(ulysses_comm.fc, "all_to_all_single") as nccl, + pytest.raises(RuntimeError, match="cannot safely fall back"), + ): + ulysses_comm.ulysses_scatter_heads(tensor, MagicMock(), tag="q", barrier=False) + + nccl.assert_not_called() + + +@patch("telefuser.distributed.ulysses_comm.dist.get_world_size", return_value=4) +@patch("telefuser.distributed.ulysses_comm.dist.get_rank", return_value=0) +def test_ipc_failure_before_group_start_falls_back_to_nccl(mock_rank, mock_world_size) -> None: + del mock_rank, mock_world_size + from telefuser.distributed import ulysses_comm + + tensor = torch.randn(2, 10, 32, 64) + process_group = MagicMock() + backend = MagicMock(has_pending_group=False) + backend.all_to_all_single_4d_async.side_effect = RuntimeError("injected setup failure") + with ( + patch.dict(ulysses_comm._cuda_ipc_groups, {id(process_group): backend}, clear=True), + patch.object(ulysses_comm, "_get_cuda_ipc_group", return_value=backend), + patch.object(ulysses_comm.fc, "all_to_all_single", return_value=tensor.flatten()), + ): + wait = ulysses_comm.ulysses_scatter_heads(tensor, process_group, tag="q", barrier=False) + + backend.close.assert_called_once_with() + assert wait().shape == (2, 40, 8, 64) + + +def test_close_cuda_ipc_groups_closes_initialized_groups_and_clears_cache() -> None: + from telefuser.distributed import ulysses_comm + + first = MagicMock() + second = MagicMock() + with patch.dict(ulysses_comm._cuda_ipc_groups, {1: first, 2: None, 3: second}, clear=True): + ulysses_comm._close_cuda_ipc_groups() + assert ulysses_comm._cuda_ipc_groups == {} + + first.close.assert_called_once_with() + second.close.assert_called_once_with() diff --git a/tests/unit/service/livekit/test_worker_pool.py b/tests/unit/service/livekit/test_worker_pool.py new file mode 100644 index 00000000..130a7272 --- /dev/null +++ b/tests/unit/service/livekit/test_worker_pool.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import asyncio + +from telefuser.service.livekit.session_registry import SessionRecord +from telefuser.service.livekit.worker_pool import InProcessLiveKitWorkerPool + + +class _CooperativeWorker: + def __init__(self, *, complete_on_stop: bool) -> None: + self.complete_on_stop = complete_on_stop + self.running = asyncio.Event() + self.stop_requested = asyncio.Event() + self.completed = False + self.cancelled = False + + async def start(self, *, skip_validation: bool = False) -> None: + return None + + async def run_session(self, record: SessionRecord) -> None: + del record + self.running.set() + try: + await self.stop_requested.wait() + if not self.complete_on_stop: + await asyncio.Event().wait() + self.completed = True + except asyncio.CancelledError: + self.cancelled = True + raise + + async def stop_session(self, session_id: str) -> None: + del session_id + self.stop_requested.set() + + async def stop(self) -> None: + return None + + +def _record() -> SessionRecord: + return SessionRecord( + session_id="session-1", + room_name="room-1", + controller_identity="controller-1", + status="assigned", + worker_id="worker-0", + config={}, + created_at=0, + updated_at=0, + ) + + +def test_worker_pool_allows_cooperative_session_cleanup() -> None: + async def _run() -> None: + worker = _CooperativeWorker(complete_on_stop=True) + pool = InProcessLiveKitWorkerPool({"worker-0": worker}) + await pool.start() + pool.start_session(_record()) + await worker.running.wait() + + await pool.stop_session("session-1") + + assert worker.completed is True + assert worker.cancelled is False + + asyncio.run(_run()) + + +def test_worker_pool_cancels_session_after_cleanup_timeout(monkeypatch) -> None: + async def _run() -> None: + worker = _CooperativeWorker(complete_on_stop=False) + pool = InProcessLiveKitWorkerPool({"worker-0": worker}) + await pool.start() + pool.start_session(_record()) + await worker.running.wait() + monkeypatch.setattr("telefuser.service.livekit.worker_pool._SESSION_STOP_GRACE_SECONDS", 0.01) + + await pool.stop_session("session-1") + + assert worker.completed is False + assert worker.cancelled is True + + asyncio.run(_run()) diff --git a/tests/unit/worker/test_parallel_worker.py b/tests/unit/worker/test_parallel_worker.py index 43cc6d42..9ad858f1 100644 --- a/tests/unit/worker/test_parallel_worker.py +++ b/tests/unit/worker/test_parallel_worker.py @@ -1,5 +1,6 @@ """Tests for parallel_worker module.""" +import signal from queue import Empty from unittest.mock import MagicMock, Mock, patch @@ -331,7 +332,8 @@ class TestWorkerLoopUnit: @patch("telefuser.worker.parallel_worker.dist") @patch("telefuser.worker.parallel_worker.current_platform") - def test_worker_loop_single_process(self, mock_platform, mock_dist): + @patch("telefuser.worker.parallel_worker.signal.signal") + def test_worker_loop_single_process(self, mock_signal, mock_platform, mock_dist): """Test _worker_loop with world_size=1 (no distributed).""" import multiprocessing as mp @@ -366,6 +368,10 @@ def test_worker_loop_single_process(self, mock_platform, mock_dist): mock_stage.test_method.assert_called_once() queue_out.put.assert_called() mock_dist.init_process_group.assert_not_called() + assert mock_signal.call_args_list == [ + ((signal.SIGINT, signal.SIG_IGN),), + ((signal.SIGINT, mock_signal.return_value),), + ] @patch("telefuser.worker.parallel_worker.dist") @patch("telefuser.worker.parallel_worker.current_platform") diff --git a/tests/unit/worker/test_tensor_channel.py b/tests/unit/worker/test_tensor_channel.py index 727d9c9f..510ff5e8 100644 --- a/tests/unit/worker/test_tensor_channel.py +++ b/tests/unit/worker/test_tensor_channel.py @@ -187,6 +187,22 @@ def test_tensor_channel_validates_bindings_and_rank() -> None: WorkerTensorChannel(consumer_world_size=1, cuda_ipc_slots=0) +def test_tensor_channel_releases_producer_pools_with_local_ipc_state() -> None: + channel = WorkerTensorChannel(consumer_world_size=1, timeout=1) + channel._cuda_pools[("profile",)] = object() + channel._cuda_storage_cache[(b"storage", 0)] = object() + + try: + channel._ensure_local_cuda_ipc_finalizer() + channel.release_local_cuda_ipc() + + assert channel._cuda_pools == {} + assert channel._cuda_storage_cache == {} + assert channel._cuda_ipc_finalizer is None + finally: + channel.close() + + def test_tensor_channel_transfers_between_independent_spawned_processes() -> None: context = mp.get_context("spawn") channel = WorkerTensorChannel(consumer_world_size=1, timeout=10) diff --git a/tf-kernel/csrc/distributed/ulysses_cuda_ipc.cu b/tf-kernel/csrc/distributed/ulysses_cuda_ipc.cu index 78ec891c..848ebae1 100644 --- a/tf-kernel/csrc/distributed/ulysses_cuda_ipc.cu +++ b/tf-kernel/csrc/distributed/ulysses_cuda_ipc.cu @@ -188,13 +188,15 @@ void ulysses_stream_barrier( "cuStreamWriteValue64"); } const auto local_base = reinterpret_cast(local_barrier.data_ptr()); + // Epochs reuse one slot per peer. GEQ remains satisfied if a faster peer has + // already published the next epoch before this stream evaluates its wait. for (int64_t peer = 0; peer < world_size; ++peer) { check_driver( cuStreamWaitValue64( driver_stream, local_base + peer * sizeof(uint64_t), static_cast(epoch), - CU_STREAM_WAIT_VALUE_EQ), + CU_STREAM_WAIT_VALUE_GEQ), "cuStreamWaitValue64"); } } diff --git a/tf-kernel/tests/distributed/ulysses_correctness.py b/tf-kernel/tests/distributed/ulysses_correctness.py index f0dc3b4a..daee5ecf 100644 --- a/tf-kernel/tests/distributed/ulysses_correctness.py +++ b/tf-kernel/tests/distributed/ulysses_correctness.py @@ -36,6 +36,7 @@ def main() -> None: torch.cuda.set_device(local_rank) dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank)) rank = dist.get_rank() + world_size = dist.get_world_size() shape = (2, 12, 8, 16) group = CudaIpcUlyssesGroup(dist.group.WORLD) @@ -82,6 +83,28 @@ def main() -> None: for output, offset in zip(outputs, (80, 96, 112), strict=True): torch.testing.assert_close(output, _expected_scatter(rank, shape, dtype, offset), rtol=0, atol=0) + assert len(group._targets) <= 12 + + stress_dtype = torch.bfloat16 + stress_input = _input(rank, shape, stress_dtype) + latest_outputs: list[torch.Tensor] = [] + latest_offsets: tuple[int, int, int] = (0, 0, 0) + for step in range(64): + base_offset = (step % 8) * 3 + latest_offsets = (base_offset + 1, base_offset + 2, base_offset + 3) + handles = [ + group.all_to_all_single_4d_async(stress_input + latest_offsets[0], mode=0, tag="stress-q", barrier=False), + group.all_to_all_single_4d_async(stress_input + latest_offsets[1], mode=0, tag="stress-k", barrier=False), + group.all_to_all_single_4d_async(stress_input + latest_offsets[2], mode=0, tag="stress-v"), + ] + latest_outputs = [handle.wait() for handle in handles] + if rank == step % world_size: + torch.cuda._sleep(5_000_000) + + for output, offset in zip(latest_outputs, latest_offsets, strict=True): + torch.testing.assert_close(output, _expected_scatter(rank, shape, stress_dtype, offset), rtol=0, atol=0) + assert len(group._targets) <= 12 + torch.cuda.synchronize() dist.barrier() group.close() diff --git a/tf-kernel/tf_kernel/ulysses.py b/tf-kernel/tf_kernel/ulysses.py index 063adc06..587bf4af 100644 --- a/tf-kernel/tf_kernel/ulysses.py +++ b/tf-kernel/tf_kernel/ulysses.py @@ -3,11 +3,14 @@ from __future__ import annotations import socket +from collections import OrderedDict from dataclasses import dataclass import torch import torch.distributed as dist +_MAX_CACHED_TARGETS = 12 + @dataclass class _Completion: @@ -68,7 +71,7 @@ def __init__(self, process_group: dist.ProcessGroup, device: torch.device | None _, greatest_priority = torch.cuda.Stream.priority_range() self._comm_stream = torch.cuda.Stream(device=self.device, priority=greatest_priority) - self._targets: dict[tuple[object, ...], _Target] = {} + self._targets: OrderedDict[tuple[object, ...], _Target] = OrderedDict() self._barrier = torch.zeros(self.world_size, dtype=torch.int64, device=self.device) self._peer_barriers, self._barrier_remote_pointers = self._open_peer_handles(self._barrier) self._barrier_epoch = 0 @@ -139,14 +142,30 @@ def _target(self, input: torch.Tensor, mode: int, tag: str) -> _Target: key = (tag, mode, output_shape, input.dtype, input.device) target = self._targets.get(key) if target is not None: + self._targets.move_to_end(key) return target + if len(self._targets) >= _MAX_CACHED_TARGETS: + self._evict_oldest_target() output = torch.empty(output_shape, dtype=input.dtype, device=input.device) peer_outputs, remote_pointers = self._open_peer_handles(output) target = _Target(output=output, peer_outputs=peer_outputs, remote_pointers=remote_pointers) self._targets[key] = target return target + @staticmethod + def _close_target(target: _Target) -> None: + for pointer in target.remote_pointers: + torch.ops.tf_kernel.cuda_ipc_close_mem_handle(pointer) + + def _evict_oldest_target(self) -> None: + # Peer ranks may still be writing while the caller stream consumes an older target. + # A cache miss is infrequent, so prefer a device-wide synchronization over unsafe reuse. + torch.cuda.synchronize(self.device) + dist.barrier(group=self.process_group) + _, target = self._targets.popitem(last=False) + self._close_target(target) + def all_to_all_single_4d_async( self, input: torch.Tensor, @@ -178,6 +197,7 @@ def all_to_all_single_4d_async( torch.ops.tf_kernel.ulysses_all_to_all_ce( input, target.peer_outputs[peer], self.rank, self.world_size, mode, peer ) + input.record_stream(self._comm_stream) if barrier: self._barrier_epoch += 1 @@ -198,9 +218,10 @@ def close(self) -> None: if self._closed: return self._comm_stream.synchronize() + if dist.is_initialized(): + dist.barrier(group=self.process_group) for target in self._targets.values(): - for pointer in target.remote_pointers: - torch.ops.tf_kernel.cuda_ipc_close_mem_handle(pointer) + self._close_target(target) for pointer in self._barrier_remote_pointers: torch.ops.tf_kernel.cuda_ipc_close_mem_handle(pointer) self._targets.clear()