diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 074feae11..8945f1d49 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/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" diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 83888afae..993c04c16 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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: diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index f9cb07da8..6054f3f7a 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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( diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 30f0bc8ca..6c807c590 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/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." diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index acc8c46c6..f23519560 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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: diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index c09c0fb17..068b1a19f 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/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 diff --git a/test/registered/unit/disaggregation/test_minimax_sparse_disagg_state_kv_args.py b/test/registered/unit/disaggregation/test_minimax_sparse_disagg_state_kv_args.py new file mode 100644 index 000000000..eeb68a897 --- /dev/null +++ b/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()