PyTorch 团队用 SGLang-TorchTPU 在 TPU v7 上优化 Muse Spark Open 推理

TL;DR

我们在单台四芯片的 TPU v7 主机上，以 TorchTPU PyTorch 后端运行 SGLang，为 Meta 的 Muse Spark Open（一个 816B 参数的混合专家模型）提供服务。核心是一个 megakernel：解码或验证步骤的全部 62 层都作为一个 Pallas 程序运行，每一次权重加载、KV 读取和跨设备传输都被流水线化到计算之后。其余工作分为三组：更快的解码步骤（collective folding、无 max 的注意力循环并对其 KV 分块进行流水线化、专家循环将寄存器溢出排除在常见路径之外，以及 DFlash2 推测解码）、每芯片更多 KV 缓存（megakernel 中按 lane 序列的 KV 布局，以及 TPU 上带会话感知驱逐的 L2 HiCache），以及更快的预填充（融合 MoE 内核和序列并行）。

• 解码速度。 megakernel 将一个或两个并发请求的每输出 token 时间减半，这是它帮助最大的场景；随着并发增加，收益收窄。其验证步骤的运行时间在仅从 HBM 流式加载其权重所需时间的 1.5–1.7 倍以内。
• AgentX。 在 SemiAnalysis 的 agentic 编码基准的 1M 上下文版本上，我们的曲线从每芯片 6,101 tokens/s 时每请求 433 输出 tokens/s，到每请求 95 tokens/s 时每芯片 25,906 tokens/s（图 1）。

结果

图 1。 InferenceX AgentX，1M 上下文轨迹集：每美元 token 数与 p90 交互性。TPU v7：一台主机，四芯片，FP8 计算配 NVFP4 专家权重。GB200（FP4）作为基线，来自 InferenceX 截至 2026-11-02 的已发布结果。成本为 InferenceX 的外部持有成本：每 TPU v7 芯片小时 $1.21，每 GB200 GPU 小时 $1.86。

AgentX 由 393 个选择加入的 Claude Code 会话构建。它重放这些会话的请求长度、轮次之间的前缀共享、工具暂停和子代理结构，并用合成 token 代替原始内容。我们运行其 1M 上下文轨迹集，与 GB200 基线一样，其中提示达到约 990K tokens；上下文窗口较短的模型运行上限为 256K 的版本。吞吐量计入输入和输出 tokens。图 1 中每个 TPU 点都在 p90 下于 47 秒内完成请求。

背景

模型

Muse Spark Open 是 Meta 的开源权重混合专家（MoE）[1] 模型，总参数 816B，每 token 约 42B 活跃。

• 62 个解码层，全部为 MoE。 每层将每个 token 路由到 256 个专家中的 8 个。
• 混合注意力。 46 层对最后 2,048 个 tokens 使用滑动窗口注意力 [2]；其余 16 层关注完整上下文，最多 1M tokens。注意力有 128 个查询头和 16 个 KV 头，每个宽 64。
• NVFP4 专家。 专家权重为 4 位，每 16 个值一个缩放因子，是一种微缩放格式 [3]。

硬件

一台 TPU v7 (Ironwood) 主机有四芯片。每个芯片封装两个 chiplet，每个 chiplet 有一个 TensorCore 和自己的 HBM，框架将每个 chiplet 暴露为一个设备，因此一台主机有 8 个设备。我们用张量并行 [4] 将模型分片为 8 份，每个设备一个分片。

图 2。 一个 TPU v7 芯片，简化自 Google Cloud 的 TPU7x 文档中的架构图。数字为 Google 公布的数值。

每个设备都是单线程处理器：一次只运行一个程序，因此速度来自让矩阵单元持续获得数据，并在该程序内部将内存传输与计算重叠。这些是本篇文章中重要的单元，如Google Cloud 的 TPU7x 文档和 JAX 的Pallas TPU 指南所述：

• 标量单元。每个设备一个，带有自己的小型标量内存（SMEM）。它运行控制流：循环、分支和地址运算。
• MXU，矩阵单元。每个设备两个。它们执行矩阵乘法；芯片的峰值在 bf16 下为 2,307 TFLOPs，在 fp8 下为 4,614。在小批量下，它们大多在等待权重从 HBM 到达。
• 向量单元（VPU）。逐元素工作，例如归一化、softmax 指数和格式转换。在近期的 TPU 上，它在 8 × 128 个 32 位值的 2D 向量寄存器上计算，即 8 个子通道 × 128 个通道，因此让通道空置的数据布局会同时浪费向量工作和内存（第 7 节）。旁边的跨通道单元（XLU）执行转置和置换。
• VMEM。紧邻 MXU 和向量单元的片上 SRAM，带宽远高于 HBM。内核在此暂存其分块，因此其大小限制了内核能保持多少数据在途。
• DMA 引擎。它们在后台于 HBM、VMEM 和其他设备之间移动数据：内核启动一次复制并在稍后等待它，同时进行计算。巨型内核的流水线化（第 2 节）就建立在此之上。
• HBM。每芯片 192 GiB，带宽 7.37 TB/s，其中每个设备 96 GiB。
• 裸片间链路。它连接一个芯片的两个设备；Google 文档称其比 1D ICI 链路快六倍。第 8 节的成对交换依赖它。
• ICI。每芯片六条链路，每条 200 GB/s（双向合计），即每芯片 1,200 GB/s。主机的四个芯片通过它相互通信；64 个或更多芯片的切片构成 3D 环面。
• 主机。每个芯片通过 PCIe Gen5 ×16 连接到主机的 CPU。4 芯片主机有 960 GB DRAM，L2 HiCache 将其用作第二级缓存层（第 7 节）。
• SparseCores。每芯片四个。本工作不使用它们。

TorchTPU 文章有更长的入门介绍。

软件

服务器是 SGLang-TorchTPU：SGLang 的 [5] 调度器、KV 缓存和 PyTorch 中的模型代码，热路径上使用 Pallas 内核。我们使用推测解码 [6, 7] 进行解码：一个小型草稿模型猜测接下来的几个 token，完整模型在一次前向传播中检查它们全部，即验证步骤。验证步骤中的每个（请求，位置）对是一行，因此两个请求各验证 8 个 token 构成一个 16 行的步骤。

1. 一个巨型内核

在小批量下，生成一个 token 对每个权重所做的数学计算很少，因此只要没有东西中断流，该步骤就能以 HBM 流式传输权重的速度完成。因此，解码或验证步骤应当是从 HBM 进入 MXU 的不间断权重流。我们之前的路径（下文称“原生”）将每一层作为多个内核加集合通信来运行，每个边界都会中断流。巨型内核将一步的全部 62 层作为每个设备一个 Pallas 程序运行，因此下一批权重始终在途。整步巨型内核已在 GPU 上构建 [8, 9]，并在 TPU v7 上为 Kimi K3 和 Qwen 3.8 构建 [10]。我们的内核运行在 SGLang 的服务路径内，基于其分页 KV 缓存，同时处理多个请求的解码和 DFlash2 验证步骤。

原生路径已经由融合的 Pallas 块构建完成，但其验证步骤耗时是流式加载其权重所需时间的 2.9 倍。对我们第一个单内核原型进行性能分析后，我们找到了原因：在 batch 1 下，矩阵乘法和注意力几乎不耗时，而集合通信每层约耗时 45 µs。程序边界之间没有任何重叠，因此该步骤受限于开销，而非带宽。

图 3。 一个步骤。在 megakernel 内部，每一层都运行相同的流程（白色）；阴影框是内核自身发起的数据移动。

对于每个步骤，每台设备运行一个 pallas_call，覆盖全部 62 层和最终的 norm。SGLang 仍在其周围运行 embedding、草稿模型、一个小型 KV 写入内核、LM head 和采样。

图 4。 每个验证步骤的耗时，两条路径均使用 DFlash2 推测解码（GSM8K，在较早的 Muse Spark checkpoint 上测量）。下限是从 HBM 流式加载该步骤的不同专家和稠密权重所需的时间。

对于每个输出 token，在单个请求下，megakernel 比原生路径快 1.94 倍（1.95 ms 对 3.78 ms），在两个并发请求下快 1.85 倍。接下来的章节将展示是什么让一个程序如此之快。

2. 在 megakernel 内部流水化一切

一个步骤主要是数据移动：权重、KV 缓存以及设备间的部分和。一旦整个步骤成为一个程序，内核就可以调度每一次传输，使其恰好在需要它的计算之前到达 [11]，包括跨层边界 [10]。图 5 勾勒了一层的情况。

图 5。 一个步骤中的一层，示意图。每条 DMA 流都领先于或并行于消费它的计算运行。

权重在需要之前到达

• 用于稠密权重的环形缓冲区。 注意力、o-proj 和 router 权重通过一个 tile 环流式加载，该环可跨层边界预取。约十个 tile 在途即可维持每设备 3.34 TB/s，而单条 DMA 流仅能达到 1.13 TB/s。
• 路由完成后立即加载专家。 router 的选择只有在层中间才能确定，因此内核在路由结束的那一刻就请求每个被选中专家的权重，同时有三个专家在途。
• 专家循环期间预取下一层。 下一层注意力权重的重新填充被推迟到专家循环运行时，因此它们永远不会与当前层正在等待的专家权重竞争。

集合通信在数据一存在时就离开

张量并行每层在 8 台设备之间多次交换部分和。内核不是在每个矩阵乘法之后单独进行集合通信，而是通过从内核内部发出的远程 DMA 将部分结果直接写入所属设备的内存；一小轮约耗时 3 µs。

• 流式 o-proj。 每个 o-proj 输出 tile 一计算完成就发往其所有者，而下一个 tile 仍在进行乘法运算。

流水化 KV 写入

每个步骤将其新的 K 和 V 行写入缓存。在我们的 fp8 sequence-on-lane 缓存（第 7 节）中，一个 token 是 256-token 页的一个 lane 列，因此写入一个 token 意味着将整个页块读入 VMEM，合并新列，然后将页写回。

这作为一个小内核在 megakernel 之后立即运行，因为 XLA 会复制任何发出远程 DMA 的内核的别名操作数。它过去会依次处理每一页：读、等待、合并、写、等待，对每个请求和 62 层中的每一层都如此。

图 6. 示意图。流水线化内核一次性发出某一层的所有页读取请求，每页到达时即进行合并，并启动其写入；下一层的读取立即在第二个 VMEM bank 上开始。

新内核一次性加载所有层的新行，然后保留两个 VMEM bank：层 l 读入其中一个，同时层 l−1 的写入从另一个排空，且一个 bank 只有在早两层的写入落地后才会被复用。在 16 行时，每步的写入开销从 810 µs 降至 155 µs；在 64 行时，从 2,286 µs 降至 336 µs，且结果逐位一致。

3. 集合折叠

由于每一层被拆分到 8 个设备上，设备往往需要先收集其他设备的结果才能继续。在 MoE 之后和注意力之后，每个设备沿隐藏维度持有该层输出的一片，并且必须对另外七片执行 all-gather，下一 matvec 才能运行。集合折叠让该 matvec 在切片到达时就开始处理，而不是等待整个 gather 完成，正如分解式集合矩阵乘法所做的那样 [12]，此处跨越位于两者之间的 RMS norm。

图 7. 示意图。折叠式 matvec 在每一片 gather 到达时立即对其进行乘法，从设备自身的那片开始，并在最后应用逐行归一化。

使这成为可能的是 Muse Spark 进行归一化的位置。下一层的注意力输入是残差流 α·x + β·s·f 的 RMS norm [13]，其中 x 是该设备已持有的残差，f 是 gather 得到的 MoE 输出，s 是 f 自身的范数因子。两个范数因子 s 和最终的 r 均为每行一个标量。每行标量与矩阵乘法可交换，因此投影拆分为设备可立即计算的部分和逐片累加的部分：

这里 g 是范数的增益向量，f_d 是来自设备 d 的切片，W_d 是权重行的对应块。内核将本地部分和到达的切片堆叠为每块一个乘积，按固定的对等顺序消费切片，并重排权重环，使 W 的块以相同顺序流入。最后一片到达后，它计算 s 和 r 并对结果进行缩放。

内核以这种方式每层折叠两次 all-gather：MoE 输出折叠进下一层的 qkv 投影，注意力输出折叠进 pooler 和 router。残差流本身不变；只有矩阵乘法输入舍入到 bf16 的位置发生移动，因此结果并非逐位一致，我们的精度检查已覆盖这一点。

4. 在注意力中实现更好的流水线

在长上下文下，注意力是一步中最大的部分。内核逐块遍历 KV 缓存，关键瓶颈是 online softmax [14] 的运行最大值：它将每个块与前一个块串联起来，因此块无法流水线化。

运行最大值串联每个块

FlashAttention [15] 所使用的 online softmax 会保留迄今为止见到的最大分数。每个新块计算其分数、更新该最大值、针对它进行指数运算，并在最大值增长时重新缩放此前累积的所有内容。因此每个块都是一条串行链（等待其 DMA、QK、max、exp、sum、P·V、rescale），在前一块的最大值已知之前无法开始其指数运算。在微基准测试中，这条链每块约花费 0.17 µs，无论块是 fp8 还是 bf16。

Muse Spark 使最大值变得不必要

Muse Spark 使用无学习权重的 RMS norm 对 q 和 k 进行归一化（QK-norm [16, 17]），并为 q 赋予固定增益。任何注意力分数的绝对值都不会超过 43.8，远低于 f32 中 exp 溢出的阈值。运行最大值的存在只是为了防止这种溢出，因此内核将其丢弃，同时丢弃的还有重缩放。内核在构建时根据模型常量检查该边界。FlashDecoding++ [18] 在 GPU 上以统一最大值和针对罕见溢出的重计算路径打破了同样的链条；而这里的边界是精确的，因此内核无需回退方案。

图 8. 示意图。使用运行最大值时，每个块的链需要等待前一个块。不使用它时，块之间相互独立，其阶段可以重叠。

独立的块流水线化

去掉最大值后，块之间不再相互依赖。内核在每次循环迭代的开头发出所有 DMA 等待，将块拆分到多个独立的 softmax 流中，这些流在最后合并，正如 Flash-Decoding 在 KV 序列上所做的那样 [19]，并对循环进行软件流水线化 [11]：在将前一个块的概率乘以 V 的同时，对下一个块进行打分。

三个较小的改动进一步缩短了每条链：

• 按层类型划分块大小。 全局层以 1,024 到 2,048 个 token 的长块读取缓存，因此 100K 上下文需要的链要少得多。滑动窗口层使用 512 个 token 的块，因为其窗口仅跨越几个块。
• fp8 键直接送入 MXU。 在向量单元上将每个 fp8 块扩展为 bf16 所消耗的向量工作量超过了其供给的 MXU 工作量。混合 bf16 × fp8 点积是精确的，因此键现在以 fp8 形式输入。
• 注意力缩放放在查询中。 Muse Spark 的注意力缩放为 1/8，是 2 的幂，因此对查询乘以一次而不是对每个分数乘以一次是精确的。

图 9. 生产形状下的微基准测试。内存下限是读取 KV 字节所需的时间。

在微基准测试中，更长的块配合独立的流将注意力时间缩短了 2.4 倍，而去掉运行最大值后达到了 3.9 倍。在服务器中，流水线化循环在 500K 上下文、两个请求的情况下将每请求解码速度提升了 18%。

5. 让寄存器压力远离常见路径

处理器在其寄存器中保存的值上进行计算，寄存器是算术单元旁边一组小而固定的存储槽。当代码需要的活跃值超过寄存器容量时，编译器会将其中一些溢出到内存，稍后再加载回来 [20]，这些额外的存储和加载在每次经过代码时都会耗费时间。在具有多个分支的循环中，编译器必须将所有分支的活跃值放入相同的寄存器中，因此一个繁重的分支可能使循环的每次执行都发生溢出，即使该分支从未运行。

巨型内核的专家循环恰好存在这个问题。一个步骤在其行被路由到的专家上循环。对于每个专家，内核等待专家的权重，并根据专家服务的行数运行所选路径：少行路径（针对一到八行的小变体），以及针对更繁忙专家的反量化后相乘路径。随着批次增长，一个步骤会访问更多专家，但大多数专家仍然只服务一两行，因此轻量路径才是关键所在。

一个繁重的分支拖慢了轻量路径

8 行变体保留着大型 f32 累加器，这些累加器即使单独存放也无法容纳在向量寄存器中。由于它与单行和双行变体共用一个循环体，公共路径会把寄存器溢出到内存，并在每个专家上重新加载它们：每个专家约 0.37 µs，而流式加载该专家的权重约需 1.2 µs。这一开销与轻量路径本身所做的工作无关。

两个循环而非一个

修复方案将重型代码排除在轻量路径的循环之外。内核按行数对一步中的专家进行排序，并运行两个循环：一个用于行数至多为四的专家，另一个用于其余专家。现在每个循环体只包含自己的变体，因此少行循环的溢出指令减少了 61–64%。每个专家仍然走与之前相同的路径，因此输出逐位相同。

图 10。示意图。在一个共享循环体中，每个少行专家之后都跟着一次溢出和重新填充，循环落后于其权重流。在它们自己的循环中，少行专家与权重保持同步。

在随机路由的内核测试中，具有许多请求的解码步骤快了 13–19%，专家循环现在与其权重流保持同步。

6. 使用 DFlash2 的推测解码

在没有推测的情况下，每个输出 token 都需要对活动权重进行一次完整遍历。DFlash [21] 在小模型的一次前向传播中起草一整块 token；我们使用 Muse Spark 的 DFlash2 草稿模型。

图 11。一个推测步骤。草稿模型读取从完整模型捕获的隐藏状态。token 条带仅为示意。

草稿在其块内进行双向注意力，并一次性提出该块的所有 token。完整模型在一个 megakernel 步骤中检查整个块，并保留与其自身选择匹配的最长草稿前缀，再加上它自己的一个 token。

这很适合 TPU：MXU 时间随行瓦片和输出瓦片的数量增长，而不是随乘加运算增长，并且一个块的少数几行与一个 token 适合放在同一个行瓦片中，因此验证一个块的成本仅比验证一个 token 略高。在 GSM8K 上以真实接受率测量，DFlash2 将单请求下 megakernel 每个输出 token 的时间从 3.93 ms 降至 1.95 ms，即 2.0 倍。

7. 每设备更多上下文：KV 布局与 L2 HiCache

在智能体流量中，KV 容量就是吞吐量：每个从缓存中掉出的前缀都会以 100K token 的重新预填充形式返回。在权重之后，每个设备仅剩约 20 GiB 用于 KV，分为用于 16 个全局层的全注意力池和用于其他 46 层的滑动窗口池。

按 lane 序列布局的 fp8 KV

每个设备持有 16 个 KV 头中的 2 个，每个头 64 个值宽：即向量寄存器 128 个 lane 的一半。按 token 主序的缓存将一个头的 K 和 V 并排放置以填满 lane。这在 bf16 中可行，但 fp8 [22] 沿子 lane 将四个值打包进一个 32 位字，因此一对 K 和 V 会被填充到四个槽位，仅切换到 fp8 并不能节省任何空间。

按 lane 顺序的布局把页面的 token 沿 128 条 lane 排列，把 head 维度放在子 lane 上，因此每个词持有同一个 token 的四个 fp8 值，且没有任何填充。它把 KV 内存和每一步读取的 KV 字节数都减半。Google 工程师将此布局加入 SGLang-TorchTPU 的 ragged paged attention 内核（#344），该内核在此处负责 prefill，TorchTPU 博文 针对 Qwen 3.5 描述了它。我们为此重写了 megakernel 的 attention 读取和 KV 写入。

滑动窗口池中草稿模型的 KV

草稿模型只关注最后 2,048 个 token，但服务器为其分配了完整上下文的 KV，几乎占每个缓存 token 全部 KV 内存的三分之一。从滑动窗口池中提供这部分 KV，使同样的 HBM 能容纳多 38% 的 token。

TPU 上的 L2 HiCache

我们在 TPU 上启用了 SGLang 的分层缓存（HiCache）[23]，以主机 DRAM 作为每个 HBM 池背后的 L2：full-attention、滑动窗口和草稿。当 HBM 在工具暂停期间驱逐某个会话的前缀时，其页面保留在主机内存中，并在下一轮复制回来，而不是重新计算，正如 CachedAttention [24] 等多轮 KV 缓存系统所做的那样。要使其与投机解码和重叠调度器协同工作，需要一个新的同步栅栏。在 agent 流量上，L2 缓存将吞吐量提升了最高 77%。

会话感知驱逐

Agent 流量带来了第二个问题：缓存不知道哪些旧前缀会再次被需要。一个 agent 会话是一棵对话树：主 agent 可以在其子 agent 运行时闲置一刻钟，然后带着整个上下文回来。对于普通的 LRU 驱逐，那个闲置对话看起来像缓存中最冷的数据，因此它首先被驱逐，其下一轮就要重新计算一个长前缀。混合模型使损失更严重：重用前缀需要其 full-attention KV 和最后 2,048 个 token 的滑动窗口 KV，因此只丢失小窗口就会使大的 full-attention KV 变得无用。在此 trace 上，prefill 计算的 token 数约为最佳可能缓存的 1.2 倍。

图 12. 示意图。主 agent 在其子 agent 运行时闲置。普通 LRU 首先驱逐其 KV；会话感知驱逐将其保留在主机内存中，直到其他所有内容都被驱逐。

SGLang 的会话感知基数树缓存 [5]（sgl-project/sglang#29173）为缓存提供了这种知识。来自同一对话的请求共享一个会话 id，活跃会话将重用的 KV 只有在其他所有内容之后才会被驱逐。我们将其适配到 agent 流量：会话 id 来自客户端已经发送的每对话 id，保护仅适用于主机内存，因此 HBM 对正在运行的请求保持普通 LRU，而主机内存则持有闲置对话。在图 1 的最高吞吐运行中，prefill 现在的计算量与最佳可能缓存相差在 1–5% 以内。

8. Prefill：融合 MoE 内核与公平批处理

在请求可以开始解码之前，服务器以 4,096 个 token 的块 [25, 26] 将其新提示 token 通过模型运行：即 prefill。一个块有足够的行来保持矩阵单元忙碌，因此与解码不同，prefill 时间用于计算以及在设备之间移动结果。在 agent 流量上，单轮可能带来数十万个新 token，因此 prefill 速度决定了请求在收到第一个 token 之前等待多久。

在预填充的 MoE 层中，每台设备都会为每个（token，专家）对计算部分结果，而每个 token 的所有者必须将其中 8 个相加。以每次一小笔写入的方式搬运这些部分结果，比矩阵运算本身开销更大。

融合 MoE 内核与配对交换

专家通过张量并行进行分片：每台设备持有每个专家的 1/8。因此，每台设备都会针对一个预填充块中的每一（token，专家）行，与其分片进行计算，为每行生成一个全宽的部分结果；Muse Spark 在加权求和之前会对每个专家的输出进行归一化，因此一行的 8 个部分结果必须全部相加后，其所有者才能应用归一化并合并专家。该内核在流式加载一个专家的权重的同时预取下一个，执行 gate/up、SwiGLU [27] 和 down，并将每个完成的行直接写入其所有者的 HBM，与后续 tile 重叠。

图 13。 芯片上的两台设备通过 die-to-die 链路相连，Google 文档称其比连接到另一芯片的 1D ICI 链路快六倍。

每行的写入在发出矩阵乘法的同一单元上大约需要 50 ns 的标量工作，因此写入与计算串行化。配对交换则利用芯片的拓扑结构。每台设备通过快速片上链路将其整个 tile 发送给同一芯片上的兄弟设备，然后该对的求和结果作为每次连续行运行的一次 DMA，一次性发送给每个远程所有者。跨芯片字节数减半，每个所有者相加 5 个块而不是 8 个，MoE 内核从每层 3,247 µs 降至 2,368 µs。

较小的改动削减了路径的其余部分：

• 权重乒乓。 交替使用两个权重缓冲区，使每个专家的获取可以在前一个专家计算时开始。
• 单遍合并。 所有者的部分结果求和、按专家归一化和加权求和原本作为三个 XLA 遍运行；一个融合的 Pallas 遍只读取一次部分结果。
• 无零填充，展开写入。 每个输出行恰好写入一次，因此输出缓冲区的零填充被去掉，并且逐行 DMA 循环每次迭代发出 8 行。

结合这些较小的改动，配对交换将冷启动 100K token 提示下的 4,096 token 预填充块从 440 ms 缩短至 384 ms。

序列并行预填充

张量并行将每个矩阵乘法拆分到 8 台设备上，但每台设备仍持有整个块的激活值，因此矩阵乘法之间逐 token 的工作（如归一化、残差更新和路由）在每台设备上重复执行。序列并行 [28] 则改为拆分 token：在层之间，每台设备只保留该块 4,096 个 token 中属于自己的 512 个。归一化和残差更新随后在八分之一的 token 上运行，注意力的 all-reduce 变为 reduce-scatter，后者只向每台设备交付其自身 token 的求和结果。路由将每个 token 的专家 id 和权重打包成一次小型 all-gather。MoE 输出通过 all-to-all 到达其 token 所在的设备：每台设备只从其他七台设备接收自己 token 的列，字节数是被它取代的 all-gather 的八分之一，后者将全部 4,096 个 token 的列发送给每台设备。这在 131K 上下文下将 4,096 token 的前向传播缩短了 51 ms，约 11%。它改变了一些归约的顺序，因此结果并非逐位相同，我们的准确性检查已覆盖这一点。

块共享

一个数十万 token 的冷 prompt 以每块 4,096 token、每块约半秒的速度进行预填充。调度器过去会把每个块都分配给该 prompt，直到它完成，因此其他轮次的几千个新 token 要在它后面等待 15 到 95 秒。

图 14. 示意图。长 prompt 在每批中至少保留 1,024 个 token，等待中的轮次填充其余部分，长 prompt 稍晚一点完成。

现在，处于分块预填充中的 prompt 在每批中至少保留 1,024 个 token，等待中的轮次填充其余部分。在 AgentX 测试中，这将端到端延迟中位数降低了 58%，p90 降低了 33%，而吞吐量变化在 5% 以内。

未来工作

• 让注意力更接近内存下限，通过将 64 宽的头打包以填满 MXU，并进一步缩短每块的链。
• 拓扑感知的 MoE 布局，跨芯片进行专家并行，在每个芯片内部进行张量并行，使大部分 MoE 交换在本地完成。
• 更快的预填充，这现在是高并发下的瓶颈，通过更大的预填充图、与计算重叠的集合通信以及重构的 MoE 预填充内核来实现。
• 预填充/解码（P/D）分离，将预填充和解码放在不同主机上以优化 AgentX 性能。
• 上游合入，将这项工作合入 SGLang-TorchTPU 的主分支。

代码

这里的一切都是开源的，位于 SGLang-TorchTPU，分支 haolin/ms-v11。megakernel 位于 sglang_tpu/muse_megakernel/，融合 MoE 预填充内核位于 sglang_tpu/moe/pallas/fused_rs/。

致谢

我们感谢 Google 和 Meta 对 SGLang-TorchTPU 的内核贡献，这项工作建立在其之上，并特别感谢 Qi Zhou、Yunlong Liu、Charles Chen、Yin Lin、Jiaxin Cao 和 Anne Stern 在本文所述内核和 AgentX 优化方面的工作。我们也感谢 SemiAnalysis、Google、Meta 和 SGLang 社区提供 InferenceX、TPU v7 和 TorchTPU、Muse Spark Open 以及 SGLang。

参考文献

1.  N. Shazeer, A. Mirhoseini, K. Maziarz, et al. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017. arXiv:1701.06538
2. I. Beltagy, M. E. Peters, A. Cohan. Longformer: The Long-Document Transformer. 2020. arXiv:2004.05150
3. B. D. Rouhani, R. Zhao, A. More, et al. Microscaling Data Formats for Deep Learning. 2023. arXiv:2310.10537
4.  M. Shoeybi, M. Patwary, R. Puri, et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019. arXiv:1909.08053
5. L. Zheng, L. Yin, Z. Xie, et al. SGLang: Efficient Execution of Structured Language Model Programs. NeurIPS 2024. arXiv:2312.07104
6. Y. Leviathan, M. Kalman, Y. Matias. Fast Inference from Transformers via Speculative Decoding. ICML 2023. arXiv:2211.17192
7. C. Chen, S. Borgeaud, G. Irving, et al. Accelerating Large Language Model Decoding with Speculative Sampling. 2023. arXiv:2302.01318
8. B. Spector, J. Juravsky, et al. Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B. Hazy Research, 2025. 博客文章
9. X. Cheng, Z. Zhang, Y. Zhou, et al. MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs. OSDI 2026. 论文
10. G. Novack, X. Liu, J. Ma, W. Kwon. TPU megakernels for Kimi K3 and Qwen3.8-27B. 2026. 代码
11. M. Lam. Software Pipelining: An Effective Scheduling Technique for VLIW Machines. PLDI 1988. doi:10.1145/53990.54022
12. S. Wang、J. Wei、A. Sabne 等。Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models。ASPLOS 2023。doi:10.1145/3567955.3567959
13. B. Zhang、R. Sennrich。Root Mean Square Layer Normalization。NeurIPS 2019。arXiv:1910.07467
14. M. Milakov、N. Gimelshein。Online Normalizer Calculation for Softmax。2018。arXiv:1805.02867
15. T. Dao、D. Y. Fu、S. Ermon 等。FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness。NeurIPS 2022。arXiv:2205.14135
16. A. Henry、P. R. Dachapally、S. Pawar 等。Query-Key Normalization for Transformers。Findings of EMNLP 2020。arXiv:2010.04245
17. M. Dehghani、J. Djolonga、B. Mustafa 等。Scaling Vision Transformers to 22 Billion Parameters。ICML 2023。arXiv:2302.05442
18. K. Hong、G. Dai、J. Xu 等。FlashDecoding++: Faster Large Language Model Inference on GPUs。MLSys 2024。arXiv:2311.01282
19. T. Dao、D. Haziza、F. Massa、G. Sizov。Flash-Decoding for Long-Context Inference。Stanford CRFM，2023。博客文章
20. G. J. Chaitin。Register Allocation & Spilling via Graph Coloring。SIGPLAN Symposium on Compiler Construction，1982。doi:10.1145/800230.806984
21. J. Chen、Y. Liang、Z. Liu 等。DFlash: Block Diffusion for Flash Speculative Decoding。2026。arXiv:2602.06036
22. P. Micikevicius、D. Stosic、N. Burgess 等。FP8 Formats for Deep Learning。2022。arXiv:2209.05433
23. SGLang 团队。SGLang HiCache: Fast Hierarchical KV Caching with Your Favorite Storage Backends。LMSYS Org，2025。博客文章
24. B. Gao、Z. He、P. Sharma 等。Cost-Efficient Large Language Model Serving for Multi-turn Conversations with CachedAttention。USENIX ATC 2024。arXiv:2403.19708
25. A. Agrawal、A. Panwar、J. Mohan 等。SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills。2023。arXiv:2308.16369
26. A. Agrawal、N. Kedia、A. Panwar 等。Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve。OSDI 2024。arXiv:2403.02310
27. N. Shazeer。GLU Variants Improve Transformer。2020。arXiv:2002.05202
28. V. Korthikanti、J. Casper、S. Lym 等。Reducing Activation Recomputation in Large Transformer Models。MLSys 2023。arXiv:2205.05198

文章 Optimizing Muse Spark Open on TPU v7 with SGLang-TorchTPU 首次发表于 PyTorch。