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

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

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()