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

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

Browse Source
This commit is contained in:
Brian
2026-04-29 16:58:41 +08:00
committed by GitHub
parent 73e93bebd6
commit 6c7b242181
2 changed files with 114 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/cookbook/autoregressive/Xiaomi/MiMo-V2.5.mdx
+36
View File
@@ -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
+78 -16
View File
@@ -34,10 +34,12 @@ export const MiMoV25Deployment = () => {
name: "hardware",
title: "Hardware Platform",
items: [
{ id: "h200", label: "H200", default: true },
{ id: "h100", label: "H100", default: false },
{ id: "b200", label: "B200", default: false },
{ id: "gb300", label: "GB300", default: false },
{ id: "h200", label: "H200", default: true },
{ 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|h100": { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 16, multinode: true, nnodes: 2, blackwell: false },
"pro|b200": { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 8, multinode: false, blackwell: true },
"pro|gb300": { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 8, multinode: true, nnodes: 2, blackwell: true },
"base|h200": { slug: "XiaomiMiMo/MiMo-V2.5", tp: 8, multinode: false, blackwell: false, dp: 2 },
"base|h100": { slug: "XiaomiMiMo/MiMo-V2.5", tp: 8, multinode: false, blackwell: false, dp: 2 },
"base|b200": { slug: "XiaomiMiMo/MiMo-V2.5", tp: 4, multinode: false, blackwell: true, dp: 1 },
"base|gb300": { slug: "XiaomiMiMo/MiMo-V2.5", tp: 4, multinode: false, blackwell: true, dp: 1 },
"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|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 },
"pro|tpu-v6e": { slug: "XiaomiMiMo/MiMo-V2.5-Pro", tp: 64, multinode: true, nnodes: 16, blackwell: false, jax: true },
"base|h200": { slug: "XiaomiMiMo/MiMo-V2.5", tp: 8, multinode: false, blackwell: false, jax: false, dp: 2 },
"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";