sgl-project--sglang/python/sglang/srt/mem_cache/pure_swa_radix_cache.py

"""Radix cache for all-SWA models (every layer is sliding-window attention)."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import torch

from sglang.srt.mem_cache.base_prefix_cache import (
    EvictParams,
    EvictResult,
    InsertParams,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey

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

logger = logging.getLogger(__name__)


class PureSWARadixCache(RadixCache):
    """Radix cache for all-SWA models (no full attention layers).

    Extends RadixCache with SWA semantics. Only caches the prefill portion
    [0, evict_floor) on request completion. Window-range KV is freed.
    No tombstone mechanism needed.
    """

    def __init__(self, params: CacheInitParams):
        super().__init__(params)
        self.sliding_window_size = params.sliding_window_size

    def supports_swa(self) -> bool:
        assert (
            self.sliding_window_size is not None
        ), "sliding_window_size must be set for PureSWARadixCache"
        return True

    def swa_evictable_size(self):
        return self.evictable_size_

    def swa_protected_size(self):
        return self.protected_size_

    def full_evictable_size(self):
        return 0

    def full_protected_size(self):
        return 0

    def sanity_check(self):
        """No-op: PureSWARadixCache uses RadixCache's simple tree structure
        which doesn't need the dual-LRU sanity checks of SWARadixCache."""
        pass

    def evict(self, params: EvictParams) -> EvictResult:
        """For all-SWA models, evict_from_tree_cache passes swa_num_tokens
        (with num_tokens=0). Use whichever is non-zero."""
        num_tokens = max(params.num_tokens, params.swa_num_tokens)
        return super().evict(EvictParams(num_tokens=num_tokens))

    def cache_finished_req(self, req: Req, is_insert: bool = True):
        """Cache request when it finishes.

        Only inserts the prefill portion [0, evict_floor) into the radix tree.
        The window portion [swa_evicted_seqlen, committed_len) is freed back
        to the allocator. The range [evict_floor, swa_evicted_seqlen) was already
        freed by _evict_swa during decode — we skip it to avoid double-free.
        """
        if self.disable_finished_insert:
            is_insert = False

        kv_committed_len = req.pop_committed_kv_cache()
        if self.disable:
            kv_indices = self.req_to_token_pool.req_to_token[
                req.req_pool_idx, :kv_committed_len
            ]
            self.token_to_kv_pool_allocator.free(kv_indices)
            return

        token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
        kv_indices = self.req_to_token_pool.req_to_token[
            req.req_pool_idx, :kv_committed_len
        ]

        radix_key = RadixKey(
            token_ids, req.extra_key, is_bigram=self.is_eagle
        ).page_aligned(self.page_size)
        keys_len = len(radix_key)

        old_prefix_len = req.cache_protected_len
        swa_evict_floor = req.swa_evict_floor
        swa_evicted_seqlen = req.swa_evicted_seqlen

        if self.page_size > 1 and swa_evict_floor > 0:
            swa_evict_floor = -(-swa_evict_floor // self.page_size) * self.page_size

        if swa_evict_floor > 0:
            insert_end = min(swa_evict_floor, keys_len)
        else:
            insert_end = keys_len

        if is_insert and insert_end > 0:
            insert_values = kv_indices[:insert_end].to(dtype=torch.int64, copy=True)
            result = self.insert(
                InsertParams(key=radix_key[:insert_end], value=insert_values)
            )
            new_prefix_len = result.prefix_len
            if new_prefix_len > old_prefix_len:
                self.token_to_kv_pool_allocator.free(
                    kv_indices[old_prefix_len:new_prefix_len]
                )
            alive_start = max(swa_evicted_seqlen, insert_end)
            if alive_start < keys_len:
                self.token_to_kv_pool_allocator.free(kv_indices[alive_start:keys_len])
        else:
            free_end = (
                min(swa_evict_floor, keys_len) if swa_evict_floor > 0 else keys_len
            )
            if free_end > old_prefix_len:
                self.token_to_kv_pool_allocator.free(
                    kv_indices[old_prefix_len:free_end]
                )
            alive_start = max(swa_evicted_seqlen, old_prefix_len)
            if swa_evicted_seqlen > 0 and alive_start < keys_len:
                self.token_to_kv_pool_allocator.free(kv_indices[alive_start:keys_len])

        self.token_to_kv_pool_allocator.free(kv_indices[keys_len:])

        if req.last_node is not None:
            self.dec_lock_ref(req.last_node)

    def cache_unfinished_req(self, req: Req, chunked=False):
        """During chunked prefill, swa_evicted_seqlen is 0 and no SWA eviction
        has happened yet, so standard RadixCache logic is correct."""
        super().cache_unfinished_req(req, chunked=chunked)

    def available_and_evictable_str(self) -> str:
        allocator = self.token_to_kv_pool_allocator
        swa_available = allocator.swa_available_size()
        swa_evictable = self.swa_evictable_size()
        return (
            f"SWA available tokens: {swa_available + swa_evictable} "
            f"({swa_available=} + {swa_evictable=})\n"
        )