diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 47bec0d7686a..090ab0989354 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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 ) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 09fd8efa944e..4f6a10fd381a 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/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() diff --git a/python/sglang/srt/models/minimax_m3.py b/python/sglang/srt/models/minimax_m3.py index 5e9c0f42f0bd..50ce7c226db4 100644 --- a/python/sglang/srt/models/minimax_m3.py +++ b/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"] 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 index eeb68a897448..559998cd0912 100644 --- 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 @@ -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() diff --git a/test/registered/unit/models/test_minimax_m3_pp_cuda_graph.py b/test/registered/unit/models/test_minimax_m3_pp_cuda_graph.py new file mode 100644 index 000000000000..01610d07761d --- /dev/null +++ b/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()