From 1c7acba57937c26c459f07138be2ff4621c4100d Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 6 Jun 2026 14:58:50 -0700 Subject: [PATCH] [spec] Consolidate the per-decode KV alloc reserve into one helper (#27458) --- .../sglang/srt/debug_utils/pr_fix_toggle.py | 15 ++--- python/sglang/srt/managers/schedule_batch.py | 7 +-- python/sglang/srt/managers/utils.py | 33 ---------- python/sglang/srt/mem_cache/common.py | 62 ++++++++++++++++++- .../model_runner_kv_cache_mixin.py | 21 +------ .../sglang/srt/speculative/eagle_info_v2.py | 9 ++- 6 files changed, 74 insertions(+), 73 deletions(-) diff --git a/python/sglang/srt/debug_utils/pr_fix_toggle.py b/python/sglang/srt/debug_utils/pr_fix_toggle.py index 814a995ad..1f1edc361 100644 --- a/python/sglang/srt/debug_utils/pr_fix_toggle.py +++ b/python/sglang/srt/debug_utils/pr_fix_toggle.py @@ -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: "" """ diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index c3886d3b5..cc24f773b 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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) diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index eb21eda02..e9ede57ad 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -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. diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 11a52acdc..2139e2902 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -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, diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 3661197a4..189b4a3b5 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -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 ( diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index e90cff6d9..e458cab29 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -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