diff --git a/python/sglang/srt/mem_cache/allocation_sizing.py b/python/sglang/srt/mem_cache/allocation_sizing.py index fc3c463da..92f114113 100644 --- a/python/sglang/srt/mem_cache/allocation_sizing.py +++ b/python/sglang/srt/mem_cache/allocation_sizing.py @@ -1,12 +1,19 @@ from __future__ import annotations from sglang.srt.runtime_context import ( + get_parallel, get_schedule, get_spec, max_speculative_num_draft_tokens, ) +def get_alloc_page_size() -> int: + # Mirrors _build_token_to_kv_pool_allocator's DCP branch; the platform + # allocators that skip it page smaller, so this is an upper bound for them. + return get_schedule().page_size * get_parallel().attn_dcp_size + + def get_alloc_len_per_decode() -> int: """KV length one request may allocate in a single decode step. @@ -21,7 +28,7 @@ def get_alloc_len_per_decode() -> int: spec_steps = spec.speculative_num_steps or 1 spec_topk = spec.speculative_eagle_topk or 1 spec_tokens = max_speculative_num_draft_tokens() - page_size = get_schedule().page_size + page_size = get_alloc_page_size() from sglang.srt.speculative.spec_info import SpeculativeAlgorithm @@ -80,7 +87,7 @@ def get_req_to_token_extra_context_len() -> int: """ # FIXME(lsyin): temporary fix for the context length issue under spec decoding extra = 4 + (max_speculative_num_draft_tokens() or 0) - page_size = get_schedule().page_size + page_size = get_alloc_page_size() if get_spec().speculative_algorithm is not None and page_size > 1: # kv_allocated_len is page-aligned (eagle_prepare_for_decode), so near # the context limit the aligned reserve can overshoot by page_size - 1;