精选理由
LMSYS 和 InclusionAI 联手,用 SGLang-JAX 让 1T 参数 MoE 在 TPU 上跑得快 53%,技术细节都在博客里。
LMSYS 发表博客详解如何用 SGLang-JAX 在 TPU v7x 上优化 Ling-2.6-1T(1T 参数混合 MoE 模型)。通过 Fused MoE V2 内核将令牌和累加器留在 VMEM 中并双缓冲专家权重,MoE 预填充延迟降低 53%。混合内存池为 10 个全注意力层分配逐令牌 MLA KV,为 70 个 GLA 层分配逐请求循环状态。GLA 线性注意力采用分块并行预填充,单控制器 DP 保持分组 RMS Norm 芯片本地化,无需逐层跨芯片规约。
AI 翻译 · 中文
LMSYS 发表博客详解如何用 SGLang-JAX 在 TPU v7x 上优化 Ling-2.6-1T(1T 参数混合 MoE 模型)。通过 Fused MoE V2 内核将令牌和累加器留在 VMEM 中并双缓冲专家权重,MoE 预填充延迟降低 53%。混合内存池为 10 个全注意力层分配逐令牌 MLA KV,为 70 个 GLA 层分配逐请求循环状态。GLA 线性注意力采用分块并行预填充,单控制器 DP 保持分组 RMS Norm 芯片本地化,无需逐层跨芯片规约。
🚀 Our new blog: Optimizing Ling-2.6-1T on TPU with SGLang-JAX: Hiding MoE Data Movement Behind Compute with One Pallas Kernel Ling-2.6-1T, a 1T hybrid MoE model, now serves on TPU v7x with SGLang-JAX. The SGLang-JAX tea…