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

Check the topology identities where the layout is written, and build at the published widths (#40340)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-21 12:22:59 -07:00
committed by GitHub
parent d5fdab7022
commit 2d0e94e3a3
43 changed files with 843 additions and 261 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/hf3fs/bench_zerocopy.py
+4 -5
View File
@@ -5,10 +5,10 @@ import torch
from tqdm import tqdm
from sglang.srt.distributed import (
get_world_group,
init_distributed_environment,
initialize_model_parallel,
)
from sglang.srt.distributed.parallel_state import get_world_group
from sglang.srt.managers.cache_controller import (
HiCacheController,
PrefetchOperation,
@@ -17,6 +17,7 @@ from sglang.srt.managers.cache_controller import (
from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
from sglang.test.test_utils import publish_build_topology
init_distributed_environment(
world_size=1,
@@ -26,10 +27,8 @@ init_distributed_environment(
backend="gloo",
)
initialize_model_parallel(
tensor_model_parallel_size=1,
pipeline_model_parallel_size=1,
)
publish_build_topology()
initialize_model_parallel()
group = get_world_group().cpu_group
benchmark/kernels/all_reduce/benchmark_all_reduce.py
+3 -1
View File
@@ -21,6 +21,7 @@ from sglang.srt.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
)
from sglang.test.test_utils import publish_build_topology
def parse_args():
@@ -85,7 +86,8 @@ def init_dist(backend: str):
distributed_init_method=distributed_init_method,
local_rank=rank,
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(world_rank=rank, tp_size=world_size)
initialize_model_parallel()
return dist.group.WORLD
benchmark/kernels/all_reduce/benchmark_fused_ar_rms_amd.py
+3 -1
View File
@@ -41,6 +41,7 @@ from sglang.srt.distributed.parallel_state import (
initialize_model_parallel,
set_custom_all_reduce,
)
from sglang.test.test_utils import publish_build_topology
Shape = Tuple[int, int]
@@ -381,7 +382,8 @@ def main():
distributed_init_method="env://",
backend="nccl",
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(world_rank=rank, tp_size=world_size)
initialize_model_parallel()
prefill_shapes = parse_shapes(args.prefill_shapes)
decode_shapes = parse_shapes(args.decode_shapes)
benchmark/kernels/all_reduce/benchmark_fused_ar_rms_quant_amd.py
+3 -1
View File
@@ -47,6 +47,7 @@ from sglang.srt.distributed.parallel_state import (
initialize_model_parallel,
set_custom_all_reduce,
)
from sglang.test.test_utils import publish_build_topology
Shape = Tuple[int, int]
FP8_DTYPE = torch.float8_e4m3fnuz
@@ -400,7 +401,8 @@ def main() -> None:
distributed_init_method="env://",
backend="nccl",
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(world_rank=rank, tp_size=world_size)
initialize_model_parallel()
if rank == 0:
print(
benchmark/kernels/all_reduce/benchmark_mscclpp.py
+3 -1
View File
@@ -30,6 +30,7 @@ from sglang.srt.distributed.parallel_state import (
initialize_model_parallel,
set_mscclpp_all_reduce,
)
from sglang.test.test_utils import publish_build_topology
def torch_allreduce(torch_input: torch.Tensor, group: ProcessGroup) -> torch.Tensor:
@@ -173,7 +174,8 @@ if __name__ == "__main__":
rank=rank,
local_rank=rank % 8,
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(world_rank=rank, tp_size=world_size)
initialize_model_parallel()
group = get_tensor_model_parallel_group().device_group
cpu_group = get_tensor_model_parallel_group().cpu_group
pynccl_comm = get_tensor_model_parallel_group().pynccl_comm
benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py
+3 -1
View File
@@ -44,6 +44,7 @@ from sglang.srt.distributed.parallel_state import (
initialize_model_parallel,
set_torch_symm_mem_all_reduce,
)
from sglang.test.test_utils import publish_build_topology
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
@@ -188,7 +189,8 @@ if __name__ == "__main__":
rank=rank,
local_rank=rank % 8,
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(world_rank=rank, tp_size=world_size)
initialize_model_parallel()
group = get_tensor_model_parallel_group().device_group
cpu_group = get_tensor_model_parallel_group().cpu_group
pynccl_comm = get_tensor_model_parallel_group().pynccl_comm
benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py
+5 -2
View File
@@ -30,14 +30,16 @@ import torch.distributed as dist # type: ignore
from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype as SGLANG_FP8_DTYPE
from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
from sglang.srt.distributed import get_tp_group, tensor_model_parallel_all_reduce
from sglang.srt.distributed import tensor_model_parallel_all_reduce
from sglang.srt.distributed.parallel_state import (
cleanup_dist_env_and_memory,
get_tp_group,
graph_capture,
init_distributed_environment,
initialize_model_parallel,
)
from sglang.srt.layers.layernorm import RMSNorm # noqa
from sglang.test.test_utils import publish_build_topology
try:
from sgl_kernel import fused_add_rmsnorm as SGL_FUSED_ADD_RMS_NORM
@@ -1178,7 +1180,8 @@ def main():
local_rank=rank,
backend="nccl",
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(world_rank=rank, tp_size=world_size)
initialize_model_parallel()
# Validate world size (must be > 1 for collective operations)
if world_size <= 1:
benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py
+3 -4
View File
@@ -26,6 +26,7 @@ from sglang.srt.layers.moe.topk import (
select_experts,
)
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.test_utils import publish_build_topology
def fused_moe_triton_api(
@@ -227,10 +228,8 @@ def main():
backend="nccl" if torch.cuda.is_available() else "gloo",
)
initialize_model_parallel(
tensor_model_parallel_size=1,
expert_model_parallel_size=1,
)
publish_build_topology()
initialize_model_parallel()
model_config = get_model_config(args.model, args.tp_size, args.ep_size)
benchmark.run(
benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py
+3 -4
View File
@@ -15,6 +15,7 @@ from sglang.srt.distributed.parallel_state import (
from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import (
fused_moe as fused_moe_sglang,
)
from sglang.test.test_utils import publish_build_topology
from .common_utils import get_model_config
@@ -243,10 +244,8 @@ def main():
backend="nccl" if torch.cuda.is_available() else "gloo",
)
initialize_model_parallel(
tensor_model_parallel_size=1,
pipeline_model_parallel_size=1,
)
publish_build_topology()
initialize_model_parallel()
shape_configs = get_model_config(args.model, args.tp_size, args.ep_size)
benchmark.run(
examples/usage/modelopt_quantize_and_export.py
+3 -4
View File
@@ -21,6 +21,7 @@ from sglang.srt.distributed.parallel_state import (
initialize_model_parallel,
)
from sglang.srt.model_loader.loader import get_model_loader
from sglang.test.test_utils import publish_build_topology
def _validate_export(export_dir: str) -> bool:
@@ -113,10 +114,8 @@ def quantize_and_export_model(
local_rank=0,
backend="nccl" if device == "cuda" else "gloo",
)
initialize_model_parallel(
tensor_model_parallel_size=1,
pipeline_model_parallel_size=1,
)
publish_build_topology()
initialize_model_parallel()
# Configure model loading with ModelOpt quantization and export
model_config = ModelConfig(
python/sglang/srt/arg_groups/arg_utils.py
+1 -1
View File
@@ -110,7 +110,7 @@ class Derived(msgspec.Struct, frozen=True):
Every declaration carries ``fn`` today, the parallel quotients included:
they are a function of the configured leaves, so they are computed at
publish like the rest. What is special about them is not how they are
computed but that a stamp can move one afterwards -- an elastic scale-up
computed but that a stamp can move one afterwards -- ``initialize_dp_attention``
restamps ``attn_dp_size`` -- which ``ParallelContext`` answers above the
published leaf.
"""
python/sglang/srt/disaggregation/encoder/server.py
+25 -3
View File
@@ -589,10 +589,32 @@ class MMEncoder:
distributed_init_method=dist_init_method,
local_rank=rank,
)
initialize_model_parallel(
tensor_model_parallel_size=get_parallel().tp_size,
attention_context_model_parallel_size=get_parallel().attn_cp_size,
# The encoder serves the vision tower on a world of its own: `tp_size`
# ranks wide, with no pipeline, no expert or MoE-DP dimension and no
# decode context parallelism, whatever the generation side published.
# That has always been the layout it builds; stating it is what stops
# the context from answering with the other side's topology while these
# groups answer with this one.
parallel = get_parallel()
attn_cp_size = parallel.attn_cp_size
attn_tp_size = parallel.tp_size // attn_cp_size
attn_cp_rank, attn_tp_rank = divmod(rank, attn_tp_size)
parallel.override_permanently(
tp_rank=rank,
pp_size=1,
pp_rank=0,
attn_dp_size=1,
attn_dp_rank=0,
attn_tp_size=attn_tp_size,
attn_tp_rank=attn_tp_rank,
attn_cp_rank=attn_cp_rank,
attn_dcp_size=1,
moe_ep_size=1,
moe_ep_rank=0,
moe_dp_size=1,
moe_tp_size=parallel.tp_size,
)
initialize_model_parallel()
initialize_dp_attention(server_args, self.model_config)
self.model = load_model(
python/sglang/srt/distributed/bootstrap.py
+4 -27
View File
@@ -106,15 +106,6 @@ def init_torch_distributed(
server_args=server_args,
model_config=model_config,
gpu_id=ps.gpu_id,
tp_rank=ps.tp_rank,
tp_size=ps.tp_size,
pp_rank=ps.pp_rank,
pp_size=ps.pp_size,
attn_dp_size=ps.attn_dp_size,
attn_cp_size=ps.attn_cp_size,
moe_ep_size=ps.moe_ep_size,
moe_dp_size=ps.moe_dp_size,
dcp_size=ps.attn_dcp_size,
)
# Pre-warm NCCL/RCCL/HCCL to eliminate cold-start latency in first request
@@ -255,19 +246,13 @@ def _init_parallel_groups(
server_args: ServerArgs,
model_config: ModelConfig,
gpu_id: int,
tp_rank: int,
tp_size: int,
pp_rank: int,
pp_size: int,
attn_dp_size: int,
attn_cp_size: int,
moe_ep_size: int,
moe_dp_size: int,
dcp_size: int,
) -> None:
parallel = get_parallel()
tp_size, pp_size = parallel.tp_size, parallel.pp_size
tp_rank, pp_rank = parallel.tp_rank, parallel.pp_rank
is_ep_joiner = get_exec().moe.is_ep_joiner
is_scale_joiner = get_exec().moe.is_ep_scale_joiner
rank_offset = get_parallel().ep_join_rank_offset if is_scale_joiner else 0
rank_offset = parallel.ep_join_rank_offset if is_scale_joiner else 0
world_size = (
rank_offset + tp_size * pp_size if is_scale_joiner else tp_size * pp_size
)
@@ -285,14 +270,6 @@ def _init_parallel_groups(
max_world_size=get_parallel().max_ep_size,
)
initialize_model_parallel(
tensor_model_parallel_size=tp_size,
attention_data_parallel_size=attn_dp_size,
pipeline_model_parallel_size=pp_size,
expert_model_parallel_size=moe_ep_size,
attention_context_model_parallel_size=attn_cp_size,
moe_data_model_parallel_size=moe_dp_size,
decode_context_parallel_size=dcp_size,
shared_experts_tensor_parallel_size=get_parallel().shared_experts_tp_size,
duplicate_tp_group=get_disagg().enable_pdmux,
enable_symm_mem=get_exec().comm.enable_symm_mem,
# Only WORLD is extended during scale-up. The joiner's model-parallel
python/sglang/srt/distributed/parallel_state.py
+51 -29
View File
@@ -50,6 +50,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
)
from sglang.srt.platforms.device_mixin import _DEVICE_TO_DISTRIBUTED_BACKEND
from sglang.srt.runtime_context import (
_validate_parallel,
derive_parallel_widths,
get_global_dwdp_manager,
get_parallel,
@@ -2515,44 +2516,41 @@ def init_distributed_environment(
def initialize_model_parallel(
tensor_model_parallel_size: int = 1,
expert_model_parallel_size: int = 1,
pipeline_model_parallel_size: int = 1,
attention_data_parallel_size: int = 1,
attention_context_model_parallel_size: int = 1,
moe_data_model_parallel_size: int = 1,
decode_context_parallel_size: int = 1,
backend: Optional[str] = None,
duplicate_tp_group: bool = False,
enable_symm_mem: bool = False,
recovered_rank: bool = False,
rank_offset: int = 0,
max_world_size: Optional[int] = None,
shared_experts_tensor_parallel_size: Optional[int] = None,
) -> None:
"""
Initialize model parallel groups.
Initialize model parallel groups at the published widths.
Arguments:
tensor_model_parallel_size: number of GPUs used for tensor model
parallelism.
expert_model_parallel_size: number of GPUs used for expert model
parallelism.
pipeline_model_parallel_size: number of GPUs used for pipeline model
parallelism.
attention_data_parallel_size: number of GPUs used for attention data
parallelism.
attention_context_model_parallel_size: number of GPUs used for attention context
parallelism.
moe_data_model_parallel_size: number of GPUs used for moe data
parallelism.
decode_context_parallel_size: number of GPUs used for decode context
parallelism, which splits the KV cache across GPUs within each
tensor-parallel group during decoding. Must be a divisor of
tensor_model_parallel_size and is currently only supported on the
AMD HIP platform.
shared_experts_tensor_parallel_size: optional shared-expert TP width.
Must divide attention TP; subgroups never cross attention replicas.
Every width comes from the runtime context rather than from an argument:
the configuration already says how wide each dimension is, and a caller
that translates it again is a second place for the two to disagree. A
process that needs a narrower layout than the one it published -- the
media encoder is the case in the tree -- states that layout on the context
first, so what it builds and what it answers stay the same thing.
The remaining arguments are not topology. `backend` is decided by the
device, `duplicate_tp_group` and `enable_symm_mem` by other namespaces, and
`recovered_rank` / `rank_offset` / `max_world_size` describe this
particular join rather than the layout being joined.
The widths this reads:
tp_size: GPUs used for tensor model parallelism.
moe_ep_size: GPUs used for expert model parallelism.
pp_size: GPUs used for pipeline model parallelism.
attn_dp_size: GPUs used for attention data parallelism.
attn_cp_size: GPUs used for attention context parallelism.
moe_dp_size: GPUs used for MoE data parallelism.
attn_dcp_size: GPUs used for decode context parallelism, which splits
the KV cache across GPUs within each tensor-parallel group during
decoding. Must be a divisor of `tp_size` and is currently only
supported on the AMD HIP platform.
shared_experts_tp_size: optional shared-expert TP width. Must divide
attention TP; subgroups never cross attention replicas.
Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
@@ -2589,6 +2587,16 @@ def initialize_model_parallel(
assert torch.distributed.is_initialized()
backend = backend or torch.distributed.get_backend(get_world_group().device_group)
parallel = get_parallel()
tensor_model_parallel_size = parallel.tp_size
expert_model_parallel_size = parallel.moe_ep_size
pipeline_model_parallel_size = parallel.pp_size
attention_data_parallel_size = parallel.attn_dp_size
attention_context_model_parallel_size = parallel.attn_cp_size
moe_data_model_parallel_size = parallel.moe_dp_size
decode_context_parallel_size = parallel.attn_dcp_size
shared_experts_tensor_parallel_size = parallel.shared_experts_tp_size
# Joiners construct their local TP/PP layout in global rank space.
world_size: int = (
tensor_model_parallel_size * pipeline_model_parallel_size
@@ -2933,6 +2941,12 @@ def initialize_model_parallel(
max_world_size=max_world_size,
)
# The groups just built and the configuration they were built from are two
# accounts of one layout. Check them against each other here, where the
# disagreement is still attributable, rather than letting a collective run
# on the wrong peers.
_validate_parallel(get_parallel(), "group build")
def create_custom_parallel_group(
group_ranks: List[int], backend: str = "gloo"
@@ -3094,10 +3108,18 @@ def patch_tensor_parallel_group(tp_group: GroupCoordinator, *, owns_attention: b
narrowed.update(
attn_tp_size=tp_group.world_size,
attn_tp_rank=tp_group.rank_in_group,
attn_tp_group=tp_group,
attn_dp_size=1,
attn_dp_rank=0,
attn_cp_size=1,
attn_cp_rank=0,
attn_cp_group=None,
moe_ep_size=1,
moe_ep_rank=0,
moe_ep_group=None,
moe_dp_size=1,
moe_tp_size=tp_group.world_size,
moe_tp_rank=tp_group.rank_in_group,
)
try:
with get_parallel().override(**narrowed):
python/sglang/srt/layers/dp_attention.py
+55 -21
View File
@@ -15,11 +15,9 @@ from sglang.srt.arg_groups.model_override_base import (
)
from sglang.srt.distributed import (
GroupCoordinator,
get_attn_tensor_model_parallel_world_size,
)
from sglang.srt.distributed import get_moe_dp_group as _get_moe_dp_group
from sglang.srt.distributed import (
get_tensor_model_parallel_world_size,
tensor_model_parallel_all_reduce,
)
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -49,6 +47,33 @@ if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
def dp_gather_width() -> int:
"""How many replicas the DP sync gathers over.
The attention-DP replicas, except after an elastic-EP scale-up, when the
gather spans the expanded WORLD -- whose width is the `dp_size` the
scale-up published. Read from the context either way: a scoped width has
to reach this, which is the whole reason the name has one home.
"""
parallel = get_parallel()
return parallel.dp_size if world_dp_gather_enabled() else parallel.attn_dp_size
def dp_gather_slot() -> int:
"""This process's index in the list the DP sync just gathered.
The gather spans the attention-DP replicas, except after an elastic-EP
scale-up, when it spans the expanded WORLD and the joining cohort is
numbered from its offset. Which list was gathered is what the flag below
says, so the index is read from there rather than kept as a second name on
the topology.
"""
parallel = get_parallel()
if world_dp_gather_enabled():
return parallel.tp_rank + parallel.ep_join_rank_offset
return parallel.attn_dp_rank
def world_dp_gather_enabled() -> bool:
"""Whether DP gathers should use expanded WORLD after joiner admission."""
dp = get_flags().dp
@@ -60,9 +85,13 @@ def enable_joiner_all_gather():
def update_dp_attention_post_scale(new_dp_size: int, new_dp_rank: int):
get_parallel().override_permanently(
attn_dp_size=new_dp_size, attn_dp_rank=new_dp_rank
)
"""Point the DP gather at the expanded WORLD.
The widths themselves are not written here: the caller scales `dp_size` on
the published bag, and the gather reads its width and this process's slot
from there. The arguments are the values the caller is about to publish,
kept so the log says which scale-up this was.
"""
get_flags().dp.use_world_group_for_gather = True
logger.debug(
"[Elastic EP] dp_attention switched to WORLD: dp_size=%d dp_rank=%d",
@@ -92,7 +121,7 @@ class DpPaddingMode(IntEnum):
def get_dp_padding_mode(
cls, is_extend_in_batch, global_num_tokens: List[int]
) -> DpPaddingMode:
dp_size = get_parallel().attn_dp_size
dp_size = dp_gather_width()
# (trangdough) pplx-kernels a2a is a symmetric collective: every EP rank
# must dispatch the same number of tokens or the device-side handshake
@@ -374,14 +403,13 @@ def initialize_dp_attention(
dp.enabled = enable_dp_attention
tp_rank = get_parallel().tp_rank
tp_size = get_tensor_model_parallel_world_size()
tp_size = get_parallel().tp_size
_, _, attn_dp_rank, attn_dp_size = compute_dp_attention_world_info(
enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size
)
if get_exec().moe.elastic_ep_backend is not None and get_parallel().max_ep_size:
attn_dp_rank = tp_rank + get_parallel().ep_join_rank_offset
# Reads the resolution, not a bag: this runs under
# `initialize_dp_attention`, which the weight-cache daemon calls from
# `_init_distributed` -- and other callers reach it from processes
@@ -414,7 +442,10 @@ def is_allocation_symmetric() -> bool:
def get_dp_local_info(forward_batch: ForwardBatch) -> Tuple[torch.Tensor, torch.Tensor]:
# `get_dp_local_info` is only called in global DP gather and scatter. We use global DP rank here.
dp_rank = get_parallel().attn_dp_rank
# The slot in the list that was gathered. A scale-up widens that list
# to WORLD, and this process's index in it is not its index among the
# launch replicas.
dp_rank = dp_gather_slot()
if forward_batch.dp_local_start_pos is None:
cumtokens = torch.cumsum(forward_batch.global_num_tokens_gpu, dim=0)
@@ -438,7 +469,10 @@ def get_dp_local_slice_cpu(
# CPU (start, length) slice for DP-local data in a rank-padded buffer.
# Returns Python ints (no D2H sync) and handles the cuda-graph-padded layout.
global_num_tokens = forward_batch.global_num_tokens_cpu
dp_rank = get_parallel().attn_dp_rank
# The slot in the list that was gathered. A scale-up widens that list
# to WORLD, and this process's index in it is not its index among the
# launch replicas.
dp_rank = dp_gather_slot()
local_num_tokens = global_num_tokens[dp_rank]
if can_run_graph:
local_start_pos = dp_rank * cuda_graph_batch
@@ -514,7 +548,7 @@ def _dp_gather_via_all_reduce(
NUM_GPUS_PER_NODE = 8
if (
not local_tokens.dtype.is_floating_point
and get_tensor_model_parallel_world_size() <= NUM_GPUS_PER_NODE
and get_parallel().tp_size <= NUM_GPUS_PER_NODE
):
from sglang.srt.distributed.parallel_state import inplace_all_reduce
@@ -534,7 +568,7 @@ def _dp_gather_via_all_gather(
):
use_world = world_dp_gather_enabled()
if get_attn_tensor_model_parallel_world_size() == 1:
if get_parallel().attn_tp_size == 1:
if use_world:
torch.distributed.all_gather_into_tensor(
global_tokens,
@@ -548,9 +582,9 @@ def _dp_gather_via_all_gather(
if not is_partial:
if get_parallel().attn_tp_rank != 0:
local_tokens.fill_(0)
scattered_local_tokens = local_tokens.tensor_split(
get_attn_tensor_model_parallel_world_size()
)[get_parallel().attn_tp_rank]
scattered_local_tokens = local_tokens.tensor_split(get_parallel().attn_tp_size)[
get_parallel().attn_tp_rank
]
get_parallel().attn_tp_group.reduce_scatter_tensor(
scattered_local_tokens, local_tokens
)
@@ -721,8 +755,8 @@ def is_dp_gatherv_active() -> bool:
return (
_USE_DP_GATHERV
and not world_dp_gather_enabled()
and get_attn_tensor_model_parallel_world_size() == 1
and get_tensor_model_parallel_world_size() == get_parallel().attn_dp_size
and get_parallel().attn_tp_size == 1
and get_parallel().tp_size == get_parallel().attn_dp_size
and not _DpGatheredBufferWrapper.is_dp_max_padding()
)
@@ -881,12 +915,12 @@ def dp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
if sizes is not None:
get_parallel().tp_group.reduce_scatterv(input, output=output, sizes=sizes)
return
if get_tensor_model_parallel_world_size() == get_parallel().attn_dp_size:
if get_parallel().tp_size == get_parallel().attn_dp_size:
get_parallel().tp_group.reduce_scatter_tensor(output, input)
else:
scattered_local_tokens = input.tensor_split(
get_tensor_model_parallel_world_size()
)[get_parallel().tp_rank]
scattered_local_tokens = input.tensor_split(get_parallel().tp_size)[
get_parallel().tp_rank
]
get_parallel().tp_group.reduce_scatter_tensor(scattered_local_tokens, input)
get_parallel().attn_tp_group.all_gather_into_tensor(
output, scattered_local_tokens
python/sglang/srt/layers/sampler.py
+6 -1
View File
@@ -112,7 +112,12 @@ class Sampler(nn.Module):
self.cp_sync_group = None
if is_dp_attention_enabled():
self.tp_sync_group = get_parallel().attn_tp_group.device_group
self.cp_sync_group = get_parallel().attn_cp_group.device_group
# Only when there is more than one context shard to reconcile. The
# sync below already short-circuits on that, and a model running on
# one shard -- a speculative draft, under the scope that says so --
# has no context-parallel communicator to name.
if get_parallel().attn_cp_size > 1:
self.cp_sync_group = get_parallel().attn_cp_group.device_group
self.rl_on_policy_target = get_exec().deterministic.rl_on_policy_target
# In RL on-policy mode, deterministic inference is automatically enabled.
python/sglang/srt/managers/scheduler_components/dp_attn.py
+2 -2
View File
@@ -10,7 +10,7 @@ from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.layers.cp.utils import get_cp_strategy
from sglang.srt.layers.dp_attention import world_dp_gather_enabled
from sglang.srt.layers.dp_attention import dp_gather_width, world_dp_gather_enabled
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler_components.recv_skipper import (
@@ -58,7 +58,7 @@ def _resolve_elastic_world_dp_size(
from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
live_dp_size = get_parallel().attn_dp_size
live_dp_size = dp_gather_width()
effective_ep_size = ElasticEPStateManager.get_effective_ep_size()
# The group's own membership, not the width it was built at: this is the
# one number an out-of-process join moves, and it is the upper bound the
python/sglang/srt/model_executor/forward_batch_info.py
+3 -4
View File
@@ -45,6 +45,7 @@ from sglang.srt.kv_canary.req_to_expected_token_ids_manager import (
)
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
dp_gather_slot,
set_dp_buffer_len,
set_is_extend_in_batch,
world_dp_gather_enabled,
@@ -1165,9 +1166,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
if self.global_num_tokens_cpu is not None:
# DP / MLP-sync path: per-DP padded width.
if require_mlp_tp_gather():
num_tokens_per_dp = self.global_num_tokens_cpu[
get_parallel().attn_dp_rank
]
num_tokens_per_dp = self.global_num_tokens_cpu[dp_gather_slot()]
else:
num_tokens_per_dp = self.global_num_tokens_cpu[0]
else:
@@ -1520,7 +1519,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
buffer_len = sum(global_num_tokens)
if len(global_num_tokens) > 1:
num_tokens = global_num_tokens[get_parallel().attn_dp_rank]
num_tokens = global_num_tokens[dp_gather_slot()]
else:
num_tokens = global_num_tokens[0]
python/sglang/srt/models/inkling_common/moe.py
+6
View File
@@ -673,9 +673,15 @@ class InklingSharedFusedMoE(FusedMoE):
) -> None:
# FusedMoE.__init__ reads get_parallel() once and caches it on self, so
# scoping the override to just this call is sufficient for the module's lifetime.
# The shared experts are replicated rather than sharded, so there is no
# expert-parallel communication here and no group to name: a width of
# one with the wider group still installed would describe a layout that
# does not exist.
with get_parallel().override(
moe_ep_size=1,
moe_ep_rank=0,
moe_ep_group=None,
moe_dp_size=1,
moe_tp_size=get_parallel().tp_size,
moe_tp_rank=get_parallel().tp_rank,
):
python/sglang/srt/models/kimi_k3.py
+1 -2
View File
@@ -22,7 +22,6 @@ from sglang.srt.configs.kimi_k3 import KimiK3Config
from sglang.srt.configs.kimi_linear import KimiLinearConfig
from sglang.srt.distributed import (
divide,
get_shared_experts_tp_group,
tensor_model_parallel_all_reduce,
)
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -582,7 +581,7 @@ class KimiK3MoE(nn.Module):
shared_experts_tp_kwargs = dict(tp_rank=0, tp_size=1)
elif self._shared_experts_tp_comm:
group = (
get_shared_experts_tp_group()
parallel.shared_experts_tp_group
if requested_shared_tp is not None
else parallel.attn_tp_group
)
python/sglang/srt/runtime_context.py
+153 -9
View File
@@ -218,6 +218,7 @@ _LIVE_READS: dict = {
"moe_tp_group": "get_moe_tp_group",
"attn_tp_group": "get_attn_tp_group",
"attn_cp_group": "get_attn_cp_group",
"shared_experts_tp_group": "get_shared_experts_tp_group",
"dcp_group": "get_dcp_group",
}
@@ -448,14 +449,138 @@ class SpawnRanks(msgspec.Struct, frozen=True):
the rank cannot say which replica this is. `None` means "no controller",
which is an answer rather than an absence, and it is recorded as one.
Nothing else belongs here. A device index, for instance, is a placement
decision rather than a position -- the launcher may reindex it, and Ray
assigns it from its own allocator -- so it stays an argument to whoever
was handed it.
`gpu_id` is the device the parent picked for this process. It is not a
position in any group -- reindexing narrows the visible devices before the
spawn, and Ray allocates from its own pool -- but it is the same kind of
fact: something only the entry that spawned the process can state. `None`
for a process that runs on no device.
"""
world_rank: int
dp_rank: Optional[int] = None
gpu_id: Optional[int] = None
_RANK_AND_WIDTH = (
("tp_rank", "tp_size"),
("pp_rank", "pp_size"),
("attn_tp_rank", "attn_tp_size"),
("attn_dp_rank", "attn_dp_size"),
("attn_cp_rank", "attn_cp_size"),
("moe_ep_rank", "moe_ep_size"),
)
# `moe_dp` is absent because `initialize_model_parallel` aliases the MoE-DP
# group to the attention-CP group when the latter is wider: there the group and
# the name are two facts, which is the same reason `moe_dp_rank` is left off the
# record at publish.
_WIDTH_AND_GROUP = (
("tp_size", "tp_group"),
("pp_size", "pp_group"),
("attn_tp_size", "attn_tp_group"),
("attn_cp_size", "attn_cp_group"),
("moe_ep_size", "moe_ep_group"),
)
_UNREADABLE = object()
def _validate_parallel(parallel, source: str) -> None:
"""Fail on a topology that cannot describe a real process layout.
Every identity holds unconditionally: a width and a rank are both
plausible small integers whichever way they are wrong, so an inconsistent
set is not caught by anything downstream -- it surfaces as a hang or a
wrong answer in a collective, far from the write. Stating one leaf without
the quotients that follow from it leaves the namespace describing no real
layout, and the caller that did so is the one that has to say what it meant.
Names that cannot be read are skipped rather than treated as zero: a
process that has published nothing can still stamp a rank, and a group that
has not been built answers nothing at all.
"""
def read(name):
"""A width or a rank, or `_UNREADABLE` for anything these identities
cannot be stated about -- an absent name, `None`, or a stand-in a test
put in a group's place. Booleans are integers in Python and are not
widths, so they are out too."""
try:
value = getattr(parallel, name)
except Exception:
return _UNREADABLE
if isinstance(value, bool) or not isinstance(value, int):
return _UNREADABLE
return value
problems = []
for rank_name, size_name in _RANK_AND_WIDTH:
rank, size = read(rank_name), read(size_name)
if _UNREADABLE in (rank, size):
continue
if not 0 <= rank < size:
problems.append(
f"0 <= {rank_name} < {size_name}\n {rank} is not a rank of {size}"
)
terms = ("tp_size", "attn_tp_size", "attn_dp_size", "attn_cp_size")
tp_size, a_tp, a_dp, a_cp = (read(n) for n in terms)
if _UNREADABLE not in (tp_size, a_tp, a_dp, a_cp):
if tp_size != a_tp * a_dp * a_cp:
problems.append(
"tp_size == attn_tp_size * attn_dp_size * attn_cp_size\n"
f" {tp_size} != {a_tp} * {a_dp} * {a_cp} (= {a_tp * a_dp * a_cp})"
)
moe_terms = ("tp_size", "moe_ep_size", "moe_dp_size", "moe_tp_size")
tp_size, m_ep, m_dp, m_tp = (read(n) for n in moe_terms)
if _UNREADABLE not in (tp_size, m_ep, m_dp, m_tp):
if tp_size != m_ep * m_dp * m_tp:
problems.append(
"tp_size == moe_ep_size * moe_dp_size * moe_tp_size\n"
f" {tp_size} != {m_ep} * {m_dp} * {m_tp} (= {m_ep * m_dp * m_tp})"
)
layout_terms = (
"tp_rank",
"attn_dp_rank",
"attn_cp_rank",
"attn_tp_rank",
"attn_cp_size",
"attn_tp_size",
)
tp_rank, r_dp, r_cp, r_tp, w_cp, w_tp = (read(n) for n in layout_terms)
if _UNREADABLE not in (tp_rank, r_dp, r_cp, r_tp, w_cp, w_tp):
laid_out = (r_dp * w_cp + r_cp) * w_tp + r_tp
if tp_rank != laid_out:
problems.append(
"tp_rank == (attn_dp_rank * attn_cp_size + attn_cp_rank)"
" * attn_tp_size + attn_tp_rank\n"
f" {tp_rank} != ({r_dp} * {w_cp} + {r_cp})"
f" * {w_tp} + {r_tp} (= {laid_out})"
)
for size_name, group_name in _WIDTH_AND_GROUP:
size = read(size_name)
if size is _UNREADABLE:
continue
try:
group = getattr(parallel, group_name)
except Exception:
continue
built = getattr(group, "world_size", _UNREADABLE)
if isinstance(built, int) and not isinstance(built, bool) and built != size:
problems.append(
f"{group_name}.world_size == {size_name}\n"
f" built {built}, configured {size}"
)
if problems:
raise ValueError(
f"parallel topology is inconsistent (set by {source}):\n"
+ "\n".join(problems)
)
class ParallelContext:
@@ -561,7 +686,13 @@ class ParallelContext:
unknown = set(values) - _parallel_fields()
if unknown:
raise ValueError(f"unknown parallel field(s): {sorted(unknown)}")
saved = dict(self._stamp)
self._stamp.update(values)
try:
_validate_parallel(self, "override_permanently")
except Exception:
self._stamp = saved
raise
def clear_stamp(self) -> None:
"""Drop every stamped name, ranks included."""
@@ -576,6 +707,11 @@ class ParallelContext:
raise ValueError(f"unknown parallel field(s): {sorted(unknown)}")
saved = dict(self._overrides)
self._overrides.update(kwargs)
try:
_validate_parallel(self, "override")
except Exception:
self._overrides = saved
raise
try:
yield self
finally:
@@ -1776,6 +1912,8 @@ def publish(
),
)
_CONTEXT._publish_role = role
if ranks is not None and ranks.gpu_id is not None:
_CONTEXT.override("spawn", gpu_id=ranks.gpu_id)
if ranks is not None:
# The placement, worked out here rather than carried: the widths are on
# the bag a moment ago, and `world_rank` fixes the rest. A read of any
@@ -1804,8 +1942,13 @@ def publish(
# "no controller" rather than an absence.
placement["dp_rank"] = ranks.dp_rank
placement["launch_world_rank"] = ranks.world_rank
placement.update(_attention_ranks(parallel, placement["tp_rank"]))
# One stamp, not two: the identities are checked on every write, and a
# half-placed process satisfies none of them.
parallel.override_permanently(**placement)
_stamp_attention_ranks(parallel, placement["tp_rank"])
# Publish established the whole layout, so every identity applies here,
# not just the ones the stamp happened to name.
_validate_parallel(parallel, "publish")
if _ROLE_NS_MODE == "record":
# The '-' marker distinguishes a zero-read role from a process where
# recording never ran (signal teardown skips atexit).
@@ -1820,7 +1963,7 @@ def publish(
return _CONTEXT
def _stamp_attention_ranks(parallel, tp_rank: int) -> None:
def _attention_ranks(parallel, tp_rank: int) -> dict:
"""Place this process in the attention topology, from the configuration.
The widths are already on the bag -- `publish` computed them a moment ago --
@@ -1829,8 +1972,9 @@ def _stamp_attention_ranks(parallel, tp_rank: int) -> None:
that never initialises distributed, which is what `ParallelState` provided
by being a plain frozen record.
It is a stamp rather than a bag leaf because it is a per-process fact, and
nothing about the configuration distinguishes one rank from another.
These are stamped rather than written as bag leaves because they are
per-process facts, and nothing about the configuration distinguishes one
rank from another.
"""
attn_tp_rank, attn_dp_rank = derive_attention_ranks(
tp_rank=tp_rank,
@@ -1838,7 +1982,7 @@ def _stamp_attention_ranks(parallel, tp_rank: int) -> None:
attn_cp_size=parallel.attn_cp_size,
enable_dp_attention=parallel.enable_dp_attention,
)
parallel.override_permanently(attn_tp_rank=attn_tp_rank, attn_dp_rank=attn_dp_rank)
return {"attn_tp_rank": attn_tp_rank, "attn_dp_rank": attn_dp_rank}
def assert_published(server_args, *, role: str) -> RuntimeContext:
python/sglang/srt/weight_cache/daemon.py
+1 -10
View File
@@ -232,16 +232,7 @@ class WeightCacheDaemon:
moe_a2a_backend=self.moe_a2a_backend,
)
initialize_model_parallel(
tensor_model_parallel_size=self.tp_size,
pipeline_model_parallel_size=self.pp_size,
expert_model_parallel_size=self.ep_size,
attention_data_parallel_size=(
self.dp_size if self.enable_dp_attention else 1
),
attention_context_model_parallel_size=self.attn_cp_size,
moe_data_model_parallel_size=self.moe_dp_size,
)
initialize_model_parallel()
# Initialize DP attention state (required by some models like Qwen3 MoE)
from sglang.srt.layers.dp_attention import initialize_dp_attention
python/sglang/test/layer_ut_utils.py
+4 -6
View File
@@ -7,6 +7,8 @@ import os
import torch
from sglang.test.test_utils import publish_build_topology
def init_single_process_dist(master_port: int = 29632, backend: str = "gloo"):
"""world=1 dist + model-parallel groups; srt layers require them even
@@ -29,12 +31,8 @@ def init_single_process_dist(master_port: int = 29632, backend: str = "gloo"):
if not model_parallel_is_initialized():
# kwargs only: a positional backend would land in the
# attention_data_parallel_size slot and explode on int // str.
initialize_model_parallel(
tensor_model_parallel_size=1,
expert_model_parallel_size=1,
pipeline_model_parallel_size=1,
backend=backend,
)
publish_build_topology(tp_size=1, ep_size=1, pp_size=1)
initialize_model_parallel(backend=backend)
def make_tp1_column_parallel_linear(
python/sglang/test/test_utils.py
+22
View File
@@ -2084,6 +2084,28 @@ def published_topology(role: str = "test", *, ranks=None, **server_args_fields):
reset_context()
def publish_build_topology(*, world_rank: int = 0, **server_args_fields):
"""State the widths `initialize_model_parallel` is about to build at.
The build reads every width from the runtime context, so a test that wants
a particular topology publishes it here rather than passing it in -- the
same door production uses, which also keeps the derived widths honest.
Unlike `published_topology` this is not a scope: the groups it is about to
build outlive any block, so the configuration describing them has to as
well. Callers that tear the groups down are already resetting the process.
"""
from sglang.srt.runtime_context import SpawnRanks, publish, reset_context
from sglang.srt.server_args import ServerArgs
reset_context()
publish(
ServerArgs(model_path="dummy", **server_args_fields),
role="test",
ranks=SpawnRanks(world_rank=world_rank),
)
_GPU_IDLE_TIMEOUT_SECS = 30.0
_GPU_IDLE_POLL_INTERVAL_SECS = 2.0
_GPU_IDLE_USED_MEMORY_THRESHOLD = 2 << 30 # 2 GiB
test/manual/ep/test_flashinfer_dispatcher.py
+3 -4
View File
@@ -14,7 +14,7 @@ from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatch
from sglang.srt.layers.moe.utils import initialize_moe_config
from sglang.srt.runtime_context import get_context, publish
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.test_utils import CustomTestCase
from sglang.test.test_utils import CustomTestCase, publish_build_topology
class TestFlashinferDispatcher(CustomTestCase):
@@ -44,9 +44,8 @@ class TestFlashinferDispatcher(CustomTestCase):
publish(server_args, role="scheduler")
initialize_moe_config()
initialize_model_parallel(
tensor_model_parallel_size=world_size, expert_model_parallel_size=world_size
)
publish_build_topology(tp_size=world_size, ep_size=world_size, world_rank=rank)
initialize_model_parallel()
@classmethod
def tearDownClass(cls):
test/manual/test_custom_allreduce.py
+5 -3
View File
@@ -18,7 +18,7 @@ from sglang.srt.distributed.parallel_state import (
initialize_model_parallel,
)
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.test_utils import CustomTestCase
from sglang.test.test_utils import CustomTestCase, publish_build_topology
def get_open_port() -> int:
@@ -98,7 +98,8 @@ class TestCustomAllReduce(CustomTestCase):
distributed_init_method=distributed_init_method,
local_rank=rank,
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(tp_size=world_size, world_rank=rank)
initialize_model_parallel()
group = get_tensor_model_parallel_group().device_group
# Set global server args to avoid "Global server args is not set yet!" error
@@ -161,7 +162,8 @@ class TestCustomAllReduce(CustomTestCase):
distributed_init_method=distributed_init_method,
local_rank=rank,
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(tp_size=world_size, world_rank=rank)
initialize_model_parallel()
group = get_tensor_model_parallel_group().device_group
# Set global server args to avoid "Global server args is not set yet!" error
test/manual/test_quick_allreduce.py
+5 -3
View File
@@ -23,7 +23,7 @@ from sglang.srt.distributed.parallel_state import (
graph_capture,
initialize_model_parallel,
)
from sglang.test.test_utils import CustomTestCase
from sglang.test.test_utils import CustomTestCase, publish_build_topology
torch.manual_seed(42)
random.seed(44) # keep the deterministic seed
@@ -117,7 +117,8 @@ class TestQuickAllReduce(CustomTestCase):
distributed_init_method=distributed_init_method,
local_rank=rank,
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(tp_size=world_size, world_rank=rank)
initialize_model_parallel()
group = get_tensor_model_parallel_group().device_group
# A small all_reduce for warmup.
@@ -186,7 +187,8 @@ class TestQuickAllReduce(CustomTestCase):
distributed_init_method=distributed_init_method,
local_rank=rank,
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(tp_size=world_size, world_rank=rank)
initialize_model_parallel()
group = get_tensor_model_parallel_group().device_group
for sz in self.TEST_SIZES:
test/manual/test_weight_loader_v2_equiv.py
+3 -1
View File
@@ -24,6 +24,7 @@ import unittest
import torch
from sglang.srt.environ import envs
from sglang.test.test_utils import publish_build_topology
MODEL = "Qwen/Qwen2-0.5B"
@@ -43,7 +44,8 @@ def _init_model_parallel() -> None:
local_rank=0,
distributed_init_method="tcp://127.0.0.1:29634",
)
initialize_model_parallel(tensor_model_parallel_size=1)
publish_build_topology(tp_size=1)
initialize_model_parallel()
monkey_patch_vllm_parallel_state()
except AssertionError:
pass
test/registered/cpu/test_rank_consensus_checker.py
+7 -6
View File
@@ -23,7 +23,11 @@ from sglang.srt.utils.rank_consensus_checker import (
shutdown,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase, find_available_port
from sglang.test.test_utils import (
CustomTestCase,
find_available_port,
publish_build_topology,
)
register_cpu_ci(est_time=193, suite="stage-a-test-cpu-intel")
@@ -80,11 +84,8 @@ def run_distributed_test(
backend="gloo",
)
initialize_model_parallel(
tensor_model_parallel_size=tp_size,
pipeline_model_parallel_size=pp_size,
backend="gloo",
)
publish_build_topology(tp_size=tp_size, pp_size=pp_size, world_rank=rank)
initialize_model_parallel(backend="gloo")
fn()
except Exception as e:
test/registered/e2e/pp/test_pp_spec_embed_scan.py
+7 -2
View File
@@ -10,7 +10,11 @@ import torch
from transformers import MistralConfig, PretrainedConfig
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
CustomTestCase,
publish_build_topology,
)
register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-small")
@@ -77,7 +81,8 @@ class TestDraftEmbedScan(CustomTestCase):
init_distributed_environment(
world_size=1, rank=0, local_rank=0, distributed_init_method="env://"
)
initialize_model_parallel(tensor_model_parallel_size=1)
publish_build_topology(tp_size=1)
initialize_model_parallel()
torch.set_default_dtype(torch.bfloat16)
torch.cuda.set_device(0)
test/registered/eplb/test_lplb_distributed.py
+4 -3
View File
@@ -38,6 +38,7 @@ from sglang.srt.distributed.parallel_state import (
initialize_model_parallel,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import publish_build_topology
register_cuda_ci(est_time=18, stage="base-b", runner_config="2-gpu-large")
@@ -238,10 +239,10 @@ def _worker_main(local_rank: int, world_size: int):
init_distributed_environment(
world_size=world_size, rank=local_rank, local_rank=local_rank
)
initialize_model_parallel(
tensor_model_parallel_size=world_size,
expert_model_parallel_size=world_size,
publish_build_topology(
tp_size=world_size, ep_size=world_size, world_rank=local_rank
)
initialize_model_parallel()
from sglang.srt.eplb.lplb_solver import clear_global_lplb_solvers
test/registered/kernels/ops/communication/test_deterministic_reduce_scatter.py
+3 -1
View File
@@ -14,6 +14,7 @@ from sglang.srt.distributed import parallel_state as ps
from sglang.srt.environ import envs
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kernels.utils import multigpu_pytest_main
from sglang.test.test_utils import publish_build_topology
register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
@@ -37,7 +38,8 @@ def group():
local_rank=local_rank,
distributed_init_method="env://",
)
ps.initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(tp_size=world_size, world_rank=rank)
ps.initialize_model_parallel()
yield ps.get_tp_group()
ps.destroy_model_parallel()
ps.destroy_distributed_environment()
test/registered/kernels/ops/moe/test_fused_swiglu_epilogue.py
+3 -6
View File
@@ -19,6 +19,7 @@ import pytest
import torch
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import publish_build_topology
register_cuda_ci(est_time=12, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@@ -52,12 +53,8 @@ def _runtime_scaffolding():
if not torch.distributed.is_initialized():
init_distributed_environment(world_size=1, rank=0, local_rank=0, backend="gloo")
if not model_parallel_is_initialized():
initialize_model_parallel(
tensor_model_parallel_size=1,
expert_model_parallel_size=1,
pipeline_model_parallel_size=1,
backend="gloo",
)
publish_build_topology(tp_size=1, ep_size=1, pp_size=1)
initialize_model_parallel(backend="gloo")
def _interleave_w13_rows(w13: torch.Tensor) -> torch.Tensor:
test/registered/layers/mamba/test_mamba2_mixer.py
+3 -1
View File
@@ -17,6 +17,7 @@ from sglang.srt.distributed.parallel_state import (
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import get_device, get_device_count
from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci
from sglang.test.test_utils import publish_build_topology
register_cuda_ci(est_time=30, stage="base-b", runner_config="2-gpu-large")
register_xpu_ci(est_time=60, suite="nightly-xpu-2-gpu", nightly=True)
@@ -105,7 +106,8 @@ def mixer2_gated_norm_tensor_parallel(
local_rank=local_rank,
backend=get_default_distributed_backend(device.type),
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(tp_size=world_size, world_rank=local_rank)
initialize_model_parallel()
# create random weights an inputs
weight = torch.rand((hidden_size,), dtype=dtype, device=device)
test/registered/ops/test_aiter_allreduce_fusion_amd.py
+3 -2
View File
@@ -14,7 +14,7 @@ import torch
from sglang.srt.layers import communicator as comm
from sglang.srt.layers.communicator import LayerCommunicator, ScatterMode
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
from sglang.test.test_utils import CustomTestCase, publish_build_topology
register_amd_ci(est_time=240, suite="stage-c-test-large-8-gpu-amd")
@@ -64,7 +64,8 @@ def _run_residual_accuracy_check():
distributed_init_method="env://",
backend="nccl",
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
publish_build_topology(tp_size=world_size, world_rank=rank)
initialize_model_parallel()
dtype = torch.bfloat16
eps = 1e-6
test/registered/unit/distributed/test_parallel_state.py
+5 -11
View File
@@ -44,6 +44,7 @@ import pytest
import torch
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import publish_build_topology
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
@@ -232,11 +233,8 @@ def test_parallel_group_construction_tp8_attn_cp2():
mock_world_group.return_value = mock_world
# Call the actual function
parallel_state.initialize_model_parallel(
tensor_model_parallel_size=8,
pipeline_model_parallel_size=1,
attention_context_model_parallel_size=2,
)
publish_build_topology(tp_size=8, pp_size=1, attn_cp_size=2)
parallel_state.initialize_model_parallel()
# Verify TP groups
tp_groups = created_groups.get("tp", [])
@@ -330,12 +328,8 @@ def test_parallel_group_construction_tp8_moe_ep4_cp2():
mock_world_group.return_value = mock_world
# Call the actual function
parallel_state.initialize_model_parallel(
tensor_model_parallel_size=8,
expert_model_parallel_size=4,
pipeline_model_parallel_size=1,
moe_data_model_parallel_size=2,
)
publish_build_topology(tp_size=8, ep_size=4, pp_size=1, moe_dp_size=2)
parallel_state.initialize_model_parallel()
# Verify TP groups
tp_groups = created_groups.get("tp", [])
test/registered/unit/layers/attention/test_vision_strided_qkv.py
+3 -1
View File
@@ -16,6 +16,7 @@ from sglang.srt.distributed.parallel_state import (
from sglang.srt.layers.attention import vision
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import publish_build_topology
register_cpu_ci(est_time=12, suite="base-a-test-cpu")
@@ -52,7 +53,8 @@ def gloo_world():
distributed_init_method=f"tcp://127.0.0.1:{port}",
backend="gloo",
)
initialize_model_parallel(tensor_model_parallel_size=1, backend="gloo")
publish_build_topology(tp_size=1)
initialize_model_parallel(backend="gloo")
yield
destroy_model_parallel()
destroy_distributed_environment()
test/registered/unit/mem_cache/test_dsa_layer_split_broadcast.py
+3 -5
View File
@@ -20,7 +20,7 @@ import torch
import torch.multiprocessing as mp
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
from sglang.test.test_utils import CustomTestCase, publish_build_topology
register_cuda_ci(est_time=28, stage="base-c", runner_config="4-gpu-b200")
@@ -54,10 +54,8 @@ def _run(rank: int, world: int, port: int):
distributed_init_method=f"tcp://127.0.0.1:{port}",
backend="nccl",
)
initialize_model_parallel(
tensor_model_parallel_size=world,
attention_context_model_parallel_size=world,
)
publish_build_topology(tp_size=world, attn_cp_size=world, world_rank=rank)
initialize_model_parallel()
from sglang.srt.mem_cache.dsa_cache_layer_split import (
LayerSplitDSATokenToKVPool,
test/registered/unit/mem_cache/test_page_interleave_shard.py
+3 -5
View File
@@ -84,7 +84,7 @@ from sglang.srt.runtime_context import get_parallel, publish
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import ceil_div
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
from sglang.test.test_utils import CustomTestCase, publish_build_topology
register_cpu_ci(est_time=30, suite="base-a-test-cpu")
@@ -1215,10 +1215,8 @@ def _dist_init(rank, world, port, attn_cp_size):
ServerArgs(model_path="dummy", tp_size=world, attn_cp_size=attn_cp_size),
role="scheduler",
)
initialize_model_parallel(
tensor_model_parallel_size=world,
attention_context_model_parallel_size=attn_cp_size,
)
publish_build_topology(tp_size=world, attn_cp_size=attn_cp_size, world_rank=rank)
initialize_model_parallel()
def _gather_make_spec(shard_rank, max_prefix_groups=16, chunk_groups=4):
test/registered/unit/model_loader/test_presharded_loader.py
+3 -1
View File
@@ -844,7 +844,9 @@ class TestShardConfig(unittest.TestCase):
override.install()
self.addCleanup(override.restore)
with (
get_parallel().override(tp_size=8, pp_size=1, moe_dp_size=2, moe_ep_size=4),
get_parallel().override(
tp_size=8, pp_size=1, moe_dp_size=2, moe_ep_size=4, moe_tp_size=1
),
mock.patch(
"sglang.srt.layers.dp_attention.get_moe_cp_size",
return_value=2,
test/registered/unit/models/test_glm5_next_bfg_fusion.py
+12 -1
View File
@@ -83,8 +83,19 @@ class TestGlm5NextBfgFusion(unittest.TestCase):
for attn_tp, rank in ((1, 0), (2, 0), (2, 1)):
with (
self.subTest(route=expected_route, attn_tp=attn_tp, rank=rank),
# A width is a whole topology: the attention triple has
# to factor `tp_size`, and this process has to sit where
# the triple puts it.
get_parallel().override(
tp_size=4, tp_rank=3, attn_tp_size=attn_tp, attn_tp_rank=rank
tp_size=4,
tp_rank=rank,
attn_tp_size=attn_tp,
attn_tp_rank=rank,
attn_dp_size=4 // attn_tp,
attn_dp_rank=0,
attn_cp_size=1,
attn_cp_rank=0,
moe_tp_size=4,
),
):
quant = MockFp8Config(ignored)
test/registered/unit/models/test_shared_experts_fusion_gates.py
+9 -1
View File
@@ -61,7 +61,15 @@ class _FusionGateCase(CustomTestCase):
def _reason(self, model_class, hf_config, quant_config=None, moe_ep_size=1):
# The gates consult the live EP size; without a group installed the
# canonical getter asserts, so every case states a topology.
with get_parallel().override(moe_ep_size=moe_ep_size):
with get_parallel().override(
tp_size=moe_ep_size,
attn_tp_size=moe_ep_size,
attn_dp_size=1,
attn_cp_size=1,
moe_ep_size=moe_ep_size,
moe_dp_size=1,
moe_tp_size=1,
):
return model_class.shared_experts_fusion_disable_reason(
hf_config, quant_config
)
test/registered/unit/test_runtime_context.py
+392 -65
View File
@@ -41,6 +41,7 @@ from sglang.srt.runtime_context import (
RuntimeContext,
SpawnRanks,
_FlagGroupBase,
_validate_parallel,
assert_published,
derive_parallel_widths,
get_context,
@@ -59,6 +60,56 @@ from sglang.srt.server_args import ServerArgs
from sglang.test.test_utils import CustomTestCase
_SRT = _pathlib.Path(next(iter(_sglang.__path__))).resolve() / "srt"
_PACKAGE = _pathlib.Path(next(iter(_sglang.__path__))).resolve()
def _sources():
"""Every Python file this checkout ships, package and siblings alike.
The package alone is the wrong subject set for anything about entries or
public names: `benchmark/`, `examples/` and the top-level `test/` call the
same doors and are not covered by any suite that would notice them break.
An installed package has no siblings, and then this is the package alone."""
roots = [_PACKAGE]
checkout = _PACKAGE.parents[1]
roots += [
checkout / name
for name in ("benchmark", "examples", "scripts", "test")
if (checkout / name).is_dir()
]
for root in roots:
for path in root.rglob("*.py"):
yield path
def _scope_entries_that_say_nothing(paths):
"""Draft-scope entries that do not state `owns_attention`, as `path:line`.
The scope either narrows the draft's attention and expert identity or
leaves the target's in place, and only the worker knows which -- so the
keyword has no default. Omitting it is a `TypeError`, but only on the path
that runs, and those paths want a GPU and a draft model.
"""
import ast
missing = []
for path in paths:
try:
tree = ast.parse(path.read_text(encoding="utf-8"))
except (SyntaxError, UnicodeDecodeError):
continue
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
func = node.func
name = getattr(func, "attr", None) or getattr(func, "id", None)
if name not in ("draft_tp_context", "patch_tensor_parallel_group"):
continue
if not any(kw.arg == "owns_attention" for kw in node.keywords):
missing.append(f"{path}:{node.lineno}")
return missing
_PS = "sglang.srt.distributed.parallel_state"
_DP = "sglang.srt.layers.dp_attention"
@@ -347,9 +398,8 @@ class TestStampedRanks(_IsolatedOverrides):
"""`attn_dp_rank` comes from the stamp, and says so when there is none.
It is the one rank no group answers with: `initialize_dp_attention`
computes it from this process's `tp_rank`, and an elastic scale-up
replaces it with a rank in the expanded WORLD. Falling back to anything
would be inventing a placement for this process.
computes it from this process's `tp_rank`. Falling back to anything would
be inventing a placement for this process.
"""
def setUp(self):
@@ -407,21 +457,71 @@ class TestStampedRanks(_IsolatedOverrides):
)
self.assertIs(mode, DpPaddingMode.SUM_LEN)
def test_a_scale_up_stamps_the_width_and_the_rank_together(self):
"""The two describe one topology; a reader that saw only one moved
would place this process in a group it is not in."""
def test_the_gather_slot_follows_the_list_that_was_gathered(self):
"""The DP sync gathers over the attention-DP replicas, or over the
expanded WORLD once a scale-up has moved the gather there. The index
into that list is a property of the gather, so it is read beside the
flag that says which one happened rather than kept on the topology."""
from sglang.srt.layers.dp_attention import dp_gather_slot
self.addCleanup(reset_context)
dp_flags = get_flags().dp
saved = (
dp_flags.use_world_group_for_gather,
dp_flags.joiner_skip_all_gather,
)
def restore():
(
dp_flags.use_world_group_for_gather,
dp_flags.joiner_skip_all_gather,
) = saved
self.addCleanup(restore)
publish(
ServerArgs(
model_path="dummy", tp_size=8, dp_size=8, enable_dp_attention=True
),
role="test",
ranks=SpawnRanks(world_rank=3),
)
parallel = get_parallel()
dp_flags.use_world_group_for_gather = False
self.assertEqual(dp_gather_slot(), parallel.attn_dp_rank)
# After a scale-up the gather spans the expanded WORLD, and the joining
# cohort is numbered from its offset.
dp_flags.use_world_group_for_gather = True
dp_flags.joiner_skip_all_gather = False
parallel.override_permanently(ep_join_rank_offset=8)
self.assertEqual(dp_gather_slot(), 8 + parallel.tp_rank)
# and the topology it was read off is untouched
self.assertEqual(parallel.attn_dp_size, 8)
self.assertEqual(parallel.tp_size, 8)
def test_a_scale_up_writes_no_width(self):
"""The identities stay unconditional because nothing overrides them:
the scale-up only points the gather at the expanded WORLD."""
from sglang.srt.layers.dp_attention import update_dp_attention_post_scale
# It also flips a process-wide gather flag; put it back, or every
# later test in this process runs as if a scale-up had happened.
dp_flags = get_flags().dp
saved_gather = dp_flags.use_world_group_for_gather
self.addCleanup(setattr, dp_flags, "use_world_group_for_gather", saved_gather)
self.addCleanup(reset_context)
publish(
ServerArgs(
model_path="dummy", tp_size=8, dp_size=8, enable_dp_attention=True
),
role="test",
ranks=SpawnRanks(world_rank=3),
)
parallel = get_parallel()
before = (parallel.attn_dp_size, parallel.attn_dp_rank, parallel.tp_size)
update_dp_attention_post_scale(new_dp_size=16, new_dp_rank=11)
self.assertEqual(parallel.attn_dp_size, 16)
self.assertEqual(parallel.attn_dp_rank, 11)
self.assertTrue(dp_flags.use_world_group_for_gather)
self.assertEqual(
(parallel.attn_dp_size, parallel.attn_dp_rank, parallel.tp_size), before
)
class TestEveryDeclaredParallelNameIsStatable(_IsolatedOverrides):
@@ -1898,14 +1998,26 @@ class TestDerivedWidths(_IsolatedOverrides):
def test_a_topology_is_stated_by_naming_the_width(self):
"""Overriding a leaf does not move the quotient -- the quotient is not
recomputed on read. Naming it is how a test states one."""
recomputed on read. Naming it is how a caller states one, and naming
only some of them is refused: the caller owns the arithmetic, the
context only checks it."""
reset_context()
self.addCleanup(reset_context)
publish(ServerArgs(model_path="dummy", tp_size=8), role="test")
self.assertEqual(get_parallel().attn_tp_size, 8)
with get_parallel().override(tp_size=2):
self.assertEqual(get_parallel().attn_tp_size, 8)
with get_parallel().override(attn_tp_size=4):
with self.assertRaises(ValueError) as caught:
with get_parallel().override(tp_size=2):
pass
self.assertIn(
"tp_size == attn_tp_size * attn_dp_size * attn_cp_size",
str(caught.exception),
)
with get_parallel().override(tp_size=2, attn_tp_size=2, moe_tp_size=2):
self.assertEqual(get_parallel().tp_size, 2)
self.assertEqual(get_parallel().attn_tp_size, 2)
with get_parallel().override(attn_tp_size=4, tp_size=4, moe_tp_size=4):
self.assertEqual(get_parallel().attn_tp_size, 4)
def test_an_unstated_topology_still_fails(self):
@@ -2127,17 +2239,11 @@ class TestDerivedWidths(_IsolatedOverrides):
)
self.assertEqual(published, recomputed)
def test_initialize_model_parallel_no_longer_touches_the_bag(self):
"""`initialize_model_parallel` used to recompute and
permanently override the six derived widths on `get_parallel()`
after building its groups; that call is gone. Publish a placeholder
config (tp_size defaults to 1), then build real groups at a
different width -- the published leaf must now stay exactly what it
was, because nothing corrects it. This is the behavior a caller
relies on being told about, loudly, the first time it publishes and
builds inconsistently -- see
`test_recomputing_from_published_leaves_matches_the_publish_bag`
for why every real caller must not do that.
def test_initialize_model_parallel_builds_at_the_published_widths(self):
"""The build takes every width from the context rather than from an
argument, so "published one width, built another" is no longer a state
a caller can reach -- there is nothing left to translate, and nothing
to correct afterwards either.
"""
from unittest.mock import Mock
@@ -2145,11 +2251,11 @@ class TestDerivedWidths(_IsolatedOverrides):
reset_context()
self.addCleanup(reset_context)
publish(ServerArgs(model_path="dummy"), role="test")
self.assertEqual(get_parallel().attn_tp_size, 1)
self.assertEqual(get_parallel().moe_ep_size, 1)
world_size = 8
publish(ServerArgs(model_path="dummy", tp_size=world_size), role="test")
self.assertEqual(get_parallel().attn_tp_size, world_size)
built_at = []
with (
patch.object(parallel_state, "_WORLD", None),
patch.object(parallel_state, "_TP", None),
@@ -2168,25 +2274,20 @@ class TestDerivedWidths(_IsolatedOverrides):
patch.object(
parallel_state,
"init_model_parallel_group",
return_value=Mock(device_group=Mock()),
side_effect=lambda group_ranks, *a, **k: (
built_at.append(group_ranks),
Mock(device_group=Mock()),
)[1],
),
patch.object(parallel_state, "get_world_group") as mock_world_group,
):
mock_world_group.return_value = Mock(device_group=Mock(), local_rank=0)
parallel_state.initialize_model_parallel(
tensor_model_parallel_size=world_size,
expert_model_parallel_size=world_size,
)
parallel_state.initialize_model_parallel()
self.addCleanup(parallel_state.destroy_model_parallel)
self.assertEqual(
get_parallel().attn_tp_size,
1,
"initialize_model_parallel must not touch the published leaf -- "
"a caller that needs it corrected must publish a config that "
"already matches the width it is about to build",
)
self.assertEqual(get_parallel().moe_ep_size, 1)
# The first group built is TP, one group spanning the published width.
self.assertEqual(built_at[0], [list(range(world_size))])
self.assertEqual(get_parallel().attn_tp_size, world_size)
class TestTheDerivedHalfIsDeclared(CustomTestCase):
@@ -2294,6 +2395,225 @@ class TestAnEntryThatBuildsARunnerHandsOverItsPlacement(CustomTestCase):
)
class TestTheAccessorsHaveNoCallersOutsideTheirPackage(CustomTestCase):
"""`parallel_state`'s getters are the definition, not a second spelling.
Business code asks `get_parallel()`; a call that goes straight to the getter
is a read the context cannot redirect, which is what a scope needs it to be
able to do. The package that defines them is exempt -- a read there would
go through the context back into itself -- and so is `multimodal_gen`, which
has its own parallel state.
"""
#: Not topology. `get_self_pp_group` builds the single-rank group a draft
#: pipeline scope installs, so there is nothing for the context to answer
#: with until the scope has installed it.
ALLOWED = {
"get_self_pp_group",
"get_default_distributed_backend",
"get_mooncake_transfer_engine",
}
def _accessors(self):
"""Derived from the source, not listed here: a guard whose subject set
is written by hand stops watching whatever gets added next."""
from sglang.srt.distributed import parallel_state as parallel_state_module
source = _pathlib.Path(parallel_state_module.__file__).read_text().splitlines()
return {
line[len("def ") : line.index("(")]
for line in source
if line.startswith("def get_") or line.startswith("def is_")
}
def _callers(self, name):
import re
from sglang.srt.distributed import parallel_state as parallel_state_module
root = _pathlib.Path(parallel_state_module.__file__).parents[2]
pattern = re.compile(rf"(?<![.\w]){re.escape(name)}\(")
hits = []
for path in root.rglob("*.py"):
rel = path.relative_to(root).as_posix()
if rel.startswith(("srt/distributed/", "multimodal_gen/", "test/")):
continue
for number, line in enumerate(path.read_text().splitlines(), 1):
if line.lstrip().startswith(("def ", "#")):
continue
if pattern.search(line):
hits.append(f"{rel}:{number}")
return hits
def test_no_business_code_calls_them(self):
offenders = {}
for name in sorted(self._accessors() - self.ALLOWED):
callers = self._callers(name)
if callers:
offenders[name] = callers
self.assertEqual(
offenders,
{},
"read these through get_parallel() instead, or say here why the "
"context cannot answer them",
)
def test_the_guard_would_notice_a_caller(self):
"""The subject set is derived, so this checks the search finds a real
call rather than that the list happens to be empty: `get_self_pp_group`
is exempt and does have one caller."""
self.assertTrue(self._callers("get_self_pp_group"))
class TestTheTopologyIdentities(CustomTestCase):
"""One set of identities, checked wherever the layout is written.
Each one is injected in the direction that breaks it and in the direction
that keeps it: a guard that only ever fires is as uninformative as one that
never does. They hold unconditionally -- a caller that states one leaf owes
the quotients that follow from it, because a namespace describing no real
layout is what the guard exists to refuse.
"""
def _publish_square(self):
"""tp=4 over two attention-DP replicas of two: every identity holds."""
reset_context()
self.addCleanup(reset_context)
publish(
ServerArgs(
model_path="dummy", tp_size=4, dp_size=2, enable_dp_attention=True
),
role="scheduler",
ranks=SpawnRanks(world_rank=3, dp_rank=1),
)
def test_a_published_topology_is_consistent(self):
"""The quiet direction, and the reason publish can check everything:
it is the one write that establishes the whole layout."""
self._publish_square()
parallel = get_parallel()
self.assertEqual(parallel.tp_size, 4)
self.assertEqual(parallel.attn_dp_size, 2)
self.assertEqual(parallel.attn_tp_size, 2)
self.assertEqual(parallel.tp_rank, 3)
def test_a_width_that_does_not_factor_is_refused(self):
self._publish_square()
with self.assertRaises(ValueError) as caught:
with get_parallel().override(
tp_size=4, attn_tp_size=3, attn_dp_size=1, attn_cp_size=1
):
pass
message = str(caught.exception)
self.assertIn("set by override", message)
self.assertIn("attn_tp_size * attn_dp_size * attn_cp_size", message)
self.assertIn("4 != 3 * 1 * 1", message)
def test_a_rank_at_its_width_is_refused(self):
self._publish_square()
with self.assertRaises(ValueError) as caught:
with get_parallel().override(tp_rank=4, tp_size=4):
pass
self.assertIn("0 <= tp_rank < tp_size", str(caught.exception))
def test_a_moe_width_that_does_not_factor_is_refused(self):
self._publish_square()
with self.assertRaises(ValueError) as caught:
with get_parallel().override(
tp_size=4, moe_ep_size=1, moe_dp_size=1, moe_tp_size=3
):
pass
message = str(caught.exception)
self.assertIn("moe_ep_size * moe_dp_size * moe_tp_size", message)
self.assertIn("4 != 1 * 1 * 3", message)
def test_a_rank_the_attention_layout_cannot_produce_is_refused(self):
"""`tp_rank` is not free of the attention ranks: the layout derives one
from the other, so a set that does not satisfy it places this process
in two different seats at once."""
self._publish_square()
with self.assertRaises(ValueError) as caught:
with get_parallel().override(
tp_rank=0,
attn_dp_rank=1,
attn_cp_rank=0,
attn_tp_rank=0,
attn_cp_size=1,
attn_tp_size=2,
):
pass
self.assertIn("attn_tp_size + attn_tp_rank", str(caught.exception))
def test_the_published_ranks_satisfy_the_layout(self):
"""The quiet direction for the same identity: publish derives the
attention ranks from `tp_rank` through that very equation, so a
published process always sits in one seat."""
self._publish_square()
parallel = get_parallel()
self.assertEqual(
parallel.tp_rank,
(parallel.attn_dp_rank * parallel.attn_cp_size + parallel.attn_cp_rank)
* parallel.attn_tp_size
+ parallel.attn_tp_rank,
)
def test_a_refused_write_leaves_nothing_behind(self):
"""The scope never opened, so the value it tried to state must not be
readable afterwards -- a half-applied override is the state this guard
exists to prevent."""
self._publish_square()
with self.assertRaises(ValueError):
with get_parallel().override(
tp_size=4, attn_tp_size=3, attn_dp_size=1, attn_cp_size=1
):
pass
self.assertEqual(get_parallel().attn_tp_size, 2)
def test_a_group_built_at_another_width_is_refused(self):
"""The other end of the same identity: what the configuration says and
what the coordinators were actually built at, checked where the
disagreement is still attributable to the build."""
from sglang.srt.distributed import parallel_state
from sglang.srt.distributed.parallel_state import GroupCoordinator
self._publish_square()
wrong = GroupCoordinator.__new__(GroupCoordinator)
wrong.world_size = 8
wrong.rank_in_group = 0
with patch.object(parallel_state, "_TP", wrong):
with self.assertRaises(ValueError) as caught:
_validate_parallel(get_parallel(), "group build")
message = str(caught.exception)
self.assertIn("set by group build", message)
self.assertIn("tp_group.world_size == tp_size", message)
self.assertIn("built 8, configured 4", message)
def test_a_group_built_at_the_configured_width_is_quiet(self):
from sglang.srt.distributed import parallel_state
from sglang.srt.distributed.parallel_state import GroupCoordinator
self._publish_square()
right = GroupCoordinator.__new__(GroupCoordinator)
right.world_size = 4
right.rank_in_group = 3
with patch.object(parallel_state, "_TP", right):
_validate_parallel(get_parallel(), "group build")
def test_a_draft_scope_states_a_consistent_topology(self):
"""The scope narrows four names at once, so the identity applies to it
-- and holds, which is what step lets the guard stay on."""
from sglang.srt.distributed import parallel_state
from sglang.srt.distributed.parallel_state import GroupCoordinator
self._publish_square()
group = GroupCoordinator.__new__(GroupCoordinator)
group.world_size = 2
group.rank_in_group = 1
with patch.object(parallel_state, "_TP", group):
with parallel_state.patch_tensor_parallel_group(group, owns_attention=True):
self.assertEqual(get_parallel().attn_tp_size, 2)
class TestWhoAnswersDuringADraftScope(CustomTestCase):
"""A draft worker runs in one process with the target, under a scope.
@@ -2397,29 +2717,36 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase):
worker states it, and a caller that forgets is the bug this catches --
`owns_attention` has no default, but a missing one is a TypeError only
on the path that runs, and these paths need a GPU and a draft model."""
import ast
self.assertEqual(
_scope_entries_that_say_nothing(_sources()),
[],
"these enter the draft scope without saying",
)
package = _pathlib.Path(next(iter(_sglang.__path__))).resolve()
checkout = package.parents[1]
roots = [package] + [
checkout / name for name in ("test",) if (checkout / name).is_dir()
]
missing = []
for path in (q for root in roots for q in root.rglob("*.py")):
try:
tree = ast.parse(path.read_text(encoding="utf-8"))
except (SyntaxError, UnicodeDecodeError):
continue
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
func = node.func
name = getattr(func, "attr", None) or getattr(func, "id", None)
if name not in ("draft_tp_context", "patch_tensor_parallel_group"):
continue
if not any(kw.arg == "owns_attention" for kw in node.keywords):
missing.append(f"{path}:{node.lineno}")
self.assertEqual(missing, [], "these enter the draft scope without saying")
def test_the_census_would_notice_one(self):
"""Two ways for it to report zero and still be wrong: the matcher does
not recognise the call, or the walk never reaches the file. A draft
scope entered from `benchmark/` breaks the same way as one in the
package and no suite covers it, so the roots are part of the check."""
trees = {
part
for path in _sources()
for part in ("benchmark", "examples", "scripts", "test")
if f"/{part}/" in path.as_posix()
}
self.assertEqual(
trees,
{"benchmark", "examples", "scripts", "test"},
"the walk misses a tree that can enter the scope",
)
with tempfile.TemporaryDirectory() as tmp:
probe = _pathlib.Path(tmp) / "probe.py"
probe.write_text(
"with self.draft_tp_context(runner.tp_group):\n pass\n",
encoding="utf-8",
)
self.assertEqual(_scope_entries_that_say_nothing([probe]), [f"{probe}:1"])
def test_a_full_width_swap_leaves_the_attention_layout_alone(self):
"""The other caller. A draft built outside any scope carries the