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

config: route DCP topology reads through get_parallel() (#33925)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
2026-08-07 14:53:54 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent df3aa20d89
commit 07297049e9
14 changed files with 31 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/common/conn.py
+2 -2
View File
@@ -169,8 +169,8 @@ class CommonKVManager(BaseKVManager):
self.attn_tp_rank = parallel.attn_tp_rank
self.attn_cp_size = parallel.attn_cp_size
self.attn_cp_rank = parallel.attn_cp_rank
self.dcp_size = server_args.dcp_size
self.dcp_rank = parallel.dcp_rank if self.dcp_size > 1 else 0
self.dcp_size = parallel.attn_dcp_size
self.dcp_rank = parallel.attn_dcp_rank
self.attn_dp_size = get_attention_dp_size()
self.attn_dp_rank = get_attention_dp_rank()
self.system_dp_size = (
python/sglang/srt/layers/attention/attention_registry.py
+2 -2
View File
@@ -13,7 +13,7 @@ from sglang.srt.configs.linear_attn_model_registry import (
get_linear_attn_config,
import_backend_class,
)
from sglang.srt.runtime_context import get_context
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.srt.utils import get_device_capability, is_hip, is_musa, is_npu
_is_musa = is_musa()
@@ -71,7 +71,7 @@ def create_trtllm_mla_backend(runner):
if not runner.use_mla_backend:
raise ValueError("trtllm_mla backend can only be used with MLA models.")
if (
runner.server_args.dcp_size > 1
get_parallel().dcp_enabled
and runner.server_args.speculative_algorithm is not None
):
_, decode_backend = runner.server_args.get_attention_backends()
python/sglang/srt/layers/attention/dsa/utils.py
+1 -1
View File
@@ -78,7 +78,7 @@ def should_remap_pd_dsa_seed_to_local_slots(server_args: "ServerArgs") -> bool:
and envs.SGLANG_DSA_FUSE_TOPK.get()
and server_args.disaggregation_mode == "decode"
and not server_args.enable_hisparse
and server_args.dcp_size == 1
and not get_parallel().dcp_enabled
)
python/sglang/srt/managers/scheduler.py
+3 -3
View File
@@ -2031,7 +2031,8 @@ class Scheduler(
enable_hisparse=self.enable_hisparse,
full_tokens_per_layer=self.full_tokens_per_layer,
swa_tokens_per_layer=self.swa_tokens_per_layer,
max_total_num_tokens=self.max_total_num_tokens * self.server_args.dcp_size,
max_total_num_tokens=self.max_total_num_tokens
* get_parallel().attn_dcp_size,
get_last_batch=lambda: self.last_batch,
get_running_batch=lambda: self.running_batch,
)
@@ -2045,7 +2046,6 @@ class Scheduler(
full_tokens_per_layer=self.full_tokens_per_layer,
swa_tokens_per_layer=self.swa_tokens_per_layer,
max_total_num_tokens=self.max_total_num_tokens,
server_args=self.server_args,
tree_cache=self.tree_cache,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
req_to_token_pool=self.req_to_token_pool,
@@ -2164,7 +2164,7 @@ class Scheduler(
min(
max_new_tokens,
self.max_req_len - input_len - 1,
self.max_total_num_tokens * self.server_args.dcp_size
self.max_total_num_tokens * get_parallel().attn_dcp_size
- paged_input_len
- self.page_size
- 1,
python/sglang/srt/managers/scheduler_components/invariant_checker.py
+3 -8
View File
@@ -23,7 +23,7 @@ from sglang.srt.managers.scheduler_components.pool_stats_observer import (
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.server_args import ServerArgs
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils.common import (
ceil_align,
raise_error_or_warn,
@@ -49,7 +49,6 @@ class SchedulerInvariantChecker:
full_tokens_per_layer: Optional[int]
swa_tokens_per_layer: Optional[int]
max_total_num_tokens: int
server_args: ServerArgs
tree_cache: BasePrefixCache
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
req_to_token_pool: ReqToTokenPool
@@ -106,7 +105,7 @@ class SchedulerInvariantChecker:
total = self.max_total_num_tokens
full_evictable_size = ps.full_evictable_size
allocator = self.token_to_kv_pool_allocator
if getattr(self.server_args, "dcp_size", 1) > 1 and allocator.page_size > 1:
if get_parallel().dcp_enabled and allocator.page_size > 1:
# DCP stores logical tokens in widened physical pages. Prefix cache
# counters are logical-token based, while the allocator frees whole
# physical pages, so round cached tokens up to physical page units.
@@ -124,11 +123,7 @@ class SchedulerInvariantChecker:
total,
uncached,
)
if (
leak
and getattr(self.server_args, "dcp_size", 1) > 1
and allocator.page_size > 1
):
if leak and get_parallel().dcp_enabled and allocator.page_size > 1:
# Radix/Mamba cache accounting is logical-token based while DCP full
# KV allocation is physical-page based. Partial physical pages can
# leave a small page-level slack even when all pages are owned by
python/sglang/srt/mem_cache/allocation.py
+3 -3
View File
@@ -26,7 +26,7 @@ from sglang.srt.mem_cache.common import (
evict_from_tree_cache,
)
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
from sglang.srt.runtime_context import get_exec, get_server_args
from sglang.srt.runtime_context import get_exec, get_parallel
from sglang.srt.utils import (
is_cpu,
is_cuda,
@@ -292,10 +292,10 @@ def alloc_req_slots(
def _alloc_page_size(batch: ScheduleBatch) -> int:
# DCP swaps in an allocator whose page_size is server_args.page_size *
# DCP swaps in an allocator whose page_size is the configured page_size *
# dcp_size, so it can be > 1 even when tree_cache.page_size is 1; branch on
# the real allocator's page_size there. Elsewhere the two are equal.
if (_is_hip or _is_cuda) and get_server_args().dcp_size > 1:
if (_is_hip or _is_cuda) and get_parallel().dcp_enabled:
return batch.tree_cache.token_to_kv_pool_allocator.page_size
return batch.tree_cache.page_size
python/sglang/srt/mem_cache/kv_cache_configurator.py
+4 -4
View File
@@ -277,7 +277,7 @@ class KVCacheConfigurator:
# 2. A pool must page as its allocator does, or its last page falls short.
@property
def loc_space_scale(self) -> int:
dcp_size = self.server_args.dcp_size
dcp_size = get_parallel().attn_dcp_size
return dcp_size if (self.is_draft_worker and dcp_size > 1) else 1
@property
@@ -1550,7 +1550,7 @@ class KVCacheConfigurator:
host_to_device_ratio=hisparse_cfg.host_to_device_ratio,
)
elif (
get_schedule().page_size == 1 and self.server_args.dcp_size == 1
get_schedule().page_size == 1 and not get_parallel().dcp_enabled
):
token_to_kv_pool_allocator = TokenToKVPoolAllocator(
sizes.max_total_num_tokens,
@@ -1561,9 +1561,9 @@ class KVCacheConfigurator:
)
else:
token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
sizes.max_total_num_tokens * self.server_args.dcp_size,
sizes.max_total_num_tokens * get_parallel().attn_dcp_size,
page_size=get_schedule().page_size
* self.server_args.dcp_size,
* get_parallel().attn_dcp_size,
dtype=self.kv_cache_dtype,
device=self.device,
kvcache=token_to_kv_pool,
python/sglang/srt/model_executor/model_runner.py
+4 -3
View File
@@ -306,8 +306,6 @@ class ModelRunner:
self.memory_pool_config = memory_pool_config
self.device = server_args.device
self.gpu_id = gpu_id
self.dcp_size = server_args.dcp_size
self.dcp_rank = ps.tp_rank % self.dcp_size
self.ps = ps
self.model_config = model_config
self.dist_port = nccl_port
@@ -399,6 +397,9 @@ 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()
@@ -938,7 +939,7 @@ class ModelRunner:
self.decode_attn_backend = backends.decode_attn_backend
self.decode_attn_backend_group = backends.decode_attn_backend_group
if self.server_args.dcp_size > 1 and get_parallel().dcp_replicate_q_proj:
if get_parallel().dcp_enabled and get_parallel().dcp_replicate_q_proj:
self._prepare_replicated_q_proj()
def _prepare_replicated_q_proj(self) -> None:
python/sglang/srt/model_executor/pool_configurator.py
+2 -1
View File
@@ -177,7 +177,8 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
self._cell_size = scale_kv_cell_size_per_token_for_dflash(
target_cell_size_per_token=self._cell_size,
target_num_layers=int(num_layers),
draft_num_layers=int(draft_num_layers) * kvc.server_args.dcp_size,
draft_num_layers=int(draft_num_layers)
* get_parallel().attn_dcp_size,
)
def _compute_cell_size(self, kvc: KVCacheConfigurator, num_layers: int) -> int:
python/sglang/srt/model_executor/runner/base_runner.py
+4 -2
View File
@@ -282,8 +282,10 @@ class BaseRunner(ABC):
comm backend; must run before CG capture (it syncs the stream + barriers
cross-rank, uncapturable) and raises early on non-MNNVL platforms.
"""
mr = self.model_runner
if mr.server_args.dcp_size <= 1 or mr.server_args.dcp_comm_backend != "fi_a2a":
if (
not get_parallel().dcp_enabled
or get_parallel().dcp_comm_backend != "fi_a2a"
):
return
from sglang.srt.layers.dcp import init_fi_a2a_workspace