5.0
关注
AI SCORE
技术实践2026-07-30 21:18
千问 3.5 MoE 在 TPU 上的系统优化实战
Google Developers Blog#MoE模型#系统优化#TPU
Editor brief · 编辑速览
讲解 397B 参数 MoE 模型在 Ironwood TPU 上的混合数据并行部署方案，实现 4.7x 加速。系统工程深度强但适用范围极窄。
在专用硬件加速器上部署和服务像 Qwen3.5-397B 这样的混合专家（MoE）模型呈现了重大的系统工程挑战。将 400 GB 的权重加载到高带宽内存（HBM）并最大化硬件利用率需要采用严格的第一性原理工程方法，而非经验式的试错修改。
至关重要的是，随着开源权重模型的景观日益复杂，工程团队再也无法承受花费数个月时间分别优化每个新模型架构的局面。为了解决这一可扩展性挑战，我们的性能团队开创了一套模块化、与模型无关的优化策略。与其将模型视为整体系统来处理，我们将其分解为自包含、独立的构建块（例如 Batched RPA、Grouped GEMMs 和 SparseCore 反置换），并配合硬件感知的成本模型。当新的架构出现时，这些预优化的模块可以以近零工程摩擦进行移植。这使我们的工程师能够提前远超预期交付最先进的服务性能，将重点从本地化的模型优化转向全局的平台级可扩展性。
本技术报告详细说明了我们如何系统地将这一全局优化手册应用于 Ironwood（TPU v7x）平台上的 Qwen 3.5 MoE。通过利用我们的可复用 JAX/Pallas 内核库，并专注于 Qwen 3.5 的新颖组件——如 Gated DeltaNet（GDN）线性注意力和注意力数据并行——我们的团队为解码密集和预填充密集工作负载都实现了显著的性能提升。
下述优化使我们在 2026 年 4 月至 6 月期间实现了性能提升：解码密集工作负载约 3.1 倍，预填充密集工作负载（512 并发度）约 4.7 倍。此外，通过将这些模块化优化原生集成到 vLLM 和 SGLang 等开源服务框架中，我们消除了传统软件障碍，为全球企业工作负载大规模部署提供了一条无缝的、生产就绪的迁移路径。
该模型包含 397 亿总参数，但采用了高度稀疏的路由方案，在每次前向传播中每个令牌仅激活 170 亿个参数。这种稀疏配置代表 4.3% 的路由激活比，使模型能够在保持推理占用和执行速度与小得多的 20B 级系统相当的同时，传递 400B 级模型的表达能力和智能。
官方模型权重和配置可直接通过 Qwen3.5-397B-A17B Hugging Face 存储库访问。关于 Qwen 3.5 混合线性注意力和门控组件的全面结构分析，请参阅 Qwen3.5：没有人对注意力达成一致（Hugging Face 博客）和 Gated DeltaNet 线性注意力（Sebastian Raschka，PhD）上的技术深潜。
网络总共包含 60 层，隐藏维度 D=4096，以及填充词汇表大小为 248,320 个令牌。Qwen 3.5 并未采用统一的 Transformer 层堆栈，而是采用了由 15 个重复结构块组成的高度定制的混合布局。每个块按 3:1 比例排列：
这个重复序列可以表示为：
[GDN, GDN, GDN, GQA] × 15
该模型的混合性质整合了三个不同的数学表述：
Gated DeltaNet 线性注意力：标准自注意力机制的计算复杂度随序列长度呈二次方增长（O(S²)），这为长上下文生成创造了计算瓶颈。GDN 通过计算线性注意力来解决这一问题，对值（V）使用 64 个线性注意力头，对查询和键（QK）使用 16 个头，头维度为 128。GDN 不是构造成对的 softmax 注意力矩阵，而是为每个头维护一个常数大小的隐状态矩阵（与 d_k 和 d_v 键值维度相匹配），该矩阵充当循环内存。
在每个令牌步骤 t，状态矩阵使用增量规则更新：
h_t = α_t ⊙ h_{t-1} + k_t ⊗ v_t
其中 q_t、k_t 和 v_t 分别是查询、键和值向量，α_t 是一个学习的门控参数。此循环更新前面有一个因果的一维卷积（K=4）以捕获局部空间依赖关系。这种循环表述允许上下文窗口在内存方面呈线性扩展（O(S)），保持循环状态占用恒定。
分组查询注意力（GQA）：为了锚定线性注意力检索，该模型在其 25% 的层中使用标准 GQA。GQA 全局使用 32 个查询头（Nq=32）和恰好 2 个键值（KV）头（Nkv=2），头维度为 256，旋转位置嵌入（RoPE）维度为 64。这种极端的 GQA 布局在生成过程中压缩 KV 缓存占用，但施加了严格的硬件级分片约束，详见第 3 节。
路由稀疏混合专家：前馈网络（FFN）层分片成 512 个小专家，中间专家维度为 1024。执行过程中，路由门投影令牌表示并通过 softmax 概率分布选择前 10 个路由专家。至关重要的是，该模型还加入了一个总是执行的共享专家路径，充当通用表示层：
y = Router(x) + SharedExpert(x)
这种原生多模态 MoE 架构通过早期融合训练范式原生处理文本、图像和视频输入，在数万亿多模态令牌上进行训练。上下文窗口支持 262,144 个令牌的原生上下文长度，可使用 YaRN RoPE 缩放扩展到超过 1,010,000 个令牌。
为了系统地隔离、分析和解决编译器和内核瓶颈，系统工程团队基于实际的非对称工作负载建立了严格的多维度评估矩阵。
我们的基准测试跨越设计用于强调不同硬件执行子系统的非对称工作负载进行：
预填充密集工作负载（8K 输入 / 1K 输出）：以长输入提示序列和短令牌生成输出为特征。这些工作负载是计算受限的，严重强调 TPU 张量核矩阵执行单元（MXUs）的物理浮点矩阵乘法容量。
解码密集工作负载（1K 输入 / 8K 输出）：以短输入提示和延长的逐令牌生成阶段为特征。这些工作负载是内存受限的，因为系统必须不断将所有 400 GB 参数从高带宽内存（HBM）流式传输到执行核心以为每个请求生成单个令牌。
并发度层级：为了观察系统扩展曲线并在负载下识别硬件排队/内存瓶颈，这两种工作负载都在四个并发度层级上进行评估：64、128、256 和 512 个并发请求。
基准测试在企业级单主机集群上执行：
加速器拓扑：一台物理主机，容纳 4 个物理 Ironwood 芯片。每个物理芯片由 2 个逻辑芯片组成，暴露 8 个不同执行核心（设备）的逻辑拓扑，这些核心通过高速、亚微秒级片间互连（ICI）平面相连接。
推理服务器引擎：使用 vllm-project/tpu-inference 设计实现。对于采用注意力 DP 的最终优化运行，服务器执行循环配置为 --max-num-batched-tokens=1024 和 --max-num-seqs=64 每核心（与早期张量并行基准中使用的 --max-num-batched-tokens=8192 和 --max-num-seqs=512 相比）。
指标追踪：性能被追踪并报告为每芯片令牌吞吐量（TPS/chip），计算方式为处理的总令牌数（输入 + 输出）除以执行时长和物理芯片数量（4）。
Qwen 3.5 的特定架构约束——即 GQA 层中恰好 2 个 KV 头，以及 MoE 层中 512 个专家——使传统统一分片方法失效。
在标准注意力张量并行（TP）+ 专家 MoE 配置中，注意力权重在设备维度上被切片和分片。然而，试图以张量并行大小 8（TP=8）对 GQA 层进行分片会强制分数头分片（2/8 = 0.25 头每设备），这在硬件上物理上是不可能的。
在本地跨 8 个核心复制注意力头会在每个设备上增加物理 KV 缓存的内存占用，抵消了 GQA 的内存节省优势。这种内存冗余严重限制了在高负载工作负载下可用于活跃 KV 缓存的 HBM 空间。这一容量限制迫使服务器引擎将实际达到的并发数限制在远低于预期的水平——仅能支持约 200 个并发请求，而不是计划的 512 个。
为了消除这一瓶颈，我们共同设计了一种混合分片方案（PR #2577）：8 路注意力批分片（数据并行性，DP=8）结合 MoE 层中的 8 路专家并行性（EP=8）。
在所有 8 个设备上复制 GQA 和 GDN 权重，允许每个核心使用完整的 2 个 KV 头在本地处理注意力，保留了本地 KV 缓存的一致性，消除了注意力内部分片通信。在前馈 MoE 层中，我们改用专家并行性（EP=8）。512 个路由专家均匀分布（每个设备 64 个专家），这既避免了在所有节点间复制 400 GB 的参数占用，又保持了集体传输负载的可管理性。
在注意力 DP 和 MoE EP 之间转换需要跨设备令牌路由。在设计我们的混合专家模型（MoE）路由层时，我们评估了两种主要的结构方法来处理这种跨设备转换：
方案 A（All-to-All 洗牌）：此方法利用 All-to-All → 本地 MoE → All-to-All 流程。令牌通过网络动态洗牌到托管其目标专家的特定芯片，在本地计算，然后洗牌回。虽然这最小化了冗余计算，但由于在变量工作负载下进行全局 All-to-All 步骤，会产生巨大且不可预测的网络路由开销。
方案 B（完整令牌复制）：此方法利用 All-Gather → 本地 MoE → Reduce-Scatter 流程。All-Gather 将所有令牌向量跨所有设备复制。每个芯片随后仅过滤和计算其本地专家的输入，稍后通过 Reduce-Scatter 聚合输出。这完全绕过了不可预测的 All-to-All 路由开销，代价是增加了本地内存消耗。
由于确定性延迟对实际服务至关重要，我们选择了方案 B，随后开发了低级通信融合来优化其集体路径。
在朴素的方案 B 实现下，为本地 MoE 计算做准备需要跨集群向每个设备秩广播三个不同的数据片段。假设令牌隐藏维度的本地张量切片形状为 [1024,4096]，我们通常必须执行三个独立的集体操作：
All-Gather 1：令牌隐藏维度（[1024,4096]）。
All-Gather 2：所选专家索引（[1024,10]，假设 topk=10）。
All-Gather 3：门控 topk 权重（[1024,10]）。
每次集体通信调用都会在 TPU 上产生固定的内核启动和网络同步延迟开销。为了优化专家并行性（EP）效率，我们在 PR #2836 中将这三个 All-Gather 合并为两个。由于专家索引（整数）和 topk 权重（浮点）共享相同的张量形状（[1024,10]），我们将它们堆叠、按位转换并将其打包到沿新维度的单个稠密 32 位整数数组（blob）中。这使我们能够在数据维度（ShardingAxisName.MLP_DATA）上运行单个 All-Gather，用于两个路由元数据块，在本地解包它们，从而将路由元数据集体延迟减半。
专家执行后，令牌输出必须返回到其数据并行秩。标准的 All-Reduce 操作在 8 设备网格上效率非常低。我们将其替换为用 Pallas/Mosaic 编写的自定义、原生 TPU 分层 Reduce-Scatter（见 PR #2679）。该集体操作在两个流水线阶段中运行：
芯片内 Reduce-Scatter：同一物理芯片上的逻辑芯粒使用快速、本地共享内存传输交换并求和其数据（速度比芯片间 ICI 带宽快 6 倍）。
跨芯片 Reduce-Scatter：部分归约的数据使用 TPU 物理 ICI 链路上的递归倍增超立方体算法在物理芯片间交换。
为了防止 VMEM 内存不足（OOM）错误，数据被分片为 2 到 4 个微批次。内核在 TensorCore 对微批次 i-1 执行向量加法时，对微批次 i 的远程 DMA 传输进行流水线操作，将通信延迟隐藏在计算后面。
为了识别我们系统工程的理论边界，并理解执行停顿发生的位置，我们在 64 并发的标准 8K/1K 配置下对 Qwen 3.5 工作负载进行了第一性原理的屋顶线分析。
Tensor Core（TC）频率：2.2 GHz
每个芯片的 Tensor Core 数：2
每个 TC 的 MXU（矩阵执行单元）数：2（每个芯片共 4 个 MXU）
峰值 BF16 性能：2,307 TFLOPS/芯片（(262,144 FLOP/cycle/MXU × 2.2 GHz × 4 MXU = 2,307 TFLOPS)）
峰值 FP8 性能：4,614 TFLOPS/芯片
在 prefill 阶段，64 个提示词批次，每个包含 8,192 个输入令牌，共产生 524,288 个令牌并行处理。
算术强度：投影层中的 GEMM 操作随序列长度和批大小二次方缩放。算术强度（FLOPs/Byte）极高，将执行放在屋顶线模型的计算密集区域深处。
操作边界：受 TPU v7 TensorCore MXU 峰值浮点执行能力限制（FP8 中 4,614 TFLOPS）。
系统瓶颈：MXU 利用不足主要源于令牌在专家间的分布不均。如果一个专家在给定批次中收到明显更多令牌，对应的设备成为滞后者。最小化分组 GEMM 内核中的填充对于缩小实际 TFLOPS 与理论峰值之间的差距至关重要。
在 decode 阶段，模型每步处理 64 个令牌（每个活跃请求 1 个令牌）。
算术强度：为生成一个令牌，系统必须将所有 400 GB 的模型权重从 HBM 流到处理器。算术强度接近 1（~1 FLOP/Byte），将工作负载直接放在内存密集区域。
操作边界：受 HBM 内存带宽限制。
系统瓶颈：主要延迟贡献者是模型参数的 HBM 传输延迟、稀疏 KV 缓存检索期间的 VPU 索引停顿，以及 Gated DeltaNet（GDN）层中的循环状态更新往返。
为了将这些第一性原理硬件约束转化为可规划的软件工程指标，我们使用端到端屋顶线模型对标准评估工作负载（64 并发，采用 8K/1K 预填充密集型和 1K/8K 解码密集型序列长度布局）进行了建模。该分析建立了每个服务阶段的绝对应用级吞吐量边界（每个物理芯片的令牌数/秒）：
Prefill 密集阶段（8K 输入 / 1K 输出）：由于 prefill 阶段是计算密集的，其受到 TensorCore 矩阵执行单元（MXU）峰值浮点执行能力的限制（每个芯片 4,614 TFLOPS FP8）。考虑到 GQA 注意力操作在 8,192 个令牌上的二次方缩放以及标准硬件执行开销，我们的模型建立了估计的最大理论屋顶线吞吐量为 5,170 令牌/秒/芯片（未折扣），以及在标准调度折扣因子下 4,500 令牌/秒/芯片。
Decode 密集阶段（1K 输入 / 8K 输出）：由于为每个活跃流生成一个令牌是内存密集的，性能严格受 HBM 接口带宽限制。跨所有 60 层的总执行延迟计算为每个令牌步长 16.36 毫秒。这导致峰值理论吞吐量为 978 令牌/秒/芯片（未折扣），以及现实的、折扣后的服务屋顶线限制为 850 令牌/秒/芯片。
通过使用 JAX 自定义内核语言 Pallas 编写自定义内核，我们绕过了标准 XLA 降级路径，直接控制 VMEM 布局、寄存器和内存调度。
注意力轨道：PR #1820（RPA v3）和 PR #1961（批处理 RPA）
MoE 轨道：PR #1688（GMM v2）和 PR #2137（SparseCore 不规则 Gather）
GDN 轨道：PR #2149（分块 GDN）和 PR #3016（完全融合 Conv1D 和循环 / 分块 GDN）
管理 25% GQA 层的 KV 缓存需要动态内存分配。我们采用不规则页注意力（RPA）来索引 HBM 中的非连续内存块（见 PR #2632）。
过去通常使用 16 个 token 的块大小，以尽量减少内存碎片。然而在 TPU 上，较小的块会产生巨大的索引开销，导致向量处理单元（VPU）在解码阶段停顿。我们通过将索引粗粒度化，把 KV 页大小调整为 256（通过服务器命令 --block-size=256 启用），解决了这个问题。在并发数为 512 时，这种粗粒度索引将解码步骤的延迟从 428µs 降低至 283µs，实现了 33.8% 的内核级加速。
为了进一步跑满内存总线，我们设计了批处理 RPA 内核。该设计将多个解码流组合到单个已编译的 Pallas 内核中（#PR 2632），从而摊薄 VPU 指令分派延迟、打破顺序请求造成的数据依赖停顿，并改善内存对齐。
Qwen 3.5 中 top_k=10 的细粒度路由因子会产生非 2 的幂次张量维度。此前，在 TensorCore 上对这些数组执行排列与逆排列，会导致大量填充且未对齐的 HBM 内存写入。我们通过 SparseCore-TensorCore 协同设计流程解决了这个问题：
我们编写了一个自定义 Pallas/Mosaic 内核，将 token 路由卸载到 TPU 的 SparseCore（SC）上；SparseCore 是针对间接寻址优化的硬件单元（参见 PR #2137）。SC 读取路由索引，直接从 HBM 对 token 嵌入执行间接 DMA 聚集，并将其写入连续的虚拟缓冲区。这避免了在 HBM 中物化大量填充且未对齐的中间张量，节省了大量内存带宽。
在 GMM V2 内核中，我们将 SwiGLU 激活函数直接融合到主矩阵乘法循环中（门控投影和升维投影被打包，并通过双路 DMA 读取在单个分块中处理），从而避免寄存器内容溢出到 HBM。此外，我们实现了动态有界切片，以尽可能少的填充处理每个专家大小可变的 token 载荷。对于 FP8 运算，我们改用 512 子通道激活量化，以消除 VREG 溢出和内存加载停顿，使 VPU 上的向量算术吞吐量翻倍。
将 token 逆排列和局部归约操作完全卸载到 SparseCore。通过直接在 SC 上执行间接聚集和局部归约，我们避免了在 HBM 中物化经过填充的中间激活张量，将 HBM 读取次数从 20 次降至 10 次、写入次数从 15 次降至 5 次，大幅削减了 MoE 开销。
为最大限度提高硬件效率，我们的实现采用分块级流水线架构，而不是在完整的 [81920,4096] 张量上依次执行局部归约和 8 设备 Reduce-Scatter。工作负载被划分为 4 个不同的分块。分块 1 一旦在 SparseCore 上完成局部逆排列和聚集归约，就会立即通过物理 ICI 链路异步启动其 Reduce-Scatter 集合通信。与此同时，SparseCore 开始为分块 2 执行局部聚集归约。这种严格的分块级流水线有效地将 Reduce-Scatter 的跨设备网络延迟与后续分块的局部计算重叠，并将其隐藏起来。
由于需要持续更新循环状态，占比 75% 的 Gated DeltaNet（GDN）层中的循环状态更新极易遭遇内存带宽瓶颈。
为优化这条路线，性能团队实现了一系列算法融合和精度协同设计：
GDN 循环更新之前有一个因果一维卷积（K=4）。最初，它被编译为独立操作，迫使中间卷积输出写入 HBM 后再从中读回。我们设计了一种寄存器级滑动窗口算法，直接在 TPU 的 VPU 寄存器中缓存历史 token 状态。将一维卷积与 GDN 循环状态更新融合到单个执行块中，消除了 6 次冗余的 HBM 往返访问（参见 PR #2823）。
我们重构了线性注意力更新方程，以利用代数恒等式。通过在数学上重新排列运算，我们在融合的 GDN 内核中完全跳过了昂贵的秩 1 更新后矩阵乘法，从而减少了计算量（参见 PR #2498）。
此外，为进一步跑满向量处理单元（VPU），我们将循环状态空间模型（SSM）的状态变量精度从 Float32 改为 BFloat16。这使 VPU 上的向量算术吞吐量翻倍，同时不会影响数值收敛或输出质量。
为防止批量预填充执行期间的填充开销浪费 MXU FLOPs，我们优化了 JAX 原生分块布局，并在 PR #2218 中引入了专门的序列处理例程，以原生方式处理不规则输入，确保长度可变的序列不会产生拖慢整体处理的尾部任务。
我们没有依赖彼此分离的执行阶段，而是在 PR #3016 中设计并合入了一个完全融合的 Pallas 内核，将因果一维卷积和整个 GDN 循环线性注意力块编译为 VPU 上单一、统一的执行单元。通过直接在本地寄存器中缓存中间序列状态和循环状态，该内核完全不需要从 VMEM 或 HBM 读取中间激活张量，也不需要将其写入其中。
这种寄存器级融合消除了寄存器到内存的同步延迟，并为两个服务阶段都带来了关键的性能提升：
预填充阶段：处理长输入序列提示词时，它显著降低了内存带宽占用，最大限度提高 TensorCore MXU 的浮点运算效率。
解码阶段：它消除了长时间逐 token 生成期间受内存限制的往返访问停顿。
为 Qwen 3.5 提供服务需要管理两种异构注意力状态结构：Gated DeltaNet（GDN）固定大小的循环线性注意力状态，以及 Grouped Query Attention（GQA）动态增长的标准注意力键值（KV）缓存。由于 TPU v7 每颗芯片具有 192GB HBM 容量（例如，相比 Blackwell GB300 GPU 提供的 288GB，容量差异约为 50%），高并发下优化 HBM 占用是一项严峻的系统约束。在 PR #2416 中，我们引入了一种自定义内存布局，用于在 HBM 中对齐并共同存储这些混合注意力状态。该布局最大限度减少填充并防止内存碎片，直接释放了至关重要的 HBM 余量。这项优化