晶圆级GPU架构优化MoE模型推理加速技术解析

1. 晶圆级GPU架构与MoE模型推理加速概述

在人工智能计算领域，混合专家模型（Mixture of Experts, MoE）已成为扩展大语言模型容量的关键技术路径。MoE模型通过动态激活不同专家子网络，实现了模型参数规模的指数级增长，同时保持计算量的线性增长。然而，这种架构特性也带来了独特的数据移动挑战，特别是在多芯片系统上部署时，专家选择的不均衡性和跨芯片数据通信成为性能瓶颈。

晶圆级GPU（Wafer-Scale GPU）作为新一代计算架构，通过系统级封装技术将数十个计算芯片集成在单一晶圆上，提供了前所未有的计算密度和内存带宽。但这种架构也面临着跨芯片通信延迟、负载均衡和内存一致性等挑战。本文提出的优化方案正是针对这些痛点，通过硬件/软件协同设计方法，显著提升MoE模型在晶圆级GPU上的推理效率。

关键创新点：我们的方案通过重构命令处理器架构，实现了专家位置感知的任务分配和智能数据预取，在保持现有编程模型兼容性的同时，显著降低了跨芯片通信开销。

2. 系统架构设计与核心组件

2.1 晶圆级GPU基础架构

现代晶圆级GPU采用多芯片模块（Multi-Chiplet Module, MCM）设计，如图10(a)所示。典型配置包含：

• 计算芯片（GPU Die）：每个芯片包含流式多处理器（SM）、本地HBM内存和命令处理器
• 高带宽内存（HBM）：通过硅中介层（Interposer）与计算芯片连接
• 片上网络（NoC）：实现芯片间通信
• 全局命令处理器（Global CP）：协调整个晶圆的操作

与传统多GPU系统不同，晶圆级GPU的芯片间通过超短距离的硅中介层互连，提供高达1.7TB/s的带宽（如Tesla Dojo架构），但通信延迟和能耗仍显著高于芯片内访问。

2.2 改进的全局命令处理器架构

我们提出的两级命令处理器架构如图10(c)所示：

1. 全局命令处理器（Global CP）：

   • 维护专家分布表（Expert Distribution Table）：记录每个专家在各芯片的存储位置
   • 跨令牌热图（Cross-token Heatmap）：跟踪专家激活的历史模式
   • 任务分配算法：动态分配计算任务到各芯片

2. 本地命令处理器（Local CP）：

   • 接收Global CP分配的子任务
   • 管理本地预测表（Prediction Table）
   • 协调SM执行和本地HBM管理

这种分层设计实现了全局资源协调与本地执行的解耦，既保证了系统级的负载均衡，又维持了芯片级的执行效率。

3. 关键算法与数据流优化

3.1 专家位置感知的任务分配算法

算法1展示了我们的启发式任务分配机制，其核心思想是将MoE计算任务智能分配到存储相关专家的芯片上。算法执行流程如下：

1. 输入处理：

   • expert_reqs_dict：各专家的请求数量
   • expert_die_map：专家分布信息

2. 候选芯片生成：

   def GenCandidateList(expert_id, dis=1): local_die_list = expert_die_map[expert_id] # 存储专家权重的芯片 remote_die_list = FindNearDies(local_die_list, dis) # 相邻芯片 return local_die_list + remote_die_list

3. 分块分配策略：

   • 将专家请求分成50个token的块（平衡效率与准确性）
   • 使用成本模型选择最优芯片：
     • DRAM访问成本
     • 计算成本
     • 芯片间通信成本

4. 分配计划合并：

   • 合并分配到同一芯片的任务块
   • 生成最终分配计划

该算法通过限制候选芯片数量和采用分块分配策略，在多项式时间内获得近似最优解，避免了NP难问题的计算复杂度。

3.2 数据驱动的预测器设计

如图11(b)所示，预测器算法基于时间局部性原理工作：

1. 从热图中识别当前专家选择对应的行
2. 从每行选择top-n热门专家
3. 预测下一令牌可能使用的专家

预测结果被传输到各芯片的PDU，指导硬件管理的本地HBM缓存策略。例如，当某芯片当前计算专家1和4时，预测可能接下来需要专家2、4和6，则预先将专家4（已在本地）保留在本地DRAM中。

3.3 数据流优化机制

我们通过扩展D2D控制器实现两种数据访问路径：

非复制数据读取（绿色路径）：

1. SM发起远程数据读取请求
2. D2D控制器常规路由请求
3. PDU检查预测表决定是否复制
4. 数据返回SM，必要时写入本地HBM

本地化数据读取（蓝色路径）：

1. SM请求已缓存的远程数据
2. ATU将远程地址转换为本地地址
3. 从LLC和本地HBM获取数据
4. 直接返回给SM，避免跨芯片通信

4. 实现细节与硬件开销

4.1 关键数据结构实现

全局CP数据结构：

• 专家分布表：每个条目包含专家初始芯片ID和分布状态（n位二进制码）
• 跨令牌热图：记录专家激活随时间变化的模式

PDU预测表：

• cp_en位：指示专家是否应缓存在本地
• is_local位：跟踪专家是否已在本地HBM

4.2 硬件开销分析

表II总结了各模块的面积和功耗开销：

在25芯片的晶圆上，总面积和功耗开销仅为0.04%，实现了高性能与低开销的平衡。

5. 性能评估与结果分析

5.1 实验设置

仿真平台：

• 基于Python开发定制多芯片GPU模拟器
• 验证：与8×H100 DGX服务器的实测数据误差<5%
• 支持配置芯片数量、布局和连接性

硬件配置（表I）：

• Tesla Dojo：5×5 2D网格
• TSMC SoW：8×3 2D网格
• 每个芯片：1000 TFLOPS FP16算力，80GB HBM

基准模型：

• DeepSeek-V3、Kimi-K2、Llama-4、Qwen3
• 批量大小：4,096至16,384

5.2 吞吐量提升

图12显示了我们方案在不同配置下的吞吐量提升：

专家数量更多的模型（如DeepSeek和Kimi的256专家）受益更大，因其选择模式更复杂。TSMC的矩形布局因芯片间距更大，从我们的策略中获得更高收益。

5.3 通信跳数减少

图12底部展示了通信跳数减少比率：

结果表明，随着优化策略的应用，通信不再是主要瓶颈，工作负载均衡成为关键因素。

5.4 内存访问分析

图14展示了DRAM访问模式的变化：

• 基准方案：大部分为远程读取
• Pred Only：部分远程读取转为本地读取
• Allo+Pred：极少远程读取，主要处理极热门专家

这种转变显著降低了芯片间流量，提升了整体效率。

6. 实际部署考量

6.1 与主机CPU方案的对比

图15比较了GPU命令处理器与主机CPU实现的分配开销：

随着GPU性能提升，主机CPU实现的相对开销显著增加，凸显了我们的GPU集成方案的必要性。

6.2 专家放置策略案例研究

针对现有GPU系统，我们提出了基于预填充信息的专家放置算法（算法2）：

1. 重映射策略：

   • 保持每GPU专家数量不变
   • 根据负载重新分配专家

2. 复制策略：

   • 每GPU保留额外专家槽位
   • 复制热门专家减少拥塞

图17显示，这两种策略在8×H100系统上实现了12.5%-15.5%的速度提升，接近理论最优解的90%。

7. 技术延伸与应用前景

本文提出的见解和技术不仅适用于晶圆级GPU，还可扩展到：

1. 多GPU集群：通过专家感知的任务分配减少节点间通信
2. CXL内存系统：优化专家在异构内存层级中的放置
3. 闪存存储系统：利用预测器减少存储层级间的数据移动

特别是Insight 3（专家位置感知的任务分配）和Insight 1/2（时间局部性预测）构成了通用优化原则，可应用于各种大规模MoE服务系统。