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

在 TPU 上用 SGLang-JAX 优化 Ling-2.6-1T：用一个 Pallas kernel 把 MoE 数据搬运隐藏在计算之下

Prayer、JamesBrianD、Haolin Fu、Haoguang Cai、Qinghan Chen

SGLang-JAX 现已支持在 TPU v7x 上高效服务 inclusionAI 的 Ling-2.6-1T。在建立起可用的基线之后，性能剖析指出 Mixture-of-Experts（MoE）路径是主要瓶颈：每一层都会把 token 分散到 32 个 JAX 设备上（每块 v7x 芯片两个），运行专家 FFN，再把输出收集回来。本文首先聚焦 Fused MoE V2，这是一个新的 Pallas kernel，它将 scatter、专家 FFN 和 gather 融合在一起，同时让 TPU 的计算与数据搬运相互重叠。

借助 Fused MoE V2，MoE prefill 延迟从 5.16 ms 降至 2.42 ms，而在同一 SGLang decode 基准测试中，16 块 TPU v7x 芯片的输出吞吐量达到 16 块 H200 GPU 的 1.29×–1.77×。完整数据见下文。

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

TL;DR

• Fused MoE V2：MoE prefill 延迟相比 Fused MoE V1 下降 53%（5.16 → 2.42 ms）；decode kernel 延迟下降约 15%（0.249 → 0.211 ms）。
• 端到端收益：仅替换 MoE kernel，就将 prefill 吞吐量提升了 24.8%，decode 吞吐量提升了 18.5%–35.3%。
• TPU 与 H200 解码对比：TPU v7x-16 提供1.29×H200×16 的解码输出吞吐量，达到mc=128，以及1.77×达到mc=512.
• 超越 MoE：Ling-2.6-1T 的完整搭建还包括混合 KV/循环记忆池、GLA 线性注意力以及单控制器数据并行。

Ling-2.6-1T 一览：一个 1T 稀疏 MoE 模型，每个 token 激活 63B 参数，256 个路由专家，采用 top-8 路由外加一个共享专家，逐通道 fp8 MoE 权重，以及混合 MLA + Lightning Linear 主干。MoE 结构驱动了本文前半部分的 kernel 工作；混合主干则引出了后续的内存池与 GLA 启动章节。

搭建环境：优化融合 MoE Kernel

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

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

1. MoE kernel 成本模型

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

scatter tokens -> local expert FFN -> gather results

在这种结构下，MoE 的运行成本不只是 GEMM FLOPs。该 kernel 必须让数据经过三条昂贵的路径：跨芯片的 token 路由、专家权重从 HBM 读入 VMEM，以及矩阵乘法单元（MXU）周围的 fp8 布局 / scale 处理。

共享专家是一条本地稠密路径。它增加了本地 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 公开规格列出每芯片约 4.614 PFLOP/s 的 fp8 算力。在此部署中，每颗芯片被暴露为两个设备，因此每设备 fp8 峰值大致为 2.307 PFLOP/s。理想计算下界为：

824.6 GFLOP / 2307 TFLOP/s = 0.36 ms

这是一个理想下界，不包含数据搬运、fp8 打包/解包，以及向量处理单元（VPU）上的 scale 处理。实测的 2.42 ms 生产环境 trace 仍比该下界高出约 7×，因此纯 GEMM FLOPs 无法解释这一延迟。

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 的 torus 使每颗芯片拥有 4 条有效链路，因此有效的单向芯片带宽大约为 4 × 100 GB/s = 400 GB/s。由于两个设备共享一颗芯片，每设备单向注入带宽大致约为 200 GB/s。

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

但 all-to-all 并不是单链路带宽测试。在 2×2×4 的 torus 上，随机目的地的平均距离约为两跳：x 方向约 0.5 跳，y 方向约 0.5 跳，z 方向约 1.0 跳。考虑这一跳数因子后，经拓扑调整的下界更接近：

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

HBM 权重搬运下界

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

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

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

共享专家额外增加了一组本地 FFN 权重，大小大致相当于一个本地专家，但它不会引入 all-to-all 通信流量。下面的估算聚焦于路由专家路径。

TPU v7x 的 HBM 带宽约为每芯片 7.38 TB/s，即每设备约 3.69 TB/s。将所有 8 个本地专家各读取一次的下界为：

402 MB / 3.69 TB/s = 0.11 ms

在实践中，内核会针对每个 token 暂存 tile 重新读取一次权重。tile 大小由 bts 决定，即块 token 暂存大小：为单个专家 FFN tile 引入 VMEM 的路由 token 行数。Ling 预填充使用 bts=160。由于每个专家大约处理 512 行，预填充需要 ceil(512 / 160) = 4 个 token 暂存 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每 tile，或约0.44 ms使用bts=160;
• fp8 打包、scale 广播和布局重排仍然消耗 VPU 和 VMEM 带宽。

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

2. 为什么这需要一个 Pallas 融合 kernel

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

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

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

纯 JAX 原生的 MoE 可以正确表达路由、专家 FFN 和输出聚合。但它无法暴露的是单个 MoE 层内部的细粒度调度。一旦 scatter、专家 FFN、HBM 权重搬运、fp8 布局处理以及 gather 跨越多个 JAX 算子或集合通信边界，XLA 就无法可靠地将 ICI-DMA、HBM-DMA、MXU 和 VPU 的工作放置到同一条手工调度的流水线上。

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

图 3. 朴素融合流水线，通信与计算阶段串行执行。语义正确，但各引擎未按细粒度重叠进行调度。

理想的稳态是：当 MXU 计算专家 i 时，HBM-DMA 预取专家 i+1 的权重，ICI-out 发送下一批路由 token，ICI-in 接收上一批输出，而 VPU 处理来自上一次 matmul 的 scale 和 layout 工作。

要表达这种调度，scatter、专家 FFN 和 gather 需要位于同一个 Pallas kernel 内。融合主要不是为了减少算子数量；它创造了一个调度空间，使相互依赖的阶段可以手动安排在 MXU、VPU、HBM-DMA 和 ICI-DMA 上执行。

3. V1：已融合，但采用碎片化的 hidden 分块

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

问题出在专家内部。一个 MoE 专家需要的不仅仅是输入 token tile 和一次 GEMM 输出。为了重叠通信和计算，kernel 还需要权重暂存缓冲区、中间激活、输出累加器以及 DMA 双缓冲。在 Ling 的 hidden size 为 8192 的情况下，保持完整 hidden 维度常驻会很快耗尽 VMEM，尤其是对于 f32 累加器和 W1/W3/W2 暂存而言。

因此，V1 采取了保守路径：对隐藏维度进行切分，并将更小的工作集流式传输通过 VMEM。

对于 Ling 16,384 prefill，V1 配置为：

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

这个块配置回答了一个放置问题：哪些 token 行、中间通道和隐藏通道留在 VMEM 中，哪些从 HBM 流式传入。

这些参数可以理解为沿 GEMM 各轴的 tile 大小：

因此，bf/bd1/bd2 主要控制特征 / 隐藏维度，而 bts/btc 控制每个专家的 token 行数。二者共同决定了一个分块能否装进 64 MB 的 VMEM 预算，以及围绕 MXU 能重叠多少 HBM-DMA / VPU 工作。

V1 付出三项结构性代价：

V1 有一些微重叠，但隐藏维度切片使重叠窗口很小。预取一次只能覆盖一个小切片，且 FFN2 的部分输出仍需经由 HBM 往返。V1 的 prefill 延迟为 5.16 ms。

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

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

这会在长生命周期张量上消耗更多 VMEM，但它消除了大部分隐藏切片暂存，并几乎消除了 FFN2 的 HBM 读-改-写路径。

Ling 16,384 prefill V2 生产配置为：

bf=512 / bts=160 / btc=80

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

这也解释了为什么在 V1 中单纯增大 bd1 / bd2 并不够。在 V1 中，更大的隐藏分块也会同时增大权重缓冲区、token 暂存缓冲区和部分输出暂存，很快就会触及 64 MB 的 VMEM 上限。更重要的是，V1 仍然要循环遍历隐藏切片；它并没有让 token 和输出累加器常驻。

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

Decode 遵循同样的逻辑，但余量更小。在 512-token 的 decode 批次下，内核延迟从 0.249 ms 降至 0.211 ms，约 15%。每个专家的有效 M 维度都很小，因此 MXU 的 tile 无法很好地摊销固定开销；这条路径也更接近专家权重 HBM 读取的下界，decode 的 trace 已经达到约 80% 的 HBM 带宽利用率。所以 V2 对 decode 仍有帮助，但它无法像 prefill 那样充分实现 VMEM 驻留和 routed-window 带来的收益。

图 4. V1 与 V2 融合 MoE 的概念时间线。V1 只创建了很小的重叠窗口，因为 hidden 维度的切片频繁轮转；V2 将 token 和累加器常驻在 VMEM 中，对专家权重进行双缓冲，并将大部分 scatter/gather 流量隐藏在 routed 计算窗口之后。

5. 针对性的 V2 优化

逐通道 direct_scaled_dot

fp8 权重量化的 scale 粒度决定了 MXU 看到的是一个大的 GEMM 还是一连串小的 GEMM。

在逐块量化下，scale 取决于 K 块：

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

scale 无法从归约中提取出来，因此 K 必须被拆分成块。每个块执行一次小的 fp8 点积，乘以该块的 scale，然后累加。一个大的 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 子字打包、缩放广播和 lane 重排。逐块量化则会在此基础上增加 K 分段和块间缩放处理。

激活量化

V2 在 scatter 之前将激活从 bf16 量化为 fp8，直接将路由 token 的载荷减半。在 Ling 16,384 prefill 上，核内 scatter 阶段从 1.39 ms 降至 0.65 ms。

这与上文 ICI 下界的计算相符：当载荷从 bf16 的 67 MB 降至 fp8 的 33.5 MB 时，通信下界几乎减半。

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

核内共享专家

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

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

6. 收益从何而来

下面的分解展示了在启用激活量化和 kernel 内共享专家的情况下，prefill 16,384 的关键路径。阴影区域是隐藏在其他阶段之下的实际工作。

图 5. Fused MoE V2 的实测重叠结构。大部分 scatter/gather 流量隐藏在路由专家窗口之下；只有 scatter 前导和 gather 尾部仍然可见。

元数据块是路由的簿记工作：token 到专家/设备的映射、每个专家的偏移/计数，以及 scatter/gather 索引。它只搬运少量元数据，耗时几十微秒，不是 prefill 的核心开销。

对同一个 V2 kernel 做消融，可以看清关键路径上还有哪些暴露的部分：

这与成本模型分析相符。即使把共享专家也算进去，理想的计算下界也只有约 0.36 ms，而且去掉 matmul 几乎不改变总延迟。scatter/gather 的工作量接近 2.4 ms，但其中约 1.8 ms 被隐藏在路由计算窗口之下。

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

• token 和累加器常驻 VMEM，减少了 token 暂存和 HBM 的读-改-写；
• 专家权重采用双缓冲，使 HBM 读取隐藏在 MXU 工作之下；
• scatter/gather 使用分 bank 缓冲区和出站/入站 ICI 通道，与路由计算重叠。

7. V2 之后还剩下什么

重叠之后，图 5 中剩余的最长片段是路由计算窗口，约占 2.42 ms 总时长的 68%。这并不意味着问题又回到了纯粹的 FLOPs：Mosaic LLO dump 显示，剩余的瓶颈主要是 fp8 打包 / lane 重排 / scale 广播，以及 VMEM 对 tile 大小的限制。

通信受限于拓扑

在我们的测量中，扁平 all-to-all 优于分层 all-to-all。在扁平配置中，发送/接收分区直接根据最终专家归属构建，一次 32 路 all-to-all 将路由后的 token 载荷从源设备发送到最终目标设备。

我们还测量了一种分层配置：将 32 设备的交换沿 2×2×4 ICI torus 拆分，先在本地维度内重新洗牌，再沿下一个维度中继，直到每个 token 到达目标专家的设备。每一轮通信的范围更小，但同样的路由 token 负载要经过多个中继阶段，增加了暂存缓冲区、同步边界，并且总传输字节数几乎翻倍。两种模式都是在融合内核之外作为独立的 all-to-all 基准测量的，因此这些数字无法与内核内追踪结果直接比较。

因此，通信方面实际可用的杠杆并不是更复杂的路由算法，而是更少的字节和更好的重叠。激活量化正是如此。

路由计算受限于 VPU / VMEM

Routed FFN1（W1+W3）的耗时约为 0.72 ms，而理想的稠密 fp8 GEMM 下界约为 0.12 ms。这一差距并非由激活量化造成：FFN1 在开启激活量化时约为 0.74 ms，关闭时约为 0.71 ms。

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

Mosaic LLO dump 解释了原因。整个 kernel 只有 4096 条真正的 vmatmul 指令，而 fp8 布局和向量侧准备工作却占据了指令流的主导地位：

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

总结

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

• ICI all-to-all 受限于 torus 拓扑和争用。
• HBM 权重读取必须通过双缓冲来隐藏。
• fp8 打包和 scale 处理让 MXU 一直等待数据成形。
• VMEM 容量限制了 tile 大小以及可共存的 overlap 缓冲区数量。

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

• 内核侧：减少 fp8 打包/解包和 scale 处理，但这越来越依赖于将模型量化与 TPU 原生执行格式对齐：对 TPU 友好的 scale 粒度、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 prefill 和 decode，并映射 DP/TP 以使分组 RMSNorm 保持在芯片本地。

混合内存池

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

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

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

JAX 还多了一层约束：运行时无法像 CUDA 路径那样原地更新这些缓冲区。SGLang-JAX 将 KV 池和循环池包装进一个 MemoryPools pytree，并作为 donated JIT 参数传入模型。每次前向传播都会返回更新后的池缓冲区，运行时再通过 replace_all() 将它们写回。这样就把缓冲区捐赠、TP/DP 分片以及未来的池扩展都保持在容器层面，而不是把各种特殊情况散落到前向循环中。

GLA（门控线性注意力）

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

St=γt St−1+kt⊤vt,ot=qt St

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

预填充：让循环足够并行以适应 TPU。 按字面理解，上述循环是串行的：token t 依赖于 token t−1 经过衰减和更新后的状态。以这种方式运行预填充，会把 16K 或 256K 的提示词变成漫长的逐 token 扫描，而这恰恰是 TPU 最不擅长的形态。

SGLang-JAX 采用数学上等价的按块（chunk-wise）形式。序列被切分为固定大小的块，每块 64 个 token。跨块之间，一个块的最终状态成为下一个块的初始状态，因此长程依赖仍然随时间向前传递。然而在块内部，递归被重新组织为针对该 token 块的稠密矩阵运算。只有块边界仍然是串行的；每个块内部的工作以块并行的 TPU 数学运算方式运行。

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

服务集成：让 GLA 留在同一条运行时路径中。 GLA 作为层级后端选择被集成，而不是作为一种独立的调度器模式。全注意力层读写 KV cache；GLA 层读写递归状态；两者都通过相同的预填充和解码批次推进。调度器仍然只看到一种生命周期：接纳、预填充、解码、释放。

这一集成在功能上已经完整，但预填充 kernel 尚未调到与 Fused MoE V2 相同的程度。GLA 的数学部分无需改动；需要改的是执行调度。

单控制器数据并行支持

Ling-2.6-1T 的分组后注意力 RMSNorm 对张量并行施加了硬性约束。每个 norm 组包含 8 个 head。如果一个组跨越了多块芯片，方差计算就会在每一层变成跨芯片 reduce，直接落在 decode 关键路径上。因此纯 TP 没有好的配置：tp ≤ 8 能让 norm 组保持在芯片本地，但对这个万亿参数模型来说并行度不足；而 tp > 8 会拆分 norm 组，并付出 all-reduce 的代价。

单控制器 DP 通过将数据并行视为另一条 mesh 轴来化解这一矛盾。mesh 被划分为若干 DP 组；每个组使用足够小的 TP，使分组 RMSNorm 保持在芯片本地，同时请求被划分到各个 DP rank 上。权重在每个 DP 组内仍按 TP 分片。逐层的 norm reduce 消失了，释放出的 ICI/HBM 预算可以转而用于更高的并发度。

重要的设计选择在于，DP 是 SPMD 运行时的一部分，而不是一组相互独立的服务器副本。SGLang-JAX 运行一个逻辑调度器，dp_rank 被附加到请求、KV 分配和前缀缓存键上。这样一来，就能基于单一负载快照进行全局准入控制，跨主机确定性地构建 batch，并拥有一个全局前缀缓存结构，其条目以 (dp_rank, prefix) 为键。

这也能与混合运行时的其余部分干净地组合。将 mesh 扩展到更大的配置，例如在每个数据并行组内加入张量并行，只是 mesh 形状的变更，而不是调度器的分叉，因此内存池、批处理路径和注意力后端都保持同一套心智模型。

实验与基准测试

所有 TPU 结果均使用 SGLang-JAX 在单个 TPU v7x slice 上服务 Ling-2.6-1T；该设置在 V1/V2 消融实验中完全一致；仅 MoE kernel 配置不同。

基准测试配置

• 硬件： TPU v7x，16 芯片（2×2×4 ICI torus）→ 32 设备
• 并行策略： tp = ep = 32，dp = 8
• 模型： Ling-2.6-1T，bf16 激活值，per-channel fp8 MoE 权重
• 数据集： SGLang 默认的 random 基准数据集（从 ShareGPT 采样）
• 运行时： SGLang-JAX（JAX 0.8.1），dvfs p_state=7
• 输入长度： 16384
• Prefill： 输出 1，并发 128
• Decode： 输出 1024，并发 128 / 512

在 16384 token 输入、mc=128 下的 Prefill 输入吞吐量。设置完全相同，仅 MoE kernel 配置不同：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）对比 GPU H200×16（2 个节点，tp8·pp2），相同模型和 SGLang bench 工作负载，双方各 16 个加速器。关于 prefill 差距的说明见下方注释。

关于端到端 prefill 与 MoE kernel 加速比的说明：Fused MoE V2 kernel 将 MoE 层的 prefill 延迟降低了约 53%（device trace），但端到端 prefill 吞吐量仅提升约 25%（v1 → v2）。MoE 层已不再是 prefill 的主要开销：GLA（gated linear attention）prefill kernel 目前是主要瓶颈，且尚未优化到同等程度，因此稀释了端到端 prefill 的加速效果。同样的瓶颈也是图 6 中 TPU v7x-16 在 prefill 一列落后于 H200×16、却在两个 decode 数据点上领先的原因。将 GLA prefill kernel 提升到同等水平是正在进行的工作，我们预计这将释放更大的端到端 prefill 增益。

局限性与未来工作

我们对 Ling-2.6-1T 的支持在本版本中有意限定了范围；若干事项仍作为后续工作，我们正在积极推进：

• GLA / 线性注意力 prefill kernel。 正如基准测试部分所指出的，GLA（Lightning Linear）prefill kernel 目前已成为 prefill 阶段的主要开销。通过考虑更好的分块/平铺、融合门控与循环状态更新，以及采用与 MoE kernel 相同的 MXU/VPU/DMA 重叠处理等方法将其提升到同等水平，是端到端 prefill 剩余最直接的优化手段。
• 动态专家并行负载均衡（EPLB）。 当前 FusedEPMoE 路径使用静态的专家到设备放置策略，但真实工作负载在 256 个路由专家上的命中率并不均匀。一个动态 EPLB 过程，能够根据观测到的流量周期性地重新平衡专家到 rank 的映射，将缩小每设备峰值利用率与平均利用率之间的差距，尤其是在更大的 batch size 下。
• 混合内存池上的 Radix 缓存。 SGLang 的 RadixAttention [9] 前缀缓存假设存在单一的 per-token KV 池，而 Ling-2.6-1T 将 per-token KV 与 per-request 循环状态混合在一起，因此简单的 prefix-share 会在线性层上静默地跨请求混合状态。我们正在设计一种扩展方案，按 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 基准工作负载：prefill（out 1，mc 128）· decode（out 1024，mc 128）· decode（out 1024，mc 512）。

TPU：SGLang-JAX（Fused 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 base = fused_v2 --no-moe-fused-act-quant --no-moe-fused-shared-experts。

v2 +act-quant 用例增加了 --moe-fused-act-quant；v2 +act +SE-overlap 将两者都开启。两个 external-shared-expert 配置使用 --mem-fraction-static 0.85，因为它们在 0.88 时 OOM。

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

性能运行的完整基准命令见 SGLang-JAX cookbook。

服务器启动与精度复现

AIME 2026 检查使用 MathArena/aime_2026，30 道题，pass@1：26 / 30 = 86.7%。该运行零请求错误，所有响应均正常终止（finish_reason=stop，在 32768 tokens 处无截断）。这表明 fp8 fused-MoE 服务路径没有明显的精度回退。

完整的启动服务器命令、请求和工具调用示例，以及 AIME 2026 精度复现，均见同一个 SGLang-JAX cookbook。

参考文献

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

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

[3] Ragged Paged Attention

[4] Fused MoE V1 kernel，tpu-inference

[5] Fused MoE V1 kernel 适配于 SGLang-JAX

[6] DeepSeek-V2（MLA）

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

[8] MiniMax-01（Lightning Attention）

[9] SGLang（RadixAttention）

致谢

蚂蚁集团 ASystem 核心团队： Zhenxuan Pan、Guowei Wang、Yuhong Guo、Shuo Wan

SGLang-JAX 团队： jimoosciuc、Prayer、aolemila、neo、leos、pathfinder-pf、Haolin Fu、Qinghan Chen、JamesBrianD、Haoguang Cai、Yuhao Hu、cjx0709、Zhengke Zhou、Yuxin Wei、Lianfang Wang、0xaskr