单主机承载8160亿参数模型：SGLang-TorchTPU如何压榨TPU v7性能

导语

大参数模型推理的瓶颈，往往不只在矩阵乘法本身。尤其在小批量解码阶段，计算单元可能很快完成工作，真正拖慢响应的是权重从HBM搬运、KV缓存读取、跨设备通信，以及多个内核之间无法连续衔接的调度开销。PyTorch博客披露的一项实践，正是围绕这些“计算之外”的成本展开。

团队使用SGLang的调度、KV缓存和模型服务能力，并以TorchTPU作为PyTorch后端，在一台包含四颗TPU v7芯片的主机上部署Meta的Muse Spark Open。该模型共有8160亿参数，每个token约激活420亿参数；62个解码层全部采用混合专家结构，每层从256个专家中选择8个。模型还结合滑动窗口注意力与最长约100万token的全上下文注意力，因此同时考验计算、显存和长上下文缓存管理能力。

核心要点

• 用一个Megakernel覆盖完整步骤。 过去一次解码或验证步骤要经过多个内核和集合通信，边界会中断权重流。新的Pallas程序把62层及最终归一化放进同一次设备调用，让后续权重、KV数据和跨设备传输可以在计算进行时提前启动。
• 把通信隐藏在计算之后。 方案通过DMA流水线，让权重在真正使用前到达，把KV分块读取、部分和传递以及设备间交换穿插到层间执行中。这样做尤其适合小批量场景，因为此时计算量不足以掩盖内存访问延迟。
• 针对解码细节继续优化。 实现包含集体通信折叠、无需显式最大值的注意力循环、减少寄存器溢出的专家循环，以及DFlash2推测解码。推测模型先提出多个token，再由完整模型一次验证，减少逐token执行的次数。
• 扩大可用KV缓存。 Megakernel采用按序列和设备通道组织的KV布局；同时启用TPU上的L2 HiCache，把主机DRAM作为更高层缓存，并结合会话感知的淘汰策略，以支持更长上下文和前缀复用。
• 改善预填充阶段。 融合MoE内核和序列并行用于提升prefill效率，使长提示词处理不完全受限于单一设备的执行路径。

性能与意义

在较早的Muse Spark检查点上，Megakernel配合DFlash2后，单请求每输出token耗时为1.95毫秒，原生路径为3.78毫秒，速度提升约1.94倍；两路并发时提升约1.85倍。随着并发增加，优势会收窄，因为批量化本身能够提高硬件利用率，内核边界不再是唯一主要开销。验证步骤的耗时也被压缩到接近仅从HBM读取相关权重所需时间的1.5至1.7倍，说明系统逐渐逼近内存带宽限制。

在SemiAnalysis的AgentX百万上下文轨迹测试中，该方案覆盖从高交互性低吞吐到高吞吐低单请求速率的不同工作点。测试使用四芯片TPU v7主机，计算采用FP8，专家权重采用NVFP4；结果以每美元token数和p90交互延迟观察服务效率。需要注意的是，这些数字来自特定模型检查点、硬件配置和成本口径，不能直接等同于所有模型或云环境的通用结论。

这项工作的价值在于，它把TPU推理优化从“写出更快算子”推进到“重新组织整个执行步骤”。对于MoE和超长上下文模型，权重流、KV缓存、跨设备通信与调度边界往往共同决定体验。将这些环节放进同一可流水化程序，再配合分层缓存和推测解码，可能比单独优化某个矩阵乘法更接近真实服务瓶颈。其代价则是更复杂的内核编排、布局设计和模型适配；未来能否扩展到更多模型与动态负载，仍取决于编译器、运行时和服务框架的协同能力。

来源：PyTorch Blog