From e789ca24a73613ade58244d534fa2fce667df6ae Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Wed, 15 Jul 2026 14:49:37 +0800 Subject: [PATCH] Lightweight extract allocation logic from mem_cache/common.py to more clearly show nearly parallel variants (#29431) --- .../sglang/srt/debug_utils/pr_fix_toggle.py | 2 +- .../npu/dsv4/dsv4_allocator.py | 2 +- python/sglang/srt/managers/schedule_batch.py | 8 +- python/sglang/srt/mem_cache/allocation.py | 469 ++++++++++++++++ .../sglang/srt/mem_cache/allocation_sizing.py | 60 +++ python/sglang/srt/mem_cache/common.py | 507 +----------------- .../srt/mem_cache/kv_cache_configurator.py | 2 +- .../srt/model_executor/pool_configurator.py | 2 +- .../sglang/srt/speculative/dflash_info_v2.py | 2 +- python/sglang/srt/speculative/eagle_utils.py | 4 +- .../spec/test_decode_bookkeeping_ownership.py | 4 +- 11 files changed, 546 insertions(+), 516 deletions(-) create mode 100644 python/sglang/srt/mem_cache/allocation.py create mode 100644 python/sglang/srt/mem_cache/allocation_sizing.py diff --git a/python/sglang/srt/debug_utils/pr_fix_toggle.py b/python/sglang/srt/debug_utils/pr_fix_toggle.py index 4f7e06185..77bdd4919 100644 --- a/python/sglang/srt/debug_utils/pr_fix_toggle.py +++ b/python/sglang/srt/debug_utils/pr_fix_toggle.py @@ -66,7 +66,7 @@ patches: _PR_REVERT_YAML_26972 = """ patches: - - target: sglang.srt.mem_cache.common.get_req_to_token_extra_context_len + - target: sglang.srt.mem_cache.allocation_sizing.get_req_to_token_extra_context_len edits: - match: | if ( diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py index 4e702bd4f..7060f90ce 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py @@ -33,8 +33,8 @@ from sglang.srt.hardware_backend.npu.allocator_npu import NPUPagedTokenToKVPoolA from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( maybe_write_dsv4_extend, ) +from sglang.srt.mem_cache.allocation import alloc_paged_token_slots_extend from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator -from sglang.srt.mem_cache.common import alloc_paged_token_slots_extend from sglang.srt.model_executor.forward_batch_info import DSV4OutCacheLoc, DSV4StateLens if TYPE_CHECKING: diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 7177b5787..1132e6811 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -77,6 +77,11 @@ from sglang.srt.managers.embed_types import PositionalEmbeds from sglang.srt.managers.scheduler_components.new_token_ratio_tracker import ( NewTokenRatioTracker, ) +from sglang.srt.mem_cache.allocation import ( + alloc_for_decode, + alloc_for_extend, +) +from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, @@ -85,11 +90,8 @@ from sglang.srt.mem_cache.base_prefix_cache import ( zero_match_result, ) from sglang.srt.mem_cache.common import ( - alloc_for_decode, - alloc_for_extend, evict_from_tree_cache, free_swa_out_of_window_slots, - get_alloc_reserve_per_decode, release_kv_cache, ) from sglang.srt.mem_cache.memory_pool import ReqToTokenPool diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py new file mode 100644 index 000000000..f8a0a4add --- /dev/null +++ b/python/sglang/srt/mem_cache/allocation.py @@ -0,0 +1,469 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernels.ops.memory.common import ( + get_last_loc_triton, + get_last_loc_triton_safe, + write_req_to_token_pool_triton, +) +from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( + maybe_write_dsv4_decode, + maybe_write_dsv4_extend, +) +from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams +from sglang.srt.mem_cache.common import ( + MAMBA_STATE_PER_REQ_NO_CACHE, + MAMBA_STATE_PER_REQ_PREFIX_CACHE, + MAMBA_STATE_PER_REQ_PREFIX_CACHE_LAZY, + available_and_evictable_str, + evict_from_tree_cache, +) +from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool +from sglang.srt.runtime_context import get_server_args +from sglang.srt.utils import is_cuda, is_hip, is_npu, support_triton +from sglang.srt.utils.common import is_pin_memory_available + +_is_hip = is_hip() +_is_npu = is_npu() +_is_cuda = is_cuda() + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req, ScheduleBatch + from sglang.srt.model_executor.forward_batch_info import DSV4StateLens + +logger = logging.getLogger(__name__) + + +def write_cache_indices( + out_cache_loc: torch.Tensor, + req_pool_indices_tensor: torch.Tensor, + req_pool_indices_cpu: torch.Tensor, + prefix_lens_tensor: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens_tensor: torch.Tensor, + seq_lens_cpu: torch.Tensor, + extend_lens_tensor: torch.Tensor, + extend_lens_cpu: torch.Tensor, + prefix_tensors: list[torch.Tensor], + req_to_token_pool: ReqToTokenPool, +): + if support_triton(get_server_args().attention_backend): + prefix_pointers = torch.tensor( + [t.data_ptr() for t in prefix_tensors], + dtype=torch.uint64, + pin_memory=is_pin_memory_available(req_to_token_pool.device), + ).to(req_to_token_pool.device, non_blocking=True) + # TODO: some tensors can be reused for ForwardBatchInfo (e.g., extend_lens, cumsum_start) + write_req_to_token_pool_triton[(req_pool_indices_tensor.shape[0],)]( + req_to_token_pool.req_to_token, + req_pool_indices_tensor, + prefix_pointers, + prefix_lens_tensor, + seq_lens_tensor, + extend_lens_tensor, + out_cache_loc, + req_to_token_pool.req_to_token.shape[1], + ) + else: + pt = 0 + for i in range(req_pool_indices_cpu.shape[0]): + req_idx = req_pool_indices_cpu[i].item() + prefix_len = prefix_lens_cpu[i].item() + seq_len = seq_lens_cpu[i].item() + extend_len = extend_lens_cpu[i].item() + + req_to_token_pool.write( + (req_idx, slice(0, prefix_len)), + prefix_tensors[i], + ) + req_to_token_pool.write( + (req_idx, slice(prefix_len, seq_len)), + out_cache_loc[pt : pt + extend_len], + ) + pt += extend_len + + +def get_last_loc( + req_to_token: torch.Tensor, + req_pool_indices_tensor: torch.Tensor, + prefix_lens_tensor: torch.Tensor, +) -> torch.Tensor: + attn_backend = get_server_args().attention_backend + uses_triton_dispatch = attn_backend not in ("ascend", "torch_native") + + if _is_hip and uses_triton_dispatch: + # HIP-only: the legacy get_last_loc_triton kernel emits a + # mixed-width int32->int64 store that Triton mis-compiles on HIP, + # producing out-of-range last_loc values under EAGLE + + # page_size>1 (e.g. with aiter unified attention or the triton + # attention backend). The bug is in the Triton HIP codegen, not + # in any particular attention backend, so route every HIP path + # that would otherwise use get_last_loc_triton through the + # int32-safe variant. Non-HIP hardware keeps the original + # dispatcher below. + return get_last_loc_triton_safe( + req_to_token, req_pool_indices_tensor, prefix_lens_tensor + ) + + if uses_triton_dispatch: + impl = get_last_loc_triton + else: + impl = get_last_loc_torch + + return impl(req_to_token, req_pool_indices_tensor, prefix_lens_tensor) + + +def get_last_loc_torch( + req_to_token: torch.Tensor, + req_pool_indices_tensor: torch.Tensor, + prefix_lens_tensor: torch.Tensor, +) -> torch.Tensor: + return torch.where( + prefix_lens_tensor > 0, + req_to_token[req_pool_indices_tensor, prefix_lens_tensor - 1], + torch.full_like(prefix_lens_tensor, -1), + ) + + +def alloc_token_slots( + tree_cache: BasePrefixCache, + num_tokens: int, + backup_state: bool = False, +): + allocator = tree_cache.token_to_kv_pool_allocator + evict_from_tree_cache(tree_cache, num_tokens) + + state = None + if backup_state: + state = allocator.backup_state() + + out_cache_loc = allocator.alloc(num_tokens) + + if out_cache_loc is None: + error_msg = ( + f"Out of memory. Try to lower your batch size.\n" + f"Try to allocate {num_tokens} tokens.\n" + f"{available_and_evictable_str(tree_cache)}" + ) + logger.error(error_msg) + if tree_cache is not None: + tree_cache.pretty_print() + raise RuntimeError(error_msg) + + return (out_cache_loc, state) if backup_state else out_cache_loc + + +def _compute_dsv4_state_lens(batch, *, is_decode: bool): + """Per-req c{4,128}_state pool alloc lens (``DSV4StateLens``) for this step. + None on CUDA / non-V4 paths (allocator has no ``compute_dsv4_state_lens_*``). + """ + allocator = batch.token_to_kv_pool_allocator + if not hasattr(allocator, "compute_dsv4_state_lens_extend"): + return None + if is_decode: + return allocator.compute_dsv4_state_lens_decode(batch.reqs) + return allocator.compute_dsv4_state_lens_extend( + batch.reqs, batch.seq_lens_cpu.tolist() + ) + + +def alloc_paged_token_slots_extend( + tree_cache: BasePrefixCache, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + backup_state: bool = False, + req_pool_indices: Optional[torch.Tensor] = None, + dsv4_state_lens: Optional[DSV4StateLens] = None, + batch=None, +): + # Over estimate the number of tokens: assume each request needs a new page. + allocator = tree_cache.token_to_kv_pool_allocator + num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size + evict_from_tree_cache(tree_cache, num_tokens) + + state = None + if backup_state: + state = allocator.backup_state() + + is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c4_attn_allocator") + extra_alloc_kwargs = {} + if is_dsv4: + extra_alloc_kwargs["req_pool_indices"] = req_pool_indices + # Per-call per-req tables for the c-pool / state last_loc lookup. + if batch is not None: + extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool + if dsv4_state_lens is not None: + extra_alloc_kwargs["dsv4_state_lens"] = dsv4_state_lens + + out = allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + **extra_alloc_kwargs, + ) + + if is_dsv4: + bundle = out + out_cache_loc = None if bundle is None else bundle.out_full_loc + if batch is not None: + batch.out_cache_loc_dsv4 = bundle + else: + out_cache_loc = out + + if out_cache_loc is None: + error_msg = ( + f"Prefill out of memory. Try to lower your batch size.\n" + f"Try to allocate {extend_num_tokens} tokens.\n" + f"{available_and_evictable_str(tree_cache)}" + ) + logger.error(error_msg) + if tree_cache is not None: + tree_cache.pretty_print() + raise RuntimeError(error_msg) + + return (out_cache_loc, state) if backup_state else out_cache_loc + + +def alloc_req_slots( + req_to_token_pool: ReqToTokenPool, + reqs: list[Req], + tree_cache: BasePrefixCache | None, +) -> list[int]: + """Allocate request slots from the pool. + + Fail-loud: raises ``RuntimeError`` if the pool can't satisfy the batch. An + alloc failure here means the admission budget (``PrefillAdder``) was wrong + and should surface rather than be masked. + """ + num_reqs = len(reqs) + if isinstance(req_to_token_pool, HybridReqToTokenPool): + # Byte-coordinated for the shared allocator (accounts for the peer full + # sub-pool's bytes); plain slot free count for the non-shared one. + mamba_available_size = ( + req_to_token_pool.mamba_allocator.schedulable_available_size() + ) + # Eviction headroom factor: 3x (or lazy variant) for radix COW, 1x for chunk. + if tree_cache.supports_mamba(): + factor = ( + MAMBA_STATE_PER_REQ_PREFIX_CACHE_LAZY + if req_to_token_pool.enable_mamba_extra_buffer_lazy + else MAMBA_STATE_PER_REQ_PREFIX_CACHE + ) + else: + factor = MAMBA_STATE_PER_REQ_NO_CACHE + mamba_state_needed = num_reqs * factor + if mamba_available_size < mamba_state_needed: + if tree_cache is not None and tree_cache.supports_mamba(): + mamba_num = max(0, mamba_state_needed - mamba_available_size) + tree_cache.evict(EvictParams(num_tokens=0, mamba_num=mamba_num)) + req_pool_indices = req_to_token_pool.alloc(reqs) + if req_pool_indices is None: + raise RuntimeError( + "alloc_req_slots runs out of memory. " + "Please set a smaller number for `--max-running-requests`. " + f"{req_to_token_pool.available_size()=}, {num_reqs=}, " + ) + return req_pool_indices + + +def _alloc_page_size(batch: ScheduleBatch) -> int: + # DCP swaps in an allocator whose page_size is server_args.page_size * + # dcp_size, so it can be > 1 even when tree_cache.page_size is 1; branch on + # the real allocator's page_size there. Elsewhere the two are equal. + if (_is_hip or _is_cuda) and get_server_args().dcp_size > 1: + return batch.tree_cache.token_to_kv_pool_allocator.page_size + return batch.tree_cache.page_size + + +def alloc_for_extend( + batch: ScheduleBatch, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Allocate KV cache for extend batch and write to req_to_token_pool. + + Returns ``(out_cache_loc, req_pool_indices_device, req_pool_indices_cpu)`` + (the last is the host/CPU mirror). ``alloc_req_slots`` raises ``RuntimeError`` + if the pool can't satisfy the batch (fail-loud — see its docstring). + """ + # free out-of-window swa tokens + batch.maybe_evict_swa() + + prefix_tensors = [r.prefix_indices for r in batch.reqs] + + # Create tensors for allocation + prefix_lens_cpu = torch.tensor(batch.prefix_lens, dtype=torch.int64) + extend_lens_cpu = torch.tensor(batch.extend_lens, dtype=torch.int64) + prefix_lens_device = prefix_lens_cpu.to(batch.device, non_blocking=True) + extend_lens_device = extend_lens_cpu.to(batch.device, non_blocking=True) + + # Allocate req slots (raises RuntimeError if the pool is exhausted) + req_pool_indices = alloc_req_slots( + batch.req_to_token_pool, batch.reqs, batch.tree_cache + ) + req_pool_indices_cpu = torch.tensor(req_pool_indices, dtype=torch.int64) + req_pool_indices_device = req_pool_indices_cpu.to(batch.device, non_blocking=True) + + # Allocate KV cache (throws exception on failure) + if _alloc_page_size(batch) == 1: + out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens) + else: + # Paged allocation - build last_loc + last_loc = [ + (t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device)) + for t in prefix_tensors + ] + out_cache_loc = alloc_paged_token_slots_extend( + tree_cache=batch.tree_cache, + prefix_lens=prefix_lens_device, + prefix_lens_cpu=prefix_lens_cpu, + seq_lens=batch.seq_lens, + seq_lens_cpu=batch.seq_lens_cpu, + last_loc=torch.cat(last_loc), + extend_num_tokens=batch.extend_num_tokens, + req_pool_indices=req_pool_indices_device, + dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=False), + batch=batch, + ) + + # Write to req_to_token_pool + write_cache_indices( + out_cache_loc, + req_pool_indices_device, + req_pool_indices_cpu, + prefix_lens_device, + prefix_lens_cpu, + batch.seq_lens, + batch.seq_lens_cpu, + extend_lens_device, + extend_lens_cpu, + prefix_tensors, + batch.req_to_token_pool, + ) + + # DSV4-NPU hook: no-op on non-DSV4 paths. + if _is_npu: + maybe_write_dsv4_extend( + batch, + req_pool_indices_cpu, + prefix_lens_cpu, + batch.seq_lens_cpu, + ) + + return out_cache_loc, req_pool_indices_device, req_pool_indices_cpu + + +def alloc_paged_token_slots_decode( + tree_cache: BasePrefixCache, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + token_per_req: int = 1, + req_pool_indices: Optional[torch.Tensor] = None, + dsv4_state_lens: Optional[DSV4StateLens] = None, + batch=None, +) -> torch.Tensor: + """Allocate paged KV cache for decode batch.""" + allocator = tree_cache.token_to_kv_pool_allocator + # Over estimate the number of tokens: assume each request needs a new page. + num_tokens = len(seq_lens) * allocator.page_size + evict_from_tree_cache(tree_cache, num_tokens) + + # DSV4-NPU allocator also needs req_pool_indices + per-req state lens and + # returns a DSV4OutCacheLoc bundle; hasattr-gated so others stay unchanged. + is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c4_attn_allocator") + extra_alloc_kwargs = {} + if is_dsv4: + extra_alloc_kwargs["req_pool_indices"] = req_pool_indices + # Per-call per-req tables for the last_loc lookup. + if batch is not None: + extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool + if dsv4_state_lens is not None: + extra_alloc_kwargs["dsv4_state_lens"] = dsv4_state_lens + + out = allocator.alloc_decode(seq_lens, seq_lens_cpu, last_loc, **extra_alloc_kwargs) + + if is_dsv4: + bundle = out + out_cache_loc = None if bundle is None else bundle.out_full_loc + if batch is not None: + batch.out_cache_loc_dsv4 = bundle + else: + out_cache_loc = out + + if out_cache_loc is None: + error_msg = ( + f"Decode out of memory. Try to lower your batch size.\n" + f"Try to allocate {len(seq_lens) * token_per_req} tokens.\n" + f"{available_and_evictable_str(tree_cache)}" + ) + logger.error(error_msg) + if tree_cache is not None: + tree_cache.pretty_print() + raise RuntimeError(error_msg) + + return out_cache_loc + + +def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor: + """ + Allocate KV cache for decode batch and write to req_to_token_pool. + + Returns: + out_cache_loc: allocated cache locations + """ + + batch.maybe_evict_swa() + + seq_lens_gpu = batch.seq_lens + bs = seq_lens_gpu.shape[0] + + if _alloc_page_size(batch) == 1: + # Non-paged allocation + out_cache_loc = alloc_token_slots(batch.tree_cache, bs * token_per_req) + else: + # Paged allocation + last_loc = batch.req_to_token_pool.req_to_token[ + batch.req_pool_indices, seq_lens_gpu - 1 + ] + seq_lens_next = seq_lens_gpu + token_per_req + out_cache_loc = alloc_paged_token_slots_decode( + tree_cache=batch.tree_cache, + seq_lens=seq_lens_next, + seq_lens_cpu=batch.seq_lens_cpu + token_per_req, + last_loc=last_loc, + token_per_req=token_per_req, + req_pool_indices=batch.req_pool_indices, + dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=True), + batch=batch, + ) + + # Write to req_to_token_pool + if batch.model_config.is_encoder_decoder: + locs = batch.encoder_lens + seq_lens_gpu + else: + locs = seq_lens_gpu.clone() + + batch.req_to_token_pool.write( + (batch.req_pool_indices, locs), out_cache_loc.to(torch.int32) + ) + + # DSV4-NPU hook: no-op on non-DSV4 paths. + if _is_npu: + maybe_write_dsv4_decode( + batch, + batch.seq_lens_cpu + token_per_req, + token_per_req, + ) + + return out_cache_loc diff --git a/python/sglang/srt/mem_cache/allocation_sizing.py b/python/sglang/srt/mem_cache/allocation_sizing.py new file mode 100644 index 000000000..b1fcf293a --- /dev/null +++ b/python/sglang/srt/mem_cache/allocation_sizing.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from typing import Optional + +from sglang.srt.runtime_context import get_server_args +from sglang.srt.server_args import ServerArgs + + +def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int: + if server_args is None: + server_args = get_server_args() + + if server_args.speculative_algorithm is None: + return 1 + + # Spec decoding allocates max(topk * num_steps, num_draft_tokens) per decode step. + 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 + + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + spec_algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm) + if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv(): + return max(spec_steps * spec_topk, spec_tokens) + else: + # spec v2 tree (page>1, topk>1): worst-case page-aligned footprint per + # topk branch is ceil((page_size-1 + num_steps) / page) pages, each branch + # duplicated -- 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_utils.eagle_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; the spec v2 page>1 topk>1 holey + draft footprint can outgrow the default num_draft_tokens headroom. + """ + # FIXME(lsyin): temporary fix for the context length issue under spec 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 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index a5d00d6a0..8b98cc547 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, Optional +from typing import TYPE_CHECKING import numpy as np import torch @@ -10,34 +10,18 @@ from sglang.kernels.ops.memory.common import ( _get_last_loc_safe_kernel as _get_last_loc_safe_kernel, ) from sglang.kernels.ops.memory.common import get_last_loc_kernel as get_last_loc_kernel -from sglang.kernels.ops.memory.common import ( - get_last_loc_triton, - get_last_loc_triton_safe, - write_req_to_token_pool_triton, -) from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( maybe_evict_dsv4_state_on_swa, - maybe_write_dsv4_decode, - maybe_write_dsv4_extend, ) from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool from sglang.srt.runtime_context import get_server_args -from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import is_cuda, is_hip, is_npu, support_triton -from sglang.srt.utils.common import ceil_align, is_pin_memory_available - -_is_npu = is_npu() - -_is_hip = is_hip() - -_is_cuda = is_cuda() +from sglang.srt.utils.common import ceil_align if TYPE_CHECKING: - from sglang.srt.managers.schedule_batch import Req, ScheduleBatch + from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator - from sglang.srt.model_executor.forward_batch_info import DSV4StateLens # Needs 2 + 1 slots for mamba request with prefix cache. 2 for ping pong cache, 1 for running mamba state. MAMBA_STATE_PER_REQ_PREFIX_CACHE = 3 @@ -118,179 +102,6 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): tree_cache.cache_unfinished_req(req, **kwargs) -def write_cache_indices( - out_cache_loc: torch.Tensor, - req_pool_indices_tensor: torch.Tensor, - req_pool_indices_cpu: torch.Tensor, - prefix_lens_tensor: torch.Tensor, - prefix_lens_cpu: torch.Tensor, - seq_lens_tensor: torch.Tensor, - seq_lens_cpu: torch.Tensor, - extend_lens_tensor: torch.Tensor, - extend_lens_cpu: torch.Tensor, - prefix_tensors: list[torch.Tensor], - req_to_token_pool: ReqToTokenPool, -): - if support_triton(get_server_args().attention_backend): - prefix_pointers = torch.tensor( - [t.data_ptr() for t in prefix_tensors], - dtype=torch.uint64, - pin_memory=is_pin_memory_available(req_to_token_pool.device), - ).to(req_to_token_pool.device, non_blocking=True) - # TODO: some tensors can be reused for ForwardBatchInfo (e.g., extend_lens, cumsum_start) - write_req_to_token_pool_triton[(req_pool_indices_tensor.shape[0],)]( - req_to_token_pool.req_to_token, - req_pool_indices_tensor, - prefix_pointers, - prefix_lens_tensor, - seq_lens_tensor, - extend_lens_tensor, - out_cache_loc, - req_to_token_pool.req_to_token.shape[1], - ) - else: - pt = 0 - for i in range(req_pool_indices_cpu.shape[0]): - req_idx = req_pool_indices_cpu[i].item() - prefix_len = prefix_lens_cpu[i].item() - seq_len = seq_lens_cpu[i].item() - extend_len = extend_lens_cpu[i].item() - - req_to_token_pool.write( - (req_idx, slice(0, prefix_len)), - prefix_tensors[i], - ) - req_to_token_pool.write( - (req_idx, slice(prefix_len, seq_len)), - out_cache_loc[pt : pt + extend_len], - ) - pt += extend_len - - -def get_last_loc( - req_to_token: torch.Tensor, - req_pool_indices_tensor: torch.Tensor, - prefix_lens_tensor: torch.Tensor, -) -> torch.Tensor: - attn_backend = get_server_args().attention_backend - uses_triton_dispatch = attn_backend not in ("ascend", "torch_native") - - if _is_hip and uses_triton_dispatch: - # HIP-only: the legacy get_last_loc_triton kernel emits a - # mixed-width int32->int64 store that Triton mis-compiles on HIP, - # producing out-of-range last_loc values under EAGLE + - # page_size>1 (e.g. with aiter unified attention or the triton - # attention backend). The bug is in the Triton HIP codegen, not - # in any particular attention backend, so route every HIP path - # that would otherwise use get_last_loc_triton through the - # int32-safe variant. Non-HIP hardware keeps the original - # dispatcher below. - return get_last_loc_triton_safe( - req_to_token, req_pool_indices_tensor, prefix_lens_tensor - ) - - if uses_triton_dispatch: - impl = get_last_loc_triton - else: - impl = get_last_loc_torch - - return impl(req_to_token, req_pool_indices_tensor, prefix_lens_tensor) - - -def get_last_loc_torch( - req_to_token: torch.Tensor, - req_pool_indices_tensor: torch.Tensor, - prefix_lens_tensor: torch.Tensor, -) -> torch.Tensor: - return torch.where( - prefix_lens_tensor > 0, - req_to_token[req_pool_indices_tensor, prefix_lens_tensor - 1], - torch.full_like(prefix_lens_tensor, -1), - ) - - -def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int: - if server_args is None: - server_args = get_server_args() - - if server_args.speculative_algorithm is None: - return 1 - - # Spec decoding allocates max(topk * num_steps, num_draft_tokens) per decode step. - 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 - - from sglang.srt.speculative.spec_info import SpeculativeAlgorithm - - spec_algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm) - if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv(): - return max(spec_steps * spec_topk, spec_tokens) - else: - # spec v2 tree (page>1, topk>1): worst-case page-aligned footprint per - # topk branch is ceil((page_size-1 + num_steps) / page) pages, each branch - # duplicated -- 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_utils.eagle_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; the spec v2 page>1 topk>1 holey - draft footprint can outgrow the default num_draft_tokens headroom. - """ - # FIXME(lsyin): temporary fix for the context length issue under spec 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, - backup_state: bool = False, -): - allocator = tree_cache.token_to_kv_pool_allocator - evict_from_tree_cache(tree_cache, num_tokens) - - state = None - if backup_state: - state = allocator.backup_state() - - out_cache_loc = allocator.alloc(num_tokens) - - if out_cache_loc is None: - error_msg = ( - f"Out of memory. Try to lower your batch size.\n" - f"Try to allocate {num_tokens} tokens.\n" - f"{available_and_evictable_str(tree_cache)}" - ) - logger.error(error_msg) - if tree_cache is not None: - tree_cache.pretty_print() - raise RuntimeError(error_msg) - - return (out_cache_loc, state) if backup_state else out_cache_loc - - def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if tree_cache is None: return @@ -317,318 +128,6 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): tree_cache.evict(EvictParams(num_tokens=num_tokens)) -def _compute_dsv4_state_lens(batch, *, is_decode: bool): - """Per-req c{4,128}_state pool alloc lens (``DSV4StateLens``) for this step. - None on CUDA / non-V4 paths (allocator has no ``compute_dsv4_state_lens_*``). - """ - allocator = batch.token_to_kv_pool_allocator - if not hasattr(allocator, "compute_dsv4_state_lens_extend"): - return None - if is_decode: - return allocator.compute_dsv4_state_lens_decode(batch.reqs) - return allocator.compute_dsv4_state_lens_extend( - batch.reqs, batch.seq_lens_cpu.tolist() - ) - - -def alloc_paged_token_slots_extend( - tree_cache: BasePrefixCache, - prefix_lens: torch.Tensor, - prefix_lens_cpu: torch.Tensor, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, - extend_num_tokens: int, - backup_state: bool = False, - req_pool_indices: Optional[torch.Tensor] = None, - dsv4_state_lens: Optional[DSV4StateLens] = None, - batch=None, -): - # Over estimate the number of tokens: assume each request needs a new page. - allocator = tree_cache.token_to_kv_pool_allocator - num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size - evict_from_tree_cache(tree_cache, num_tokens) - - state = None - if backup_state: - state = allocator.backup_state() - - is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c4_attn_allocator") - extra_alloc_kwargs = {} - if is_dsv4: - extra_alloc_kwargs["req_pool_indices"] = req_pool_indices - # Per-call per-req tables for the c-pool / state last_loc lookup. - if batch is not None: - extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool - if dsv4_state_lens is not None: - extra_alloc_kwargs["dsv4_state_lens"] = dsv4_state_lens - - out = allocator.alloc_extend( - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - last_loc, - extend_num_tokens, - **extra_alloc_kwargs, - ) - - if is_dsv4: - bundle = out - out_cache_loc = None if bundle is None else bundle.out_full_loc - if batch is not None: - batch.out_cache_loc_dsv4 = bundle - else: - out_cache_loc = out - - if out_cache_loc is None: - error_msg = ( - f"Prefill out of memory. Try to lower your batch size.\n" - f"Try to allocate {extend_num_tokens} tokens.\n" - f"{available_and_evictable_str(tree_cache)}" - ) - logger.error(error_msg) - if tree_cache is not None: - tree_cache.pretty_print() - raise RuntimeError(error_msg) - - return (out_cache_loc, state) if backup_state else out_cache_loc - - -def alloc_req_slots( - req_to_token_pool: ReqToTokenPool, - reqs: list[Req], - tree_cache: BasePrefixCache | None, -) -> list[int]: - """Allocate request slots from the pool. - - Fail-loud: raises ``RuntimeError`` if the pool can't satisfy the batch. An - alloc failure here means the admission budget (``PrefillAdder``) was wrong - and should surface rather than be masked. - """ - num_reqs = len(reqs) - if isinstance(req_to_token_pool, HybridReqToTokenPool): - # Byte-coordinated for the shared allocator (accounts for the peer full - # sub-pool's bytes); plain slot free count for the non-shared one. - mamba_available_size = ( - req_to_token_pool.mamba_allocator.schedulable_available_size() - ) - # Eviction headroom factor: 3x (or lazy variant) for radix COW, 1x for chunk. - if tree_cache.supports_mamba(): - factor = ( - MAMBA_STATE_PER_REQ_PREFIX_CACHE_LAZY - if req_to_token_pool.enable_mamba_extra_buffer_lazy - else MAMBA_STATE_PER_REQ_PREFIX_CACHE - ) - else: - factor = MAMBA_STATE_PER_REQ_NO_CACHE - mamba_state_needed = num_reqs * factor - if mamba_available_size < mamba_state_needed: - if tree_cache is not None and tree_cache.supports_mamba(): - mamba_num = max(0, mamba_state_needed - mamba_available_size) - tree_cache.evict(EvictParams(num_tokens=0, mamba_num=mamba_num)) - req_pool_indices = req_to_token_pool.alloc(reqs) - if req_pool_indices is None: - raise RuntimeError( - "alloc_req_slots runs out of memory. " - "Please set a smaller number for `--max-running-requests`. " - f"{req_to_token_pool.available_size()=}, {num_reqs=}, " - ) - return req_pool_indices - - -def _alloc_page_size(batch: ScheduleBatch) -> int: - # DCP swaps in an allocator whose page_size is server_args.page_size * - # dcp_size, so it can be > 1 even when tree_cache.page_size is 1; branch on - # the real allocator's page_size there. Elsewhere the two are equal. - if (_is_hip or _is_cuda) and get_server_args().dcp_size > 1: - return batch.tree_cache.token_to_kv_pool_allocator.page_size - return batch.tree_cache.page_size - - -def alloc_for_extend( - batch: ScheduleBatch, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Allocate KV cache for extend batch and write to req_to_token_pool. - - Returns ``(out_cache_loc, req_pool_indices_device, req_pool_indices_cpu)`` - (the last is the host/CPU mirror). ``alloc_req_slots`` raises ``RuntimeError`` - if the pool can't satisfy the batch (fail-loud — see its docstring). - """ - # free out-of-window swa tokens - batch.maybe_evict_swa() - - prefix_tensors = [r.prefix_indices for r in batch.reqs] - - # Create tensors for allocation - prefix_lens_cpu = torch.tensor(batch.prefix_lens, dtype=torch.int64) - extend_lens_cpu = torch.tensor(batch.extend_lens, dtype=torch.int64) - prefix_lens_device = prefix_lens_cpu.to(batch.device, non_blocking=True) - extend_lens_device = extend_lens_cpu.to(batch.device, non_blocking=True) - - # Allocate req slots (raises RuntimeError if the pool is exhausted) - req_pool_indices = alloc_req_slots( - batch.req_to_token_pool, batch.reqs, batch.tree_cache - ) - req_pool_indices_cpu = torch.tensor(req_pool_indices, dtype=torch.int64) - req_pool_indices_device = req_pool_indices_cpu.to(batch.device, non_blocking=True) - - # Allocate KV cache (throws exception on failure) - if _alloc_page_size(batch) == 1: - out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens) - else: - # Paged allocation - build last_loc - last_loc = [ - (t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device)) - for t in prefix_tensors - ] - out_cache_loc = alloc_paged_token_slots_extend( - tree_cache=batch.tree_cache, - prefix_lens=prefix_lens_device, - prefix_lens_cpu=prefix_lens_cpu, - seq_lens=batch.seq_lens, - seq_lens_cpu=batch.seq_lens_cpu, - last_loc=torch.cat(last_loc), - extend_num_tokens=batch.extend_num_tokens, - req_pool_indices=req_pool_indices_device, - dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=False), - batch=batch, - ) - - # Write to req_to_token_pool - write_cache_indices( - out_cache_loc, - req_pool_indices_device, - req_pool_indices_cpu, - prefix_lens_device, - prefix_lens_cpu, - batch.seq_lens, - batch.seq_lens_cpu, - extend_lens_device, - extend_lens_cpu, - prefix_tensors, - batch.req_to_token_pool, - ) - - # DSV4-NPU hook: no-op on non-DSV4 paths. - if _is_npu: - maybe_write_dsv4_extend( - batch, - req_pool_indices_cpu, - prefix_lens_cpu, - batch.seq_lens_cpu, - ) - - return out_cache_loc, req_pool_indices_device, req_pool_indices_cpu - - -def alloc_paged_token_slots_decode( - tree_cache: BasePrefixCache, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, - token_per_req: int = 1, - req_pool_indices: Optional[torch.Tensor] = None, - dsv4_state_lens: Optional[DSV4StateLens] = None, - batch=None, -) -> torch.Tensor: - """Allocate paged KV cache for decode batch.""" - allocator = tree_cache.token_to_kv_pool_allocator - # Over estimate the number of tokens: assume each request needs a new page. - num_tokens = len(seq_lens) * allocator.page_size - evict_from_tree_cache(tree_cache, num_tokens) - - # DSV4-NPU allocator also needs req_pool_indices + per-req state lens and - # returns a DSV4OutCacheLoc bundle; hasattr-gated so others stay unchanged. - is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c4_attn_allocator") - extra_alloc_kwargs = {} - if is_dsv4: - extra_alloc_kwargs["req_pool_indices"] = req_pool_indices - # Per-call per-req tables for the last_loc lookup. - if batch is not None: - extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool - if dsv4_state_lens is not None: - extra_alloc_kwargs["dsv4_state_lens"] = dsv4_state_lens - - out = allocator.alloc_decode(seq_lens, seq_lens_cpu, last_loc, **extra_alloc_kwargs) - - if is_dsv4: - bundle = out - out_cache_loc = None if bundle is None else bundle.out_full_loc - if batch is not None: - batch.out_cache_loc_dsv4 = bundle - else: - out_cache_loc = out - - if out_cache_loc is None: - error_msg = ( - f"Decode out of memory. Try to lower your batch size.\n" - f"Try to allocate {len(seq_lens) * token_per_req} tokens.\n" - f"{available_and_evictable_str(tree_cache)}" - ) - logger.error(error_msg) - if tree_cache is not None: - tree_cache.pretty_print() - raise RuntimeError(error_msg) - - return out_cache_loc - - -def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor: - """ - Allocate KV cache for decode batch and write to req_to_token_pool. - - Returns: - out_cache_loc: allocated cache locations - """ - - batch.maybe_evict_swa() - - seq_lens_gpu = batch.seq_lens - bs = seq_lens_gpu.shape[0] - - if _alloc_page_size(batch) == 1: - # Non-paged allocation - out_cache_loc = alloc_token_slots(batch.tree_cache, bs * token_per_req) - else: - # Paged allocation - last_loc = batch.req_to_token_pool.req_to_token[ - batch.req_pool_indices, seq_lens_gpu - 1 - ] - seq_lens_next = seq_lens_gpu + token_per_req - out_cache_loc = alloc_paged_token_slots_decode( - tree_cache=batch.tree_cache, - seq_lens=seq_lens_next, - seq_lens_cpu=batch.seq_lens_cpu + token_per_req, - last_loc=last_loc, - token_per_req=token_per_req, - req_pool_indices=batch.req_pool_indices, - dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=True), - batch=batch, - ) - - # Write to req_to_token_pool - if batch.model_config.is_encoder_decoder: - locs = batch.encoder_lens + seq_lens_gpu - else: - locs = seq_lens_gpu.clone() - - batch.req_to_token_pool.write( - (batch.req_pool_indices, locs), out_cache_loc.to(torch.int32) - ) - - # DSV4-NPU hook: no-op on non-DSV4 paths. - if _is_npu: - maybe_write_dsv4_decode( - batch, - batch.seq_lens_cpu + token_per_req, - token_per_req, - ) - - return out_cache_loc - - def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True): # the two resources currently have the same lifecycle, thus simplify logic below assert (req.req_pool_idx is None) == (req.kv is None) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index b0047f885..c5cf3a993 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -21,6 +21,7 @@ from sglang.srt.configs.model_config import ( ) from sglang.srt.distributed.parallel_state import get_world_group from sglang.srt.environ import envs +from sglang.srt.mem_cache.allocation_sizing import get_req_to_token_extra_context_len from sglang.srt.mem_cache.allocator import ( BaseTokenToKVPoolAllocator, PagedTokenToKVPoolAllocator, @@ -34,7 +35,6 @@ from sglang.srt.mem_cache.allocator.swa import ( PureSWATokenToKVPoolAllocator, 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 HiSparseDSATokenToKVPool from sglang.srt.mem_cache.memory_pool import ( diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 81d0c9eab..58ac874b7 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -30,7 +30,7 @@ from sglang.srt.configs.model_config import ( is_minimax_sparse, ) from sglang.srt.environ import envs -from sglang.srt.mem_cache.common import get_alloc_len_per_decode +from sglang.srt.mem_cache.allocation_sizing import get_alloc_len_per_decode from sglang.srt.mem_cache.deepseek_v4_memory_pool import get_compress_state_ring_size from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool from sglang.srt.runtime_context import get_parallel diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 4c966a0db..5f4db3149 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/python/sglang/srt/speculative/dflash_info_v2.py @@ -8,7 +8,7 @@ import torch from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ScheduleBatch -from sglang.srt.mem_cache.common import ( +from sglang.srt.mem_cache.allocation import ( alloc_paged_token_slots_extend, alloc_token_slots, get_last_loc, diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index dad18a053..a0fb64938 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -18,12 +18,12 @@ from sglang.srt.hardware_backend.npu.dsv4.dsv4_allocator import ( from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( maybe_build_dsv4_verify_bundle, ) -from sglang.srt.mem_cache.common import ( +from sglang.srt.mem_cache.allocation import ( alloc_paged_token_slots_extend, alloc_token_slots, - get_alloc_reserve_per_decode, get_last_loc, ) +from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import ( is_cpu, diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index 191d441ea..b41de7fca 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -54,8 +54,8 @@ _OWNER_SITES = { (_SB, "ScheduleBatch.prepare_for_extend", "extend_batch_idx"): 1, (_SB, "ScheduleBatch.prepare_for_extend", "kv_committed_len"): 1, (_SB, "ScheduleBatch.prepare_for_extend", "kv_allocated_len"): 1, - ("mem_cache/common.py", "alloc_for_extend", "evict"): 1, - ("mem_cache/common.py", "alloc_for_decode", "evict"): 1, + ("mem_cache/allocation.py", "alloc_for_extend", "evict"): 1, + ("mem_cache/allocation.py", "alloc_for_decode", "evict"): 1, # spec v2: no pre-claim; resolve commits the full accepted run uniformly. (*_EAGLE_DECODE, "decode_batch_idx"): 1, (*_EAGLE_DECODE, "evict"): 1,