minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Fix MiniMax M3 state transfer in Mori PD (#29756)

Browse Source 
Co-authored-by: Duyi-Wang <duyi.wang@amd.com>
This commit is contained in:
YukioZzz
2026-07-02 15:02:10 -07:00
committed by GitHub
co-authored by Duyi-Wang
parent ae1f0c6d07
commit caf2e5da2d
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/mori/conn.py
+10 -1
View File
@@ -1092,7 +1092,7 @@ class MoriKVManager(CommonKVManager):
dst_dims, dst_dims,
) )
) )
elif st in ("swa", "dsa", "swa_ring", "c128_state"): elif st in ("swa", "dsa", "swa_ring", "c128_state", "minimax_index_k"):
statuses.extend( statuses.extend(
self._send_swa_dsa_state( self._send_swa_dsa_state(
peer_info, peer_info,
@@ -1219,6 +1219,15 @@ class MoriKVManager(CommonKVManager):
f"PD state transfer does not support TP-mismatched non-MLA SWA models " f"PD state transfer does not support TP-mismatched non-MLA SWA models "
f"(prefill_tp_size={self.attn_tp_size}, decode_tp_size={peer_info.decode_tp_size})" f"(prefill_tp_size={self.attn_tp_size}, decode_tp_size={peer_info.decode_tp_size})"
) )
if state_type == "minimax_index_k":
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 peer_info.decode_tp_size != self.attn_tp_size:
raise RuntimeError(
"PD disagg: heterogeneous TP not supported for MiniMax sparse index yet."
)
common_len = min(src_state_indices.size, dst_state_indices.size) common_len = min(src_state_indices.size, dst_state_indices.size)
if ( if (