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

[HiCache & HybridModel] mooncake backend support DSA & mamba model (#21259)

Browse Source 
Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com>
Co-authored-by: hzh0425 <hzh0425@apache.org>
Co-authored-by: pansicheng <sicheng.pan.chn@gmail.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: Vladislav Nosivskoy <vladnosiv@gmail.com>
This commit is contained in:
huangtingwei
2026-04-13 18:47:36 -07:00
committed by GitHub
co-authored by hzh0425 pansicheng ispobock Vladislav Nosivskoy
parent cc449ac4e5
commit e9d6b9eb2d
8 changed files with 761 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_nsa_pool_host_unit.py
+19 -5
View File
@@ -5,7 +5,8 @@ import torch
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import (
ALLOC_MEMORY_FUNCS,
NSATokenToKVPoolHost,
MLATokenToKVPoolHost,
NSAIndexerPoolHost,
alloc_with_pin_memory,
)
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
@@ -58,7 +59,7 @@ class TestNSAHiCacheTransfer(unittest.TestCase):
if pin_memory:
ALLOC_MEMORY_FUNCS["cuda"] = alloc_with_pin_memory
try:
host_pool = NSATokenToKVPoolHost(
mla_host = MLATokenToKVPoolHost(
device_pool=device_pool,
host_to_device_ratio=2.0,
host_size=0,
@@ -66,6 +67,16 @@ class TestNSAHiCacheTransfer(unittest.TestCase):
layout="layer_first",
pin_memory=pin_memory,
device="cpu",
allocator_type="default",
override_kv_cache_dim=device_pool.kv_cache_dim,
)
indexer_host = NSAIndexerPoolHost(
device_pool=device_pool,
anchor_host=mla_host,
layout="layer_first",
pin_memory=pin_memory,
device="cpu",
allocator_type="default",
)
finally:
ALLOC_MEMORY_FUNCS["cuda"] = original_alloc
@@ -97,7 +108,10 @@ class TestNSAHiCacheTransfer(unittest.TestCase):
device="cuda" if io_backend == "kernel" else "cpu",
)
host_pool.backup_from_device_all_layer(
mla_host.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend
)
indexer_host.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend
)
@@ -105,14 +119,14 @@ class TestNSAHiCacheTransfer(unittest.TestCase):
for host_page, device_page in zip(
host_pages.tolist(), device_pages.tolist()
):
got = host_pool.index_k_with_scale_buffer[layer_id][host_page].cpu()
got = indexer_host.index_k_with_scale_buffer[layer_id][host_page].cpu()
expected = device_pool.index_k_with_scale_buffer[layer_id][
device_page
].cpu()
self.assertTrue(torch.equal(got, expected))
host_start = host_page * page_size
device_start = device_page * page_size
got_kv = host_pool.kv_buffer[layer_id][
got_kv = mla_host.kv_buffer[layer_id][
host_start : host_start + page_size
].cpu()
expected_kv = device_pool.kv_buffer[layer_id][