[mickqian]

Branch: mick/diffusion-ulysses-pipeline-same-device-copies (on main efd2ce8), commit 2aec3e1

Motivation

The head-group pipeline (IpcA2AMultiState.pipelined_attention) overlaps the Ulysses exchange with attention: group g's peer copies run on the copy engine while group g-1 attends. Two parts of it were same-device copies, and a same-device cuMemcpy2DAsync (or copy_) runs on the SMs, so it waits out whatever attention kernel is running:

1. Own rows on the copy stream. This rank's rows of each group were copied into its own receive slot on the copy stream, interleaved with the peer sends. Each group's own copy waited for the previous group's attention to finish, and the peer sends queued behind it waited with it. Wan2.2 A14B 720p on 4xB300 (CFG 2 x Ulysses 2, 5 head groups) showed a 0.28-0.46 ms gap before every group: 1.76 ms per block between RoPE and the output projection.
2. Merges after the last group. Every group's output was copied into the merged output (own rows) and the peers' blocks were copied in after their flags (D2D on the group streams). With the gaps gone these all run after the last attention: 0.63 ms per block.

A microbenchmark on B300 (FA4 or cuDNN attention on one stream, six 38.7 MB pitched copies issued on another): D2D copies finish 0.2 ms after the attention ends; P2P copies finish 0.4 ms after they start.

Modifications

• kernels/ops/diffusion/layout/head_group_copy_triton.py: copy_head_groups(src, dst) copies [group, batch, rows, width] views with arbitrary strides (last dim contiguous) in one Triton pass.
• ipc_a2a_multi.pipelined_attention:
  • this rank's rows of every group go into its slots in one copy_head_groups pass per part on the main stream, before the first group; the copy stream carries only peer sends;
  • new attend_into(q, k, v, out): when the backend can write through a strided output, each group attends straight into a token-major [batch, world * s_local, heads, head_dim] buffer. This rank's rows of it are the returned output; the peers' rows are sent from it on the copy engine; the peers' blocks come back in one copy_head_groups pass per peer. No own merge, one gather per peer.
• AttentionImpl.forward_into(q, k, v, out, attn_metadata) -> bool (default False). FlashAttentionImpl serves it on FA4 (same call as forward, with out=); DynamicCudnnSDPAImpl delegates when it would take FA4 for those inputs. out=None only asks.
• USPAttention's dense Ulysses path passes attend_into.

Masked, replicated-row, fill (MiniMax-H3) and cuDNN-served calls keep the previous output path.

Accuracy Tests

The pipeline checks itself against the sequential exchange on first sight and keeps it only on byte equality; with this change it still passes on Wan2.2 (no "differs from the sequential exchange" in the log). Decoded videos are byte-identical (mp4 md5) to the run without the change for the same requests.

• test_usp_pipelined_attention_multi_gpu.py: every Ulysses layout, FA and cuDNN backends, pipelined vs sequential, bitwise, 2 and 4 ranks: passes (the FA4 cases now take attend_into).
• test_ipc_a2a_multi_gpu.py: passes (2 and 4 ranks, expandable segments without pidfd_getfd).
• test_layout.py::test_copy_head_groups_*: both directions vs per-group copies, bitwise; input validation.
• test_cuda_attention_backend.py (forward_into cases): FA4 forward_into into one head group of a wider buffer equals forward bitwise and writes nothing else; backends without it decline.

Speed Tests and Profiling

Wan2.2 T2V A14B 720p, 81 frames, 40 steps, 4xB300, CFG 2 x Ulysses 2, eager, FA4 attention (needs "Prefer FA4 over cuDNN SDPA for long self-attention on B300" for the default backend).

Per transformer block (torch profiler, rank 0):

	RoPE -> first attention	between groups	after last group	block
main	0.36 ms	0.28-0.46 ms x4	0.18 ms (merges in the gaps)	56.6 ms
own rows in one pass	0.23 ms	0	0.63 ms (merges)	54.9 ms
+ attend into token-major output	0.23 ms	0	0.14 ms

E2E, harness single_e2e (client latency, two measured requests per fresh server, two interleaved rounds), on top of the other B300 Wan changes: without attend_into 90.820 / 90.797 s, with it 90.683 / 90.680 s (-0.13 s). The own-rows half was measured together with the FA4 selection change (92.72 -> 91.34 s median, probe client, six requests each).

All four GPUs sit at the 1100 W power cap through the denoise, so idle gaps cost less than their length: removing them raises power and the clocks drop. The e2e gain is smaller than the timeline gain.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #40062976150
Latest PR Test (Extra): ❌ Run #40062976164
Latest PR Test (AMD ROCm 10): ⏳ Run #40062976135

[mickqian]

/tag-and-rerun-ci