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
test/srt/test_swa_unittest.py
+11 -2
View File
@@ -21,6 +21,7 @@ class TestSWA(unittest.TestCase):
def test_swa_memory_pool(self):
size = 16
size_swa = 16
page_size = 1
head_num = 8
head_dim = 128
num_layers = 48
@@ -35,6 +36,7 @@ class TestSWA(unittest.TestCase):
pool = SWAKVPool(
size=size,
size_swa=size_swa,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
@@ -46,6 +48,7 @@ class TestSWA(unittest.TestCase):
alloc = SWATokenToKVPoolAllocator(
size=size,
size_swa=size_swa,
page_size=page_size,
dtype=dtype,
device=device,
kvcache=pool,
@@ -69,6 +72,7 @@ class TestSWA(unittest.TestCase):
max_context_len = 128
kv_size = 128
kv_size_swa = 64
page_size = 1
sliding_window_size = 4
head_num = 8
head_dim = 128
@@ -92,6 +96,7 @@ class TestSWA(unittest.TestCase):
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,
@@ -104,6 +109,7 @@ class TestSWA(unittest.TestCase):
allocator = SWATokenToKVPoolAllocator(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
device=device,
kvcache=kv_pool,
@@ -115,7 +121,7 @@ class TestSWA(unittest.TestCase):
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
disable=False,
page_size=1,
page_size=page_size,
),
sliding_window_size=sliding_window_size,
)
@@ -202,6 +208,7 @@ class TestSWA(unittest.TestCase):
max_context_len = 128
kv_size = 128
kv_size_swa = 64
page_size = 1
sliding_window_size = 4
head_num = 8
head_dim = 128
@@ -225,6 +232,7 @@ class TestSWA(unittest.TestCase):
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,
@@ -237,6 +245,7 @@ class TestSWA(unittest.TestCase):
allocator = SWATokenToKVPoolAllocator(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
device=device,
kvcache=kv_pool,
@@ -247,7 +256,7 @@ class TestSWA(unittest.TestCase):
params=CacheInitParams(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
page_size=1,
page_size=page_size,
disable=False,
is_eagle=True,
),