Prhub

#23936 mimo v2.5 pro sglang-jax cookbook

原始 PR 作者 JamesBrianD 合并时间 2026-04-29 16:58 文件变更 2 提交数 4 评论 11 代码增减 +114 / -16

feature documentation tpu

执行摘要

为 MiMo-V2.5-Pro 新增 TPU 部署指南与交互面板

Add a TPU deployment recipe for MiMo-V2.5-Pro to the cookbook, served via sgl-jax. This documents the verified v7x (16 chips / 4 nodes) and v6e (64 chips / 16 nodes) topologies and lets users generate the launch command from the existing variant/hardware panel instead of copy-pasting from a separate sglang-jax doc.

值得精读，特别是对部署多硬件方案感兴趣的读者。该 PR 展示了如何通过有限改动扩展部署面板支持新硬件，并处理了不同运行时（CUDA vs JAX）的差异。Review 中的讨论对文档精确性有很好的启发。

讨论亮点
  1. TPU v7x tp 值疑问：gemini-code-assist[bot] 指出 tp: 32 与文档中 16 chips 不符。作者解释 v7x 每芯片暴露 2 个 JAX 设备，并在文档中增加 JAX Devices/Chip 和 Total JAX Devices 列明确定义。
  2. 端口混淆：文档提及内部 JAX 端口 8471 但生成命令使用 --port 30271 造成困惑。作者将启动端口统一为 30000，与 CUDA 路径保持一致，消除单独说明。
  3. 代码重复：--chunked-prefill-size 4096 和 --max-running-requests 512 在 v7x 和 v6e 分支重复。作者将其提取到条件外部，减少冗余。

实现拆解

实现拆解

  1. 扩展硬件选项：在 mimo-v25-deployment.jsx 的 hardware 选项中新增 tpu-v7x 和 tpu-v6e，并标注 sgl-jax, Pro only。
  2. 补充拓扑配置：在 HW_VARIANT_SPEC 中添加 pro|tpu-v7x 和 pro|tpu-v6e，设置 jax: true 标志；tp 值分别为 32 和 64（对应 JAX 设备总数）。
  3. 命令生成逻辑：在 generateCommand 函数中检测 jax 属性，切换为 python -m sgl_jax.launch_server，强制启用 EP，禁用 EAGLE MTP 和 DeepEP，并设置 TPU 特有参数（mem-fraction, swa 等）。公共 flag 如 --chunked-prefill-size 和 --max-running-requests 被提取到条件外部。
  4. 新增文档章节：在 MiMo-V2.5.mdx 中添加 3.3 TPU Deployment，包含拓扑表、JAX 容器镜像、sgl-jax 安装指南，并解释端口用途。
文件 模块 状态 重要度
docs_new/src/snippets/autoregressive/mimo-v25-deployment.jsx 部署面板 modified 6.95
docs_new/cookbook/autoregressive/Xiaomi/MiMo-V2.5.mdx 文档 modified 3.81

关键符号

MiMoV25Deployment

关键源码片段

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

核心逻辑文件：添加 TPU 硬件选项、拓扑配置及命令生成路径。

// 硬件选项新增 TPU 条目
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: "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" },
],
​
// 拓扑配置扩展，jax: true 表示使用 sgl-jax 运行时
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|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 variants 保持不变
};

评论区精华

TPU v7x tp 值与芯片数不一致 正确性

gemini-code-assist[bot] 指出 spec 中 tp=32 与文档中 16 chips 矛盾，要求澄清。

结论：作者解释 v7x 每芯片 2 个 JAX 设备，因此 tp=32 正确，并在文档中增加 JAX Devices/Chip 列明确说明。 · 已解决

文档端口说明混淆 documentation

gemini-code-assist[bot] 提出内部 JAX 端口 8471 与生成命令端口 30271 不一致，易混淆。

结论：作者将启动端口统一为 30000，与 CUDA 路径一致，不再需要单独解释。 · 已解决

TPU 分支存在重复 flag style

gemini-code-assist[bot] 指出 --chunked-prefill-size 和 --max-running-requests 在两个 TPU 分支重复。

结论：作者将公共 flag 提取到 if 条件外部，消除重复。 · 已解决

风险与影响

风险较低。主要风险包括：

  • 配置错误：TPU 拓扑表或 tp 值若标注错误可能导致启动失败，需严格按验证值填写。
  • 外部依赖：sgl-jax 与 sglang 主版本分离，若 sgl-jax 更新不兼容，文档可能过时。
  • 用户误解：若用户不理解 JAX 设备与物理芯片区别，可能设置错误 tp 值。

对用户：TPU 用户现在可直接通过 cookbook 部署 MiMo-V2.5-Pro，无需自行查阅 sgl-jax 文档。对项目：增加 TPU 部署路径的可见度和信任度，有助于吸引更多 TPU 用户。对团队：需维护 sgl-jax 版本的兼容性。

文档与配置变更 依赖外部 sgl-jax 项目 无测试覆盖

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。

完整报告

Markdown PDF

参与讨论