这次我们来看一个重量级的技术突破：Google 工程师团队在最新的 Ironwood（TPU 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.7x（512并发） |
| 关键技术 | 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路张量并行（TP=8），GQA层会出现分数头分片（2/8=0.25头/设备），这在硬件上无法实现。
2.2 内存带宽瓶颈
TPU v7 每个芯片具有192GB HBM容量，相比Blackwell GB300 GPU的288GB有约50%的容量差距。在高并发场景下，HBM占用优化成为关键系统约束。模型需要同时管理两种不同的注意力状态结构：Gated DeltaNet的固定大小循环线性注意力状态和GQA的动态增长KV缓存。
3. 混合分片策略突破
为了解决架构约束带来的分片难题，工程团队设计了创新的混合分片方案（PR #2577）：8路注意力批处理分片（数据并行，DP=8）与8路专家并行（EP=8）结合。
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-size=256启用），在并发512下的解码步骤延迟从428µs降低到283µs，实现33.8%的内核级加速。
批处理RPA：为了进一步饱和内存总线，设计了批处理RPA内核，将多个解码流分组到单个编译的Pallas内核中，分摊VPU指令分发延迟，打破顺序请求的数据依赖停滞。
4.2 MoE轨迹：稀疏核心与张量核心协同设计
Qwen 3.5中top_k=10的细粒度路由因子引入了非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卷积（K=4）。团队设计了寄存器级滑动窗口算法，在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-tokens=1024 --max-num-seqs=64 --block-size=2569.2 性能监控指标
在实际部署中，需要重点监控的关键指标包括：
- 令牌吞吐量每芯片：总处理令牌数除以执行时间和物理芯片数
- HBM利用率：确保内存带宽得到有效利用
- 专家负载平衡：监控各专家节点的令牌分布均匀性
- 集体通信延迟：关注All-Gather和Reduce-Scatter操作耗时
9.3 容错与扩展性
由于模型规模巨大，需要实现健壮的故障恢复机制：
- 检查点保存：定期保存模型状态和路由信息
- 优雅降级：在部分专家节点故障时保持基本服务能力
- 动态扩展：根据负载动态调整并发数和批处理大小
10. 技术影响与行业意义
这项优化工作的重要意义不仅在于提升了Qwen 3.5在TPU上的性能，更重要的是建立了一套可重用的系统工程方法论。通过将模型分解为自包含的独立构建块，并配备硬件感知成本模型，工程团队能够以接近零的工程摩擦移植预优化模块。
这种模块化方法彻底改变了大规模模型优化的游戏规则。传统上，每个新模型家族都需要数月独立优化，而现在团队可以专注于模型的新颖组件，快速交付最先进的服务性能。
优化策略的成功证明，面对日益复杂的开放权重模型景观，基于第一原理的系统工程方法比经验性的试错修改更加有效。通过屋顶线建模映射理论硬件限制，通过性能分析跟踪系统性地隔离瓶颈，并通过手动调度的自定义Pallas内核和硬件感知分片拓扑桥接剩余差距，这套方法论为下一代稀疏混合专家架构建立了坚固、可重用的开源软件栈。
对于企业和研究机构来说，这项工作的最大价值在于提供了生产就绪的迁移路径。通过将优化原生集成到vLLM和SGLang等开源服务框架中，团队消除了传统软件障碍，使全球企业工作负载能够无缝扩展到前所未有的规模。
这项技术突破标志着大规模AI模型服务进入了新的时代，硬件和软件的协同优化将成为释放万亿参数模型潜力的关键所在。
