sglang/python/sglang/srt/mem_cache/common.py

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import torch
import triton
import triton.language as tl

from sglang.srt.mem_cache.base_prefix_cache import (
    BasePrefixCache,
    EvictParams,
    EvictResult,
)
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
    build_in_seq_page_compute_owners,
    get_in_seq_page_compute_owner_unavailable_reason,
)
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import support_triton
from sglang.srt.utils.common import ceil_align

if TYPE_CHECKING:
    from sglang.srt.managers.schedule_batch import Req, ScheduleBatch

# 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
MAMBA_STATE_PER_REQ_NO_CACHE = 1

logger = logging.getLogger(__name__)


class KVCapacityWaitError(RuntimeError):
    """Recoverable KV-capacity pressure with CP owner-lane diagnostics."""

    def __init__(
        self,
        message: str,
        *,
        required_by_owner: list[int] | None = None,
        available_by_owner: list[int] | None = None,
        deficit_by_owner: list[int] | None = None,
    ):
        self.required_by_owner = required_by_owner
        self.available_by_owner = available_by_owner
        self.deficit_by_owner = deficit_by_owner
        super().__init__(
            f"{message}; required_by_owner={required_by_owner} "
            f"available_by_owner={available_by_owner} "
            f"deficit_by_owner={deficit_by_owner}"
        )


def _log_cp_shared_kv_alloc_fallback(
    reason: str,
    message: str,
    *args,
) -> None:
    logger.warning(
        "[CP_SHARED_KV_FALLBACK][compute_owner_alloc] reason=%s " + message,
        reason,
        *args,
    )


@triton.jit
def write_req_to_token_pool_triton(
    req_to_token_ptr,  # [max_batch, max_context_len]
    req_pool_indices,
    prefix_tensors,
    pre_lens,
    seq_lens,
    extend_lens,
    out_cache_loc,
    req_to_token_ptr_stride: tl.constexpr,
):
    BLOCK_SIZE: tl.constexpr = 512
    pid = tl.program_id(0)

    req_pool_index = tl.load(req_pool_indices + pid)
    pre_len = tl.load(pre_lens + pid)
    seq_len = tl.load(seq_lens + pid)
    prefix_tensor = tl.load(prefix_tensors + pid).to(tl.pointer_type(tl.int64))

    # write prefix
    num_loop = tl.cdiv(pre_len, BLOCK_SIZE)
    for i in range(num_loop):
        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
        mask = offset < pre_len
        value = tl.load(prefix_tensor + offset, mask=mask)
        tl.store(
            req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + offset,
            value,
            mask=mask,
        )

    # NOTE: This can be slow for large bs
    cumsum_start = tl.cast(0, tl.int64)
    for i in range(pid):
        cumsum_start += tl.load(extend_lens + i)

    num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE)
    for i in range(num_loop):
        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
        mask = offset < (seq_len - pre_len)
        value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask)
        tl.store(
            req_to_token_ptr
            + req_pool_index * req_to_token_ptr_stride
            + offset
            + pre_len,
            value,
            mask=mask,
        )


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_global_server_args().attention_backend):
        prefix_pointers = torch.tensor(
            [t.data_ptr() for t in prefix_tensors],
            device=req_to_token_pool.device,
            dtype=torch.uint64,
        )
        # 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:
    if (
        get_global_server_args().attention_backend != "ascend"
        and get_global_server_args().attention_backend != "torch_native"
    ):
        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),
    )


@triton.jit
def get_last_loc_kernel(
    req_to_token,
    req_pool_indices_tensor,
    prefix_lens_tensor,
    result,
    num_tokens,
    req_to_token_stride,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE
    mask = offset < num_tokens

    prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0)
    req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0)

    token_mask = prefix_lens > 0
    token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1)
    tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1)

    tl.store(result + offset, tokens, mask=mask)


def get_last_loc_triton(
    req_to_token: torch.Tensor,
    req_pool_indices_tensor: torch.Tensor,
    prefix_lens_tensor: torch.Tensor,
) -> torch.Tensor:
    BLOCK_SIZE = 256
    num_tokens = prefix_lens_tensor.shape[0]
    result = torch.empty_like(prefix_lens_tensor)
    grid = (triton.cdiv(num_tokens, BLOCK_SIZE),)

    get_last_loc_kernel[grid](
        req_to_token,
        req_pool_indices_tensor,
        prefix_lens_tensor,
        result,
        num_tokens,
        req_to_token.stride(0),
        BLOCK_SIZE,
    )
    return result


def alloc_token_slots(
    tree_cache: BasePrefixCache,
    num_tokens: int,
    backup_state: bool = False,
):
    allocator = tree_cache.token_to_kv_pool_allocator
    evict_result = 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)}"
            f"{_evict_result_str(evict_result)}\n"
        )
        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_result_str(evict_result: EvictResult | None) -> str:
    if evict_result is None:
        return "evict_result=None"
    return (
        "evict_result=("
        f"num_tokens_evicted={evict_result.num_tokens_evicted}, "
        f"swa_num_tokens_evicted={evict_result.swa_num_tokens_evicted}, "
        f"mamba_num_evicted={evict_result.mamba_num_evicted})"
    )


def evict_from_tree_cache(
    tree_cache: BasePrefixCache | None, num_tokens: int
) -> EvictResult:
    if tree_cache is None:
        return EvictResult()

    if tree_cache.is_chunk_cache():
        return EvictResult()

    allocator = tree_cache.token_to_kv_pool_allocator

    if isinstance(allocator, SWATokenToKVPoolAllocator):
        # Hybrid allocator
        full_available_size = allocator.full_available_size()
        swa_available_size = allocator.swa_available_size()

        if full_available_size < num_tokens or swa_available_size < num_tokens:
            full_num_tokens = max(0, num_tokens - full_available_size)
            swa_num_tokens = max(0, num_tokens - swa_available_size)
            return tree_cache.evict(
                EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
            )
        return EvictResult()
    else:
        # Standard allocator
        available = allocator.available_size()
        if available < num_tokens:
            logger.debug(
                "[MemCache-evict] evict_from_tree_cache: available=%d < num_tokens=%d deficit=%d, triggering eviction",
                available,
                num_tokens,
                num_tokens - available,
            )
            return tree_cache.evict(EvictParams(num_tokens=num_tokens))
        return EvictResult()


def _evict_for_compute_owner_lanes(
    *,
    tree_cache: BasePrefixCache | None,
    allocator,
    page_compute_owners: list[int],
) -> None:
    if tree_cache is None or tree_cache.is_chunk_cache():
        return

    compute_owner_lane_stats = getattr(allocator, "compute_owner_lane_stats", None)
    if compute_owner_lane_stats is None:
        return

    max_attempts = max(2, min(8, int(getattr(allocator, "cp_size", 1))))
    for attempt in range(max_attempts):
        _required, _available, deficits = compute_owner_lane_stats(page_compute_owners)
        deficit_pages = sum(deficits)
        if deficit_pages <= 0:
            return

        try:
            evictable_size = tree_cache.evictable_size()
        except Exception:
            evictable_size = allocator.page_size
        if isinstance(evictable_size, tuple):
            evictable_size = evictable_size[0]
        if evictable_size <= 0:
            logger.debug(
                "[MemCache-evict] _evict_for_compute_owner_lanes: evictable_size=%d <= 0, giving up",
                evictable_size,
            )
            return

        # ``deficits`` is already expressed in physical owner-lane pages.  The
        # old ``* cp_size`` multiplier turned a one-page lane deficit into one
        # full CP stripe of eviction and caused large L1 churn under HiCache
        # load-back pressure.
        evict_tokens = max(allocator.page_size, deficit_pages * allocator.page_size)
        before_available = allocator.available_size()
        logger.debug(
            "[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d: deficit_pages=%d evict_tokens=%d before_available=%d evictable_size=%d",
            attempt,
            deficit_pages,
            evict_tokens,
            before_available,
            evictable_size,
        )
        evict_result = tree_cache.evict(
            EvictParams(
                num_tokens=evict_tokens,
                owner_lane_deficits=[int(v) for v in deficits],
            )
        )
        after_available = allocator.available_size()
        evicted_tokens = getattr(evict_result, "num_tokens_evicted", 0)
        logger.debug(
            "[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d result: evicted=%d after_available=%d",
            attempt,
            evicted_tokens,
            after_available,
        )
        if after_available <= before_available and evicted_tokens <= 0:
            return


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,
):
    # 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_result = None

    # logger.info(
    #     "[MemCache-alloc] alloc_paged_token_slots_extend: extend_num_tokens=%d batch_size=%d num_tokens=%d page_size=%d "
    #     "available_size=%d evicted=%d",
    #     extend_num_tokens,
    #     len(seq_lens_cpu),
    #     num_tokens,
    #     allocator.page_size,
    #     allocator.available_size(),
    #     getattr(evict_result, "num_tokens_evicted", 0),
    # )

    alloc_extend_compute_owner = getattr(
        allocator, "alloc_extend_compute_owner", None
    )
    page_compute_owners = None
    compute_owner_unavailable_reason = None
    if alloc_extend_compute_owner is not None and len(prefix_lens_cpu) == 1:
        try:
            server_args = get_global_server_args()
        except ValueError:
            server_args = None
        if (
            server_args is not None
            and server_args.enable_nsa_prefill_cp_shared_kv
            and server_args.enable_nsa_prefill_context_parallel
            and server_args.nsa_prefill_cp_mode == "in-seq-split"
        ):
            extend_len = int(seq_lens_cpu[0].item() - prefix_lens_cpu[0].item())
            page_compute_owners = build_in_seq_page_compute_owners(
                extend_len=extend_len,
                extend_prefix_len=int(prefix_lens_cpu[0].item()),
                page_size=int(allocator.page_size),
                cp_size=int(allocator.cp_size),
            )
            if page_compute_owners is None:
                compute_owner_unavailable_reason = (
                    get_in_seq_page_compute_owner_unavailable_reason(
                        extend_len=extend_len,
                        extend_prefix_len=int(prefix_lens_cpu[0].item()),
                        page_size=int(allocator.page_size),
                        cp_size=int(allocator.cp_size),
                    )
                    or "unknown"
                )
        else:
            compute_owner_unavailable_reason = "server_args_not_enabled"
    elif alloc_extend_compute_owner is not None:
        compute_owner_unavailable_reason = (
            "multi_batch" if len(prefix_lens_cpu) != 1 else "unknown"
        )

    if page_compute_owners is not None:
        state = None
        if backup_state:
            state = allocator.backup_state()

        out_cache_loc = alloc_extend_compute_owner(
            prefix_lens,
            prefix_lens_cpu,
            seq_lens,
            seq_lens_cpu,
            last_loc,
            extend_num_tokens,
            page_compute_owners,
        )
        if out_cache_loc is None:
            _evict_for_compute_owner_lanes(
                tree_cache=tree_cache,
                allocator=allocator,
                page_compute_owners=page_compute_owners,
            )
            if backup_state:
                state = allocator.backup_state()
            out_cache_loc = alloc_extend_compute_owner(
                prefix_lens,
                prefix_lens_cpu,
                seq_lens,
                seq_lens_cpu,
                last_loc,
                extend_num_tokens,
                page_compute_owners,
            )
        if out_cache_loc is None:
            required = available = deficits = None
            compute_owner_lane_stats = getattr(
                allocator, "compute_owner_lane_stats", None
            )
            if compute_owner_lane_stats is not None:
                required, available, deficits = compute_owner_lane_stats(
                    page_compute_owners
                )
            logger.warning(
                "[CP_SHARED_KV_FAIL_FAST][owner_lane_exhausted] "
                "failed to allocate pages from compute-owner lanes after eviction; "
                "raising recoverable capacity wait. extend_num_tokens=%s page_size=%s "
                "required_by_owner=%s available_by_owner=%s deficit_by_owner=%s",
                extend_num_tokens,
                allocator.page_size,
                required,
                available,
                deficits,
            )
            raise KVCapacityWaitError(
                "owner_lane_exhausted: failed to allocate pages from CP "
                "compute-owner lanes after owner-lane eviction",
                required_by_owner=required,
                available_by_owner=available,
                deficit_by_owner=deficits,
            )
    else:
        evict_result = evict_from_tree_cache(tree_cache, num_tokens)
        state = None
        if backup_state:
            state = allocator.backup_state()

        if alloc_extend_compute_owner is not None:
            _log_cp_shared_kv_alloc_fallback(
                compute_owner_unavailable_reason or "compute_owner_not_available",
                "page-aligned compute-owner page assignment is unavailable; "
                "falling back to legacy page allocation. batch_size=%s extend_num_tokens=%s page_size=%s reason=%s",
                len(prefix_lens_cpu),
                extend_num_tokens,
                allocator.page_size,
                compute_owner_unavailable_reason or "compute_owner_not_available",
            )
        out_cache_loc = allocator.alloc_extend(
            prefix_lens,
            prefix_lens_cpu,
            seq_lens,
            seq_lens_cpu,
            last_loc,
            extend_num_tokens,
        )

    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)}"
            f"{_evict_result_str(evict_result)}\n"
        )
        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."""
    num_reqs = len(reqs)
    if isinstance(req_to_token_pool, HybridReqToTokenPool):
        mamba_available_size = req_to_token_pool.mamba_pool.available_size()
        factor = (
            MAMBA_STATE_PER_REQ_PREFIX_CACHE
            if tree_cache.supports_mamba()
            else 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()=}, "
            f"{num_reqs=}, "
        )
    return req_pool_indices


def alloc_for_extend(
    batch: ScheduleBatch,
) -> tuple[torch.Tensor, torch.Tensor, list[int]]:
    """
    Allocate KV cache for extend batch and write to req_to_token_pool.

    Returns:
        out_cache_loc: allocated cache locations
        req_pool_indices_device: request pool indices at a device tensor
        req_pool_indices: request pool indices as list
    """
    # 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
    newly_allocated_reqs = [req for req in batch.reqs if req.req_pool_idx is None]
    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)
    try:
        if batch.tree_cache.page_size == 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,
            )
    except KVCapacityWaitError:
        for req in newly_allocated_reqs:
            if req.req_pool_idx is not None:
                batch.req_to_token_pool.free(req)
        raise

    # 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,
    )

    return out_cache_loc, req_pool_indices_device, req_pool_indices


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,
) -> 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)

    out_cache_loc = allocator.alloc_decode(seq_lens, seq_lens_cpu, last_loc)

    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()

    bs = batch.seq_lens.shape[0]

    if batch.tree_cache.page_size == 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, batch.seq_lens - 1
        ]
        seq_lens_next = batch.seq_lens + 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,
        )

    # Write to req_to_token_pool
    if batch.model_config.is_encoder_decoder:
        locs = batch.encoder_lens + batch.seq_lens
    else:
        locs = batch.seq_lens.clone()

    batch.req_to_token_pool.write(
        (batch.req_pool_indices, locs), out_cache_loc.to(torch.int32)
    )

    return out_cache_loc


def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True):
    # MambaRadixCache may alloc mamba state before alloc KV cache
    if req.req_pool_idx is None:
        assert (
            tree_cache.supports_mamba()
        ), "Only MambaRadixCache allow freeing before alloc"
        # TODO (csy, hanming): clean up this early allocation logic
        if req.mamba_pool_idx is not None:
            tree_cache.req_to_token_pool.mamba_pool.free(
                req.mamba_pool_idx.unsqueeze(-1)
            )
            req.mamba_pool_idx = None
        return

    tree_cache.cache_finished_req(req, is_insert=is_insert)

    # FIXME: SessionAwareCache.cache_finished_req sets req_pool_idx = None to
    # transfer KV ownership to the SessionSlot, so we skip the remaining
    # cleanup (overalloc free + pool slot free). This means over-allocated
    # tokens from speculative decoding are NOT freed between turns.
    if req.req_pool_idx is None:
        return

    start_p, end_p = req.pop_overallocated_kv_cache()

    global_server_args = get_global_server_args()
    page_size = global_server_args.page_size
    spec_algo = global_server_args.speculative_algorithm

    if spec_algo is None:
        assert (
            start_p == end_p
        ), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv_allocated_len=}"

    if page_size > 1:
        start_p = ceil_align(start_p, page_size)

    if start_p < end_p:
        indices_to_free = tree_cache.req_to_token_pool.req_to_token[req.req_pool_idx][
            start_p:end_p
        ]
        tree_cache.token_to_kv_pool_allocator.free(indices_to_free)
    # If the prefix cache doesn't manage mamba states, we must free them here.
    if isinstance(tree_cache.req_to_token_pool, HybridReqToTokenPool) and (
        not tree_cache.supports_mamba()
    ):
        assert (
            req.mamba_pool_idx is not None
        ), "mamba state is freed while the tree cache does not manage mamba states"
        tree_cache.req_to_token_pool.free_mamba_cache(req)
    tree_cache.req_to_token_pool.free(req)


def available_and_evictable_str(tree_cache: BasePrefixCache) -> str:
    return tree_cache.available_and_evictable_str()