From a25c330eb1514deff1fa39dc620ccd9b0ccb64ef Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 8 Aug 2026 10:56:28 +0800 Subject: [PATCH] [diffusion] feat: cross-node sequence parallelism (Ulysses x Ring) (#33327) Co-authored-by: Claude Fable 5 --- .../cookbook/diffusion/MiniMax/MiniMax-H3.mdx | 74 ++++++- .../sglang-diffusion/ring_sp_performance.mdx | 129 ++++++++++++ .../snippets/configs/MiniMaxAI/minimax-h3.jsx | 23 +++ .../runtime/distributed/parallel_state.py | 15 ++ .../runtime/distributed/sp_shard_utils.py | 3 - .../multimodal_gen/runtime/launch_server.py | 50 +++-- .../runtime/layers/attention/layer.py | 192 +++++++++++++----- .../multimodal_gen/runtime/layers/usp.py | 164 ++++++++++++++- .../runtime/managers/gpu_worker.py | 21 +- .../runtime/models/dits/minimax_h3.py | 168 ++++++++------- .../runtime/models/dits/zimage.py | 31 +-- .../minimax_h3/denoise_loop.py | 39 ++-- .../runtime/server_args/server_args.py | 40 ++++ .../server/consistency_thresholds/h100.json | 6 - .../sglang/multimodal_gen/test/test_utils.py | 2 +- .../test/unit/test_minimax_h3_denoise_loop.py | 2 +- .../test/unit/test_minimax_h3_dit_contract.py | 28 ++- .../multimodal_gen/test/unit/test_sp_shard.py | 13 +- .../test/unit/test_usp_ring_replicated.py | 191 +++++++++++++++++ .../test/unit/test_usp_ring_tail_pad.py | 86 ++++++++ 20 files changed, 1061 insertions(+), 216 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_usp_ring_tail_pad.py diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index 83ad5b723..d6eaf3eb0 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -65,7 +65,8 @@ behavior such as Cache-DiT is documented separately below. **Deployment Profile** exposes resident and FSDP placement on B200, B300, H200, and H100. Resident is the latency-oriented default; FSDP reduces DiT -weight residency at the cost of per-block parameter collectives. **Online +weight residency at the cost of per-block parameter collectives. On H200 it +also selects the verified 2-node cross-node topology. **Online Quantization** appears only on B200 and B300. AMD keeps its resident AITER recipe, while RTX 5090 uses its dedicated layerwise-offload profile. @@ -564,7 +565,8 @@ listed hardware and topology; it is not inherited by a similar GPU family. | Feature | Validation status | Notes | | --- | --- | --- | -| Ulysses sequence parallelism | Verified: 8× B200, 4× H200, 4× H100, and Ulysses1/2/4/8 on MI300X and MI355X | Use `--ulysses-degree`; Ring is not compatible with H3's packed multi-segment attention. | +| Ulysses sequence parallelism | Verified: 8× B200, 4× H200, 4× H100, and Ulysses1/2/4/8 on MI300X and MI355X | Use `--ulysses-degree`. Combine with Ring for cross-node scaling; see the next row. | +| Ring sequence parallelism (cross-node) | Verified: 2 nodes of 8× H200 each (Ulysses8 × Ring2) | Use `--ring-degree` together with `--nnodes`/`--node-rank`/`--dist-init-addr`. Ring shards the sequence across nodes while Ulysses shards heads within a node; H3's packed multi-segment attention only supports Ring across the node boundary, not within a single node's Ulysses group. Requires `--encoder-parallel replicate` — `auto`'s fold decision is not node-boundary aware. See the benchmark section below. | | SageAttention | Supported | Use `--attention-backend sage_attn` to select the native packed varlen path; install the SageAttention dependency first. | | Tensor parallelism | Verified: B200 TP2 + Ulysses4; H100 TP2 + Ulysses2 and TP4 + Ulysses1 | `--tp-size` may be combined with Ulysses when the TP-local head count remains divisible by the Ulysses degree. On 4×H100, TP2 + Ulysses2 is the measured speed default. | | FSDP inference | Verified: 4× B200 and 4× H100 + Ulysses4 | Preserves H3's mixed BF16/FP32 parameter policy. B200 completed the exact eager comparison; H100 completed consecutive real requests at about 57 GB peak memory per GPU. | @@ -673,7 +675,7 @@ the configurations with collected measurements: | --- | --- | --- | | B300 | 8× Ulysses8 resident | 8× FSDP + Ulysses8; the 8-GPU sweep is not a minimum-GPU claim. | | B200 | 8× Ulysses8 resident | 4× FSDP + Ulysses4 | -| H200 | 4× Ulysses4 resident | 4× FSDP + Ulysses4; 4× TP2 + Ulysses2 | +| H200 | 4× Ulysses4 resident | 4× FSDP + Ulysses4; 4× TP2 + Ulysses2; 2 nodes × 8× Ulysses8×Ring2 cross-node | | H100 | 4× TP2 + Ulysses2 resident | 4× TP4 + Ulysses1; 4× FSDP + Ulysses4 | | MI300X / MI355X | 8× Ulysses8 resident | 1×, 2×, and 4× scaling runs | | RTX 5090 | 2× TP2 + layerwise offload | — | @@ -779,6 +781,72 @@ GPU lower, which is why it remains the 80 GB H100 recipe. Matching the warmup request to the served resolution removes the cold first-request cost on both topologies (about 10 s end-to-end on this workload). +### H200 cross-node scaling + +Long references and long durations grow the packed sequence length, and +Ulysses alone cannot scale sequence parallelism past the GPU count of one +node without either violating head-count divisibility or exposing +all-to-all traffic across the slower inter-node link. H3 combines +node-local Ulysses with cross-node Ring: Ring's point-to-point KV rotation +is designed to overlap with attention compute, which fits a slower +cross-node link better than an all-to-all does. + +**Hardware.** 2 nodes × 8× NVIDIA H200 SXM, same cluster, InfiniBand +between nodes. + +**Serve command.** The cross-node cell the picker emits for H200, run +identically on both nodes with `--node-rank` set to 0 and 1: + +```bash Command +sglang serve \ + --model-path MiniMaxAI/MiniMax-H3 \ + --model-variant ref2va \ + --num-gpus 16 \ + --nnodes 2 \ + --node-rank {{NODE_RANK}} \ + --dist-init-addr {{NODE0_IP}}:20000 \ + --sp-degree 16 \ + --ulysses-degree 8 \ + --ring-degree 2 \ + --encoder-parallel replicate \ + --performance-mode speed \ + --host 0.0.0.0 \ + --port 30010 +``` + +**What was measured.** A controlled denoise-stage comparison on identical +hardware: 8× H200 single-node (Ulysses8, no Ring) versus the same 16-GPU +cross-node command above (Ulysses8 × Ring2), holding prompt, seed, and +step count fixed: + +| Task | Single-node (Ulysses8) | Cross-node (Ulysses8 × Ring2) | Change | +| --- | ---: | ---: | ---: | +| T2VA denoise/step | 0.749 s | 0.477 s | −36.3% | +| Ref2VA/V2V denoise/step | 2.572 s | 1.494 s | −41.9% | + +The gain grows with sequence length because Ring's per-hop communication +cost stays roughly constant while attention compute grows quadratically +with sequence length, so V2V's longer packed sequence benefits more than +T2VA's shorter one. With the point-to-point KV rotation pipelined against +attention compute, one V2V request's full denoise stage completed in +68.1–68.3 seconds versus 128.6 seconds on the single-node 8-GPU baseline +(−47.0%), with byte-identical output to the unpipelined cross-node path. + +Cross-node determinism was confirmed separately: the same request run +twice against the same cross-node deployment produced byte-identical +output. A cross-node run's output is not expected to bit-match a +single-node run of the same prompt and seed — Ring's online-softmax merge +across hops accumulates floating-point operations in a different order +than single-node attention, which is an expected source of bit-level +difference, not a correctness regression. + + +`--encoder-parallel auto`'s fold decision is not yet node-boundary aware +and attempts to fold the text encoder across nodes, which crashes the +Ref2VA reference-conditioned encoder. Always pass +`--encoder-parallel replicate` explicitly for cross-node H3 deployments. + + ### H100 topology comparison The same four-card H100 host completed three lossless placements. TP2 with diff --git a/docs/docs/sglang-diffusion/ring_sp_performance.mdx b/docs/docs/sglang-diffusion/ring_sp_performance.mdx index 0d337c0cb..c81c025d0 100644 --- a/docs/docs/sglang-diffusion/ring_sp_performance.mdx +++ b/docs/docs/sglang-diffusion/ring_sp_performance.mdx @@ -245,6 +245,82 @@ sglang serve \ +## Cross-Node Sequence Parallelism + +Ulysses alone cannot scale sequence parallelism past the GPU count of one +node: going wider either violates head-count divisibility or exposes an +all-to-all across the slower inter-node link. Ring's point-to-point KV +rotation is designed to overlap with attention compute, which tolerates a +slower cross-node link far better than an all-to-all does — so the pattern +for scaling SP across nodes is **node-local Ulysses × cross-node Ring**, not +Ulysses alone. + +Cross-node launches add three flags on top of the usual SP degrees: + +- `--nnodes`: number of nodes. `--num-gpus` stays the *total* GPU count + across every node; each node runs `num_gpus // nnodes` local workers. +- `--node-rank`: this node's rank, `0` on the head node (which keeps the + HTTP/TokenizerManager surface) and `1..nnodes-1` on the others (worker-only). +- `--dist-init-addr`: a `host:port` rendezvous address reachable from every + node — typically the head node's address. + +Run the same command on every node, changing only `--node-rank`: + +```bash +# node 0 (head) +sglang serve \ + --model-path MiniMaxAI/MiniMax-H3 \ + --model-variant ref2va \ + --num-gpus 16 \ + --nnodes 2 \ + --node-rank 0 \ + --dist-init-addr :23456 \ + --sp-degree 16 \ + --ulysses-degree 8 \ + --ring-degree 2 \ + --encoder-parallel replicate \ + --port 30010 + +# node 1 (worker) +sglang serve \ + --model-path MiniMaxAI/MiniMax-H3 \ + --model-variant ref2va \ + --num-gpus 16 \ + --nnodes 2 \ + --node-rank 1 \ + --dist-init-addr :23456 \ + --sp-degree 16 \ + --ulysses-degree 8 \ + --ring-degree 2 \ + --encoder-parallel replicate \ + --port 30010 +``` + +`--encoder-parallel replicate` is required for cross-node deployments today: +the `auto` fold decision is not yet node-boundary aware and will try to fold +the text encoder across nodes, which crashes reference-conditioned encoders. +See [Encoder Parallelism](/docs/sglang-diffusion/encoder_parallel). + + +Cross-node ring support is model-specific, not a property of the launch +flags alone — see [Supported Models and Optimization Compatibility](/docs/sglang-diffusion/compatibility_matrix) +for which models have it. Passing `--ring-degree > 1` for a model that only +has single-node Ulysses may either raise or, in some cases, silently +compute incorrect output; check the model's cookbook page before assuming +cross-node scaling is supported. + + +### Numerics across node boundaries + +Ring's online-softmax merge across P2P hops accumulates floating-point +operations in a different order than single-node attention, so a cross-node +run is **not** expected to bit-match a single-node run of the same prompt and +seed — this is the same class of difference as choosing a different +attention backend, not a correctness regression. What *is* expected: the same +request run twice against the same cross-node deployment must produce +byte-identical output. Use that repeat-request check, not a cross-topology +comparison, to validate a cross-node deployment's determinism. + ## Benchmarking Guidance When benchmarking SP, compare the same model, precision, resolution, frame count, step count, scheduler settings, prompt type, and output path. Report both stage latency and peak GPU memory; SP can reduce per-GPU memory while adding communication overhead. @@ -375,3 +451,56 @@ The following numbers are a reference measurement for one setup. They are not a In this setup, end-to-end latency improved from `90.63s` to `63.74s` (`1.42x`) and peak GPU memory dropped by `7.33GB`. The overhead ratio increased, so future tuning should still check communication and runtime overhead on the target hardware. + +## Cross-Node Reference Benchmark + +The following numbers are a reference measurement for MiniMax-H3's cross-node +Ulysses × Ring deployment. They are not a general promise for every +model or topology — see each model's cookbook page for its own verified +cross-node status. + +- Model: `MiniMaxAI/MiniMax-H3` +- Hardware: 2 nodes × 8× NVIDIA H200 SXM, same cluster, InfiniBand between nodes +- Cross-node config: `--num-gpus 16 --sp-degree 16 --ulysses-degree 8 --ring-degree 2` +- Single-node baseline: `--num-gpus 8 --sp-degree 8 --ulysses-degree 8 --ring-degree 1` + +Denoise-stage-only comparison, holding prompt, seed, and step count fixed: + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
TaskSingle-node (s/step)Cross-node (s/step)Change
T2VA denoise0.7490.477-36.3%
Ref2VA / V2V denoise2.5721.494-41.9%
+ +The gain grows with sequence length: ring's per-hop communication cost stays +roughly constant while attention compute grows quadratically with sequence +length, so V2V's longer packed sequence benefits more than T2VA's shorter +one. With the point-to-point KV rotation pipelined against attention compute +(rather than a blocking `all_gather`), one V2V request's full denoise stage +completed in 68.1-68.3s versus 128.6s on the single-node 8-GPU baseline +(-47.0%), with byte-identical output to the unpipelined cross-node path. diff --git a/docs/src/snippets/configs/MiniMaxAI/minimax-h3.jsx b/docs/src/snippets/configs/MiniMaxAI/minimax-h3.jsx index fc2579ecd..74b902fef 100644 --- a/docs/src/snippets/configs/MiniMaxAI/minimax-h3.jsx +++ b/docs/src/snippets/configs/MiniMaxAI/minimax-h3.jsx @@ -44,6 +44,11 @@ export const config = { label: "Layerwise offload", showWhen: (s) => s.hw === "rtx5090", }, + { + id: "cross_node", + label: "Cross-node (2 nodes)", + showWhen: (s) => s.hw === "h200", + }, ], }, ], @@ -512,6 +517,24 @@ export const config = { warn: "FSDP reduces resident DiT memory but adds per-block parameter collectives. Prefer Resident when the full pipeline fits.", }, + { + match: { hw: "h200", profile: "cross_node" }, + nnodes: 2, + verified: true, + flags: [ + "--model-path {{MODEL_NAME}}", + "--num-gpus 16", + "--sp-degree 16", + "--ulysses-degree 8", + "--ring-degree 2", + "--encoder-parallel replicate", + "--performance-mode speed", + "--host {{HOST_IP}}", + "--port {{PORT}}", + ], + warn: + "Verified on 2 nodes of 8× H200 each (Ulysses8 within a node, Ring2 across nodes). Requires --encoder-parallel replicate: --encoder-parallel auto's fold decision is not yet node-boundary aware and will crash across nodes.", + }, { match: { hw: "b200", profile: "fsdp" }, nnodes: 1, diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index e5165b81e..2200c28bd 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -774,6 +774,21 @@ def get_ring_parallel_rank() -> int: return get_sp_group().ring_rank +def get_ulysses_ctx() -> tuple[int, int]: + """(world_size, rank) of the Ulysses group; (1, 0) when uninitialized + (unit tests / single-process debug paths).""" + if not model_parallel_is_initialized(): + return 1, 0 + return get_ulysses_parallel_world_size(), get_ulysses_parallel_rank() + + +def get_ring_ctx() -> tuple[int, int]: + """(world_size, rank) of the Ring group; (1, 0) when uninitialized.""" + if not model_parallel_is_initialized(): + return 1, 0 + return get_ring_parallel_world_size(), get_ring_parallel_rank() + + # PP def get_pp_group() -> PipelineGroupCoordinator: assert _PP is not None, "pipeline model parallel group is not initialized" diff --git a/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py b/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py index 23b62e586..bd603d245 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py +++ b/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py @@ -19,7 +19,6 @@ from sglang.multimodal_gen.runtime.distributed.communication_op import ( sequence_model_parallel_all_gather, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( - get_ring_parallel_world_size, get_sp_parallel_rank, get_sp_world_size, ) @@ -223,8 +222,6 @@ def plan_text_strategy(txt_len: int) -> str: # padding must fit in the final shard to remain one global-tail block if num_pad > local_len: return "replicate" - if txt_len % sp_size != 0 and get_ring_parallel_world_size() > 1: - return "replicate" if txt_len < _TEXT_SHARD_MIN: return "replicate" return "shard" diff --git a/python/sglang/multimodal_gen/runtime/launch_server.py b/python/sglang/multimodal_gen/runtime/launch_server.py index e299925ef..85fcbeddb 100644 --- a/python/sglang/multimodal_gen/runtime/launch_server.py +++ b/python/sglang/multimodal_gen/runtime/launch_server.py @@ -178,39 +178,48 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True): # Start a new server with multiple worker processes logger.info("Starting server...") + # num_gpus is the total world size across every node; each node runs + # its own num_gpus // nnodes local workers, offset by node_rank into the + # global rank space (mirrors srt's tp_size_per_node convention). With + # nnodes == 1 this is exactly the prior single-node arithmetic. num_gpus = server_args.num_gpus + nnodes = server_args.nnodes + node_rank = server_args.node_rank + local_num_gpus = num_gpus // nnodes + rank_offset = node_rank * local_num_gpus processes = [] - # Pipes for master to talk to slaves + # Pipes for master to talk to slaves (local to this node) task_pipes_to_slaves_w = [] task_pipes_to_slaves_r = [] - for _ in range(num_gpus - 1): + for _ in range(local_num_gpus - 1): r, w = mp.Pipe(duplex=False) task_pipes_to_slaves_r.append(r) task_pipes_to_slaves_w.append(w) - # Pipes for slaves to talk to master + # Pipes for slaves to talk to master (local to this node) result_pipes_from_slaves_w = [] result_pipes_from_slaves_r = [] - for _ in range(num_gpus - 1): + for _ in range(local_num_gpus - 1): r, w = mp.Pipe(duplex=False) result_pipes_from_slaves_r.append(r) result_pipes_from_slaves_w.append(w) - # Launch all worker processes + # Launch this node's local worker processes master_port = server_args.master_port scheduler_pipe_readers = [] scheduler_pipe_writers = [] - for i in range(num_gpus): + for i in range(local_num_gpus): + rank = rank_offset + i reader, writer = mp.Pipe(duplex=False) scheduler_pipe_writers.append(writer) - if i == 0: # Master worker + if i == 0: # This node's local pipe master process = mp.Process( target=run_scheduler_process, args=( i, # local_rank - i, # rank + rank, master_port, server_args, writer, @@ -219,7 +228,7 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True): task_pipes_to_slaves_w, result_pipes_from_slaves_r, ), - name=f"sglang-diffusionWorker-{i}", + name=f"sglang-diffusionWorker-{rank}", daemon=True, ) else: # Slave workers @@ -227,7 +236,7 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True): target=run_scheduler_process, args=( i, # local_rank - i, # rank + rank, master_port, server_args, writer, @@ -236,7 +245,7 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True): task_pipes_to_slaves_r[i - 1], result_pipes_from_slaves_w[i - 1], ), - name=f"sglang-diffusionWorker-{i}", + name=f"sglang-diffusionWorker-{rank}", daemon=True, ) scheduler_pipe_readers.append(reader) @@ -263,7 +272,8 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True): data = reader.recv() except EOFError: logger.error( - f"Rank {i} scheduler is dead. Please check if there are relevant logs." + f"Rank {rank_offset + i} scheduler is dead. Please check if " + "there are relevant logs." ) processes[i].join() logger.error(f"Exit code: {processes[i].exitcode}") @@ -278,6 +288,22 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True): logger.debug("All workers are ready") + if node_rank != 0: + # The TokenizerManager / HTTP surface lives on the node that owns + # global rank 0; this node only hosts local workers, which tear + # down together with the distributed group on shutdown. + logger.info( + "Node %d ready with %d local worker(s); no local HTTP surface.", + node_rank, + local_num_gpus, + ) + try: + for p in processes: + p.join() + finally: + shutdown_scheduler_processes(None, processes, request_shutdown=False) + return processes + if launch_http_server: if server_args.pipeline_config.task_type.is_action_gen(): logger.info( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 811409b89..f7e7ba813 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -25,11 +25,13 @@ from sglang.multimodal_gen.runtime.distributed.communication_op import ( sequence_model_parallel_all_to_all_4D, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_ring_parallel_rank, get_ring_parallel_world_size, get_sequence_parallel_world_size, get_sp_group, get_sp_parallel_rank, get_sp_world_size, + get_ulysses_parallel_rank, get_ulysses_parallel_world_size, ) from sglang.multimodal_gen.runtime.layers.attention.backends import ( @@ -45,6 +47,8 @@ from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import ( ) from sglang.multimodal_gen.runtime.layers.usp import ( _ipc_input_a2a_qkv, + _merge_attention_partials, + _ring_attention_varlen, _usp_input_all_to_all, _usp_input_all_to_all_qkv, _usp_input_all_to_all_varlen, @@ -786,6 +790,13 @@ class USPAttention(nn.Module): ), "Varlen USPAttention does not support masks or replicated tokens" if effective_skip_sp or get_sequence_parallel_world_size() == 1: return self.attn_impl.forward(q, k, v, ctx_attn_metadata) + if get_ring_parallel_world_size() > 1: + # The varlen all-to-all spans the combined SP group and is not + # ring-aware; it would shuffle rows across ring ranks instead + # of rotating KV, corrupting the output silently. + raise NotImplementedError( + "Varlen USPAttention does not support ring parallelism yet." + ) qkv = torch.cat([q, k, v], dim=0) qkv = _usp_input_all_to_all_varlen(qkv, seq_lens, head_dim=2) qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata) @@ -955,8 +966,15 @@ class USPAttention(nn.Module): ).transpose(1, 2) if get_ring_parallel_world_size() > 1: + if ( + meta_only_pad + and q.shape[0] == 1 + and self.backend == AttentionBackendEnum.FA + ): + return self._forward_ring_tail_pad(q, k, v, attn_mask_meta) raise NotImplementedError( - "USPAttention masked path does not support ring parallelism yet." + "USPAttention masked path supports ring parallelism only " + "for batch-1 tail-pad metadata on the FA backend." ) if attn_mask is not None and attn_mask.dim() != 2: raise NotImplementedError( @@ -1132,15 +1150,19 @@ class USPAttention(nn.Module): raise ValueError( "USPAttention supports at most one replicated-token mode per call." ) - if sp_size > 1 and num_replicated_prefix > 0: + # Replicated-token handling is keyed on the full SP group: with u=1, + # r>1 the plain ring path would rotate the replicated tokens as if + # they were sharded rows, double-counting them. + sp_ws = get_sequence_parallel_world_size() + if sp_ws > 1 and num_replicated_prefix > 0: return self._forward_with_replicated_prefix( q, k, v, ctx_attn_metadata, num_replicated_prefix ) - if sp_size > 1 and num_replicated_suffix > 0: + if sp_ws > 1 and num_replicated_suffix > 0: return self._forward_with_replicated_suffix( q, k, v, ctx_attn_metadata, num_replicated_suffix ) - if sp_size > 1 and num_replicated_kv_prefix > 0: + if sp_ws > 1 and num_replicated_kv_prefix > 0: return self._forward_with_replicated_kv_prefix( q, k, v, ctx_attn_metadata, num_replicated_kv_prefix ) @@ -1183,6 +1205,39 @@ class USPAttention(nn.Module): return out + def _forward_ring_tail_pad( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attn_mask_meta: dict, + ) -> torch.Tensor: + """Ring attention for a tail-padded shard. + + After the Ulysses all-to-all each ring rank holds one contiguous block + of the gathered sequence, and the tail-pad invariant keeps all padding + at the global tail — exactly the ring kernel's real-length clamp, so + no masks or repacking are needed. Pad rows receive garbage output, + which the tail-pad consumers already trim. + """ + q, k, v = _usp_input_all_to_all_qkv(q, k, v) + out = _ring_attention_varlen( + q.squeeze(0), + k.squeeze(0), + v.squeeze(0), + softmax_scale=self.softmax_scale, + real_seq_len=int(attn_mask_meta["pad_start"]), + ring_ws=get_ring_parallel_world_size(), + ) + # Match the Ulysses tail path: masked query rows read as zeros. This + # rank's chunk covers global rows [rank*chunk, (rank+1)*chunk). + pad_from = ( + int(attn_mask_meta["pad_start"]) - get_ring_parallel_rank() * out.shape[0] + ) + if pad_from < out.shape[0]: + out[max(pad_from, 0) :].zero_() + return _usp_output_all_to_all(out.unsqueeze(0), head_dim=2) + @staticmethod def _gather_sharded_sequence( tensor: torch.Tensor, @@ -1389,13 +1444,8 @@ class USPAttention(nn.Module): 4. Concatenate [prefix_h_local, gathered_suffix] and run attention. 5. Split output, all-to-all back the suffix, all-gather prefix heads. """ - if get_ring_parallel_world_size() > 1: - raise NotImplementedError( - "USPAttention replicated-prefix/suffix path does not support " - "ring parallelism yet." - ) sp_size = get_ulysses_parallel_world_size() - sp_rank = get_sp_parallel_rank() + u_rank = get_ulysses_parallel_rank() q_rep, q_shard = q[:, :num_rep], q[:, num_rep:] k_rep, k_shard = k[:, :num_rep], k[:, num_rep:] @@ -1410,33 +1460,73 @@ class USPAttention(nn.Module): # For MHA (kv heads == q heads) this is identical to the q shard. h_local = q_shard.shape[2] kv_h_local = k_shard.shape[2] - h_start = sp_rank * h_local - kv_h_start = sp_rank * kv_h_local + h_start = u_rank * h_local + kv_h_start = u_rank * kv_h_local q_rep = q_rep[:, :, h_start : h_start + h_local, :].contiguous() k_rep = k_rep[:, :, kv_h_start : kv_h_start + kv_h_local, :].contiguous() v_rep = v_rep[:, :, kv_h_start : kv_h_start + kv_h_local, :].contiguous() q = torch.cat([q_rep, q_shard], dim=1) - k = torch.cat([k_rep, k_shard], dim=1) - v = torch.cat([v_rep, v_shard], dim=1) - - out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + out = self._replicated_kv_attention( + q, k_shard, v_shard, k_rep, v_rep, ctx_attn_metadata + ) out_rep = out[:, :num_rep] out_shard = out[:, num_rep:] out_shard = _usp_output_all_to_all(out_shard, head_dim=2) - gathered = [torch.empty_like(out_rep) for _ in range(sp_size)] - torch.distributed.all_gather( - gathered, - out_rep.contiguous(), - group=get_sp_group().ulysses_group, - ) - out_rep = torch.cat(gathered, dim=2) + if sp_size > 1: + gathered = [torch.empty_like(out_rep) for _ in range(sp_size)] + torch.distributed.all_gather( + gathered, + out_rep.contiguous(), + group=get_sp_group().ulysses_group, + ) + out_rep = torch.cat(gathered, dim=2) return torch.cat([out_rep, out_shard], dim=1) + def _replicated_kv_attention( + self, + q: torch.Tensor, + k_shard: torch.Tensor, + v_shard: torch.Tensor, + k_rep: torch.Tensor, + v_rep: torch.Tensor, + ctx_attn_metadata, + rep_first: bool = True, + ) -> torch.Tensor: + """Attention of q against replicated + ring-sharded KV. + + Without ring parallelism the two KV parts concatenate into one local + kernel call, replicated part first unless `rep_first=False` (the + suffix path keeps KV in tail order for bitwise stability). Under ring + parallelism the sharded KV rotates around the ring while the + replicated KV contributes one extra local partial, LSE-merged with + the ring result (exact up to float reordering). + """ + if get_ring_parallel_world_size() > 1: + out_ring, lse_ring = ring_attn( + q, k_shard, v_shard, self.attn_impl, return_softmax_lse=True + ) + out_rep, lse_rep, *_ = self.attn_impl.forward( + q, k_rep, v_rep, attn_metadata=None, return_softmax_lse=True + ) + merged = _merge_attention_partials(out_ring, lse_ring, out_rep, lse_rep) + return merged.to(q.dtype) + kv_parts = ( + ([k_rep, k_shard], [v_rep, v_shard]) + if rep_first + else ( + [k_shard, k_rep], + [v_shard, v_rep], + ) + ) + k = torch.cat(kv_parts[0], dim=1) + v = torch.cat(kv_parts[1], dim=1) + return self.attn_impl.forward(q, k, v, ctx_attn_metadata) + def forward_with_replicated_kv_prefix( self, q: torch.Tensor, @@ -1461,7 +1551,10 @@ class USPAttention(nn.Module): v = torch.cat([v_prefix, v_suffix], dim=1) return self.attn_impl.forward(q, k, v, ctx_attn_metadata) - if get_ulysses_parallel_world_size() == 1: + if ( + get_ulysses_parallel_world_size() == 1 + and get_ring_parallel_world_size() == 1 + ): k = torch.cat([k_prefix, k_suffix], dim=1) v = torch.cat([v_prefix, v_suffix], dim=1) return self(q, k, v) @@ -1508,14 +1601,9 @@ class USPAttention(nn.Module): ctx_attn_metadata, ) -> torch.Tensor: """split form avoids materializing full K/V before Ulysses all-to-all""" - if get_ring_parallel_world_size() > 1: - raise NotImplementedError( - "USPAttention replicated-kv-prefix path does not support " - "ring parallelism yet." - ) - sp_rank = get_sp_parallel_rank() + u_rank = get_ulysses_parallel_rank() - if q.device.type == "cuda": + if q.device.type == "cuda" and get_ulysses_parallel_world_size() > 1: q, k_shard, v_shard = async_a2a_communicate( [q, k_shard, v_shard], get_ulysses_parallel_world_size(), @@ -1532,15 +1620,14 @@ class USPAttention(nn.Module): v_shard = _usp_input_all_to_all(v_shard, head_dim=2) h_kv_local = k_shard.shape[2] - h_start = sp_rank * h_kv_local + h_start = u_rank * h_kv_local h_end = h_start + h_kv_local k_rep = k_rep[:, :, h_start:h_end, :].contiguous() v_rep = v_rep[:, :, h_start:h_end, :].contiguous() - k = torch.cat([k_rep, k_shard], dim=1) - v = torch.cat([v_rep, v_shard], dim=1) - - out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + out = self._replicated_kv_attention( + q, k_shard, v_shard, k_rep, v_rep, ctx_attn_metadata + ) return _usp_output_all_to_all(out, head_dim=2) def _forward_with_replicated_suffix( @@ -1561,12 +1648,7 @@ class USPAttention(nn.Module): """ if num_rep <= 0: raise ValueError("num_rep must be positive for replicated suffix.") - if get_ring_parallel_world_size() > 1: - raise NotImplementedError( - "USPAttention replicated-prefix/suffix path does not support " - "ring parallelism yet." - ) - sp_rank = get_sp_parallel_rank() + u_rank = get_ulysses_parallel_rank() q_shard, q_rep = q[:, :-num_rep], q[:, -num_rep:] k_shard, k_rep = k[:, :-num_rep], k[:, -num_rep:] @@ -1578,17 +1660,16 @@ class USPAttention(nn.Module): h_local = q_shard.shape[2] kv_h_local = k_shard.shape[2] - h_start = sp_rank * h_local - kv_h_start = sp_rank * kv_h_local + h_start = u_rank * h_local + kv_h_start = u_rank * kv_h_local q_rep = q_rep[:, :, h_start : h_start + h_local, :].contiguous() k_rep = k_rep[:, :, kv_h_start : kv_h_start + kv_h_local, :].contiguous() v_rep = v_rep[:, :, kv_h_start : kv_h_start + kv_h_local, :].contiguous() q = torch.cat([q_shard, q_rep], dim=1) - k = torch.cat([k_shard, k_rep], dim=1) - v = torch.cat([v_shard, v_rep], dim=1) - - out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + out = self._replicated_kv_attention( + q, k_shard, v_shard, k_rep, v_rep, ctx_attn_metadata, rep_first=False + ) out_shard = out[:, :-num_rep] out_rep = out[:, -num_rep:] @@ -1596,13 +1677,14 @@ class USPAttention(nn.Module): out_shard = _usp_output_all_to_all(out_shard, head_dim=2) sp_size = get_ulysses_parallel_world_size() - gathered = [torch.empty_like(out_rep) for _ in range(sp_size)] - torch.distributed.all_gather( - gathered, - out_rep.contiguous(), - group=get_sp_group().ulysses_group, - ) - out_rep = torch.cat(gathered, dim=2) + if sp_size > 1: + gathered = [torch.empty_like(out_rep) for _ in range(sp_size)] + torch.distributed.all_gather( + gathered, + out_rep.contiguous(), + group=get_sp_group().ulysses_group, + ) + out_rep = torch.cat(gathered, dim=2) return torch.cat([out_shard, out_rep], dim=1) diff --git a/python/sglang/multimodal_gen/runtime/layers/usp.py b/python/sglang/multimodal_gen/runtime/layers/usp.py index 461dcad59..e40313cc2 100644 --- a/python/sglang/multimodal_gen/runtime/layers/usp.py +++ b/python/sglang/multimodal_gen/runtime/layers/usp.py @@ -8,15 +8,20 @@ import torch.distributed as dist import torch.distributed._functional_collectives as ft_c from torch.distributed.tensor.experimental._attention import _cp_options +from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func from sglang.kernels.ops.diffusion.triton.ulysses_qkv import ( pack_qkv_destination_major, ) from sglang.kernels.ops.diffusion.usp_relayout import usp_merge_heads from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_ring_ctx, get_sp_group, get_ulysses_parallel_rank, get_ulysses_parallel_world_size, ) +from sglang.multimodal_gen.runtime.layers.attention.backends import ( + flash_attn as _fa_backend, +) from sglang.srt.utils.common import torch_release _cp_options.enable_load_balance = False @@ -675,6 +680,7 @@ def ring_attn( attn_impl: "AttentionImpl", is_causal: bool = False, dropout_p: float = 0.0, + return_softmax_lse: bool = False, ): """ Ring Attention implementation. @@ -748,15 +754,169 @@ def ring_attn( if use_segment_id: # For torch >= 2.6, segment_id is required. The value '1' is a placeholder # as we are not using complex segmentation features. - out, *_ = _templated_ring_attention( + out, lse, *_ = _templated_ring_attention( seq_dim=1, # segment_id **attn_kwargs, ) else: - out, *_ = _templated_ring_attention( + out, lse, *_ = _templated_ring_attention( **attn_kwargs, ) # Permute the output back to [B, S, H, D] layout. output = torch.permute(out, [0, 2, 1, 3]) + if return_softmax_lse: + return output, lse return output + + +def _merge_attention_partials( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, +) -> torch.Tensor: + """Merge two attention partials computed over disjoint KV sets. + + `out_*` are `[B, S, H, D]`; `lse_*` are the dense-FA LSE layout `[B, H, S]`. + Each partial is self-normalized over its own KV, so the exact combine is a + two-term logsumexp reweighting; done in fp32 for stability. + """ + lse_a = lse_a.transpose(1, 2).unsqueeze(-1).to(torch.float32) + lse_b = lse_b.transpose(1, 2).unsqueeze(-1).to(torch.float32) + new_lse = torch.logaddexp(lse_a, lse_b) + return out_a.to(torch.float32) * torch.exp(lse_a - new_lse) + out_b.to( + torch.float32 + ) * torch.exp(lse_b - new_lse) + + +def _ring_merge_attention( + out_acc: torch.Tensor | None, + lse_acc: torch.Tensor | None, + step_out: torch.Tensor, + step_lse: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Online-softmax combine of one more ring step's partial attention. + + `step_out` is `[T, H, D]`; `step_lse` is FlashAttention's varlen LSE + layout `[H, T]`. Both are already self-normalized over their own KV + chunk, so combining chunks is the standard two-term logsumexp merge + (exact up to float rounding); done in fp32 for stability. + """ + step_lse = step_lse.transpose(0, 1).unsqueeze(-1).to(torch.float32) + step_out = step_out.to(torch.float32) + if out_acc is None: + return step_out, step_lse + new_lse = torch.logaddexp(lse_acc, step_lse) + out_acc = out_acc * torch.exp(lse_acc - new_lse) + step_out * torch.exp( + step_lse - new_lse + ) + return out_acc, new_lse + + +def _ring_attention_varlen( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + softmax_scale: float, + real_seq_len: int, + ring_ws: int, +) -> torch.Tensor: + """Ring-rotated varlen attention over one rank's local packed chunk. + + `q, k, v` are this rank's full local ring chunk (`ring_chunk_len` rows, + real rows followed by however many of this chunk's rows are padding). + KV is P2P-rotated around the ring one hop per step (send this step's + buffer to the next rank, receive the following step's buffer from the + previous rank) so the hop for step+1 overlaps this step's attention + compute -- unlike a single blocking all_gather, no step waits on the + full ring's transfer before its own compute can start. Each hop's + *real* prefix is attended locally and merged via online softmax. + Padding rows never contribute to any KV chunk (their output is unused + downstream, but must not be corrupted by attending across the padding + boundary), and a chunk that is entirely padding is skipped outright. + """ + ring_pg = get_sp_group().ring_group + assert ring_pg is not None, "Ring process group is not initialized." + ring_chunk_len = q.shape[0] + _, ring_rank = get_ring_ctx() + + # `isend`/`irecv` (unlike collectives) address peers by global rank even + # under a sub-group, so resolve this ring's neighbors once up front. + next_global_rank = torch.distributed.get_global_rank( + ring_pg, (ring_rank + 1) % ring_ws + ) + prev_global_rank = torch.distributed.get_global_rank( + ring_pg, (ring_rank - 1) % ring_ws + ) + + # K and V travel as one stacked buffer so each hop is a single P2P + # send/recv pair instead of two; stacking also makes the buffer + # contiguous, so no separate .contiguous() call is needed. + kv0 = torch.stack((k, v)) + kv_bufs = [kv0, torch.empty_like(kv0)] + cur = 0 + + q_cu = torch.tensor([0, ring_chunk_len], dtype=torch.int32, device=q.device) + out_acc: torch.Tensor | None = None + lse_acc: torch.Tensor | None = None + pending_ops = None + for step in range(ring_ws): + nxt = 1 - cur + if step < ring_ws - 1: + # kick off next hop before this step's compute so the transfer + # overlaps it; wait for completion only after issuing compute. + pending_ops = torch.distributed.batch_isend_irecv( + [ + torch.distributed.P2POp( + torch.distributed.isend, + kv_bufs[cur], + next_global_rank, + group=ring_pg, + ), + torch.distributed.P2POp( + torch.distributed.irecv, + kv_bufs[nxt], + prev_global_rank, + group=ring_pg, + ), + ] + ) + + src_rank = (ring_rank - step) % ring_ws + remote_used = min( + max(real_seq_len - src_rank * ring_chunk_len, 0), ring_chunk_len + ) + if remote_used > 0: + k_cu = torch.tensor([0, remote_used], dtype=torch.int32, device=q.device) + result = flash_attn_varlen_func( + q, + kv_bufs[cur][0, :remote_used], + kv_bufs[cur][1, :remote_used], + cu_seqlens_q=q_cu, + cu_seqlens_k=k_cu, + max_seqlen_q=ring_chunk_len, + max_seqlen_k=remote_used, + softmax_scale=softmax_scale, + causal=False, + ver=_fa_backend.fa_ver, + return_softmax_lse=True, + ) + if not isinstance(result, tuple): + raise RuntimeError( + "flash_attn_varlen_func did not return softmax_lse; ring " + "parallelism requires a backend that supports " + "return_softmax_lse=True." + ) + step_out, step_lse, *_ = result + out_acc, lse_acc = _ring_merge_attention( + out_acc, lse_acc, step_out, step_lse + ) + + if pending_ops is not None: + for op in pending_ops: + op.wait() + pending_ops = None + cur = nxt + return out_acc.to(q.dtype) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 7654f8815..8b665082f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -223,12 +223,21 @@ class GPUWorker(GPUWorkerPostTrainingMixin): def init_device_and_model(self) -> None: """Initialize the device and load the model.""" current_platform.set_device(current_platform.get_device(self.local_rank)) - intra_op_threads = _worker_cpu_intra_op_threads(self.server_args.num_gpus) + # num_gpus is the total world size across every node; the co-located, + # CPU-contending worker count on THIS host is num_gpus // nnodes. + local_num_gpus = self.server_args.num_gpus // self.server_args.nnodes + intra_op_threads = _worker_cpu_intra_op_threads(local_num_gpus) if intra_op_threads is not None: torch.set_num_threads(intra_op_threads) - # Set environment variables for distributed initialization - os.environ["MASTER_ADDR"] = "localhost" - os.environ["MASTER_PORT"] = str(self.master_port) + # Set environment variables for distributed initialization. Single + # node rendezvous stays on loopback; cross-node rendezvous must use + # an address every node can reach, so --dist-init-addr takes over. + if self.server_args.nnodes > 1: + rendezvous_addr = NetworkAddress.parse(self.server_args.dist_init_addr) + else: + rendezvous_addr = NetworkAddress("127.0.0.1", self.master_port) + os.environ["MASTER_ADDR"] = rendezvous_addr.host + os.environ["MASTER_PORT"] = str(rendezvous_addr.port) os.environ["LOCAL_RANK"] = str(self.local_rank) os.environ["RANK"] = str(self.rank) os.environ["WORLD_SIZE"] = str(self.server_args.num_gpus) @@ -241,9 +250,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin): ring_degree=self.server_args.ring_degree, sp_size=self.server_args.sp_degree, dp_size=self.server_args.dp_size, - distributed_init_method=NetworkAddress( - "127.0.0.1", self.master_port - ).to_tcp(), + distributed_init_method=rendezvous_addr.to_tcp(), dist_timeout=self.server_args.dist_timeout, ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 5f5709dd1..25b5ae31b 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -36,6 +36,10 @@ from sglang.multimodal_gen.runtime.distributed import ( get_tp_world_size, tensor_model_parallel_all_gather, ) +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_ring_ctx, + get_ulysses_ctx, +) from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionRequirements, ) @@ -48,6 +52,7 @@ from sglang.multimodal_gen.runtime.layers.linear import ( from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) +from sglang.multimodal_gen.runtime.layers.usp import _ring_attention_varlen from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, is_layerwise_offloaded_module, @@ -118,34 +123,6 @@ _FORWARD_SUPPORTED_KWARGS = frozenset( ) -def _ulysses_ctx() -> tuple[int, int]: - """(world_size, rank) of the Ulysses sequence-parallel group. - - Returns (1, 0) when model parallelism is not initialized (unit tests / - single-process debug paths init tp=1 sp=1 which also yields ws=1). - """ - from sglang.multimodal_gen.runtime.distributed.parallel_state import ( - get_ulysses_parallel_rank, - get_ulysses_parallel_world_size, - model_parallel_is_initialized, - ) - - if not model_parallel_is_initialized(): - return 1, 0 - return get_ulysses_parallel_world_size(), get_ulysses_parallel_rank() - - -def _ring_world_size() -> int: - from sglang.multimodal_gen.runtime.distributed.parallel_state import ( - get_ring_parallel_world_size, - model_parallel_is_initialized, - ) - - if not model_parallel_is_initialized(): - return 1 - return get_ring_parallel_world_size() - - def _reorder_grouped_qkv_to_qkv( weight: torch.Tensor, *, @@ -467,8 +444,9 @@ def _minimax_h3_attention_core_impl( cu_seqlens_host: tuple[int, ...] | None, max_seqlen: int, ulysses_active: bool, + ring_active: bool = False, ) -> torch.Tensor: - """Dynamic varlen attention and Ulysses collectives. + """Dynamic varlen attention and Ulysses/Ring collectives. This is the narrow BCG break point: projections, normalization, RoPE, residuals, and MLPs remain captured while the dynamic packed attention @@ -491,14 +469,33 @@ def _minimax_h3_attention_core_impl( attention_requirements=AttentionRequirements(packed_varlen=True), ) ) - out = attention._attention_impl.forward_varlen( - q, - k, - v, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - cu_seqlens_host=cu_seqlens_host, - ) + + if ring_active: + ring_ws, _ = get_ring_ctx() + if attention._attention_backend_enum is not AttentionBackendEnum.FA: + raise NotImplementedError( + "MiniMax H3 ring parallelism requires the FlashAttention " + "backend (matches --ring-degree's general restriction)." + ) + # max_seqlen is cu_seqlens[1] (`used`) by construction -- the real, + # non-padding row count ring needs, already a host int here. + out = _ring_attention_varlen( + q, + k, + v, + softmax_scale=attention.softmax_scale, + real_seq_len=max_seqlen, + ring_ws=ring_ws, + ) + else: + out = attention._attention_impl.forward_varlen( + q, + k, + v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + cu_seqlens_host=cu_seqlens_host, + ) if ulysses_active: out = _usp_output_all_to_all(out[None], head_dim=2)[0] return out @@ -531,6 +528,7 @@ class MiniMaxH3Attention(nn.Module): self.local_inner_dim = self.num_heads * self.head_dim self.softmax_scale = self.head_dim**-0.5 self._attention_impl = None + self._attention_backend_enum: AttentionBackendEnum | None = None # The checkpoint stores one fused qkv tensor. Each logical Q/K/V # matrix must be sharded independently; a plain ColumnParallelLinear # would instead slice across the concatenated tensor and is incorrect @@ -579,6 +577,10 @@ class MiniMaxH3Attention(nn.Module): softmax_scale=self.softmax_scale, num_kv_heads=self.num_heads, ) + # Ring only supports FA (see _minimax_h3_attention_core_impl); keep + # the resolved enum alongside the impl instance instead of a second + # get_attn_backend() call at the ring gate. + self._attention_backend_enum = backend.get_enum() def _install_qkv_weight_loader(self, arch: MiniMaxH3DiTArchConfig) -> None: weight = self.qkv_proj.weight @@ -620,6 +622,7 @@ class MiniMaxH3Attention(nn.Module): cu_seqlens_host: tuple[int, ...] | None = None, max_seqlen: int, ulysses_active: bool = False, + ring_active: bool = False, ) -> torch.Tensor: """x: [T, hidden] packed thd rows -> [T, hidden]. @@ -686,6 +689,7 @@ class MiniMaxH3Attention(nn.Module): cu_seqlens_host=cu_seqlens_host, max_seqlen=max_seqlen, ulysses_active=ulysses_active, + ring_active=ring_active, ) out = out.reshape(total, self.num_heads * self.head_dim) out, _ = self.out_proj(out) @@ -905,6 +909,7 @@ class MiniMaxH3DiTBlock(nn.Module): cu_seqlens_host: tuple[int, ...] | None = None, max_seqlen: int, ulysses_active: bool = False, + ring_active: bool = False, adaln_params: tuple[torch.Tensor, ...] | None = None, ) -> torch.Tensor: """x: [T, H]; adaln_input: [M, t_dim]; combined_indices: [T] @@ -930,6 +935,7 @@ class MiniMaxH3DiTBlock(nn.Module): cu_seqlens_host=cu_seqlens_host, max_seqlen=max_seqlen, ulysses_active=ulysses_active, + ring_active=ring_active, ) x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE) @@ -1070,12 +1076,8 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): ) -> None: if ulysses_size <= 0: raise ValueError("MiniMax H3 Ulysses size must be positive.") - if ring_size != 1: - raise NotImplementedError( - "MiniMax H3 packed multi-segment attention does not support " - "Ring or mixed USP. Set --ring-degree 1 and use Ulysses " - "sequence parallelism." - ) + if ring_size <= 0: + raise ValueError("MiniMax H3 ring size must be positive.") local_heads = arch.num_attention_heads // tp_size if local_heads % ulysses_size: raise ValueError( @@ -1083,13 +1085,19 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): f"Ulysses size {ulysses_size} (total heads=" f"{arch.num_attention_heads}, TP={tp_size})." ) - if MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT % ulysses_size: + # ring never shards heads (only rows), so it has no head-divisibility + # constraint; the packed sequence alignment constant must still + # divide the *combined* sequence-parallel size, since ring adds an + # outer row split on top of Ulysses's inner one (see forward()). + sp_size = ulysses_size * ring_size + if MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT % sp_size: raise ValueError( "MiniMax H3 packed sequence alignment " f"{MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT} must be divisible by " - f"Ulysses size {ulysses_size}. Choose a Ulysses size that " - "divides both the TP-local attention heads and the packed " - "sequence alignment." + f"the combined sequence-parallel size {sp_size} " + f"(ulysses={ulysses_size} x ring={ring_size}). Choose degrees " + "whose product divides both the TP-local attention heads and " + "the packed sequence alignment." ) def __init__( @@ -1105,13 +1113,13 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): self.num_attention_heads = arch.num_attention_heads self.num_channels_latents = arch.latents_dim tp_size = get_tp_world_size() - ulysses_size, _ = _ulysses_ctx() + ulysses_size, _ = get_ulysses_ctx() self._validate_tp_config(arch=arch, tp_size=tp_size) self._validate_sequence_parallel_config( arch=arch, tp_size=tp_size, ulysses_size=ulysses_size, - ring_size=_ring_world_size(), + ring_size=get_ring_ctx()[0], ) self.video_patch_proj = ColumnParallelLinear( @@ -1266,20 +1274,31 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): *, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: - """Build request-static RoPE inputs for this Ulysses rank.""" + """Build request-static RoPE inputs for this rank's row shard. + + Same 2D row split as forward(): ring first (outer, contiguous + ring_chunk_len slice), Ulysses second (inner slice within that + chunk) -- see forward()'s row_start derivation for the identity + this must stay in sync with. + """ if img_position_ids.dim() != 3 or img_position_ids.shape[0] != 1: raise ValueError( "img_position_ids must be [1, S, 3], got " f"{list(img_position_ids.shape)}" ) seq_len = int(img_position_ids.shape[1]) - sp_ws, sp_rank = _ulysses_ctx() + ulysses_ws, ulysses_rank = get_ulysses_ctx() + ring_ws, ring_rank = get_ring_ctx() + sp_ws = ulysses_ws * ring_ws if seq_len % sp_ws: raise ValueError( - f"packed seq_len {seq_len} not divisible by ulysses world size {sp_ws}" + f"packed seq_len {seq_len} not divisible by the combined " + f"sequence-parallel world size {sp_ws} " + f"(ulysses={ulysses_ws} x ring={ring_ws})" ) local_seq_len = seq_len // sp_ws - row_start = sp_rank * local_seq_len + ring_chunk_len = local_seq_len * ulysses_ws + row_start = ring_rank * ring_chunk_len + ulysses_rank * local_seq_len rope_freqs = self.rope( img_position_ids[:, row_start : row_start + local_seq_len] ).to(device) @@ -1505,6 +1524,9 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): else raw_cu_seqlens_host ) ) + # max_seqlen_q is set to cu_seqlens[1] (`used`, the real/non-padding + # row count) by construction -- already a plain host int here, so + # ring can reuse it as real_seq_len below with no new device sync. max_seqlen = int(self._psp_field(psp, "packed_seq_params", "max_seqlen_q")) refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params") refiner_cu = self._psp_field( @@ -1528,29 +1550,34 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): ) device = x.device self._resolve_attention_backend_once() - if _ring_world_size() != 1: - raise NotImplementedError( - "MiniMax H3 packed multi-segment attention requires " - "--ring-degree 1; Ring and mixed USP are unsupported." - ) - sp_ws, sp_rank = _ulysses_ctx() + # Row split is 2D: ring first (an outer, contiguous ring_chunk_len + # slice of the packed sequence), Ulysses second (an inner slice + # within this rank's ring chunk). Only Ulysses shards heads inside + # attention -- ring instead ring-rotates each rank's local KV chunk + # and online-softmax merges partial outputs (see + # _minimax_h3_attention_core_impl), so it has no head constraint. + ulysses_ws, ulysses_rank = get_ulysses_ctx() + ring_ws, ring_rank = get_ring_ctx() + sp_ws = ulysses_ws * ring_ws local_seq_len = seq_len if sp_ws > 1: if seq_len % sp_ws: raise ValueError( - f"packed seq_len {seq_len} not divisible by ulysses " - f"world size {sp_ws}" + f"packed seq_len {seq_len} not divisible by the combined " + f"sequence-parallel world size {sp_ws} " + f"(ulysses={ulysses_ws} x ring={ring_ws})" ) local_heads = self.num_attention_heads // get_tp_world_size() - if local_heads % sp_ws: + if local_heads % ulysses_ws: raise ValueError( f"TP-local heads {local_heads} not divisible by Ulysses " - f"world size {sp_ws} (total heads={self.num_attention_heads}, " - f"TP={get_tp_world_size()})" + f"world size {ulysses_ws} (total heads=" + f"{self.num_attention_heads}, TP={get_tp_world_size()})" ) local_seq_len = seq_len // sp_ws - row_start = sp_rank * local_seq_len + ring_chunk_len = local_seq_len * ulysses_ws + row_start = ring_rank * ring_chunk_len + ulysses_rank * local_seq_len row_stop = row_start + local_seq_len # RoPE and latent projections are row-local before Ulysses exchanges @@ -1622,9 +1649,11 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): block.adaln_proj.split_output(output) for block, output in zip(self.blocks, gathered_adaln) ) - # With Ulysses sequence parallelism, shard rows across the group for - # the block stack. Attention trades sequence for heads internally; - # everything else, including the final layer, is row-local. + # With sequence parallelism, shard rows across the group for the + # block stack. Attention trades sequence for heads internally + # (Ulysses) and/or ring-rotates KV across ring ranks; everything + # else, including the final layer, is row-local. Only the narrow + # video/audio logits are gathered after the final layer. for index, block in enumerate(self.blocks): hidden = block( hidden, @@ -1634,7 +1663,8 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): cu_seqlens=cu_seqlens, cu_seqlens_host=cu_seqlens_host, max_seqlen=max_seqlen, - ulysses_active=sp_ws > 1, + ulysses_active=ulysses_ws > 1, + ring_active=ring_ws > 1, adaln_params=( None if block_adaln_params is None else block_adaln_params[index] ), diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 6cc1501ac..c413edc27 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -6,13 +6,8 @@ import torch.nn as nn from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig from sglang.multimodal_gen.runtime.distributed import ( - get_sp_parallel_rank, get_sp_world_size, get_tp_world_size, - sequence_model_parallel_all_gather, -) -from sglang.multimodal_gen.runtime.distributed.parallel_state import ( - get_ring_parallel_world_size, ) from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul from sglang.multimodal_gen.runtime.layers.attention import ( @@ -1575,24 +1570,6 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): ) cap_seq_len = cap_feats.shape[1] - use_full_unified_sequence = ( - get_sp_world_size() > 1 and get_ring_parallel_world_size() > 1 - ) - if use_full_unified_sequence: - # Ring support for this attention layout is not implemented; the - # full-sequence gather is correct but gives up ring's memory and - # overlap benefits. - logger.warning_once( - "zimage under ring_degree > 1 falls back to a full-sequence " - "K/V gather" - ) - x_local_seq_len = x.shape[1] - if use_full_unified_sequence: - x = sequence_model_parallel_all_gather(x.contiguous(), dim=1) - x_freqs_cis = ( - sequence_model_parallel_all_gather(x_freqs_cis[0].contiguous(), dim=0), - sequence_model_parallel_all_gather(x_freqs_cis[1].contiguous(), dim=0), - ) unified = torch.cat([x, cap_feats], dim=1) unified_freqs_cis = ( torch.cat([x_freqs_cis[0], cap_freqs_cis[0]], dim=-2), @@ -1608,7 +1585,7 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): unified_rope_cos_sin_cache, unified_rope_positions = self._get_rope_cache( "_cached_unified_rope_cache", unified_freqs_cis ) - num_replicated_suffix = cap_seq_len if not use_full_unified_sequence else 0 + num_replicated_suffix = cap_seq_len for layer in self.layers: unified = layer( @@ -1620,17 +1597,11 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): attn_mask=unified_attn_mask, attn_mask_meta=unified_attn_mask_meta, num_replicated_suffix=num_replicated_suffix, - skip_sequence_parallel_override=use_full_unified_sequence, ) unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( unified, adaln_input ) - if use_full_unified_sequence: - sp_rank = get_sp_parallel_rank() - start = sp_rank * x_local_seq_len - end = start + x_local_seq_len - unified = unified[:, start:end] x = list(unified.unbind(dim=0)) x = self.unpatchify(x, x_size, patch_size, f_patch_size) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py index 54715b2b8..147340b61 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py @@ -16,6 +16,10 @@ import torch from sglang.multimodal_gen.configs.models.dits.minimax_h3 import ( MINIMAX_H3_ADALN_MODALITY_NUM, ) +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_ring_ctx, + get_ulysses_ctx, +) MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999 # ref2va audio reference anchor timestep @@ -46,18 +50,6 @@ def _minimax_h3_update_target_rows_( torch.add(state, velocity, out=state) -def _ulysses_ctx() -> tuple[int, int]: - from sglang.multimodal_gen.runtime.distributed.parallel_state import ( - get_ulysses_parallel_rank, - get_ulysses_parallel_world_size, - model_parallel_is_initialized, - ) - - if not model_parallel_is_initialized(): - return 1, 0 - return get_ulysses_parallel_world_size(), get_ulysses_parallel_rank() - - def _build_local_embedding_layout( *, seq_len: int, @@ -70,8 +62,8 @@ def _build_local_embedding_layout( ) -> dict[str, torch.Tensor | int]: if seq_len % world_size: raise ValueError( - f"packed seq_len {seq_len} not divisible by Ulysses world size " - f"{world_size}" + f"packed seq_len {seq_len} not divisible by the combined " + f"sequence-parallel world size {world_size}" ) local_seq_len = seq_len // world_size row_start = rank * local_seq_len @@ -180,10 +172,19 @@ class MiniMaxH3DenoiseBranch: 1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device ) text_pos_dev = text_pos.to(device) - ulysses_world_size, ulysses_rank = _ulysses_ctx() + ulysses_world_size, ulysses_rank = get_ulysses_ctx() + ring_world_size, ring_rank = get_ring_ctx() + # Combined SP-local rank/world size: the group coordinator lays out + # ring as the outer (slower-varying) dimension and Ulysses as the + # inner one (see set_seq_parallel_pg_by_sp_groups), so this matches + # minimax_h3.py's row_start = ring_rank*ring_chunk_len + + # ulysses_rank*local_seq_len exactly -- `_build_local_embedding_layout` + # below needs the same combined rank, not just the Ulysses component. + sp_world_size = ulysses_world_size * ring_world_size + sp_rank = ring_rank * ulysses_world_size + ulysses_rank token_tags_host = token_tags.view(-1).to(dtype=torch.long) - local_seq_len = seq_len // ulysses_world_size - local_row_start = ulysses_rank * local_seq_len + local_seq_len = seq_len // sp_world_size + local_row_start = sp_rank * local_seq_len local_row_stop = local_row_start + local_seq_len self.local_row_slice = slice(local_row_start, local_row_stop) self.block_token_tags = ( @@ -210,8 +211,8 @@ class MiniMaxH3DenoiseBranch: text_pos=text_pos, img_pos=self.img_pos, audio_pos=self.audio_pos, - world_size=ulysses_world_size, - rank=ulysses_rank, + world_size=sp_world_size, + rank=sp_rank, device=device, ), "packed_seq_params": { diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 82db9d504..0b99108d6 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -224,6 +224,12 @@ class ServerArgs(DisaggServerArgsMixin): performance_mode: str = "auto" base_gpu_id: int = 0 gpu_ids: list[int] | None = None + # cross-node: num_gpus is the total world size across all nodes; each + # node runs num_gpus // nnodes local GPU workers (mirrors srt's + # tp_size_per_node convention) + nnodes: int = 1 + node_rank: int = 0 + dist_init_addr: str | None = None tp_size: Optional[int] = None sp_degree: Optional[int] = None # sequence parallelism @@ -1467,6 +1473,27 @@ class ServerArgs(DisaggServerArgsMixin): help="The starting GPU ID for this instance. Used with --disagg-role " "to place role instances on specific GPUs without CUDA_VISIBLE_DEVICES.", ) + parser.add_argument( + "--nnodes", + type=int, + default=ServerArgs.nnodes, + help="The number of nodes for cross-node parallelism. --num-gpus is " + "the total GPU count across all nodes; each node runs " + "num_gpus // nnodes local workers.", + ) + parser.add_argument( + "--node-rank", + type=int, + default=ServerArgs.node_rank, + help="The rank of this node among --nnodes nodes, in [0, nnodes).", + ) + parser.add_argument( + "--dist-init-addr", + type=str, + default=ServerArgs.dist_init_addr, + help="The host:port distributed rendezvous address, reachable from " + "every node. Required when --nnodes > 1.", + ) parser.add_argument( "--gpu-ids", nargs="+", @@ -2496,6 +2523,19 @@ class ServerArgs(DisaggServerArgsMixin): "sequence parallelism after dp/tp/cfg" ) + if self.nnodes < 1: + raise ValueError("--nnodes must be a natural number") + if not (0 <= self.node_rank < self.nnodes): + raise ValueError( + f"--node-rank ({self.node_rank}) must be in [0, nnodes={self.nnodes})" + ) + if self.nnodes > 1 and self.dist_init_addr is None: + raise ValueError("--dist-init-addr is required when --nnodes > 1") + if self.num_gpus % self.nnodes != 0: + raise ValueError( + f"num_gpus ({self.num_gpus}) must be divisible by nnodes ({self.nnodes})" + ) + if self.sp_degree > self.num_gpus or self.num_gpus % self.sp_degree != 0: raise ValueError( f"num_gpus ({self.num_gpus}) must be >= and divisible by sp_degree ({self.sp_degree})" diff --git a/python/sglang/multimodal_gen/test/server/consistency_thresholds/h100.json b/python/sglang/multimodal_gen/test/server/consistency_thresholds/h100.json index 620792e44..0f349b55c 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_thresholds/h100.json +++ b/python/sglang/multimodal_gen/test/server/consistency_thresholds/h100.json @@ -55,12 +55,6 @@ "psnr_threshold": 13.5, "mean_abs_diff_threshold": 17.2 }, - "qwen_image_t2i_2_gpus": { - "clip_threshold": 0.98, - "ssim_threshold": 0.79, - "psnr_threshold": 15.7, - "mean_abs_diff_threshold": 17.2 - }, "flux_2_image_t2i": { "clip_threshold": 0.98, "ssim_threshold": 0.95, diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 1c73d1dca..86b796974 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -39,7 +39,7 @@ logger = init_logger(__name__) # NPU/ascend) is read from sgl-project/ci-data-diffusion, where the GT-gen workflows # publish. SGL_TEST_FILES_CI_DATA_REPO = "sgl-project/ci-data-diffusion" -SGL_TEST_FILES_CI_DATA_REVISION = "dc0e1bb34f2776313a259bcfab3e30daed85160e" +SGL_TEST_FILES_CI_DATA_REVISION = "eccf85dcebaaded92df8b0fce3064ebea910c6d4" # The NPU pin is kept as a separate branch so ascend GT can be bumped independently # when it's regenerated on its own cadence. diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py index 8ca111133..1c5e60f02 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py @@ -182,7 +182,7 @@ def test_rank_local_token_tags_match_reference_slice(): for rank in range(world_size): with patch( "sglang.multimodal_gen.runtime.pipelines_core.stages." - "model_specific_stages.minimax_h3.denoise_loop._ulysses_ctx", + "model_specific_stages.minimax_h3.denoise_loop.get_ulysses_ctx", return_value=(world_size, rank), ): branch = _branch(mode, token_tags=token_tags) diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py index 8bb1d0857..d1409179b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py @@ -120,12 +120,36 @@ def test_tp_and_ulysses_admission_uses_tp_local_shapes(): ulysses_size=4, ring_size=1, ) - with pytest.raises(NotImplementedError): + # ring is implemented now: it splits rows, not heads, so it carries no + # head-divisibility constraint of its own + MiniMaxH3DiTModel._validate_sequence_parallel_config( + arch=arch, + tp_size=1, + ulysses_size=1, + ring_size=2, + ) + MiniMaxH3DiTModel._validate_sequence_parallel_config( + arch=arch, + tp_size=1, + ulysses_size=8, + ring_size=2, + ) + # what ring does constrain is the packed-sequence alignment, which has to + # divide by the *combined* degree because ring adds an outer row split on + # top of Ulysses's inner one + with pytest.raises(ValueError): + MiniMaxH3DiTModel._validate_sequence_parallel_config( + arch=arch, + tp_size=1, + ulysses_size=8, + ring_size=3, + ) + with pytest.raises(ValueError): MiniMaxH3DiTModel._validate_sequence_parallel_config( arch=arch, tp_size=1, ulysses_size=1, - ring_size=2, + ring_size=0, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_sp_shard.py b/python/sglang/multimodal_gen/test/unit/test_sp_shard.py index 9d7b29fc4..b467ffb6a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sp_shard.py +++ b/python/sglang/multimodal_gen/test/unit/test_sp_shard.py @@ -13,10 +13,9 @@ from sglang.multimodal_gen.runtime.distributed.sp_shard_utils import ( ) -def _fake_sp(monkeypatch, sp_size, sp_rank=0, ring=1): +def _fake_sp(monkeypatch, sp_size, sp_rank=0): monkeypatch.setattr(sps, "get_sp_world_size", lambda: sp_size) monkeypatch.setattr(sps, "get_sp_parallel_rank", lambda: sp_rank) - monkeypatch.setattr(sps, "get_ring_parallel_world_size", lambda: ring) # --- build_shard_plan math -------------------------------------------------------- @@ -150,10 +149,12 @@ def test_strategy_replicates_when_padding_spans_multiple_shards(monkeypatch): assert sps.plan_text_strategy(14) == "shard" -def test_strategy_ring_blocks_padded_shard(monkeypatch): - _fake_sp(monkeypatch, 2, ring=2) - assert sps.plan_text_strategy(15) == "replicate" # padded shard needs mask - assert sps.plan_text_strategy(16) == "shard" # divisible: no mask needed +def test_strategy_shards_padded_text_under_ring(monkeypatch): + # Tail-padded shards ride the ring kernel now; the strategy no longer + # consults the ring degree at all. + _fake_sp(monkeypatch, 2) + assert sps.plan_text_strategy(15) == "shard" + assert sps.plan_text_strategy(16) == "shard" def test_strategy_min_len_threshold(monkeypatch): diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py b/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py new file mode 100644 index 000000000..d3a985cb1 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py @@ -0,0 +1,191 @@ +"""Replicated-token paths under ring parallelism: dispatch, merge, KV order.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.multimodal_gen.runtime.layers.attention.layer import USPAttention +from sglang.multimodal_gen.runtime.layers.usp import _merge_attention_partials +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum + +_LAYER = "sglang.multimodal_gen.runtime.layers.attention.layer" + + +def _sdpa(q, k, v, scale): + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), + k.transpose(1, 2).float(), + v.transpose(1, 2).float(), + scale=scale, + ).transpose(1, 2) + + +class _LseImpl: + """SDPA with explicit LSE so merge math can be verified for real.""" + + def __init__(self, scale): + self.scale = scale + self.seen_k = [] + + def forward(self, q, k, v, attn_metadata=None, return_softmax_lse=False): + self.seen_k.append(k) + out = _sdpa(q, k, v, self.scale) + if not return_softmax_lse: + return out.to(q.dtype) + logits = torch.einsum("bshd,bthd->bhst", q.float(), k.float()) * self.scale + lse = torch.logsumexp(logits, dim=-1) # [B, H, S] + return out, lse + + +def _ring_pair_via_impl(impl, q, k_shard, v_shard): + """Stand-in for ring_attn on a 1-chunk ring: one local partial + LSE.""" + return impl.forward( + q, k_shard, v_shard, attn_metadata=None, return_softmax_lse=True + ) + + +class RingReplicatedBase(unittest.TestCase): + B, S_SHARD, REP, H, D = 1, 6, 3, 2, 8 + + def _attn(self): + obj = USPAttention.__new__(USPAttention) + obj.skip_sequence_parallel = False + obj.sp_attention_mode = "ulysses" + obj.sp_attention_mode_is_auto = False + obj.softmax_scale = self.D**-0.5 + obj.backend = AttentionBackendEnum.FA + obj.causal = False + obj.dropout_p = 0.0 + obj.attn_impl = _LseImpl(obj.softmax_scale) + return obj + + def _patches(self, ring_ws=2): + return ( + patch( + f"{_LAYER}.get_forward_context", + return_value=SimpleNamespace(attn_metadata=None), + ), + patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=ring_ws), + patch(f"{_LAYER}.get_ulysses_parallel_world_size", return_value=1), + patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=ring_ws), + patch(f"{_LAYER}.get_ulysses_parallel_rank", return_value=0), + patch(f"{_LAYER}._usp_input_all_to_all", side_effect=lambda x, head_dim: x), + patch( + f"{_LAYER}._usp_output_all_to_all", side_effect=lambda x, head_dim: x + ), + patch( + f"{_LAYER}.ring_attn", + side_effect=lambda q, k, v, impl, return_softmax_lse: _ring_pair_via_impl( + impl, q, k, v + ), + ), + ) + + def _rand(self, s): + return torch.randn(self.B, s, self.H, self.D) + + +class TestRingReplicatedPrefix(RingReplicatedBase): + def test_u1_r2_dispatches_and_merges_exactly(self): + obj = self._attn() + q = self._rand(self.REP + self.S_SHARD) + k = self._rand(self.REP + self.S_SHARD) + v = self._rand(self.REP + self.S_SHARD) + + ps = self._patches() + with ps[0], ps[1], ps[2], ps[3], ps[4], ps[5], ps[6], ps[7]: + out = obj.forward(q, k, v, num_replicated_prefix=self.REP) + + # One local ring chunk + rep partial merged == full attention. + ref = _sdpa(q, k, v, obj.softmax_scale) + self.assertEqual(out.shape, q.shape) + torch.testing.assert_close(out.float(), ref, atol=1e-5, rtol=1e-5) + + def test_kv_prefix_u1_r2_matches_full_attention(self): + obj = self._attn() + q = self._rand(self.S_SHARD) + k = self._rand(self.REP + self.S_SHARD) + v = self._rand(self.REP + self.S_SHARD) + + ps = self._patches() + with ps[0], ps[1], ps[2], ps[3], ps[4], ps[5], ps[6], ps[7]: + out = obj.forward(q, k, v, num_replicated_kv_prefix=self.REP) + + ref = _sdpa(q, k, v, obj.softmax_scale) + torch.testing.assert_close(out.float(), ref, atol=1e-5, rtol=1e-5) + + +class TestRingReplicatedSuffix(RingReplicatedBase): + def test_u1_r2_dispatches_and_merges_exactly(self): + obj = self._attn() + q = self._rand(self.S_SHARD + self.REP) + k = self._rand(self.S_SHARD + self.REP) + v = self._rand(self.S_SHARD + self.REP) + + ps = self._patches() + with ps[0], ps[1], ps[2], ps[3], ps[4], ps[5], ps[6], ps[7]: + out = obj.forward(q, k, v, num_replicated_suffix=self.REP) + + ref = _sdpa(q, k, v, obj.softmax_scale) + torch.testing.assert_close(out.float(), ref, atol=1e-5, rtol=1e-5) + + def test_non_ring_path_keeps_kv_tail_order(self): + obj = self._attn() + q = self._rand(self.S_SHARD + self.REP) + k = self._rand(self.S_SHARD + self.REP) + v = self._rand(self.S_SHARD + self.REP) + + def _fake_gather(out_list, t, group=None): + for o in out_list: + o.copy_(t) + + ps = self._patches(ring_ws=1) + with ps[0], ps[4], ps[5], ps[6], patch( + f"{_LAYER}.get_ring_parallel_world_size", return_value=1 + ), patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=2), patch( + f"{_LAYER}.get_ulysses_parallel_world_size", return_value=2 + ), patch( + f"{_LAYER}.get_sp_group", return_value=SimpleNamespace(ulysses_group=None) + ), patch( + "torch.distributed.all_gather", side_effect=_fake_gather + ): + # Identity-mocked collectives don't reproduce head-shard shapes, + # so the final concat may fail — the kernel K order is recorded + # before that and is all this test asserts. + try: + obj.forward(q, k, v, num_replicated_suffix=self.REP) + except RuntimeError: + pass + + # Bitwise contract: suffix KV stays at the tail in the kernel call. + kernel_k = obj.attn_impl.seen_k[-1] + torch.testing.assert_close(kernel_k[:, -self.REP :], k[:, -self.REP :]) + + +class TestMergePartials(unittest.TestCase): + def test_two_disjoint_halves_merge_to_full_attention(self): + torch.manual_seed(7) + B, S, T, H, D = 1, 5, 8, 2, 16 + scale = D**-0.5 + q = torch.randn(B, S, H, D) + k = torch.randn(B, T, H, D) + v = torch.randn(B, T, H, D) + + def part(ks, vs): + logits = torch.einsum("bshd,bthd->bhst", q, ks) * scale + lse = torch.logsumexp(logits, dim=-1) + out = torch.softmax(logits, dim=-1) + return torch.einsum("bhst,bthd->bshd", out, vs), lse + + out_a, lse_a = part(k[:, :3], v[:, :3]) + out_b, lse_b = part(k[:, 3:], v[:, 3:]) + merged = _merge_attention_partials(out_a, lse_a, out_b, lse_b) + + ref, _ = part(k, v) + torch.testing.assert_close(merged, ref, atol=1e-6, rtol=1e-6) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_ring_tail_pad.py b/python/sglang/multimodal_gen/test/unit/test_usp_ring_tail_pad.py new file mode 100644 index 000000000..a8be07e6a --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_usp_ring_tail_pad.py @@ -0,0 +1,86 @@ +"""Ring + tail-pad dispatch: a2a within Ulysses, ring clamped to pad_start.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.multimodal_gen.runtime.layers.attention.layer import USPAttention +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum + +_LAYER = "sglang.multimodal_gen.runtime.layers.attention.layer" + + +class TestRingTailPadDispatch(unittest.TestCase): + def _attn(self): + obj = USPAttention.__new__(USPAttention) + obj.skip_sequence_parallel = False + obj.sp_attention_mode = "ulysses" + obj.sp_attention_mode_is_auto = False + obj.softmax_scale = 0.5 + obj.backend = AttentionBackendEnum.FA + obj.causal = False + obj.dropout_p = 0.0 + return obj + + def test_tail_pad_meta_reaches_the_ring_kernel(self): + obj = self._attn() + q = torch.randn(1, 4, 2, 8) + meta = {"pad_start": 13, "pad_end": 16, "local_pad": 3} + seen = {} + + def fake_ring(qc, kc, vc, *, softmax_scale, real_seq_len, ring_ws): + seen.update( + shape=tuple(qc.shape), + real=real_seq_len, + ws=ring_ws, + scale=softmax_scale, + ) + return torch.ones_like(qc) + + with ( + patch( + f"{_LAYER}.get_forward_context", + return_value=SimpleNamespace(attn_metadata=None), + ), + patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=4), + patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=2), + patch( + f"{_LAYER}._usp_input_all_to_all_qkv", + side_effect=lambda q, k, v: (q, k, v), + ), + patch( + f"{_LAYER}._usp_output_all_to_all", side_effect=lambda t, head_dim: t + ), + patch(f"{_LAYER}._ring_attention_varlen", side_effect=fake_ring), + patch(f"{_LAYER}.get_ring_parallel_rank", return_value=3), + ): + out = obj.forward(q, q, q, attn_mask_meta=meta) + + self.assertEqual(out.shape, q.shape) + self.assertEqual(seen["real"], 13) + self.assertEqual(seen["ws"], 2) + self.assertEqual(seen["shape"], (4, 2, 8)) + # Last ring rank holds global rows [12, 16): row 13 onward is pad. + self.assertTrue(torch.all(out[0, 1:] == 0)) + self.assertTrue(torch.all(out[0, :1] == 1)) + + def test_explicit_mask_under_ring_still_refuses(self): + obj = self._attn() + q = torch.randn(1, 4, 2, 8) + mask = torch.ones(1, 4, dtype=torch.bool) + with ( + patch( + f"{_LAYER}.get_forward_context", + return_value=SimpleNamespace(attn_metadata=None), + ), + patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=4), + patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=2), + ): + with self.assertRaisesRegex(NotImplementedError, "ring"): + obj.forward(q, q, q, attn_mask=mask) + + +if __name__ == "__main__": + unittest.main()