Repository navigation
[Feature] Compose Ulysses with AllGather-KV as a 2D diffusion SP - #7641
Conversation
2846773 to
547bfa7
Compare
|
This PR appears to belong to: docs/design/module/diffusion/index.md, docs/design/module/diffusion/parallelism.md. Module owners: @princepride @xuechendi @david6666666 Routing: @princepride via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @xuechendi via module of the changed files, module named in the PR description, semantic router; @david6666666 via module of the changed files, module named in the PR description @Bounty-hunter, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer. Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment. |
Omni ReviewBot: supersededThe CI failure noted on |
Omni ReviewBot: no human activity for 14 days@Bounty-hunter this pull request has had no human commit, comment or review since 2026-09-16. Please consider marking this PR as draft until work can resume. The author or a maintainer decides whether to change the PR state. To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline. |
|
@Bounty-hunter this PR is labeled Could you please take a look and push an update to get CI green? Once the checks pass we can proceed with review/merge. Thanks! |
Omni ReviewBot routing recordAssigned Strict on zcode (GLM-5.3-Flash) under experiment |
|
@Bounty-hunter gentle ping — this high-priority PR has seen no updates for 22 days and review feedback is waiting. Could you share a status or ETA when you get a chance? Thanks! |
547bfa7 to
7d7b908
Compare
Omni ReviewBot: supersededThe CI failure noted on |
…logy
AllGather-KV was mutually exclusive with Ulysses and Ring. Allow it to be
composed with Ulysses, forming two orthogonal process-group dimensions:
sequence_parallel_size = ulysses_degree * ring_degree * allgather_degree
The initial scope requires ring_degree == 1 and non-causal attention; composing
all three dimensions at once stays unsupported.
Per rank the tensors start at [B, S/(U*A), H, D] and the attention is:
Ulysses all-to-all (U group): Q, K, V -> [B, S/A, H/U, D]
K/V AllGather (A group): K, V -> [B, S, H/U, D]
local-Q / global-KV attention
reverse Ulysses all-to-all
The collective order is load-bearing: gathering the pre-transform shards would
concatenate the strided head layout of different regions and permute the
sequence. The rank layout mirrors it -- Ulysses groups are contiguous blocks of
U ranks and AllGather-KV groups stride by U -- so the flat contiguous sharding
produced by sp_shard makes each Ulysses group hold one contiguous region and the
AllGather rebuild the global sequence in order. The attention backend therefore
keeps seeing a plain local-Q/global-KV problem, which is what sparse-attention
backends such as FastH3 VSA need; Ring keeps K/V block-local and owns
distributed attention itself.
UlyssesAllGatherKVParallelAttention reuses both halves instead of reimplementing
them. It runs the Ulysses strategy with defer_joint=True, which reshards the
image Q/K/V and records the head-sliced joint query/key/value in the attention
metadata rather than concatenating them into K/V, then delegates to the
AllGather-KV path to gather K/V and re-attach the joint tensors, and finishes
with the Ulysses reverse transform.
Unsupported combinations fail closed with actionable errors rather than falling
back to dense attention: ring_degree > 1, causal attention, 2D key masks, uneven
per-rank regions, Scheduler-paged KV, and TRTLLM_ATTN.
Signed-off-by: dengyunyang <584797741@qq.com>
7d7b908 to
fb7b945
Compare
Address review feedback on the composed 2D SP topology: the 2D-mask guard only inspected attn_mask, so models that carry text padding solely in joint_attn_mask (Qwen-Image with unequal prompt lengths, where attn_mask can be None) passed the check. defer_joint=True skips the Ulysses mask merge and the AllGather path only consumes attn_mask, so padded joint K/V attended unmasked and changed the result. Merge the 2D joint mask against the gathered key layout after the K/V AllGather (mirroring the non-deferred Ulysses merge), and fail closed when the joint mask is not 2D or cannot be merged with a non-2D image mask. Also apply two review suggestions on the region-length check: skip the per-forward length collective under advanced_uaa when every active SP boundary is framework-managed auto_pad (sp_rank_local_seq_lens_equal), and keep the no-op early returns outside the torch.compiler.disable'd checker so strict-mode compiled forwards do not unconditionally split the graph. Adds regression tests for the joint-only mask merge (front/rear), the 4D-image-mask + joint-mask rejection, and the auto_pad fast path. Signed-off-by: dengyunyang <584797741@qq.com>
Address the follow-up review on the joint-mask merge: the helper returned early when joint_attn_mask was None, so it did not fully replace Ulysses's mask handling, and the blanket 2D image-mask rejection in pre_attention stayed. Remove that guard and make _merge_joint_attn_mask the single owner of 2D key-mask handling against the final, post-gather key layout: - Fill either missing side with True: the joint segment is sized from joint_key (joint_query can be empty during cached-KV reuse), the image segment from the remaining gathered keys. - A provided 2D attn_mask must cover the gathered global image keys (replicated across SP ranks, the same contract as plain AllGather-KV); batch and length are validated in every case, including no-joint inputs, and the result is normalized to a [B, K_final] bool mask. - 4D masks keep the inherited AllGather query-range slicing; combining one with a 2D joint_attn_mask still fails closed. This also unblocks mask_sp_padding=True under the composed topology: model code builds that padding mask against the global padded sequence, which is exactly the validated layout. Docs updated accordingly. Adds numerical tests for image-only, joint-only, both masks, cached-KV (joint K without joint Q), both layouts, and invalid shapes (length, batch, joint mask without joint K/V). Signed-off-by: dengyunyang <584797741@qq.com>
Restructure the mask merge so each concern reads as one block: normalize both sides' ranks up front (an absent side counts as 2D and gets True-filled later), validate ranks joint-first -- which also makes the error accurate when both masks have unexpected ranks -- then check joint K/V presence, compare each side's shape as one tuple, and finally fill-and-concatenate. Shortens the docstring to the single mental model (2D masks are key-validity masks, so merging is concatenation; 4D masks belong to the inherited query-range slicing). No behavior change except tightened error messages. Signed-off-by: dengyunyang <584797741@qq.com>
|
I completed a self-review of this PR and updated the composed Ulysses × AllGather-KV path to merge text and image masks after gathering K/V, preserving text padding for prompts of different lengths. I removed the blanket rejection of 2D image masks, added batch and length validation, and filled missing masks using the actual joint K length to support cached KV reuse. I also avoided redundant length collectives for strict mode and inputs with guaranteed equal shard lengths, while keeping the fast path outside the disabled checker. Regression tests and documentation were updated accordingly. |
|
please add perf comparasion and accuracy comparasion |
david6666666
left a comment
david6666666
left a comment
There was a problem hiding this comment.
The mask fix looks correct. The packed H3 path still needs its query metadata adapted to the new topology.
Address review feedback: producers like MiniMax-H3 pass global packed cu_seqlens_q/max_seqlen_q through AttentionMetadata.extra, and packed backends (CUDA FlashAttention varlen) consume those boundaries directly without reading query_ranges. After the K/V gather the query is local (S/allgather_degree rows), so the global boundaries pointed past the buffer and mis-attributed request rows (rank 1's second request landing inside the first request's query span). Pure Ulysses restores the full query and was unaffected; pure AllGather-KV carried the gap and the composed strategy inherited it. Remap the packed query metadata in the shared AllGather-KV metadata slicing so both topologies benefit: intersect every global query document with this rank's span [rank * S/A, (rank+1) * S/A) with device tensor ops (no per-layer .item() host sync), keeping exactly one -- possibly zero-length -- local entry per document so the one-to-one pairing with the global cu_seqlens_k documents survives. Key-side boundaries stay global (K/V are gathered); max_seqlen_q keeps the producer value, a safe upper bound of every local segment. packed_padding (the single-document [real, pad] view) narrows q_length to the local valid prefix while its key length keeps the global valid prefix, so Q/KV-length consumers fall back instead of truncating the gathered K/V. Packed metadata combined with joint attention is not defined by any producer yet and fails closed. Tests: the review's reproduction (docs [0,8,12] @ A=2 -> local cu [0,6,6]/[0,2,6]), a document with zero local rows keeping its entry, single-request [real, pad] padding, and the joint rejection at mock level; plus a real-CUDA equivalence test running FlashAttention varlen on the 2-rank AllGather-KV shard path against the unsplit packed sequence. Signed-off-by: dengyunyang <584797741@qq.com>
Address review feedback: the timing payload is dicts/lists/numbers, so drop the pickle round-trip in bench_ulysses_cp_topology.py and remove the benchmark's entry from the pickle allowlist. Signed-off-by: dengyunyang <584797741@qq.com>
f5f7201 to
eebd7a4
Compare
Signed-off-by: dengyunyang <584797741@qq.com>
show in pr description |
david6666666
added
ready
…m-project#7641) Signed-off-by: dengyunyang <584797741@qq.com>
Purpose
Support composing Ulysses with AllGather-KV as an orthogonal 2D sequence-parallel topology
for diffusion attention (
sequence_parallel_size = ulysses × ring × allgather). Implements thedense path of #7589.
Test Plan
cp2usp2andcp4vs golden baseline (pixel + SSIM/PSNR).Test Result
Setup: 4 × NVIDIA L20X, HunyuanImage-3.0 (DiT-only pipeline), t2i 1024×1024, 50 steps,
guidance 2.5, seed 42, enforce-eager, bf16.
Accuracy vs golden diffusers baseline
Thresholds: SSIM ≥ 0.97, PSNR ≥ 30 dB, pixel mean diff ≤ 3e-2, p99 ≤ 3e-1.
allgather_degree=4ulysses=2 × allgather=2Both composed topologies stay within thresholds; cp2usp2 on the rebased head reproduces the
pre-rebase result bit-exactly (same seed, eager). The 4-rank SP-vs-single-rank equivalence
test also passes element-wise as part of the unit suite (335 passed).
Perf — attention forward (isolates the UAA head-padding scenario)
python benchmarks/diffusion/bench_ulysses_cp_topology.py --heads 14 --head-dim 128 --seqs 16384,32768,65536heads=14 is the faithful 4-GPU analogue of MiniMax-H3's 56 heads on 16 GPUs: both leave
heads / U = 3.5, so usp4+uaa pads 14 → 16 heads (+14.3% attention work) whilecp2usp2's U=2 divides evenly and pays an orthogonal K/V all-gather instead.
The margin grows with sequence length and converges to the +14.3% padding cost. At 64K the
composed topology is even cheaper on communication: 0.99 ms (K/V all-gather) vs 1.66 ms
(padded Ulysses all-to-all).
Perf — end-to-end generation
allgather_degree=4)U=2 × A=2)H=32/16,H % 2 == 0)Interpretation: HunyuanImage-3.0's head counts (32/16) divide both 2 and 4, so neither
topology pays any UAA padding and the E2E comparison reduces to the collective structure:
cp2usp2 stacks two communication stages per attention layer (head all-to-all, then the
orthogonal K/V all-gather) against cp4's single all-gather, which costs ~0.12 s/step here.
This is the expected shape of the trade-off: on models whose heads divide the Ulysses
degree, plain AllGather-KV is the lighter 4-card option. The composed topology is the winner
precisely when
heads % U != 0— the case the attention benchmark isolates (heads=14:8.9% → 13.8% reduction from 16K to 64K seq, converging to the +14.3% UAA padding cost) and
the case it was built for (MiniMax-H3-class 56-head models,
usp8+cp2on 16 cards), whileadditionally preserving the head-sharded local-Q/global-KV form that sparse-attention
backends (e.g. FastH3 VSA) require.
Reproduction