chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
wxiwnd
2026-04-08 20:12:47 +08:00
parent a371eacd86
commit 53a04a9a97
3 changed files with 264 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
python/sglang/srt/disaggregation/decode.py
View File

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

62
python/sglang/srt/mem_cache/memory_pool.py
View File

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