扩展前沿的混合专家（MoE）模型不仅仅靠试错调优

探索如何通过混合注意力数据并行 + MoE 专家并行、优化的路由元数据集合以及 JAX/Pallas 内核融合，在 Google Cloud TPU（Ironwood v7x）上优化 Qwen 3.5-397B，实现解码性能提升约 3.1 倍，预填充性能提升约 4.7 倍。

阅读系统工程手册 → https://goo.gle/4x5hdv5

系统工程手册：在 Ironwood (TPU7x) 上优化 Qwen 3.5-397B MoE

来源：https://developers.googleblog.com/systems-engineering-playbook-optimizing-qwen-35-397b-moe-on-ironwood-tpu7x/

执行摘要

在专用硬件加速器上部署和服务像 Qwen3.5-397B 这样的混合专家（MoE）模型，带来了重大的系统工程挑战。将 400 GB 的权重占用加载到高带宽内存（HBM）中并最大化硬件利用率，需要一套严谨、基于第一性原理的工程方法，而不是依靠经验性的试错修改。

至关重要的是，随着开放权重模型格局日益复杂，工程团队再也无法花费数月时间孤立地优化每一个新模型家族。为了解决这一可扩展性挑战，我们的性能团队开创了一种模块化、与模型无关的优化策略。我们不再将模型视为单一系统，而是将其分解为自包含、独立的构建模块（如批处理 RPA、分组 GEMM 和 SparseCore 反排列），并配备硬件感知的成本模型。当新架构出现时，这些预优化的模块几乎零工程摩擦即可移植。这使得我们的工程师能够在初始预期之前就提供最先进的服务性能，将我们的重点从局部模型优化转移到全局、平台级的可扩展性上。

本技术报告详细介绍了我们如何系统地将这一全局优化手册应用于 Ironwood（TPU v7x）平台上的 Qwen 3.5 MoE。通过利用我们可重用的 JAX/Pallas 内核库，并仅针对 Qwen 3.5 的新颖组件（如门控 DeltaNet（GDN）线性注意力和注意力数据并行），我们的团队在解码密集型和工作负载密集的预填充方面实现了显著的性能提升。

下面讨论的优化使我们能够在 2026 年 4 月至 6 月期间，将解码密集型工作负载的推理性能提升约 3.1 倍，将预填充密集型工作负载（512 并发层）的推理性能提升约 4.7 倍。此外，通过将这些模块化优化原生集成到 vLLM 和 SGLang 等开源服务框架中，我们消除了传统软件障碍，为大规模企业工作负载提供了无缝、生产就绪的迁移路径。

Qwen 3.5 497B 吞吐量图表

1. Qwen 3.5 架构概述与模型配置

该模型共有 3970 亿个参数，但利用了一种高度稀疏的路由方案，每次前向传递每个 token 仅激活恰好 170 亿个参数。这种稀疏配置实现了 4.3% 的路由激活率，使模型能够在保持约 200 亿参数级系统的推理占用和执行速度的同时，提供 4000 亿参数级模型的表达能力和智能。

官方模型权重和配置可直接通过 Qwen3.5-397B-A17B Hugging Face 仓库 (https://huggingface.co/Qwen/Qwen3.5-397B-A17B) 访问。有关 Qwen 3.5 混合线性注意力和门控组件的全面结构分析，请参阅技术深度文章：Qwen3.5：没人再对注意力机制有共识了（Hugging Face 博客）(https://huggingface.co/blog/mlabonne/qwen35) 和用于线性注意力的门控 DeltaNet（Sebastian Raschka 博士）(https://sebastianraschka.com/llms-from-scratch/ch04/08_deltanet/)。

架构布局与交错

网络总共包含 60 层，隐藏维度 D=4096，填充词汇表大小为 248,320 个 token。Qwen 3.5 没有采用统一的 Transformer 层堆栈，而是采用了一种高度定制化的混合布局，由 15 个重复的结构块组成。每个块以 3:1 的比例排列：

• 门控 DeltaNet 层（占 75% 的层）：3 个连续层，结合了门控 DeltaNet（线性注意力）与路由稀疏 MoE。
• 分组查询注意力（GQA）层（占 25% 的层）：1 层标准 GQA，结合路由稀疏 MoE。

这个重复序列可以表示为：

重复序列公式

图片 (7)

关键数学元素

该模型的混合特性融合了三种不同的数学公式：

1. 门控 DeltaNet（GDN）

标准的自注意力机制随序列长度呈二次方扩展（O(S2)），为长上下文生成造成了计算瓶颈。GDN 通过计算线性注意力解决了这一问题，利用 64 个线性注意力头处理值（V），16 个头处理查询和键（QK），头维度为 128。GDN 不构建成对的 softmax 注意力矩阵，而是为每个头维护一个恒定大小的隐藏状态矩阵（匹配 dkd**v 键-值维度），作为循环记忆。

在每个 token 步骤 t，使用 delta 规则更新状态矩阵：

Delta 规则公式

其中 q**t、k**t 和 v**t 是查询、键和值向量，而 τ 是学习到的门控参数。这种循环更新之前会有一个因果一维卷积（K=4）来捕获局部空间依赖关系。这种循环公式允许上下文窗口在内存中以线性方式扩展（O(S)），使循环状态占用保持恒定。

2. 分组查询注意力（GQA）

为了锚定线性注意力检索，该模型在 25% 的层中使用标准 GQA。GQA 全局使用 32 个查询头（N**q=32）和恰好 2 个键-值（KV）头（N**kv=2），头维度为 256，旋转位置嵌入（RoPE）维度为 64。这种极端的 GQA 布局压缩了生成过程中的 KV 缓存占用，但施加了严格的硬件级分片约束，详见第 3 节。

3. 稀疏混合专家（MoE）

前馈网络（FFN）层被分片为 512 个小专家，中间专家维度为 1024。在执行过程中，路由器门控将 token 表示投影，并通过 softmax 概率分布选择前 10 个路由专家。关键的是，该模型还包含一个始终执行的共享专家路径，作为公共表示层：

表示层公式

这种原生多模态 MoE 架构通过早期融合训练范式，在数万亿多模态 token 上原生处理文本、图像和视频输入。上下文窗口支持原生 262,144 个 token 的上下文长度，可使用 YaRN RoPE 缩放扩展到超过 1,010,000 个 token。

2. 基准测试设置与工作负载配置

为了系统地隔离、分析和解决编译器与内核瓶颈，系统工程团队基于真实世界的不对称工作负载，建立了一个严格的多维度评估矩阵。

多维度评估矩阵

我们的基准测试扫描了旨在给不同硬件执行子系统施压的不对称工作负载：

1. 预填充密集型工作负载（8K 输入 / 1K 输出）： 特点是长输入提示序列和短 token 生成输出。这些工作负载是计算密集型的，严重依赖 TPU TensorCore 矩阵执行单元（MXU）的物理浮点矩阵乘法能力。
2. 解码密集型工作负载（1K 输入 / 8K 输出）： 特点是短输入提示和长时间的逐 token 生成阶段。这些工作负载是内存密集型的，因为系统必须持续将所有 400 GB 参数从高带宽内存（HBM）流式传输到执行核心，以生成每个请求的单个 token。
3. 并发层级： 为了观察系统扩展曲线并识别负载下的硬件排队/内存瓶颈，两种工作负载均在四个并发层级上进行评估：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 相比）。
• 指标跟踪：性能跟踪并报告为每芯片 token 吞吐量（TPS/chip），计算方法为处理的 token 总数（输入 + 输出）除以执行持续时间和物理芯片数量（4）。

3. 分片策略与分布式集合

Qwen 3.5 特定的架构约束——即 GQA 层中恰好有 2 个 KV 头，MoE 层中有 512 个专家——使得传统的统一分片方法失效。

并行性权衡：张量并行 vs. 数据并行

在标准的注意力张量并行（TP）+ 专家 MoE 配置中，注意力权重被切片并沿设备维度分片。然而，尝试使用大小为 8 的张量并行（TP=8）对 GQA 层进行分片，会导致分数头分片（2/8 = 每个设备 0.25 个头），这在硬件上是物理上不可能的。

在 8 个核心上本地复制头，会在每个设备上重复物理 KV 缓存内存占用，从而抵消 GQA 节省内存的好处。这种内存冗余严重限制了高负载工作负载下可用于活动 KV 缓存的 HBM 空间。这种容量限制迫使服务引擎将实际达到的并发数远低于预期目标——将系统限制在大约 ~200 个并发请求，而不是计划的 512 个。

为了消除这一瓶颈，我们共同设计了一种混合分片方案（PR #2577 (https://github.com/vllm-project/tpu-inference/pull/2577)）：8 路注意力批次分片（数据并行，DP=8）结合 MoE 层中的 8 路专家并行（EP=8）。

注意力层

在所有 8 个设备上复制 GQA 和 GDN 权重，允许每个核心在本地处理注意力，拥有完整的 2 个 KV 头，从而保持本地 KV 缓存一致性，并消除了注意力内部的分片通信。在前馈 MoE 层中，我们切换到专家并行（EP=8）。512 个路由专家均匀分布（每个设备 64 个专家），这避免了在所有节点上重复 400 GB 参数占用，同时保持集合负载大小可控。

深入分析分布式集合序列

在注意力 DP 和 MoE EP 之间切换需要跨设备 token 路由。在设计我们的混合专家（MoE）路由层时，我们评估了两种主要的处理这种跨设备转换的结构性方法：

• 选项 A（All-to-All 洗牌）： 该方法采用 All-to-All -> 本地 MoE -> All-to-All 流水线。token 动态地在网络中洗牌，到达托管其目标专家的特定芯片，在本地计算，然后洗牌返回。虽然这最小化了冗余计算，但由于全局 All-to-All 步骤在可变工作负载下引入了巨大且不可预测的网络路由开销。
• 选项 B（完整 token 复制）： 该方法采用 All-Gather -> 本地 MoE -> Reduce-Scatter 流水线。All-Gather 将所有 token 向量复制到所有设备上。然后每个芯片仅过滤并计算其本地专家的输入，稍后通过 Reduce-Scatter 聚合输出。这完全避免了不可预测的 All-to-All 路由惩罚，但代价是更高的本地内存消耗。

由于确定性延迟对于现实世界的服务至关重要，我们选择了选项 B，随后开发了底层通信融合来优化其集合路径。

1. 3 到 2 的 All-Gather 优化

在朴素的选项 B 实现下，为本地 MoE 计算做准备需要将三个独立的数据片段广播到集群中的每个设备等级。假设本地张量切片的形状为 [1024,4096]（用于 token 隐藏维度），我们通常需要执行三个独立的集合操作：

1. All-Gather 1： token 隐藏维度（[1024,4096]）。
2. All-Gather 2： 选中的专家索引（[1024,10]，假设 topk=10）。
3. All-Gather 3： 门控 topk 权重（[1024,10]）。

每次集合通信调用都会在 TPU 上带来固定的内核启动和网络同步延迟开销。为了优化专家并行（EP）效率，我们将这三个 All-Gather 合并为两个，详见 PR #2836 (https://github.com/vllm-project/tpu-inference/pull/2836)。由于专家索引（整数）和 topk 权重（浮点数）共享相同的张量形状（[1024,10]），我们将它们沿新维度堆叠、位转换并打包到一个密集的 32 位整数数组（blob）中。这允许我们在数据维度（ShardingAxisName.MLP_DATA）上对两个路由元数据块运行单个 All-Gather，在本地解包，从而将路由元数据集合延迟减半。

2. 分层 Reduce-Scatter

专家执行后，token 输出必须返回到它们的数据并行等级。在 8 设备网格上进行标准 All-Reduce 效率极低。我们将其替换为用 Pallas/Mosaic 编写的自定义、TPU 原生分层 Reduce-Scatter（参见 PR #2679 (https://github.com/vllm-project/tpu-inference/pull/2679)）。该集合在两个流水线阶段中运行：

• 芯片内 Reduce-Scatter： 同一物理芯片上的逻辑芯粒使用快速本地共享内存传输（比芯片到芯片 ICI 带宽快 6 倍）交换并求和它们的数据。
• 芯片间 Reduce-Scatter： 部分减少的数据通过 TPU 物理 ICI 链路上的递归加倍超立方体算法在物理芯片间交换。

为防止 VMEM 内存不足（OOM）错误，数据被切分为 2 到 4 个微批次。内核将微批次 i 的远程 DMA 传输与微批次 i-1 的 TensorCore 向量加法进行流水线处理，从而将通信延迟隐藏在计算之后。

4. 预填充与解码 Roofline 分析

为了确定系统工程的极限并了解执行停滞发生的位置，我们对 Qwen 3.5 工作负载在标准 8K/1K 配置和 64 并发下进行了第一性原理的 roofline 分析。

Ironwood 硬件规格

• Tensor Core（TC）频率：2.2 GHz
• 每芯片 Tensor Core 数量：2
• 每个 TC 的 MXU（矩阵执行单元）数量：2（每芯片共 4 个 MXU）
• 峰值 BF16 性能：2,307 TFLOPS/芯片（(262,144 * 每周期每 MXU 的 FLOP × 2.2 GHz × 4 MXUs = 2,307 TFLOPS)）
• 峰值 FP8 性能：4,614 TFLOPS/芯片

理论界限

1. 预填充阶段（计算密集型）

在预填充阶段，一批 64 个提示，每个 8,192 个输入 token，共并行处理 524,288 个 token。

• 算术强度： 投影层中的 GEMM 操