Skip to content

[diffusion] feat: cross-node sequence parallelism (Ulysses x Ring) - #33327

Merged
mickqian merged 23 commits into
sgl-project:mainfrom
mickqian:codex/minimax-h3-cross-node-ring
Aug 8, 2026
Merged

mickqian merged 23 commits into
sgl-project:mainfrom
mickqian:codex/minimax-h3-cross-node-ring

Conversation

@mickqian

@mickqian mickqian commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add cross-node sequence parallelism for MiniMax-H3 by combining node-local Ulysses with cross-node Ring attention, so long Ref2VA/V2V sequences can scale past one node's GPU count.
  • Add --nnodes/--node-rank/--dist-init-addr launch infra (worker-only mode on non-head nodes, head node keeps the HTTP/TokenizerManager surface).
  • Pipeline the ring KV rotation's point-to-point hop with attention compute (double-buffered batch_isend_irecv) instead of a blocking all_gather.
  • Fix two ring-unaware call sites (denoise_loop.py's local embedding layout, minimax_h3.py's RoPE cache) found only under real cross-node re-validation after rebasing onto the current tree.
  • Document the verified 2-node H200 topology in the deployment picker and cookbook, replacing a stale "Ring is not compatible" claim.

Validation

  • Rebased cleanly onto this PR's tree with zero conflicts (verified: identical diffstat before/after rebase across the 5 touched Python files).
  • Real hardware: 2 nodes × 8× H200, --num-gpus 16 --nnodes 2 --sp-degree 16 --ulysses-degree 8 --ring-degree 2 --encoder-parallel replicate.
    • Functional: 4/4 Ref2VA/V2V requests (warmup + 3 cold) completed successfully end-to-end.
    • Determinism: the same request run twice against the same cross-node deployment produced byte-identical output.
    • A cross-node run is not expected to bit-match a single-node run of the same prompt/seed — Ring's online-softmax hop merge changes floating-point reduction order, same as flash vs. naive attention. This is expected, not a regression.
  • Performance (measured on this exact ring implementation prior to this rebase, same code confirmed unchanged by the zero-conflict rebase): denoise-stage-only comparison, 8× H200 single-node (Ulysses8) vs. the same 16-GPU cross-node command (Ulysses8×Ring2):
    • T2VA: 0.749s/step → 0.477s/step (−36.3%)
    • Ref2VA/V2V: 2.572s/step → 1.494s/step (−41.9%)
    • With KV-rotation pipelining: V2V full denoise stage 68.1–68.3s vs. 128.6s single-node baseline (−47.0%)

Known limitation

--encoder-parallel auto's fold decision is not yet node-boundary aware and will try to fold the text encoder across nodes, crashing the Ref2VA reference encoder. Cross-node deployments must pass --encoder-parallel replicate explicitly (documented in the cookbook and enforced nowhere in code yet — tracked as follow-up).

USP ring coverage (follow-up commits on this PR)

The H3 ring components were generalized (usp.py / parallel_state.py) and the remaining USPAttention branches were made ring-capable, closing silent-corruption holes found during validation (details in the validation comment):

USPAttention path production users ring status
plain wan-family, flux (divisible text) supported (pre-existing ring_attn)
packed varlen (H3) minimax_h3 supported (this PR's kernel)
tail-pad masked qwen, flux, flux_2 (non-divisible text) added — rides the H3 ring kernel via real_seq_len = pad_start (batch-1, FA)
replicated prefix / suffix / kv-prefix qwen short-text, zimage, hunyuan, mova added — sharded KV rotates, replicated KV LSE-merged as one local partial
generic varlen hunyuan packed multi-segment loud NotImplementedError (was silently corrupt)
explicit masks, batch>1, non-FA masked ideogram, CFG batches loud NotImplementedError

Also fixed (pre-existing on main): with u=1, r>1 the replicated-token dispatch keyed on the Ulysses world size, silently double-counting replicated text on the ring (qwen e2e MAE 15/255 → 0.58–0.74/255 after the fix).

Test plan

  • py_compile on all changed Python files
  • Real 2-node H200 functional smoke test (Ref2VA/V2V, warmup + 3 cold)
  • Real 2-node H200 determinism check (repeat request, byte-identical hash)
  • CI (opened as draft/ready per reviewer preference)

CI States

Latest PR Test (Base): ✅ Run #31231999230
Latest PR Test (Extra): ✅ Run #31231998922

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@github-actions github-actions Bot added documentation Improvements or additions to documentation diffusion SGLang Diffusion labels Aug 3, 2026
@mickqian mickqian changed the title [diffusion] MiniMax-H3: cross-node sequence parallelism (Ulysses x Ring) [diffusion] feat: cross-node sequence parallelism (Ulysses x Ring) Aug 3, 2026
@mickqian
mickqian force-pushed the codex/minimax-h3-cross-node-ring branch from 67eeaaa to a917ca5 Compare August 5, 2026 06:58
@mickqian

mickqian commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci label Aug 5, 2026
@mickqian

mickqian commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

Follow-up plan for which parallelism axes to take cross-node once this lands, ordered by expected ease (the per-axis communication analysis is now a section in #33704):

  1. DP, one replica per node — zero request-path traffic between replicas, so it needs no fast fabric at all. Blocker is small and concrete: [diffusion] feat: data-parallel serving (--dp-size) #33725's replica ingress binds on the local host, so multi-host placement needs per-replica host addressing (endpoint list gains a host per replica). Cheapest true cross-node capacity scaling for every model, including those ring cannot serve yet.
  2. CFG cross-node — one latent-sized branch-combine per step, deadline-tolerant by nature. First experiment after merge: a true-CFG wan case, 2 nodes × one branch each vs single-node cfg=2, ABAB-interleaved, comparing denoise ms/step (the combine cost is directly visible there) and e2e. If the delta is small this is the easiest cross-node latency win for CFG models — no attention changes at all.
  3. Ring per-model rollout — already the plan of record here: migrate the three legacy UlyssesAttention paths, generalize the H3-style admission ((H/T) % U, alignment % U×R) to launch-time checks for all USPAttention models, then per-family 2-node consistency + perf validation.
  4. Ulysses / TP stay intra-node — documented as such; no work intended.

@mickqian
mickqian force-pushed the codex/minimax-h3-cross-node-ring branch 3 times, most recently from 5a4a523 to 12a3341 Compare August 7, 2026 10:09
mickqian and others added 9 commits August 7, 2026 21:40
num_gpus stays the total world size across every node (srt's
tp_size-per-node convention); nnodes splits it evenly into local worker
counts, node_rank offsets into the global rank space. Node 0 keeps the
TokenizerManager/HTTP surface; other nodes run worker-only (mirrors
srt's node_rank>=1 path) and block until their local workers exit.
Rendezvous stays on loopback for nnodes==1 (unchanged single-node
behavior); --dist-init-addr provides a node-reachable address
otherwise. CPU intra-op thread budget now divides by the co-located
(local) worker count, not the global total.
Row split becomes 2D: ring first (an outer, contiguous ring_chunk_len
slice of the packed sequence), Ulysses second (an inner slice within
that chunk) -- only Ulysses shards heads, ring never does. Inside
attention, ring gathers every rank's local KV chunk once, attends each
chunk's real (non-padding) prefix locally against the full local query
set, and combines the per-chunk partial outputs via the standard
online-softmax (logsumexp) merge -- exact up to floating-point
reassociation, done in fp32. The real/padding boundary (used) is read
from the already-host-synced cu_seqlens_q_host tuple, never via a new
device sync, so the ring path doesn't regress the existing hot-loop
no-sync invariant.

ring_degree==1 is byte-identical to before this change (same row-split
arithmetic, same non-ring FA call) -- zero regression risk for the
production hash gate. ring_degree>1 changes accumulation order and is
therefore NOT expected to be bit-identical to a ring_degree==1 run of
the same request; validation is a same-config-determinism + SSIM/PSNR
quality gate instead, pending on-box confirmation.

Only a first-cut data-movement strategy (one all_gather of every rank's
KV chunk per attention call); a pipelined P2P rotation that overlaps
transfer with compute is a follow-up once correctness is established.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replace the correctness-first all_gather (blocking full-ring transfer
before any compute) with a rotating P2P double-buffer: each hop sends
this step's KV to the next ring rank and receives the following hop
from the previous rank before this step's attention compute is issued,
so the transfer overlaps compute instead of gating it. ring_degree=1
is unaffected (this path is only reachable when ring_active).
Stack K and V into a single buffer before each ring hop so the P2P
rotation is one send/recv pair instead of two, and drop the now-
redundant .contiguous() calls (torch.stack already materializes a
contiguous buffer). Same data moved, same overlap, no behavior change.
_build_local_embedding_layout's call site computed row_start/row_stop
from Ulysses world_size/rank alone (both here and in the preceding
block_token_tags slice) -- correct when ring_degree=1 (ring_rank is
always 0, so the combined formula degenerates to the Ulysses-only one,
which is why single-node testing didn't catch this), wrong once ring
actually shards rows across nodes. Combine world_size = ulysses*ring
and rank = ring_rank*ulysses_world_size + ulysses_rank (matching
minimax_h3.py's row_start formula and the process-group's ring-outer/
ulysses-inner rank layout) before both call sites.

Found via cross-node re-validation: a v2v request crashed with
'size of tensor a (5100) must match tensor b (6131)' inside
_embed's text-row copy, isolated to the 2-node (ring_degree=2)
config -- an identical request on 8 GPUs (ring_degree=1) with the
same code was unaffected, pointing straight at this ulysses-only
formula.
Same class of bug as the previous commit's local-embedding-layout fix:
build_rope_cache computed row_start from Ulysses world_size/rank alone
(named sp_ws/sp_rank but only ever holding the Ulysses component),
degenerating to the right answer only when ring_degree=1. Use the same
combined ring-outer/Ulysses-inner formula as forward()'s own row
split. Found by the same cross-node v2v re-test, one layer further in
after the previous fix: 'query, key and positions must have the same
number of tokens' inside _apply_rope_qk, i.e. the RoPE position cache
built for a larger (Ulysses-only) local_seq_len than the Q/K rows
actually ring-sharded down to.
Adds the verified 2-node H200 (Ulysses8 x Ring2) topology to the
deployment picker and cookbook: a new H200 "Cross-node" profile cell,
a feature-table row replacing the stale "Ring is not compatible" claim,
and a benchmark subsection with the measured per-step and pipelined
denoise-stage numbers plus the determinism/precision caveats.
Extends the general Sequence Parallelism doc with a cross-node section
(--nnodes/--node-rank/--dist-init-addr, the required --encoder-parallel
replicate workaround, and the cross-topology numerics caveat) plus a
reference benchmark using MiniMax-H3's measured cross-node numbers,
alongside the existing single-node Wan2.2 example.
The contract test asserted that ring_size=2 raises NotImplementedError, which is
exactly the limitation this branch removes, so it failed the moment it ran
against the implementation. It now checks what ring actually constrains: nothing
head-wise, since ring splits rows rather than heads, but the packed-sequence
alignment must divide the combined ulysses x ring degree, because ring adds an
outer row split on top of Ulysses's inner one.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

Ring numeric validation — done (2×B200)

Promised validation for the new ring paths, plus one pre-existing silent-corruption bug this surfaced and fixed.

The bug (exists on main, not introduced here)

With --ulysses-degree 1 --ring-degree 2, the replicated-token dispatch keyed on the Ulysses world size, so replicated text fell through to the plain ring path and was rotated as if sharded — every ring rank double-counted it. Measured on qwen-image e2e: MAE 15.2/255, SSIM 0.894 vs single GPU, silently. Three holes of the same shape: the dispatch gate, the u==1 concat shortcut in forward_with_replicated_kv_prefix, and plan_text_strategy's ring guard that routed non-divisible text into the hole. All fixed in ce8626712 (gate on the full SP group + ring execution for the replicated paths).

Ring execution added

path mechanism
tail-pad masked (qwen/flux/flux_2, batch-1, FA) padded shard rides the H3 ring kernel, real_seq_len = pad_start; pad rows zeroed like the Ulysses path
replicated prefix / suffix / kv-prefix sharded KV rotates via ring_attn (now exposes merged LSE); replicated KV joins as one local partial, two-term LSE merge in fp32
varlen (hunyuan-style) loud NotImplementedError (was silently corrupt)

Non-ring suffix keeps its bitwise tail-order contract; unsupported combos (batch>1, non-FA masked) raise loudly.

Validation (u1r2 vs single GPU / u2r1, qwen-image 1024², 20 steps, FA, seed 42)

2-rank kernel-level (vs fp32 SDPA reference, bf16): tail-pad + prefix + suffix + kv-prefix all MAE ~2e-4, max ~2.4e-3

e2e images (probe counters confirm each path engaged all 2400 block-calls):

comparison path MAE SSIM
u1r2 vs single (long prompt) tail-pad ring 0.580/255 0.9995
u2r1 vs single (reference) ulysses tail-pad 0.622/255 0.9994
u1r2 vs single (forced replicate) replicated-prefix ring 0.742/255 0.9977
u1r2 vs single — before fix (silent double-count) 15.2/255 0.894

Ring parity is at the same FP-reorder level as the long-validated Ulysses path. Unit tests (7) cover dispatch, merge math, and the tail-order contract.

/tag-and-rerun-ci

@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

The ring>1 fallback gathered the full unified sequence and skipped SP
attention, but its output is corrupt (silently: single vs u1r2 MAE
67/255, SSIM 0.25, deterministic on main). With the replicated-suffix
path now ring-capable, zimage takes the standard suffix route under
ring — validated at MAE 0.7/255 vs single GPU on qwen's twin path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

Per-model ring validation, round 2: zimage

Found another silent corruption on main: zimage's ring fallback (use_full_unified_sequence — full-sequence K/V gather + skip-SP attention) produces corrupt output: single vs u1r2 MAE 67/255, SSIM 0.25, bit-identical failure on main (12de7fb1f) and this branch, deterministic. The fallback had never been validated end-to-end.

Fix (ba94c9233): removed the fallback — with the replicated-suffix path now ring-capable, zimage takes the standard suffix route under ring.

After (Z-Image-Turbo 1024², 9 steps, seed 42, FA, 2×B200): probe confirms the suffix-ring path engaged all 540 layer-calls; single vs u1r2 MAE 1.343/255, SSIM 0.9986 — normal FP-reorder for a 9-step turbo model (few-step models amplify float reordering; the Ulysses-migration acceptance for zimage sat at the same level).

Per-model status so far: qwen ✅ (tail-pad ring 0.58/255 + replicated-prefix ring 0.74/255), zimage ✅ (0.79→suffix ring 1.34/255, main corruption fixed), wan (plain ring) running.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

Per-model ring validation — final matrix (2×B200, u1r2 vs single GPU, fixed seed)

model path exercised probe (layer-calls) MAE vs single verdict
qwen-image (long prompt) tail-pad masked ring 2400/2400 0.580/255 (SSIM 0.9995)
qwen-image (forced replicate) replicated-prefix ring 2400/2400 0.742/255 (SSIM 0.9977)
Z-Image-Turbo replicated-suffix ring 540/540 1.343/255 (SSIM 0.9986) ✅ (main's broken fallback removed, was 67/255)
Wan2.1-T2V-1.3B (33 frames) plain ring 2400/2400 2.700/255 all frames
minimax_h3 packed-varlen ring byte-identical repeat, 2-node H200 ✅ (original PR validation)

kv-prefix (mova) shares the same helper + merge mechanism as prefix/suffix and is covered by the 2-rank kernel test (MAE ~2e-4 vs fp32 SDPA); no e2e run — the weights are large and the code path is identical up to the rep/shard split point. hunyuan-style generic varlen and explicit-mask/batch>1 combos raise loudly by design.

Every MAE above is FP-reorder territory (few-step turbo amplifies it, consistent with the zimage Ulysses-migration acceptance). Unit suite: 30 passed.

CI note: the earlier multimodal-gen-unit-test red was test debt from lifting the ring text-shard guard (test_sp_shard mocked the removed symbol and asserted the removed behavior) — fixed and re-triggered.

@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci extra

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci extra

… era

The qwen_image_t2i_2_gpus GT was generated while u=1,r>1 silently
double-counted replicated text, so it froze the corrupt trajectory —
the correctness fix now reads as a consistency failure against it.
New GT published from this branch's CI output (ci-data-diffusion
211b2602f9, verified byte-identical), revision pin bumped, and the
override calibrated against the old GT dropped back to defaults.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci extra

@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

CI red root-caused: the GT froze the bug

The qwen_image_t2i_2_gpus failure (SSIM 0.70 vs threshold 0.79) is the correctness fix colliding with a ground truth generated during the double-count era — that CI case runs exactly the u1r2 + replicated-text combination this PR fixes, so its pinned GT is the corrupt trajectory. Both images in the failure artifact are coherent Doraemon renders (sampling-trajectory divergence, no corruption); the branch output is the one that matches a single GPU (0.74/255, validated upthread), the GT is the one that doesn't (main measures 15/255 vs single GPU on this path).

Handled per the standard GT-refresh flow, scoped to this PR so main stays green until merge:

  • New GT published to ci-data-diffusion 211b2602f9 from this branch's own CI output (byte-verified),
  • SGL_TEST_FILES_CI_DATA_REVISION bumped in this branch — other PRs keep reading the old pin/GT untouched,
  • the h100 threshold override for this case (calibrated against the buggy GT) dropped back to defaults.

The earlier multimodal-gen-test-1-b200 13s red was fast-fail collateral of this same case. CI re-triggered.

The consistency lookup prefers <platform>/<file> over the root file, so
the refreshed root-level GT was shadowed by the stale h100 copy.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci extra

@mickqian

mickqian commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

CI: green on NV ✅

All NVIDIA jobs pass, including qwen_image_t2i_2_gpus against the refreshed GT (byte-exact match, SSIM 1.0) — one wrinkle from the first refresh attempt: the consistency lookup prefers <platform>/<file> over the root file, so the GT had to be published under sglang_generated/h100/ (root copy was shadowed by the stale platform copy; pin now at eccf85dceb).

Remaining red is the AMD mi300 / NPU pool, which is failing identically on unrelated concurrent PRs (RL, HiCache) — pre-existing, not from this change.

Summary of what this PR now carries beyond the original H3 cross-node work:

  1. ring support for tail-pad masked + replicated prefix/suffix/kv-prefix paths (validated per-model upthread),
  2. two silent-corruption fixes that ship with it: the u=1,r>1 replicated dispatch hole (qwen 15/255) and zimage's broken ring fallback (67/255),
  3. the qwen u1r2 GT refresh, scoped to this branch via the revision pin so main stays green until merge.

Ready for review.

…helpers

usp.py's _ring_attention_varlen local-imported flash_attn_varlen_func,
parallel_state, and attention.backends every call — none of the three
create a cycle with usp.py, so hoisted to module level. denoise_loop.py's
_ulysses_ctx/_ring_ctx were a local-import duplicate of parallel_state's
get_ulysses_ctx/get_ring_ctx (added earlier this session for exactly this
purpose); deleted the duplicates and call the shared versions directly.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@mickqian

mickqian commented Aug 8, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci extra

attention.flash_attention sorts before diffusion.triton.ulysses_qkv
alphabetically; my earlier lint pass used bare 'ruff check' instead of
this repo's actual pre-commit scope (isort==7.0.0 is the import-sort
hook here, not ruff -- ruff is scoped to F401/F821/UP037 only), which
missed it. Verified against the exact CI tool versions this time.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@mickqian

mickqian commented Aug 8, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci extra

@mickqian
mickqian merged commit a25c330 into sgl-project:main Aug 8, 2026
123 of 131 checks passed
mickqian added a commit to mickqian/sglang that referenced this pull request Aug 8, 2026
…-node merge

Three corrections, all traced to PRs merged since this page was drafted:

- Ring now supports USPAttention's masked/tail-pad text path and its
  replicated-prefix/-suffix/-kv-prefix paths (sgl-project#33327) -- covers the joint
  text+image attention most models use. The page previously said these
  paths didn't support ring; that's no longer true.
- Cross-node launch (--nnodes/--node-rank/--dist-init-addr) is merged
  (sgl-project#33327), not pending. Kept the honest caveat: only H3 is validated
  end-to-end across nodes; other models' new ring support is same-node-
  validated so far.
- The legacy stacked-QKV UlyssesAttention set shrank to one user (Wan's VSA
  variant) after sgl-project#33923 migrated zimage/hunyuanvideo onto USPAttention;
  reworded from a migration-debt framing to the current, narrower and more
  precise one (a kernel without softmax LSE, not an unmigrated path).
- Fixed a since-reversed default: the page said Q/K/V packing was opt-in
  (enable_packed_qkv_input_a2a); sgl-project#33667 made single-collective packing the
  default, and that flag now selects the older stream-pipelined
  alternative instead.
- Ring's backend gate is now a declared capability
  (supports_ring_rotation(), sgl-project#33928) rather than a hardcoded enum check;
  reworded to match.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
saturn-acc pushed a commit to saturn-acc/sglang that referenced this pull request Aug 16, 2026
Atituiset pushed a commit to Atituiset/sglang that referenced this pull request Sep 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

diffusion SGLang Diffusion documentation Improvements or additions to documentation run-ci run-ci-extra

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant