mimo v2.5 pro sglang-jax cookbook (#23936)

docs_new/cookbook/autoregressive/Xiaomi/MiMo-V2.5.mdx

@@ -65,6 +65,8 @@ Refer to the [official SGLang installation guide](../../../docs/get-started/inst
 
 > Pull the image matching your GPU's CUDA driver. `lmsysorg/sglang:latest` will not load either checkpoint.
 
+**TPU (sgl-jax):** MiMo-V2.5-Pro can also be served on TPU via the JAX-based [sgl-jax](https://github.com/sgl-project/sglang-jax) runtime. The container image and `pip install` steps are listed in [§3.3 TPU Deployment](#33-tpu-deployment-mimo-v25-pro-sgl-jax).
+
 ## 3. Model Deployment
 
 ### 3.1 Basic Configuration
@@ -95,6 +97,40 @@ import { MiMoV25Deployment } from '/src/snippets/autoregressive/mimo-v25-deploym
 - Enabling adds `--moe-a2a-backend deepep` + `--moe-dense-tp-size 1` (and `--ep <tp>` for Pro) plus `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256` env to cap the dispatch buffer. Requires `pip install deep_ep` (not part of the default sglang install).
 - On Blackwell (B200, GB300) the verified MoE backend is `flashinfer_trtllm`; the DeepEP toggle is a no-op there.
 
+### 3.3 TPU Deployment (MiMo-V2.5-Pro, sgl-jax)
+
+MiMo-V2.5-Pro can also be served on TPU via [sgl-jax](https://github.com/sgl-project/sglang-jax). The runtime is a separate JAX-based stack (`sgl_jax.launch_server`); pick **TPU v7x** or **TPU v6e** in the panel above to generate the launch command. Verified topologies:
+
+| TPU Type | Topology | Chips/Node | Nodes | Total Chips | JAX Devices/Chip | Total JAX Devices (= `--tp-size`) |
+| --- | --- | --- | --- | --- | --- | --- |
+| **v7x** | 2×2×4 | 4 | 4 | 16 | 2 | 32 |
+| **v6e** | 4×4×4 | 4 | 16 | 64 | 1 | 64 |
+
+> v7x exposes **2 logical JAX devices per chip**, so `--tp-size = 16 chips × 2 = 32`. v6e exposes 1 device per chip, so `--tp-size = 64`. Always set `--tp-size` to the total JAX device count across all nodes, not the chip count.
+
+All nodes must sit in the same TPU slice and reach each other on the JAX init port (`20000`) and the TPU process port (`8471`).
+
+**Step 1 — Launch the JAX TPU container on every node:**
+
+```shell Command
+docker run -it --privileged \
+  --shm-size=32g \
+  --ipc=host \
+  --network=host \
+  -v /dev:/dev \
+  us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:jax0.8.1-rev1 bash
+```
+
+> The image is pinned to `jax0.8.1-rev1` to keep the JAX runtime aligned with sgl-jax's TPU extras.
+
+**Step 2 — Clone and install sgl-jax (inside the container):**
+
+```shell Command
+git clone https://github.com/sgl-project/sglang-jax.git
+cd sglang-jax
+pip install -e "python[tpu]"
+```
+
 ## 4. Model Invocation
 
 ### 4.1 Basic Usage

docs_new/src/snippets/autoregressive/mimo-v25-deployment.jsx

@@ -38,6 +38,8 @@ export const MiMoV25Deployment = () => {
         { id: "h100",     label: "H100",     default: false },
         { id: "b200",     label: "B200",     default: false },
         { id: "gb300",    label: "GB300",    default: false },
+        { id: "tpu-v7x",  label: "TPU v7x",  default: false, subtitle: "sgl-jax, Pro only" },
+        { id: "tpu-v6e",  label: "TPU v6e",  default: false, subtitle: "sgl-jax, Pro only" },
       ],
     },
     eagleMtp: {
@@ -93,15 +95,19 @@ 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.
+  // 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 },
+    "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 },
+    "pro|h100":     { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 16, multinode: true,  nnodes: 2,  blackwell: false, jax: false },
-    "pro|b200":   { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 8,  multinode: false,            blackwell: true  },
+    "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  },
+    "pro|gb300":    { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 8,  multinode: true,  nnodes: 2,  blackwell: true,  jax: false },
-    "base|h200":  { slug: "XiaomiMiMo/MiMo-V2.5",     tp: 8,  multinode: false,            blackwell: false, dp: 2 },
+    "pro|tpu-v7x":  { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 32, multinode: true,  nnodes: 4,  blackwell: false, jax: true  },
-    "base|h100":  { slug: "XiaomiMiMo/MiMo-V2.5",     tp: 8,  multinode: false,            blackwell: false, dp: 2 },
+    "pro|tpu-v6e":  { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 64, multinode: true,  nnodes: 16, blackwell: false, jax: true  },
-    "base|b200":  { slug: "XiaomiMiMo/MiMo-V2.5",     tp: 4,  multinode: false,            blackwell: true,  dp: 1 },
+    "base|h200":    { slug: "XiaomiMiMo/MiMo-V2.5",     tp: 8,  multinode: false,             blackwell: false, jax: false, dp: 2 },
-    "base|gb300": { slug: "XiaomiMiMo/MiMo-V2.5",     tp: 4,  multinode: false,            blackwell: true,  dp: 1 },
+    "base|h100":    { slug: "XiaomiMiMo/MiMo-V2.5",     tp: 8,  multinode: false,             blackwell: false, jax: false, dp: 2 },
+    "base|b200":    { slug: "XiaomiMiMo/MiMo-V2.5",     tp: 4,  multinode: false,             blackwell: true,  jax: false, dp: 1 },
+    "base|gb300":   { slug: "XiaomiMiMo/MiMo-V2.5",     tp: 4,  multinode: false,             blackwell: true,  jax: false, dp: 1 },
   };
 
   const multiNodeFlags = (nnodes) => [
@@ -124,6 +130,7 @@ export const MiMoV25Deployment = () => {
     const isPro = variant === "pro";
     const spec = HW_VARIANT_SPEC[`${variant}|${hardware}`];
     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
@@ -139,15 +146,25 @@ export const MiMoV25Deployment = () => {
       // for sm_100 in sglang and the verified Blackwell stack uses flashinfer_trtllm.
       c.deepep = { force: "disabled", reason: "Blackwell uses flashinfer_trtllm; DeepEP is Hopper / Ampere only." };
     }
+    if (jax) {
+      // sgl-jax stack: only V2.5-Pro is supported on TPU today; speculative
+      // decoding and the DeepEP CUDA backend do not apply to the JAX runtime.
+      // EP is always on (both verified launch commands set --ep-size = --tp-size).
+      c.modelVariant = { force: "pro", reason: "sgl-jax TPU runtime only supports MiMo-V2.5-Pro today." };
+      c.eagleMtp = { force: "disabled", reason: "EAGLE MTP is not supported on the sgl-jax TPU runtime." };
+      c.deepep = { force: "disabled", reason: "DeepEP is a CUDA-only backend; sgl-jax uses the fused Pallas MoE kernel." };
+      c.expertParallelism = { force: "enabled", reason: "sgl-jax TPU recipes always use EP = TP." };
+    }
     return c;
   };
 
   const resolveItems = (option, constraints) => {
     const c = constraints[option.name];
     if (!c) return option.items;
-    const grayId = c.force === "enabled" ? "disabled" : "enabled";
+    // Gray out every item that doesn't match the forced choice. Works for both
+    // binary (enabled/disabled) toggles and N-way options like modelVariant.
     return option.items.map((item) =>
-      item.id === grayId ? { ...item, disabled: true, disabledReason: c.reason } : item,
+      item.id !== c.force ? { ...item, disabled: true, disabledReason: c.reason } : item,
     );
   };
 
@@ -205,9 +222,54 @@ export const MiMoV25Deployment = () => {
     const { modelVariant, hardware, eagleMtp, dpAttention, expertParallelism, deepep, reasoningParser, toolcall } = values;
     const specKey = `${modelVariant}|${hardware}`;
     const spec = HW_VARIANT_SPEC[specKey];
-    const { slug, tp, multinode, nnodes, blackwell } = spec;
+    const { slug, tp, multinode, nnodes, blackwell, jax } = spec;
     const isPro = modelVariant === "pro";
-    // Toggles. EP / DeepEP / DP-attn are gated by hardware + variant
+
+    // ---------------- sgl-jax (TPU) branch ----------------
+    if (jax) {
+      // Recipe sources:
+      //   v7x: tp=ep=32, dp=4, omits --attention-backend, mem-frac 0.95, swa 0.25
+      //   v6e: tp=ep=64, dp=8, --attention-backend fa,    mem-frac 0.92, swa 0.15
+      const isV7x = hardware === "tpu-v7x";
+      const useEp = expertParallelism === "enabled";
+      const useDpAttn = dpAttention === "enabled";
+      const dpSize = isV7x ? 4 : 8;
+      const flags = [];
+      flags.push(`  --model-path ${slug}`);
+      flags.push("  --trust-remote-code");
+      flags.push(`  --tp-size ${tp}`);
+      if (useEp) flags.push(`  --ep-size ${tp}`);
+      if (useDpAttn) flags.push(`  --dp-size ${dpSize}`, "  --enable-dp-attention");
+      flags.push("  --moe-backend fused");
+      if (!isV7x) flags.push("  --attention-backend fa");
+      flags.push("  --host 0.0.0.0");
+      flags.push("  --port 30000");
+      flags.push("  --page-size 256");
+      flags.push("  --context-length 262144");
+      flags.push("  --chunked-prefill-size 4096");
+      flags.push("  --max-running-requests 512");
+      if (isV7x) {
+        flags.push("  --dtype bfloat16");
+        flags.push("  --mem-fraction-static 0.95");
+        flags.push("  --swa-full-tokens-ratio 0.25");
+        flags.push("  --log-level info");
+      } else {
+        flags.push("  --max-seq-len 4096");
+        flags.push("  --max-prefill-tokens 16384");
+        flags.push("  --mem-fraction-static 0.92");
+        flags.push("  --swa-full-tokens-ratio 0.15");
+      }
+      if (reasoningParser === "enabled") flags.push("  --reasoning-parser mimo");
+      if (toolcall === "enabled") flags.push("  --tool-call-parser mimo");
+      flags.push(`  --nnodes ${nnodes}`);
+      flags.push("  --node-rank <node-rank>");
+      flags.push("  --dist-init-addr <node0-ip>:20000");
+      const cmd = `JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python -m sgl_jax.launch_server \\\n${flags.join(" \\\n")}`;
+      return prependMultiNodeNote(cmd, nnodes);
+    }
+
+    // ---------------- CUDA (sglang serve) branch ----------------
+    // Toggles. EAGLE MTP / EP / DeepEP / DP-attn are gated by hardware + variant
     // through computeConstraints; here we just read the (already-snapped) value.
     const useMtp = eagleMtp === "enabled";
     const useDeepep = !blackwell && deepep === "enabled";