Optimize SWA memory preallocation for disaggregated decode (#24857)

Co-authored-by: weireweire <weiliangl@nvidia.com>
Co-authored-by: Cheng Wan <chwan@rice.edu>

python/sglang/srt/disaggregation/decode.py

@@ -62,11 +62,13 @@ from sglang.srt.mem_cache.common import (
     page_align_floor,
     release_kv_cache,
 )
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.mem_cache.memory_pool import (
     HybridReqToTokenPool,
     KVCache,
     ReqToTokenPool,
 )
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.observability.req_time_stats import (
     set_schedule_time_batch,
     set_time_batch,
@@ -175,7 +177,6 @@ class DecodeReqToTokenPool:
 
 
 class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
-
     def __init__(
         self,
         size: int,
@@ -316,13 +317,55 @@ class DecodePreallocQueue:
         if self.enable_staging:
             self.transfer_queue._init_staging_handler(self.kv_manager)
 
-        if self.scheduler.tp_worker.is_hybrid_swa:
+        if (
-            # FIXME: current SWA allocation allocate full kv cache size in prefill
+            self.scheduler.tp_worker.is_hybrid_swa
+            and not self._uses_swa_tail_prealloc()
+        ):
+            # Fallback for SWA allocators that still allocate the SWA pool at
+            # full prompt length.
             self.max_total_num_tokens = min(
                 self.max_total_num_tokens,
                 self.scheduler.tp_worker.model_runner.swa_max_total_num_tokens,
             )
 
+    def _uses_swa_tail_prealloc(self) -> bool:
+        return (
+            isinstance(self.token_to_kv_pool, (SWAKVPool, DeepSeekV4TokenToKVPool))
+            and self.token_to_kv_pool_allocator.page_size > 1
+            and hasattr(self.token_to_kv_pool_allocator, "alloc_extend_swa_tail")
+        )
+
+    def _swa_tail_len(self, seq_len: int) -> int:
+        if not self._uses_swa_tail_prealloc() or seq_len <= 0:
+            return max(seq_len, 0)
+
+        window_size = self.scheduler.sliding_window_size
+        if window_size is None or window_size <= 0:
+            return seq_len
+
+        page_size = self.token_to_kv_pool_allocator.page_size
+        window_start = max(0, seq_len - window_size)
+        window_start = (window_start // page_size) * page_size
+        return seq_len - window_start
+
+    def _swa_retractable_len(self, req: Req) -> int:
+        if not self._uses_swa_tail_prealloc():
+            return len(req.origin_input_ids) + len(req.output_ids)
+        return self._swa_tail_len(len(req.origin_input_ids)) + len(req.output_ids)
+
+    def _prealloc_kv_lens(self, req: Req) -> Tuple[int, int]:
+        allocated_kv_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
+        if self._uses_swa_tail_prealloc():
+            return allocated_kv_len, self._swa_tail_len(allocated_kv_len)
+        return allocated_kv_len, allocated_kv_len
+
+    def _prealloc_required_tokens(self, req: Req) -> Tuple[int, int]:
+        full_len, swa_len = self._prealloc_kv_lens(req)
+        return (
+            full_len + self.num_reserved_decode_tokens,
+            swa_len + self.num_reserved_decode_tokens,
+        )
+
     def _init_kv_manager(self) -> CommonKVManager:
         kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS)
         kv_args = kv_args_class()
@@ -486,6 +529,18 @@ class DecodePreallocQueue:
             prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST)
             self.scheduler.stream_output([req], req.return_logprob)
             return True
+        if self._uses_swa_tail_prealloc():
+            _, swa_required = self._prealloc_required_tokens(req)
+            swa_capacity = self.token_to_kv_pool_allocator.size_swa
+            if swa_required > swa_capacity:
+                message = (
+                    f"Request {req.rid} requires too many SWA KV tokens for "
+                    f"decode preallocation: {swa_required} > {swa_capacity}"
+                )
+                logger.error(message)
+                prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST)
+                self.scheduler.stream_output([req], req.return_logprob)
+                return True
         return False
 
     def extend(self, reqs: List[Req], is_retracted: bool = False) -> None:
@@ -501,7 +556,15 @@ class DecodePreallocQueue:
         # allocate memory
         resumed_reqs = []
         indices_to_remove = set()
-        allocatable_tokens = self._allocatable_tokens(count_retracted=False)
+        uses_swa_tail_prealloc = self._uses_swa_tail_prealloc()
+        if uses_swa_tail_prealloc:
+            full_allocatable_tokens, swa_allocatable_tokens = (
+                self._swa_aware_allocatable_token_budgets(count_retracted=False)
+            )
+        else:
+            full_allocatable_tokens = self._allocatable_token_budgets(
+                count_retracted=False
+            )
 
         for i, req in enumerate(self.retracted_queue):
             if rids_to_check is not None and req.rid not in rids_to_check:
@@ -510,19 +573,19 @@ class DecodePreallocQueue:
             if self.req_to_token_pool.available_size() <= 0:
                 break
 
-            required_tokens_for_request = (
+            full_required, swa_required = self._prealloc_required_tokens(req)
-                len(req.origin_input_ids)
+            if full_required > full_allocatable_tokens:
-                + len(req.output_ids)
+                break
-                + self.num_reserved_decode_tokens
+            if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens:
-            )
-            if required_tokens_for_request > allocatable_tokens:
                 break
 
             resumed_reqs.append(req)
             indices_to_remove.add(i)
             req.is_retracted = False
             self._pre_alloc(req)
-            allocatable_tokens -= required_tokens_for_request
+            full_allocatable_tokens -= full_required
+            if uses_swa_tail_prealloc:
+                swa_allocatable_tokens -= swa_required
 
             # load from cpu, release the cpu copy
             req.load_kv_cache(self.req_to_token_pool, self.token_to_kv_pool_allocator)
@@ -677,9 +740,24 @@ class DecodePreallocQueue:
             len(r.origin_input_ids) + len(r.output_ids)
             for r in self.scheduler.running_batch.reqs
         )
-        allocatable_tokens = self._allocatable_tokens(
+        uses_swa_tail_prealloc = self._uses_swa_tail_prealloc()
-            retractable_tokens=retractable_tokens, count_retracted=True
+        swa_allocatable_tokens = 0
-        )
+        if uses_swa_tail_prealloc:
+            retractable_swa_tokens = sum(
+                self._swa_retractable_len(r) for r in self.scheduler.running_batch.reqs
+            )
+            full_allocatable_tokens, swa_allocatable_tokens = (
+                self._swa_aware_allocatable_token_budgets(
+                    retractable_tokens=retractable_tokens,
+                    retractable_swa_tokens=retractable_swa_tokens,
+                    count_retracted=True,
+                )
+            )
+        else:
+            retractable_swa_tokens = 0
+            full_allocatable_tokens = self._allocatable_token_budgets(
+                retractable_tokens=retractable_tokens, count_retracted=True
+            )
         # First, remove all failed requests from the queue
         for i, decode_req in enumerate(self.queue):
             if rids_to_check is not None and decode_req.req.rid not in rids_to_check:
@@ -746,7 +824,7 @@ class DecodePreallocQueue:
                 # Matching may lock previously-evictable radix pages, so refresh
                 # the admission budget against the post-lock pool state before we
                 # decide whether this request still fits.
-                allocatable_tokens = self._allocatable_tokens(
+                full_allocatable_tokens = self._allocatable_token_budgets(
                     retractable_tokens=retractable_tokens,
                     count_retracted=True,
                     extra_reserved_reqs=len(preallocated_reqs),
@@ -771,25 +849,47 @@ class DecodePreallocQueue:
                     )
                     - retractable_tokens,
                 )
-                > allocatable_tokens
+                > full_allocatable_tokens
             ):
                 if prefix_len > 0:
                     self.tree_cache.dec_lock_ref(decode_req.req.last_node)
                 break
-            if required_tokens_for_request > allocatable_tokens:
+            if required_tokens_for_request > full_allocatable_tokens:
                 if prefix_len > 0:
                     self.tree_cache.dec_lock_ref(decode_req.req.last_node)
                 break
 
+            if uses_swa_tail_prealloc:
+                _, swa_required = self._prealloc_required_tokens(decode_req.req)
+                _, swa_len = self._prealloc_kv_lens(decode_req.req)
+                max_new_tokens = min(
+                    decode_req.req.sampling_params.max_new_tokens,
+                    CLIP_MAX_NEW_TOKEN,
+                )
+                if (
+                    max(
+                        swa_required,
+                        swa_len + max_new_tokens - retractable_swa_tokens,
+                    )
+                    > swa_allocatable_tokens
+                ):
+                    if prefix_len > 0:
+                        self.tree_cache.dec_lock_ref(decode_req.req.last_node)
+                    break
+
             dst_kv_indices = self._pre_alloc(decode_req.req, prefix_indices, prefix_len)
             hisparse_req_budget -= 1
             # Recompute from actual pool state for the next queue entry.
             # This accounts for page rounding and newly locked evictable cache.
-            allocatable_tokens = self._allocatable_tokens(
+            full_allocatable_tokens = self._allocatable_token_budgets(
                 retractable_tokens=retractable_tokens,
                 count_retracted=True,
                 extra_reserved_reqs=len(preallocated_reqs) + 1,
             )
+            if uses_swa_tail_prealloc:
+                # SWA budget uses simple decrement (no radix cache eviction in
+                # the SWA pool, so page-rounding drift is negligible).
+                swa_allocatable_tokens -= swa_required
             decode_req.req.cache_protected_len = prefix_len
 
             if self.scheduler.enable_hisparse:
@@ -896,11 +996,8 @@ class DecodePreallocQueue:
             len(decode_req.req.fill_ids) for decode_req in self.transfer_queue.queue
         )
 
-    def _allocatable_tokens(
+    def _need_space_for_single_req(
-        self,
+        self, retractable_tokens: Optional[int] = None
-        retractable_tokens: Optional[int] = None,
-        count_retracted: bool = True,
-        extra_reserved_reqs: int = 0,
     ) -> int:
         need_space_for_single_req = (
             max(
@@ -915,12 +1012,69 @@ class DecodePreallocQueue:
             and len(self.scheduler.running_batch.reqs) > 0
             else 0
         )
+        return need_space_for_single_req
+
+    def _active_req_count(self, extra_reserved_reqs: int = 0) -> int:
+        return (
+            len(self.scheduler.running_batch.reqs)
+            + len(self.transfer_queue.queue)
+            + len(self.scheduler.waiting_queue)
+            + extra_reserved_reqs
+        )
+
+    def _active_reserved_tokens(
+        self, n_active: Optional[int] = None, extra_reserved_reqs: int = 0
+    ) -> int:
+        if n_active is None:
+            n_active = self._active_req_count(extra_reserved_reqs)
+        return self.num_reserved_decode_tokens * n_active
+
+    def _swa_aware_allocatable_token_budgets(
+        self,
+        retractable_tokens: Optional[int] = None,
+        retractable_swa_tokens: Optional[int] = None,
+        count_retracted: bool = True,
+    ) -> Tuple[int, int]:
+        n_active = self._active_req_count()
+        reserved_tokens = self._active_reserved_tokens(n_active)
+
+        full_allocatable_tokens = self._allocatable_token_budgets(
+            retractable_tokens=retractable_tokens,
+            count_retracted=count_retracted,
+            reserved_tokens=reserved_tokens,
+        )
+
+        return full_allocatable_tokens, self._swa_tail_allocatable_token_budget(
+            retractable_tokens=retractable_tokens,
+            retractable_swa_tokens=retractable_swa_tokens,
+            count_retracted=count_retracted,
+            n_active=n_active,
+            reserved_tokens=reserved_tokens,
+        )
+
+    def _allocatable_token_budgets(
+        self,
+        retractable_tokens: Optional[int] = None,
+        count_retracted: bool = True,
+        extra_reserved_reqs: int = 0,
+        reserved_tokens: Optional[int] = None,
+    ) -> int:
+        need_space_for_single_req = self._need_space_for_single_req(retractable_tokens)
+        if reserved_tokens is None:
+            reserved_tokens = self._active_reserved_tokens(
+                extra_reserved_reqs=extra_reserved_reqs
+            )
+
         if self.scheduler.enable_hisparse:
             # HiSparse pre-alloc only allocates logical indices (alloc_logical_only),
             # so the logical pool is the binding constraint for admission control.
             available_size = (
                 self.token_to_kv_pool_allocator.logical_attn_allocator.available_size()
             )
+        elif self._uses_swa_tail_prealloc():
+            available_size = self.token_to_kv_pool_allocator.full_available_size()
+            if self.scheduler.server_args.disaggregation_decode_enable_radix_cache:
+                available_size += self.tree_cache.evictable_size()
         else:
             available_size = self.token_to_kv_pool_allocator.available_size()
             # Include evictable decode-radix cache entries in the budget -- they
@@ -928,16 +1082,7 @@ class DecodePreallocQueue:
             if self.scheduler.server_args.disaggregation_decode_enable_radix_cache:
                 available_size += self.tree_cache.evictable_size()
         allocatable_tokens = available_size - max(
-            # preserve some space for future decode
+            reserved_tokens, need_space_for_single_req
-            self.num_reserved_decode_tokens
-            * (
-                len(self.scheduler.running_batch.reqs)
-                + len(self.transfer_queue.queue)
-                + len(self.scheduler.waiting_queue)
-                + extra_reserved_reqs
-            ),
-            # make sure each request can finish if reach max_tokens with all other requests retracted
-            need_space_for_single_req,
         )
 
         # Note: if the last prebuilt extend just finishes, and we enter `pop_preallocated` immediately in the next iteration
@@ -951,16 +1096,75 @@ class DecodePreallocQueue:
             )
 
         if count_retracted:
-            allocatable_tokens -= sum(
+            for req in self.retracted_queue:
-                [
+                full_required, _ = self._prealloc_required_tokens(req)
-                    len(req.origin_input_ids)
+                allocatable_tokens -= full_required
-                    + len(req.output_ids)
+
-                    + self.num_reserved_decode_tokens
-                    for req in self.retracted_queue
-                ]
-            )
         return allocatable_tokens
 
+    def _swa_tail_allocatable_token_budget(
+        self,
+        retractable_tokens: Optional[int] = None,
+        retractable_swa_tokens: Optional[int] = None,
+        count_retracted: bool = True,
+        n_active: Optional[int] = None,
+        reserved_tokens: Optional[int] = None,
+    ) -> int:
+        need_swa_space_for_single_req = self._need_space_for_single_req(
+            retractable_tokens
+        )
+        if (
+            retractable_swa_tokens is not None
+            and len(self.scheduler.running_batch.reqs) > 0
+        ):
+            need_swa_space_for_single_req = max(
+                self._swa_tail_len(len(x.origin_input_ids))
+                + min(x.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKEN)
+                - retractable_swa_tokens
+                for x in self.scheduler.running_batch.reqs
+            )
+
+        if n_active is None:
+            n_active = self._active_req_count()
+        if reserved_tokens is None:
+            reserved_tokens = self._active_reserved_tokens(n_active)
+
+        # SWA growth is bounded by the sliding window: once a req's SWA
+        # footprint reaches `sliding_window_size`, further decode tokens
+        # evict old ones and net growth is zero. The linear reservation
+        # `num_reserved_decode_tokens * n_active` (correct for the full
+        # pool) over-reserves SWA in steady state. Cap by the actual
+        # remaining headroom up to per-req window cap.
+        window_size = self.scheduler.sliding_window_size or 0
+        swa_total = self.token_to_kv_pool_allocator.size_swa
+        swa_used = swa_total - self.token_to_kv_pool_allocator.swa_available_size()
+        swa_growth_potential = max(0, n_active * window_size - swa_used)
+        swa_reserved_tokens = min(reserved_tokens, swa_growth_potential)
+        swa_allocatable_tokens = (
+            self.token_to_kv_pool_allocator.swa_available_size()
+            - max(swa_reserved_tokens, need_swa_space_for_single_req)
+        )
+
+        # Note: if the last prebuilt extend just finishes, and we enter `pop_preallocated` immediately in the next iteration
+        #       the extend batch is not in any queue, so we need to explicitly add the tokens slots here
+        if (
+            self.scheduler.last_batch
+            and self.scheduler.last_batch.forward_mode.is_prebuilt()
+        ):
+            prebuilt_reserved_tokens = self.num_reserved_decode_tokens * len(
+                self.scheduler.last_batch.reqs
+            )
+            prebuilt_n = len(self.scheduler.last_batch.reqs)
+            prebuilt_swa_growth = max(0, prebuilt_n * window_size - swa_used)
+            swa_allocatable_tokens -= min(prebuilt_reserved_tokens, prebuilt_swa_growth)
+
+        if count_retracted:
+            for req in self.retracted_queue:
+                _, swa_required = self._prealloc_required_tokens(req)
+                swa_allocatable_tokens -= swa_required
+
+        return swa_allocatable_tokens
+
     def _required_alloc_tokens(self, *, fill_len: int, prefix_len: int) -> int:
         page_size = self.token_to_kv_pool_allocator.page_size
         if page_size == 1:
@@ -1062,16 +1266,31 @@ class DecodePreallocQueue:
                 if prefix_len > 0
                 else torch.tensor([-1], dtype=torch.int64, device=device)
             )
-            kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
+            if self._uses_swa_tail_prealloc() and prefix_len == 0:
-                prefix_lens=torch.tensor(
+                # Tail-only SWA allocation: only valid when prefix_len == 0.
-                    [prefix_len], dtype=torch.int64, device=device
+                # When prefix_len > 0 (radix cache hit), we fall back to
-                ),
+                # alloc_extend which allocates SWA at full page count; the
-                prefix_lens_cpu=torch.tensor([prefix_len], dtype=torch.int64),
+                # SWA budget in that case may slightly under-estimate.
-                seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
+                kv_loc = self.token_to_kv_pool_allocator.alloc_extend_swa_tail(
-                seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
+                    prefix_lens=torch.tensor([0], dtype=torch.int64, device=device),
-                last_loc=last_loc,
+                    prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
-                extend_num_tokens=delta_len,
+                    seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
-            )
+                    seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
+                    last_loc=last_loc,
+                    extend_num_tokens=fill_len,
+                    swa_tail_len=self._swa_tail_len(fill_len),
+                )
+            else:
+                kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
+                    prefix_lens=torch.tensor(
+                        [prefix_len], dtype=torch.int64, device=device
+                    ),
+                    prefix_lens_cpu=torch.tensor([prefix_len], dtype=torch.int64),
+                    seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
+                    seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
+                    last_loc=last_loc,
+                    extend_num_tokens=delta_len,
+                )
 
         assert kv_loc is not None, (
             f"KV cache is full! Bug in memory estimation. "
@@ -1333,7 +1552,6 @@ class DecodeTransferQueue:
 
 
 class SchedulerDisaggregationDecodeMixin:
-
     @torch.no_grad()
     def event_loop_normal_disagg_decode(self: Scheduler):
         """A normal scheduler loop for decode worker in disaggregation mode."""

python/sglang/srt/mem_cache/swa_memory_pool.py

@@ -214,20 +214,55 @@ class SWAKVPool(BaseSWAKVPool):
         src_loc_swa = self.translate_loc_from_full_to_swa(src_loc)
         self.swa_kv_pool.move_kv_cache(tgt_loc_swa, src_loc_swa)
 
+    def _filter_swa_cpu_copy(self, swa_kv_cpu, row_mask: torch.Tensor):
+        if swa_kv_cpu is None:
+            return None
+        if row_mask is None or bool(torch.all(row_mask).item()):
+            return swa_kv_cpu
+
+        chunk_size = getattr(
+            self.swa_kv_pool, "cpu_offloading_chunk_size", len(row_mask)
+        )
+        filtered = []
+        for layer_chunks in swa_kv_cpu:
+            if len(layer_chunks) == 0:
+                filtered.append([])
+                continue
+
+            k_cpu = torch.cat([chunk[0] for chunk in layer_chunks], dim=0)
+            v_cpu = torch.cat([chunk[1] for chunk in layer_chunks], dim=0)
+            k_cpu = k_cpu[row_mask]
+            v_cpu = v_cpu[row_mask]
+
+            filtered_layer = []
+            for i in range(0, len(k_cpu), chunk_size):
+                filtered_layer.append(
+                    [k_cpu[i : i + chunk_size], v_cpu[i : i + chunk_size]]
+                )
+            filtered.append(filtered_layer)
+        return filtered
+
     def get_cpu_copy(self, indices, mamba_indices=None):
         # For SWA, we need to copy KV cache from both full and SWA pools
         # The indices are for the full pool, and we use mapping to get SWA indices
         full_kv_cpu = self.full_kv_pool.get_cpu_copy(indices)
 
-        # Get SWA indices through the mapping
+        swa_mask = None
-        # Note: SWA allocation always creates 1:1 mapping, so no need to filter
         if self.full_to_swa_index_mapping is not None:
             swa_indices = self.full_to_swa_index_mapping[indices]
-            swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices)
+            # Slot 0 is reserved as a dummy slot. Tail-only SWA allocations leave
+            # the out-of-window full KV indices unmapped, so only copy mapped SWA
+            # tokens and keep their positions for load_cpu_copy().
+            swa_mask = swa_indices > 0
+            if torch.any(swa_mask):
+                swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices[swa_mask])
+                swa_mask = swa_mask.cpu()
+            else:
+                swa_kv_cpu = None
         else:
             swa_kv_cpu = None
 
-        return {"full": full_kv_cpu, "swa": swa_kv_cpu}
+        return {"full": full_kv_cpu, "swa": swa_kv_cpu, "swa_mask": swa_mask}
 
     def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
         # Load KV cache back from CPU to both full and SWA pools
@@ -241,6 +276,20 @@ class SWAKVPool(BaseSWAKVPool):
         # Load SWA KV cache if it exists
         if swa_kv_cpu is not None and self.full_to_swa_index_mapping is not None:
             swa_indices = self.full_to_swa_index_mapping[indices]
+            new_swa_mask = swa_indices > 0
+            old_swa_mask = kv_cache_cpu.get("swa_mask")
+            if old_swa_mask is not None:
+                old_swa_mask = old_swa_mask.to(indices.device)
+                row_mask = new_swa_mask[old_swa_mask].cpu()
+                swa_indices = swa_indices[old_swa_mask][row_mask.to(indices.device)]
+            else:
+                row_mask = new_swa_mask.cpu()
+                swa_indices = swa_indices[new_swa_mask]
+
+            if swa_indices.numel() == 0:
+                return
+
+            swa_kv_cpu = self._filter_swa_cpu_copy(swa_kv_cpu, row_mask)
             self.swa_kv_pool.load_cpu_copy(swa_kv_cpu, swa_indices)
 
 
@@ -437,6 +486,71 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
 
         return alloc_full_indices
 
+    def alloc_extend_swa_tail(
+        self,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,  # last_loc for full layers
+        extend_num_tokens: int,
+        swa_tail_len: int,
+    ):
+        """Allocate full KV for the whole extend and SWA KV only for the tail.
+
+        This is used by disaggregated decode preallocation: decode receives full
+        prompt KV for full-attention layers, but only the sliding-window state is
+        transferred for SWA layers.
+        """
+        assert self.page_size > 1
+        assert len(seq_lens_cpu) == 1, "SWA tail allocation currently supports bs=1"
+        assert len(prefix_lens_cpu) == 1
+        assert 0 <= swa_tail_len <= extend_num_tokens
+
+        num_full_pages = get_num_new_pages(
+            seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu
+        )
+        num_swa_pages = (swa_tail_len + self.page_size - 1) // self.page_size
+        if num_full_pages > self.full_attn_allocator.available_size() // self.page_size:
+            return None
+        if num_swa_pages > self.swa_attn_allocator.available_size() // self.page_size:
+            return None
+
+        alloc_full_indices = self.full_attn_allocator.alloc_extend(
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc,
+            extend_num_tokens,
+        )
+        assert alloc_full_indices is not None
+
+        if swa_tail_len == 0:
+            return alloc_full_indices
+
+        device = self.device
+        swa_prefix_lens = torch.zeros((1,), dtype=torch.int64, device=device)
+        swa_prefix_lens_cpu = torch.zeros((1,), dtype=torch.int64)
+        swa_seq_lens = torch.tensor([swa_tail_len], dtype=torch.int64, device=device)
+        swa_seq_lens_cpu = torch.tensor([swa_tail_len], dtype=torch.int64)
+        swa_last_loc = torch.tensor([-1], dtype=torch.int64, device=device)
+
+        alloc_swa_indices = self.swa_attn_allocator.alloc_extend(
+            swa_prefix_lens,
+            swa_prefix_lens_cpu,
+            swa_seq_lens,
+            swa_seq_lens_cpu,
+            swa_last_loc,
+            swa_tail_len,
+        )
+        assert alloc_swa_indices is not None
+
+        self.full_to_swa_index_mapping[alloc_full_indices[-swa_tail_len:]] = (
+            alloc_swa_indices
+        )
+        return alloc_full_indices
+
     def alloc_decode(
         self,
         seq_lens: torch.Tensor,

test/registered/unit/mem_cache/test_decode_radix_lock_ref.py

@@ -319,6 +319,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
         queue.req_to_token_pool.available_size.return_value = 1
         queue.req_to_metadata_buffer_idx_allocator = MagicMock()
         queue.req_to_metadata_buffer_idx_allocator.available_size.return_value = 1
+        queue.token_to_kv_pool = MagicMock()
         queue.token_to_kv_pool_allocator = MagicMock()
         queue.token_to_kv_pool_allocator.page_size = 4
 
@@ -336,7 +337,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
         queue.scheduler = scheduler
 
         # Initial budget says the request fits; post-lock budget says it does not.
-        queue._allocatable_tokens = MagicMock(side_effect=[8, 3])
+        queue._allocatable_token_budgets = MagicMock(side_effect=[8, 3])
 
         preallocated, failed = queue.pop_preallocated()
 
@@ -344,7 +345,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
         self.assertEqual(failed, [])
         queue._pre_alloc.assert_not_called()
         queue.tree_cache.dec_lock_ref.assert_called_once_with(req.last_node)
-        self.assertEqual(queue._allocatable_tokens.call_count, 2)
+        self.assertEqual(queue._allocatable_token_budgets.call_count, 2)
 
     def test_repeated_incremental_no_leak(self):
         """Multiple incremental transfers shouldn't leak lock_refs."""