Google TPU v7x深度优化Qwen 3.5-397B MoE模型，推理性能提升3-4倍

这次我们来看一个重量级的技术突破Google 工程师团队在最新的 IronwoodTPU v7x平台上对 Qwen 3.5-397B MoE 模型进行了深度优化。这个项目不是简单的模型微调而是从系统工程的底层出发通过一系列创新的并行策略和内核优化让这个拥有 3970 亿参数的巨型模型在专用硬件上实现了接近理论极限的推理性能。从实际效果来看优化后的系统在 512 并发层级下解码密集型工作负载性能提升了约 3.1 倍预填充密集型工作负载性能提升了约 4.7 倍。更重要的是这些优化已经被集成到 vLLM 和 SGLang 等开源推理框架中为企业级大规模工作负载提供了生产就绪的迁移路径。1. 核心能力速览能力项技术规格模型规模3970 亿总参数激活参数 170 亿/令牌4.3%路由激活率硬件平台Ironwood TPU v7x单主机4物理芯片8逻辑核心优化重点注意力数据并行 专家并行混合分片策略性能提升解码密集型 3.1x预填充密集型 4.7x512并发关键技术JAX/Pallas 自定义内核稀疏核心协同设计集成框架vLLM、SGLang 开源推理框架适用场景企业级大规模模型服务长上下文多模态推理2. Qwen 3.5 架构特点与优化挑战Qwen 3.5-397B 采用混合专家MoE架构虽然总参数量达到 3970 亿但通过稀疏路由机制每个令牌前向传播仅激活 170 亿参数。这种设计让模型在保持 400B 级别智能表现的同时推理计算量仅相当于 20B 级别的稠密模型。2.1 混合注意力架构模型的网络结构包含 60 层隐藏维度为 4096采用独特的 3:1 混合布局门控DeltaNet层75%3个连续层结合门控DeltaNet线性注意力与路由稀疏MoE分组查询注意力层25%1个标准GQA层与路由稀疏MoE结合这种混合设计带来了显著的优化挑战。GQA层只有2个KV头而MoE层有512个专家传统的均匀分片策略完全失效。如果采用8路张量并行TP8GQA层会出现分数头分片2/80.25头/设备这在硬件上无法实现。2.2 内存带宽瓶颈TPU v7 每个芯片具有192GB HBM容量相比Blackwell GB300 GPU的288GB有约50%的容量差距。在高并发场景下HBM占用优化成为关键系统约束。模型需要同时管理两种不同的注意力状态结构Gated DeltaNet的固定大小循环线性注意力状态和GQA的动态增长KV缓存。3. 混合分片策略突破为了解决架构约束带来的分片难题工程团队设计了创新的混合分片方案PR #25778路注意力批处理分片数据并行DP8与8路专家并行EP8结合。3.1 注意力数据并行通过在8个设备上复制GQA和GDN权重每个核心都能本地处理完整的2个KV头注意力保持本地KV缓存一致性消除注意力内分片通信。这种设计完全保留了GQA的内存节省优势避免了KV缓存在所有节点上的冗余复制。3.2 专家并行优化在MoE前馈层切换到专家并行将512个路由专家均匀分布每个设备64个专家避免了400GB参数在所有节点上的重复同时保持集体负载大小可控。3.3 分布式集体通信优化在注意力DP和MoE EP之间过渡需要跨设备令牌路由。团队评估了两种主要结构方法方案A全到全洗牌使用All-to-All - 本地MoE - All-to-All流水线令牌动态洗牌到目标专家所在芯片。虽然计算冗余最小但全局All-to-All步骤在可变工作负载下产生巨大网络路由开销。方案B全令牌复制使用All-Gather - 本地MoE - Reduce-Scatter流水线。All-Gather在所有设备上复制所有令牌向量每个芯片过滤并计算本地专家输入最后通过Reduce-Scatter聚合输出。这完全避免了不可预测的All-to-All路由惩罚代价是更高的本地内存消耗。由于确定性延迟对实际服务至关重要团队选择方案B并开发了底层通信融合优化。4. 核心内核优化技术4.1 注意力轨迹粗糙页面注意力管理25% GQA层的KV缓存需要动态内存分配。团队采用粗糙页面注意力来索引HBM中的非连续内存块。KV页面大小调优历史上使用16令牌的块大小来最小化内存碎片。但在TPU上较小的块大小会导致巨大的索引开销在解码阶段导致向量处理单元停滞。通过将索引粗粒化为KV页面大小256通过服务器命令--block-size256启用在并发512下的解码步骤延迟从428µs降低到283µs实现33.8%的内核级加速。批处理RPA为了进一步饱和内存总线设计了批处理RPA内核将多个解码流分组到单个编译的Pallas内核中分摊VPU指令分发延迟打破顺序请求的数据依赖停滞。4.2 MoE轨迹稀疏核心与张量核心协同设计Qwen 3.5中top_k10的细粒度路由因子引入了非2的幂次张量维度。在张量核心上置换和反置换这些数组之前会导致重度填充、未对齐的HBM内存写入。自定义稀疏核心粗糙收集内核编写自定义Pallas/Mosaic内核将令牌路由卸载到TPU的稀疏核心SC——一个为间接寻址优化的硬件单元。SC读取路由索引直接从HBM执行间接DMA收集令牌嵌入写入连续虚拟缓冲区绕过重度填充中间张量在HBM中的物化。分组GEMM V2与融合激活在GMM V2内核中将SwiGLU激活函数直接融合到主矩阵乘法循环中避免寄存器溢出到HBM。实现动态有界切片以最小填充处理每个专家的可变令牌负载。4.3 GDN轨迹门控DeltaNet优化75%门控DeltaNet层中的循环状态更新由于恒定循环状态更新而极易受内存带宽瓶颈影响。因果Conv1D融合GDN循环更新之前有一个因果1D卷积K4。团队设计了寄存器级滑动窗口算法在TPU的VPU寄存器内直接缓存历史令牌状态将1D卷积和GDN循环状态更新融合到单个执行块消除6个冗余HBM往返。代数恒等式优化重新构建线性注意力更新方程以利用代数恒等式通过数学重排操作完全跳过融合GDN内核中昂贵的后秩1矩阵乘法。5. 性能基准测试与结果验证5.1 测试工作负载配置团队建立了基于真实世界非对称工作负载的严格多维评估矩阵预填充密集型工作负载8K输入/1K输出以长输入提示序列和短令牌生成为特征这些工作负载是计算受限的重度压力TPU张量核心矩阵执行单元的物理浮点矩阵乘法容量。解码密集型工作负载1K输入/8K输出以短输入提示和 prolonged 令牌逐令牌生成阶段为特征这些工作负载是内存受限的系统必须持续从HBM流式传输所有400GB参数到执行核心。5.2 实际性能成果在严格的实证基准测试场景下验证系统工程优化在4个并发层级上呈现原始吞吐量结果在基准并发64层级下将实证服务结果与第一原理屋顶线限制并排比较预填充密集型效率在8K/1K预填充密集型工作负载下JAX服务栈提供实际吞吐量3,707令牌/秒/芯片。相比估计的预填充屋顶线限制4,500令牌/秒/芯片折扣后自定义稀疏核心和张量核心协同设计的GEMM成功提取了TPU v7张量核心绝对计算容量的82.4%。解码密集型效率在1K/8K解码密集型工作负载下栈提供实际吞吐量677令牌/秒/芯片。相比内存受限的解码屋顶线限制850令牌/秒/芯片折扣后粗糙页面注意力和门控DeltaNet融合成功实现了理论HBM带宽限制的79.6%。这种紧密对齐表明底层编译器和内核融合将TPU硬件推近其物理执行极限留下最小的剩余空间证明了开源软件栈的极端效率。6. 数值验证与正确性保证大规模运行混合专家模型不仅需要原始吞吐量还需要严格的数学正确性。在高并发下门控和路由矩阵对低精度累积误差高度敏感。在设计和实现自定义JAX/Pallas门控内核时系统工程团队加入了专用数值验证层来审计跨FP8缩放块的累积精度。通过持续监控softmax分布范围和专家负载平衡验证了Pallas降低的门控权重保持与高精度Float32参考路径的零偏差保证高吞吐量同时严格输出质量。7. 资源占用与性能观察7.1 内存布局优化服务Qwen 3.5需要管理两种异构注意力状态结构Gated DeltaNet的固定大小循环线性注意力状态和Grouped Query Attention的动态增长Key-Value缓存。在PR #2416中团队引入了自定义内存布局在HBM中对齐和存储这些混合注意力状态。这种布局最小化填充并防止内存碎片直接回收关键的HBM空间。7.2 分块流水线架构为了最大化硬件效率实现采用块级流水线架构而不是在完整[81920,4096]张量上顺序执行本地减少和8设备Reduce-Scatter。工作负载被划分为4个不同块。一旦块1在稀疏核心上完成其本地反置换和收集减少它就异步启动跨物理ICI链路的Reduce-Scatter集体。同时稀疏核心开始块2的本地收集减少。这种严格的块级流水线有效地重叠并隐藏了Reduce-Scatter的跨设备网络延迟在后面块的本地计算之后。8. 未来优化路线图为了消除剩余瓶颈工程团队构建了活跃的优化路线图分为两个主要技术轨道8.1 集体优化轨道低带宽FP8 All-Gather集体为令牌/元数据All-Gather步骤设计低带宽FP8集体。在跨节点传输之前将路由元量化为FP8将通信量减半直接降低路由延迟障碍。分层Reduce-Scatter调优继续在自定义分层Reduce-Scatter内核内优化块大小和微批流水线参数。特别地旨在实现动态的、令牌相关的微批大小以优化可变路由分布下的带宽利用率。8.2 内核与门控融合轨道路由器门控和Top-K融合计划将路由门控投影和后续top_k选择内核直接在VPU上融合。当前路由logits在张量核心上计算并传输到VPU进行top_k选择引入序列化瓶颈。融合这些操作将保持路由管道本地到VPU。9. 部署实践与工程建议9.1 环境配置要点对于计划在类似硬件上部署Qwen 3.5的团队建议关注以下配置参数# vLLM TPU推理服务器配置示例 --max-num-batched-tokens1024 --max-num-seqs64 --block-size2569.2 性能监控指标在实际部署中需要重点监控的关键指标包括令牌吞吐量每芯片总处理令牌数除以执行时间和物理芯片数HBM利用率确保内存带宽得到有效利用专家负载平衡监控各专家节点的令牌分布均匀性集体通信延迟关注All-Gather和Reduce-Scatter操作耗时9.3 容错与扩展性由于模型规模巨大需要实现健壮的故障恢复机制检查点保存定期保存模型状态和路由信息优雅降级在部分专家节点故障时保持基本服务能力动态扩展根据负载动态调整并发数和批处理大小10. 技术影响与行业意义这项优化工作的重要意义不仅在于提升了Qwen 3.5在TPU上的性能更重要的是建立了一套可重用的系统工程方法论。通过将模型分解为自包含的独立构建块并配备硬件感知成本模型工程团队能够以接近零的工程摩擦移植预优化模块。这种模块化方法彻底改变了大规模模型优化的游戏规则。传统上每个新模型家族都需要数月独立优化而现在团队可以专注于模型的新颖组件快速交付最先进的服务性能。优化策略的成功证明面对日益复杂的开放权重模型景观基于第一原理的系统工程方法比经验性的试错修改更加有效。通过屋顶线建模映射理论硬件限制通过性能分析跟踪系统性地隔离瓶颈并通过手动调度的自定义Pallas内核和硬件感知分片拓扑桥接剩余差距这套方法论为下一代稀疏混合专家架构建立了坚固、可重用的开源软件栈。对于企业和研究机构来说这项工作的最大价值在于提供了生产就绪的迁移路径。通过将优化原生集成到vLLM和SGLang等开源服务框架中团队消除了传统软件障碍使全球企业工作负载能够无缝扩展到前所未有的规模。这项技术突破标志着大规模AI模型服务进入了新的时代硬件和软件的协同优化将成为释放万亿参数模型潜力的关键所在。