3D堆叠NMP与Systolic Array优化LLM解码性能

1. 3D堆叠NMP与Systolic Array在LLM解码中的架构革新

在大型语言模型(LLM)推理过程中，解码阶段往往成为性能瓶颈。与预填充(prefill)阶段不同，解码以逐个令牌(token-by-token)的方式进行，权重复用有限，导致算术强度(arithmetic intensity)较低。这种特性使得解码性能对内存带宽极度敏感——当系统无法快速提供足够数据时，强大的计算单元将处于闲置状态。

传统解决方案依赖片外存储器接口，但这类接口的带宽提升面临物理限制。3D堆叠近内存处理(3D-Stacked Near-Memory Processing, NMP)技术通过将计算逻辑与DRAM堆叠在同一封装内，利用硅通孔(TSV)等垂直互连技术，实现了比传统方案高出一个数量级的本地内存带宽。这种架构特别适合LLM解码场景，因为：

• 高带宽缓解了数据供给瓶颈
• 低延迟减少了计算单元等待时间
• 能效比显著优于传统方案

然而，我们的分析揭示了一个有趣现象：当内存带宽不再是主要限制时，许多解码算子会重新进入计算受限状态。在3D-NMP环境下，逻辑芯片(die)的面积预算极为有限（通常仅约120mm²），这使得计算基板的设计成为新的挑战。

2. 从MAC树到Systolic Array的范式转变

2.1 传统MAC树的局限性

现有3D-NMP设计普遍采用基于MAC树(Multiply-Accumulate Tree)的计算单元，其架构特点包括：

• 并行乘法器阵列
• 多级加法器树进行结果归约
• 全局操作数分发网络

但随着规模扩大，MAC树暴露出明显缺点：

1. 面积效率低下：RTL实现显示，相同功能下MAC树需要8.23倍于Systolic Array的面积
2. 能耗问题：宽位宽操作数广播消耗大量能量
3. 数据复用有限：难以有效利用矩阵乘法的数据局部性

2.2 Systolic Array的架构优势

Systolic Array(脉动阵列)采用完全不同的设计理念：

PE(0,0) -> PE(0,1) -> PE(0,2) | | | v v v PE(1,0) -> PE(1,1) -> PE(1,2) | | | v v v PE(2,0) -> PE(2,1) -> PE(2,2)

数据在相邻PE间规律流动，实现高效计算

关键优势包括：

• 规则互联：仅需邻近PE间的连接，布线复杂度低
• 数据复用：输入数据在阵列中流动时被多次使用
• 计算密度：90%以上的面积用于实际计算逻辑

在LLM解码场景中，当batch size=8时，Systolic Array相比MAC树可实现：

• 2.91倍速度提升
• 2.40倍能效改善

3. 面向解码的Systolic Array微架构优化

3.1 动态重构设计

LLM解码算子的矩阵形状呈现显著多样性，特别是当采用分组查询注意力(GQA)或多查询注意力(MQA)时。我们的测量显示，OPT-66B模型中：

• 投影层(Projection)的M维度通常为8-32
• 注意力计算中K/V维度可达4096
• MoE层专家FFN的N维度变化剧烈

为此，我们提出SNAKE映射技术，使物理PE阵列能动态重构为不同逻辑形状：

重构过程仅需1个周期，通过可配置交叉开关实现数据路径切换。

3.2 数据流自适应

针对不同算子特征，我们支持两种高效数据流：

1. 输出静止(Output Stationary, OS)：

   • 适合K>N的情况
   • 输出数据保留在PE中
   • 沿K维度时间展开

2. 输入静止(Input Stationary, IS)：

   • 适合N>K的情况
   • 输入数据保留在PE中
   • 沿N维度时间展开

数据流选择算法：

def select_dataflow(N, K): if K >= 1.5 * N: return OS elif N >= 1.5 * K: return IS else: return OS if estimated_cycles(OS) < estimated_cycles(IS) else IS

3.3 存储层次优化

3D-NMP的高带宽特性改变了传统存储设计原则：

优化后的存储层次：

1. 权重缓存：32KB (左右边界各16KB)
2. 输入缓存：8KB (顶部边界)
3. 输出缓存：16KB (底部边界，与向量核共享)

4. 多核调度框架

4.1 空间-时空分区策略

在16个处理单元(PU)的系统中，我们开发了创新的调度方案：

1. 纯空间分区(IS-S/OS-S)：

   • 沿K或N维度分割工作负载
   • 各PU独立处理分配到的分区
   • 适合高度不均衡的算子

2. 时空混合分区(IS-ST/OS-ST)：

   • 同时沿空间和时间维度分割
   • 每个PU处理部分空间分区和部分时间步
   • 适合规整的大矩阵运算

4.2 算子感知调度

针对不同类型算子采用特定策略：

4.3 轻量级互连设计

采用两种逻辑拓扑适应不同调度模式：

1. 1x16链式：用于纯空间分区
2. 4x4网格：用于时空混合分区

互连仅需支持三种基本操作：

• 广播(Broadcast)
• 归约(Reduce)
• 移位(Shift)

5. 实际部署考量

5.1 热设计考虑

在3D堆叠环境中，热管理至关重要。我们的测量显示：

• Systolic Array的功耗密度为0.8W/mm²
• 峰值温度出现在阵列中心区域
• 采用动态频率调节可保持结温<85°C

温度管理策略：

1. 每PU集成温度传感器
2. 超过阈值时自动降频
3. 任务迁移至较冷PU

5.2 与主机处理器协同

典型部署模式：

1. 主机GPU/TPU处理预填充阶段
2. 3D-NMP专精解码阶段
3. 通过PCIe或专用接口通信

数据流示例：

主机 -> 权重预取 -> 3D-NMP -> 生成令牌 -> 主机

5.3 实际性能数据

在DeepSeek 236B模型(batch=8)上的测试结果：

6. 开发者实践指南

6.1 硬件描述语言实现

建议的Verilog编码风格：

module PE #(parameter WIDTH=16) ( input clk, rst, input [WIDTH-1:0] in_left, in_top, output [WIDTH-1:0] out_right, out_bottom ); reg [WIDTH-1:0] partial_sum; always @(posedge clk) begin if (rst) partial_sum <= 0; else partial_sum <= partial_sum + in_left * in_top; end assign out_right = in_left; assign out_bottom = in_top; endmodule

6.2 编译器优化要点

1. 算子融合：将相邻的线性/非线性算子合并
2. 数据布局转换：提前将权重转为适合Systolic Array的格式
3. 双缓冲管理：隐藏数据传输延迟

6.3 性能调优检查表

1. [ ] 确认数据流选择与算子形状匹配
2. [ ] 检查PE利用率是否>80%
3. [ ] 验证存储带宽利用率在70-90%之间
4. [ ] 监控温度曲线是否符合预期
5. [ ] 分析算子间流水线空隙

7. 未来演进方向

1. 光互连集成：利用硅光子技术进一步提升带宽
2. 存内计算：结合FeFET等新型器件实现更高能效
3. 3D集成扩展：向更高堆叠层数发展
4. 自适应精度：动态调整计算精度平衡质量与效率

这种架构创新不仅适用于LLM，也可扩展至：

• 推荐系统矩阵运算
• 科学计算中的稀疏代数
• 计算机视觉中的特征变换

在实际芯片设计中，我们验证了这种方法的可行性：在TSMC 7nm工艺下，单个计算核心面积仅3.2mm²，运行频率可达1.2GHz。通过将16个这样的核心集成在逻辑芯片上，配合4层HBM3内存堆叠，实现了突破性的解码性能。