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

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
+              server_args.speculative_algorithm is not None
-              and self.server_args.page_size > 1
+              and server_args.page_size > 1
-              and (self.server_args.speculative_eagle_topk or 1) > 1
+              and (server_args.speculative_eagle_topk or 1) > 1
           ):
-              from sglang.srt.managers.utils import get_alloc_len_per_decode
+              extra = max(extra, get_alloc_reserve_per_decode(server_args))
-
-              extra_max_context_len = max(
-                  extra_max_context_len,
-                  2 * get_alloc_len_per_decode(self.server_args),
-              )
         replacement: ""
 """
 

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
+        reserve = get_alloc_reserve_per_decode()
-
-        alloc_len = get_alloc_len_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

@@ -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

@@ -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

@@ -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
+            extra_max_context_len = get_req_to_token_extra_context_len(self.server_args)
-            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),
-                )
 
             if self.server_args.disaggregation_mode == "decode":
                 from sglang.srt.disaggregation.decode import (

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 = get_alloc_reserve_per_decode()
-        double_alloc = alloc_len_per_decode + alloc_len_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