@coderabbitai summary

Description

Adds an opt-in async Ulysses pipeline for sequence parallelism on diffusion video models (LTX-2 + WAN 2.2). The pipeline overlaps per-rank V/Q/K compute (GEMM → RMSNorm → RoPE) with cross-rank data movement (PyTorch _SymmetricMemory CUDA-IPC peer access + symm-mem barrier release fence) on a single dedicated per-device side stream. SMs stay free for the next V/Q/K while the GPU Copy Engine drains the prior V/Q/K's peer push.

Opt-in gate: parallel_config.async_ulysses (default False). Both LTX-2 (LTX2Attention.__init__(async_ulysses=...)) and WAN 2.2 self-attn read the same gate. Non-Ulysses, audio, cross-attn, and ulysses_size == 1 paths are unchanged.

Requires PyTorch ≥ 2.5 for torch._C._distributed_c10d._SymmetricMemory.empty_strided_p2p / rendezvous / barrier. No NCCL device-API dependency.

Enabling

YAML config (trtllm-serve --config foo.yaml):

parallel_config:
  ulysses_size: 4         # any P > 1
  async_ulysses: true     # default False — flip to enable the async pipeline

Python API (VisualGenArgs):

from tensorrt_llm._torch.visual_gen.config import VisualGenArgs, VisualGenParallelConfig

args = VisualGenArgs(
    parallel_config=VisualGenParallelConfig(
        ulysses_size=4,
        async_ulysses=True,
    ),
    ...
)

A/B comparison against the baseline blocking-A2A path needs zero code change — flip the bool.

Design

Stream lanes per attention call

default      [V_GEMM+norm+rope+permute]   [Q_GEMM+norm+rope+permute]   [K_GEMM+norm+rope+permute]                            [SDPA]
                      │ev.record                  │ev.record                  │ev.record
                      ▼                           ▼                           ▼                                                ▲
side_stream      [V_push]                    [Q_push]                    [K_push]   [barrier × 3] ─► ev_done.record ──────────┘
                                                                                                     default waits ev_done

Issue order V → Q → K. V's CE push overlaps with Q's compute; Q's push overlaps with K's compute. Three barriers fire only at _join_async, after which the default stream waits the tail event before SDPA. Side stream is a per-device singleton (UlyssesAttention._side_stream_by_device) so all transformer layers on the same device share one comm lane — avoids per-layer stream proliferation under cuda_graph capture / inductor.

C++ ops surface

// Phase 1 — default stream: lazy-alloc slot ring + fused permute+scatter kernel.
std::tuple<Tensor, SendHandle> ulysses_a2a_async_prepare(Tensor input, ProcessGroup pg);

// Phase 2a — side stream: CE peer push only (no barrier).
void ulysses_a2a_async_push(SendHandle send_h, ProcessGroup pg);

// Phase 2b — side stream: symm-mem barrier release fence (channel=0).
//   Called N times at join, once per deferred push.
void ulysses_a2a_async_barrier(ProcessGroup pg);

AsyncUlyssesOp caches a canonical SymmetricMemory handle on the first slot allocation (mCanonicalHandle); emitBarrier() uses it directly instead of scanning slots each call.

forward_async API (closure-based)

@torch.compiler.disable(recursive=False)
def _issue_async(self, perm_4d):
    recv, send_h = torch.ops.trtllm.ulysses_a2a_async_prepare(perm_4d, self._pg_boxed)
    ev = torch.cuda.Event(); ev.record()
    with torch.cuda.stream(self._async_side_stream):
        ev.wait()
        torch.ops.trtllm.ulysses_a2a_async_push(send_h, self._pg_boxed)
    self._pending_barriers += 1
    return recv

@torch.compiler.disable(recursive=False)
def _join_async(self):
    with torch.cuda.stream(self._async_side_stream):
        for _ in range(self._pending_barriers):
            torch.ops.trtllm.ulysses_a2a_async_barrier(self._pg_boxed)
        ev_done = torch.cuda.Event(); ev_done.record()
    self._pending_barriers = 0
    torch.cuda.current_stream().wait_event(ev_done)

def forward_async(self, compute_q, compute_k, compute_v, **attn_kwargs):
    v_4d = compute_v(); v_5d = self._issue_async(v_4d)
    q_4d = compute_q(); q_5d = self._issue_async(q_4d)
    k_4d = compute_k(); k_5d = self._issue_async(k_4d)
    self._join_async()
    # ulyssesPostUnscatterKernel ▶ SDPA ▶ reverse a2a — all in the outer compile region.

@torch.compiler.disable(recursive=False) on _issue_async / _join_async is the only stream-switch boundary. The caller's compute_q/k/v closures execute on the default stream between successive _issue_async graph breaks, so their bodies sit inside the outer block's inductor graph and get full GEMM+norm+RoPE fusion (verified nvjet_sm100_ootst_..._Avec16UE4M3_Bvec16UE4M3_TNT kernels emitted under NVFP4).

_SymmetricMemory slot ring

AsyncUlyssesOp owns a 3-slot ring (V/Q/K per layer — intra-layer slot-reuse hazard requires ≥ 3, cross-layer reuse is safe because _join_async syncs the side stream before next layer's V starts). First call at (numel, dtype) triggers empty_strided_p2p + rendezvous; subsequent calls reuse the cached (handle, peer_ptrs). Ops cached by pg->getGroupName() so multiple PGs (e.g. ulysses subgroup × cfg subgroup) coexist cleanly. Allocation is not graph-capture-safe (rendezvous / cudaMalloc inside) — forward_async requires one out-of-capture warmup pass before cuda_graph capture starts.

Key Optimizations

The four mechanisms below are what make this PR's async path actually faster than the baseline blocking A2A. Listed bottom-up from hardware-level resource choice to host-side launch reduction.

1. Copy-Engine peer push (avoid SM contention)

V/Q/K peer-to-peer data movement runs entirely on the GPU Copy Engines, never on SM-resident NCCL kernels:

• Eager path: one cudaMemcpyBatchAsync over the P-1 per-peer chunks (driver fans out across CEs).
• Capture path: per-peer cudaMemcpyAsync loop (because cudaMemcpyBatchAsync isn't graph-capture-safe until CUDA 13.5 — see Follow-ups Bump onnx from 1.12.0 to 1.13.0 #1).

The Copy Engine handles all P-1 peer pushes, leaving SMs fully free for the next V/Q/K GEMM → RMSNorm → RoPE. Compared to NCCL all-to-all (Send/Recv kernels on SMs competing for the same compute resource as the GEMMs), the CE-based push has disjoint hardware resources from compute — no channel multiplexing or ring-topology hacks needed.

2. Fuse permute with self-chunk scatter (bypass CE local-D2D bandwidth cap)

ulyssesPermuteScatterKernel does two things in one launch:

• Permute [B, S, H, D] → [P, B, S, H/P, D]
• Write the local rank's self chunk directly to the recv slot's self position (slot.basePtr + my_rank * chunk_bytes); only the P-1 peer chunks go to slot.sendBuf.

Without the fuse, the alternative is: SM permute → sendBuf; then CE runs P memcpys (P-1 peer pushes + 1 local sendBuf → recv self-slot D2D copy). The CE's local D2D channel is hardware-capped at ~100 GB/s (independent from the NVLink/NVSwitch peer channel at ~400+ GB/s — they don't share bandwidth, but local D2D's own ceiling is low). Since the self chunk is the same size as one peer chunk, the local copy would become the critical path of the entire a2a — the P-1 peer pushes finish 4× faster while we wait for one local D2D.

Fusing moves the self-chunk write into the SM kernel, using HBM bandwidth (~5–10 TB/s) — one to two orders of magnitude above CE local-D2D. The alternative of green-context isolation for the local copy (separate CE queue) requires complex multi-stream/queue coordination; fusing bypasses that whole subsystem.

3. Deferred barriers (push×N FIFO, barrier×N drain at join)

V/Q/K's a2a is split into two C++ ops, ulysses_a2a_async_push (CE push only) and ulysses_a2a_async_barrier (symm-mem barrier release fence). Side-stream sequence:

[push_V, push_Q, push_K, barrier, barrier, barrier]

Originally the op was push+barrier fused, producing [push, barrier, push, barrier, push, barrier] where each barrier kernel cut the CE FIFO and forced the next push to wait. Each barrier kernel drops from ~164 µs → ~93 µs (-43%) under the deferred scheme, and per-step CE PTOP throughput improves by ~4 ms (WAN 8-GPU step capture). Python state tracks _pending_barriers on UlyssesAttention; _join_async drains exactly N barriers on the side stream before the default stream waits the tail event.

4. Dedup fp4_quantize across QKV self-attn

Under async ulysses the video self-attn is forced to SEPARATE_QKV (so V/Q/K can stream-pipeline independently), but with NVFP4 static quant the three Linears all share input_scale:

# Verified bit-equal across WAN 36+36 / LTX-2 42 self-attn layers (modelopt invariant)
to_q.input_scale == to_k.input_scale == to_v.input_scale

Pre-quantize hidden_states once and share the resulting Fp4QuantizedTensor across all three Linears via the existing fp4-shortcut path in NVFP4LinearMethod._input_prepare (which just unpacks .fp4_tensor / .scaling_factor for GEMM — no internal swizzle so the SF layout is the same as Linear's own quantize path). Eliminates 2 of 3 quantize launches per self-attn block (verified −368 launches in WAN A14B 8-GPU step capture; identical kernel count between async-OFF FUSE_QKV and async-ON SEPARATE_QKV+dedup).

Gating:

• Structural (_maybe_share_qkv_quantize, set in __init__): SEPARATE_QKV + NVFP4 + not force_dynamic_quantization.
• Runtime: getattr(to_q, "input_scale", None) is not None — guards against checkpoints that exclude individual Linears from NVFP4 (e.g. LTX-2 transformer_blocks.10.attn1, bf16-only — falls back to per-Linear quantize).

Benchmarks

Isolated all-to-all (B200, single rank)

bench_alltoall_ab.py — single-tensor a2a latency (ms), B=2, S_total=6144, H=32, D=128, warmup=30, bench=100. naive_fused = dist.all_to_all_single on 5D fused QKV; naive_split = 3 separate 4D a2a; symm_ce = the new symm-mem CE pipeline (per V/Q/K call).

symm_ce is 1.16–1.29× faster than the fused baseline in every cell except P=8 graph (where cuda_graph capture amortizes Python launch overhead in the naive_fused 1-collective path enough that the merged baseline wins).

End-to-end on B200 (production scale)

3 timed E2E runs/cell, median reported. Per-cell range is well below the OFF→ON delta in every case (non-overlap). WAN's larger relative gain reflects more compute-per-A2A and a longer overlap window per layer (no cuda_graph overhead to amortize the baseline path).

Files

Follow-ups

1. cudaMemcpyBatchAsync under CUDA Graph capture (CUDA 13.5). Today the capture branch falls back to a P-1-iteration per-peer cudaMemcpyAsync loop because cudaMemcpyBatchAsync is not graph-capture-safe. Per NVBugs 5853376 and NVBugs 5760690, targeted at CUDA 13.5 (Q4 2026). We tried cudaGraphAddMemcpyNode1D with sibling-DAG independence as a stop-gap; nsys confirmed runtime serializes the sibling memcpys on the same HW queue (~4 µs gaps, no CE fanout), so the per-peer loop wins today.

2. Collapse N barriers at join → 1 barrier. Current _join_async drains N symm-mem barriers in a loop (one per deferred push) — kept 1:1 with the original protocol for risk-free refactor. The barriers are semantically identical (all channel-0 PG-level fences); a single barrier after [push_V, push_Q, push_K] should sync the recv buffers just as well (side-stream FIFO + cross-rank barrier covers all prior pushes). Saves 2 of 3 barrier_kernel per layer (~186 µs × N_layers per step on B200). Requires verifying that all ranks symmetrically issue 1 barrier per join (slot rotation is independent of barrier count — nextSlotIdx ticks in _prepare, not _barrier).

3. Drop @torch.compiler.disable, native multi-stream torch.compile. Stream-switch boundaries in _issue_async/_join_async are guarded by @torch.compiler.disable(recursive=False) because inductor doesn't yet model with cuda.stream(...) + event.record/wait_event as first-class graph ops (4 graph breaks per layer × N layers). Inline once torch.compile's multi-stream support matures.

4. Fuse pre-A2A permute into the RMSNorm+RoPE kernel. We already own the fused RMSNorm+RoPE kernel (see fuse_qk_norm_rope=True path). Passing the slot's send_buf pointer + permute index map directly into the kernel epilogue lets it write the permuted layout straight into the symm-mem slot, eliminating the separate ulyssesPermuteScatterKernel launch on the critical path. Leaves the post-unscatter kernel untouched.

Test Coverage

• Kernel unit tests (single-GPU) under tests/unittest/_torch/thop/parallel_hw_agnostic/: test_ulysses_permute_scatter.py and test_ulysses_post_unscatter.py validate the two new fused CUDA kernels against torch references in isolation.
• Multi-rank pipeline test (ws=2) under tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py: three torch.multiprocessing.spawn tests covering the production prepare/push/barrier path end-to-end —
  • test_slot_ring_wraparound: eager loop > kNumSlots*2 iterations, byte-exact vs all_to_all_4d. Catches off-by-one slot-reuse bugs.
  • test_capture_smoke: warm slot out-of-capture, then torch.cuda.CUDAGraph replay 8× with fresh inputs, byte-exact vs all_to_all_4d. Exercises the per-peer cudaMemcpyAsync path used under production cuda_graph.
  • test_multi_pg: two distinct PGs spanning the same ranks, alternate calls, byte-exact vs each PG's all_to_all_4d. Exercises PG-name caching in getOrCreateOp and set_group_info re-registration across groups.
• Recommended CI bot stages: DGX_B200-4_GPUs-PyTorch-1, DGX_B200-8_GPUs-PyTorch-1 (covers Ulysses ws=4 and ws=8 in both LTX-2 and WAN end-to-end).

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.