Lightweight extract allocation logic from mem_cache/common.py to more clearly show nearly parallel variants (#29431)

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 (

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:

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

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

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

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.common import ceil_align
-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()
 
 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)

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 (

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

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,

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,

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/allocation.py", "alloc_for_extend", "evict"): 1,
-    ("mem_cache/common.py", "alloc_for_decode", "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,