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

[minimax-m3] Split 2/4: mem-cache / HiCache / sparse KV pool (#28713)

Browse Source 
Co-authored-by: hzh0425 <hzh0425@apache.org>
This commit is contained in:
Xinyuan Tong
2026-06-28 00:19:38 +08:00
committed by GitHub
co-authored by hzh0425
parent cfd911ad6e
commit 592f6c849b
10 changed files with 1715 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_minimax_sparse_pool_host_unit.py
+410
View File
@@ -0,0 +1,410 @@
import unittest
import psutil
import torch
from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName
from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
HybridCacheController,
)
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
from sglang.srt.mem_cache.memory_pool_host import (
HICACHE_HOST_MEMORY_RESERVE_BYTES,
MHATokenToKOnlyPoolHost,
MHATokenToKVPoolHost,
)
from sglang.srt.mem_cache.pool_host.common import (
ALLOC_MEMORY_FUNCS,
alloc_with_pin_memory,
)
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small")
def _cuda_major() -> int:
cuda = getattr(torch.version, "cuda", None)
try:
return int(cuda.split(".")[0]) if cuda else 0
except ValueError:
return 0
# direct+page_first_direct routes to transfer_kv_all_layer_direct_lf_pf, which on
# CUDA 13 throws (cudaErrorInvalidValue) instead of falling back. M3 uses kernel+layer_first.
_DIRECT_PF_BATCHCOPY_BROKEN_CUDA13 = _cuda_major() >= 13
class _FakeLayerTransferCounter:
def __init__(self):
self.waited_layers = []
def wait_until(self, layer_id: int):
self.waited_layers.append(layer_id)
def _make_cpu_minimax_sparse_pool(start_layer: int = 4) -> MiniMaxSparseKVPool:
end_layer = start_layer + 4
return MiniMaxSparseKVPool(
size=8,
page_size=4,
dtype=torch.float32,
head_num=1,
head_dim=2,
idx_head_dim=3,
dense_layer_ids=[start_layer, start_layer + 2],
sparse_layer_ids=[start_layer + 1, start_layer + 3],
disable_value_sparse_layer_ids=[start_layer + 1, start_layer + 3],
device="cpu",
start_layer=start_layer,
end_layer=end_layer,
)
class TestMiniMaxSparseHiCacheIntegration(unittest.TestCase):
def test_hiradix_extra_pools_include_minimax_indexer(self):
pool = _make_cpu_minimax_sparse_pool()
cache = object.__new__(HiRadixCache)
cache.cache_controller = object.__new__(HybridCacheController)
cache.kv_cache = pool
extra = HiRadixCache._get_extra_pools(cache)
transfers = extra["extra_pools"]
self.assertEqual(len(transfers), 1)
self.assertEqual(transfers[0].name, PoolName.INDEXER)
self.assertEqual(transfers[0].indices_from_pool, PoolName.KV)
self.assertEqual(transfers[0].hit_policy, PoolHitPolicy.ALL_PAGES)
def test_index_k_waits_for_full_local_layer(self):
pool = _make_cpu_minimax_sparse_pool()
counter = _FakeLayerTransferCounter()
pool.register_layer_transfer_counter(counter)
pool.get_index_k_buffer(7)
self.assertEqual(counter.waited_layers, [3])
self.assertIsNone(pool.main_pool.layer_transfer_counter)
self.assertIsNone(pool.index_k_pool.layer_transfer_counter)
def test_main_kv_waits_on_minimax_wrapper(self):
pool = _make_cpu_minimax_sparse_pool()
counter = _FakeLayerTransferCounter()
pool.register_layer_transfer_counter(counter)
pool.get_kv_buffer(6)
self.assertEqual(counter.waited_layers, [2])
self.assertIsNone(pool.main_pool.layer_transfer_counter)
self.assertIsNone(pool.index_k_pool.layer_transfer_counter)
def test_k_only_host_pool_layout_contracts(self):
if psutil.virtual_memory().available <= HICACHE_HOST_MEMORY_RESERVE_BYTES:
self.skipTest("Not enough spare host memory for HiCache host pool tests.")
for layout in ("layer_first", "page_first", "page_first_direct"):
with self.subTest(layout=layout):
pool = _make_cpu_minimax_sparse_pool(start_layer=0)
kv_host = MHATokenToKVPoolHost(
device_pool=pool.main_pool,
host_to_device_ratio=2.0,
host_size=0,
page_size=pool.page_size,
layout=layout,
pin_memory=False,
device="cpu",
allocator_type="default",
)
index_host = MHATokenToKOnlyPoolHost(
pool.index_k_pool,
kv_host,
layout=layout,
pin_memory=False,
device="cpu",
allocator_type="default",
)
if layout == "layer_first":
self.assertEqual(
index_host.k_buffer.shape,
(
index_host.layer_num,
index_host.size,
index_host.head_num,
index_host.head_dim,
),
)
elif layout == "page_first":
self.assertEqual(
index_host.k_buffer.shape,
(
index_host.size,
index_host.layer_num,
index_host.head_num,
index_host.head_dim,
),
)
else:
self.assertEqual(
index_host.k_buffer.shape,
(
index_host.page_num,
index_host.layer_num,
index_host.page_size,
index_host.head_num,
index_host.head_dim,
),
)
page_start = pool.page_size
flat_page = torch.arange(
index_host.layer_num
* index_host.page_size
* index_host.head_num
* index_host.head_dim,
dtype=index_host.dtype,
)
index_host.set_from_flat_data_page(page_start, flat_page)
self.assertTrue(
torch.equal(index_host.get_data_page(page_start), flat_page)
)
self.assertEqual(
index_host.get_dummy_flat_data_page().numel(), flat_page.numel()
)
self.assertIs(
index_host.get_hybrid_pool_buffer()[0], index_host.k_buffer
)
indices = torch.arange(
page_start,
page_start + pool.page_size,
dtype=torch.int64,
)
if layout == "layer_first":
with self.assertRaisesRegex(ValueError, "layer_first"):
index_host.get_page_buffer_meta(indices)
continue
ptrs, sizes = index_host.get_page_buffer_meta(indices)
self.assertEqual(len(ptrs), 1)
expected_size = (
index_host.layer_num
* index_host.page_size
* index_host.head_num
* index_host.head_dim
* index_host.dtype.itemsize
)
self.assertEqual(sizes, [expected_size] * len(ptrs))
class TestMiniMaxSparseHiCacheTransfer(unittest.TestCase):
def setUp(self):
if not torch.cuda.is_available():
self.skipTest("CUDA is required for MiniMax sparse host transfer tests.")
if is_npu() or is_xpu():
self.skipTest("MiniMax sparse host transfer tests only support CUDA/ROCm.")
if not (is_cuda() or is_hip()):
self.skipTest("CUDA/ROCm not available.")
@staticmethod
def _token_indices_for_pages(pages: torch.Tensor, page_size: int, device: str):
parts = [
torch.arange(
int(page_id) * page_size,
(int(page_id) + 1) * page_size,
device=device,
dtype=torch.int64,
)
for page_id in pages.tolist()
]
return torch.cat(parts, dim=0)
@staticmethod
def _host_k_page(host_pool, layer_id: int, page_id: int, page_size: int):
start = page_id * page_size
if host_pool.layout == "layer_first":
return host_pool.k_buffer[layer_id][start : start + page_size]
if host_pool.layout == "page_first":
return host_pool.k_buffer[start : start + page_size, layer_id]
if host_pool.layout == "page_first_direct":
return host_pool.k_buffer[page_id, layer_id]
raise ValueError(f"Unsupported layout: {host_pool.layout}")
@staticmethod
def _host_v_page(host_pool, layer_id: int, page_id: int, page_size: int):
start = page_id * page_size
if host_pool.layout == "layer_first":
return host_pool.v_buffer[layer_id][start : start + page_size]
if host_pool.layout == "page_first":
return host_pool.v_buffer[start : start + page_size, layer_id]
if host_pool.layout == "page_first_direct":
return host_pool.v_buffer[page_id, layer_id]
raise ValueError(f"Unsupported layout: {host_pool.layout}")
def _run_device_to_host_copy(self, io_backend: str, layout: str):
page_size = 64
layer_num = 4
size = page_size * 4
dense_layer_ids = [0, 1]
sparse_layer_ids = [2, 3]
device_pool = MiniMaxSparseKVPool(
size=size,
page_size=page_size,
dtype=torch.bfloat16,
head_num=4,
head_dim=64,
idx_head_dim=128,
dense_layer_ids=dense_layer_ids,
sparse_layer_ids=sparse_layer_ids,
disable_value_sparse_layer_ids=sparse_layer_ids,
device="cuda",
start_layer=0,
end_layer=layer_num,
)
assert device_pool.index_kv_pool is None
assert device_pool.index_k_pool is not None
pin_memory = io_backend == "kernel"
original_alloc = ALLOC_MEMORY_FUNCS["cuda"]
if pin_memory:
ALLOC_MEMORY_FUNCS["cuda"] = alloc_with_pin_memory
try:
kv_host = MHATokenToKVPoolHost(
device_pool=device_pool.main_pool,
host_to_device_ratio=2.0,
host_size=0,
page_size=page_size,
layout=layout,
pin_memory=pin_memory,
device="cpu",
allocator_type="default",
)
index_host = MHATokenToKOnlyPoolHost(
device_pool.index_k_pool,
kv_host,
layout=layout,
pin_memory=pin_memory,
device="cpu",
allocator_type="default",
)
finally:
ALLOC_MEMORY_FUNCS["cuda"] = original_alloc
for layer_id in range(layer_num):
k_main, v_main = device_pool.get_kv_buffer(layer_id)
k_main.copy_(torch.randn_like(k_main) + float(layer_id))
v_main.copy_(torch.randn_like(v_main) + float(layer_id))
for local_id, global_id in enumerate(sparse_layer_ids):
idx_k = device_pool.index_k_pool.k_buffer[local_id]
idx_k.copy_(torch.randn_like(idx_k) + float(global_id) + 100.0)
device_pages = torch.tensor([1, 2, 3], device="cuda", dtype=torch.int64)
host_pages = torch.tensor(
[0, 1, 2],
device="cuda" if io_backend == "kernel" else "cpu",
dtype=torch.int64,
)
device_indices = self._token_indices_for_pages(
device_pages, page_size, device="cuda"
)
host_indices = self._token_indices_for_pages(
host_pages,
page_size,
device="cuda" if io_backend == "kernel" else "cpu",
)
# page_first main-KV backup (staged_write_back.cuh) needs CPU dst_indices,
# index-k backup (hicache.cuh) needs CUDA indices — feed a CPU copy to main only.
kv_host_indices = (
host_indices.cpu()
if (io_backend, layout) == ("kernel", "page_first")
else host_indices
)
kv_host.backup_from_device_all_layer(
device_pool.main_pool, kv_host_indices, device_indices, io_backend
)
index_host.backup_from_device_all_layer(
device_pool.index_k_pool, host_indices, device_indices, io_backend
)
for layer_id in range(layer_num):
for host_page, device_page in zip(
host_pages.tolist(), device_pages.tolist()
):
device_start = device_page * page_size
got_k = self._host_k_page(kv_host, layer_id, host_page, page_size).cpu()
expected_k = device_pool.main_pool.k_buffer[layer_id][
device_start : device_start + page_size
].cpu()
self.assertTrue(torch.equal(got_k, expected_k))
got_v = self._host_v_page(kv_host, layer_id, host_page, page_size).cpu()
expected_v = device_pool.main_pool.v_buffer[layer_id][
device_start : device_start + page_size
].cpu()
self.assertTrue(torch.equal(got_v, expected_v))
for local_id, global_id in enumerate(sparse_layer_ids):
for host_page, device_page in zip(
host_pages.tolist(), device_pages.tolist()
):
got = self._host_k_page(
index_host, local_id, host_page, page_size
).cpu()
expected = device_pool.index_k_pool.k_buffer[local_id][
device_page * page_size : (device_page + 1) * page_size
].cpu()
self.assertTrue(torch.equal(got, expected))
# Round-trip H2D for one sparse index layer.
reload_pages = torch.tensor([0, 1], device="cuda", dtype=torch.int64)
host_device = "cuda" if io_backend == "kernel" else "cpu"
reload_host_pages = torch.tensor([3, 0], device=host_device, dtype=torch.int64)
reload_device_indices = self._token_indices_for_pages(
reload_pages, page_size, device="cuda"
)
reload_host_indices = self._token_indices_for_pages(
reload_host_pages, page_size, device=host_device
)
device_pool.index_k_pool.k_buffer[0].zero_()
index_host.load_to_device_per_layer(
device_pool.index_k_pool,
reload_host_indices,
reload_device_indices,
0,
io_backend,
)
for host_page, device_page in zip(
reload_host_pages.tolist(), reload_pages.tolist()
):
got = device_pool.index_k_pool.k_buffer[0][
device_page * page_size : (device_page + 1) * page_size
].cpu()
expected = self._host_k_page(index_host, 0, host_page, page_size).cpu()
self.assertTrue(torch.equal(got, expected))
def test_device_to_host_kernel_layer_first(self):
self._run_device_to_host_copy(io_backend="kernel", layout="layer_first")
def test_device_to_host_kernel_page_first(self):
self._run_device_to_host_copy(io_backend="kernel", layout="page_first")
def test_device_to_host_direct_layer_first(self):
self._run_device_to_host_copy(io_backend="direct", layout="layer_first")
@unittest.skipIf(
_DIRECT_PF_BATCHCOPY_BROKEN_CUDA13,
"direct+page_first_direct host transfer hits cudaMemcpyBatchAsync "
"cudaErrorInvalidValue on CUDA 13 (sgl-kernel transfer_kv_all_layer_direct_lf_pf "
"throws instead of falling back to per-page copy); M3 production uses "
"io_backend=kernel + layer_first, not this combo.",
)
def test_device_to_host_direct_page_first_direct(self):
self._run_device_to_host_copy(io_backend="direct", layout="page_first_direct")
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_minimax_sparse_pool_pd_unit.py
+58
View File
@@ -0,0 +1,58 @@
import unittest
import torch
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _make_k_only_pool(start_layer: int = 0) -> MiniMaxSparseKVPool:
"""Mirror the released MiniMax-M3 config shape: all sparse layers K-only."""
dense_layer_ids = [start_layer, start_layer + 1, start_layer + 2]
sparse_layer_ids = [start_layer + 3 + i for i in range(4)]
end_layer = sparse_layer_ids[-1] + 1
return MiniMaxSparseKVPool(
size=8,
page_size=4,
dtype=torch.float32,
head_num=2,
head_dim=8,
idx_head_dim=16,
dense_layer_ids=dense_layer_ids,
sparse_layer_ids=sparse_layer_ids,
disable_value_sparse_layer_ids=sparse_layer_ids,
device="cpu",
start_layer=start_layer,
end_layer=end_layer,
)
class TestMiniMaxSparsePoolPD(unittest.TestCase):
def test_contiguous_buf_infos_main_only(self):
pool = _make_k_only_pool()
ptrs, lens, item_lens = pool.get_contiguous_buf_infos()
# Main K/V only: 2 entries per main layer (K then V), no index buffers.
n = pool.main_pool.layer_num
self.assertEqual(len(ptrs), 2 * n)
self.assertEqual(len(lens), 2 * n)
self.assertEqual(len(item_lens), 2 * n)
self.assertEqual(ptrs, pool.main_pool.get_contiguous_buf_infos()[0])
def test_index_k_state_buf_infos(self):
pool = _make_k_only_pool()
ptrs, lens, item_lens = pool.get_index_k_state_buf_infos()
n = pool.index_k_pool.layer_num
self.assertEqual(len(ptrs), n)
self.assertEqual(len(lens), n)
self.assertEqual(len(item_lens), n)
for i in range(n):
buf = pool.index_k_pool.k_buffer[i]
self.assertEqual(ptrs[i], buf.data_ptr())
self.assertEqual(lens[i], buf.nbytes)
self.assertEqual(item_lens[i], buf[0].nbytes * pool.page_size)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py
+52 -1
View File
@@ -1,5 +1,5 @@
import unittest
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
from sglang.srt.mem_cache.hicache_storage import PoolName, SidecarPoolSpec
from sglang.srt.mem_cache.hybrid_cache import hybrid_pool_assembler
@@ -11,6 +11,7 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
_DeepSeekV4Strategy,
_DsaStrategy,
_MambaStrategy,
_MiniMaxSparseStrategy,
_PlainKvStrategy,
_select_strategy,
_SwaStrategy,
@@ -36,6 +37,9 @@ class TestUnifiedRadixHiCacheDispatch(unittest.TestCase):
order = [type(s) for s in _STRATEGIES]
# DeepSeekV4 inherits from SWAKVPool, so it must resolve before _SwaStrategy.
self.assertLess(order.index(_DeepSeekV4Strategy), order.index(_SwaStrategy))
self.assertLess(
order.index(_MiniMaxSparseStrategy), order.index(_PlainKvStrategy)
)
self.assertEqual(order[-1], _PlainKvStrategy)
def test_deepseek_v4_full_swa(self):
@@ -68,6 +72,53 @@ class TestUnifiedRadixHiCacheDispatch(unittest.TestCase):
strategy = _select_strategy(kvcache, {FULL})
self.assertIsInstance(strategy, _DsaStrategy)
def test_minimax_sparse(self):
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
kvcache = _mock_kvcache(MiniMaxSparseKVPool)
strategy = _select_strategy(kvcache, {FULL})
self.assertIsInstance(strategy, _MiniMaxSparseStrategy)
def test_minimax_sparse_build_registers_indexer_sidecar(self):
strategy = _MiniMaxSparseStrategy()
host_pool_group = MagicMock()
kv_host_pool = object()
host_pool_group.get_pool.return_value = kv_host_pool
cache_controller = MagicMock()
cache = MagicMock(page_size=4)
kvcache = MagicMock()
kvcache.index_k_pool = object()
kvcache.main_pool.layer_num = 8
params = MagicMock()
params.tp_cache_group = None
params.pp_rank = 0
params.pp_size = 1
server_args = MagicMock()
with patch.object(
hybrid_pool_assembler,
"build_minimax_sparse_hicache_stack",
return_value=(host_pool_group, cache_controller),
) as build_stack:
result = strategy.build(
cache=cache,
kvcache=kvcache,
params=params,
server_args=server_args,
load_cache_event=object(),
)
build_stack.assert_called_once()
self.assertIs(build_stack.call_args.kwargs["sparse_pool"], kvcache)
self.assertIs(result.host_pool_group, host_pool_group)
self.assertIs(result.cache_controller, cache_controller)
self.assertIs(result.component_host_pools[FULL], kv_host_pool)
self.assertEqual(result.pools_desc, "KV + INDEXER(k-only)")
self.assertEqual(result.transfer_layer_num, 8)
self.assertEqual(len(result.sidecars), 1)
self.assertEqual(result.sidecars[0].pool_name, PoolName.INDEXER)
self.assertEqual(result.sidecars[0].indices_from_pool, PoolName.KV)
def test_plain_kv_fallback(self):
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool