From ce1b9f88b6f56bae6edc67a6e09324b925cc8331 Mon Sep 17 00:00:00 2001 From: Khoa Pham Date: Sun, 9 Aug 2026 01:18:24 -0700 Subject: [PATCH] config: derive the runner's DCP topology from its ParallelState (#34133) Co-authored-by: Claude Opus 5 (1M context) --- python/sglang/benchmark/one_batch.py | 3 ++- python/sglang/srt/distributed/bootstrap.py | 2 +- python/sglang/srt/distributed/parallel_state_wrapper.py | 6 ++++-- python/sglang/srt/managers/scheduler.py | 3 ++- python/sglang/srt/managers/tp_worker.py | 6 ++---- python/sglang/srt/model_executor/forward_batch_info.py | 5 +++-- python/sglang/srt/model_executor/model_runner.py | 3 --- python/sglang/srt/model_executor/runner/eager_runner.py | 2 +- 8 files changed, 15 insertions(+), 15 deletions(-) diff --git a/python/sglang/benchmark/one_batch.py b/python/sglang/benchmark/one_batch.py index 642f4d863..b94396c37 100644 --- a/python/sglang/benchmark/one_batch.py +++ b/python/sglang/benchmark/one_batch.py @@ -322,13 +322,14 @@ def load_model(server_args, port_args, gpu_id, tp_rank): attn_tp_size=attn_tp_size, attn_cp_rank=0, attn_cp_size=server_args.attn_cp_size, + attn_dcp_rank=tp_rank % server_args.dcp_size, + attn_dcp_size=server_args.dcp_size, attn_dp_rank=attn_dp_rank, attn_dp_size=attn_dp_size, moe_ep_rank=moe_ep_rank, moe_ep_size=server_args.ep_size, moe_dp_rank=None, moe_dp_size=server_args.moe_dp_size, - dcp_size=server_args.dcp_size, gpu_id=gpu_id, ) runner_kwargs = dict( diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index e0c37793d..f21844415 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/python/sglang/srt/distributed/bootstrap.py @@ -93,7 +93,7 @@ def init_torch_distributed( attn_cp_size=ps.attn_cp_size, moe_ep_size=ps.moe_ep_size, moe_dp_size=ps.moe_dp_size, - dcp_size=ps.dcp_size, + dcp_size=ps.attn_dcp_size, ) # Pre-warm NCCL/RCCL/HCCL to eliminate cold-start latency in first request diff --git a/python/sglang/srt/distributed/parallel_state_wrapper.py b/python/sglang/srt/distributed/parallel_state_wrapper.py index 6edb43eff..82a7ece32 100644 --- a/python/sglang/srt/distributed/parallel_state_wrapper.py +++ b/python/sglang/srt/distributed/parallel_state_wrapper.py @@ -14,13 +14,14 @@ class ParallelState: attn_tp_size: int attn_cp_rank: int attn_cp_size: int + attn_dcp_rank: int + attn_dcp_size: int attn_dp_rank: int attn_dp_size: int moe_ep_rank: int moe_ep_size: int moe_dp_rank: Optional[int] moe_dp_size: int - dcp_size: int gpu_id: int @staticmethod @@ -36,13 +37,14 @@ class ParallelState: attn_tp_size=1, attn_cp_rank=0, attn_cp_size=1, + attn_dcp_rank=0, + attn_dcp_size=1, attn_dp_rank=0, attn_dp_size=1, moe_ep_rank=0, moe_ep_size=1, moe_dp_rank=0, moe_dp_size=1, - dcp_size=1, gpu_id=0, ) kwargs.update(overrides) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 13570db3f..7954fd44e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -460,13 +460,14 @@ class Scheduler( attn_tp_size=attn_tp_size, attn_cp_rank=attn_cp_rank, attn_cp_size=server_args.attn_cp_size, + attn_dcp_rank=tp_rank % server_args.dcp_size, + attn_dcp_size=server_args.dcp_size, attn_dp_rank=attn_dp_rank, attn_dp_size=attn_dp_size, moe_ep_rank=moe_ep_rank, moe_ep_size=server_args.ep_size, moe_dp_rank=moe_dp_rank, moe_dp_size=server_args.moe_dp_size, - dcp_size=server_args.dcp_size, gpu_id=gpu_id, ) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index f6eec4b37..cf71bc64c 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -409,8 +409,7 @@ class TpModelWorker(BaseTpWorker): assert self.model_runner.max_running_requests > 0, "max_running_request is zero" max_req_len = min( self.model_config.context_len - 1, - self.model_runner.effective_max_total_num_tokens - * self.model_runner.dcp_size + self.model_runner.effective_max_total_num_tokens * self.ps.attn_dcp_size - 1, ) assert max_req_len > 0, "Memory pool size is too small" @@ -513,8 +512,7 @@ class TpModelWorker(BaseTpWorker): def get_worker_info(self): max_req_len = min( self.model_config.context_len - 1, - self.model_runner.effective_max_total_num_tokens - * self.model_runner.dcp_size + self.model_runner.effective_max_total_num_tokens * self.ps.attn_dcp_size - 1, ) return ( diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 14ebe2e83..2d01ef258 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -962,12 +962,13 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): model_runner.lora_manager.prepare_lora_batch(ret) if ( - getattr(model_runner, "dcp_size", 1) > 1 + model_runner.ps.attn_dcp_size > 1 and ret.out_cache_loc is not None and is_hip() ): ret.dcp_kv_mask = ( - ret.positions % model_runner.dcp_size == model_runner.dcp_rank + ret.positions % model_runner.ps.attn_dcp_size + == model_runner.ps.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 8fe718a9b..9788acf62 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -397,9 +397,6 @@ class ModelRunner: # Stored for later use by alloc_memory_pool(). self.init_torch_distributed() - self.dcp_size = get_parallel().attn_dcp_size - self.dcp_rank = get_parallel().attn_dcp_rank - # Init forward stream for overlap schedule self.forward_stream = torch.get_device_module(self.device).Stream() diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 5f0483140..e3e72d4ae 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -264,7 +264,7 @@ class EagerRunner(BaseRunner): prepare_cp_forward(forward_batch) if forward_batch.needs_forward_metadata_init() or cp_v2_active: - if model_runner.dcp_size > 1 and hasattr( + if model_runner.ps.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