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

Deepseek V4 (#23882)

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: fzyzcjy <ch271828n@outlook.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
Co-authored-by: yueming-yuan <yym022502@gmail.com>
Co-authored-by: DarkSharpness <2040703891@qq.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: yhyang201 <yhyang201@users.noreply.github.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>
Co-authored-by: Qiaolin Yu <90088090+qiaolin-yu@users.noreply.github.com>
Co-authored-by: Ethan (Yusheng) Su <11704492+yushengsu-thu@users.noreply.github.com>
Co-authored-by: Mingyi <27337995+wisclmy0611@users.noreply.github.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: Yihao Wang <42559837+againstentropy@users.noreply.github.com>
This commit is contained in:
Liangsheng Yin
2026-05-07 18:32:21 -07:00
committed by GitHub
co-authored by Baizhou Zhang Claude Opus 4.7 fzyzcjy ispobock Zhiqiang Xie yueming-yuan DarkSharpness Yuhao Yang yhyang201 yhyang201 Qiaolin Yu Ethan Su Mingyi Cheng Wan Yihao Wang
parent 55224fff08
commit 35870d55ac
154 changed files with 24970 additions and 759 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_swa_eviction_boundary.py
+1
View File
@@ -110,6 +110,7 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
extra_key=None,
last_node=tree.root_node,
swa_uuid_for_lock=None,
swa_prefix_lock_released=False,
prefix_indices=torch.tensor([], dtype=torch.int64, device=tree.device),
_kv_committed_len=len(token_ids),
)
test/registered/unit/mem_cache/test_swa_unittest.py
+213 -71
View File
@@ -2,7 +2,9 @@ import unittest
import torch
from sglang.srt.environ import envs
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams,
EvictResult,
InsertParams,
@@ -16,80 +18,113 @@ from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllo
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=9, suite="stage-b-test-1-gpu-large")
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
class _DummyReq:
def __init__(self):
self._kv_committed_len = 0
self.swa_prefix_lock_released = False
def pop_committed_kv_cache(self):
return self._kv_committed_len
def _build_swa_tree(
is_eagle: bool,
page_size: int = 1,
req_size: int = 8,
max_context_len: int = 64,
kv_size: int = 64,
kv_size_swa: int = 32,
sliding_window_size: int = 4,
):
head_num = 8
head_dim = 128
num_layers = 24
global_interval = 4
dtype = torch.bfloat16
device = get_device()
full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
full_attention_layer_ids_set = set(full_attention_layer_ids)
swa_attention_layer_ids = [
i for i in range(num_layers) if i not in full_attention_layer_ids_set
]
req_to_token_pool = ReqToTokenPool(
size=req_size,
max_context_len=max_context_len,
device=device,
enable_memory_saver=False,
)
kv_pool = SWAKVPool(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
swa_attention_layer_ids=swa_attention_layer_ids,
full_attention_layer_ids=full_attention_layer_ids,
enable_kvcache_transpose=False,
device=device,
)
allocator = SWATokenToKVPoolAllocator(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
device=device,
kvcache=kv_pool,
need_sort=False,
)
tree = SWARadixCache(
params=CacheInitParams(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
page_size=page_size,
disable=False,
is_eagle=is_eagle,
sliding_window_size=sliding_window_size,
),
)
return tree, allocator, req_to_token_pool
def _swa_alloc(allocator, need_size):
"""SWA-pool alloc that also works for page_size > 1 (built-in alloc asserts page_size == 1)."""
if allocator.page_size == 1:
return allocator.alloc(need_size)
assert need_size % allocator.page_size == 0
full_indices = allocator.full_attn_allocator.alloc(need_size)
swa_indices = allocator.swa_attn_allocator.alloc(need_size)
assert full_indices is not None and swa_indices is not None
allocator.full_to_swa_index_mapping[full_indices] = swa_indices
return full_indices
def _insert(tree, allocator, token_ids):
indices = _swa_alloc(allocator, len(token_ids))
assert indices is not None
tree.insert(InsertParams(key=RadixKey(token_ids), value=indices))
def _insert_chain(tree, allocator, token_ids):
_insert(tree, allocator, token_ids)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(token_ids)))
return match.last_device_node
def _expected_tail_size(window: int, page_size: int) -> int:
"""Mirror of _maybe_split_leaf_for_swa_lock's tail_size formula."""
return (window + page_size - 1) // page_size * page_size
class TestSWA(unittest.TestCase):
class _DummyReq:
def __init__(self):
self._kv_committed_len = 0
def pop_committed_kv_cache(self):
return self._kv_committed_len
def _build_swa_tree(
self,
is_eagle: bool,
page_size: int = 1,
req_size: int = 8,
max_context_len: int = 64,
kv_size: int = 64,
kv_size_swa: int = 32,
sliding_window_size: int = 4,
):
head_num = 8
head_dim = 128
num_layers = 24
global_interval = 4
dtype = torch.bfloat16
device = get_device()
full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
full_attention_layer_ids_set = set(full_attention_layer_ids)
swa_attention_layer_ids = [
i for i in range(num_layers) if i not in full_attention_layer_ids_set
]
req_to_token_pool = ReqToTokenPool(
size=req_size,
max_context_len=max_context_len,
device=device,
enable_memory_saver=False,
)
kv_pool = SWAKVPool(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
swa_attention_layer_ids=swa_attention_layer_ids,
full_attention_layer_ids=full_attention_layer_ids,
enable_kvcache_transpose=False,
device=device,
)
allocator = SWATokenToKVPoolAllocator(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
device=device,
kvcache=kv_pool,
need_sort=False,
)
tree = SWARadixCache(
params=CacheInitParams(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
page_size=page_size,
disable=False,
is_eagle=is_eagle,
sliding_window_size=sliding_window_size,
),
)
return tree, allocator, req_to_token_pool
@classmethod
def setUpClass(cls):
pass
@@ -475,10 +510,10 @@ class TestSWA(unittest.TestCase):
self.assertEqual(list(last_node.key), [(5, 60), (60, 70)])
def test_swa_cache_finished_req_eagle_uses_cache_protected_len_and_bigram_key(self):
tree, allocator, req_to_token_pool = self._build_swa_tree(is_eagle=True)
tree, allocator, req_to_token_pool = _build_swa_tree(is_eagle=True)
# Case 1: is_insert=True should pass bigram key and use cache_protected_len.
req = self._DummyReq()
req = _DummyReq()
req.req_pool_idx = 0
req.origin_input_ids = [1, 2, 3, 4, 5, 6]
req.output_ids = []
@@ -513,7 +548,7 @@ class TestSWA(unittest.TestCase):
# Case 2: is_insert=False should free [cache_protected_len:page_aligned_len]
# even when len(prefix_indices) is intentionally larger.
req2 = self._DummyReq()
req2 = _DummyReq()
req2.req_pool_idx = 1
req2.origin_input_ids = [11, 12, 13, 14, 15, 16]
req2.output_ids = []
@@ -546,5 +581,112 @@ class TestSWA(unittest.TestCase):
self.assertEqual(freed_lens, [4, 1])
# Optimization: SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.
# Splits a freshly-inserted leaf at the (page-aligned) sliding-window
# boundary so a future inc_lock_ref protects only ~sliding_window_size SWA
# tokens instead of the whole chunked-prefill chain.
class TestSWASplitLeafOnInsert(CustomTestCase):
def _insert_and_lock(self, *, window, page_size, leaf_len, flag_on):
tree, allocator, _ = _build_swa_tree(
is_eagle=False,
kv_size=128,
kv_size_swa=64,
sliding_window_size=window,
page_size=page_size,
)
token_ids = list(range(leaf_len))
with envs.SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.override(flag_on):
leaf = _insert_chain(tree, allocator, token_ids)
result = tree.inc_lock_ref(leaf)
return tree, leaf, result
def test_flag_off_protects_full_leaf(self):
tree, leaf, _ = self._insert_and_lock(
window=4, page_size=1, leaf_len=12, flag_on=False
)
self.assertEqual(len(leaf.value), 12)
self.assertEqual(tree.swa_protected_size_, 12)
def test_flag_on_caps_protection_at_window(self):
# (window, page_size, leaf_len, expected_tail_size); leaf_len picked
# > tail_size and page-aligned for page_size > 1.
cases = [
(4, 1, 12, 4),
(4, 1, 5, 4),
(1, 1, 5, 1),
(4, 2, 12, 4),
(8, 2, 12, 8),
(4, 4, 12, 4),
# window NOT page-aligned -> tail rounds up to page boundary.
(3, 2, 12, 4),
(5, 4, 12, 8),
(3, 4, 12, 4),
]
for window, page_size, leaf_len, expected_tail in cases:
with self.subTest(window=window, page_size=page_size, leaf_len=leaf_len):
self.assertEqual(_expected_tail_size(window, page_size), expected_tail)
tree, leaf, _ = self._insert_and_lock(
window=window,
page_size=page_size,
leaf_len=leaf_len,
flag_on=True,
)
self.assertEqual(len(leaf.value), expected_tail)
self.assertEqual(tree.swa_protected_size_, expected_tail)
def test_flag_on_no_split_when_leaf_within_window(self):
# leaf_len <= tail_size: split must no-op.
cases = [
(4, 1, 4),
(4, 1, 3),
(4, 2, 4),
(3, 2, 4),
(8, 2, 4),
(4, 4, 4),
]
for window, page_size, leaf_len in cases:
with self.subTest(window=window, page_size=page_size, leaf_len=leaf_len):
tree, leaf, _ = self._insert_and_lock(
window=window,
page_size=page_size,
leaf_len=leaf_len,
flag_on=True,
)
self.assertEqual(len(leaf.value), leaf_len)
self.assertEqual(tree.swa_protected_size_, leaf_len)
def test_match_prefix_returns_full_chain_after_split(self):
tree, allocator, _ = _build_swa_tree(
is_eagle=False,
kv_size=128,
kv_size_swa=64,
sliding_window_size=4,
page_size=1,
)
token_ids = list(range(12))
with envs.SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.override(True):
inserted_leaf = _insert_chain(tree, allocator, token_ids)
self.assertEqual(len(inserted_leaf.value), 4)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(token_ids)))
self.assertEqual(match.device_indices.shape[0], 12)
self.assertIs(match.last_device_node, inserted_leaf)
def test_dec_lock_ref_after_split_balances_to_zero(self):
tree, leaf, result = self._insert_and_lock(
window=4, page_size=1, leaf_len=12, flag_on=True
)
self.assertEqual(tree.swa_protected_size_, 4)
self.assertEqual(tree.full_protected_size_, 12)
tree.dec_lock_ref(
leaf,
params=DecLockRefParams(swa_uuid_for_lock=result.swa_uuid_for_lock),
)
self.assertEqual(tree.swa_protected_size_, 0)
self.assertEqual(tree.full_protected_size_, 0)
tree.sanity_check()
if __name__ == "__main__":
unittest.main()