fix: support MiniMax-M3 pipeline parallelism in PD disaggregation

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

@@ -1387,12 +1387,8 @@ def maybe_send_extra(
                     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."
-                    )
+                # Equal attention-TP only. Sparse sub-pools are compacted per
+                # pipeline stage, so pair their entries by global layer id.
                 if (
                     target_rank_registration_info is not None
                     and self.attn_tp_size
@@ -1418,6 +1414,8 @@ def maybe_send_extra(
                         dst_data_indices=np.array(dst_indices_local, dtype=np.int32),
                         executor=executor,
                         force_flat=True,
+                        src_layer_ids=src_state_layer_ids,
+                        dst_layer_ids=dst_state_layer_ids,
                     )
                     or rc
                 )

python/sglang/srt/disaggregation/utils.py

@@ -1011,7 +1011,14 @@ def setup_state_kv_args(
             )
         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)
+            append_state_component(
+                kv_args,
+                StateType.MINIMAX_INDEX_K,
+                dp,
+                dl,
+                il,
+                layer_ids=list(token_to_kv_pool.index_k_layer_id_mapping),
+            )
     elif hasattr(token_to_kv_pool, "get_state_buf_infos"):
         data_ptrs, data_lens, item_lens = token_to_kv_pool.get_state_buf_infos()
 

python/sglang/srt/models/minimax_m3.py

@@ -230,6 +230,38 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self._qm.apply(self, x, None)
 
 
+def _normalize_scattered_pp_proxy_tensors_for_cuda_graph(
+    pp_proxy_tensors: PPProxyTensors,
+    *,
+    positions: torch.Tensor,
+    attn_tp_size: int,
+) -> PPProxyTensors:
+    """Slice graph dummy PP inputs to the token-local shape used by EP."""
+    global_num_tokens = positions.shape[0]
+    if attn_tp_size <= 1 or global_num_tokens % attn_tp_size != 0:
+        return pp_proxy_tensors
+
+    token_keys = ("hidden_states", "residual")
+    if not all(
+        key in pp_proxy_tensors.tensors
+        and pp_proxy_tensors[key].shape[0] == global_num_tokens
+        for key in token_keys
+    ):
+        return pp_proxy_tensors
+
+    local_num_tokens = global_num_tokens // attn_tp_size
+    return PPProxyTensors(
+        {
+            key: (
+                value[:local_num_tokens]
+                if key in token_keys
+                else value
+            )
+            for key, value in pp_proxy_tensors.tensors.items()
+        }
+    )
+
+
 def build_minimax_fused_qkv_index(model: nn.Module) -> None:
     for module in model.modules():
         if isinstance(module, MiniMaxM3Attention):
@@ -1488,6 +1520,18 @@ def forward(
             residual = None
         else:
             assert pp_proxy_tensors is not None
+            first_layer = self.layers[self.start_layer]
+            if (
+                first_layer.layer_scatter_modes.layer_input_mode
+                == ScatterMode.SCATTERED
+            ):
+                pp_proxy_tensors = (
+                    _normalize_scattered_pp_proxy_tensors_for_cuda_graph(
+                        pp_proxy_tensors,
+                        positions=positions,
+                        attn_tp_size=get_parallel().attn_tp_size,
+                    )
+                )
             hidden_states = pp_proxy_tensors["hidden_states"]
             residual = pp_proxy_tensors["residual"]
 

test/registered/unit/disaggregation/test_minimax_sparse_disagg_state_kv_args.py

@@ -61,6 +61,7 @@ def test_setup_state_kv_args_single_minimax_component(self):
         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)
+        self.assertEqual(kv_args.state_layer_ids, [[3, 4, 5, 6]])
 
     def test_index_kv_pool_raises(self):
         pool = _make_kv_pool()

test/registered/unit/models/test_minimax_m3_pp_cuda_graph.py

@@ -0,0 +1,50 @@
+import unittest
+
+import torch
+
+from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
+from sglang.srt.models.minimax_m3 import (
+    _normalize_scattered_pp_proxy_tensors_for_cuda_graph,
+)
+
+
+class TestMiniMaxM3PPCudaGraph(unittest.TestCase):
+    def test_slices_graph_dummy_proxy_to_local_tp_tokens(self):
+        proxy = PPProxyTensors(
+            {
+                "hidden_states": torch.arange(24).reshape(8, 3),
+                "residual": torch.arange(24).reshape(8, 3),
+            }
+        )
+
+        normalized = _normalize_scattered_pp_proxy_tensors_for_cuda_graph(
+            proxy,
+            positions=torch.zeros(8, dtype=torch.int64),
+            attn_tp_size=4,
+        )
+
+        self.assertEqual(normalized["hidden_states"].shape, (2, 3))
+        self.assertEqual(normalized["residual"].shape, (2, 3))
+        torch.testing.assert_close(
+            normalized["residual"], proxy["residual"][:2]
+        )
+
+    def test_keeps_runtime_proxy_that_is_already_scattered(self):
+        proxy = PPProxyTensors(
+            {
+                "hidden_states": torch.zeros((2, 3)),
+                "residual": torch.zeros((2, 3)),
+            }
+        )
+
+        normalized = _normalize_scattered_pp_proxy_tensors_for_cuda_graph(
+            proxy,
+            positions=torch.zeros(8, dtype=torch.int64),
+            attn_tp_size=4,
+        )
+
+        self.assertIs(normalized, proxy)
+
+
+if __name__ == "__main__":
+    unittest.main()