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

Support swa allocator page size > 1 (#16296)

Browse Source
This commit is contained in:
Ke Bao
2026-01-03 08:49:47 +08:00
committed by GitHub
parent 31ed68e7c1
commit 7c1b4b1c4c
4 changed files with 205 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

168
python/sglang/srt/mem_cache/allocator.py
View File

@@ -20,6 +20,7 @@ Page-aligned memory pool.
"""
import abc
import weakref
from typing import TYPE_CHECKING
import torch
@@ -179,39 +180,77 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self,
size: int,
size_swa: int,
page_size: int,
dtype: torch.dtype,
device: str,
kvcache: SWAKVPool,
need_sort: bool,
):
super().__init__(size, 1, dtype, device, kvcache, need_sort)
assert isinstance(kvcache, SWAKVPool)
self._size_full = size
self._size_swa = size_swa
self.full_attn_allocator = TokenToKVPoolAllocator(
size,
dtype,
device,
kvcache.full_kv_pool,
need_sort,
)
self.swa_attn_allocator = TokenToKVPoolAllocator(
size_swa,
dtype,
device,
kvcache.swa_kv_pool,
need_sort,
)
self.full_to_swa_index_mapping = torch.empty(
size + size_swa + 1,
dtype=torch.int64,
device=device,
)
self.clear()
self.dtype = dtype
self.device = device
self.page_size = page_size
self._kvcache.full_to_swa_index_mapping = self.full_to_swa_index_mapping
if page_size == 1:
self.full_attn_allocator = TokenToKVPoolAllocator(
size,
dtype,
device,
kvcache.full_kv_pool,
need_sort,
)
self.swa_attn_allocator = TokenToKVPoolAllocator(
size_swa,
dtype,
device,
kvcache.swa_kv_pool,
need_sort,
)
else:
self.full_attn_allocator = PagedTokenToKVPoolAllocator(
size,
page_size,
dtype,
device,
kvcache.full_kv_pool,
need_sort,
)
self.swa_attn_allocator = PagedTokenToKVPoolAllocator(
size_swa,
page_size,
dtype,
device,
kvcache.swa_kv_pool,
need_sort,
)
# Note: append one more item of value -1 in the end so -1 maps to -1.
# It is needed for the last_loc in alloc_extend, where the first full_last_loc
# is -1, and we need to map it to swa_last_loc -1 as well.
self.full_to_swa_index_mapping = torch.cat(
[
torch.zeros(
size + self.page_size,
dtype=torch.int64,
device=device,
),
torch.tensor([-1], dtype=torch.int64, device=device),
]
)
self.need_sort = need_sort
self.free_pages = None
self.release_pages = None
self.is_not_in_free_group = True
self.free_group = []
self.clear()
self._kvcache = kvcache
self._kvcache.register_mapping(weakref.proxy(self.full_to_swa_index_mapping))
def available_size(self):
# Note: use full_available_size() and swa_available_size() instead.
raise NotImplementedError()
def full_available_size(self):
@@ -221,13 +260,17 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
return self.swa_attn_allocator.available_size()
@property
def size_full(self):
return self._size_full
def size(self):
return min(self._size_full, self._size_swa)
@property
def size_swa(self):
return self._size_swa
@property
def size_full(self):
return self._size_full
def debug_print(self) -> str:
msg = ""
msg += f"#swa-available-size: {self.swa_attn_allocator.available_size()}, "
@@ -240,10 +283,11 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
return self._kvcache
def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
assert self.full_to_swa_index_mapping is not None
return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
assert self._kvcache.full_to_swa_index_mapping is not None
return self._kvcache.translate_loc_from_full_to_swa(kv_indices)
def alloc(self, need_size: int):
assert self.page_size == 1
if need_size > self.full_attn_allocator.available_size():
return None
if need_size > self.swa_attn_allocator.available_size():
@@ -251,12 +295,81 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
alloc_full_indices = self.full_attn_allocator.alloc(need_size)
alloc_swa_indices = self.swa_attn_allocator.alloc(need_size)
assert alloc_full_indices is not None
assert alloc_swa_indices is not None
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
return alloc_full_indices
def alloc_extend(
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,
):
assert self.page_size > 1
num_tokens = extend_num_tokens + len(seq_lens) * self.page_size
if num_tokens > self.full_attn_allocator.available_size():
return None
if num_tokens > self.swa_attn_allocator.available_size():
return None
swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
alloc_full_indices = self.full_attn_allocator.alloc_extend(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
last_loc,
extend_num_tokens,
)
alloc_swa_indices = self.swa_attn_allocator.alloc_extend(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
swa_last_loc,
extend_num_tokens,
)
assert alloc_full_indices is not None
assert alloc_swa_indices is not None
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
return alloc_full_indices
def alloc_decode(
self,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor, # last_loc for full layers
):
assert self.page_size > 1
swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
alloc_full_indices = self.full_attn_allocator.alloc_decode(
seq_lens, seq_lens_cpu, last_loc
)
alloc_swa_indices = self.swa_attn_allocator.alloc_decode(
seq_lens, seq_lens_cpu, swa_last_loc
)
if alloc_full_indices is None or alloc_swa_indices is None:
return None
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
return alloc_full_indices
def free(self, free_index: torch.Tensor):
if free_index.numel() == 0:
return
# NOTE: the API is not idempotent.
if self.is_not_in_free_group:
self.full_attn_allocator.free(free_index)
self.free_swa(free_index)
@@ -287,7 +400,8 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def clear(self):
self.swa_attn_allocator.clear()
self.full_attn_allocator.clear()
self.full_to_swa_index_mapping.fill_(0)
# Note: the last item is -1, we don't clear it, see the comment in __init__
self.full_to_swa_index_mapping[:-1].fill_(0)
self.is_not_in_free_group = True
self.free_group = []

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

@@ -1284,6 +1284,7 @@ class SWAKVPool(KVCache):
self,
size: int,
size_swa: int,
page_size: int,
dtype: torch.dtype,
head_num: int,
head_dim: int,
@@ -1303,9 +1304,10 @@ class SWAKVPool(KVCache):
self.swa_layer_nums = len(swa_attention_layer_ids)
self.full_layer_nums = len(full_attention_layer_ids)
self.start_layer = 0
self.page_size = 1
self.page_size = page_size
self.swa_loc = None
kwargs["page_size"] = 1
kwargs["page_size"] = page_size
kwargs["enable_memory_saver"] = False
kwargs["head_num"] = head_num
kwargs["head_dim"] = head_dim
@@ -1347,6 +1349,9 @@ class SWAKVPool(KVCache):
f"SWAKVPool mem usage: {self.mem_usage:.2f} GB, swa size: {self.size_swa}, full size: {self.size}"
)
def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
self.full_to_swa_index_mapping = full_to_swa_index_mapping
def get_kv_size_bytes(self):
k_size, v_size = self.full_kv_pool.get_kv_size_bytes()
k_size_swa, v_size_swa = self.swa_kv_pool.get_kv_size_bytes()
@@ -1356,12 +1361,11 @@ class SWAKVPool(KVCache):
full_kv_data_ptrs, full_kv_data_lens, full_kv_item_lens = (
self.full_kv_pool.get_contiguous_buf_infos()
)
kv_data_ptrs = full_kv_data_ptrs
kv_data_lens = full_kv_data_lens
kv_item_lens = full_kv_item_lens
return kv_data_ptrs, kv_data_lens, kv_item_lens
return (
full_kv_data_ptrs,
full_kv_data_lens,
full_kv_item_lens,
)
def get_state_buf_infos(self):
swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens = (
@@ -1391,8 +1395,14 @@ class SWAKVPool(KVCache):
else:
return self.full_kv_pool.get_kv_buffer(layer_id_pool)
def set_swa_loc(self, loc: torch.Tensor):
self.swa_loc = loc
def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
assert self.full_to_swa_index_mapping is not None
# Note: kv_indices could have -1 values (from alloc_extend), which will be mapped to -1
# since the last item of full_to_swa_index_mapping is -1.
return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
def set_kv_buffer(
@@ -1408,8 +1418,12 @@ class SWAKVPool(KVCache):
layer_id = layer.layer_id
layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
if is_swa_layer:
if self.full_to_swa_index_mapping is not None:
loc = self.translate_loc_from_full_to_swa(loc)
if self.swa_loc is not None:
loc = self.swa_loc
else:
if self.full_to_swa_index_mapping is not None:
loc = self.translate_loc_from_full_to_swa(loc)
self.swa_kv_pool.set_kv_buffer(
None,
loc,
@@ -1430,6 +1444,12 @@ class SWAKVPool(KVCache):
layer_id_override=layer_id_pool,
)
def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
self.full_kv_pool.move_kv_cache(tgt_loc, src_loc)
tgt_loc_swa = self.translate_loc_from_full_to_swa(tgt_loc)
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 get_cpu_copy(self, indices):
# 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

44
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
View File

@@ -517,6 +517,7 @@ class ModelRunnerKVCacheMixin:
self.token_to_kv_pool = SWAKVPool(
size=self.full_max_total_num_tokens,
size_swa=self.swa_max_total_num_tokens,
page_size=self.page_size,
dtype=self.kv_cache_dtype,
head_num=self.model_config.get_num_kv_heads(
get_attention_tp_size()
@@ -614,17 +615,18 @@ class ModelRunnerKVCacheMixin:
need_sort=need_sort,
)
else:
if self.page_size == 1:
if self.is_hybrid_swa:
self.token_to_kv_pool_allocator = SWATokenToKVPoolAllocator(
self.full_max_total_num_tokens,
self.swa_max_total_num_tokens,
dtype=self.kv_cache_dtype,
device=self.device,
kvcache=self.token_to_kv_pool,
need_sort=need_sort,
)
else:
if self.is_hybrid_swa:
self.token_to_kv_pool_allocator = SWATokenToKVPoolAllocator(
self.full_max_total_num_tokens,
self.swa_max_total_num_tokens,
page_size=self.page_size,
dtype=self.kv_cache_dtype,
device=self.device,
kvcache=self.token_to_kv_pool,
need_sort=need_sort,
)
else:
if self.page_size == 1:
self.token_to_kv_pool_allocator = TokenToKVPoolAllocator(
self.max_total_num_tokens,
dtype=self.kv_cache_dtype,
@@ -632,16 +634,16 @@ class ModelRunnerKVCacheMixin:
kvcache=self.token_to_kv_pool,
need_sort=need_sort,
)
else:
assert not self.is_hybrid_swa
self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
self.max_total_num_tokens,
page_size=self.page_size,
dtype=self.kv_cache_dtype,
device=self.device,
kvcache=self.token_to_kv_pool,
need_sort=need_sort,
)
else:
self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
self.max_total_num_tokens,
page_size=self.page_size,
dtype=self.kv_cache_dtype,
device=self.device,
kvcache=self.token_to_kv_pool,
need_sort=need_sort,
)
else:
assert self.is_draft_worker
if self.is_hybrid_swa: