[mickqian]

Motivation

#43164, #43361 and #43373 pipeline MiniMax-H3's Ulysses all-to-all against its dense attention, one head group at a time, over the copy engine:

• while group g attends, group g + 1's q/k/v moves;
• finished groups send their output back while later groups still attend;
• the result is bit for bit that of the sequential exchange.

It took 2–4% off H3 on H200 and B300.

No other model got it: every USPAttention Ulysses path still exchanges all heads first and attends afterwards. That covers Wan2.1/2.2, LTX-2, the FLUX family, Qwen-Image, Z-Image, HunyuanVideo, LongCat, GLM-Image, Helios, Krea2, Joy-Image, Kandinsky 6, Cosmos3 and more. This PR turns the pipeline into a building block that every Ulysses path of the layer uses, and that a custom attention module can call directly.

Modifications

The transport: ulysses_pipelined_attention / IpcA2AMultiState.pipelined_attention

• Layouts. It covers every layout USPAttention exchanges:

  • q/k/v as [batch, s_local, heads, head_dim] at any batch size;
  • GQA, where k/v have fewer heads than q;
  • text rows every rank holds (replicated=, per part): a prefix, a suffix (replicated_first=False), or a K/V-only prefix.

  Replicated rows stay out of the exchange. Each rank places its own head groups of them into its receive slot on the copy engine. Their output reaches every rank with the shard rows, so the merged output comes back in the caller's row order: replicated rows (all heads) and this rank's shard.

• Receive slots. Per head group and part, the slot is one contiguous [batch, rows, group_heads, head_dim] tensor in exactly the order attention reads. attend(q, k, v) runs on views of it.

• Buffers: one pool for every shape. One pool per dtype holds the receive slots and output blocks of every call, grown to the largest call seen. The first version kept a set per call shape (up to 16), each twice over, alternating per call. One set is enough:

  • a rank sends a group's output only after attending over that group's slot, and a peer writes the next call's rows only after merging all of that output;
  • a peer's next output block needs this rank's rows of the next call, which this rank sends only after merging the previous call.

  The pool takes about 4 × rows per rank × heads × head_dim. That is half of what the sequential exchange's staging keeps at Ulysses 2 (two slots of q/k/v and output, per shape).

• First-sight check (sequential=, signature=). The first call per plan and signature also runs the caller's sequential path. It keeps the pipeline only if every rank sees the same bytes (one all-reduce); otherwise that shape keeps the sequential exchange for good. A backend that would not compute a head group exactly as it computes those heads in the full call can therefore never change a result.

  • The check runs the sequential path first, with the pool released, and over NCCL on per-call buffers (uncached_a2a_staging()). So the call never holds both exchanges' buffers, and no staging sized for it stays allocated. An all-to-all is a permutation, so its bytes do not depend on the transport.

• Automatic group count. It is now the largest of 7, 5, 4, 3 and 2 that divides the q and k/v heads per rank while each group's call still fills four waves of 128-row tiles. 5 and 3 are new candidates for Wan's 40 and 12 heads. The waves floor comes from B300, where cuDNN SDPA calls too small to fill the SMs cost more than the overlap saves:

  • Wan2.1 1.3B at Ulysses 4 in three one-head groups ran 44% slower than sequentially;
  • an LTX-2-shaped layer in four two-head groups ran 17% slower.

  Short sequences, such as a 1024² image, keep the sequential exchange. H3 keeps 7 groups on H200. On B300, T2VA at Ulysses 4 moves from 7 to 2 groups (3.99 waves at 7); the two were within run-to-run noise there in [diffusion] perf: pipeline MiniMax-H3's Ulysses exchange against cuDNN SDPA on Blackwell #43373.

• Host cost. Copies are issued from precomputed addresses instead of tensor views. A view per copy cost more host time than the copy at Ulysses 4. The decision per call shape is cached, so a shape that cannot pipeline steps aside at once.

• H3 unchanged. The H3 fill path (batch 1, no replicated rows) issues the same copies as before. Its buffers drop from about 10 to 6 × rows per rank × heads × head_dim (Ref2VA at Ulysses 2: 3.6 → 2.1 GB). The cookbook page says so.

The layer: USPAttention._pipelined_ulysses

Every Ulysses path describes its layout, hands over its attention core as attend and its own exchange-attend-exchange as sequential. The sequential code is kept as written and still runs whenever the pipeline does not. The paths covered:

Path	Users	Attention core per head group
dense	Wan2.1/2.2, LTX-2 (unmasked)	attn_impl.forward
key mask / tail padding	Kandinsky 6, Anima, MOVA; FLUX family and Qwen-Image when the text is sharded	varlen FA over cu_seqlens_tail, the dense gap, or the gathered mask (FA varlen or SDPA)
replicated prefix (± key mask)	FLUX.1/2 with replicated text, Qwen-Image, LongCat, GLM-Image, Helios, Krea2	attn_impl.forward, or the masked local attention
replicated suffix	Z-Image, Joy-Image, HunyuanVideo	attn_impl.forward (K/V in tail order, as before)
replicated K/V prefix	Cosmos3, Qwen-Image 2.1	attn_impl.forward

• The masked path's gathered mask and its varlen metadata are computed once before the exchange. Attention then has no collective or host sync per group.
• Ring attention, uneven sequence shards (seq_lens), the K/V-gather mode and the legacy UlyssesAttention keep the sequential exchange.
• SGLANG_DIFFUSION_ULYSSES_PIPELINE_GROUPS=0 turns the pipeline off. The parallelism page and the environment-variable table describe it.

Accuracy Tests

• test_ipc_a2a_multi_gpu.py passes at 2 and 4 ranks (4×B300). New cases:

  • batch 2;

  • replicated prefix, suffix and K/V-only prefix;

  • GQA;

  • automatic groups at a shape that fills the GPU, and declining at one that does not;

  • a verified signature that keeps pipelining;

  • a mismatching first call that returns the sequential result and stops pipelining that shape.

  • one pool serving every plan of the test, sized for the largest;

  • six back-to-back calls of two alternating plans with rank 0 attending about 10 ms slower, compared only after all were queued.

  Each case is checked byte for byte against attention over the global layout, with every head at once.

• New test_usp_pipelined_attention_multi_gpu.py passes at 2 and 4 ranks. It runs 13 USPAttention cases sequentially and with the pipeline and requires identical bytes:

  • dense at batch 1 and 2;
  • replicated prefix, with and without a key mask;
  • replicated suffix;
  • replicated K/V prefix;
  • gathered key mask;
  • GQA.

  It also checks that every case actually pipelined and verified, and that the first call of a new shape leaves no sequential staging behind. On B300 FlashAttention resolves to Torch SDPA, so these ran under cuDNN SDPA. The tail-padding varlen-FA branch needs Hopper and is covered there by the first-sight check.

• End to end, the output bytes do not change:

  • Wan2.1-T2V-1.3B 480p 81f: four alternating runs give the same MP4 (0b3acb06debb);
  • Z-Image-Turbo 2048²: four alternating runs give the same PNG (34e0d5d0fb90);
  • MiniMax-H3 T2VA at Ulysses 2 (its own fill pipeline through the rewritten transport): the MP4 is the same as on main (639d5a2f219d), before and after the shared pool;
  • Wan2.1-T2V-1.3B 1280×720 49f with the CI's server flags: six servers across main and both versions of this PR give the same MP4 (10e8e381ddbc).

Speed Tests and Profiling

One attention layer (exchange + attention + output exchange), cuDNN SDPA, 4×B300, CUDA-event median. Every pipelined output equals the sequential one:

Shape (global tokens × heads, replicated rows)	Ulysses 2: sequential → pipelined	Ulysses 4: sequential → pipelined
Wan2.1 14B 480p (32760 × 40)	7.91 → 7.01 ms (−11.4%)	4.10 → 3.61 ms (−11.8%)
Wan2.1 14B 720p (75600 × 40)	40.98 → 39.13 ms (−4.5%)	20.35 → 18.95 ms (−6.9%)
Wan2.1 14B 480p, batch 2	17.06 → 14.67 ms (−14.0%)	9.07 → 7.76 ms (−14.5%)
Wan2.1 1.3B 480p (32760 × 12)	2.65 → 2.31 ms (−12.6%)	1.434 → 1.435 ms (keeps the sequential exchange)
LTX-2-like (24576 × 32)	3.85 → 3.21 ms (−16.7%)	2.17 → 1.88 ms (−13.3%)
Z-Image 2048² (16384 × 30, 128-row caption suffix)	2.17 → 1.69 ms (−22.1%)	— (30 heads)
FLUX 2048² (16384 × 24, 512-row text prefix)	1.90 → 1.44 ms (−24.0%)	1.150 → 1.138 ms (keeps the sequential exchange)
HunyuanVideo-like (61440 × 24, 256-row text suffix)	16.62 → 15.37 ms (−7.5%)	8.39 → 7.80 ms (−7.1%)

Shapes too short to split at Ulysses 4 keep the sequential exchange, at no cost: they skip straight to it once the first call turns them down.

End to end on B300, client latency per request, alternating sequential (SGLANG_DIFFUSION_ULYSSES_PIPELINE_GROUPS=0) and pipelined runs:

Model, layout	sequential	pipelined
Wan2.1-T2V-1.3B 480p 81f, --num-gpus 4 --ulysses-degree 2 --enable-cfg-parallel (the cookbook's multi-GPU layout for 14B, here on 1.3B)	6.89 / 6.89 / 6.89 / 6.90 s	6.49 / 6.50 / 6.49 / 6.49 s (−5.8%)
Z-Image-Turbo 2048², --num-gpus 2 --ulysses-degree 2 (the cookbook serves it on one GPU; Ulysses 2 here exercises the replicated caption suffix)	1.53 / 1.53 / 1.55 / 1.55 s	1.41 / 1.44 / 1.43 / 1.43 s (−7.1%)

MiniMax-H3, which main already pipelines, does not regress through the rewritten transport. T2VA 5 s at Ulysses 2, 50 steps: 65.37 / 64.86 s on main, 64.41 / 64.40 s with this PR, same MP4. With the shared pool, in another session: 65.32 / 64.83 s against main's 65.38 / 64.90 s, same MP4.

On the 2×H100 CI runners (the FlashAttention path), Wan2.1 I2V 14B at Ulysses 2 denoises 3.0% (720P), 2.7% (720P with LoRA) and 3.1% (480P) faster than the H100 baselines.

Memory. The first version failed CI's peak guard on Wan2.1 I2V 14B 720P at Ulysses 2 on H100: 46440 MiB runtime allocated against a 44992 MiB baseline (+3.2%). The measured request's shape differs from every warmup shape, so the first-sight check ran inside it. It left the sequential exchange's staging allocated next to the pipeline's buffers, and a warmup shape's buffer set stayed allocated too. The shared pool and the uncached first-sight check above remove both.

One attention layer, per rank, at the Wan2.1 14B 720P 81f shape at Ulysses 2 (37800 rows × 40 heads × 128), 4×B300:

	first call: peak	between calls: held	later calls: peak
sequential exchange	4434 MiB	2954 MiB	4434 MiB
this PR, first version	7831 MiB	5907 MiB	6425 MiB
this PR	2475 MiB	1477 MiB	1994 MiB

The first version's held memory assumes nothing released the staging between requests. The worker does release it after a warmup probe, and the first version then held 2953 MiB.

End to end, Wan2.1-T2V-1.3B with the CI's server flags (--num-gpus 2 --ulysses-degree=2 --warmup-sampling-params '{"quality": "exact"}'), two requests per server, two servers per arm in ABCCBA order. The request is 1280×720 at 49 frames: at 25 frames this model's groups fall just short of four waves on B300's 148 SMs, so that shape keeps the sequential exchange.

	runtime allocated peak per request	second request
main	15.21–15.22 GB	22.04 / 22.08 s
this PR, first version	15.75 GB	20.92 / 20.71 s
this PR	15.08–15.09 GB	20.86 / 20.87 s

The layer timings did not move with the shared pool. Against the first version, ABBA at Ulysses 2 and 4, every shape of the table above stays within ±1%; shapes that keep the sequential exchange in both trees vary by ±0.7%. The outputs are identical.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #38053870955
Latest PR Test (Extra): ❌ Run #38053870819
Latest PR Test (AMD ROCm 10): ⏳ Run #38053870936

[mintlify]

Preview deployment for your docs. Learn more about Mintlify Previews.

Project	Status	Preview	Updated
sglang-doc	🟢 Ready	View Preview	Oct 10, 2026, 12:57 PM

💡 Tip: Enable Automations to automatically generate PRs for you.

[mickqian]

/tag-and-rerun-ci