# 用SGLang-JAX在TPU上优化Ling-2.6-1T：一个Pallas核将MoE数据移动隐藏在计算中

- 来源：LMSYS：Blog（Chatbot Arena 团队）
- 发布时间：2026-06-17 00:00
- AIHOT 分数：69
- AIHOT 标记：精选
- AIHOT 链接：https://aihot.virxact.com/items/cmqiddvsx07xcslf0wcefjhm6
- 原文链接：https://www.lmsys.org/blog/2026-06-17-ling-2-6-tpu

## 精选理由

这是针对TPU上MoE推理的硬核优化复盘，用成本模型定位瓶颈，通过单内核融合把延迟砍半，对做大规模推理工程的团队是高质量参考。

## AI 摘要

SGLang-JAX现已支持inclusionAI的Ling-2.6-1T（1T稀疏MoE，63B激活参数，256路由专家，top-8路由加共享专家）在TPU v7x上高效推理。团队开发了Fused MoE V2——一个融合scatter、专家FFN和gather的Pallas核，通过将MoE数据移动隐藏在计算中，使MoE预填充延迟从5.16ms降至2.42ms（降幅53%），解码核延迟从0.249ms降至0.211ms（降幅约15%）。仅替换MoE核即提升预填充吞吐量24.8%，解码吞吐量18.5%–35.3%。在SGLang解码基准测试中，16块TPU v7x芯片输出吞吐量达16块H200 GPU的1.29倍（mc=128）至1.77倍（mc=512）。完整上线还包含混合KV/循环内存池、GLA线性注意力和单控制器数据并行支持。

## 正文

目录

摘要

方案：优化融合 MoE 内核

1. MoE 内核成本模型

计算下界

ICI token 路由下界

HBM 权重搬运下界

要点

2. 为何需要 Pallas 融合内核

3. V1 版本：融合，但隐藏层分块零散

4. V2 版本：VMEM 驻留与权重双缓冲

5. V2 专项优化

逐通道 directscaleddot

激活量化

内核内共享专家

6. 性能提升来源

7. V2 版本之后的待改进

通信受拓扑限制

路由计算受 VPU/VMEM 限制

总结

Ling-2.6-1T 启动

混合内存池

GLA（门控线性注意力）

单控制器数据并行支持

实验与基准测试

基准测试配置

局限性与未来工作

附录

成本模型中使用的 TPU v7x 规格

性能复现

服务启动与精度复现

参考文献

致谢

在 TPU 上用 SGLang-JAX 优化 Ling-2.6-1T 服务：用一个 Pallas 内核将 MoE 数据搬运隐藏在计算背后

作者：Prayer, JamesBrianD, Haolin Fu, Haoguang Cai, Qinghan Chen。2026 年 6 月 17 日

SGLang-JAX 现已在 TPU v7x 上支持 inclusionAI 的 Ling-2.6-1T 高效推理。在建立可用基线后，性能分析指出混合专家（MoE）路径是主要瓶颈：每一层将 token 分散到 32 个 JAX 设备（每个 v7x 芯片两个），运行专家 FFN，再收集输出结果。本文首先聚焦于 Fused MoE V2，这是一个新的 Pallas 内核，它融合了分散、专家 FFN 和收集过程，同时将 TPU 计算与数据搬运重叠。

借助 Fused MoE V2，MoE 预填延迟从 5.16 ms 降至 2.42 ms；在同一 SGLang 解码基准测试中，16 个 TPU v7x 芯片的输出吞吐量达到 16 块 H200 GPU 的 1.29–1.77 倍。完整数据如下。

图 1. TPU v7x-16 与 H200×16 上 Ling-2.6-1T 的解码吞吐量对比，使用 SGLang 默认的 `random` 基准数据集（从 ShareGPT 采样），输入 16,384 token，输出 1,024 token。

摘要

Fused MoE V2：相比 Fused MoE V1，MoE 预填延迟下降 53%（从 5.16 ms 降至 2.42 ms）；解码内核延迟下降约 15%（从 0.249 ms 降至 0.211 ms）。

端到端收益：仅替换 MoE 内核即可使预填吞吐量提升 24.8%，解码吞吐量提升 18.5%–35.3%。

TPU vs H200 解码：TPU v7x-16 在 mc=128 时的解码输出吞吐量是 H200×16 的 1.29 倍，在 mc=512 时则达到了 1.77 倍。

MoE 之外：Ling-2.6-1T 的完整启动还包括混合 KV/循环记忆池、GLA 线性注意力以及单控制器数据并行。

Ling-2.6-1T 概览：一个拥有 1T 参数的稀疏 MoE 模型，每个 token 激活 63B 参数，256 个路由专家采用 top-8 路由机制并包含一个共享专家，使用每通道 fp8 MoE 权重，以及混合 MLA + Lightning Linear 骨干网络。MoE 结构驱动了本文前半部分的内核工作；混合骨干网络则引出了后续关于记忆池和 GLA 启动的章节。

配置：优化融合 MoE 内核

本节中所有 MoE 数据均来自 jax.profiler 设备轨迹，除非另有说明。配置环境为 16 芯片 TPU v7x 切片：ep=32，一个 2×2×4 的 ICI 环面，每芯片两个设备。工作负载为 Ling-2.6-1T，包含 16,384 个 token 的预填充和 512 个 token 的解码批次，使用每通道 fp8 MoE 权重。本节中的所有下界均按每设备计算，约为芯片级计算和带宽的一半；详细芯片规格见附录。

融合 MoE V2 通过改变路由 token、专家权重和累加器在 VMEM、HBM 和 ICI 之间的移动方式来实现这一目标。

1. MoE 内核成本模型

Ling-2.6-1T 每层拥有 256 个路由专家和一个共享专家，采用 top-8 路由。使用 ep=32 时，每个设备拥有 8 个本地路由专家。一个 token 选中的 8 个专家通常分布在不同的设备上，因此每一层路由路径的形状都相同：

scatter tokens -> local expert FFN -> gather results

采用这种结构，MoE 的运行成本远不止 GEMM 的 FLOPs 那么简单。内核需要通过三个高成本路径来移动数据：跨芯片的 token 路由、从 HBM 到 VMEM 的专家权重读取，以及矩阵乘法单元（MXU）周围围绕 fp8 布局/缩放因子的处理。

共享专家是一条本地的稠密路径。它增加了本地的 FFN 计算量，但不参与路由的 all-to-all 通信，因此对 token 路由的数据量影响很小。

计算下界

在预填充 16,384、top-8 路由和 ep=32 的情况下，每个设备处理：

16384 * 8 / 32 = 4096 routed rows / device

平均而言，8个本地路由专家各处理约512行。共享专家不通过top-k路由扇出；它在本地4096行上运行一次。路由+共享FFN的计算量为：

FFN1: 8 experts * 2 matrices * (2 * 512 * 8192 * 2048) = 274.9 GFLOP
FFN2: 8 experts * 1 matrix * (2 * 512 * 2048 * 8192) = 137.4 GFLOP
Routed total: 412.3 GFLOP / device
Shared expert: 3 matrices * (2 * 4096 * 8192 * 2048) = 412.3 GFLOP
Total: 824.6 GFLOP / device

TPU v7x 的公开规格显示每颗芯片的 fp8 算力约为 4.614 PFLOP/s。在此部署中，每颗芯片暴露为两个设备，因此每个设备的粗略 fp8 峰值算力为 2.307 PFLOP/s。理想计算下界为：

824.6 GFLOP / 2307 TFLOP/s = 0.36 ms

这是一个理想下界，未考虑数据搬运、fp8 打包/解包以及在向量处理单元（VPU）上的缩放处理。实测的生产追踪记录为 2.42 毫秒，仍比该下界高约 7 倍，因此纯粹的 GEMM FLOP 无法解释延迟成因。

ICI token 路由下界

每个设备的散列（scatter）载荷为：

4096 rows * 8192 hidden = 33,554,432 elements
bf16: 67.1 MB
fp8 : 33.5 MB

TPU v7x 每颗芯片具备 1.2 TB/s 的双向 ICI 带宽，折算下来每条链路上每个方向约 100 GB/s。2×2×4 环面拓扑使每颗芯片拥有 4 条有效链路，因此有效单向芯片带宽约为 4 × 100 GB/s = 400 GB/s。由于两颗设备共享一颗芯片，每台设备的粗略单向注入带宽约为 200 GB/s。

仅考虑注入带宽，忽略跳数和争用，下界为：

载荷单次散列散列 + 收集

fp8 33.5 MB0.17 毫秒0.34 毫秒

bf16 67.1 MB0.34 毫秒0.67 毫秒

但 all-to-all 并非单链路带宽测试。在 2×2×4 环面上，随机目的地的平均跳数约为两跳：x 轴约 0.5 跳，y 轴约 0.5 跳，z 轴约 1.0 跳。考虑此跳数因子后，拓扑调整后的下界更接近：

载荷单次散列，平均跳数调整后散列 + 收集

fp80.34 毫秒0.67 毫秒

bf160.67 毫秒1.34 毫秒

这仍未计入链路争用、DMA 小粒度开销、运行时开销以及 fp8 布局处理。即便如此，token 路由的时间范围已与路由+共享的理想计算下界处于同一量级，且远高于纯路由的计算下界。

HBM 权重搬运下界

现在考虑路由专家权重的 HBM 读取成本。如果权重预取未被流水线隐藏，该成本将立即显现。

一个本地专家的 fp8 权重为：

W1 + W3 + W2 = 3 * 8192 * 2048 bytes = 50.3 MB
8 local experts = 402 MB

共享专家增加了另一组本地 FFN 权重集，大小约等于一个本地专家，但不会引入全互联流量。以下估算聚焦于路由专家路径。

TPU v7x 的 HBM 带宽约为每芯片 7.38 TB/s，即每设备约 3.69 TB/s。一次性读取全部 8 个本地专家的下限为：

402 MB / 3.69 TB/s = 0.11 ms

在实践中，内核每处理一个 token staging tile 会重新读取一次权重。tile 大小由 bts（block token staging size）决定，即每次为一个专家 FFN tile 带入 VMEM 的路由 token 行数。Ling prefill 使用 bts=160。由于每个专家处理约 512 行，prefill 需要 ceil(512 / 160) = 4 个 token staging tile。V2 跨这些 tile 流水线预取权重，因此 HBM 读取下限约为：

4 * 402 MB / 3.69 TB/s = 0.44 ms

权重读取不必处于关键路径上。V2 通过双缓冲将其隐藏在 MXU 窗口之后。这些数字解释了为什么需要这样的调度：如果 HBM 读取在 GEMM 之前串行化，它们已经超过了纯计算下限。

要点总结

TPU 上的 MoE 主要是一个数据传输和重叠问题：

路由 + 共享专家 FFN 计算下限：约 0.36 ms；

fp8 scatter + gather 拓扑下限：约 0.67 ms；

专家权重 HBM 读取下限：约 0.11 ms per tile，或约 0.44 ms（bts=160 时）；

fp8 打包、缩放广播和布局重排仍消耗 VPU 和 VMEM 带宽。

优化目标不是减少 FFN FLOPs，而是将 token 路由、权重预取和 fp8 重排隐藏在路由计算窗口之后。

2. 为何需要 Pallas 融合内核

本节其余部分使用一些 TPU 术语。简化的示意图是：TensorCore 包含 MXU、VPU 和 VMEM；HBM 位于芯片外部；芯片通过 ICI 通信。

图 2. 本节使用的简化 TPU 执行模型，改编自 JAX Scaling Book 的 TPU 概述。

在 MoE 内核中，这些单元对应以下工作：

硬件单元TPU 角色MoE 中的工作

MXU矩阵乘法单元路由专家的 W1/W3 门控-升维 GEMM 和 W2 降维 GEMM

VPU向量数学、归约、布局操作SiLU、门控乘法、缩放乘法、fp8 打包/解包、通道重排

VMEM靠近 MXU/VPU 的芯片内暂存器路由 token 分片、专家中间结果、输出累加器、预取权重分片

HBM每个芯片连接的大容量片外存储器专家权重、token 暂存缓冲区、大容量中间缓冲区

HBM-DMAHBM ↔ VMEM 数据传输将当前/下一个专家权重预取到 VMEM；在需要时移动暂存缓冲区

ICI / ICI-DMATPU slice 内部的直接芯片间网络在源芯片和目标芯片之间移动路由 token 数据；将 token 分散给专家所有者，并将输出聚合回 token 顺序

纯 JAX 原生的 MoE 可以正确表达路由、专家 FFN 和输出聚合。但它无法暴露单个 MoE 层内的细粒度调度。一旦 scatter、专家 FFN、HBM 权重移动、fp8 布局工作和 gather 跨越多个 JAX op 或 collective 边界，XLA 就无法可靠地将 ICI-DMA、HBM-DMA、MXU 和 VPU 工作安排到一条手动调度的流水线中。

这条路径也不能被当作独立的稀疏查找并卸载到 SparseCore：由 scatter 产生的局部专家 token 布局、每个专家的偏移量、专家输出以及最终的 token 顺序都相互依赖。有用的优化空间在 MoE 内核内部。

图 3. 具有串行通信和计算阶段的朴素融合流水线。语义正确，但引擎未以细粒度重叠方式调度。

理想的稳态是：当 MXU 计算专家 i 时，HBM-DMA 预取专家 i+1 的权重，ICI-out 发送下一批路由 token，ICI-in 接收上一批输出，而 VPU 处理来自前一个矩阵乘法的缩放和布局工作。

为了表达这种调度，scatter、专家 FFN 和 gather 需要位于同一个 Pallas 内核中。融合的主要目的不是减少操作数量；它创造了一个调度空间，使得依赖阶段可以在 MXU、VPU、HBM-DMA 和 ICI-DMA 之间手动排列。

3. V1：融合，但带有碎片化的隐藏分块

我们的起点是融合 MoE V1，最初由 Jevin Jiang、Kyuyeun Kim 等人在 tpu-inference 项目 [4] 中提出并优化，后经部分修改以 FusedEPMoE 的形式适配到 SGLang-JAX [5] 中。V1 已经将 scatter、专家 FFN 和 gather 合并到一次 Pallas 调用中，并在每个设备上执行 8 个本地专家。这满足了核内通信/计算调度的前提条件，但 V1 仍未达到上述理想稳态。

问题出在专家内部。一个 MoE 专家需要的远不止输入 token 分片和一个 GEMM 输出。为了实现通信与计算重叠，核还需要权重暂存缓冲区、中间激活值、输出累加器以及 DMA 双缓冲。当 Ling 的隐藏层维度为 8192 时，将完整隐藏维度常驻内存会很快耗尽 VMEM，尤其是 f32 累加器和 W1/W3/W2 权重的暂存。

因此 V1 采用了一条保守路径：对隐藏层维度进行切片，通过 VMEM 流式处理更小的工作集。

对于 Ling 16,384 的预填充，V1 的配置如下：

bf=1024 / bd1=512 / bd2=512 / bts=128 / btc=128

该分块配置回答了一个布局问题：哪些 token 行、中间通道和隐藏通道留在 VMEM 中，哪些从 HBM 流式读取。

这些参数可以理解为沿 GEMM 轴方向的分片尺寸：

参数控制性能含义

bts为一个专家分片暂存到 VMEM 中的路由 token 行数控制 M；如果过小，DMA / VPU / MXU 的固定开销无法被摊平

btcbts 内部一次计算循环馈入的 token 行数内部 M 计算分片；不能超过 bts，通常整除 bts

bfW1/W3/W2 的中间通道数控制 FFN 中间分片；越大通常 MXU 窗口越长，但 VMEM 占用也越大

bd1FFN1 隐藏层 K 的归约切片V1 对隐藏层 K 进行切片；bd1 越小，FFN1 点积次数越多、每次更小

bd2FFN2 隐藏层输出 N 的切片V1 对输出隐藏层进行切片；bd2 越小，部分输出会更频繁地经 HBM 往返

因此，bf/bd1/bd2 主要控制特征/隐藏层维度，而 bts/btc 控制每个专家的 token 行数。这些参数共同决定了分片能否放入 64 MB 的 VMEM 预算，以及 MXU 周围能重叠多少 HBM-DMA / VPU 工作。

V1 承受了三种结构性代价：

代价V1 行为为何有害

FFN1 点积过小bd1=512；经过 fp8 打包后，有效 K 约为 256，因此 V1 扫描 16 个隐藏维度切片。vmatmul 的固定开销难以被充分平摊。

token 分段过于频繁。num_bf * num_bd1 * num_token_tiles = 2 * 16 * 4 = 128 次 HBM→VMEM 分段。大量小型 DMA 和布局步骤。

FFN2 部分结果溢出到 HBM。部分输出被写入 a2a_s_acc_x2_hbm，随后被读回用于后续的 bf 累加。HBM 读-改-写操作割裂了关键路径。

V1 存在一些微观重叠，但隐藏维度切片使得重叠窗口很小。预取每次只覆盖一个小切片，而 FFN2 的部分输出仍需经过 HBM 往返。V1 的预填充延迟为 5.16 毫秒。

4. V2：VMEM 驻留与权重双缓冲

V2 不仅仅是 V1 的更大分块。它改变了张量的生命周期。V1 的循环遍历隐藏维度切片；V2 则在 FFN 循环期间将路由后的 token、门控/中间层结果以及输出累加器驻留在 VMEM 中，同时 W1/W3/W2 通过双缓冲从 HBM 流式加载。

这会在长寿命张量上消耗更多 VMEM，但消除了大部分隐藏切片分段，并几乎完全消除了 FFN2 的 HBM 读-改-写路径。

Ling 16,384 预填充 V2 生产配置如下：

bf=512 / bts=160 / btc=80

V2 没有 bd1 或 bd2，因为它不再对隐藏维度进行切片。结构上的变化是：

每个专家V1V2效果

FFN1 点积每次硬件点积的有效 K 约为 256fp8 分块 K 约为 2048；4 个分块覆盖整个隐藏维度K 增大约 8 倍

W2 输出bd2=512，每次产生一个窄的隐藏切片输出分块约 4096 个隐藏通道N 增大约 16 倍

token 分段128 次小分段约 4 次全隐藏分段分段次数减少约 32 倍

FFN2 累加器部分结果通过 HBM 溢出/重新加载b_y_acc_vmem 在 VMEM 内部跨 bf 累加HBM 读-改-写基本消失

这也解释了为什么单纯在 V1 中增大 bd1/bd2 是不够的。在 V1 中，更大的隐藏分块同样会增大权重缓冲区、token 分段缓冲区和部分输出分段缓冲区，从而迅速触及 64 MB 的 VMEM 上限。更重要的是，V1 仍然会遍历隐藏切片，它并没有让 token 和输出累加器驻留。

通过这种 VMEM 驻留工作集，V2 获得了更大的 MXU 瓦片、更少的 HBM 溢出以及更长的路由计算窗口。在激活量化之前，V2 在设备跟踪中已将预填充延迟从 5.16 毫秒降至 3.02 毫秒。启用激活量化和内核内共享专家重叠后，生产环境跟踪达到 2.42 毫秒，比 V1 降低了约 53%。

解码遵循相同的逻辑，但空间更小。对于 512 token 的解码批次，内核延迟从 0.249 毫秒降至 0.211 毫秒，降幅约 15%。每个专家的有效 M 维度较小，因此 MXU 瓦片无法很好地分摊固定开销；该路径也更接近专家权重 HBM 读取下限，解码跟踪已达到约 80% 的 HBM 带宽利用率。因此 V2 仍然有助于解码，但无法像预填充那样实现完整的 VMEM 驻留和路由窗口增益。

图 4. V1 和 V2 融合 MoE 的概念时间线。V1 由于隐藏维度切片频繁循环，只形成较小的重叠窗口；V2 将 token 和累加器驻留在 VMEM 中，双缓冲专家权重，并将大多数 scatter/gather 流量隐藏在路由计算窗口之后。

5. 针对性的 V2 优化

逐通道 direct_scaled_dot

fp8 权重量化的缩放粒度决定了 MXU 是看到一个大的 GEMM 还是一系列小的 GEMM。

采用逐块量化时，缩放因子取决于 K 块：

out[m,n] = sum_k A[m,k] * W[k,n] * scale[block(k),n]

缩放因子无法从归约中提取出来，因此必须将 K 拆分为多个块。每个块执行一个小型 fp8 点积，乘以该块的缩放因子，然后累加。一个大的 GEMM 变成了许多较小的 GEMM，并在它们之间插入 VPU 工作。

采用逐通道量化时，缩放因子仅取决于输出通道：

out[m,n] = (sum_k A[m,k] * W[k,n]) * scale[n]

缩放因子可以在归约之后应用。V2 的 direct_scaled_dot 将 fp8 token 和 fp8 权重直接送入 MXU，得到 f32 部分和，然后才应用逐 token / 逐通道缩放。Ling 的 MoE 权重使用逐通道缩放因子，因此该路径可用。

这保留了完整的 K 点积，并避免将大型 GEMM 切分为缩放块。剩余的开销是 fp8 子字打包、缩放广播和通道重排。逐块量化则会额外增加 K 分段和块间缩放处理。

激活量化

V2 在分散之前将激活值从 bf16 量化到 fp8，直接将路由 token 的数据量减半。在 Ling 16,384 预填充中，内核内分散阶段从 1.39 毫秒降至 0.65 毫秒。

这与上文 ICI 下界计算一致：当数据量从 bf16 67 MB 降至 fp8 33.5 MB 时，通信下界几乎减半。

Ling-2.6-1T 支持激活量化，因此 V2 使用动态逐 token fp8，我们在评估中未观察到精度下降（参见附录中的 AIME 2026 检查）。

内核内共享专家

Ling 每层还有一个共享专家。如果作为独立的稠密 MLP 运行，它会增加自身的关键路径段。V2 将共享专家移入同一内核，复用路由专家的 token/权重 VMEM 缓冲区，并将其调度到分散窗口内。

共享专家自身的计算约为 0.159 毫秒，但仅给关键路径增加 0.068 毫秒（约 2.7%）。原因很简单：共享专家不需要跨芯片 token 分发；所有必需的 token 都在本地，因此它可以与路由 FFN 之前的分散阶段重叠。

6. 增益从何而来

下面的分解图展示了启用激活量化和内核内共享专家后在预填充 16,384 时的关键路径。阴影区域是被其他阶段隐藏的实际工作。

图5. Fused MoE V2 的实测重叠结构。大部分分散/收集流量被隐藏在路由专家窗口下；只有分散头部和收集尾部可见。

元数据块是路由记账：token 到专家/设备映射、每专家偏移量/计数以及分散/收集索引。它仅传送小量元数据，耗时数十微秒，并非核心预填充开销。

对同一 V2 内核进行消融实验，展示关键路径上仍暴露出的部分：

消融 / 组件结果解读

完整 V2 生产版本2.42 毫秒本节使用的标准 MoE 预填充延迟

禁用所有专家矩阵乘法与完整版本相比 -2.2%纯 MXU 计算并未暴露

可见的 scatter 操作0.42 毫秒通信前导部分仍处于关键路径上

可见的 gather 操作0.18 毫秒gather 尾部仍处于关键路径上

无计算可隐藏的 scatter + gather 操作约 2.4 毫秒在重叠之前，实际通信延迟接近完整内核延迟

这与成本模型分析结果一致。即便计入共享专家，理论计算下限也仅为约 0.36 毫秒，而移除矩阵乘法后总延迟几乎没有变化。scatter/gather 操作耗时接近 2.4 毫秒，但其中约 1.8 毫秒被路由计算的窗口隐藏掉了。

因此 V2 的收益来自三种机制：

token 和累加器常驻 VMEM，减少了 token 暂存和 HBM 读-改-写操作；

专家权重采用双缓冲，因此 HBM 读取可以隐藏在 MXU 计算之后；

scatter/gather 利用片内缓冲器和出站/入站 ICI 通道，与路由计算重叠进行。

7. V2 之后还剩什么

经过重叠之后，图 5 中最长的剩余段是路由计算窗口，约占 2.42 毫秒总延迟的 68%。这并不意味着问题又回到了纯算力瓶颈：Mosaic LLO 转储显示，剩余瓶颈主要是 fp8 打包/通道重排/缩放广播，以及 VMEM 对 tile 尺寸的限制。

通信受拓扑限制

在我们的测量中，扁平 all-to-all 优于分层 all-to-all。在扁平配置下，发送/接收分区直接由最终专家设备决定，一次 32 路 all-to-all 即可将路由 token 载荷从源设备发送到最终目标设备。

我们还测量了分层配置：将 32 设备的交换沿 2×2×4 ICI 环面展开，先在本地维度内重排，然后沿下一维度中继，直至每个 token 到达目标专家设备。每轮通信范围更小，但相同的路由 token 载荷要经过多个中继阶段，从而增加了暂存缓冲区、同步边界，并且总移动字节数几乎翻倍。两种模式均在融合内核之外作为独立的 all-to-all 基准测试进行测量，因此这些数字无法直接与内核内的轨迹数据相比较。

模式（载荷 = 16384 × 8192 × dtype_size）bf16fp8

扁平 all-to-all2.09 毫秒1.34 毫秒

层级化全对全3.12 毫秒1.88 毫秒

因此，通信侧的实际杠杆并非更复杂的路由算法，而是更少的字节数和更好的重叠。激活量化正是实现这一目标的手段。

路由计算受限于 VPU / VMEM

路由 FFN1（W1+W3）的测量耗时约为 0.72 毫秒，而理想稠密 fp8 GEMM 下限约为 0.12 毫秒。这一差距并非由激活量化导致：开启激活量化时 FFN1 约为 0.74 毫秒，关闭时约为 0.71 毫秒。

分块扫描也表明当前配置接近局部最优：

bts / btc内核延迟VMEM

160 / 802.42 毫秒47 MB

160 / 1602.44 毫秒47 MB

128 / 1283.12 毫秒44 MB

256 / 1283.19 毫秒54 MB

256 / 2563.23 毫秒54 MB

384 / 128内存溢出62 MB

Mosaic LLO 转储解释了原因。整个内核只有 4096 条真正的 vmatmul 指令，而 fp8 布局和向量侧准备工作主导了指令流：

LLO 指令计数角色

vselect50880子通道选择 / 混合

vbitcast46566fp8 子字重解释

vcombine36380子通道合并

vpack_format34368MXU 输入打包

slane29960子通道移动

vunpack25600fp8 解包

matmul_data_format25600MXU 前的格式转换

vrot21524通道旋转

vmatres / vmatprep17408 / 10240MXU 排空 / 喂入

vslreplicate6032缩放因子广播

vmatmul4096实际矩阵乘法

V2 避免了逐块量化的 K 切片操作，但 fp8 子字打包、缩放因子广播以及 MXU 的喂入/排空仍然消耗了大量的 VPU / 布局工作。由于 VMEM 上限为 64 MB，bts 无法持续增长；在分块较小的情况下，这些固定成本无法被摊薄。

总结

在 V2 隐藏了大部分显式通信和 HBM 权重移动之后，剩余瓶颈仍然是数据移动，只是换了一种形式：fp8 布局工作、VMEM 容量压力以及保持 MXU 持续被喂入。

ICI 全对全受限于环形拓扑和争用。

HBM 权重读取必须通过双缓冲来隐藏。

fp8 打包和缩放因子处理导致 MXU 等待数据成形。

VMEM 容量限制了分块大小以及可同时存在的重叠缓冲区数量。

下一步必须改变约束条件本身。

内核侧：减少 fp8 的打包/解包与缩放处理，但这越来越依赖于将模型量化与 TPU 原生执行格式对齐：TPU 友好的缩放粒度、fp8 布局，或未来 MXU 原生的低精度格式（如 FP4 或 MXFP8）。

工作负载侧：跨批次重叠，使得路由窗口能与其他层的工作并行执行。

硬件侧：提供更优支持全连接（all-to-all）的互连拓扑，或提供更大的 VMEM / 更高的 ICI 带宽。

关于未来的 TPU 硬件，请参阅 Google Cloud 的 TPU 8t 和 TPU 8i 技术深度解读。

Ling-2.6-1T 的上线调试

MoE 融合只是让 Ling-2.6-1T 在 TPU 上高效服务的其中一环。其余的上线调试工作在于将运行时与模型的混合骨干架构匹配：为全注意力层和线性注意力层分别分配不同状态，通过 TPU 友好的内核运行 GLA 的预填充和解码阶段，并映射 DP/TP 使得分组的 RMSNorm 保持在芯片本地。

混合内存池

Ling-2.6-1T 并未向运行时暴露单一统一的注意力状态。其 10 个 MLA 全注意力层写入基于 token 索引的 KV 缓存，而 70 个 Lightning / GLA 层则携带基于请求索引的循环状态。因此内存分配器必须同时管理两种不同的容量：MLA 的常驻历史 token，以及线性注意力层的活跃请求槽。

这两个单位的对比容易误读。在 TP=4 的情况下，使用 bf16 KV 和 fp32 循环状态，MLA KV 缓存每设备每 token 在 10 个全注意力层上的成本约为 12.5 KiB。Lightning 循环状态在 70 个线性层上每设备每请求的成本约为 70 MiB。只有将这些数字放回到一个请求中时才有意义：一个 16K token 的提示词每请求大约需要 200 MiB 的 MLA KV，一个 256K token 的提示词大约需要 3.1 GiB，而循环状态则保持在 70 MiB 左右。循环状态是固定的并发成本；KV 缓存是与上下文长度线性增长的 token 容量成本。

SGLang-JAX 将这些状态类型分开，同时保持单个请求的生命周期：HybridLinearKVPool 仅存储 10 个全注意力层的 KV（那 70 个线性层不消耗 KV 槽位），RecurrentStatePool 为每个活跃请求存储一个 fp32 的循环槽位，而 HybridReqToTokenPool 将它们关联在一起：一个请求在准入时同时获取两者，在结束时同时释放。分块预填充和解码继续使用同一个循环槽位，而不是为每个块或每个 token 分配新状态。HBM 预算也以相同方式划分：可配置的一部分保留给循环槽位，这限制了并发数；其余部分用于 KV 缓存，这限制了驻留 token 数。

JAX 增加了一个额外约束：运行时不能像 CUDA 路径那样原地更新这些缓冲区。SGLang-JAX 将 KV 池和循环池封装在一个 MemoryPools pytree 中，并将其作为捐赠的 JIT 参数传递给模型。每次前向传播返回更新后的池缓冲区，运行时通过 `replace_all()` 将其写回。这使得缓冲区捐赠、TP/DP 分片以及未来的池扩展保持在容器级别，而不是在前向循环中散布特殊情况。

GLA（门控线性注意力，Gated Linear Attention）

每个 GLA 层 [7] 将历史保存在固定大小的循环状态中，而不是为每个过去的 token 存储一个 KV 条目。其更新可以写作：

这将注意力历史从逐 token 增长的形式转变为每个活跃请求一个状态张量。在长上下文场景下，这是主要优势：携带历史在计算上保持线性，在状态上保持固定大小，而不是物化并读取不断增长的 KV 历史。

预填充：使循环过程足够并行以适配 TPU。从字面上看，上述循环是串行的：token t 依赖于来自 token t-1 的衰减和更新后的状态。以此方式运行预填充会将 16K 或 256K 的提示词转化为漫长的逐 token 扫描，这恰好是 TPU 不擅长的形态。

SGLang-JAX 使用数学上等价的逐块（chunk-wise）形式。序列被切分为固定大小为 64 个 token 的数据块。跨数据块时，一个数据块的最终状态成为下一个数据块的初始状态，因此长程依赖仍然随时间向前传递。但在数据块内部，递归被重新排列为对 token 块进行的密集矩阵运算。只有数据块边界保持串行；块内的工作则以块并行（block-parallel）TPU 计算的方式运行。

解码：递归的自然形式。解码更简单：预填充已将提示词折叠到循环状态中，因此每个新 token 读取请求的当前状态，执行一次循环更新，输出注意力结果，并写回新状态。问题从长序列并行转变为高效的少量状态更新。

服务集成：将 GLA 保留在相同的运行时路径中。GLA 作为层级的后端选择进行集成，而非独立的调度器模式。全注意力层读写 KV 缓存；GLA 层读写循环状态；两者都在相同的预填充和解码批次中推进。调度器仍然只看到一个生命周期：准入、预填充、解码、释放。

这种集成本身功能上是完整的，但预填充内核尚未达到与 Fused MoE V2 同等的调优程度。GLA 的数学运算无需改变，需要改变的是执行调度。

单控制器数据并行支持

Ling-2.6-1T 的分组后注意力 RMSNorm 对张量并行施加了硬性约束。每个归一化组包含 8 个注意力头。如果一个组跨芯片分布，每一层的方差计算都会变成跨芯片的归约（all-reduce），直接位于解码关键路径上。因此纯 TP 没有好的设置：tp ≤ 8 可以保持归一化组位于芯片内部，但对于万亿参数模型而言并行度不足；而 tp > 8 则会拆分归一化组并付出全归约（all-reduce）代价。

单控制器数据并行通过将数据并行视为另一条网格轴来解决这一矛盾。网格被划分为若干数据并行组；每个组使用规模足够小的张量并行，使得分组后的 RMSNorm 保持在芯片本地，并且请求被分配到不同的数据并行秩上。权重在每个数据并行组内保持张量并行分片。逐层的归一化归并操作被消除，释放出的 ICI/HBM 预算可用于更高的并发度。

重要的设计选择在于，数据并行是 SPMD 运行时的一部分，而不是一组独立的服务器副本。SGLang-JAX 运行一个逻辑调度器，而 dp_rank 被附加到请求、KV 分配和前缀缓存键上。这使得系统能够从一个负载快照实现全局准入控制、跨主机的确定性批构建，以及一个以 (dp_rank, prefix) 为键条目的全局前缀缓存结构。

这种方式还能与混合运行时的其余部分干净地组合。将网格扩展到更大的配置（例如在每个数据并行组内增加张量并行）只是网格形状的变化，而非调度器的分叉，因此内存池、批处理路径和注意力后端的思维模型保持不变。

实验与基准测试

所有 TPU 结果均使用 SGLang-JAX 在单个 TPU v7x 切片上服务 Ling-2.6-1T 模型；V1/V2 消融实验的设置完全相同，仅 MoE 内核配置不同。

基准测试配置

硬件：TPU v7x，16 芯片（2×2×4 ICI 环面）→ 32 设备

并行度：tp = ep = 32, dp = 8

模型：Ling-2.6-1T，bf16 激活值，每通道 fp8 MoE 权重

数据集：SGLang 默认随机基准数据集（从 ShareGPT 采样）

运行时：SGLang-JAX (JAX 0.8.1)，dvfs p_state=7

输入长度：16384

预填充：输出 1，并发数 128

解码：输出 1024，并发数 128 / 512

预填充输入吞吐量，输入长度 16384 个 token，mc=128。设置完全相同，仅 MoE 内核配置不同：Fused v1 → v2 base → v2 +act-quant → v2 +act +SE-overlap（提升 24.8%）。

峰值输出（解码）吞吐量，输入长度 16384 个 token，输出 1024，针对 np=512/mc=128 和 np=2048/mc=512。百分比表示相对于 Fused v1 的增益。

图 6. TPU 与 GPU 完整对比：TPU v7x-16（fused_v2） vs GPU H200×16（2 节点，tp8·pp2），相同模型和 SGLang 基准工作负载，每侧 16 个加速器。关于预填充差距的说明见下方注释。

关于端到端预填充与 MoE 核心加速的说明：Fused MoE V2 核心可将 MoE 层预填充延迟降低约 53%（设备追踪数据），但端到端预填充吞吐量仅提升约 25%（v1 到 v2）。MoE 层已不再是预填充的主要成本：GLA（门控线性注意力）预填充核心目前是主要瓶颈，且尚未得到同等程度的优化，因此它稀释了端到端预填充的加速效果。同样的瓶颈也解释了为什么在图 6 中，TPU v7x-16 在预填充指标上落后于 H200×16，却在两个解码指标上领先。将 GLA 预填充核心优化到同等水平是正在进行的工作，我们预计这将释放更大的端到端预填充收益。

局限性与未来工作

我们对 Ling-2.6-1T 的支持是针对本次发布特意设定的范围；以下几项内容作为后续工作，我们正在积极推动：

GLA / 线性注意力预填充核心。正如基准测试部分所指出，GLA（Lightning Linear）预填充核心目前已成为预填充的主要成本。通过考虑更好的分块/平铺方法、融合门控与循环状态更新，以及应用与 MoE 核心相同的 MXU/VPU/DMA 重叠处理手段来将其优化到同等水平，是提升端到端预填充性能最直接有效的剩余手段。

动态专家并行负载均衡（EPLB）。当前的 FusedEPMoE 路径采用静态的专家到设备映射，但实际工作负载中，256 个路由专家的命中率并不均匀。一个动态的 EPLB 过程，能够根据观测到的流量定期重新平衡专家到进程的映射，将缩小设备峰值利用率与平均利用率之间的差距，尤其是在较大批量下。

混合内存池上的 Radix 缓存。SGLang 的 RadixAttention [9] 前缀缓存假设存在一个单一的每 token KV 池，而 Ling-2.6-1T 将每 token KV 与每请求循环状态混合在一起，因此简单的基于前缀共享会在线性层上错误地跨请求混合状态。我们正在设计一种扩展方案，该方案按 token 前缀共享 MLA KV，同时为每个共享前缀对循环状态进行快照并重新生成键值，这样共享的系统提示词和较长的智能体轨迹就可以被重用而不损失正确性。

MTP / EAGLE 推测解码。Ling-2.6-1T 检查点内置了一个 EAGLE 风格的 MTP 头部（3 步推测，4 个草稿 token，top-k 1）。我们当前的路径仅运行基础模型解码；将 MTP 头部与 SGLang-JAX 的推测解码运行时集成是解码吞吐量的下一个里程碑。混合内存池层已经考虑了草稿步状态，因此剩余工作集中在验证器和草稿接受核函数上。

附录

成本模型中使用的 TPU v7x 规格

TPU v7x 公开规格列出每芯片约 4.614 PFLOP/s 的 fp8 算力、7.38 TB/s 的 HBM 带宽以及 1.2 TB/s 的双向 ICI 带宽。在此部署中，每个芯片被暴露为两个设备，因此成本模型部分中的每设备下限使用了大约一半的芯片级算力和带宽。关于 TPU 内存层次结构和执行单元（MXU、VPU、VMEM、HBM、ICI）的背景信息，请参阅 Google Cloud 的 TPU 系统架构。

性能复现

双方运行相同的模型和相同的 SGLang 基准测试工作负载：预填充（out 1, mc 128）· 解码（out 1024, mc 128）· 解码（out 1024, mc 512）。

TPU：SGLang-JAX（融合 MoE V1 / V2）。TPU v7x，16 芯片（2×2×4 ICI 环面 → 32 设备），tp = ep = 32，dp = 8，每通道 fp8 MoE 权重。

TPU 运行使用 sgl-jax 分支 fused-moe-v2-with-sp-rs @ 49c2ed1 和镜像 jax-ai-image/tpu:jax0.8.1。

V1/V2 消融实验仅更改 MoE 标志：Fused v1 = --moe-backend fused；v2 基础 = fused_v2 --no-moe-fused-act-quant --no-moe-fused-shared-experts。

v2 +act-quant 情况添加 --moe-fused-act-quant；v2 +act +SE-overlap 同时开启两者。两个外部共享专家配置使用 --mem-fraction-static 0.85，因为它们在 0.88 时内存不足（OOM）。

GPU：SGLang（H200×16，参考）。2 节点 × 8× H200，tp = 8，pp = 2；与 TPU 运行相同的模型和基准测试工作负载。

性能运行的完整基准测试命令见 SGLang-JAX 食谱。

服务器启动与准确率复现

AIME 2026 测试使用了 MathArena/aime_2026 数据集，共 30 道题，pass@1 结果为 26/30 = 86.7%。该运行过程中零请求错误，所有响应均正常终止（finish_reason=stop，未在 32768 token 处截断）。这表明 fp8 融合 MoE 推理路径没有明显的精度回归。

完整的启动服务命令、请求与工具调用示例，以及 AIME 2026 精度复现方法，均收录于同一份 SGLang-JAX 烹饪指南中。

参考文献

[1] Ling-2.6-1T 模型卡

[2] 混合模型与 SGLang（博客）

[3] 参差不齐的分页注意力

[4] 融合 MoE V1 内核，tpu-inference

[5] 适配于 SGLang-JAX 的融合 MoE V1 内核

[6] DeepSeek-V2（MLA）

[7] 门控线性注意力（GLA）

[8] MiniMax-01（闪电注意力）

[9] SGLang（RadixAttention）

致谢

蚂蚁集团-系统核心技术团队：潘振轩、王国伟、郭宇红、万烁

SGLang-JAX 团队：jimoosciuc、Prayer、aolemila、neo、leos、pathfinder-pf、郝琳付、陈庆涵、JamesBrianD、蔡浩光、胡宇豪、cjx0709、周正科、魏宇鑫、王联芳、0xaskr