[Fix] Scale the req_to_token row headroom by attn_dcp_size (#35424)
This commit is contained in:
@@ -1,12 +1,19 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from sglang.srt.runtime_context import (
|
from sglang.srt.runtime_context import (
|
||||||
|
get_parallel,
|
||||||
get_schedule,
|
get_schedule,
|
||||||
get_spec,
|
get_spec,
|
||||||
max_speculative_num_draft_tokens,
|
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:
|
def get_alloc_len_per_decode() -> int:
|
||||||
"""KV length one request may allocate in a single decode step.
|
"""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_steps = spec.speculative_num_steps or 1
|
||||||
spec_topk = spec.speculative_eagle_topk or 1
|
spec_topk = spec.speculative_eagle_topk or 1
|
||||||
spec_tokens = max_speculative_num_draft_tokens()
|
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
|
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
|
# FIXME(lsyin): temporary fix for the context length issue under spec decoding
|
||||||
extra = 4 + (max_speculative_num_draft_tokens() or 0)
|
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:
|
if get_spec().speculative_algorithm is not None and page_size > 1:
|
||||||
# kv_allocated_len is page-aligned (eagle_prepare_for_decode), so near
|
# kv_allocated_len is page-aligned (eagle_prepare_for_decode), so near
|
||||||
# the context limit the aligned reserve can overshoot by page_size - 1;
|
# the context limit the aligned reserve can overshoot by page_size - 1;
|
||||||
|
|||||||
Reference in New Issue
Block a user