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:
+
+
+
+
+
+
+
+
+
+
+ | Task |
+ Single-node (s/step) |
+ Cross-node (s/step) |
+ Change |
+
+
+
+
+ | T2VA denoise |
+ 0.749 |
+ 0.477 |
+ -36.3% |
+
+
+ | Ref2VA / V2V denoise |
+ 2.572 |
+ 1.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()