From 53a04a9a979bc44341dcfd773eff444e95cd7cd5 Mon Sep 17 00:00:00 2001 From: wxiwnd Date: Wed, 8 Apr 2026 20:12:47 +0800 Subject: [PATCH] perf(disaggregation): reuse req pool freelists and alloc_extend tensors --- python/sglang/srt/disaggregation/decode.py | 76 ++++++-- python/sglang/srt/mem_cache/memory_pool.py | 62 +++---- .../unit/mem_cache/test_req_to_token_pool.py | 167 ++++++++++++++++++ 3 files changed, 264 insertions(+), 41 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_req_to_token_pool.py diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index b78154453..90ea6117c 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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] - self.free_slots = self.free_slots[need_size:] + select_index = [self.free_slots.popleft() for _ in range(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_cpu=torch.tensor([0], 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=torch.tensor([-1], dtype=torch.int64, device=device), + prefix_lens=prefix_lens, + prefix_lens_cpu=prefix_lens_cpu, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + last_loc=last_loc, extend_num_tokens=fill_len, ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 51530e865..8a61142bf 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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] - self.free_slots = self.free_slots[need_size:] + select_index = [self.free_slots.popleft() for _ in range(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), ( - torch.cuda.use_mem_pool(self.custom_mem_pool) - if self.enable_custom_mem_pool - else nullcontext() + with ( + self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE), + ( + 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 ( - 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)=}" + assert 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)=}" + ) 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 ( - req.mamba_ping_pong_track_buffer is not None - ), "Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio." + assert req.mamba_ping_pong_track_buffer is not None, ( + "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( - mamba_indices - ), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size." + assert len(select_index) == len(mamba_indices), ( + 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( - mamba_ping_pong_track_buffers - ), f"Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio." + assert len(select_index) == len(mamba_ping_pong_track_buffers), ( + 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 ( - self._kv_copy_config is not None - ), "KV copy not initialized. Set enable_kv_cache_copy=True in __init__" + assert self._kv_copy_config is not None, ( + "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 ( diff --git a/test/registered/unit/mem_cache/test_req_to_token_pool.py b/test/registered/unit/mem_cache/test_req_to_token_pool.py new file mode 100644 index 000000000..752eb0801 --- /dev/null +++ b/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()