Repository navigation
[Feature][Diffusion] Compose Ulysses with AllGather-KV as a 2D SP topology, with MiniMax-H3 NPU bring-up fixes - #7937
brandneway wants to merge 4 commits into
Conversation
… world works The encoder TP group was built as a single group covering only the first text_encoder_tp_size DiT ranks via init_world_group. GroupCoordinator asserts that every constructing rank belongs to one of its groups, so with tp_size smaller than the DiT world (e.g. tp=8 on 16 NPUs, the only legal topologies for Qwen3-VL's 64 Q / 8 KV heads) every rank outside [0, tp) failed the assert at startup. Tile the DiT world into dit_world/tp_size contiguous groups instead. Every rank holds membership in exactly one group; only the first group encodes, and its output still reaches all DiT ranks through the existing world-group broadcast. tp_size == world yields the same single group as before, so the existing deployments are unchanged. Signed-off-by: brandneway <gyuan4892@gmail.com>
Under AllGather-KV sequence parallelism (allgather_degree > 1), the attention backend sees local Q against global K/V, so the square packed contract (used_q <= total_q) declines and the mask fallback then rejects the global valid_kv_length against the local query length, refusing every packed request with 'packed metadata is unusable'. _resolve_packed_seq_npu now accepts the query_ranges recorded by the AllGather-KV strategy: a single packed range means the local Q rows are a contiguous slice of the same global real document, so seq_q becomes the slice's [local_real, local_total] pair while seq_k keeps the global [real, pad] pair — the same 'pad is a strict suffix, pad rows are never consumed' contract, no mask materialized. Both packed NPU paths (varlen, laser prefix-slice) forward the ranges; their square behavior is unchanged. The mask rebuild fallback keeps working as the safety net for multi-range layouts by checking valid_kv_length against the K/V domain (key.shape[1]) instead of the local query length. Signed-off-by: brandneway <gyuan4892@gmail.com>
…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>
(cherry picked from commit 547bfa7)
…ather The default diffusion stage config still derived sequence_parallel_size under the v1 exclusivity rule (allgather_degree when allgather_degree>1, else ulysses_degree*ring_degree). With the composed Ulysses x AllGather-KV topology this contradicts the triple-product validation and fails startup (e.g. usp8+cp2: asserted 16, computed 8). Use the triple product here too. Signed-off-by: brandneway <gyuan4892@gmail.com>
|
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, model owner @brandneway, 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 triage noteAutomated triage of commit
These are automated triage suggestions only — the final decision belongs to the maintainers. |
hsliuustc0106
added
diffusion
Omni ReviewBot: resolvedThe CI failure noted on |
|
Thanks @brandneway and @Bounty-hunter. #7641 introduces the shared dense Ulysses × AllGather-KV topology and reports HunyuanImage3.0 DiT-only validation. #7937 cherry-picks its feature commit (547bfa7 → 00d2aaf), including the tests and benchmark, and adds MiniMax-H3 NPU fixes and validation. The two PRs therefore carry one shared implementation with distinct model evidence; #7937 explicitly supersedes #7889, but I have not seen a decision about #7641. Could you and the maintainers confirm which PR should carry the shared topology and the merge order? If #7641 lands first, #7937 could rebase to its MiniMax-H3/NPU fixes and tests. If #7937 is the integration PR, the cherry-picked commit already preserves #7641's authorship and tests; please confirm how its HunyuanImage3.0 validation should be referenced and what should happen to #7641. One separate code detail: 7973c35 edits an earlier |
Omni ReviewBot: no human activity for 7 days@brandneway this pull request has had no human commit, comment or review since 2026-09-23. Please confirm the current plan and next step. 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. |
|
@brandneway 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 cursor (cursor-grok-4.6-high) under experiment |
| value: torch.Tensor, | ||
| attn_metadata: AttentionMetadata | None, | ||
| ): | ||
| if attn_metadata is not None and attn_metadata.attn_mask is not None: |
There was a problem hiding this comment.
Qwen-Image passes its text padding mask in joint_attn_mask. With defer_joint=True, the Ulysses mask merge is skipped and AllGather ignores that field, so padded text keys become visible. Could we merge it after gathering, or reject it here too?
| # inherited AllGather-KV path slice the attention metadata, prepend | ||
| # the joint tensors, and build the local-Q/global-KV view. | ||
| self._assert_equal_region_lengths(key.shape[1], key.device) | ||
| query, key, value, attn_metadata, _ = super().pre_attention(query, key, value, attn_metadata) |
There was a problem hiding this comment.
MiniMax-H3 on CUDA still passes global cu_seqlens_q to flash_attn_varlen_func. Here Q has only S / allgather_degree rows, but the final query offset remains S. Please remap the packed query boundaries using query_ranges, or reject this backend/topology combination before calling the kernel.
Purpose
Bring up the Ulysses × AllGather-KV 2D diffusion SP topology and the MiniMax-H3 NPU serving path on top of it. This PR contains one feature commit cherry-picked from @dengyunyang's
cpanduspbranch (authorship preserved) plus three bugfixes discovered while bringing up the topology on Ascend NPU (16×950PR, MiniMax-H3 FL2VA).1. Ulysses × AllGather-KV 2D topology (cherry-picked from @dengyunyang's
cpandusp)Composes Ulysses (head sharding via all-to-all) with AllGather-KV (sequence gathering) as an orthogonal 2D SP topology:
The kernel always runs a local-Q/global-KV problem with a single, ordered global K/V tensor — which matters for attention backends that need stable global K/V indexing (e.g. sparse-mask generation), unlike Ring Attention. It also avoids the head-padding tax of large pure-Ulysses degrees: MiniMax-H3 has 56 attention heads, so usp16 needs advanced_uaa padding 56→64 (+14.3% attention work) while usp8+cp2 divides evenly.
UlyssesAllGatherKVParallelAttentionreuses the Ulysses strategy withdefer_joint=True(joint tensors stay head-sliced in metadata and are re-attached after the K/V gather — attaching earlier would replicate joint K/V once per AllGather rank).build_ulysses_allgather_rank_groupsbuilds the orthogonal groups (contiguous Ulysses blocks × stride-U AllGather groups; position p → ulysses_rank = p % U, allgather_rank = p // U).sequence_parallel_size = ulysses × ring × allgather; allgather composes with ulysses but not with ring.benchmarks/diffusion/bench_ulysses_cp_topology.py.2. [Bugfix] Default SP size derivation (
omni_engine_base.py)The default diffusion stage config still derived
sequence_parallel_sizeunder the v1 exclusivity rule (allgather_degreewhen allgather>1, elseulysses×ring). With the composed topology this contradicts the triple-product validation and fails startup (e.g. usp8+cp2 on 16 devices: asserted 16, computed 8).3. [Bugfix][MiniMax-H3] Text-encoder TP group tiling
_build_text_encoder_groupbuilt a single group covering only the firsttext_encoder_tp_sizeDiT ranks;GroupCoordinator.__init__asserts every constructing rank belongs to a group, sotp_size < world(e.g. tp=8 on 16 devices — the only legal topology for Qwen3-VL's 64 Q / 8 KV heads) crashed the remaining ranks at startup. The DiT world is now tiled intoworld / tp_sizecontiguous groups; only the first group encodes, and its output still reaches all DiT ranks through the existing world-group broadcast.tp_size == worldis unchanged.4. [Bugfix][NPU] Packed varlen contract generalized to local-Q/global-KV
Under AllGather-KV SP the backend sees local Q against global K/V, so the square packed contract (
used_q <= total_q) declined and the mask-rebuild fallback rejected the globalvalid_kv_lengthagainst the local query length — every packed H3 request failed with "packed metadata is unusable"._resolve_packed_seq_npunow accepts thequery_rangesrecorded by the AllGather-KV strategy: a single packed range means the local Q rows are a contiguous slice of the same global real document, so seq_q becomes the slice's[local_real, local_total]pair while seq_k keeps the global[real, pad]pair — same "pad is a strict suffix" contract, no mask materialized. Square (non-SP) behavior is unchanged; the mask fallback now checksvalid_kv_lengthagainst the K/V domain (key.shape[1]) as the safety net.Test Plan
tests/diffusion/distributed/test_parallel_state_sp_groups.py(orthogonal layout),tests/diffusion/attention/test_attention_sp.py(collective ordering mocks),tests/config/test_omni_config.py(triple-product validation).tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_group.py(encoder group tiling),tests/diffusion/attention/test_flash_attn.pygroup A2 (per-rank packed boundary resolution under AllGather-KV, unsupported-range rejection, square-contract preservation).Test Result
pytest tests/diffusion/distributed/test_parallel_state_sp_groups.py: 28 passed.pytest tests/diffusion/attention/test_flash_attn.py -k "packed_seq or cp_": 16 passed.Supersedes #7889.