perf(disaggregation): reuse req pool freelists and alloc_extend tensors

python/sglang/srt/disaggregation/decode.py

@@ -125,7 +125,7 @@ class DecodeReqToTokenPool:
                 device=device,
             )
 
-        self.free_slots = list(range(size + pre_alloc_size))
+        self.free_slots = deque(range(size + pre_alloc_size))
 
     def write(self, indices, values):
         self.req_to_token[indices] = values
@@ -147,8 +147,7 @@ class DecodeReqToTokenPool:
         need_size = len(reqs) - len(reusing)
         if need_size > len(self.free_slots):
             return None
-        select_index = self.free_slots[:need_size]
+        select_index = [self.free_slots.popleft() for _ in range(need_size)]
-        self.free_slots = self.free_slots[need_size:]
         offset = 0
         for r in reqs:
             if r.req_pool_idx is None:
@@ -162,7 +161,7 @@ class DecodeReqToTokenPool:
         req.req_pool_idx = None
 
     def clear(self):
-        self.free_slots = list(range(self.size + self.pre_alloc_size))
+        self.free_slots = deque(range(self.size + self.pre_alloc_size))
 
 
 class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
@@ -204,7 +203,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
         )
 
     def clear(self):
-        self.free_slots = list(range(self.size + self.pre_alloc_size))
+        self.free_slots = deque(range(self.size + self.pre_alloc_size))
         self.mamba_pool.clear()
 
 
@@ -275,6 +274,11 @@ class DecodePreallocQueue:
         self._ensure_last_attempt_time: Dict[str, float] = {}
         self._ensure_retry_interval: float = 1.0  # seconds
         self.kv_manager = self._init_kv_manager()
+        self._alloc_extend_prefix_lens: Optional[torch.Tensor] = None
+        self._alloc_extend_prefix_lens_cpu: Optional[torch.Tensor] = None
+        self._alloc_extend_seq_lens: Optional[torch.Tensor] = None
+        self._alloc_extend_seq_lens_cpu: Optional[torch.Tensor] = None
+        self._alloc_extend_last_loc: Optional[torch.Tensor] = None
 
         if self.scheduler.tp_worker.is_hybrid_swa:
             # FIXME: current SWA allocation allocate full kv cache size in prefill
@@ -788,6 +792,50 @@ class DecodePreallocQueue:
             )
         return allocatable_tokens
 
+    def _get_alloc_extend_args(
+        self, fill_len: int
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        device = self.token_to_kv_pool_allocator.device
+
+        if self._alloc_extend_prefix_lens is None:
+            self._alloc_extend_prefix_lens = torch.zeros(
+                1, dtype=torch.int64, device=device
+            )
+            self._alloc_extend_prefix_lens_cpu = torch.zeros(1, dtype=torch.int64)
+            self._alloc_extend_seq_lens = torch.empty(
+                1, dtype=torch.int64, device=device
+            )
+            self._alloc_extend_seq_lens_cpu = torch.empty(1, dtype=torch.int64)
+            self._alloc_extend_last_loc = torch.empty(
+                1, dtype=torch.int64, device=device
+            )
+
+        prefix_lens = self._alloc_extend_prefix_lens
+        prefix_lens_cpu = self._alloc_extend_prefix_lens_cpu
+        seq_lens = self._alloc_extend_seq_lens
+        seq_lens_cpu = self._alloc_extend_seq_lens_cpu
+        last_loc = self._alloc_extend_last_loc
+
+        assert prefix_lens is not None
+        assert prefix_lens_cpu is not None
+        assert seq_lens is not None
+        assert seq_lens_cpu is not None
+        assert last_loc is not None
+
+        prefix_lens.zero_()
+        prefix_lens_cpu.zero_()
+        seq_lens.fill_(fill_len)
+        seq_lens_cpu.fill_(fill_len)
+        last_loc.fill_(-1)
+
+        return (
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc,
+        )
+
     def _pre_alloc(self, req: Req) -> torch.Tensor:
         """Pre-allocate the memory for req_to_token and token_kv_pool"""
         req_pool_indices = self.req_to_token_pool.alloc([req])
@@ -803,13 +851,19 @@ class DecodePreallocQueue:
         if self.token_to_kv_pool_allocator.page_size == 1:
             kv_loc = self.token_to_kv_pool_allocator.alloc(fill_len)
         else:
-            device = self.token_to_kv_pool_allocator.device
+            (
+                prefix_lens,
+                prefix_lens_cpu,
+                seq_lens,
+                seq_lens_cpu,
+                last_loc,
+            ) = self._get_alloc_extend_args(fill_len)
             kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
-                prefix_lens=torch.tensor([0], dtype=torch.int64, device=device),
+                prefix_lens=prefix_lens,
-                prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
+                prefix_lens_cpu=prefix_lens_cpu,
-                seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
+                seq_lens=seq_lens,
-                seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
+                seq_lens_cpu=seq_lens_cpu,
-                last_loc=torch.tensor([-1], dtype=torch.int64, device=device),
+                last_loc=last_loc,
                 extend_num_tokens=fill_len,
             )
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -27,6 +27,7 @@ KVCache actually holds the physical kv cache.
 import abc
 import dataclasses
 import logging
+from collections import deque
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass, fields
 from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
@@ -148,7 +149,7 @@ class ReqToTokenPool:
             self.req_to_token = torch.zeros(
                 (size, max_context_len), dtype=torch.int32, device=device
             )
-        self.free_slots = list(range(size))
+        self.free_slots = deque(range(size))
 
     def write(self, indices, values):
         self.req_to_token[indices] = values
@@ -173,8 +174,7 @@ class ReqToTokenPool:
         need_size = len(reqs) - len(reusing)
         if need_size > len(self.free_slots):
             return None
-        select_index = self.free_slots[:need_size]
+        select_index = [self.free_slots.popleft() for _ in range(need_size)]
-        self.free_slots = self.free_slots[need_size:]
         offset = 0
         for r in reqs:
             if r.req_pool_idx is None:
@@ -188,7 +188,7 @@ class ReqToTokenPool:
         req.req_pool_idx = None
 
     def clear(self):
-        self.free_slots = list(range(self.size))
+        self.free_slots = deque(range(self.size))
 
 
 class MambaPool:
@@ -248,10 +248,13 @@ class MambaPool:
             maybe_init_custom_mem_pool(device=self.device)
         )
 
-        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE), (
+        with (
-            torch.cuda.use_mem_pool(self.custom_mem_pool)
+            self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE),
-            if self.enable_custom_mem_pool
+            (
-            else nullcontext()
+                torch.cuda.use_mem_pool(self.custom_mem_pool)
+                if self.enable_custom_mem_pool
+                else nullcontext()
+            ),
         ):
             conv_state = [
                 torch.zeros(
@@ -531,9 +534,9 @@ class HybridReqToTokenPool(ReqToTokenPool):
                 mid = req.mamba_pool_idx
             else:
                 mid = self.mamba_pool.alloc(1)
-                assert (
+                assert mid is not None, (
-                    mid is not None
+                    f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size. {mid=}, {self.mamba_pool.size=}, {self.mamba_pool.available_size()=}, {len(reqs)=}"
-                ), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size. {mid=}, {self.mamba_pool.size=}, {self.mamba_pool.available_size()=}, {len(reqs)=}"
+                )
                 mid = mid[0]
                 req.mamba_pool_idx = mid
             mamba_indices.append(mid)
@@ -542,18 +545,18 @@ class HybridReqToTokenPool(ReqToTokenPool):
                     req.mamba_ping_pong_track_buffer = self.mamba_pool.alloc(
                         self.mamba_ping_pong_track_buffer_size
                     )
-                    assert (
+                    assert req.mamba_ping_pong_track_buffer is not None, (
-                        req.mamba_ping_pong_track_buffer is not None
+                        "Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio."
-                    ), "Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio."
+                    )
                     req.mamba_next_track_idx = 0
                 mamba_ping_pong_track_buffers.append(req.mamba_ping_pong_track_buffer)
-        assert len(select_index) == len(
+        assert len(select_index) == len(mamba_indices), (
-            mamba_indices
+            f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size."
-        ), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size."
+        )
         if self.enable_mamba_extra_buffer:
-            assert len(select_index) == len(
+            assert len(select_index) == len(mamba_ping_pong_track_buffers), (
-                mamba_ping_pong_track_buffers
+                f"Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio."
-            ), f"Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio."
+            )
         mamba_index_tensor = torch.stack(mamba_indices).to(dtype=torch.int32)
         self.req_index_to_mamba_index_mapping[select_index] = mamba_index_tensor
         if self.enable_mamba_extra_buffer:
@@ -597,7 +600,9 @@ class HybridReqToTokenPool(ReqToTokenPool):
                 assert mamba_ping_pong_track_buffer_to_keep in [
                     0,
                     1,
-                ], f"mamba_ping_pong_track_buffer_to_keep must be 0 or 1, {mamba_ping_pong_track_buffer_to_keep=}"
+                ], (
+                    f"mamba_ping_pong_track_buffer_to_keep must be 0 or 1, {mamba_ping_pong_track_buffer_to_keep=}"
+                )
                 # Avoid Python-list advanced indexing on a device tensor.
                 # The ping-pong buffer size is either 2 (normal) or 1 (spec decode).
                 if self.mamba_ping_pong_track_buffer_size == 2:
@@ -728,7 +733,6 @@ class KVCache(abc.ABC):
 
 
 class MHATokenToKVPool(KVCache):
-
     def __init__(
         self,
         size: int,
@@ -763,7 +767,9 @@ class MHATokenToKVPool(KVCache):
         self.v_head_dim = (
             swa_v_head_dim
             if swa_v_head_dim is not None
-            else v_head_dim if v_head_dim is not None else head_dim
+            else v_head_dim
+            if v_head_dim is not None
+            else head_dim
         )
 
         self._create_buffers()
@@ -1029,9 +1035,9 @@ class MHATokenToKVPool(KVCache):
         if N == 0:
             return
 
-        assert (
+        assert self._kv_copy_config is not None, (
-            self._kv_copy_config is not None
+            "KV copy not initialized. Set enable_kv_cache_copy=True in __init__"
-        ), "KV copy not initialized. Set enable_kv_cache_copy=True in __init__"
+        )
 
         cfg = self._kv_copy_config
         cap = int(cfg.get("num_locs_upper", 256))
@@ -1071,7 +1077,6 @@ class MHATokenToKVPool(KVCache):
 
 
 class MHATokenToKVPoolFP4(MHATokenToKVPool):
-
     def _create_buffers(self):
         with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
             with (
@@ -1247,7 +1252,6 @@ class HybridLinearKVPool(KVCache):
         assert not enable_kvcache_transpose
         self.use_mla = use_mla
         if not use_mla:
-
             TokenToKVPoolClass = MHATokenToKVPool
 
             if _is_npu:
@@ -1268,7 +1272,6 @@ class HybridLinearKVPool(KVCache):
                 enable_memory_saver=enable_memory_saver,
             )
         else:
-
             TokenToKVPoolClass = MLATokenToKVPool
 
             if _is_npu:
@@ -1632,7 +1635,6 @@ class MLATokenToKVPool(KVCache):
 
 
 class MLATokenToKVPoolFP4(MLATokenToKVPool):
-
     def _create_buffers(self):
         with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
             with (

test/registered/unit/mem_cache/test_req_to_token_pool.py

@@ -0,0 +1,167 @@
+"""Unit tests for req-pool and decode prealloc optimizations."""
+
+import unittest
+from collections import deque
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from sglang.srt.disaggregation.decode import DecodePreallocQueue, DecodeReqToTokenPool
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=6, suite="stage-a-test-cpu")
+
+
+class TestReqToTokenPool(CustomTestCase):
+    def _make_req(self, *, req_pool_idx=None, is_chunked=0, kv_committed_len=0):
+        req = MagicMock(spec=Req)
+        req.req_pool_idx = req_pool_idx
+        req.is_chunked = is_chunked
+        req.kv_committed_len = kv_committed_len
+        return req
+
+    def test_req_to_token_pool_uses_deque_backed_free_slots(self):
+        pool = ReqToTokenPool(
+            size=4,
+            max_context_len=8,
+            device="cpu",
+            enable_memory_saver=False,
+        )
+
+        self.assertIsInstance(pool.free_slots, deque)
+
+        req = self._make_req()
+        self.assertEqual(pool.alloc([req]), [0])
+        pool.free(req)
+        pool.clear()
+
+        self.assertIsInstance(pool.free_slots, deque)
+        self.assertEqual(list(pool.free_slots), [0, 1, 2, 3])
+
+    def test_decode_req_to_token_pool_preserves_reuse_and_uses_deque(self):
+        pool = DecodeReqToTokenPool(
+            size=2,
+            max_context_len=8,
+            device="cpu",
+            enable_memory_saver=False,
+            pre_alloc_size=2,
+        )
+
+        self.assertIsInstance(pool.free_slots, deque)
+        self.assertEqual(pool.available_size(), 4)
+
+        reused_req = self._make_req(is_chunked=1)
+        self.assertEqual(pool.alloc([reused_req]), [0])
+
+        fresh_req = self._make_req()
+        self.assertEqual(pool.alloc([fresh_req, reused_req]), [1, 0])
+
+        pool.clear()
+
+        self.assertIsInstance(pool.free_slots, deque)
+        self.assertEqual(list(pool.free_slots), [0, 1, 2, 3])
+
+
+class TestDecodePreallocQueue(CustomTestCase):
+    def _make_req(self, origin_input_ids, output_ids):
+        req = MagicMock(spec=Req)
+        req.req_pool_idx = None
+        req.origin_input_ids = origin_input_ids
+        req.output_ids = output_ids
+        req.kv_allocated_len = 0
+        req.kv_committed_len = 0
+        req.is_chunked = 0
+        req.set_extend_input_len = MagicMock()
+        return req
+
+    def _build_prealloc_queue(self):
+        req_to_token_pool = DecodeReqToTokenPool(
+            size=4,
+            max_context_len=16,
+            device="cpu",
+            enable_memory_saver=False,
+            pre_alloc_size=2,
+        )
+
+        allocator_calls = []
+
+        def alloc_extend(**kwargs):
+            allocator_calls.append(kwargs)
+            return torch.arange(kwargs["extend_num_tokens"], dtype=torch.int64)
+
+        token_to_kv_pool_allocator = MagicMock()
+        token_to_kv_pool_allocator.page_size = 16
+        token_to_kv_pool_allocator.device = "cpu"
+        token_to_kv_pool_allocator.alloc_extend.side_effect = alloc_extend
+        token_to_kv_pool_allocator.get_kvcache.return_value = MagicMock()
+
+        scheduler = SimpleNamespace(tp_worker=SimpleNamespace(is_hybrid_swa=False))
+
+        with (
+            patch.object(
+                DecodePreallocQueue, "_init_kv_manager", return_value=MagicMock()
+            ),
+            patch(
+                "sglang.srt.disaggregation.decode.is_mla_backend", return_value=False
+            ),
+        ):
+            queue = DecodePreallocQueue(
+                req_to_token_pool=req_to_token_pool,
+                token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+                draft_token_to_kv_pool=None,
+                req_to_metadata_buffer_idx_allocator=MagicMock(),
+                metadata_buffers=MagicMock(),
+                scheduler=scheduler,
+                transfer_queue=SimpleNamespace(queue=[]),
+                tree_cache=MagicMock(),
+                gloo_group=MagicMock(),
+                tp_rank=0,
+                tp_size=1,
+                dp_size=1,
+                gpu_id=0,
+                bootstrap_port=1234,
+                max_total_num_tokens=256,
+                pp_rank=0,
+                num_reserved_decode_tokens=1,
+                transfer_backend=MagicMock(),
+            )
+
+        return queue, allocator_calls
+
+    def test_pre_alloc_reuses_alloc_extend_tensors(self):
+        queue, allocator_calls = self._build_prealloc_queue()
+
+        req1 = self._make_req(origin_input_ids=[1, 2, 3], output_ids=[4, 5])
+        req2 = self._make_req(origin_input_ids=[1, 2, 3, 4], output_ids=[5, 6])
+
+        queue._pre_alloc(req1)
+        queue._pre_alloc(req2)
+
+        self.assertEqual(len(allocator_calls), 2)
+        self.assertIs(
+            allocator_calls[0]["prefix_lens"], allocator_calls[1]["prefix_lens"]
+        )
+        self.assertIs(
+            allocator_calls[0]["prefix_lens_cpu"],
+            allocator_calls[1]["prefix_lens_cpu"],
+        )
+        self.assertIs(allocator_calls[0]["seq_lens"], allocator_calls[1]["seq_lens"])
+        self.assertIs(
+            allocator_calls[0]["seq_lens_cpu"], allocator_calls[1]["seq_lens_cpu"]
+        )
+        self.assertIs(allocator_calls[0]["last_loc"], allocator_calls[1]["last_loc"])
+
+        self.assertEqual(allocator_calls[1]["prefix_lens"].item(), 0)
+        self.assertEqual(allocator_calls[1]["prefix_lens_cpu"].item(), 0)
+        self.assertEqual(allocator_calls[1]["seq_lens"].item(), 5)
+        self.assertEqual(allocator_calls[1]["seq_lens_cpu"].item(), 5)
+        self.assertEqual(allocator_calls[1]["last_loc"].item(), -1)
+        self.assertEqual(allocator_calls[1]["extend_num_tokens"], 5)
+
+
+if __name__ == "__main__":
+    unittest.main()