[mickqian]

Motivation

With Ulysses, every MiniMax-H3 block exchanges q/k/v before dense attention and the output after it. Both exchanges run on SMs: an NCCL send/recv for q/k/v, and SM copies for the output. They sit between kernels that keep every SM busy (FlashAttention is ~63% of a denoise step on 2×H200), so their time is fully exposed.

Modifications

Four commits. Each one is bit-identical to the sequential exchange.

1. Copy-engine all-to-all with head-group pipelining (ipc_a2a_multi.py, _minimax_h3_pipelined_dense_attention).

   • The transport is the two-rank ipc_a2a design widened to N same-host ranks: each rank writes its chunk straight into its peers' staging buffers with device-to-device copies on the copy engine, and a per-peer sequence flag orders data before flag.
   • Pipelining: group g's copies run while group g−1 attends; each group attends on its own stream; stream memory operations (cuStreamWriteValue32 / cuStreamWaitValue32) order copies and kernels without holding an SM.
   • It falls back to the sequential exchange when any rank lacks peer access or stream memory operations (agreed across the group before any handle exchange), during graph capture, under torch.compile, and for other attention backends, VSA and ring.

2. QK-norm + RoPE writes straight into the exchange. The out-of-place QK-norm kernel writes each destination block's normalized q/k, plus v, in place of the in-place norm followed by a pack. Blocks for this rank go straight into its own receive slot, so they skip a device-to-device copy. Groups are produced in order, so the norm of later groups overlaps the attention of earlier ones.

3. Each head group is merged into the output as it arrives. Group streams write this rank's rows, and then each peer's, directly into the final tensor. This replaces the trailing merge kernel and the own-rank output copy.

4. On by default. SGLANG_DIFFUSION_ULYSSES_PIPELINE_GROUPS now defaults to -1:

   • it picks 7 groups for MiniMax-H3's 56 heads at Ulysses 2, 4 and 8;
   • when the buffers (11 × rows per rank × heads × head_dim; 3.9 GB for Ref2VA at Ulysses 2) would not fit on some rank, every rank keeps the sequential exchange;
   • 0 disables it, and N ≥ 2 forces N groups.

   SGLANG_DIFFUSION_IPC_A2A_MULTI (the plain all-to-all over the copy engine) stays opt-in. The cookbook's configuration notes describe the default.

Accuracy Tests

• The MP4 md5 matches main for T2VA (Ulysses 2 and 4) and Ref2VA (Ulysses 2) after every commit.
• python/sglang/multimodal_gen/test/single_test_file/test_ipc_a2a_multi_gpu.py (new; passes at 2 and 4 ranks on the final commit) checks, byte for byte:
  • the pipelined attention against the sequential exchange, over repeated calls (slot alternation);
  • the fill path;
  • the automatic group count;
  • a head count the groups cannot split, which must decline;
  • the N-rank exchange against NCCL.

Unrelated to this PR: test_ipc_a2a_2_gpu.py fails the same way on plain main ("AllToAll4D never took the IPC path").

Speed Tests and Profiling

Client end-to-end latency, 4×H200, 50 steps, --performance-mode speed, eager, median of 2 requests, G=7. Arms are interleaved on one box: adjacent columns ran back to back. The box drifts ~0.3% across sessions, so compare within a row only.

Case	main	1: pipeline	2: + fused QK-norm	3: + per-group merge
T2VA 5 s, 2 GPUs, Ulysses 2	144.33 s	142.24 s (−1.45%)	−0.62% vs 1	141.33 s (−0.42% vs 2, 141.92 s)
Ref2VA 5 s, 2 GPUs, Ulysses 2 (with #43117)	223.19 s	221.11 s (−0.93%)	−0.79% vs 1	219.97 s (one arm, not paired with 2)
T2VA 5 s, 4 GPUs, Ulysses 4	73.81 s	72.55 s (−1.71%)	−0.78% vs 1	72.23 s (flat vs 2, 72.18 s)

• 4: on by default. With no environment variable set, the pipeline picks G=7 and runs as fast as the explicit G=7 arms: T2VA Ulysses 2 141.27 s (141.33 s with G=7), Ulysses 4 72.20 s (72.23 s). Outputs are md5-equal to main.

• All four commits against main, one interleaved pass, no env: T2VA Ulysses 2 144.80 → 141.32 s (−2.4%), Ulysses 4 74.19 → 72.23 s (−2.6%). For Ref2VA the chain of paired steps above gives about −1.7% on top of [diffusion] perf: decode plain base64 MiniMax-H3 materials without a per-character loop #43117.

• Per-layer budget, Ref2VA at Ulysses 2 after commit 1 (ms): wall 88.38, made up of the compute union 87.45, exposed exchange 0.58 and idle 0.36. The copy engine moves 2.53 ms/layer underneath the attention.

• Why the gains are below the exposed time: the H200s sit at their 700 W cap during denoise, so part of the freed SM time comes back as a lower clock.

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 #37897253385
Latest PR Test (Extra): ❌ Run #37897253153
Latest PR Test (AMD ROCm 10): ❌ Run #37897253236

[mickqian]

/tag-and-rerun-ci