From aecc187869eea84b47d7a66638b2fe437d865a61 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Tue, 29 Sep 2026 13:49:39 -0700 Subject: [PATCH] Keep only the TP and PP groups on the model runner `ModelRunner.init_torch_distributed` copied fifteen placement values onto the runner. Only two of them have to outlive the scope the runner is built in: speculative workers re-enter a draft's scope through its TP group, and draft forwards read its PP group without the pipeline scope. Keep those two and drop the other thirteen (`attention_tp_group`, `tp_rank`, `tp_size`, `dp_size`, `attn_dp_size`, `pp_rank`, `pp_size`, `attn_cp_rank`, `attn_cp_size`, `attn_dcp_rank`, `attn_dcp_size`, `moe_ep_size`, `dp_rank`). Each reader now takes the value from where it runs: - readers that run while the runner is built, during a forward or a graph capture, or only on the target read `get_parallel()`, which in those scopes is what the runner saw when it copied the value; - readers that can run on a draft runner outside its scope read the runner's own groups: the KV cache configurator (built while pools are allocated) and the PP proxy hidden size (also reached when graphs are recaptured) use `pp_group.world_size` / `rank_in_group`, and the DFLASH / DSpark rank-0 log gates inside the draft scope read the target runner's `tp_group.rank_in_group`; - the HiSparse coordinator, built with the pools, reads `get_parallel().attn_tp_group`. Outside the scope that is the target's attention-TP group, which is the group an attention-owning draft runs on; a draft that does not own attention never narrows it. The runner-owned `tp_size` / `tp_rank` arguments to `build_load_config`, `dist_barrier_after_load`, the tensor-dump hook, `apply_torch_tp` and the quantized-MoE check now come from the context too. All of them run in the runner's constructor, except the barrier after an overlapped startup load, which the scheduler runs on the target's runners outside any draft scope. No value a reader sees changes. Tests that stood in for the retired fields follow: the MLX runner stub tests publish their attention-DP width, the DFLASH sampler and phase-1 tests give the fake a TP group and publish a topology, fake runners in the attention test kits and backend tests drop placement fields nothing reads, and the placement-freeze census asserts it finds the TP and PP groups. --- .../hardware_backend/mlx/model_runner_stub.py | 14 +- .../npu/attention/ascend_backend.py | 2 +- .../attention/flashattention_backend.py | 2 +- .../srt/layers/attention/hpc_ops_backend.py | 9 +- .../srt/layers/attention/wave_backend.py | 2 +- .../srt/layers/attention/xpu_backend.py | 3 +- .../layers/deep_gemm_wrapper/compile_utils.py | 2 +- python/sglang/srt/managers/tp_worker.py | 12 +- .../srt/model_executor/forward_batch_info.py | 9 +- .../sglang/srt/model_executor/model_runner.py | 122 +++++++++--------- .../srt/model_executor/runner/base_runner.py | 4 +- .../srt/model_executor/runner/eager_runner.py | 3 +- .../runner/flashinfer_autotune.py | 8 +- .../srt/speculative/dflash_worker_v2.py | 32 ++--- .../dspark_components/dspark_worker_v2.py | 14 +- python/sglang/srt/speculative/uno_lora.py | 6 +- .../attention_methods/dense_attention.py | 4 - .../attention_methods/dsa_attention.py | 4 - .../attention_methods/dsv4_attention.py | 4 - .../attention_methods/gdn_attention.py | 1 - .../attention_methods/kda_attention.py | 1 - .../attention_methods/lightning_attention.py | 1 - .../attention_methods/mamba2_attention.py | 1 - .../attention_methods/mla_attention.py | 4 - .../attention/test_flashattn_backend.py | 1 - .../mlx/test_attn_dp_request_capacity.py | 7 +- .../mlx/test_max_running_requests.py | 8 +- ...flashattention_pa_swa_prefill_lens_size.py | 3 - .../test_flashattention_return_lse.py | 2 - .../test_mamba_prefill_track_metadata.py | 1 - .../test_startup_weight_load.py | 14 +- .../unit/spec/test_dflash_logits.py | 18 ++- test/registered/unit/test_runtime_context.py | 4 +- 33 files changed, 163 insertions(+), 159 deletions(-) diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py index 185e98004e84..d5df916479e7 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py @@ -19,7 +19,13 @@ from sglang.srt.model_executor.model_runner_components.layer_setup import ( ModelLayerInfo, ) -from sglang.srt.runtime_context import get_exec, get_memory, get_model, get_schedule +from sglang.srt.runtime_context import ( + get_exec, + get_memory, + get_model, + get_parallel, + get_schedule, +) logger = logging.getLogger(__name__) @@ -179,7 +185,7 @@ def _explicit_aux_state_size_per_worker(self) -> int | None: aux_state_size = get_schedule().max_mamba_cache_size if aux_state_size is None: return None - return aux_state_size // self.attn_dp_size + return aux_state_size // get_parallel().attn_dp_size def _resolve_max_running_requests(self) -> int: """Concurrency cap handed to the scheduler. @@ -204,7 +210,7 @@ def _resolve_max_running_requests(self) -> int: requested_per_worker = None resolved = min(capacity_cap, 4096) else: - requested_per_worker = requested // self.attn_dp_size + requested_per_worker = requested // get_parallel().attn_dp_size resolved = min(requested_per_worker, capacity_cap) aux_state_size = self._explicit_aux_state_size_per_worker() @@ -216,7 +222,7 @@ def _resolve_max_running_requests(self) -> int: resolved = min(resolved, aux_state_size // ratio) if resolved <= 0: global_aux_state_size = get_schedule().max_mamba_cache_size - min_global_aux_state_size = ratio * self.attn_dp_size + min_global_aux_state_size = ratio * get_parallel().attn_dp_size raise RuntimeError( f"MLX auxiliary-state cache is too small to serve any " f"requests: max_mamba_cache_size={global_aux_state_size} " diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index e0dcf0008ca3..dd13e2164151 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -446,7 +446,7 @@ def __init__(self, model_runner: ModelRunner, speculative_step_id: int = 0): self.is_dllm_model = True self.dllm_block_size = self.dllm_config.block_size - self.attn_cp_size = model_runner.attn_cp_size + self.attn_cp_size = get_parallel().attn_cp_size def _is_swa_layer(self, layer: RadixAttention) -> bool: return ( diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 867bcf9367e8..64562b9de0a1 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -204,7 +204,7 @@ def __init__( self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA self.kv_index_translator = model_runner.kv_index_translator self.skip_prefill = skip_prefill - self.attn_cp_size = model_runner.attn_cp_size + self.attn_cp_size = get_parallel().attn_cp_size self._verify_mask = None # The worker fetches the tree-mask scratch from the target backend # only; draft-side instances must not allocate it. diff --git a/python/sglang/srt/layers/attention/hpc_ops_backend.py b/python/sglang/srt/layers/attention/hpc_ops_backend.py index 22f40a8073f8..bf6aa687ce6f 100644 --- a/python/sglang/srt/layers/attention/hpc_ops_backend.py +++ b/python/sglang/srt/layers/attention/hpc_ops_backend.py @@ -52,6 +52,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( get_tc_piecewise_forward_context, ) +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: @@ -162,8 +163,8 @@ def __init__(self, model_runner: ModelRunner): self.use_fp8 = model_runner.kv_cache_dtype == torch.float8_e4m3fn if self.use_fp8: heads = ( - model_runner.model_config.num_attention_heads // model_runner.tp_size, - model_runner.model_config.get_num_kv_heads(model_runner.tp_size), + model_runner.model_config.num_attention_heads // get_parallel().tp_size, + model_runner.model_config.get_num_kv_heads(get_parallel().tp_size), ) if heads not in FP8_ROPE_SUPPORTED_HEAD_CONFIGS: raise ValueError( @@ -176,8 +177,8 @@ def __init__(self, model_runner: ModelRunner): config = model_runner.model_config head_dim = config.head_dim - num_q_heads = config.num_attention_heads // model_runner.tp_size - num_kv_heads = config.get_num_kv_heads(model_runner.tp_size) + num_q_heads = config.num_attention_heads // get_parallel().tp_size + num_kv_heads = config.get_num_kv_heads(get_parallel().tp_size) gqa_group_size = num_q_heads // num_kv_heads if head_dim != _SUPPORTED_HEAD_DIM or gqa_group_size not in ( _SUPPORTED_GQA_GROUP_SIZES diff --git a/python/sglang/srt/layers/attention/wave_backend.py b/python/sglang/srt/layers/attention/wave_backend.py index 03ac9a4d36cf..1a5fb755aa6d 100644 --- a/python/sglang/srt/layers/attention/wave_backend.py +++ b/python/sglang/srt/layers/attention/wave_backend.py @@ -66,7 +66,7 @@ def __init__( import wave_lang.kernel.wave.cache as cache base_cache_dir = cache.CACHE_BASE_DIR - new_dir = base_cache_dir / f"worker_{model_runner.tp_rank}" + new_dir = base_cache_dir / f"worker_{get_parallel().tp_rank}" logger.info(f"Setting Wave cache dir: {new_dir}") cache.CACHE_BASE_DIR = new_dir diff --git a/python/sglang/srt/layers/attention/xpu_backend.py b/python/sglang/srt/layers/attention/xpu_backend.py index 2fdcb73d2166..83e456c674b7 100644 --- a/python/sglang/srt/layers/attention/xpu_backend.py +++ b/python/sglang/srt/layers/attention/xpu_backend.py @@ -18,6 +18,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import ( get_exec, + get_parallel, get_schedule, get_spec, ) @@ -67,7 +68,7 @@ def __init__( self.num_attention_heads = ( model_runner.model_config.hf_text_config.num_attention_heads ) - self.tp_size = model_runner.tp_size + self.tp_size = get_parallel().tp_size assert self.num_attention_heads % self.tp_size == 0 self.num_local_heads = self.num_attention_heads // self.tp_size self.device = model_runner.device diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index b1f282d90ab9..2db04ae05d98 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -510,7 +510,7 @@ def _align(bs: int) -> int: "PP-parallel DeepGEMM warmup start " "(pp_rank=%d, tp_rank=%d, batch_sizes=%s, disagg=%s).", get_parallel().pp_rank, - model_runner.tp_rank, + get_parallel().tp_rank, batch_sizes, disagg_mode, ) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 6adfbe043792..b11ee1f9958a 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -195,7 +195,7 @@ def deserialize_own_rank(self, serialized_named_tensors): refcounting.""" monkey_patch_torch_reductions() return MultiprocessingSerializer.deserialize( - serialized_named_tensors[self.model_runner.tp_rank] + serialized_named_tensors[get_parallel().tp_rank] ) def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput): @@ -250,12 +250,12 @@ def load_lora_adapter_from_tensors( extra = [n for n in tensors if n not in exp] if mismatch or missing or extra: raise RuntimeError( - f"[LORA-CHECK] rank{self.model_runner.tp_rank} adapter sync MISMATCH of {len(exp)} expected: " + f"[LORA-CHECK] rank{get_parallel().tp_rank} adapter sync MISMATCH of {len(exp)} expected: " f"{len(mismatch)} value-diff {mismatch[:5]}, {len(missing)} missing {missing[:5]}, " f"{len(extra)} extra {extra[:5]}" ) logger.info( - f"[LORA-CHECK] rank{self.model_runner.tp_rank} adapter sync OK: {len(exp)}/{len(exp)} tensors match (sha256)" + f"[LORA-CHECK] rank{get_parallel().tp_rank} adapter sync OK: {len(exp)}/{len(exp)} tensors match (sha256)" ) result = self.model_runner.load_lora_adapter_from_tensors( recv_req.to_ref(), @@ -369,15 +369,15 @@ def __init__( tp_group = self.model_runner.tp_group self.random_seed = broadcast_pyobj( [get_device().random_seed], - tp_group.ranks[self.model_runner.tp_rank], + tp_group.ranks[tp_group.rank_in_group], tp_group.cpu_group, src=tp_group.ranks[0], )[0] else: self.random_seed = broadcast_pyobj( [get_device().random_seed], - self.model_runner.tp_size * get_parallel().pp_rank - + self.model_runner.tp_rank, + get_parallel().tp_size * get_parallel().pp_rank + + get_parallel().tp_rank, self.world_group.cpu_group, src=self.world_group.ranks[0], )[0] diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index cc6af6309d33..6733618299ed 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -1195,13 +1195,10 @@ def init_new( model_runner.lora_manager.prepare_lora_batch(ret) - if ( - model_runner.attn_dcp_size > 1 - and ret.out_cache_loc is not None - and is_hip() - ): + parallel = get_parallel() + if parallel.attn_dcp_size > 1 and ret.out_cache_loc is not None and is_hip(): ret.dcp_kv_mask = ( - ret.positions % model_runner.attn_dcp_size == model_runner.attn_dcp_rank + ret.positions % parallel.attn_dcp_size == parallel.attn_dcp_rank ) return ret diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index bc1d618af480..1fc285c6974c 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -495,18 +495,19 @@ def _initialize_elastic_ep_joiner(self) -> None: ): return - join_effective_ep_size = get_parallel().ep_join_rank_offset + self.tp_size + parallel = get_parallel() + join_effective_ep_size = parallel.ep_join_rank_offset + parallel.tp_size dist.barrier(group=self.tp_group.cpu_group) - if self.tp_rank == 0: + if parallel.tp_rank == 0: register_scale_cohort( - get_parallel().ep_join_rank_offset, + parallel.ep_join_rank_offset, join_effective_ep_size, self._elastic_cuda_graph_enabled(), ) join_scale_process_group() get_context().override("elastic_ep.scale_join", ep_size=join_effective_ep_size) - global_ep_rank = self.tp_rank + get_parallel().ep_join_rank_offset + global_ep_rank = parallel.tp_rank + parallel.ep_join_rank_offset broadcast_global_expert_location_metadata( model_config=self.model_config, moe_ep_rank=global_ep_rank, @@ -549,7 +550,7 @@ def init_msprobe(self): def init_weight_updater(self): self.weight_updater = WeightUpdater( - tp_rank=self.tp_rank, + tp_rank=get_parallel().tp_rank, device=self.device, gpu_id=self.gpu_id, model_config=self.model_config, @@ -572,8 +573,8 @@ def init_spec_aux_hidden_state(self): def init_weight_exporter(self): self.weight_exporter = WeightExporter( - tp_rank=self.tp_rank, - tp_size=self.tp_size, + tp_rank=get_parallel().tp_rank, + tp_size=get_parallel().tp_size, gpu_id=self.gpu_id, get_model_path=lambda: self.model_config.model_path, get_model=lambda: self.model, @@ -596,12 +597,13 @@ def init_ngram_embedding_manager(self): def init_kv_cache_configurator(self): # The replica count that shares this KV budget is the deployment's, not - # the one a draft scope reports; the pool is allocated outside any scope. + # the one a draft scope reports; the pool is allocated outside any scope, + # so the PP width comes from the runner's own group. self.kv_cache_configurator = KVCacheConfigurator( device=self.device, gpu_id=self.gpu_id, attn_dp_size=get_parallel().attn_dp_size, - pp_size=self.pp_size, + pp_size=self.pp_group.world_size, pp_group=self.pp_group, model=self.model, model_config=self.model_config, @@ -634,9 +636,10 @@ def init_mindspore_runner(self): if get_model().model_impl.lower() == ModelImpl.MINDSPORE and _is_npu: from sglang.srt.model_executor.mindspore_runner import init_ms_distributed + parallel = get_parallel() init_ms_distributed( - world_size=self.tp_size * get_parallel().pp_size, - rank=self.tp_size * get_parallel().pp_rank + self.tp_rank, + world_size=parallel.tp_size * parallel.pp_size, + rank=parallel.tp_size * parallel.pp_rank + parallel.tp_rank, local_rank=self.gpu_id, port=self.dist_port, ) @@ -696,7 +699,10 @@ def maybe_init_expert_location_metadata(self): moe_ep_rank=expert_rank, ) ) - if self.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get(): + if ( + get_parallel().tp_rank == 0 + and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get() + ): logger.info( "Initial expert_location_metadata:\n%s", format_expert_location_layout(get_global_expert_location_metadata()), @@ -754,7 +760,7 @@ def maybe_init_expert_backup_client(self): def maybe_apply_post_load_model_transforms(self): supports_torch_tp = getattr(self.model, "supports_torch_tp", False) - if self.tp_size > 1 and supports_torch_tp: + if get_parallel().tp_size > 1 and supports_torch_tp: self.apply_torch_tp() def maybe_init_lora_manager(self): @@ -775,10 +781,12 @@ def maybe_enable_batch_invariant_mode(self): enable_batch_invariant_mode() def get_pp_proxy_dspark_hidden_size(self) -> int: + # Graph recapture reaches this outside the draft PP scope, so read the + # runner's own PP group. return misc_utils.resolve_pp_proxy_dspark_hidden_size( model=self.model, - pp_size=self.pp_size, - pp_rank=self.pp_rank, + pp_size=self.pp_group.world_size, + pp_rank=self.pp_group.rank_in_group, ) def get_aux_hidden_states_width(self) -> int: @@ -956,7 +964,7 @@ def maybe_init_hisparse_coordinator(self): device_buffer_size=hisparse_cfg.device_buffer_size, device=self.device, tp_group=( - self.attention_tp_group.cpu_group + get_parallel().attn_tp_group.cpu_group if get_parallel().enable_dp_attention else self.tp_group.cpu_group ), @@ -992,7 +1000,7 @@ def post_capture_resize_kv_pool(self, *, draft_runners=()): def post_capture_elastic_ep_recover(self): join_process_groups() - global_ep_rank = self.tp_rank + get_parallel().ep_join_rank_offset + global_ep_rank = get_parallel().tp_rank + get_parallel().ep_join_rank_offset broadcast_global_expert_location_metadata( model_config=self.model_config, moe_ep_rank=global_ep_rank, @@ -1124,7 +1132,7 @@ def init_cuda_graphs(self, capture_decode_cuda_graph: bool = True): self._recapture_elastic_cuda_graphs() finalize_cuda_graph_capture(self) # Scheduler startup calls this path even when CUDA graphs are disabled. - target_size = get_parallel().ep_join_rank_offset + self.tp_size + target_size = get_parallel().ep_join_rank_offset + get_parallel().tp_size self._finalize_elastic_ep_joiner(target_size) def init_routed_experts_capturer(self): @@ -1161,7 +1169,7 @@ def init_indexer_capturer(self): def check_quantized_moe_compatibility(self): check_quantized_moe_compatibility( model_config=self.model_config, - tp_size=self.tp_size, + tp_size=get_parallel().tp_size, moe_ep_size=get_parallel().moe_ep_size, moe_dp_size=get_parallel().moe_dp_size, ) @@ -1170,23 +1178,13 @@ def init_torch_distributed(self): self.pre_model_load_memory = bootstrap.measure_pre_model_load_memory( device=self.device, is_draft_worker=self.is_draft_worker ) - # Capture draft placement at construction; the runner outlives the scope. + # A draft runner is used outside the scope it is built in: speculative + # workers re-enter that scope through its TP group, and draft forwards + # read its PP group without the pipeline scope. Keep both groups; read + # every other placement value from the context where it is used. parallel = get_parallel() self.tp_group = parallel.tp_group self.pp_group = parallel.pp_group - self.attention_tp_group = parallel.attn_tp_group - self.tp_rank = parallel.tp_rank - self.tp_size = parallel.tp_size - self.dp_size = parallel.dp_size - self.attn_dp_size = parallel.attn_dp_size - self.pp_rank = parallel.pp_rank - self.pp_size = parallel.pp_size - self.attn_cp_rank = parallel.attn_cp_rank - self.attn_cp_size = parallel.attn_cp_size - self.attn_dcp_rank = parallel.attn_dcp_rank - self.attn_dcp_size = parallel.attn_dcp_size - self.moe_ep_size = parallel.moe_ep_size - self.dp_rank = parallel.dp_rank def load_model(self): tic_total = time.perf_counter() @@ -1207,7 +1205,7 @@ def load_model(self): self.load_config = build_load_config( server_args=self.server_args, load_format=draft_load_format, - tp_rank=self.tp_rank, + tp_rank=get_parallel().tp_rank, remote_instance_weight_transporter_engine=self.remote_instance_weight_transporter.engine, remote_instance_weight_transporter_session_id=self.remote_instance_weight_transporter.session_id, draft_model_idx=self.draft_model_idx, @@ -1220,11 +1218,11 @@ def load_model(self): ) if self.device == "cpu": self.model_config = adjust_config_with_unaligned_cpu_tp( - self.model_config, self.load_config, self.tp_size + self.model_config, self.load_config, get_parallel().tp_size ) maybe_trigger_remote_instance_nccl_send_group( - tp_rank=self.tp_rank, + tp_rank=get_parallel().tp_rank, load_format=draft_load_format, ) @@ -1300,8 +1298,8 @@ def load_model(self): model=self.model, spec_algorithm=self.spec_algorithm, is_draft_worker=self.is_draft_worker, - tp_size=self.tp_size, - tp_rank=self.tp_rank, + tp_size=get_parallel().tp_size, + tp_rank=get_parallel().tp_rank, pp_rank=get_parallel().pp_rank, ) @@ -1319,7 +1317,7 @@ def load_model(self): if self.startup_weight_load is None: dist_barrier_after_load( elastic_ep_backend=get_exec().moe.elastic_ep_backend, - tp_rank=self.tp_rank, + tp_rank=get_parallel().tp_rank, is_ep_joiner=get_exec().moe.is_ep_joiner, ) @@ -1342,7 +1340,7 @@ def finalize_startup_weight_load(self) -> None: self.startup_weight_load.finalize() dist_barrier_after_load( elastic_ep_backend=get_exec().moe.elastic_ep_backend, - tp_rank=self.tp_rank, + tp_rank=get_parallel().tp_rank, is_ep_joiner=get_exec().moe.is_ep_joiner, ) self.startup_weight_load = None @@ -1370,8 +1368,8 @@ def init_lora_manager(self): dtype=self.dtype, server_args=self.server_args, lora_backend=get_lora().lora_backend, - tp_size=self.tp_size, - tp_rank=self.tp_rank, + tp_size=get_parallel().tp_size, + tp_rank=get_parallel().tp_rank, max_lora_rank=get_lora().max_lora_rank, target_modules=get_lora().lora_target_modules, lora_paths=get_lora().lora_paths, @@ -1599,7 +1597,7 @@ def init_prefill_cuda_graph(self, force_for_draft_worker: bool = False): def apply_torch_tp(self): model_parallel.apply_torch_tp( - model=self.model, device=self.device, tp_size=self.tp_size + model=self.model, device=self.device, tp_size=get_parallel().tp_size ) def update_decode_attn_backend(self, stream_idx: int): @@ -1725,11 +1723,7 @@ def forward( # Try msprob debugger if self.msprobe_debugger is not None: - rank_id = ( - self.gpu_id - if self.attn_dp_size is not None and self.attn_dp_size > 1 - else None - ) + rank_id = self.gpu_id if get_parallel().attn_dp_size > 1 else None self.msprobe_debugger.start(model=self.model, rank_id=rank_id) # Step span @@ -2123,7 +2117,7 @@ def _expand_eplb_metadata_for_scale( set_global_expert_location_metadata(new_metadata, allow_overwrite=True) def _elastic_global_rank(self) -> int: - return self.tp_rank + get_parallel().ep_join_rank_offset + return get_parallel().tp_rank + get_parallel().ep_join_rank_offset def _rearm_eplb_after_elastic_scale(self) -> None: if self.eplb_manager is None: @@ -2145,7 +2139,7 @@ def _reset_eplb_after_elastic_scale_failure(self) -> None: self._rearm_eplb_after_elastic_scale() def _report_elastic_scale_failure(self, error: str, effective_size: int) -> None: - if self.tp_rank != 0 or get_exec().moe.is_ep_scale_joiner: + if get_parallel().tp_rank != 0 or get_exec().moe.is_ep_scale_joiner: return from sglang.srt.managers.io_struct import ElasticScaleUpdateReq @@ -2156,7 +2150,7 @@ def _report_elastic_scale_failure(self, error: str, effective_size: int) -> None ) def _elastic_scale_ready_barrier(self, target_size: int, log_tag: str) -> None: - if self.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.debug( "[Elastic EP][scale] %s entering post-scale WORLD barrier " "(target_ep_size=%d)", @@ -2164,7 +2158,7 @@ def _elastic_scale_ready_barrier(self, target_size: int, log_tag: str) -> None: target_size, ) dist.barrier(group=dist.group.WORLD) - if self.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.debug( "[Elastic EP][scale] %s passed post-scale WORLD barrier " "(target_ep_size=%d)", @@ -2239,7 +2233,7 @@ def _finalize_scale_up( ElasticEPStateManager.commit_scale() self._rearm_eplb_after_elastic_scale() - if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: + if get_parallel().tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: from sglang.srt.managers.io_struct import ElasticScaleUpdateReq self._pending_elastic_scale_update = ElasticScaleUpdateReq( @@ -2315,7 +2309,10 @@ def maybe_join_ep_ranks(self) -> None: ) ElasticEPStateManager.fail_recovery(error) self._report_elastic_scale_failure(error, effective_size) - if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: + if ( + get_parallel().tp_rank == 0 + and not get_exec().moe.is_ep_scale_joiner + ): logger.error("[Elastic EP] %s", error) return @@ -2341,7 +2338,7 @@ def maybe_join_ep_ranks(self) -> None: ElasticEPStateManager.fail_scale(error) self._reset_eplb_after_elastic_scale_failure() self._report_elastic_scale_failure(error, effective_size) - if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: + if get_parallel().tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: logger.error("[Elastic EP] %s", error) return @@ -2357,7 +2354,10 @@ def maybe_join_ep_ranks(self) -> None: ElasticEPStateManager.fail_scale(error) self._reset_eplb_after_elastic_scale_failure() self._report_elastic_scale_failure(error, effective_size) - if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: + if ( + get_parallel().tp_rank == 0 + and not get_exec().moe.is_ep_scale_joiner + ): logger.error("[Elastic EP] %s", error) return cuda_graph_enabled = self._elastic_cuda_graph_enabled() @@ -2369,7 +2369,10 @@ def maybe_join_ep_ranks(self) -> None: ElasticEPStateManager.fail_scale(error) self._reset_eplb_after_elastic_scale_failure() self._report_elastic_scale_failure(error, effective_size) - if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: + if ( + get_parallel().tp_rank == 0 + and not get_exec().moe.is_ep_scale_joiner + ): logger.error("[Elastic EP] %s", error) return try: @@ -2379,7 +2382,10 @@ def maybe_join_ep_ranks(self) -> None: ElasticEPStateManager.fail_scale(error) self._reset_eplb_after_elastic_scale_failure() self._report_elastic_scale_failure(error, effective_size) - if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner: + if ( + get_parallel().tp_rank == 0 + and not get_exec().moe.is_ep_scale_joiner + ): logger.error("[Elastic EP] %s", error) return if not ElasticEPStateManager.begin_scale(): diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 51fc3ac930a2..03fcf0917030 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -599,9 +599,9 @@ def _dummy_run( if ( capture_forward_mode == ForwardMode.EXTEND and get_parallel().pp_rank != 0 - and mr.attn_cp_size > 1 + and get_parallel().attn_cp_size > 1 ): - pp_hidden_tokens = num_tokens // mr.attn_cp_size + pp_hidden_tokens = num_tokens // get_parallel().attn_cp_size pp_proxy_tensors = PPProxyTensors( {k: v[:pp_hidden_tokens] for k, v in buffers.pp_proxy_tensors.items()} ) diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index acb438eb5770..2323061dbdca 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -59,6 +59,7 @@ ) from sglang.srt.runtime_context import ( get_exec, + get_parallel, get_spec, max_prefill_buffer_tokens, max_speculative_num_draft_tokens, @@ -295,7 +296,7 @@ def _execute_extend( or cp_active or forward_batch.forward_mode.is_target_verify() ): - if model_runner.attn_dcp_size > 1 and hasattr( + if get_parallel().attn_dcp_size > 1 and hasattr( model_runner.model, "prepare_context_parallel_metadata_for_dcp" ): # prepare kv cache buffer for dcp to gather kv cache diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 6ef4acbb32b0..70e4fdf4cb32 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -147,10 +147,10 @@ def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path: str(mr.dtype), str(get_model().quantization), str(get_exec().moe.moe_runner_backend), - str(mr.tp_size), + str(get_parallel().tp_size), str(get_parallel().pp_size), - str(mr.attn_dp_size), - str(mr.moe_ep_size), + str(get_parallel().attn_dp_size), + str(get_parallel().moe_ep_size), str(mr.model_config.hf_config.__class__.__name__), ] # A different skip policy must not reuse previously tuned tactics. @@ -171,7 +171,7 @@ def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path: cache_dir.mkdir(parents=True, exist_ok=True) return ( cache_dir - / f"rank_tp{mr.tp_rank}_pp{get_parallel().pp_rank}_dp{mr.dp_rank or 0}.json" + / f"rank_tp{get_parallel().tp_rank}_pp{get_parallel().pp_rank}_dp{get_parallel().dp_rank or 0}.json" ) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index f6d34c02bb4c..7604f0207ce3 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -464,7 +464,7 @@ def __init__( validate_domino_runtime( device=torch.device(self.device), tp_size=int(get_parallel().tp_group.world_size), - tp_rank=int(self.model_runner.tp_rank), + tp_rank=int(get_parallel().tp_rank), target_vocab_size=int(self.model_runner.model_config.vocab_size), draft_vocab_size=int(self.draft_model_runner.model_config.vocab_size), hidden_size=int(self.draft_model.config.hidden_size), @@ -511,7 +511,7 @@ def __init__( ) self._maybe_merge_trained_mask_embedding() self._cache_full_embed_weight() - if self.model_runner.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.info( "Initialized DFLASH draft runner. attention_backend=%s, model=%s, block_size=%s, draft_window_size=%s, compact_cache=%s", bundle.resolved_attention_backend, @@ -661,7 +661,7 @@ def init_cuda_graphs(self): # shared graph capture/replay; keep the draft eager under dp # attention. capture_decode_cuda_graph = False - if self.model_runner.tp_rank == 0: + if self.model_runner.tp_group.rank_in_group == 0: logger.warning( "Disable DFLASH draft cuda graph because dp attention " "is enabled (draft runs eager)." @@ -806,7 +806,7 @@ def prewarm_sampling(self) -> SamplingPrewarmResult: def _maybe_build_draft_sampler(self): def _eager(reason): - if self.model_runner.tp_rank == 0: + if self.model_runner.tp_group.rank_in_group == 0: logger.info("DFLASH draft greedy head kept eager (reason=%s).", reason) return None @@ -832,7 +832,7 @@ def _eager(reason): if not head_supported: return _eager("unsupported quantized lm_head") self.draft_model.lm_head = lm_head - if self.model_runner.tp_rank == 0: + if self.model_runner.tp_group.rank_in_group == 0: logger.info( "DFLASH selector decode folded into the draft cuda graph " "(sampling_enabled=%s).", @@ -867,7 +867,7 @@ def _eager(reason): embed_proj = self.draft_model.embed_proj if prefix_gru is None or embed_proj is None: return _eager("Domino projector modules are unavailable") - if self.model_runner.tp_rank == 0: + if self.model_runner.tp_group.rank_in_group == 0: logger.info( "DFLASH Domino rollout folded into the draft cuda graph (tp=%s).", int(tp_group.world_size), @@ -906,7 +906,7 @@ def _eager(reason): return _eager("added vocab") num_org = int(shard.num_org_elements) org_vocab_start = int(shard.org_vocab_start_index) - if self.model_runner.tp_rank == 0: + if self.model_runner.tp_group.rank_in_group == 0: logger.info( "DFLASH draft greedy head folded into the draft cuda graph (tp=%d).", tp_group.world_size, @@ -932,7 +932,7 @@ def _init_fused_kv_helper(self) -> None: fused_disable_reason = "draft model does not support fused context KV" if fused_disable_reason is not None: - if self.model_runner.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.info( "DFLASH fused KV materialization disabled: %s", fused_disable_reason, @@ -975,7 +975,7 @@ def _init_fused_kv_helper(self) -> None: break if fused_disable_reason is not None: - if self.model_runner.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.info( "DFLASH fused KV materialization disabled: %s", fused_disable_reason, @@ -997,7 +997,7 @@ def _init_fused_kv_helper(self) -> None: max_position_hint=self.target_worker.model_runner.model_config.context_len + int(self.block_size), ) - if self.model_runner.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.info( "DFLASH fused KV materialization enabled. " "n_layers=%d, num_kv_heads=%d, head_dim=%d", @@ -1290,7 +1290,7 @@ def _maybe_merge_trained_mask_embedding(self) -> None: embedding_tensor.to(embed_module.weight.dtype) ) - if self.model_runner.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.info( "Merged trained mask embedding into target model " "(mask_token_id=%s, source=%s)", @@ -1325,7 +1325,7 @@ def _cache_full_embed_weight(self) -> None: parts = [torch.empty_like(shard_t) for _ in range(tp_size)] dist.all_gather(parts, shard_t, group=tp_group.device_group) self._full_embed_gpu = torch.cat(parts, dim=0)[:vocab_size] - if self.model_runner.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.info( "DFLASH cached full embed on GPU for dp attention: shape=%s", list(self._full_embed_gpu.shape), @@ -1388,7 +1388,7 @@ def _resolve_mask_token_id( if resolved_id is None: resolved_id = tokenizer.convert_tokens_to_ids(mask_token) - if added and self.model_runner.tp_rank == 0: + if added and get_parallel().tp_rank == 0: logger.info( "Added DFLASH mask token to tokenizer. token=%s, mask_token_id=%s, tokenizer_len=%s, model_vocab_size=%s", mask_token, @@ -2211,7 +2211,7 @@ def _validate_phase1_sampling_support(self, batch: ScheduleBatch) -> None: if self.selector is not None: if self._selector_sampling_enabled: return - if not self._warned_sampling_fallback and self.model_runner.tp_rank == 0: + if not self._warned_sampling_fallback and get_parallel().tp_rank == 0: logger.warning( "DFLASH non-greedy verification is unavailable on this " "build/device; falling back to greedy argmax verification. " @@ -2228,7 +2228,7 @@ def _validate_phase1_sampling_support(self, batch: ScheduleBatch) -> None: if ( not is_dflash_sampling_verify_available() and not self._warned_sampling_fallback - and self.model_runner.tp_rank == 0 + and get_parallel().tp_rank == 0 ): logger.warning( "DFLASH non-greedy verification is unavailable on this build/device; " @@ -2653,7 +2653,7 @@ def forward_batch_generation( if ( self._lilicorr_sampling_enabled and not self._warned_lilicorr_eager - and self.model_runner.tp_rank == 0 + and get_parallel().tp_rank == 0 ): logger.warning( "LiLiCorr sampled draft ran the eager head on a decode step " diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 4ba8f59e132a..1751d53c80d7 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -238,7 +238,7 @@ def __init__( else parallel.tp_group ) - if self.model_runner.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.info( "Initialized DSpark draft runner. attention_backend=%s, model=%s, " "gamma=%s, verify_num_draft_tokens=%s, query_token_num=%s, " @@ -261,7 +261,7 @@ def __init__( ) if getattr(self.draft_model, "uses_own_vocab_modules", False): - if self.model_runner.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.info( "DSpark draft uses its checkpoint-local embedding and LM head." ) @@ -285,7 +285,7 @@ def __init__( gamma=self.gamma, model_runner=self.model_runner, device=self.device, - tp_rank=self.model_runner.tp_rank, + tp_rank=get_parallel().tp_rank, verify_num_draft_tokens=self.verify_num_draft_tokens, tp_sync=self._tp_sync, ) @@ -333,7 +333,7 @@ def __init__( and self._verify_planner.mode_value == "static" and self._draft_is_moe and not get_parallel().enable_dp_attention - and self.model_runner.pp_size == 1 + and get_parallel().pp_size == 1 ) # ROCm (V4.1 target only): inside a HIP graph the accept-site TP broadcasts # need the group's pynccl communicator @@ -399,7 +399,7 @@ def __init__( planner=self._verify_planner, gamma=self.gamma, verify_num_draft_tokens=self.verify_num_draft_tokens, - tp_rank=self.model_runner.tp_rank, + tp_rank=get_parallel().tp_rank, device=self.device, simulate_acc_len=self._simulate_acc_len, ) @@ -459,7 +459,7 @@ def init_attention_backends(self): draft_model=self.draft_model, is_deepseek_v4_draft=self._draft_is_moe, ) - if self._target_hidden_projection_enabled and self.model_runner.tp_rank == 0: + if self._target_hidden_projection_enabled and get_parallel().tp_rank == 0: logger.info( "DSpark prefill target-hidden projection runs before " "sequence-parallel gather." @@ -516,7 +516,7 @@ def _maybe_build_draft_sampler(self, *, available_memory_gb: float): gamma=self.gamma, max_bs=max(get_exec().graph.cuda_graph_config.decode.bs), device=self.device, - tp_rank=self.model_runner.tp_rank, + tp_rank=self.model_runner.tp_group.rank_in_group, tp_sync=self._tp_sync, available_memory_gb=available_memory_gb, confidence_fn=( diff --git a/python/sglang/srt/speculative/uno_lora.py b/python/sglang/srt/speculative/uno_lora.py index 397c9c5506bc..132c53186c60 100644 --- a/python/sglang/srt/speculative/uno_lora.py +++ b/python/sglang/srt/speculative/uno_lora.py @@ -6,7 +6,7 @@ from sglang.srt.lora.lora_manager import LoRAManager from sglang.srt.lora.lora_registry import LoRARef -from sglang.srt.runtime_context import get_spec +from sglang.srt.runtime_context import get_parallel, get_spec if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner @@ -45,8 +45,8 @@ def init_uno_lora_manager( dtype=model_runner.dtype, server_args=model_runner.server_args, lora_backend="uno_cublas", # fast path - tp_size=model_runner.tp_size, - tp_rank=model_runner.tp_rank, + tp_size=get_parallel().tp_size, + tp_rank=get_parallel().tp_rank, # Infer these from the one trained adapter. max_lora_rank=None, target_modules=None, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index 0bff8187972f..e08dbb0d1adb 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py @@ -331,9 +331,6 @@ def __init__( self.canary_manager = None self.page_size = case.page_size self.model_config = model_config - self.tp_size = 1 - self.dp_size = 1 - self.pp_size = 1 self.is_draft_worker = False self.max_running_requests = pool_batch_size # trtllm_mha __init__ scans model.modules() for ENCODER_ONLY layers; @@ -410,7 +407,6 @@ def __init__( get_kvcache=lambda: self.token_to_kv_pool, ) self.init_kv_index_translator() - self.attn_cp_size = 1 self.attention_chunk_size = None self.hisparse_coordinator = None self.init_new_workspace = False diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index 440e497e87d1..0458f1f265e9 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py @@ -317,10 +317,7 @@ def __init__( self.canary_manager = None self.page_size = case.page_size self.model_config = model_config - self.tp_size = 1 self._kernel_warmed_up = True - self.dp_size = 1 - self.pp_size = 1 self._server_args_override = get_context().override_server_args( attention_backend=case.backend, chunked_prefill_size=-1, @@ -403,7 +400,6 @@ def __init__( ) self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size) self.init_kv_index_translator() - self.attn_cp_size = 1 self.attention_chunk_size = None self.hisparse_coordinator = None self.init_new_workspace = False diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index 83b4161c9348..e07aa633ffde 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py @@ -345,9 +345,6 @@ def __init__( self.canary_manager = None self.page_size = case.page_size self.model_config = model_config - self.tp_size = 1 - self.dp_size = 1 - self.pp_size = 1 self._server_args_override = get_context().override_server_args( attention_backend=case.backend, chunked_prefill_size=-1, @@ -423,7 +420,6 @@ def __init__( identity = torch.arange(swa_size, dtype=torch.int64, device=device) self.token_to_kv_pool.register_mapping(identity) self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size) - self.attn_cp_size = 1 self.attention_chunk_size = None self.hisparse_coordinator = None self.init_new_workspace = False diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py index e7a3d4b61c41..af2f7d7005c5 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py @@ -323,7 +323,6 @@ def __init__( get_kvcache=lambda: self.token_to_kv_pool, ) self.init_kv_index_translator() - self.attn_cp_size = 1 self.attention_chunk_size = None self.hisparse_coordinator = None self.init_new_workspace = False diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py index 4254e7cd228a..ddd0fb9f98e0 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py @@ -319,7 +319,6 @@ def __init__( ) self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size) self.init_kv_index_translator() - self.attn_cp_size = 1 self.attention_chunk_size = None self.hisparse_coordinator = None self.init_new_workspace = False diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py index 42c285f2d00b..eafef40af80e 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py @@ -328,7 +328,6 @@ def __init__( ) self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size) self.init_kv_index_translator() - self.attn_cp_size = 1 self.attention_chunk_size = None self.hisparse_coordinator = None self.init_new_workspace = False diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py index 7b8c3675f0e3..76cf4fa4732b 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py @@ -454,7 +454,6 @@ def __init__( ) self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size) self.init_kv_index_translator() - self.attn_cp_size = 1 self.attention_chunk_size = None self.hisparse_coordinator = None self.init_new_workspace = False diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py index 0a381ca0ee3f..ae558bda2ea7 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py @@ -243,9 +243,6 @@ def __init__( self.canary_manager = None self.page_size = case.page_size self.model_config = model_config - self.tp_size = 1 - self.dp_size = 1 - self.pp_size = 1 self.spec_algorithm = SpeculativeAlgorithm.NONE speculative_num_draft_tokens = ( max(case.input_lens) @@ -312,7 +309,6 @@ def __init__( ) self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size) self.init_kv_index_translator() - self.attn_cp_size = 1 self.attention_chunk_size = None self.hisparse_coordinator = None self.init_new_workspace = False diff --git a/test/manual/attention/test_flashattn_backend.py b/test/manual/attention/test_flashattn_backend.py index 0122ef1f2a1c..02d6e6356ac4 100644 --- a/test/manual/attention/test_flashattn_backend.py +++ b/test/manual/attention/test_flashattn_backend.py @@ -62,7 +62,6 @@ def __init__( "enable_deterministic_inference": False, }, ) - self.attn_cp_size = 1 # Create a large enough req_to_token_pool to fit the test usage. self.req_to_token_pool = type( "TokenPool", diff --git a/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py b/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py index 20e83a183a99..e77137506d76 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py +++ b/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py @@ -48,8 +48,12 @@ def _stub_for_initialize( pool_size: int = 64, ): # ``initialize`` reads the config namespaces, so the config has to be - # published rather than stubbed onto the runner. + # published rather than stubbed onto the runner. One attention replica + # per TP rank gives the requested attention-DP width. override = get_context().override_server_args( + tp_size=attn_dp_size, + dp_size=dp_size, + enable_dp_attention=attn_dp_size > 1, enable_memory_saver=False, max_running_requests=max_running_requests, max_mamba_cache_size=max_mamba_cache_size, @@ -61,7 +65,6 @@ def _stub_for_initialize( stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub) stub._mlx_pool_size = pool_size stub.device = "cpu" - stub.attn_dp_size = attn_dp_size stub.server_args = server_args stub.model_config = SimpleNamespace( is_hybrid_swa=False, diff --git a/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py b/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py index 91a286724fff..974988edf30c 100644 --- a/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py +++ b/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py @@ -58,6 +58,10 @@ def _published(stub): """Publish the config the resolver reads (get_schedule() / get_memory()).""" return get_context().override_server_args( + # One attention replica per TP rank gives the stub's attention-DP width. + tp_size=stub._attn_dp_size, + dp_size=stub._attn_dp_size, + enable_dp_attention=stub._attn_dp_size > 1, max_running_requests=stub._max_running_requests, max_mamba_cache_size=stub._max_mamba_cache_size, disable_radix_cache=stub._disable_radix_cache, @@ -84,7 +88,7 @@ def _stub( stub._max_mamba_cache_size = max_mamba_cache_size stub._disable_radix_cache = disable_radix_cache stub.max_total_num_tokens = max_total_num_tokens - stub.attn_dp_size = dp_size + stub._attn_dp_size = dp_size return stub @@ -94,7 +98,7 @@ def _hybrid_stub_for_initialize( """A stub carrying what the real initialize() reads (hybrid path).""" stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub) stub._mlx_pool_size = pool - stub.attn_dp_size = 1 + stub._attn_dp_size = 1 stub.device = "cpu" # read by init_ngram_embedding_manager # Evaluated as a call argument in init_ngram_embedding_manager before # the use_ngram_embedding short-circuit; never read. diff --git a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py index 43e66379db70..774484717a0c 100644 --- a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py +++ b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py @@ -91,9 +91,6 @@ def _make_prefill_aware_swa_runner( kv_cache_dtype=torch.float16, kv_cache_dtype_str="auto", page_size=1, - attn_cp_size=1, - tp_size=1, - ps=SimpleNamespace(attn_cp_size=1, tp_size=1), is_draft_worker=False, server_args=server_args, attention_chunk_size=None, diff --git a/test/registered/unit/layers/attention/test_flashattention_return_lse.py b/test/registered/unit/layers/attention/test_flashattention_return_lse.py index 7dea741fe253..4e416f388613 100644 --- a/test/registered/unit/layers/attention/test_flashattention_return_lse.py +++ b/test/registered/unit/layers/attention/test_flashattention_return_lse.py @@ -54,8 +54,6 @@ def test_fa3_lse_and_sink_match_joint_softmax(self): kv_cache_dtype=torch.bfloat16, kv_cache_dtype_str="auto", page_size=1, - attn_cp_size=1, - tp_size=1, is_draft_worker=False, kv_index_translator=KVIndexTranslator( req_to_token=req_pool.req_to_token, diff --git a/test/registered/unit/layers/test_mamba_prefill_track_metadata.py b/test/registered/unit/layers/test_mamba_prefill_track_metadata.py index d8201c343ef3..6dc5098465d2 100644 --- a/test/registered/unit/layers/test_mamba_prefill_track_metadata.py +++ b/test/registered/unit/layers/test_mamba_prefill_track_metadata.py @@ -229,7 +229,6 @@ def test_forward_snapshot_and_padding_do_not_mutate_scheduler_lists(self): prefill_attention_backend_str="torch_native", ngram_embedding_manager=SimpleNamespace(enabled=False), lora_manager=None, - attn_dcp_size=1, attn_backend=SimpleNamespace( get_cpu_graph_seq_len_fill_value=lambda: 1, get_cuda_graph_seq_len_fill_value=lambda: 1, diff --git a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py index fa383ff034b9..2163f5a3d238 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py +++ b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py @@ -10,7 +10,12 @@ from torch import nn from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel +from sglang.test.test_utils import ( + CustomTestCase, + enter_scope, + maybe_stub_sgl_kernel, + published_topology, +) maybe_stub_sgl_kernel() @@ -658,10 +663,8 @@ def finalize(self): class TestModelRunnerStartupWeightLoadOwnership(CustomTestCase): def setUp(self): - # The code under test reads its config from the bags. - reset_context() - self.addCleanup(reset_context) - publish(ServerArgs(model_path="dummy"), role="tokenizer") + # The code under test reads its config and its rank from the context. + enter_scope(self, published_topology()) @staticmethod def _runner(manager): @@ -671,7 +674,6 @@ def _runner(manager): elastic_ep_backend=None, is_ep_joiner=False, ) - runner.tp_rank = 0 return runner def test_success_releases_ownership_after_the_barrier(self): diff --git a/test/registered/unit/spec/test_dflash_logits.py b/test/registered/unit/spec/test_dflash_logits.py index 079fb2fae3df..acfc8a585852 100644 --- a/test/registered/unit/spec/test_dflash_logits.py +++ b/test/registered/unit/spec/test_dflash_logits.py @@ -11,10 +11,18 @@ ) from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import published_topology register_cpu_ci(est_time=38, suite="base-a-test-cpu") +@pytest.fixture +def topology(): + # The worker reads its rank from the context. + with published_topology(): + yield + + def test_dflash_unary_logit_transform(): logits = torch.tensor([[-100.0, 0.0, 100.0]], dtype=torch.bfloat16) for fields in ({}, {"output_multiplier": 0.2, "final_logit_softcapping": 20.0}): @@ -248,7 +256,7 @@ def build_sampler(**kwargs): worker = SimpleNamespace( block_size=8, selector=object(), - model_runner=SimpleNamespace(tp_rank=0), + model_runner=SimpleNamespace(tp_group=SimpleNamespace(rank_in_group=0)), ps=SimpleNamespace(tp_rank=0), draft_model=SimpleNamespace(lm_head=None), device="cpu", @@ -272,7 +280,9 @@ def build_sampler(**kwargs): assert worker.draft_model.lm_head is None -def test_lilicorr_sampling_falls_back_when_the_device_cannot_accept_it(monkeypatch): +def test_lilicorr_sampling_falls_back_when_the_device_cannot_accept_it( + monkeypatch, topology +): from sglang.srt.speculative import dflash_worker_v2 as worker_mod warnings = [] @@ -287,7 +297,6 @@ def test_lilicorr_sampling_falls_back_when_the_device_cannot_accept_it(monkeypat lilicorr=object(), _lilicorr_sampling_enabled=False, _warned_sampling_fallback=False, - model_runner=SimpleNamespace(tp_rank=0), ) batch = SimpleNamespace(sampling_info=SimpleNamespace(is_all_greedy=False)) @@ -301,7 +310,7 @@ def test_lilicorr_sampling_falls_back_when_the_device_cannot_accept_it(monkeypat assert len(warnings) == 1 -def test_worker_warns_once_when_selector_sampling_is_disabled(monkeypatch): +def test_worker_warns_once_when_selector_sampling_is_disabled(monkeypatch, topology): from sglang.srt.speculative import dflash_worker_v2 as worker_mod warnings = [] @@ -312,7 +321,6 @@ def test_worker_warns_once_when_selector_sampling_is_disabled(monkeypatch): selector=object(), _selector_sampling_enabled=False, _warned_sampling_fallback=False, - model_runner=SimpleNamespace(tp_rank=0), ps=SimpleNamespace(tp_rank=0), ) batch = SimpleNamespace(sampling_info=SimpleNamespace(is_all_greedy=False)) diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 64058ba6aae6..b089fe91ff6c 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -3262,7 +3262,9 @@ def _calls_before_the_freeze(self, methods): def test_no_method_called_before_the_freeze_reads_what_it_freezes(self): methods = self._model_runner() frozen = self._frozen_names(methods) - self.assertGreater(len(frozen), 5, "found no frozen names; census is broken") + self.assertLessEqual( + {"tp_group", "pp_group"}, frozen, "found no frozen groups; census is broken" + ) offenders = [] for lineno, name in self._calls_before_the_freeze(methods): fn = methods.get(name)