1. 3D NAND闪存PIM架构概述在大型语言模型(LLM)推理过程中传统计算架构面临的主要瓶颈在于数据搬运带来的延迟和能耗。计算存储一体化(PIM)技术通过将计算单元直接嵌入存储设备从根本上解决了这一痛点。3D NAND闪存凭借其垂直堆叠的存储单元结构在单位面积内实现了极高的存储密度成为PIM架构的理想载体。我们提出的架构采用QLC(四层单元)与SLC(单层单元)混合存储方案其中QLC用于大容量KV缓存存储SLC则用于需要快速访问的计算中间结果。这种混合设计在容量和性能之间取得了最佳平衡。每个NAND平面被配置为256×2048×128的阵列结构经过设计空间探索验证这一尺寸能够在PIM延迟和计算吞吐量之间达到最优折衷。关键设计选择采用H-tree网络拓扑连接256个平面这种对称分布结构确保了信号传输延迟的一致性同时支持多平面并行计算。实测显示H-tree的布线面积仅占平面总面积的0.39%几乎不会带来额外的面积开销。2. LLM推理中的关键计算映射2.1 QKT计算的数据流优化在Transformer架构的注意力机制中QKT(Query-Key乘积)计算占据了主要计算量。我们的架构将其分解为多个向量-向量乘法(VVM)的并行执行。具体实现上将查询向量q广播到键矩阵K的所有行K矩阵保持非转置状态存放在页缓冲区内每个NAND平面对应处理一个或两个注意力头的计算(根据模型规模调整)RPU(行处理单元)在H-tree网络中并行执行乘积运算这种映射方式充分利用了3D NAND的并行特性。以OPT-30B模型为例当head维度dh128时单个QKT计算可分解为128个独立的VVM操作分布在多个平面上同时执行。2.2 SV计算的创新映射方法SV(Score-Value乘积)的计算面临独特挑战——随着生成令牌数量L的增加数据分布会动态变化。我们创新性地采用向量-标量乘法(VSM)方案将注意力分数S向量按元素分散到不同平面每个S元素与值矩阵V的对应行进行标量乘法通过RPU中的移位加法器实现行内累加数学表达为SV Σ(S[i] * V[i,:])。这种行式乘积方法完美适应了生成式任务中序列长度可变的特点。图13(f)展示了如何通过RPU-Stream数据流实现动态长度的并行计算。3. 核心硬件架构设计3.1 行处理单元(RPU)微架构RPU是执行计算的核心部件其主要特性包括250MHz时钟频率平衡了计算延迟和功耗8个INT16乘法器和9个INT32加法器支持两种工作模式RPU-ALU用于密集矩阵运算RPU-Stream处理动态长度数据流特别设计的64位和256位寄存器文件支持不同精度的数据暂存。在QLC平面中RPU通过修改页缓冲区的移位加法器实现原位计算无需额外数据搬运。3.2 混合精度计算方案为兼顾计算精度和能效我们采用W8A8(权重8位/激活8位)量化策略KV缓存使用8位整数量化中间累加采用32位精度防止溢出层归一化(LN)和softmax在ARM核上执行FP16计算实测表明这种混合精度方案相比FP16全精度计算可将能效比提升2.1倍而模型准确率损失小于1%。4. 性能优化关键技术4.1 头级并行处理MHA(多头注意力)的天然并行性通过以下方式实现每个die处理1-2个注意力头8个die堆叠实现8-16头并行H-tree网络提供384GB/s的总带宽在OPT-175B模型(96头)中采用6个BGA封装(共48die)即可满足全部头的并行需求。这种设计使得计算时间基本不随头数增加而延长。4.2 流水线调度策略三级流水线最大化硬件利用率输入I/O阶段通过PCIe 5.0×4接口接收数据(16GB/s)PIM计算阶段与I/O阶段重叠执行输出I/O阶段H-tree网络并行传输结果对于dMVM(解码阶段矩阵乘)PIM阶段替换为KV缓存读取通过预取机制隐藏延迟。实测流水线效率达到92.3%显著优于GPU的离散内存架构。5. 实测性能分析5.1 延迟对比在OPT系列模型上的测试显示相比4块RTX4090(使用vLLM)平均加速2.4倍与4块A100(AttAcc模拟)相比仅有4.9%的延迟差距单令牌生成延迟稳定在10ms左右(输入/输出长度1K)特别值得注意的是随着序列长度增加由于创新的VSM数据流设计SV计算时间仅呈线性增长而非传统架构的平方关系。5.2 能效比优势得益于PIM架构的数据本地化特性计算能效达到58.3TOPS/W是A100的6.2倍主要功耗来自NAND阵列(72%)逻辑部分仅占28%采用peri-under-array(PUA)结构将低压外围电路制程缩放到7nm进一步降低功耗6. 实际部署考量6.1 存储容量规划QLC-SLC混合方案的具体配置2个SLC die存储模型参数和运行时中间结果6个QLC die作为KV缓存可支持175B模型的1K上下文总有效容量达2TB满足大多数LLM需求初始KV缓存写入时间约120ms(4.8GB/s带宽)在生成超过12个令牌后即可摊销这部分开销。6.2 热管理策略为防止NAND单元因持续计算发热导致数据错误动态频率调节根据温度传感器反馈调整RPU时钟计算负载均衡轮换活跃平面分布采用WARM(Write-hotness Aware Retention Management)算法管理QLC单元的保持特性实测表明这些措施可使芯片温度稳定在85°C以下确保长期可靠运行。7. 架构扩展性分析该设计具有良好的技术演进路径制程缩放外围电路可继续向5/3nm节点迁移堆叠层数当前采用128层NAND未来可升级至200层计算密度通过增加RPU数量或提升并行度实现在实验室原型中我们已经验证了256层堆叠版本计算密度可再提升1.8倍。随着3D NAND技术的持续发展这种PIM架构的性能优势还将进一步扩大。关键经验在实际部署中发现RPU时钟频率设置在250MHz时能够完美隐藏H-tree网络的累积延迟。过高频率会导致时序违例而过低频率则无法充分利用硬件并行性。这个参数需要根据具体的NAND制程特性进行精细调优。