当Google工程师团队宣布在Ironwood（TPU v7）平台上将Qwen 3.5-397B MoE模型的推理性能提升3.1-4.7倍时，很多人可能只看到了最终的数字。但真正值得关注的是，这背后揭示了一个关键转折点：大模型优化正在从"针对特定模型的修修补补"转向"系统级的工程方法论"。

如果你正在部署百亿参数级别的MoE模型，可能会遇到这样的困境：每次模型架构更新，都需要重新投入数月时间进行性能优化；或者在高并发场景下，硬件利用率始终无法突破某个瓶颈。Google团队这次的技术突破，实际上提供了一套可复用的优化框架，而不仅仅是针对Qwen 3.5的特化方案。

本文将深入解析这套系统级优化方案的技术细节，从架构分析到内核优化，从分布式策略到性能验证，为你展示如何在现代AI加速器上高效部署超大规模MoE模型。

1. Qwen 3.5-397B MoE架构深度解析

1.1 模型参数规模与稀疏激活机制

Qwen 3.5-397B MoE是一个总参数量达到3970亿的巨型模型，但其核心创新在于高效的稀疏激活机制。每个token在前向传播过程中仅激活170亿参数，激活比例仅为4.3%。这意味着模型在保持400B级别智能能力的同时，实际推理计算量仅相当于一个20B规模的稠密模型。

这种设计对硬件部署提出了双重挑战：一方面需要管理400GB的权重存储，另一方面要高效处理高度稀疏的计算图。模型权重可通过Hugging Face仓库直接获取，为后续优化提供了基础。

1.2 混合注意力架构设计

Qwen 3.5采用了独特的混合层结构，整个网络包含60层，划分为15个重复的结构块，每个块遵循3:1的比例：

• Gated DeltaNet层（75%）：3个连续层，结合线性注意力与稀疏MoE路由
• Grouped Query Attention层（25%）：1个标准GQA层，同样配备MoE路由

这种混合设计不是随意的堆叠，而是针对不同计算特性的精心安排。Gated DeltaNet（GDN）解决了传统注意力机制在长序列下的O(S²)复杂度问题，而GQA层则确保在关键位置保持精确的注意力机制。

1.3 核心数学原理剖析

Gated DeltaNet线性注意力：使用64个值头（V）和16个查询键头（QK），头维度为128。与传统softmax注意力不同，GDN维护一个恒定大小的隐藏状态矩阵作为循环记忆体。在每个token步骤t，状态矩阵通过delta规则更新：

state_t = gate_t * state_{t-1} + (1 - gate_t) * (k_t^T v_t)

其中gate_t是学习得到的门控参数。这种循环公式使上下文窗口的内存占用保持恒定，实现线性复杂度。

分组查询注意力：使用32个查询头，但仅有2个KV头，头维度256，配合RoPE位置编码。这种极端的GQA设计大幅压缩了生成过程中的KV缓存，但对硬件分片提出了严格要求。

MoE专家路由：FFN层被分片为512个专家，中间维度1024。路由器为每个token选择top-10专家，同时包含一个始终执行的共享专家路径，确保基础表示的一致性。

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

2.1 硬件平台规格

优化工作在Ironwood TPU v7x平台上进行，单台物理主机包含4个物理芯片。每个物理芯片由2个逻辑小芯片组成，形成8个逻辑执行核心的拓扑结构，通过高速芯片间互联（ICI）平面连接。

关键硬件参数：

• Tensor Core频率：2.2 GHz
• 每芯片Tensor Core数：2
• 每Tensor Core的MXU数：2（总计每芯片4个MXU）
• 峰值BF16性能：2,307 TFLOPS/芯片
• 峰值FP8性能：4,614 TFLOPS/芯片
• HBM容量：192GB/芯片

2.2 工作负载设计哲学

为了系统性地识别瓶颈，团队设计了不对称的工作负载组合：

预填充密集型（Prefill-Heavy）：8K输入token，1K输出token。这种工作负载是计算受限的，主要压力在TPU的TensorCore矩阵执行单元（MXU）上。

解码密集型（Decode-Heavy）：1K输入token，8K输出token。这种工作负载是内存受限的，系统需要持续从HBM流式传输400GB参数来生成单个token。

并发层级测试覆盖了64、128、256和512并发请求，以观察系统扩展曲线和识别硬件排队/内存瓶颈。

2.3 推理服务引擎配置

使用vllm-project/tpu-inference作为推理服务器引擎。优化后的配置参数：

--max-num-batched-tokens=1024 --max-num-seqs=64 per core

相比早期的张量并行基线配置（--max-num-batched-tokens=8192，--max-num-seqs=512），这些参数反映了对硬件特性更深入的理解。

3. 分片策略与分布式集体操作

3.1 传统分片方案的局限性

Qwen 3.5的特定架构约束（GQA层仅有2个KV头，MoE层有512个专家）使得传统的均匀分片方法失效。尝试使用张量并行度8（TP=8）对GQA层分片会导致分数头分片（2/8 = 0.25头/设备），这在硬件上无法实现。

如果跨8个核心本地复制头，会在每个设备上复制物理KV缓存内存占用，抵消GQA的内存节省优势。这种内存冗余严重限制了高负载工作负载下可用于活跃KV缓存的HBM空间。

3.2 混合分片方案设计

团队共同设计了混合分片方案（PR #2577）：

• 8路注意力批处理分片（数据并行，DP=8）
• MoE层中的8路专家并行（EP=8）

在注意力层，跨所有8个设备复制GQA和GDN权重，使每个核心能够本地处理完整的2个KV头，保持本地KV缓存一致性。在MoE层切换到专家并行，将512个路由专家均匀分布（每个设备64个专家），避免跨所有节点复制400GB参数占用。

3.3 分布式集体操作优化

在注意力DP和MoE EP之间转换需要跨设备token路由。团队评估了两种主要结构方法：

方案A（All-to-All混洗）：All-to-All → 本地MoE → All-to-All管道。虽然最小化冗余计算，但在可变工作负载下产生巨大的不可预测网络路由开销。

方案B（完整token复制）：All-Gather → 本地MoE → Reduce-Scatter管道。完全绕过不可预测的All-to-All路由惩罚，代价是更高的本地内存消耗。

基于确定性延迟的考虑，选择了方案B并进行了深度优化。

3.3.1 3到2 All-Gather优化

原始实现需要三个独立的集体操作：

1. token隐藏维度（[1024,4096]）
2. 选定的专家索引（[1024,10]）
3. 门控top-k权重（[1024,10]）

通过将专家索引（整数）和top-k权重（浮点数）堆叠、位转换并打包到单个密集32位整数数组中，将三个All-Gather合并为两个，将路由元数据集体延迟减半。

3.3.2 分层Reduce-Scatter

使用自定义的Pallas/Mosaic分层Reduce-Scatter内核（PR #2679），分两个流水线阶段运行：

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

通过将数据切片为2-4个微批次，实现远程DMA传输与TensorCore计算的流水线并行，隐藏通信延迟。

4. 预填充与解码阶段的屋顶线分析

4.1 计算受限的预填充阶段

在预填充阶段，处理64个提示的批次，每个包含8,192个输入token，总共并行处理524,288个token。投影层中的GEMM操作随序列长度和批次大小呈二次缩放，算术强度极高，使执行处于屋顶线模型的计算受限区域。

系统瓶颈主要源于跨专家的不规则token分布。如果一个专家在给定批次中接收的token明显多于其他专家，相应设备就会成为落后节点。通过Grouped GEMM内核中的最小化填充来缩小实际TFLOPS与理论峰值之间的差距。

4.2 内存受限的解码阶段

在解码阶段，模型每步处理64个token（每个活跃请求1个token）。生成一个token需要从HBM流式传输所有400GB模型权重，算术强度接近单位1（~1FLOP/Byte），使工作负载处于内存受限区域。

主要延迟贡献者包括模型参数的HBM传输延迟、稀疏KV缓存检索期间的VPU索引停顿，以及Gated DeltaNet层中的循环状态更新往返。

4.3 屋顶线边界量化

基于64并发的工作负载模型分析：

预填充阶段屋顶线：计算受限，受限于TensorCore MXU的峰值FP8性能（4,614 TFLOPS/芯片）。考虑8,192个token上的GQA注意力操作二次缩放和标准硬件执行开销，估计最大理论屋顶线吞吐量为5,170 tokens/s/芯片（未折扣），标准调度降额因子下为4,500 tokens/s/芯片。

解码阶段屋顶线：内存受限，严格受HBM接口带宽限制。所有60层的总执行延迟计算为每token步骤16.36毫秒，峰值理论吞吐量为978 tokens/s/芯片（未折扣），实际服务屋顶线限制为850 tokens/s/芯片。

5. 内核级优化技术深度剖析

5.1 注意力轨道：Ragged Page Attention (RPA)

管理25% GQA层的KV缓存需要动态内存分配。使用Ragged Page Attention索引HBM中的非连续内存块。

KV页面大小调优：历史上使用16个token的块大小来最小化内存碎片。但在TPU上，较小的块大小导致大量索引开销。通过粗粒度索引将KV页面大小调整为256（通过--block-size=256启用），在并发512下将解码步骤延迟从428µs减少到283µs，实现33.8%的内核级加速。

批处理RPA：设计批处理RPA内核，将多个解码流分组到单个编译的Pallas内核中，分摊VPU指令分发延迟，打破顺序请求的数据依赖停顿，改善内存对齐。

5.2 MoE轨道：SparseCore与TensorCore协同设计

Qwen 3.5中top_k=10的细粒度路由因子引入非2的幂次张量维度。通过SparseCore-TensorCore协同设计流解决这个问题：

自定义SparseCore Ragged Gather内核：编写自定义Pallas/Mosaic内核，将token路由卸载到TPU的SparseCore（SC）——一个针对间接寻址优化的硬件单元。SC读取路由索引，直接从HBM执行token嵌入的间接DMA收集，并将其写入连续虚拟缓冲区，绕过HBM中重度填充的未对齐中间张量的物化。

Grouped GEMM V2与融合激活：在GMM V2内核中，将SwiGLU激活函数直接融合到主矩阵乘法循环中（门控和上投影通过双DMA读取在单个tile中打包和处理）。实现动态有界切片以最小填充处理每个专家的可变token负载。转换为FP8操作的512子通道激活量化，消除VREG溢出和内存加载停顿。

融合Ragged Gather Reduce内核：将token反排列和本地归约操作完全卸载到SparseCore。通过在SC上执行间接收集和本地归约，绕过HBM中中间激活张量的物化，将HBM读取需求从20减少到10，写入从15减少到5。

5.3 GDN轨道：Gated DeltaNet优化

75% Gated DeltaNet层中的循环状态更新由于恒定循环状态更新而极易受内存带宽瓶颈影响。

因果Conv1D融合：GDN循环更新之前是因果1D卷积（K=4）。设计寄存器级滑动窗口算法，将历史token状态直接缓存在TPU的VPU寄存器中。将1D卷积和GDN循环状态更新融合到单个执行块中，消除6个冗余的HBM往返。

代数恒等式优化：重新构建线性注意力更新方程以利用代数恒等式。通过数学重排操作，完全跳过融合GDN内核中昂贵的后秩1矩阵乘法，减少计算占用。

不规则序列处理与分块GDN：优化JAX原生分块布局并引入专门序列处理例程，原生处理不规则输入，确保可变序列长度不会引入处理落后节点。

完全融合的Conv1D和GDN内核：设计完全融合的Pallas内核，将因果1D卷积和整个GDN循环线性注意力块编译为VPU上的统一执行单元。通过将中间序列和循环状态直接缓存在本地寄存器中，完全绕过向VMEM或HBM读写中间激活张量的需求。

5.4 内存轨道：混合注意力KV布局优化

服务Qwen 3.5需要管理两种异构注意力状态结构：Gated DeltaNet的固定大小循环线性注意力状态和Grouped Query Attention的动态增长标准注意力KV缓存。由于TPU v7每个芯片具有192GB HBM容量（相比Blackwell GB300 GPU的288GB，容量差异约50%），高并发下的HBM占用优化是严重系统约束。

引入自定义内存布局，在HBM中对齐和存储这些混合注意力状态在一起。这种布局最小化填充并防止内存碎片，直接回收关键的HBM空间，增加最大可支持批次大小，使TPU能够平稳扩展并在重客户端并发下维持高服务吞吐量。

6. 性能结果与数值验证

6.1 实际吞吐量与屋顶线限制对比

在基准并发64层级上，将经验服务结果与第一原理屋顶线限制并排比较：

预填充密集型效率：在8K/1K预填充密集型工作负载下，JAX服务栈提供3,707 tokens/s/芯片的实际吞吐量。相比估计的预填充屋顶线限制4,500 tokens/s/芯片（折扣后），自定义SparseCore和TensorCore协同设计的GEMM成功提取了TPU v7 TensorCore绝对计算能力的82.4%。

解码密集型效率：在1K/8K解码密集型工作负载下，服务栈提供677 tokens/s/芯片的实际吞吐量。相比内存受限的解码屋顶线限制850 tokens/s/芯片（折扣后），Ragged Page Attention和Gated DeltaNet融合成功实现了理论HBM带宽限制的79.6%。

这种接近的对齐表明，低层级编译器和内核融合将TPU硬件推近其物理执行限制，留下最小的剩余空间，证明了开源软件栈的极高效率。

6.2 严格数值验证与正确性

在高并发下运行大规模MoE模型不仅需要原始吞吐量，还需要严格的数学正确性。门控和路由矩阵对低精度累积误差高度敏感。在设计自定义JAX/Pallas门控内核时，系统工程师团队集成了专用数值验证层，审计跨FP8缩放块的累积精度。

通过持续监控softmax分布范围和专家负载平衡，验证Pallas降级的门控权重保持与高精度Float32参考路径的零偏差，保证高吞吐量同时确保严格的输出质量。

7. 未来优化路线图

7.1 集体操作优化轨道

低带宽FP8 All-Gather集体操作：为Token/元数据All-Gather步骤设计低带宽FP8集体操作。在跨节点传输之前将路由元量化为FP8，将通过物理ICI链路的通信量减半，直接减少路由延迟障碍。

分层Reduce-Scatter调优：继续在自定义分层Reduce-Scatter内核内优化块大小和微批次流水线参数。特别目标是实现动态的、token相关的微批次大小调整，以优化可变路由分布下的带宽利用率。

7.2 内核与门控融合轨道

路由器门控和Top-K融合：计划将路由门控投影和后续top_k选择内核直接在VPU上融合。当前，路由logits在TensorCore上计算并传输到VPU进行top_k选择，引入串行化瓶颈。融合这些操作将使路由管道保持在VPU本地。

8. 工程实践与部署建议

8.1 环境配置最佳实践

对于希望在类似硬件上部署MoE模型的团队，建议遵循以下配置原则：

内存布局优化：根据工作负载特性动态调整KV缓存布局。对于长序列生成任务，优先考虑内存效率；对于短序列高并发任务，优化计算密度。

# 示例：动态内存布局配置 memory_config = { "kv_cache_layout": "hybrid_blocked", "block_size": 256, # 根据硬件特性调整 "max_sequence_length": 262144, "hbm_optimization": True }

8.2 性能监控与调优指标

建立全面的性能监控体系，关键指标包括：

• 令牌吞吐量/芯片：核心性能指标
• HBM带宽利用率：识别内存瓶颈
• MXU利用率：评估计算效率
• 路由平衡度：监控专家负载分布
• 集体操作延迟：优化通信开销

8.3 常见问题与解决方案

9. 技术影响与行业启示

Google这次在Ironwood TPU上优化Qwen 3.5-397B MoE的实践，标志着大模型推理优化进入了系统级工程化的新阶段。其核心价值不在于针对特定模型的性能提升，而在于建立了一套可复用的优化方法论：

模块化优化策略：将复杂模型分解为独立构建块，每个块配备硬件感知成本模型。当新架构出现时，这些预优化模块可以近乎零工程摩擦地移植。

硬件软件协同设计：深度理解TPU架构特性（如SparseCore的间接寻址优势），在算法层面进行针对性优化，而非简单适配。

开源软件栈成熟：将优化集成到vLLM和SGLang等开源服务框架中，为全球企业工作负载提供生产就绪的迁移路径。

对于广大AI工程师和研究人员而言，这项工作的启示在于：面对日益复杂的大模型架构，需要从系统级视角出发，将硬件特性、算法设计和工程实现深度融合，才能充分发挥现代AI加速器的潜力。

随着MoE架构成为大规模模型的主流选择，这套优化方法论将为整个行业提供宝贵的技术积累和实践参考。