minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Mick
2026-08-08 10:56:28 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent 52afe87a08
commit a25c330eb1
20 changed files with 1061 additions and 216 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
+15
View File
@@ -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"
python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py
-3
View File
@@ -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"
python/sglang/multimodal_gen/runtime/launch_server.py
+38 -12
View File
@@ -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(
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+137 -55
View File
@@ -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)
python/sglang/multimodal_gen/runtime/layers/usp.py
+162 -2
View File
@@ -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)
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
+14 -7
View File
@@ -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,
)
python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+99 -69
View File
@@ -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]
),
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+1 -30
View File
@@ -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)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py
+20 -19
View File
@@ -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": {
python/sglang/multimodal_gen/runtime/server_args/server_args.py
+40
View File
@@ -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})"
python/sglang/multimodal_gen/test/server/consistency_thresholds/h100.json
-6
View File
@@ -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,
python/sglang/multimodal_gen/test/test_utils.py
+1 -1
View File
@@ -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.
python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py
+1 -1
View File
@@ -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)
python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py
+26 -2
View File
@@ -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,
)
python/sglang/multimodal_gen/test/unit/test_sp_shard.py
+7 -6
View File
@@ -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):
python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py
+191
View File
@@ -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()
python/sglang/multimodal_gen/test/unit/test_usp_ring_tail_pad.py
+86
View File
@@ -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()