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

[spec] Consolidate the per-decode KV alloc reserve into one helper (#27458)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-06 14:58:50 -07:00
committed by GitHub
parent 4989d6691d
commit 1c7acba579
6 changed files with 74 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/debug_utils/pr_fix_toggle.py
+5 -10
View File
@@ -89,20 +89,15 @@ patches:
_PR_REVERT_YAML_26972 = """
patches:
- target: sglang.srt.model_executor.model_runner_kv_cache_mixin.ModelRunnerKVCacheMixin._init_pools
- target: sglang.srt.mem_cache.common.get_req_to_token_extra_context_len
edits:
- match: |
if (
self.server_args.speculative_algorithm is not None
and self.server_args.page_size > 1
and (self.server_args.speculative_eagle_topk or 1) > 1
server_args.speculative_algorithm is not None
and server_args.page_size > 1
and (server_args.speculative_eagle_topk or 1) > 1
):
from sglang.srt.managers.utils import get_alloc_len_per_decode
extra_max_context_len = max(
extra_max_context_len,
2 * get_alloc_len_per_decode(self.server_args),
)
extra = max(extra, get_alloc_reserve_per_decode(server_args))
replacement: ""
"""
python/sglang/srt/managers/schedule_batch.py
+3 -4
View File
@@ -83,6 +83,7 @@ from sglang.srt.mem_cache.common import (
alloc_for_decode,
alloc_for_extend,
evict_from_tree_cache,
get_alloc_reserve_per_decode,
release_kv_cache,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
@@ -2261,12 +2262,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
def _new_tokens_required_next_decode_spec_v2(self, requests, page_size):
"""Tight estimate matching eagle_info_v2.prepare_for_decode allocation."""
from sglang.srt.managers.utils import get_alloc_len_per_decode
alloc_len = get_alloc_len_per_decode()
reserve = get_alloc_reserve_per_decode()
total = 0
for r in requests:
x = max(0, r.kv_committed_len + 2 * alloc_len - r.kv_allocated_len)
x = max(0, r.kv_committed_len + reserve - r.kv_allocated_len)
cur = r.kv_allocated_len
nxt = cur + x
total += ceil_align(nxt, page_size) - ceil_align(cur, page_size)
python/sglang/srt/managers/utils.py
-33
View File
@@ -12,7 +12,6 @@ from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
from sglang.srt.server_args import ServerArgs
from sglang.srt.state_capturer.base import TopkCaptureOutput
if TYPE_CHECKING:
@@ -224,38 +223,6 @@ def get_logprob_from_pp_outputs(
return logits_output, extend_input_len_per_req, extend_logprob_start_len_per_req
def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
if server_args is None:
from sglang.srt.server_args import get_global_server_args
server_args = get_global_server_args()
if server_args.speculative_algorithm is None:
return 1
# Spec v1:
# 1) alloc topk * num_steps when draft decoding and then restore the allocation
# 2) alloc num_draft_tokens when verifying the drafts
# Sepc v2: allocate max(topk * num_steps, num_draft_tokens)
spec_steps = server_args.speculative_num_steps or 1
spec_topk = server_args.speculative_eagle_topk or 1
spec_tokens = server_args.max_speculative_num_draft_tokens
page_size = server_args.page_size
if page_size == 1 or spec_topk == 1:
return max(spec_steps * spec_topk, spec_tokens)
else:
# page_size > 1 + topk > 1 (spec v2 tree): worst-case page-aligned tree
# footprint. Per topk branch needs ceil((last_page_len + num_steps) / page)
# pages; the partial tail page can be up to page_size - 1, and each branch
# gets its own (duplicated) copy -- so reserve for all topk branches.
num_new_pages_per_topk = (
(page_size - 1) + spec_steps + page_size - 1
) // page_size
return max(num_new_pages_per_topk * page_size * spec_topk, spec_tokens)
@dataclass
class EmbeddingBatchResult:
"""Result from an embedding/classification forward pass.
python/sglang/srt/mem_cache/common.py
+60 -2
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Optional
import numpy as np
import torch
@@ -20,7 +20,7 @@ from sglang.srt.mem_cache.triton_ops.common import (
get_last_loc_triton_safe,
write_req_to_token_pool_triton,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.server_args import ServerArgs, get_global_server_args
from sglang.srt.utils import is_hip, support_triton
from sglang.srt.utils.common import ceil_align
@@ -152,6 +152,64 @@ def get_last_loc_torch(
)
def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
if server_args is None:
server_args = get_global_server_args()
if server_args.speculative_algorithm is None:
return 1
# Spec v1:
# 1) alloc topk * num_steps when draft decoding and then restore the allocation
# 2) alloc num_draft_tokens when verifying the drafts
# Sepc v2: allocate max(topk * num_steps, num_draft_tokens)
spec_steps = server_args.speculative_num_steps or 1
spec_topk = server_args.speculative_eagle_topk or 1
spec_tokens = server_args.max_speculative_num_draft_tokens
page_size = server_args.page_size
if page_size == 1 or spec_topk == 1:
return max(spec_steps * spec_topk, spec_tokens)
else:
# page_size > 1 + topk > 1 (spec v2 tree): worst-case page-aligned tree
# footprint. Per topk branch needs ceil((last_page_len + num_steps) / page)
# pages; the partial tail page can be up to page_size - 1, and each branch
# gets its own (duplicated) copy -- so reserve for all topk branches.
num_new_pages_per_topk = (
(page_size - 1) + spec_steps + page_size - 1
) // page_size
return max(num_new_pages_per_topk * page_size * spec_topk, spec_tokens)
def get_alloc_reserve_per_decode(server_args: Optional[ServerArgs] = None) -> int:
"""KV length reserved per request at each decode step.
The 2x is a double-buffer that absorbs the kv_committed_len lag in overlap
mode; see eagle_info_v2.prepare_for_decode.
"""
return 2 * get_alloc_len_per_decode(server_args)
def get_req_to_token_extra_context_len(server_args: ServerArgs) -> int:
"""req_to_token row headroom beyond the model context length.
Sized to hold the decode over-allocation (kv_committed_len +
get_alloc_reserve_per_decode). The spec v2 page>1 topk>1 holey draft footprint
can outgrow the default num_draft_tokens headroom (PR #26972).
"""
# FIXME(lsyin): this is the temporary fix for the context length issue when
# using speculative decoding
extra = 4 + (server_args.max_speculative_num_draft_tokens or 0)
if (
server_args.speculative_algorithm is not None
and server_args.page_size > 1
and (server_args.speculative_eagle_topk or 1) > 1
):
extra = max(extra, get_alloc_reserve_per_decode(server_args))
return extra
def alloc_token_slots(
tree_cache: BasePrefixCache,
num_tokens: int,
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+2 -19
View File
@@ -18,6 +18,7 @@ from sglang.srt.mem_cache.allocator import (
TokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.common import get_req_to_token_extra_context_len
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.mem_cache.hisparse_memory_pool import (
DeepSeekV4HiSparseTokenToKVPoolAllocator,
@@ -294,26 +295,8 @@ class ModelRunnerKVCacheMixin:
# Initialize req_to_token_pool
if self.req_to_token_pool is None:
# FIXME(lsyin): this is the temporary fix for the context length issue when using speculative decoding
max_spec_draft_tokens = self.server_args.max_speculative_num_draft_tokens
extra_max_context_len = 4
if max_spec_draft_tokens is not None:
extra_max_context_len += max_spec_draft_tokens
# page>1 + topk>1 reserves a holey draft footprint (2 * get_alloc_len_per_decode
# = topk * num_new_pages * page) far beyond the default num_draft_tokens
# headroom; widen the row to hold it, else free leaks KV and the holey gather OOBs.
if (
self.server_args.speculative_algorithm is not None
and self.server_args.page_size > 1
and (self.server_args.speculative_eagle_topk or 1) > 1
):
from sglang.srt.managers.utils import get_alloc_len_per_decode
extra_max_context_len = max(
extra_max_context_len,
2 * get_alloc_len_per_decode(self.server_args),
)
extra_max_context_len = get_req_to_token_extra_context_len(self.server_args)
if self.server_args.disaggregation_mode == "decode":
from sglang.srt.disaggregation.decode import (
python/sglang/srt/speculative/eagle_info_v2.py
+4 -5
View File
@@ -16,10 +16,10 @@ from sglang.srt.managers.schedule_batch import (
ScheduleBatch,
set_mamba_track_indices_from_reqs,
)
from sglang.srt.managers.utils import get_alloc_len_per_decode
from sglang.srt.mem_cache.common import (
alloc_paged_token_slots_extend,
alloc_token_slots,
get_alloc_reserve_per_decode,
get_last_loc,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
@@ -150,8 +150,7 @@ class EagleDraftInputV2Mixin:
)
page_size = batch.token_to_kv_pool_allocator.page_size
alloc_len_per_decode = get_alloc_len_per_decode()
double_alloc = alloc_len_per_decode + alloc_len_per_decode
double_alloc = get_alloc_reserve_per_decode()
cur_kv_lens = [0] * bs
nxt_kv_lens = [0] * bs
@@ -176,7 +175,7 @@ class EagleDraftInputV2Mixin:
nxt_kv_lens_cpu = torch.tensor(nxt_kv_lens, dtype=torch.int32, device="cpu")
# Fail fast if the page>1 + topk>1 draft over-allocation
# (2 * get_alloc_len_per_decode) outgrows the req_to_token row: the write below
# (get_alloc_reserve_per_decode) outgrows the req_to_token row: the write below
# would OOB and free would leak KV. The row is widened to hold it in _init_pools
# (PR #26972); fail here with a clear error, not on a later cryptic CUDA assert.
from sglang.srt.server_args import get_global_server_args
@@ -187,7 +186,7 @@ class EagleDraftInputV2Mixin:
assert max_alloc_len <= row_width, (
f"spec v2 page>1 topk>1 draft over-allocation ({max_alloc_len}) exceeds "
f"req_to_token row width ({row_width}); page_size={page_size}. Widen the "
f"row to hold committed + 2 * get_alloc_len_per_decode (PR #26972)."
f"row to hold committed + get_alloc_reserve_per_decode (PR #26972)."
)
# non_blocking H2D: a blocking .to() syncs the schedule stream, which the WAR