在 TPU 上使用 SGLang-JAX 优化 Ling-2.6-1T：通过一个 Pallas 内核将 MoE 数据移动隐藏在计算背后
Prayer, JamesBrianD, 傅浩林, 蔡浩光, 陈庆涵
SGLang-JAX 现在支持在 TPU v7x 上高效服务 inclusionAI 的 Ling-2.6-1T 模型。在建立了一个可工作的基线后，性能分析指出混合专家（MoE）路径是主要瓶颈：每一层都将 token 分散到 32 个 JAX 设备（每个 v7x 芯片两个）上，运行专家前馈网络，然后将输出结果收集回来。本文首先聚焦于 Fused MoE V2，这是一个新的 Pallas 内核，它融合了分散、专家 FFN 和收集操作，同时实现了 TPU 计算与数据移动的重叠。
借助 Fused MoE V2，MoE 预填充延迟从 5.16 毫秒降至 2.42 毫秒，并且在相同的 SGLang 解码基准测试中，16 个 TPU v7x 芯片达到了 16 个 H200 GPU 输出吞吐量的 1.29 倍至 1.77 倍。完整数据如下所示。
图 1. Ling-2.6-1T 在 TPU v7x-16 与 H200×16 上的解码吞吐量对比，使用 SGLang 默认的 `random` 基准测试数据集（从 ShareGPT 采样），输入为 16,384 个 token，输出为 1,024 个 token。
核心摘要
- Fused MoE V2：与 Fused MoE V1 相比，MoE 预填充延迟降低 53%（5.16 → 2.42 毫秒）；解码内核延迟降低约 15%（0.249 → 0.211 毫秒）。
- 端到端收益：仅替换 MoE 内核即可使预填充吞吐量提升 24.8%，解码吞吐量提升 18.5%–35.3%。
- TPU 与 H200 解码对比：在 mc=128 时，TPU v7x-16 的解码输出吞吐量是 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 部署部分提供了动机。
方案设置：优化 Fused MoE 内核
本节中所有 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 内核成本模型
Ling-2.6-1T 每层有 256 个路由专家和 1 个共享专家，采用 top-8 路由。在 ep=32 的情况下，每个设备拥有 8 个本地路由专家。一个 token 选中的 8 个专家通常分布在不同的设备上，因此每一层的路由路径都具有相同的结构：
scatter tokens -> local expert FFN -> gather results
在这种结构下，MoE 的运行成本不仅仅是 GEMM 浮点运算。内核需要通过三条高成本路径移动数据：跨芯片的 token 路由、从 HBM 读取专家权重到 VMEM，以及在矩阵乘法单元（MXU）周围处理 fp8 布局和缩放因子。
共享专家是一条本地密集路径。它增加了本地 FFN 计算量，但不参与路由的全对全通信，对 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 浮点运算无法解释该延迟。
ICI token 路由下界
每个设备的散射负载为：
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 MB | 0.17 毫秒 | 0.34 毫秒 |
| bf16 67.1 MB | 0.34 毫秒 | 0.67 毫秒 |
但全对全通信并非单链路带宽测试。在 2×2×4 环面拓扑上，平均随机目标距离约为两跳：x 方向约 0.5 跳，y 方向约 0.5 跳，z 方向约 1.0 跳。考虑此跳数因子后，经拓扑调整的下限更接近：
| 数据载荷 | 单次散射，经平均跳数调整 | 散射 + 收集 |
|---|---|---|
| fp8 | 0.34 毫秒 | 0.67 毫秒 |
| bf16 | 0.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 暂存分块重新读取一次权重。分块大小由 bts（块 token 暂存大小）决定，即针对一个专家 FFN 分块加载到 VMEM 中的路由 token 行数。Ling 预填充阶段使用 bts=160。由于每个专家大约处理 512 行，预填充需要 ceil(512 / 160) = 4 个 token 暂存分块。V2 流水线在这些分块之间进行权重预取，因此 HBM 读取下限大致为：
4 * 402 MB / 3.69 TB/s = 0.44 ms
权重读取不必出现在关键路径上。V2 通过双缓冲将其隐藏在 MXU 窗口之后。这些数字解释了为何需要这种调度：如果 HBM 读取在 GEMM 之前串行化，它们已经超过了纯计算下界。
要点总结
TPU 上的 MoE 主要是一个数据搬运和重叠问题：
- 路由 + 共享 FFN 计算下界：约 0.36 毫秒；
- fp8 分散 + 汇聚拓扑下界：约 0.67 毫秒；
- 专家权重 HBM 读取下界：每 tile 约 0.11 毫秒，或 bts=160 时约 0.44 毫秒；
- 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 tile、专家中间结果、输出累加器、预取的权重 tile |
| HBM | 每个芯片上的大型片外存储器 | 专家权重、token 暂存缓冲区、大型中间缓冲区 |
| HBM-DMA | HBM ↔ VMEM 数据搬运 | 将当前/下一个专家权重预取到 VMEM；按需移动暂存缓冲区 |
| ICI / ICI-DMA | TPU 切片内部的直接芯片间网络 | 在源芯片和目标芯片之间搬运路由后的 token 负载；将 token 分散到专家所有者，并将输出汇聚回 token 顺序 |
纯 JAX 原生 MoE 能够正确表达路由、专家 FFN 和输出聚合。但它无法暴露单个 MoE 层内部的细粒度调度。一旦 scatter、专家 FFN、HBM 权重移动、fp8 布局处理以及 gather 跨越多个 JAX 算子或集合通信边界，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：已融合，但存在碎片化的隐藏维度分块
我们的起点是 Fused MoE V1，最初由 Jevin Jiang、Kyuyeun Kim 等人在 tpu-inference 项目 [4] 中提出并优化，后经修改适配到 SGLang-JAX 中，成为 FusedEPMoE [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 的固定开销无法被摊还 |
| btc | 在一次计算循环中送入的、bts 内部的 token 行数 | 内部 M 计算分片；不能超过 bts，通常能整除 bts |
| bf | W1/W2/W3 的中间通道数 | 控制 FFN 中间分片；通常越大 MXU 窗口越长，但 VMEM 开销也越大 |
| bd1 | FFN1 隐藏层归约 K 切片 | V1 对隐藏层 K 进行切片；bd1 越小，FFN1 点积次数越多、每次点积越小 |
| bd2 | FFN2 隐藏层输出 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，因为它不再对隐藏维度进行切片。其结构变化是：
| 每个专家 | V1 | V2 | 效果 |
|---|---|---|---|
| FFN1 点积 | 每个硬件点积的有效 K 约为 256 | fp8 分块 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 读取的下限，解码追踪已显示 HBM 带宽利用率达到约 80%。因此 V2 仍有助于解码，但其在 VMEM 驻留和路由窗口方面的增益不如预填充阶段那样充分。
图 4. V1 和 V2 融合 MoE 的概念时间线。V1 仅产生较小的重叠窗口，因为隐藏维度切片频繁切换；V2 将 token 和累加器驻留在 VMEM 中，对专家权重进行双缓冲，并将大部分分散/汇聚流量隐藏在路由计算窗口之后。
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 缓冲区，并将其调度在 scatter 窗口内。
共享专家自身的计算量约为 0.159 毫秒，但仅给关键路径增加 0.068 毫秒，约占 2.7%。原因很简单：共享专家不需要跨芯片 token 分发；所有需要的 token 都是本地的，因此它可以与路由 FFN 之前的 scatter 阶段重叠。
6. 性能提升的来源
下面的分解图展示了在启用激活量化和内核内共享专家的情况下，prefill 16,384 的关键路径。阴影区域是隐藏在其他阶段下的实际工作。
图 5. Fused MoE V2 的实测重叠结构。大部分 scatter/gather 流量隐藏在路由专家窗口下；只有 scatter 前导和 gather 尾部是可见的。
元数据块是路由管理开销：token 到专家/设备的映射、每个专家的偏移量/计数以及 scatter/gather 索引。它只移动少量元数据，耗时几十微秒，并非核心 prefill 开销。
对同一 V2 内核进行消融实验，显示关键路径上暴露的内容如下：
| 消融 / 组件 | 结果 | 解读 |
|---|---|---|
| 完整 V2 生产环境 | 2.42 毫秒 | 本节使用的标准 MoE prefill 延迟 |
| 禁用所有专家矩阵乘法 | 相比完整版本 -2.2% | 纯 MXU 计算未暴露 |
| 可见 scatter | 0.42 毫秒 | 通信前导仍留在关键路径上 |
| 可见 gather | 0.18 毫秒 | gather 尾部仍留在关键路径上 |
| 无计算可隐藏的 scatter + gather | 约 2.4 毫秒 | 实际通信延迟接近重叠前的完整内核延迟 |
这与成本模型分析结果一致。即使计入共享专家，理想计算下限也仅为约 0.36 毫秒，而移除矩阵乘法操作对总延迟几乎没有影响。分散/聚合操作耗时接近 2.4 毫秒，但其中约 1.8 毫秒被隐藏在路由计算窗口之下。
因此，V2 的增益来自三种机制：
- token 和累加器常驻 VMEM，减少了 token 的分段搬运和 HBM 的读-修改-写操作；
- 专家权重采用双缓冲机制，使得 HBM 读取操作能够隐藏在 MXU 计算工作背后；
- 分散/聚合操作利用分库缓冲器以及出站/入站 ICI 通道，与路由计算实现重叠。
7. V2 之后还剩下什么
实现重叠之后，图 5 中剩余的最长段是路由计算窗口，约占 2.42 毫秒总时长的 68%。这并不意味着问题又回到了纯 FLOPs 瓶颈：Mosaic LLO 转储显示，剩余的瓶颈主要是 fp8 打包/通道重排/缩放广播，以及 VMEM 对 tile 尺寸的限制。
通信受拓扑结构限制
在我们的测量中，扁平全对全通信优于分层全对全通信。在扁平配置下，发送/接收分区直接根据最终专家所属设备构建，一次 32 路全对全通信将路由后的 token 负载从源设备发送到最终目标设备。
我们还测量了一种分层配置：将 32 设备的交换过程沿着 2×2×4 的 ICI 环面拆分，首先在本地维度内进行重排，然后沿着下一个维度进行中继，直到每个 token 到达目标专家所在的设备。每一轮通信的范围更小，但相同的路由 token 负载会经过多个中继阶段，这增加了分段缓冲、同步边界，并且总传输字节量几乎翻倍。两种模式都是在融合内核之外作为独立的全对全基准测试进行测量的，因此这些数字不能直接与内核内部的追踪数据进行比较。
| 模式（负载 = 16384 × 8192 × 数据类型大小） | bf16 | fp8 |
|---|---|---|
| 扁平全对全 | 2.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 / 80 | 2.42 毫秒 | 47 MB |
| 160 / 160 | 2.44 毫秒 | 47 MB |
| 128 / 128 | 3.12 毫秒 | 44 MB |
| 256 / 128 | 3.19 毫秒 | 54 MB |
| 256 / 256 | 3.23 毫秒 | 54 MB |
| 384 / 128 | 内存不足 | 62 MB |
Mosaic LLO 转储解释了原因。整个内核只有 4096 条真正的 vmatmul 指令，而 fp8 布局和向量端准备工作主导了指令流：
| LLO 指令 | 数量 | 作用 |
|---|---|---|
| vselect | 50880 | 子通道选择/混合 |
| vbitcast | 46566 | fp8 子字重新解释 |
| vcombine | 36380 | 子通道合并 |
| vpack_format | 34368 | MXU 输入打包 |
| slane | 29960 | 子通道移动 |
| vunpack | 25600 | fp8 解包 |
| matmul_data_format | 25600 | MXU 前的格式转换 |
| vrot | 21524 | 通道旋转 |
| vmatres / vmatprep | 17408 / 10240 | MXU 排出/馈入 |
| vslreplicate | 6032 | 缩放因子广播 |
| vmatmul | 4096 | 实际矩阵乘法 |
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。
- 工作负载侧：跨批次进行重叠，使路由窗口能够与其他层工作并行运行。
- 硬件方面：提供能更好支持全互联的互连拓扑结构，或提供更大的虚拟机内存（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 缓存在 10 个全注意力层上，每个设备每个 token 的成本约为 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（门控线性注意力）
每个 GLA 层 [7] 将历史信息保存在一个固定大小的循环状态中，而不是为每个过去的 token 存储一个 KV 条目。其更新可以写成：
这将注意力历史从逐 token 增长的方式，转变为每个活跃请求一个状态张量。在长上下文场景下，这是主要优势：携带历史信息在计算上保持线性，在状态上保持固定大小，而不是实例化并读取一个不断增长的 KV 历史。
预填充：使循环过程对 TPU 足够并行。从字面上看，上述循环是串行的：token t 依赖于来自 token t−1 的衰减和更新后的状态。以这种方式运行预填充会将一个 16K 或 256K 的提示词变成一个漫长的逐 token 扫描，这恰恰是 TPU 不擅长的形状。
SGLang-JAX 采用数学上等价的分块形式。序列被分割成固定大小为 64 个 token 的块。在块之间，一个块的最终状态成为下一个块的初始状态，因此长程依赖关系仍然随时间向前传递。然而，在块内部，循环被重新排列为对 token 块进行的密集矩阵运算。只有块边界保持串行；每个块内部的工作作为块并行的 TPU 数学运算执行。
解码：循环的自然形式。解码更简单：预填充已经将提示词折叠进了循环状态，因此每个新 token 读取请求的当前状态，应用一次循环更新，输出注意力结果，并写回新状态。问题从长序列并行性转变为高效的微小状态更新。
服务集成：将 GLA 保持在相同的运行时路径内。GLA 作为层级的后端选择进行集成，而非独立的调度器模式。全注意力层读写 KV 缓存；GLA 层读写循环状态；两者都通过相同的预填充和解码批次推进。调度器仍然看到一个生命周期：准入、预填充、解码、释放。
该集成在功能上已完成，但预填充内核尚未调整到与 Fused MoE V2 相同的程度。GLA 的数学运算无需改变；需要改变的是执行调度。
单控制器数据并行支持
Ling-2.6-1T 的分组后注意力 RMSNorm 对张量并行施加了硬性约束。每个归一化组包含 8 个注意力头。如果一个组跨越了芯片，方差计算就会在每一层上变成跨芯片的规约操作，且直接位于解码关键路径上。因此，纯 TP 没有好的设置：tp ≤ 8 可使归一化组保持在芯片本地，但会导致这个万亿参数模型并行不足；而 tp > 8 则会拆分归一化组并付出全规约的代价。
单控制器数据并行通过将数据并行视为另一个网格轴来解决这一矛盾。网格被划分为多个数据并行组；每个组使用足够小的张量并行，使分组后的 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）对比 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 过程，能够根据观察到的流量定期重新平衡专家到计算节点的映射，将缩小峰值与平均设备利用率之间的差距，尤其是在较大批次规模下。
- 混合内存池上的基数缓存。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 公开规格显示，每颗芯片的 fp8 算力约为 4.614 PFLOP/s，HBM 带宽为 7.38 TB/s，双向 ICI 带宽为 1.2 TB/s。在此部署中，每颗芯片被暴露为两个设备，因此成本模型部分中的每设备下限使用了大约一半的芯片级算力和带宽。有关 TPU 内存层次结构和执行单元（MXU、VPU、VMEM、HBM、ICI）的背景信息，请参阅 Google Cloud 的 TPU 系统架构。
性能复现
双方运行相同的模型和相同的 SGLang 基准测试工作负载：prefill（输出 1，mc 128）· decode（输出 1024，mc 128）· decode（输出 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 标志：融合 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 时会内存不足。
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 个 token 处无截断）。这表明 fp8 融合 MoE 推理路径未出现明显的准确率回退。
完整的启动服务器命令、请求与工具调用示例，以及 AIME 2026 准确率复现方法，均收录于同一份 SGLang-JAX 实践指南中。
参考文献
致谢
蚂蚁集团-系统部核心团队：潘振轩、王国伟、郭宇红、万硕
SGLang-JAX 团队：jimoosciuc、Prayer、aolemila、neo、leos、pathfinder-pf、傅浩林、陈庆涵、JamesBrianD、蔡浩光、胡宇浩、cjx0709、周正科、魏宇欣、王连芳、0xaskr