3D-DRAM加速器如何突破生成式推理带宽墙与KV cache瓶颈?

3D-DRAM加速器如何突破生成式推理带宽墙与KV cache瓶颈? 1. 生成式推理撞上的不是算力墙是带宽墙做生成式推理优化的人迟早会撞上同一面墙模型权重和 KV cache 都得从内存里搬出来搬到乘加器旁边。算力早就过剩带宽一直不够。我最早意识到这件事是在给一个 7B 模型做单流解码压测的时候——GPU 利用率趴在 30% 出头显存带宽却已经顶到 90% 以上SM 里的 Tensor Core 大部分时间在等数据。后来做 3D-DRAM 加速器这件事起点就是这面墙把 DRAM 堆起来、把计算单元塞进内存旁边让数据少走几毫米的线。这篇内容想讲清楚几件事为什么生成式推理的瓶颈落在访存而不是算力3D-DRAM 加速器在架构上到底动了哪些地方带宽、能效、容量这三笔账该怎么算以及从模型映射到上板调优这条链路上哪些坑是我自己踩过的。适合做推理系统、加速器架构、存算一体方向的人看也适合只是想搞清楚“为什么我的卡在 decode 阶段跑不满”的工程同学。1.1 先用算术强度把病根定位出来生成式推理分两段prefill 和 decode。prefill 一次处理整段 prompt矩阵乘的规模大算术强度高属于典型的计算密集型decode 一次只生成一个 token每个 token 都要把整份权重从头读一遍属于典型的访存密集型。真正折磨人的是 decode因为它占了整个服务生命周期里绝大部分时间。拿一个 7B 模型、FP16 权重来算。权重总量约 14 GB生成一个 token 需要把这 14 GB 全读一遍产生的计算量是每个权重参与一次乘加也就是 2 FLOPs。算术强度就是 2 FLOPs / 2 Bytes 1 FLOP/Byte。这个数字低得离谱。相比之下一个矩阵规模 4096×4096×4096 的 prefill 矩阵乘算术强度能到几百 FLOP/Byte。算术强度低意味着什么意味着你的算力利用率被带宽锁死了。假设加速器峰值 100 TFLOPS在算术强度 1 FLOP/Byte 的负载下要喂饱它需要 100 TB/s 的带宽。而当前主流的高带宽内存单栈带宽在 1 TB/s 量级八栈也就 6 到 8 TB/s差了整整一个数量级。这就是墙上那行字不是算力不够是带宽差了十倍以上。有人会说把 batch 做大不就行了。确实batch 增加会提升算术强度因为同一份权重读一次可以服务 batch 个 token。batch64 时算术强度理论上能到 64 FLOP/Byte账面上好看很多。但 batch 做大有三重代价KV cache 线性膨胀、单 token 延迟变高、显存容量先撑不住。所以在实际服务里大 batch 和低延迟是一对天生的矛盾很多时候你只能在 batch1 到 8 这个区间里工作算术强度依然低得可怜。1.2 单纯堆卡为什么救不了 decode最直觉的方案是加卡用张量并行把权重切开让每张卡只读自己那一份。这个思路在 prefill 阶段效果不错但在 decode 阶段收益会被通信吃掉。原因在于张量并行每层要做两次 all-reduce一次在前向 attention 之后一次在 FFN 之后。通信量和 batch、hidden size 成正比和切分卡数关系不大。decode 阶段每步的计算量本来就小通信延迟很难被隐藏。典型表现就是卡数翻倍吞吐涨了不到一半单 token 延迟反而涨了。因为每生成一个 token 要同步十几次每次几十微秒累加起来比计算本身还长。还有一种方案是靠片上 SRAM 缓存权重比如把一部分热权重钉在 L2 或者 SMEM 里。这个在小模型上有效但 7B 以上的模型权重体量远超片上 SRAM 容量命中率上不去最后还是得回落到 DRAM。SRAM 密度低、每 bit 成本高这是物理层面的限制不是调度能绕过去的。把这两条路都走一遍之后结论其实很清晰带宽密度和每 bit 能耗必须同时改善而且改善幅度要足够大。这正好是 3D-DRAM 能提供的两样东西。1.3 3D-DRAM 加速器真正想解决的三件事第一个是带宽密度。传统封装里DRAM 颗粒和计算 die 之间隔着 PCB 走线或者硅中介层通道数受限于引脚数量和布线密度。3D 堆叠之后DRAM die 和逻辑 die 之间靠硅通孔垂直互连互连密度提升一到两个数量级每层可以独立开通道总带宽可以堆到 TB/s 量级甚至更高。第二个是每 bit 能耗。数据搬运的能耗和走线长度强相关。跨 PCB 走线或者穿过中介层的长距离互连每 bit 能耗在几 pJ 量级垂直短距互连能压到零点几 pJ。这个差距看起来不大但乘以每秒几万亿 bit 的搬运量就是几十瓦的差距。对于功耗预算紧张的推理设备来说这几十瓦直接决定了能不能上单槽。第三个是把计算搬到数据旁边。前两项改善的是“搬得快不快、省不省”这一项改的是“要不要搬”。生成式推理的 decode 阶段做的是 GEMV矩阵向量乘向量只有一个权重全部要用到天然适合把乘加单元放到存储阵列附近在数据出来的位置就地做部分和只把归约后的结果往上传。搬运量能从 O(权重) 压到 O(输出)。至于为什么这个方向到最近才被认真做我的理解是三个条件同时成熟了异构集成的工艺成本降下来了混合键合和 TSV 的良率上来了生成式推理的负载特征被完整地测量和公开了大家终于承认 decode 是主要矛盾以及算力和带宽的剪刀差拉得足够大逼着大家往内存里面找答案。第一个针对生成式推理的 3D-DRAM 加速器就是在这个时间点上出现的。2. 3D-DRAM 加速器的架构拆解2.1 堆叠结构逻辑底座加存储层结构上我把这类加速器理解成一个“倒过来的三明治”底层是逻辑底座 die上面垂直堆叠若干层 DRAM die层与层之间用 TSV 和微凸点连接最底层通过封装基板对外提供主机接口。逻辑底座负责指令调度、数据搬运、归约和对外通信DRAM 层负责存权重和 KV cache。层数不是越多越好。每增加一层堆叠高度、散热路径长度、TSV 占用的面积都会跟着涨。我评估过的配置里8 层是一个比较舒服的区间容量能到 16 到 32 GB带宽能堆到 4 TB/s 以上堆叠高度还能控制在封装工艺的常规范围内。层数往 12 到 16 走容量和带宽继续涨但中间层的散热会变成很麻烦的事情。每层内部的通道组织也值得说。传统 DRAM 是有限数量的宽通道3D 堆叠之后每层可以独立开若干组通道每组通道再分到 bank group、bank、subarray。这种“层—通道—bank”的三级结构是后面做 bank-aware 数据布局的基础。如果布局没对齐同一时刻大量请求落到同一个 bank带宽再高也跑不满。2.2 计算单元到底放在哪一层这是架构设计里最核心的取舍。我梳理过三条路线各自的代价差别很大。方案工艺面积开销散热可编程性适用场景逻辑底座集中计算先进逻辑工艺底座面积吃紧底座好散热高通用 GEMV/GEMMDRAM die 上嵌入计算单元DRAM 工艺占用阵列面积中间层难散热中规则定点运算sense amplifier 附近做存内计算DRAM 工艺面积开销最大阵列内热点低位运算、简单累加逻辑底座集中计算的好处是工艺先进、频率高、可编程性强能把量化、归约、非线性这些复杂逻辑都放进去。缺点是所有数据还是要从 DRAM 层经过 TSV 搬到最底层搬运量没降下来只是搬得更近了。DRAM die 上嵌入计算单元的好处是搬运路径更短数据从阵列出来直接在本地做部分和。代价是要在 DRAM 工艺里塞逻辑DRAM 工艺的晶体管密度和频率都落后于逻辑工艺复杂运算做不了。所以这条路一般只做定点乘加和累加把归约往上传。sense amplifier 附近做存内计算最激进理论上搬运量最小但面积开销最大因为要在每个 subarray 旁边加算术单元DRAM 的版图规则又很紧实际能塞进去的逻辑非常有限而且阵列内部的热点很难导出。我最后倾向的是混合方案底座负责调度和归约DRAM 层里放轻量的乘加阵列。这样既压低了搬运量又不用在 DRAM 工艺里硬塞复杂逻辑。2.3 数据通路与片上互联数据从 DRAM 阵列出来到乘加单元再往上归约到输出这条路径上每一步都要算清楚。权重是驻留的一次加载常驻在 bank 的 row buffer 或者本地缓存里后续所有 token 复用激活向量是流动的每个 token 换一批部分和需要在层内累加跨层的归约放到逻辑底座做。互联设计上有两个容易忽略的点。一是 TSV 是共享资源如果计算单元分布在多层 DRAM die 上部分和的汇聚要经过多条 TSV容易在某个通道上形成热点。二是归约精度如果把部分和的位宽压得太低多层累加之后误差会超出可接受范围这个问题在 INT8 量化下尤其明显。还有一个上手时容易低估的事片上需要一个专门管调度的控制器负责把 GEMV 拆成 bank 级的任务、安排时序、躲开冲突。这个控制器的复杂度经常被低估实际上它是决定实际带宽能否逼近峰值的关键。3. 关键参数反推带宽、能效、容量怎么定3.1 从目标吞吐反推带宽做硬件方案评估我习惯从目标业务指标往回推而不是先定硬件参数再看能跑什么。假设我们要在一台设备上跑 70B 模型的 INT8 量化版本目标是单流 50 token/s。INT8 权重约 70 GB生成每个 token 要把这 70 GB 读一遍那么带宽需求就是 70 GB × 50 /s 3.5 TB/s。注意这个数字和 batch 无关——权重部分的总带宽需求只取决于步频batch 做大摊薄的是每个 token 分摊到的权重带宽。再换一个角度如果目标是 batch32、整体吞吐 1600 token/s那还是每秒 50 步权重带宽需求依然是 3.5 TB/s但每步要额外读 32 份 KV cache 分片。这时候 KV cache 的带宽开销就冒出来了而且随着上下文变长会持续上升。def bandwidth_budget(params_b, bytes_per_param, steps_per_sec, kv_bytes_per_token, batch, context_len): weight_bytes params_b * 1e9 * bytes_per_param kv_bytes kv_bytes_per_token * context_len * batch # 每步读一遍权重 batch 份 KV per_step weight_bytes kv_bytes return per_step * steps_per_sec / 1e12 # TB/s # 70B INT850 步/秒batch1上下文 4096 print(bandwidth_budget(70, 1, 50, 320 * 1024, 1, 4096)) # 约 3.57 TB/s # 同样条件 batch32 print(bandwidth_budget(70, 1, 50, 320 * 1024, 32, 4096)) # 约 5.03 TB/s这个脚本我在方案评估时改过很多遍它解释了我早期一个困惑为什么把 batch 提上来之后带宽压力反而更大。因为权重那块是常数KV cache 那块随 batch 线性增长两边加起来不会因为 batch 增大而下降太多。真正让权重带宽均价变低的方式是提升 batch 的“服务效率”而不是减少总读取次数。3.2 能效预算pJ/bit 才是真正的胜负手带宽决定能不能跑能效决定能不能装进机箱。这一块我建议一开始就把账算清楚。传统外挂内存的数据搬运含 I/O 驱动、接收、控制器开销公开资料里常见的量级是每 bit 几 pJ。经过硅中介层的 2.5D 互连会好一些但也在 1 pJ 上下。垂直短距互连的 TSV每 bit 能压到零点几 pJ。这里取 0.2 pJ/bit 做一个估算。互连形态每 bit 能耗常见量级25 W 预算下可支撑带宽外挂 PCB 走线5 pJ约 0.6 TB/s2.5D 中介层1 pJ约 3.1 TB/s3D TSV 垂直互连0.2 pJ约 15.6 TB/s换算过程很直白25 W 就是 25 J/s除以 0.2 pJ/bit 得到 1.25×10^14 bit/s折算成字节是约 15.6 TB/s。同样 25 W 预算如果每 bit 能耗是 5 pJ只能支撑 0.6 TB/s 出头。这个表基本解释了为什么做推理加速的人都往 3D 堆叠上走——不是带宽数字好看是能耗账能算得过来。当然TSV 互连只是搬运成本的一部分。DRAM 阵列本身的激活、读写操作也有能耗控制器逻辑也有开销。把这三项加在一起实际每 bit 等效能耗会落在 2 到 4 pJ 区间。即便如此相对外挂方案依然有数倍的优势而且这个优势会随着带宽需求上升而放大。下面这张表是我在方案评估阶段用的 8 层堆叠配置列出来方便对照。参数取值逻辑底座工艺先进逻辑节点DRAM 堆叠层数8每层通道数16每通道位宽64 bit每 pin 速率4 Gbps单层带宽512 GB/s总带宽4 TB/s总容量16 GBTSV 满带宽能耗约 6.4 W含阵列与控制的等效能耗约 2.5 pJ/bit单层带宽的算法是 16 通道 × 64 bit × 4 Gbps ÷ 8 512 GB/s八层加起来 4 TB/s。满带宽下 TSV 互连功耗是 3.2×10^13 bit/s × 0.2 pJ/bit等于 6.4 W。这个数字在散热设计里是可以接受的。3.3 容量账KV cache 吃掉了多少容量这块权重是静态的好规划KV cache 是动态的会把人坑到。还是拿 70B 模型举例80 层64 个注意力头每个头维度 128采用 GQA 且 KV 头数为 8。每个 token 每层的 KV 元素数是 2 × 8 × 128 2048 个FP16 存储就是 4096 字节。乘以 80 层每个 token 的 KV cache 是 327,680 字节约 320 KiB。这个数乘以上下文长度和 batch4096 上下文、batch1约 1.25 GiB32768 上下文、batch1约 10 GiB131072 上下文、batch1约 40 GiB4096 上下文、batch32约 40 GiB结论很残酷长上下文和大 batch 是两个会把 KV cache 吃爆的方向而且它们同时发生的时候是乘法关系。16 GB 的堆叠容量在 32K 上下文、batch8 的配置下就已经很紧张了。所以 3D-DRAM 加速器在容量规划上必须把 KV cache 单独算一块预算出来不能只算权重。我给自己的经验值是这样容量规划的起步配置权重占 60%KV cache 占 30%留 10% 给激活和中间结果。如果业务场景以长上下文为主KV cache 的比例还要往上提。4. 从模型到硬件的落地流程4.1 权重与 KV cache 的 bank-aware 布局模型映射的第一步不是写 kernel是确定权重的摆放位置。核心理由是 DRAM 的 row buffer 有局部性红利访问落在已经打开的 row 里成本远低于重新激活一个 row。所以布局要尽量让同一时刻被访问的权重落在已经打开的 row 里。具体做法是按输出通道维度切分权重矩阵让同一组输出通道的权重连续放置并且对齐到 bank 边界。这样 GEMV 的每一轮只需要激活少量 row其他时间都在读 row buffer。如果不做对齐每读几个数就要重新激活 row有效带宽会掉到峰值的三成以下。KV cache 的布局逻辑不同因为它是按 token 增长的追加写入和随机读取并存。我的做法是按层分 bank、按 KV 头分通道让不同层的 KV 落在不同的 bank group 里避免所有层同时抢同一组通道。写入时按 token 顺序追加读取时按 layer 和 head 并行发起。下面是我用过的一个布局示例列出来方便理解这种对齐思路。数据对象切分维度落到哪一层对齐要求FFN 权重输出通道固定层bank 边界连续存放Attention 权重head 分组固定层channel 边界KV cachelayer × head按 batch 分散至多层bank group 边界激活向量按 token底座缓存无需对齐这套布局做完之后实测有效带宽能到峰值的七成到八成。剩下的两成损失主要来自调度抖动和刷新操作这个后面再说。4.2 调度prefill 和 decode 分开走prefill 和 decode 的硬件需求差别太大我的经验是不要让它们共用同一条调度路径。prefill 是计算密集的适合把算力铺满、把矩阵乘做大decode 是访存密集的适合把 bank 并行度拉满、把每一次读的复用率做高。混在一起调度两边都做不好。我采用的方案是 chunked prefill 加连续批处理。把长 prompt 切成固定长度的 chunk和 decode 请求拼在同一个批次里让硬件在两个阶段之间平滑过渡。这样做的实际收益是平均延迟下降而且硬件利用率曲线变得平缓不会出现大段空转。还有一个细节decode 阶段的 batch 不要设太满。我试过把 batch 拉满结果单 token 延迟涨了三成因为 KV cache 读取把权重读取的带宽挤掉了。留出 20% 到 30% 的带宽余量给 KV cache整机吞吐反而更高。这个经验在 GPU 上同样适用只是数字不一样。4.3 仿真验证与上板调优的顺序我的验证顺序是先在行为级模型上跑通功能再用周期级仿真器跑带宽和能耗最后上 FPGA 原型做数值和调度验证最后才流片。跳过任何一步后面的问题都会放大。周期级仿真这一环最关键。我用的方法是把真实推理 trace 打下来喂给带 DRAM 时序模型和 NoC 模型的仿真环境重点看三个指标有效带宽、bank 冲突率、TSV 通道占用分布。这三个指标一旦有异常基本能定位到是布局问题还是调度问题。上板之后调优顺序我建议固定成先调数据布局再调调度策略最后才碰频率和电压。原因是布局带来的收益通常是百分之几十调度带来的收益是百分之十几频率只有百分之几。顺序搞反了你会发现调了半天频率收益还不如改一行对齐参数。5. 影响范围这类加速器会动到哪几块5.1 推理服务端的成本结构这是最直接的影响。当单设备带宽从 TB/s 级别往上走、每 bit 能耗下降数倍之后同样的吞吐可以用更少的设备跑出来而且功耗墙不再是硬约束。对于大规模在线推理服务来说设备数量、机房空间、供电和散热这三块的支出都会被重算。更微妙的变化在成本分布上。原来的成本大头在加速卡采购和电力现在一部分转移到内存堆叠的封装成本上。也就是说上游的堆叠封装产能会成为新的瓶颈。这个转移在近几年已经开始出现3D 加速器只会让它更明显。5.2 模型侧长上下文和小 batch 变得更划算当 KV cache 的读取带宽不再是瓶颈之后长上下文场景的性价比会明显改善。现在很多应用把上下文压在 4K 到 8K不是模型不支持更长是每多一段上下文都要多付带宽和延迟。带宽上来之后32K 甚至更长的上下文在延迟上会更可接受。另一个变化是小 batch 变得划算。现在的服务架构拼命把 batch 做大本质上是为了摊薄权重带宽。如果加速器本身带宽足够batch1 到 4 也能跑出不错的吞吐那就能做一些低延迟的交互式应用比如实时对话、代码补全这类对首字延迟敏感的场景。5.3 硬件侧内存和逻辑的边界会变模糊3D-DRAM 加速器把逻辑和存储堆到了一起这对产业链的影响不只是技术层面的。传统上内存厂商做 DRAM 颗粒逻辑厂商做计算 die两家各管一段接口标准化程度很高。堆叠之后谁来定义底座逻辑、谁来负责整体验证、良率损失怎么分摊这些问题都需要重新谈。我的判断是中间会出现一批新的角色专门做堆叠方案的设计和集成验证。这类工作既需要懂 DRAM 时序和阵列组织又需要懂逻辑设计和推理负载特性门槛不低但对做系统的人来说机会不少。6. 踩坑实录与排查清单6.1 带宽跑不满的几种典型原因这是我踩得最多的一类坑整理成表方便对照。现象可能原因排查手段处理方式有效带宽只有峰值 30%row buffer 频繁换行统计 row activate 次数调整权重对齐到 bank 边界带宽波动大bank 冲突集中看各 bank 请求分布打散热点权重到不同 bank group部分通道空闲TSV 通道分配不均抓通道占用直方图重排任务到通道的映射长时间掉到零刷新操作集中触发看刷新时序与访问的叠加调度器避让刷新窗口带宽正常但吞吐低归约成为瓶颈统计归约单元占用率拆归约层级底座并行处理其中最常见的是 row buffer 换行问题。很多团队在做布局时只考虑了连续性和对齐忘了 row buffer 容量是有限的。一个 row 打开后能装的数据量是固定的如果一轮访问需要的权重超过一个 row 的容量就会换行换行就意味着重新激活成本直接翻倍。解决办法是把一轮 GEMV 的权重切得足够小小到能装进一个 row 的容量里。6.2 热与功耗的两个坑第一个坑是热密度集中在堆叠中间层。DRAM die 上的乘加单元如果布置得太密中间层的温度会比顶层和底层高出十几度。温度上去之后DRAM 的刷新率要提升刷新又要占带宽形成负反馈。我的处理方式是把计算单元的密度降下来宁可多放几层也不要单层做太密同时在逻辑底座上加温度传感器做动态调节。第二个坑是满带宽下的功耗峰值。带宽跑到 4 TB/s 的时候TSV 互连、阵列操作、控制器三块加起来功耗会到一个不小的数字如果电源设计没有留余量会出现掉压导致误码的情况。我的做法是把峰值功耗按满载的 1.3 倍设计供电同时在固件里做带宽和功耗的联动调节接近功耗上限时主动降低并发度。6.3 精度与软件栈精度这块的问题主要在累加位宽。INT8 乘加的部分和如果只用 16 位累加累加几百次之后就会溢出或者损失有效位。多层归约之后误差会放大表现为生成文本的重复、语义漂移或者直接输出乱码。我的经验是部分和至少用 24 位跨层归约用 32 位这样在 INT8 下基本不出问题。软件栈这块的坑更隐蔽。加速器再怎么优化前端如果不做算子融合、不做量化校准收益都会被吃掉。我遇到过模型侧没做 RoPE 融合导致每层多读一次 KV cache带宽开销直接涨了两成。这类问题的排查方式是逐层统计访存量和理论值对不上就说明少了融合。最后说一个我自己的习惯任何一次方案评估我都会先写一个带宽和能耗的估算脚本把所有假设都写成参数。因为参数一改结论可能完全反过来——比如每 bit 能耗如果按 1 pJ 算25 W 只能支撑 3 TB/s方案就完全站不住了。这些数字是真的会决定方案生死的。