点击蓝字

关注我们

因为公众号平台更改了推送规则。记得点下右下角的大拇指“赞”和红心“推荐”。这样每次新文章推送，就会第一时间出现在订阅号列表里。

为什么内存层次在AI加速器中至关重要

现代AI加速器的瓶颈不再在于计算能力，而在于数据传输。你可以设计出具备千万亿次浮点运算能力的GPU，也可以将数千个张量核心集成在单个芯片上。但如果这些核心无法快速获取数据，它们就会处于空闲状态。

闲置的计算资源就是浪费硅芯片。计算吞吐量与内存带宽之间的这种失衡，是现代AI硬件面临的根本性挑战。这通常被称为AI内存瓶颈，即FLOPS的增长速度超过了内存系统向其提供数据的能力。

下面来探讨一下为何存在层次结构。

FLOPS 与带宽失衡

FLOPS 用于衡量计算能力，带宽用于衡量数据传输能力。 

在现代AI GPU 中：

• 计算量呈指数级增长。 

• 内存延迟缓慢提升。 

• 带宽在结构上可扩展（通过HBM）。

  
  

如果计算量翻倍但内存带宽未增加，利用率就会下降。张量核心会因等待操作数而陷入停滞。这就是计算饥饿。

现代GPU的延迟等级

并非所有内存都有同样的性能。

• 寄存器以单字节纳秒级运行。 

• L1缓存稍慢。 

• L2缓存更慢。 

• HBM容量大得多，但比SRAM慢。

  
  

每一层的存在，是因为物理法则阻止了单一内存类型同时满足：

• 极速 

• 超大 

• 超高能效 

  
  

设计师将内存层层叠加，形成层级结构。

张量流式行为 

AI训练并非随机内存访问，而是结构化的张量流式处理。

• 权重以块的形式移动。 

• 激活值逐层流动。 

• 梯度向后传播。

  
  

这种结构化流程受益于：

• 数据局部性 

• 缓存内的复用 

• 从HBM进行高带宽预取

   

如果没有层次结构，每次张量获取都会直接访问大容量的片外内存，从而影响性能。

计算饥饿问题 

想象一家拥有世界级厨师的餐厅厨房，但只有一名服务员负责运送食材。

这就像一块没有层级结构的GPU。

内存层次结构通过以下方式解决这个问题：

• 将常用数据保留在计算设备附近 

• 利用大带宽资源进行批量传输 

• 尽量减少昂贵的长途数据传输 

  
  

层级结构并非是选择性的。这是总体结构上必需的。

AI内存层次结构概述

在深入之前，先从整体上来看一下整个结构。

AI加速器的内存并非扁平，而是呈金字塔结构分层排列。距离计算单元越近，内存速度越快；物理位置越远，内存容量越大，但延迟也越高。

以下是简化的层级模型：

顶部（最快、最小、最接近计算）

• 寄存器

• 共享内存

• L1 缓存

• L2 缓存

  
  

底部（最大容量、最高带宽、片外）

• HBM

  
  

  
  

每一层都有其特定的功能。

寄存器（最快、最小） 

寄存器直接位于张量核心或CUDA核心内部，用于存储正在计算的运算操作数。访问延迟极低，通常仅为几个时钟周期。

但容量呢？非常小。

我们讨论的是每核的千字节。寄存器仅用于直接的算术运算，仅此而已。

共享内存 

共享内存是片上SRAM，可使同一块中的线程高效地重复使用数据，从而显著减少对更高层级内存的冗余访问。

共享内存在GPU编程模型中对程序员是可见的。

L1 缓存 

L1 缓存会自动将最近访问的数据存储在计算单元附近。它提高了时间局部性，即如果数据最近被访问过，那么它很可能再次被访问。

延迟非常低。但是容量有限。

L2缓存 

L2缓存容量更大，且在流式多处理器之间共享。它作为片上SRAM与片外HBM之间的缓冲器。

它显著减少了HBM数据传输负载。

HBM（高带宽）

HBM位于芯片外但封装在芯片上，提供每秒数TB的带宽。其速度虽低于SRAM，但容量更大。

它是以下内容的批量存储层：

• 模型参数 

• 激活函数 

• 梯度

这种层次结构的存在是因为物理上无法构建100 GB的寄存器内存。权衡取舍定义了架构。

片上内存（基于SRAM的架构） 

所有高于HBM的内容均采用SRAM构建。

SRAM 是：

• 速度极快

• 低延时

• 每次访问都非常耗电 

• 硅片面积成本非常高 

  
  

这就是为什么SRAM容量有限。

张量核中的寄存器 

寄存器用于存储张量核内部矩阵乘法的运算数。在AI训练过程中，寄存器是算术运算前的最终暂存区域。

延迟：约1–2个时钟周期 

容量：极小

它们并非用于存储，而是用于执行。

共享内存块  

共享内存允许多个线程之间进行协作式的数据共享。

示例：

如果多个线程需要相同的权重块，可以避免多次从HBM中获取：

• 一次性将数据复制到共享内存 

• 所有线程均可复用 

  
  

这减少了：

• HBM带宽压力 

• 功耗 

• 延迟 

  
  

在许多AI内核中，共享内存是显式管理的。

L1缓存的作用 

L1缓存用于存储最近使用的指令和数据。 

其目标是：

• 临时局部性捕获 

• 减少对L2的重复访问 

  
  

延迟：极低 

大小：大于寄存器，小于L2

现代GPU中的L2缓存 

L2缓存作为计算单元之间的共享缓冲区。现代AI GPU通常配备多兆字节的L2缓存，有时可达数十兆字节。

• 减少HBM流量 

• 缓冲张量块 

• 提高带宽效率 

  
  

如果没有L2缓存，每次数据复用都会影响HBM，导致外置带宽被占满。

SRAM各级延迟差异 

寄存器< 共享内存< L1  <  L2  <  HBM

每次向下都会略微增加访问时间。但向上移动会大幅增加面积成本。

优化数据复用 

AI工作负载非常适合缓存层次结构，因为：

• 矩阵运算重用块 

• 卷积层重用核 

• Transformer attention层重用查询/键/值矩阵 

  
  

内存层次结构利用了这种重用。

没有它，即使HBM的多TB带宽也远远不够。

片外高带宽内存（HBM 层） 

现在来到层次结构的最底层。HBM 并非集成在芯片上的SRAM，而是堆叠式DRAM，位于封装内部。 

为什么是外部？ 

因为如果使用SRAM，要实现数百GB 的容量将变得经济上不可行。HBM 在大规模应用中同时解决了容量和带宽的问题。

• 3D堆叠DRAM芯片 

• TSV垂直互连 

• 1024位宽接口 

• 硅基插件 

  
  

该架构可实现：现代GPU的3–5 TB/s带宽

HBM 存储：

• 完整模型权重 

• 层激活值 

• 优化器状态 

• 梯度 

  
  

但HBM比SRAM慢。

所以数据向上流动：HBM→L2 →L1 →寄存器

这种布局可避免频繁昂贵的片外访问。HBM是大容量存储器，SRAM是执行内存。

AI训练过程中的数据流动 

现在，来看看当训练一个LLM时，实际会发生什么。暂且抛开理论，想象一下Transformer层的前向传播过程：数十亿个参数存储在HBM中，张量核心正等待进行矩阵乘法运算。那么数据是如何物理移动的呢？ 

基于NVIDIA Hopper架构的现代GPU实现了先进的内存层次结构，将寄存器、共享内存、L2缓存和HBM紧密集成，以实现最大化张量的吞吐量。

以下是数据在内存中的移动路径：

步骤1：存储在HBM中的参数  

所有模型权重通常有数十甚至数百GB，都存放在HBM中。这是主要的存储层。HBM提供巨大的带宽（3–5 TB/s），但仍然比片上SRAM慢。因此数据不会永久保留在寄存器中，必须分阶段传输。

步骤2：HBM → L2缓存 

当某一层开始执行时，所需的张量块会从HBM加载到L2缓存中。L2缓存作为大容量缓冲区发挥作用。

为什么不直接去L1呢？

由于L2在计算单元之间共享，有助于减少重复的HBM数据读取。当多流多处理器需要相同的数据时，L2可作为数据重用的中心。

步骤3：L2→L1缓存  

从L2中将较小的数据块加载到L1缓存。L1更靠近计算单元，专为降低延迟而优化。 

此阶段可提升：

• 时间局部性 

• 指令重用 

• 数据块效率 

  
  

步骤4：L1 → 寄存器 

现在进入执行阶段。寄存器存储张量核心直接使用的操作数，矩阵乘加运算就在此进行。此处的延迟极低，仅有几个时钟周期。

步骤5：执行计算  

张量核心执行：

• 矩阵乘法 

• 累加 

• 激活函数应用

  
  

所有操作数均使用寄存器。

步骤6：回写周期  

计算完成后：

• 部分和可能返回寄存器 

• 最终输出传回L1 

• 然后传至L2 

• 最终写回HBM 

  
  

此循环逐层重复进行。

张量流式循环

训练期间：

• 前向传播→ 存储激活值 

• 反向传播→ 计算梯度 

• 优化步骤→ 更新权重 

  
  

每个阶段都需要在不同层级之间进行大量数据传输。如果没有层次结构，每次读写操作都将直接访问HBM，导致带宽急剧下降。整个流程的设计目的正是为了尽量减少昂贵的长距离数据传输。

不同内存层级的延迟与带宽对比 

现在将各层级进行并列比较。

注意一些有趣的现象：HBM的延迟高于SRAM，但总体带宽却高出很多。

这就是为什么存在分级结构。 

如果都是寄存器：

• 延迟 = 极其惊人 

• 容量 = 不可能 

• 成本 = 灾难性

   

如果一切都是HBM：

• 容量 = 很好 

• 带宽 = 强大 

• 延迟 = 执行速度太慢 

  
  

因此，设计师会将两者结合起来。 

层次结构平衡：

• 速度 

• 尺寸 

• 能耗

• 成本

  
  

这不是设计上的选择，而是物理上的必然。

LLM中的记忆层次优化  

LLM会放大存储设计中的任何低效之处。让我们来看看如何针对LLM进行层次结构的优化。

参数流式传输  

在训练过程中，模型参数会逐层从HBM流式传输到缓存中。 

高效的调度确保：

• 最大缓存复用 

• 最小冗余读取 

• 可预测的内存访问模式 

  
  

这能保持HBM处于饱和状态，而不会过载。

梯度累积 

梯度在寄存器和共享内存中累积，然后再写回HBM。 

这可以减少：

• 写入放大 

• 能耗 

• 带宽压力 

  
  

不再重复写入部分梯度，而是先在本地进行合并。

激活检查点

为了减少内存占用，部分激活值会被重新计算而非存储。这是一种计算与内存的权衡取舍。

层次结构使得这种策略成为可能，因为：

• 重新计算的值可以保留在SRAM中 

• 仅将必要数据写入HBM

  
  

数据局部性优化 

AI内核的设计旨在：

• 平铺矩阵 

• 重复使用缓存块 

• 将热数据保留在L1/L2缓存中

   

目标很简单： 

减少对HBM的访问次数。尽管HBM拥有巨大的带宽，但每次访问仍比SRAM消耗更多的能量。

不同内存层级的能效

这里有一个被大多数人低估的事实：数据传输通常比算术运算消耗更多的能量。

是的。移动数据比复制数据更贵。

SRAM 每次访问的功耗 

SRAM（寄存器、缓存）：

• 极低延迟 

• 每次访问能耗极低 

• 面积成本极高 

  
  

访问寄存器极为高效。

HBM每比特功耗

HBM在每瓦带宽方面经过优化，但：

• 片外信号传输消耗更多能量 

• TSV虽然缩短了走线距离，但其成本仍高于片上访问 

  
  

因此，减少HBM数据流至关重要。大规模AI集群在数据在不同内存层级之间传输时，会消耗巨大的能量。

数据流主导功耗 

在大型AI训练任务中：

• 计算功耗 = 显著 

• 数据传输功耗 = 主要因素 

  
  

• 数据传输距离越远，成本越高。 

• 内存层次结构的存在旨在缩短这一距离。
  

   

• 寄存器→ 成本最低 

• HBM→ 成本较高 

  
  

如今，节能调度与计算优化同等重要。

AI内存层次的未来  

随着AI的发展，内存层次也在不断演进。 

封装内SRAM扩展

设计师正在增加：

• L2缓存大小 

• 片上缓冲区容量 

• 共享内存的可配置性 

  
  

增加SRAM可减少对HBM的依赖，但SRAM面积成本较高。因此，仍需进行权衡取舍。

混合键合 

混合键合可降低互连电阻，并提高芯片之间的密度。

这可能会实现：

• 内存与计算的更紧密集成 

• 片外访问延迟更低 

• 每瓦更高带宽 

  
  

HBM的发展在很大程度上依赖于封装技术的创新。

计算与内存集成 

未来的加速器可能会模糊以下界限： 

• 计算 

• 内存 

  
  

内存中计算架构可大幅减少数据传输。无需将张量移动到计算单元，计算逻辑可以更接近内存数组。

迈向以内存为中心的AI架构 

历史上，系统以计算为中心。如今，它们正逐渐转向以内存为中心。因为若在不扩展内存层次的情况下增加计算能力，就会重演“内存墙”问题，这正是AI内存瓶颈讨论中所详述的限制。

AI硬件的未来不会由能够构建多少FLOPS决定。而是由如何高效地传输数据决定。

内存层次结构已不再是背景基础设施，而是AI加速的核心设计挑战。

常见问题

为什么AI加速器使用多层内存而不是单一的大容量内存？

AI加速器采用分层内存架构，是因为单一类型的内存无法同时实现超低延迟、大容量、高带宽和高能效。寄存器和SRAM缓存可为活跃计算提供快速访问，而HBM则提供大容量、高带宽的存储空间。这种分层设计可避免计算资源饥饿，并优化整体性能。

AI GPU中SRAM与HBM有何区别？ 

SRAM（用于寄存器和缓存）速度极快、延迟低，但容量有限且占用硅片面积较大。HBM容量更大，提供高达3–5 TB/s的带宽，但相比SRAM延迟更高。SRAM用于处理即时执行的数据，而HBM则存储大量模型参数和激活值。

在AI模型训练过程中，数据是如何流动的？ 

在AI训练期间，数据通常从HBM流向L2缓存，再进入L1缓存，存储到寄存器中进行张量计算，最后返回HBM进行存储。这种分阶段的数据流动方式可最大限度减少昂贵的外存访问，并确保GPU的高利用率。

如果AI加速器的内存带宽不足会发生什么？ 

当内存带宽无法跟上计算吞吐量时，张量核心会因等待数据而卡住。这被称为内存瓶颈问题。带宽不足会显著降低FLOPS利用率，并大幅增加训练时间。

为什么HBM被放置在芯片外而非芯片内？ 

HBM被放置在芯片外，是因为将数百GB的SRAM直接集成到芯片上在物理上不现实，且在经济上不可行。HBM采用3D堆叠DRAM架构，在保持与GPU通过硅基中间层接近的同时，提供大容量和高带宽。

结论

AI加速器不仅仅是庞大的计算引擎；它们是精心设计的分层内存系统，旨在避免资源饥饿。寄存器执行运算，缓存进行缓冲，HBM则以高度协调的层级结构为系统提供数据。在AI训练过程中，数据会持续地从HBM流向L2，再流向L1，进入寄存器进行计算，最终返回HBM进行存储。

每一层的存在，是因为没有任何一种单一的存储技术能够同时实现超低延迟、大容量、极高带宽和完美的能效。存储层次通过将频繁访问的数据保留在计算附近，同时使用高带宽存储器进行大规模数据存储，来平衡这些权衡关系。

随着模型参数从数十亿到数万亿级别，层次结构的重要性变得更加关键。能效经济、延迟优化和带宽扩展最终汇聚为一个架构现实：AI的性能不仅受限于计算能力，更取决于内存的智能组织方式。在现代加速器中，内存层次结构是实现AI可扩展性的真正支柱。

原文链接：

https://dailytechinsights.com/ai-accelerator-memory-hierarchy-hbm-sram-cache/

< 长按识别二维码添加好友 >

加入上述群聊

长按并关注

带你走进万物存储、万物智能、

万物互联信息革命新时代

微信号：SSDFans