[HiSparse] Add MHA hisparse support for MiniMax M3 (#31446)

Co-authored-by: Guangda Liu <bingps@users.noreply.github.com>

test/registered/kernels/ops/kvcache/test_hisparse.py

@@ -28,6 +28,7 @@ if is_xpu():
     )
 else:
     from sglang.kernels.ops.kvcache.hisparse import (
+        load_blocks_to_device_buffer_mha,
         load_cache_to_device_buffer_dsv4_mla,
         load_cache_to_device_buffer_mla,
         transfer_cache_dsv4_mla,
@@ -368,6 +369,84 @@ def test_load_cache_to_device_buffer_hits_newest_and_updates_lru() -> None:
     )
 
 
+@pytest.mark.skipif(is_xpu(), reason="MiniMax MHA block swap-in has no XPU kernel.")
+def test_load_blocks_to_device_buffer_mha_handles_partial_newest_block() -> None:
+    """A partial newest block must not consume slots for its invalid tail."""
+    sparse_block_size = 4
+    hot_buffer_size = 8
+    host_k = _host_cache()
+    host_v = _host_cache()
+    host_v.add_(1000)
+    device_k = torch.full(
+        (DEVICE_CACHE_SIZE, 1, KV_DIM), -1, dtype=DTYPE, device=DEVICE
+    )
+    device_v = torch.full_like(device_k, -1)
+    device_buffer_locs = torch.arange(
+        hot_buffer_size + 1, dtype=torch.int32, device=DEVICE
+    ).view(1, -1)
+    device_buffer_tokens = torch.tensor(
+        [[0, 1, 2, 3, -1, -1, -1, -1, -1]],
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    for slot, token in enumerate([0, 1, 2, 3]):
+        device_k[device_buffer_locs[0, slot]].copy_(host_k[token], non_blocking=True)
+        device_v[device_buffer_locs[0, slot]].copy_(host_v[token], non_blocking=True)
+    device_k[device_buffer_locs[0, hot_buffer_size]].copy_(
+        host_k[10], non_blocking=True
+    )
+    device_v[device_buffer_locs[0, hot_buffer_size]].copy_(
+        host_v[10], non_blocking=True
+    )
+
+    top_k_blocks = torch.tensor([[0, 2]], dtype=torch.int32, device=DEVICE)
+    out = torch.full(
+        (1, top_k_blocks.size(1) * sparse_block_size),
+        -1,
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    lru_slots = torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE).view(
+        1, -1
+    )
+    load_blocks_to_device_buffer_mha(
+        top_k_blocks=top_k_blocks,
+        device_buffer_tokens=device_buffer_tokens,
+        host_cache_locs=torch.arange(
+            HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE
+        ).view(1, -1),
+        device_buffer_locs=device_buffer_locs,
+        host_cache_k=host_k,
+        host_cache_v=host_v,
+        device_buffer_k=device_k,
+        device_buffer_v=device_v,
+        top_k_device_locs=out,
+        req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE),
+        seq_lens=torch.tensor([11], dtype=torch.int32, device=DEVICE),
+        lru_slots=lru_slots,
+        item_size_bytes=ITEM_SIZE_BYTES,
+        hot_buffer_size=hot_buffer_size,
+        sparse_block_size=sparse_block_size,
+        num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
+    )
+    get_device_module().synchronize()
+
+    assert torch.equal(
+        out.cpu(), torch.tensor([[0, 1, 2, 3, 4, 5, 8, -1]], dtype=torch.int32)
+    )
+    assert torch.equal(device_k[4].cpu(), host_k[8])
+    assert torch.equal(device_v[4].cpu(), host_v[8])
+    assert torch.equal(device_k[5].cpu(), host_k[9])
+    assert torch.equal(device_v[5].cpu(), host_v[9])
+    assert torch.equal(
+        device_buffer_tokens.cpu(),
+        torch.tensor([[0, 1, 2, 3, 8, 9, -1, -1, -1]], dtype=torch.int32),
+    )
+    assert torch.equal(
+        lru_slots.cpu(), torch.tensor([[6, 7, 4, 5, 0, 1, 2, 3]], dtype=torch.int16)
+    )
+
+
 def test_load_cache_to_device_buffer_miss_uses_updated_lru_slot() -> None:
     state = _long_case()
 

test/registered/unit/disaggregation/test_minimax_sparse_disagg_state_kv_args.py

@@ -53,14 +53,19 @@ def _make_kv_pool(start_layer: int = 0) -> MiniMaxSparseKVPool:
 
 
 class TestMiniMaxSparseDisaggStateKvArgs(unittest.TestCase):
-    def test_setup_state_kv_args_single_minimax_component(self):
+    def test_setup_state_kv_args_minimax_components(self):
         pool = _make_k_only_pool()
         kv_args = KVArgs()
         setup_state_kv_args(kv_args, pool)
-        self.assertEqual(kv_args.state_types, [StateType.MINIMAX_INDEX_K])
-        self.assertEqual(len(kv_args.state_data_ptrs), 1)
+        self.assertEqual(
+            kv_args.state_types,
+            [StateType.MINIMAX_INDEX_K, StateType.MINIMAX_DENSE_KV],
+        )
+        self.assertEqual(len(kv_args.state_data_ptrs), 2)
         self.assertEqual(len(kv_args.state_data_ptrs[0]), pool.index_k_pool.layer_num)
         self.assertEqual(len(kv_args.state_item_lens[0]), pool.index_k_pool.layer_num)
+        self.assertEqual(len(kv_args.state_data_ptrs[1]), 6)
+        self.assertEqual(len(kv_args.state_item_lens[1]), 6)
 
     def test_index_kv_pool_raises(self):
         pool = _make_kv_pool()

test/registered/unit/disaggregation/test_mooncake_transfer_batching.py

@@ -1,4 +1,5 @@
 import concurrent.futures
+import ctypes
 import unittest
 from threading import Event
 from types import SimpleNamespace
@@ -6,6 +7,7 @@ from unittest.mock import MagicMock, call, patch
 
 import numpy as np
 
+from sglang.srt.disaggregation.base.conn import StateType
 from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
@@ -132,6 +134,61 @@ class TestMooncakeTransferBatching(unittest.TestCase):
         )
 
 
+class TestMiniMaxStateTransfer(CustomTestCase):
+    def test_index_truncates_but_dense_rejects_mismatched_page_lists(self):
+        """Legacy index transfers copy the common prefix; incomplete dense KV must fail."""
+
+        def copy_bytes(session, sources, destinations, lengths):
+            for src, dst, length in zip(sources, destinations, lengths, strict=True):
+                ctypes.memmove(dst, src, length)
+            return 0
+
+        for state in (StateType.MINIMAX_INDEX_K, StateType.MINIMAX_DENSE_KV):
+            for src_pages, dst_pages in (([1], [0]), ([1, 2], [0]), ([1], [0, 2])):
+                with self.subTest(state=state, src=src_pages, dst=dst_pages):
+                    src = np.arange(3, dtype=np.int32)
+                    dst = np.full(3, -1, dtype=np.int32)
+                    manager = MooncakeKVManager.__new__(MooncakeKVManager)
+                    manager.kv_args = SimpleNamespace(
+                        state_types=[state],
+                        state_data_ptrs=[[src.ctypes.data]],
+                        state_item_lens=[[src.itemsize]],
+                        state_dim_per_tensor=[[]],
+                        state_layer_ids=[[]],
+                    )
+                    manager.engine = SimpleNamespace(batch_transfer_sync=copy_bytes)
+                    manager.pp_size = manager.attn_tp_size = 1
+                    manager.is_mla_backend = manager.is_hybrid_mla_backend = False
+                    manager.enable_custom_mem_pool = False
+                    manager.max_transfer_batch_indices = 0
+                    peer = SimpleNamespace(
+                        dst_state_data_ptrs=[[dst.ctypes.data]],
+                        dst_state_item_lens=[[dst.itemsize]],
+                        dst_state_dim_per_tensor=[[]],
+                        dst_state_layer_ids=[[]],
+                        dst_attn_tp_size=1,
+                    )
+                    kwargs = dict(
+                        req=SimpleNamespace(
+                            mooncake_session_id="cpu", dst_state_indices=[dst_pages]
+                        ),
+                        prefill_state_indices=[src_pages],
+                        executor=None,
+                        target_rank_registration_info=peer,
+                    )
+                    if state == StateType.MINIMAX_DENSE_KV and len(src_pages) != len(
+                        dst_pages
+                    ):
+                        with self.assertRaisesRegex(
+                            RuntimeError, "state index length mismatch"
+                        ):
+                            manager.maybe_send_extra(**kwargs)
+                        np.testing.assert_array_equal(dst, [-1, -1, -1])
+                    else:
+                        self.assertEqual(manager.maybe_send_extra(**kwargs), 0)
+                        np.testing.assert_array_equal(dst, [1, -1, -1])
+
+
 class TestDcpDraftHeadTransfer(unittest.TestCase):
     def test_transfers_draft_heads_to_logical_destination_rows(self):
         for src_tp, dst_tp in ((4, 8), (8, 4), (8, 8), (4, 32), (32, 4)):

test/registered/unit/mem_cache/test_hisparse_allocator.py

@@ -1,7 +1,7 @@
 import unittest
 from concurrent.futures import ThreadPoolExecutor
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import torch
@@ -18,6 +18,71 @@ from sglang.test.test_utils import CustomTestCase
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
 
+class TestHiSparseDecodeRemap(CustomTestCase):
+    def test_page_size_one_reclaims_temporary_device_slot(self):
+        """Decode remapping must reclaim its temporary slot without freeing the live slot."""
+        from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
+        from sglang.srt.mem_cache.allocator.hisparse import (
+            HiSparseTokenToKVPoolAllocator,
+        )
+        from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
+
+        pool = MiniMaxSparseKVPool(
+            size=8,
+            page_size=1,
+            dtype=torch.float32,
+            head_num=1,
+            head_dim=8,
+            idx_head_dim=16,
+            dense_layer_ids=[0],
+            sparse_layer_ids=[1],
+            disable_value_sparse_layer_ids=[1],
+            device="cpu",
+            start_layer=0,
+            end_layer=2,
+            enable_hisparse=True,
+        )
+        allocator = HiSparseTokenToKVPoolAllocator(
+            size=pool.size,
+            page_size=1,
+            dtype=pool.dtype,
+            device="cpu",
+            kvcache=pool,
+            need_sort=False,
+        )
+        coordinator = HiSparseCoordinator.__new__(HiSparseCoordinator)
+        coordinator.is_dsv4_hisparse = False
+        coordinator.mem_pool_device = pool.main_pool
+        coordinator.token_to_kv_pool_allocator = allocator
+        coordinator.device_buffer_size = 2
+        coordinator.req_to_device_buffer = allocator.hisparse_attn_allocator.alloc(
+            3
+        ).reshape(1, 3)
+        coordinator.req_device_buffer_size = torch.tensor([3])
+        coordinator.req_device_buffer_token_locs = torch.zeros(
+            (1, 1, 3), dtype=torch.int32
+        )
+        coordinator._skip_first_backup = [True]
+        out_loc = allocator.alloc(1)
+        with patch("sglang.srt.managers.hisparse_coordinator._is_hip", False):
+            for _ in range(2):
+                coordinator._skip_first_backup[0] = True
+                coordinator.map_last_loc_to_buffer(
+                    seq_lens=torch.tensor([3]),
+                    out_cache_loc=out_loc,
+                    req_pool_indices=torch.tensor([0]),
+                    seq_lens_cpu=torch.tensor([3]),
+                    req_pool_indices_cpu=torch.tensor([0]),
+                )
+                self.assertEqual(
+                    allocator.hisparse_attn_allocator.available_size(), pool.size - 3
+                )
+                torch.testing.assert_close(
+                    allocator.full_to_hisparse_device_index_mapping[out_loc],
+                    coordinator.req_to_device_buffer[:, 2],
+                )
+
+
 class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
     def setUp(self):
         # The code under test reads its config from the bags.

test/registered/unit/mem_cache/test_minimax_sparse_pool_pd_unit.py

@@ -1,14 +1,23 @@
 import unittest
+from unittest.mock import patch
 
 import torch
 
+from sglang.srt.disaggregation.utils import get_kv_transfer_buf_infos
 from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
+from sglang.srt.mem_cache.pool_host.mha import (
+    HiSparseMHATokenToKVPoolHost,
+    MHATokenToKVPoolHost,
+)
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
 
-def _make_k_only_pool(start_layer: int = 0) -> MiniMaxSparseKVPool:
+def _make_k_only_pool(
+    start_layer: int = 0, *, enable_hisparse: bool = False
+) -> 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)]
@@ -26,10 +35,11 @@ def _make_k_only_pool(start_layer: int = 0) -> MiniMaxSparseKVPool:
         device="cpu",
         start_layer=start_layer,
         end_layer=end_layer,
+        enable_hisparse=enable_hisparse,
     )
 
 
-class TestMiniMaxSparsePoolPD(unittest.TestCase):
+class TestMiniMaxSparsePoolPD(CustomTestCase):
     def test_contiguous_buf_infos_main_only(self):
         pool = _make_k_only_pool()
         ptrs, lens, item_lens = pool.get_contiguous_buf_infos()
@@ -53,6 +63,52 @@ class TestMiniMaxSparsePoolPD(unittest.TestCase):
             self.assertEqual(lens[i], buf.nbytes)
             self.assertEqual(item_lens[i], buf[0].nbytes * pool.page_size)
 
+    def test_hisparse_host_registration(self):
+        """PD startup must register every sparse host K/V buffer with page strides."""
+        pool = _make_k_only_pool(enable_hisparse=True)
+        host = HiSparseMHATokenToKVPoolHost.__new__(HiSparseMHATokenToKVPoolHost)
+        with patch(
+            "sglang.srt.mem_cache.pool_host.base.host_memory_budget_bytes",
+            return_value=1 << 30,
+        ):
+            MHATokenToKVPoolHost.__init__(
+                host,
+                device_pool=pool.main_pool,
+                host_to_device_ratio=2,
+                host_size=0,
+                page_size=pool.page_size,
+                layout="layer_first",
+                pin_memory=False,
+            )
+        ptrs, lens, item_lens = get_kv_transfer_buf_infos(host)
+        buffers = list(host.k_buffer.unbind()) + list(host.v_buffer.unbind())
+        self.assertEqual(len(buffers), 8)
+        self.assertEqual(ptrs, [buffer.data_ptr() for buffer in buffers])
+        self.assertEqual(lens, [buffer.nbytes for buffer in buffers])
+        self.assertEqual(
+            item_lens, [buffer[0].nbytes * pool.page_size for buffer in buffers]
+        )
+
+    def test_pd_registration_separates_dense_and_sparse_layers(self):
+        """Both PD peers must keep dense device KV out of the sparse transfer list."""
+        for hisparse in (False, True):
+            pool = _make_k_only_pool(enable_hisparse=hisparse)
+            for layers, infos in (
+                (range(3, 7), get_kv_transfer_buf_infos(pool)),
+                (range(3), pool.get_dense_kv_state_buf_infos()),
+            ):
+                with self.subTest(hisparse=hisparse, layers=layers):
+                    buffers = [pool.get_key_buffer(i) for i in layers] + [
+                        pool.get_value_buffer(i) for i in layers
+                    ]
+                    ptrs, lens, item_lens = infos
+                    self.assertEqual(ptrs, [buffer.data_ptr() for buffer in buffers])
+                    self.assertEqual(lens, [buffer.nbytes for buffer in buffers])
+                    self.assertEqual(
+                        item_lens,
+                        [buffer[0].nbytes * pool.page_size for buffer in buffers],
+                    )
+
 
 if __name__ == "__main__":
     unittest.main()