Hardware-Software Co-design for 3D-DRAM-based LLM Serving Accelerator

本文提出了名为 Helios 的基于混合键合技术的 LLM 服务加速器，通过设计细粒度的分布式注意力执行流程与空间感知的 KV 缓存分配机制，有效解决了现有近存计算架构在处理动态 LLM 服务负载时的灵活性不足问题，显著提升了加速性能与能效。

要点

这篇论文介绍了一种名为 Helios 的新型硬件加速器，专门用来让大型语言模型（LLM，比如现在的各种 AI 聊天机器人）跑得更快、更省电。

为了让你轻松理解，我们可以把运行 AI 模型想象成经营一家超级繁忙的“信息翻译工厂”。

1. 现在的痛点：工厂里的“拥堵”与“浪费”

想象一下，这家工厂（AI 模型）每天要处理成千上万个订单（用户提问）。

• 订单长短不一：有的用户只问“你好”（短文本），有的用户让 AI 写一本小说（长文本，几万个字）。
• 需求忽高忽低：有时候没人问，有时候瞬间涌入几万人。

现有的工厂（GPU 或普通近存计算芯片）遇到了两个大麻烦：

• 麻烦一：仓库管理太死板（KV Cache 管理粗糙）

  • 现状：工厂给每个订单分配了一个固定的“大仓库格子”。不管订单是写一句话还是写一本书，系统都强制预留写一本书那么大的空间。
  • 后果：如果来了个短订单，仓库大半是空的，空间浪费严重。而且，如果突然来了个超长订单，现有的格子不够大，系统还得等别的订单搬走才能腾出地方，导致排队拥堵。
  • 比喻：就像你去超市买一瓶水，超市却非要给你整个货架的位置，还规定只有买满一货架的人才能用这个货架，结果货架空着，后面想买的人却进不来。

• 麻烦二：工人干活不灵活（注意力机制执行僵化）

  • 现状：工厂里的工人（计算单元）被固定分配了特定的任务。比如，工人 A 只负责处理“前 100 个字”，工人 B 只负责“后 100 个字”。
  • 后果：如果订单很短，工人 B 就闲着没事干；如果订单很长，工人 A 累死也干不完，而工人 B 却帮不上忙。大家不能互相帮忙，导致忙闲不均。
  • 比喻：就像一群厨师，每个人只负责切一种菜。如果今天只有切土豆的活，切牛肉的厨师就只能站着发呆，效率极低。

2. Helios 的解决方案：智能的“乐高积木”工厂

Helios 提出了一套全新的硬件 + 软件协同设计，就像把工厂改造成了高度灵活的“乐高积木”系统。

核心创新一：灵活的“动态仓库” (Spatially-Aware KV Cache Allocation)

• 怎么做：Helios 不再给订单分配固定的大格子，而是把仓库切成了无数个小积木块。
• 效果：
  • 短订单只占几个小积木块，不浪费空间。
  • 长订单可以动态拼接更多积木块。
  • 智能调度：系统会像“智能物流机器人”一样，把订单的积木块均匀地分散到所有可用的仓库区域，避免有的区域爆满，有的区域空着。
• 比喻：就像玩俄罗斯方块，不管方块形状多奇怪，系统都能把它们严丝合缝地拼在一起，最大化利用空间，不再有空隙浪费。

核心创新二：流动的“流水线” (Distributed Tiled Attention)

• 怎么做：Helios 打破了工人之间的“固定岗位”。它让所有工人组成一个流动的网格。
• 效果：
  • 不管订单多长，所有工人都可以协作处理。
  • 如果订单短，大家就一起快速处理完；如果订单长，大家就接力棒式地分担工作。
  • 它还发明了一种新的“沟通语言”（通信原语），让工人之间交换数据像发微信一样快，几乎不耽误时间。
• 比喻：以前是“专人专岗”，现在是“特种部队”。不管任务多难，所有队员都能根据现场情况灵活补位，谁有空谁就干，绝不磨洋工。

核心创新三：超级高速公路 (Hybrid Bonding 技术)

• 怎么做：Helios 使用了**混合键合（Hybrid Bonding）**技术，把计算芯片和存储芯片像“三明治”一样紧紧叠在一起，中间用极细的铜柱连接。
• 效果：数据在芯片内部传输的速度极快，就像在自家客厅里拿东西，而不是要去隔壁城市取货。
• 比喻：以前的工厂，仓库和车间隔了一条河，运货要坐船（慢且贵）；Helios 直接把仓库建在车间楼上，工人伸手就能拿到材料。

3. 最终成果：快如闪电，省如节能灯

通过这套组合拳，Helios 的表现非常惊人：

• 速度快：比目前最先进的 GPU 方案快 3.25 倍。这意味着用户发问后，AI 回复的速度快了三倍多。
• 更省电：能效比提升了 3.36 倍。这意味着在同样的电量下，它能处理更多的任务，或者在同样的任务下，电费省了三分之二。
• 更稳定：即使在人流量巨大的高峰期，也能保证大部分用户（99%）的等待时间很短，不会出现“系统卡顿”的情况。

总结

简单来说，Helios 就是给 AI 模型装上了一个超级智能的“动态仓库管理系统”和“灵活协作团队”。它不再死板地分配资源，而是根据每个任务的实际情况，像变魔术一样灵活调配，让 AI 在回答你的问题时，既快又省，还能同时服务更多人。

这就好比把一家原本管理混乱、效率低下的传统工厂，升级成了全自动、智能化的现代物流枢纽，无论订单大小、多少，都能瞬间处理完毕。

技术摘要

这是一篇关于基于混合键合（Hybrid Bonding）技术的 3D-DRAM 大语言模型（LLM）服务加速器的论文，名为 HeLioS。该论文由北京大学、西安 UniIC 半导体、Houmo AI 和 HKUST 的研究团队共同完成。

以下是对该论文的详细技术总结：

1. 研究背景与问题 (Problem)

随着大语言模型（LLM）在在线生成式服务中的广泛应用，LLM 推理面临高度动态的工作负载挑战，具体表现为：

• 请求到达率波动大：高峰期与低谷期流量差异可达 1-2 个数量级。
• 请求长度差异显著：上下文长度从几百到上万 Token 不等（如 10K vs 1K）。
• 计算与访存混合：Prefill（预填充）阶段计算密集，而 Decoding（解码）阶段（特别是小 Batch 下的 Attention 和 FC 层）是访存受限的。

现有近存计算（NMP）方案的局限性：
现有的 NMP 加速器（如 AttAcc, NeuPIMs, Duplex 等）虽然通过集成处理引擎（PE）到内存模块（如 HBM）提升了带宽，但在应对动态 LLM 服务时存在两大核心瓶颈：

1. 粗粒度的 KV Cache 管理：现有方案通常将完整的 KV Cache 静态分配给固定的 NMP 模块（如 HBM 立方体）。这导致：
   • 计算负载不均衡：在低负载或短序列请求下，无法充分利用所有 NMP 模块。
   • 内存碎片化：由于请求长度动态变化，固定分配导致大量内存空间浪费（例如为短请求预留长序列空间），且无法利用分散的空闲内存块，降低了有效服务容量。
2. 僵化的 Attention 执行流：现有方案通常将 Attention 视为两个独立的 GEMM 操作，要求完整的 Key/Value 矩阵静态分布在固定模块内。这种设计无法支持细粒度的 KV Cache 分块管理（Blockwise KV Cache），限制了动态调度能力。

根本原因：NMP 架构固有的非均匀内存抽象（每个 PE 只能直接访问局部内存）与现有基于 GPU 的动态 KV Cache 管理算法（假设可任意访问所有内存）之间存在根本性不匹配。

2. 方法论与系统架构 (Methodology & Architecture)

为了解决上述问题，作者提出了 HeLioS，一种基于混合键合（Hybrid Bonding, HB）技术的 LLM 服务加速器。其核心设计理念是架构 - 系统协同设计，实现 NMP 原生的细粒度 KV Cache 管理和分布式 Tiled Attention。

A. 硬件架构 (HB-Device)

• 混合键合技术：利用 Cu-Cu 直接融合键合，将 4 层 DRAM 堆叠在逻辑晶圆上，并通过辅助层（Auxiliary Layer）聚合 I/O。
  • 优势：提供极高的 I/O 密度（110,000/mm²）和带宽，同时降低寄生电容和能耗（0.66 pJ/bit）。
  • PE 阵列：逻辑层集成 4x4 的 PE 阵列（共 16 个 PE），每个 PE 包含矩阵单元、在线 Softmax 单元、向量单元和归约单元。
  • 片上网络（NoC）：采用 Mesh 拓扑连接 PE，支持高效的片内通信。

B. 算子执行流 (Operator Execution)

• 分布式 Tiled Attention：
  • 将 KV Cache 细粒度地划分为 Block（块），并动态分配给所有 PE。
  • 采用迭代式 Tiled Attention算法（Algorithm 1），在每个 PE 上并行处理 KV Block。
  • 在线 Softmax 与归约：设计了专门的硬件单元支持在线 Softmax 和 Partial Sum 归约，通过流水线重叠（Overlap）隐藏向量操作延迟。
• 细粒度 KV Cache 传输：引入 Transfer Buffer，根据计算与访存延迟的匹配情况，动态选择细粒度或粗粒度的 KV Cache 存储策略，避免传输阻塞计算。

C. 系统调度：空间感知 KV Cache 分配 (Spatially-Aware KV Cache Allocation)

• 核心算法：提出了一种空间感知的分配策略（Algorithm 2），旨在平衡计算负载并最小化 PE 间通信开销。
• 分配逻辑：
  1. 负载均衡：优先将 Block 分配给当前 Token 数量（计算负载）最少的 PE。
  2. 通信优化：在负载相同时，根据 PE 在 Mesh 网络中的位置进行排序。
     • 对于完整 Block，优先分配给距离中心 PE 较远的 PE，保留中心 PE 给后续可能需要的 Block，减少通信跳数。
     • 对于最后一个 Block（Last Block），优先分配给距离中心 PE 较近的 PE，因为 Key/Value 向量的传输通常汇聚到中心 PE 进行归约/散射。
• 效果：实现了计算负载与通信路径的联合优化，避免了内存碎片化。

3. 关键贡献 (Key Contributions)

1. 深入分析：揭示了现有 NMP 设计在处理动态 LLM 服务负载时，因粗粒度 KV 管理和僵化执行流导致的性能瓶颈。
2. HeLioS 架构设计：提出了基于混合键合的加速器架构，支持NMP 原生的分布式 Tiled Attention，实现了从 Prefill 到 Decoding 的全流程支持。
3. 系统级创新：设计了空间感知的 KV Cache 分配机制，解决了 NMP 非均匀内存抽象与动态 KV 管理之间的不匹配问题，实现了细粒度的动态资源调度。
4. 性能突破：通过全栈优化（架构、算子、调度），显著提升了 LLM 服务的吞吐量和能效。

4. 实验结果 (Results)

作者在多种主流模型（OPT 66B, LLaMA3 70B, Mixtral 8x22B, Qwen3, DeepSeek 236B）和不同工作负载下进行了评估，对比基线包括 GPU (A100)、In-die NMP (AttAcc)、Duplex 等。

• 解码加速比：相比现有 GPU/NMP 设计，HeLioS 实现了 3.25 倍（几何平均）的加速。
• 能效提升：实现了 3.36 倍（几何平均）的能效提升。
• 服务性能：
  • P50/P99 延迟：在 Token 间隔时间（TBT）上，相比最佳基线降低了 72% / 76%。
  • 服务容量：在高动态负载下，吞吐量提升 1.19x - 7.33x。
• 碎片化与负载均衡：
  • 相比传统 NMP 方案，KV Cache 内存碎片化比率降低了 4.67x - 23.10x。
  • PE 间的计算负载不均衡度降低了 1.79x - 13.69x。
• 通信开销：PE 间通信仅占总延迟的 2.94% - 16.10%（平均 7.08%），证明了通信设计的高效性。

5. 意义与总结 (Significance)

HeLioS 的工作具有重要的学术和工业价值：

1. 填补了空白：它是首个针对动态 LLM 服务场景，专门针对NMP 架构特性进行全栈协同设计的加速器方案。它成功解决了 NMP 架构中“非均匀内存”与“动态 KV 管理”之间的根本矛盾。
2. 技术前瞻性：充分利用了**混合键合（Hybrid Bonding）**这一新兴封装技术，展示了 3D-DRAM 在 AI 推理领域的巨大潜力，特别是在解决“内存墙”问题上的优势。
3. 实用性强：提出的空间感知分配算法和细粒度 Tiled Attention 执行流，不仅适用于当前的 LLM，也为未来更长上下文、更大规模模型的服务提供了可扩展的架构范式。
4. 性能验证：实验结果表明，通过软硬协同设计，可以显著超越单纯依靠增加带宽或算力的传统方案，为下一代 LLM 服务基础设施提供了新的方向。

总之，HeLioS 通过创新的架构设计和系统调度策略，成功将 NMP 技术从静态批处理场景推向了高度动态的在线 LLM 服务场景，显著提升了服务效率和资源利用率。