From 53ed34cb882ec52d53f738b5fb808ca28ac373b7 Mon Sep 17 00:00:00 2001
From: JoyFuture <35593546+JoyFuture@users.noreply.github.com>
Date: Wed, 10 Jun 2026 22:08:28 +0800
Subject: [PATCH] Fix MiMo-V2.5-Pro DP-attention dp size in cookbook deployment
 snippet (#27668)

---
 .../autoregressive/mimo-v25-deployment.jsx    | 35 ++++++++++++-------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/docs_new/src/snippets/autoregressive/mimo-v25-deployment.jsx b/docs_new/src/snippets/autoregressive/mimo-v25-deployment.jsx
index fc46f4237..bb9b49a38 100644
--- a/docs_new/src/snippets/autoregressive/mimo-v25-deployment.jsx
+++ b/docs_new/src/snippets/autoregressive/mimo-v25-deployment.jsx
@@ -93,13 +93,17 @@ export const MiMoV25Deployment = () => {
   };
 
   // Per (variant, hardware): HF slug, tp, multinode info, Blackwell flag.
-  // V2.5 (base) checkpoint has TP=4-interleaved fused qkv_proj, so attention
-  // TP per DP group MUST be 4. Effective TP/DP = 4. With tp=8 → dp=2; tp=4 → dp=1.
+  // The attention qkv_proj is TP-interleaved, so attention-TP per DP group is
+  // fixed: V2.5 (base) is TP=4-interleaved, V2.5-Pro is TP=8-interleaved. When
+  // tp exceeds that factor, DP-attention is required with `dp = tp / factor`
+  // (Pro: tp/8, base: tp/4); when tp equals the factor there is a single
+  // attention group and no `dp` field is set. So Pro on Hopper (tp=16) needs
+  // dp=2, while Pro on Blackwell (tp=8) needs no dp-attention at all.
   // TPU rows go through the sgl-jax stack (`python -m sgl_jax.launch_server`),
   // not the CUDA `sglang serve` binary; tp == total JAX devices across nodes.
   const HW_VARIANT_SPEC = {
-    "pro|h200":     { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 16, multinode: true,  nnodes: 2,  blackwell: false, jax: false },
-    "pro|h100":     { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 16, multinode: true,  nnodes: 2,  blackwell: false, jax: false },
+    "pro|h200":     { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 16, multinode: true,  nnodes: 2,  blackwell: false, jax: false, dp: 2 },
+    "pro|h100":     { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 16, multinode: true,  nnodes: 2,  blackwell: false, jax: false, dp: 2 },
     "pro|b200":     { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 8,  multinode: false,             blackwell: true,  jax: false },
     "pro|gb300":    { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 8,  multinode: true,  nnodes: 2,  blackwell: true,  jax: false },
     "pro|tpu-v7x":  { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 32, multinode: true,  nnodes: 4,  blackwell: false, jax: true  },
@@ -132,13 +136,18 @@ export const MiMoV25Deployment = () => {
     const blackwell = spec ? spec.blackwell : false;
     const jax = spec ? spec.jax : false;
     const c = {};
-    if (!isPro) {
-      // V2.5 checkpoint is TP=4-interleaved; tp/dp must equal 4. With dp>1 we
-      // must use DP-attention; with dp=1 it must be off (single attention group).
-      if (spec && spec.dp > 1) {
-        c.dpAttention = { force: "enabled", reason: "V2.5 checkpoint is TP=4-interleaved; DP-attention is required (--dp = tp/4)." };
+    // Both checkpoints are TP-interleaved (Pro: 8, base: 4), so attention-TP per
+    // DP group must equal that factor. When the spec carries dp>1 (Pro/Hopper
+    // tp=16, base tp=8) DP-attention with `--dp = tp/factor` is required; without
+    // it a bare `--tp` gives attn_tp = tp != factor and the loader rejects the
+    // checkpoint. When no dp>1 (Pro/Blackwell tp=8, base tp=4) it's a single
+    // attention group and DP-attention must stay off.
+    if (spec && !jax) {
+      const factor = isPro ? 8 : 4;
+      if (spec.dp > 1) {
+        c.dpAttention = { force: "enabled", reason: `Checkpoint is TP=${factor}-interleaved; DP-attention is required (--dp = tp/${factor} = ${spec.dp}).` };
       } else {
-        c.dpAttention = { force: "disabled", reason: "Single attention group on this hardware (tp=4, dp=1)." };
+        c.dpAttention = { force: "disabled", reason: `Single attention group on this hardware (tp=${factor}, no dp-attention).` };
       }
     }
     if (blackwell) {
@@ -281,8 +290,10 @@ export const MiMoV25Deployment = () => {
     const useDeepep = !blackwell && deepep === "enabled";
     const useEp = isPro && !blackwell && expertParallelism === "enabled";
     const useDpAttn = dpAttention === "enabled";
-    // dp size: V2.5 picks tp/4 from spec; Pro picks tp.
-    const dpSize = !isPro ? spec.dp : tp;
+    // dp size = required DP-attention degree from the spec (tp/factor), for both
+    // variants. Only read when useDpAttn is on, which computeConstraints forces
+    // exactly for the specs that carry dp>1 (Pro/Hopper tp=16 → 2, base tp=8 → 2).
+    const dpSize = spec.dp;
 
     // ---- env (kept inline before `sglang serve`, matching the verified launch style) ----
     const envVars = [];