[Qwen3.5] mamba slice fix (Prefill TP != Decode TP & decode TP size>1) (sgl-project#20655)

YAMY1234 · ShangmingCai · web-flow · commit cfead25bbf69 · 2026-03-17T19:30:58.000+08:00
Co-authored-by: Shangming Cai &lt;csmthu@gmail.com&gt;
diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py
@@ -726,7 +726,9 @@ def _send_mamba_state_slice(
                 # Each prefill sends all its dims to the appropriate offset in decode
                 src_dim_start = 0
                 num_dims_to_send = src_dim
-                dst_dim_start = local_tp_rank_in_group * src_dim
+                writers_per_decode = self.attn_tp_size // dst_attn_tp_size
+                local_writer_idx = local_tp_rank_in_group % writers_per_decode
+                dst_dim_start = local_writer_idx * src_dim
             else:
                 # 1 prefill rank sends to multiple decode ranks
                 # Prefill sends a slice of its dims to each decode rank
