用 SGLang-TorchTPU 在 TPU v7 上优化 Meta Muse Spark Open

PyTorch：Blog· Haolin Fu, Xiaoyu Zhang, Junyan Xu, Wangyuan Zhang, Haoguang Cai, RadixArk Team·2026-11-11 00:45· 1 小时前AI 评分46

使用 SGLang-TorchTPU 在 TPU v7 上优化 Muse Spark Open

AI 导读

Meta 的 816B 参数 MoE 模型 Muse Spark Open 在单台四芯片 TPU v7 主机上通过 SGLang-TorchTPU 部署，核心是把全部 62 层的解码或验证步骤合并为一个 Pallas megakernel，并将权重加载、KV 读取和跨设备传输全部流水化。

正文 · 原文

TL;DR

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

• 解码速度。 megakernel 将一个或两个并发请求的每输出 token 时间减半，这是它帮助最大的场景；随着并发增长，收益收窄。其验证步骤的运行时间在仅从 HBM 流式加载其权重所需时间的 1.5–1.7 倍以内。
• AgentX。 在 SemiAnalysis 的 agentic 编码基准的 1M 上下文版本上，我们的曲线从每请求 433 输出 token/s、每芯片 6,101 token/s，延伸到每芯片 25,906 token/s、每请求 95 token/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 token；上下文窗口较短的模型运行上限为 256K 的版本。吞吐量统计输入和输出 token。图 1 中每个 TPU 数据点都在 p90 下于 47 秒内完成请求。

背景

模型

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

• 62 个解码层，全部为 MoE。 每一层将每个 token 路由到 256 个专家中的 8 个。
• 混合注意力。 46 层对最后 2,048 个 token 使用滑动窗口注意力 [2]；其余 16 层关注完整上下文，最多 1M token。注意力有 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 位值的二维向量寄存器上计算，即 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：PyTorch 中的 SGLang [5] 调度器、KV 缓存和模型代码，热路径上使用 Pallas 内核。我们使用推测解码 [6, 7] 进行解码：一个小型草稿模型猜测接下来的几个 token，完整模型在一次前向传播中检查它们全部，即验证步骤。验证步骤中的每个（请求，位置）对是一行，因此两个请求各验证 8 个 token 构成一个 16 行的步骤。

1. 一个巨型内核

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

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

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

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

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

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

2. 在 megakernel 内部流水化一切

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

图 5. 一步中的一层，示意图。每条 DMA 流都先于或伴随消费它的计算运行。

权重在需要之前到达

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

集合通信在数据一存在时就发出

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

• 流式 o-proj。每个 o-proj 输出分块一经计算完成，就立即发送给其所属设备，而此时下一个分块仍在进行乘法运算。

流水线式 KV 写入

每一步都会将其新的 K 和 V 行写入缓存。在我们的 fp8 序列-通道缓存（第 7 节）中，一个 token 是 256-token 页面的一个通道列，因此写入一个 token 意味着将整个页面块读入 VMEM，合并新列，再将页面写回。

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

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

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

3. 集合通信折叠

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

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

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

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

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

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

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

运行最大值将每个块链接起来

FlashAttention [15] 所使用的在线 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]，而这些额外的存储和加载在每次经过该代码时都会耗费时间。在具有多个分支的循环中，编译器必须将所有分支的活跃值放入同一组寄存器，因此一个繁重的分支可能使每次循环迭代都发生溢出，即使该分支从未执行。

megakernel 的专家循环恰好存在这个问题。一个步骤会遍历其行被路由到的专家。对每个专家，kernel 都要等待该专家的权重，并根据该专家服务的行数选择一条执行路径：少行路径（针对 1 到 8 行有小变体），以及针对更繁忙专家的反量化后相乘路径。随着 batch 增大，一个步骤会访问更多专家，但其中大多数仍只服务一两行，因此轻量路径才是关键所在。

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

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

两个循环取代一个循环

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

图 10。示意图。在共享的循环体中，每个少行专家之后都会发生一次溢出和重新填充，循环跟不上其权重流。在各自的循环中，少行专家能与权重保持同步。

在随机路由的 kernel 测试中，包含大量请求的 decode 步骤快了 13–19%，专家循环现在也能跟上其权重流。

6. 使用 DFlash2 的推测解码

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

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

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

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

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

在 agent 流量下，KV 容量就是吞吐量：每个掉出缓存的 prefix 都会以 100K token 的重新 prefill 回来。在权重之后，每个设备仅剩约 20 GiB 用于 KV，分为 16 个全局层的全注意力池和其他 46 层的滑动窗口池。

采用 sequence-on-lane 布局的 fp8 KV

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

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

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

草稿模型只关注最后 2,048 个 token，但服务器为其分配了完整上下文的 KV，几乎占每个缓存 token 全部 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 kernel 与公平批处理

在请求开始解码之前，服务器会以 4,096 个 token 为一块，将新的提示词 token 分批送入模型运行 [25, 26]：这就是预填充。一个块有足够多的行让矩阵单元保持忙碌，因此与解码不同，预填充的时间花在计算以及设备之间的结果搬运上。在智能体流量下，单轮对话就可能带来数十万个新 token，因此预填充速度决定了请求在拿到第一个 token 之前要等多久。

在预填充的 MoE 层中，每台设备都会为每个（token，专家）对计算一个部分结果，而每个 token 的归属设备必须把其中 8 个部分结果加起来。搬运这些部分结果时，每次只做一次小写入，其开销超过了矩阵运算本身。

融合 MoE 内核与配对交换

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

图 13。 芯片上的两台设备共享一条裸片到裸片链路，Google 文档称其速度比通往另一芯片的 1D ICI 链路快六倍。

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

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

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

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

序列并行预填充

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

Chunk 共享

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

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

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

未来工作

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

代码

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

致谢

我们感谢 Google 和 Meta 对 SGLang-TorchTPU 的 kernel 贡献，本工作正是建立在其之上；并特别感谢 Qi Zhou、Yunlong Liu、Charles Chen、Yin Lin、Jiaxin Cao 和 Anne Stern 在本文所述 kernel 和 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 等。Megatron-LM：使用模型并行训练数十亿参数语言模型。2019。arXiv:1909.08053
5. L. Zheng、L. Yin、Z. Xie 等。SGLang：结构化语言模型程序的高效执行。NeurIPS 2024。arXiv:2312.07104
6. Y. Leviathan、M. Kalman、Y. Matias。通过推测解码实现 Transformer 的快速推理。ICML 2023。arXiv:2211.17192
7. C. Chen、S. Borgeaud、G. Irving 等。使用推测采样加速大型语言模型解码。2023。arXiv:2302.01318
8. B. Spector、J. Juravsky 等。看，没有气泡！为 Llama-1B 设计低延迟巨型内核。Hazy Research，2025。博客文章
9. X. Cheng、Z. Zhang、Y. Zhou 等。MPK：用于张量程序巨型内核化的编译器和运行时。OSDI 2026。论文
10. G. Novack、X. Liu、J. Ma、W. Kwon。面向 Kimi K3 和 Qwen3.8-27B 的 TPU 巨型内核。2026。代码
11. M. Lam。软件流水线：一种面向 VLIW 机器的有效调度技术。PLDI 1988。doi:10.1145/53990.54022
12. S. Wang、J. Wei、A. Sabne 等。通过大型深度学习模型中的分解实现通信与依赖计算的重叠。ASPLOS 2023。doi:10.1145/3567955.3567959
13. B. Zhang、R. Sennrich。均方根层归一化。NeurIPS 2019。arXiv:1910.07467
14. M. Milakov、N. Gimelshein。Softmax 的在线归一化计算。2018。arXiv:1805.02867
15. T. Dao、D. Y. Fu、S. Ermon 等。FlashAttention：具有 IO 感知的快速且内存高效精确注意力。NeurIPS 2022。arXiv:2205.14135
16. A. Henry、P. R. Dachapally、S. Pawar 等。Transformer 的查询-键归一化。Findings of EMNLP 2020。arXiv:2010.04245
17. M. Dehghani、J. Djolonga、B. Mustafa 等。将视觉 Transformer 扩展到 220 亿参数。ICML 2023。arXiv:2302.05442
18. K. Hong、G. Dai、J. Xu 等。FlashDecoding++：在 GPU 上更快的的大型语言模型推理。MLSys 2024。arXiv:2311.01282
19. T. Dao、D. Haziza、F. Massa、G. Sizov。面向长上下文推理的 Flash-Decoding。Stanford CRFM，2023。博客文章
20. G. J. Chaitin。通过图着色的寄存器分配与溢出。SIGPLAN 编译器构造研讨会，1982。doi:10.1145/800230.806984
21. J. Chen、Y. Liang、Z. Liu 等。DFlash：用于 Flash 推测解码的块扩散。2026。arXiv:2602.06036
22. P. Micikevicius、D. Stosic、N. Burgess 等。面向深度学习的 FP8 格式。2022。arXiv:2209.05433
23. SGLang 团队。SGLang HiCache：使用你喜爱的存储后端实现快速分层 KV 缓存。LMSYS Org，2025。博客文章
24. B. Gao、Z. He、P. Sharma 等。使用 CachedAttention 为多轮对话提供高性价比的大型语言模型服务。USENIX ATC 2024。arXiv:2403.19708
25. A. Agrawal、A. Panwar、J. Mohan 等。SARATHI：通过分块预填充搭载解码实现高效 LLM 推理。2023。arXiv:2308.16369
26. A. Agrawal、N. Kedia、A. Panwar 等。使用 Sarathi-Serve 驯服 LLM 推理中的吞吐量-延迟权衡。OSDI 2024。arXiv:2403.02310
27. N. Shazeer。GLU 变体改进 Transformer。2020。arXiv:2002.05202
28. V. Korthikanti、J. Casper、S. Lym 等。减少大型 Transformer 模型中的激活重计算。MLSys 2023。arXiv:2205.05198

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

来源：PyTorch：Blog · pytorch.org