[minimax-m3] Split 3/4: disagg K-only index-K transfer (#28714)

python/sglang/srt/disaggregation/base/conn.py

@@ -18,6 +18,7 @@ class StateType(str, enum.Enum):
     MAMBA = "mamba"
     SWA = "swa"
     DSA = "dsa"
+    MINIMAX_INDEX_K = "minimax_index_k"
     # DeepSeek-V4 unified_kv SWA ring: addressed per-row by ring slot
     # (req_pool_idx * ring_stride + pos % ring_stride), needs its own component.
     SWA_RING = "swa_ring"

python/sglang/srt/disaggregation/decode.py

@@ -1058,6 +1058,10 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                     state_indices.append(_swa_payload())
                 elif st == StateType.DSA:
                     state_indices.append(_dsa_payload())
+                elif st == StateType.MINIMAX_INDEX_K:
+                    # Index rows live at the same loc as main KV on the same
+                    # page_size, so reuse the full-seq page-ids.
+                    state_indices.append(_dsa_payload())
                 elif st == StateType.SWA_RING:
                     state_indices.append(_swa_ring_payload())
                 else:

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -586,10 +586,15 @@ class MooncakeKVManager(CommonKVManager):
         prefill_data_indices: npt.NDArray[np.int32],
         dst_data_indices: npt.NDArray[np.int32],
         executor: concurrent.futures.ThreadPoolExecutor,
+        force_flat: bool = False,
     ) -> int:
         """
         Generic KV cache transfer supporting both MHA and MLA architectures.
         This method is used by both send_kvcache (full pool) and maybe_send_extra.
+
+        ``force_flat`` uses the MLA-style flat (single-buffer-per-layer) layout
+        even on a non-MLA backend, for K-only state buffers (e.g. MiniMax sparse
+        index) whose per-layer list must not be half-split into K/V.
         """
         # Group by indices for optimization
         prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
@@ -599,7 +604,7 @@ class MooncakeKVManager(CommonKVManager):
         layers_params = None
 
         # Decode pp size should be equal to prefill pp size or 1
-        if self.is_mla_backend:
+        if self.is_mla_backend or force_flat:
             src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
                 self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
             )
@@ -1036,6 +1041,41 @@ class MooncakeKVManager(CommonKVManager):
                     )
                     or rc
                 )
+            elif st == StateType.MINIMAX_INDEX_K:
+                # Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer
+                # lists, so PP>1 mis-slices and heterogeneous TP is unsupported.
+                if self.pp_size is not None and self.pp_size > 1:
+                    raise RuntimeError(
+                        "PD disagg: PP>1 not supported for MiniMax sparse index yet."
+                    )
+                if (
+                    target_rank_registration_info is not None
+                    and self.attn_tp_size
+                    != target_rank_registration_info.dst_attn_tp_size
+                ):
+                    raise RuntimeError(
+                        "PD disagg: heterogeneous TP not supported for MiniMax "
+                        "sparse index yet."
+                    )
+                src_indices = list(indices)
+                dst_indices_local = list(dst_indices)
+                if len(src_indices) > len(dst_indices_local):
+                    src_indices = src_indices[: len(dst_indices_local)]
+                elif len(src_indices) < len(dst_indices_local):
+                    dst_indices_local = dst_indices_local[: len(src_indices)]
+                rc = (
+                    self._send_kvcache_generic(
+                        mooncake_session_id=req.mooncake_session_id,
+                        src_data_ptrs=src_data_ptrs,
+                        dst_data_ptrs=dst_data_ptrs,
+                        item_lens=src_item_lens,
+                        prefill_data_indices=np.array(src_indices, dtype=np.int32),
+                        dst_data_indices=np.array(dst_indices_local, dtype=np.int32),
+                        executor=executor,
+                        force_flat=True,
+                    )
+                    or rc
+                )
         return rc
 
     def _send_mamba_state(

python/sglang/srt/disaggregation/nixl/conn.py

@@ -1279,10 +1279,14 @@ class NixlKVManager(CommonKVManager):
         notif: str,
         src_mem_kind: str = "VRAM",
         dst_mem_kind: str = "VRAM",
+        force_flat: bool = False,
     ):
         """Generic KV cache transfer supporting both MHA and MLA architectures.
-        Used by both send_kvcache and maybe_send_extra."""
+        Used by both send_kvcache and maybe_send_extra.
 
+        ``force_flat`` uses the MLA-style flat (single-buffer-per-layer) layout
+        even on a non-MLA backend, for K-only state buffers (e.g. MiniMax sparse
+        index) whose per-layer list must not be half-split into K/V."""
         # Prepped path (KV only; state transfers use the non-prepped path below).
         if (
             src_data_ptrs is self.kv_args.kv_data_ptrs
@@ -1335,7 +1339,7 @@ class NixlKVManager(CommonKVManager):
 
         logger.debug(f"sending kvcache to {peer_name} with notif {notif}")
         # Make descs
-        if self.is_mla_backend:
+        if self.is_mla_backend or force_flat:
             src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
                 self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
             )
@@ -2010,6 +2014,34 @@ class NixlKVManager(CommonKVManager):
                     dst_gpu_id=dst_gpu_id,
                     notif=comp_notif,
                 )
+            elif st == StateType.MINIMAX_INDEX_K:
+                # Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer
+                # lists, so PP>1 mis-slices and heterogeneous TP is unsupported.
+                if self.pp_size is not None and self.pp_size > 1:
+                    raise RuntimeError(
+                        "PD disagg: PP>1 not supported for MiniMax sparse index yet."
+                    )
+                if self.attn_tp_size != decode_tp_size:
+                    raise RuntimeError(
+                        "PD disagg: heterogeneous TP not supported for MiniMax "
+                        "sparse index yet."
+                    )
+                if len(src_indices) != len(dst_indices):
+                    raise RuntimeError(
+                        f"State index length mismatch at component {i}: "
+                        f"prefill={len(src_indices)}, dst={len(dst_indices)}"
+                    )
+                h = self._send_kvcache_generic(
+                    peer_name=peer_name,
+                    src_data_ptrs=src_ptrs,
+                    dst_data_ptrs=dst_ptrs,
+                    item_lens=src_lens,
+                    prefill_data_indices=np.array(src_indices, dtype=np.int32),
+                    dst_data_indices=np.array(dst_indices, dtype=np.int32),
+                    dst_gpu_id=dst_gpu_id,
+                    notif=comp_notif,
+                    force_flat=True,
+                )
             else:
                 raise RuntimeError(
                     f"PD Disaggregation via NIXL does NOT support {st} hybrid models yet."

python/sglang/srt/disaggregation/prefill.py

@@ -1070,6 +1070,10 @@ class SchedulerDisaggregationPrefillMixin:
                     state_indices.append(_swa_payload())
                 elif st == StateType.DSA:
                     state_indices.append(_dsa_payload())
+                elif st == StateType.MINIMAX_INDEX_K:
+                    # Index rows live at the same loc as main KV on the same
+                    # page_size, so reuse the full-seq page-ids.
+                    state_indices.append(_dsa_payload())
                 elif st == StateType.SWA_RING:
                     state_indices.append(_swa_ring_payload())
                 else:

python/sglang/srt/disaggregation/utils.py

@@ -643,7 +643,11 @@ def setup_state_kv_args(
     from sglang.srt.disaggregation.base.conn import StateType
     from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool
     from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
-    from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, HybridLinearKVPool
+    from sglang.srt.mem_cache.memory_pool import (
+        DSATokenToKVPool,
+        HybridLinearKVPool,
+        MiniMaxSparseKVPool,
+    )
 
     kv_args.state_types = []
     kv_args.state_data_ptrs = []
@@ -651,7 +655,16 @@ def setup_state_kv_args(
     kv_args.state_item_lens = []
     kv_args.state_dim_per_tensor = []
 
-    if hasattr(token_to_kv_pool, "get_state_buf_infos"):
+    if isinstance(token_to_kv_pool, MiniMaxSparseKVPool):
+        if token_to_kv_pool.index_kv_pool is not None:
+            raise NotImplementedError(
+                "PD disaggregation for MiniMax sparse layers with index value "
+                "(index_kv_pool) is not yet supported; only K-only sparse layers are."
+            )
+        if token_to_kv_pool.index_k_pool is not None:
+            dp, dl, il = token_to_kv_pool.get_index_k_state_buf_infos()
+            append_state_component(kv_args, StateType.MINIMAX_INDEX_K, dp, dl, il)
+    elif hasattr(token_to_kv_pool, "get_state_buf_infos"):
         data_ptrs, data_lens, item_lens = token_to_kv_pool.get_state_buf_infos()
 
         # DeepSeekV4TokenToKVPool inherits BaseSWAKVPool; its heterogeneous

test/registered/unit/disaggregation/test_minimax_sparse_disagg_state_kv_args.py

@@ -0,0 +1,74 @@
+import unittest
+
+import torch
+
+from sglang.srt.disaggregation.base.conn import KVArgs, StateType
+from sglang.srt.disaggregation.utils import setup_state_kv_args
+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,
+    )
+
+
+def _make_kv_pool(start_layer: int = 0) -> MiniMaxSparseKVPool:
+    """Sparse layers with index value (index_kv_pool != None)."""
+    dense_layer_ids = [start_layer, start_layer + 1]
+    sparse_layer_ids = [start_layer + 2, start_layer + 3]
+    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=[],
+        device="cpu",
+        start_layer=start_layer,
+        end_layer=end_layer,
+    )
+
+
+class TestMiniMaxSparseDisaggStateKvArgs(unittest.TestCase):
+    def test_setup_state_kv_args_single_minimax_component(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(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)
+
+    def test_index_kv_pool_raises(self):
+        pool = _make_kv_pool()
+        self.assertIsNotNone(pool.index_kv_pool)
+        kv_args = KVArgs()
+        with self.assertRaises(NotImplementedError):
+            setup_state_kv_args(kv_args, pool)
+
+
+if __name__ == "__main__":
+    unittest.main()