文章标题:弥合差距:张量算法的可达到数据移动和运算强度界限
作者/机构:Qijing Huang (NVIDIA), Po-An Tsai (NVIDIA), Joel S. Emer (MIT CSAIL / NVIDIA), Angshuman Parashar (NVIDIA)

A1 主要贡献

核心问题:随着技术发展,数据移动对应用的性能和能效影响日益显著,尤其是在张量代数算法中。然而,当前用于早期架构设计空间探索(DSE)的数据移动评估方法存在巨大缺陷。常用的“算法最小值”(相当于缓存的强制未命中)作为数据移动的下限,过于宽松,与实际硬件(如NVIDIA A100 GPU)上观察到的数据移动量相差数个数量级(见图1中的“Gap 0”),尤其是在靠近计算单元的存储层次上。这种不精确性可能误导架构师做出次优的设计决策。另一方面,精确的建模方法(如详尽的映射空间搜索)成本高昂且依赖于具体的硬件实现,不适用于早期、广泛的设计探索。

研究目标:本文旨在弥合上述粗略估计与精确建模之间的鸿沟,提出一种名为 Orojenesis 的方法,用于计算张量算法的数据移动界限。该方法旨在提供比算法最小值更紧凑、且独立于具体映射(tiling, parallelization, scheduling)和硬件实现的界限。其核心目标是生成一个“滑雪坡图”(ski-slope diagram),清晰地展示了片上缓冲区(如缓存或暂存器)容量与下一级存储(backing store)的最小可达到数据访问量之间的帕累托最优关系。

创新点与主要贡献
1. Orojenesis 方法论:提出了一种新颖的方法来计算张量算法的数据移动界限。该方法的核心是使用一个名为 Snowcat 的极简代理架构(单处理元件,两级存储),通过对其进行详尽的映射空间搜索,高效地推导出数据移动与缓冲区容量的关系曲线(滑雪坡图)。该界限考虑了数据复用,并且独立于具体映射,可推广至多级存储、并行架构等复杂设计。
2. “滑雪坡图”与架构洞察:生成的滑雪坡图(图1)能够帮助架构师回答关键设计问题:
* 在给定缓冲区容量下,可达到的最小访存量或最大运算强度是多少?(填补 Gap 0)
* 为达到算法最小访存量,需要多大的有效缓冲区?(揭示 Gap 1)
* 增加缓冲区容量能带来多大的收益?
3. 融合(Fusion)分析:将 Orojenesis 扩展到分析一系列张量操作的融合。通过定义“融合友好映射模板”(FFMT),该方法能够探索融合映射空间,并量化不同融合策略(如分段融合)在不同缓冲区容量下的数据移动减少效果。
4. 衍生模型与应用:基于 Orojenesis 界限,构建了“可达到运算强度(OI)平台图”和“可达到性能平台图”,这些模型将缓冲区容量作为输入,用于指导更高层次的设计决策,例如在固定芯片面积下,缓冲区与计算单元的最佳面积配比。
5. LLM 案例研究:将 Orojenesis 应用于分析 GPT-3-6.7b 大语言模型。研究揭示了:
* 对于拥有 320MB 缓冲区的情况,通过融合最多可将数据移动减少 5.6 倍。
* 在不同融合策略(如 FlashAttention 和 FLAT)之间进行量化比较。
* 确定了为支持融合 LLM 的最优硬件设计,相较于非融合设计,缓冲区面积需求降低 60%,而吞吐量提升 2.4 倍。

图1:可达到数据移动量与算法最小值之间的差异(Gap 0),以及最大有效缓冲区大小与总操作数大小之间的差异(Gap 1)。绿色曲线类似于滑雪斜坡,代表了张量操作在缓冲区大小需求方面的理想数据移动界限。
图1:可达到数据移动量与算法最小值之间的差异(Gap 0),以及最大有效缓冲区大小与总操作数大小之间的差异(Gap 1)。绿色曲线类似于滑雪斜坡,代表了张量操作在缓冲区大小需求方面的理想数据移动界限。

A3 动机

算法最小值度量的局限性:在硬件设计的早期阶段,架构师通常会进行“速度与馈送”研究,其中运算强度(OI)是关键指标,用于在屋顶线模型【76, Roofline: an insightful visual performance model for multicore architectures, Communications of the ACM, 2009】中快速评估工作负载是计算受限还是内存带宽受限。然而,计算OI所用的数据移动量通常是“算法最小值”,即所有操作数大小的总和。这个度量过于乐观。例如,图2显示,对于一个 4k x 4k x 4k 的通用矩阵乘法(GEMM)任务,在NVIDIA A100 GPU上,实际的DRAM流量是算法最小值的6.5倍。L2到L1缓存的流量差距更大,达到了32.3倍。这种巨大的差异源于映射效率低下和硬件设计约束(特别是用于数据复用的缓冲区容量)。使用一个有严重缺陷的数据移动度量可能会误导架构师做出糟糕的初始设计决策。

图2:在NVIDIA A100 GPU上,一个4k x 4k x 4k GEMM的内存和缓存访问量与算法最小值的比较。
图2:在NVIDIA A100 GPU上,一个4k x 4k x 4k GEMM的内存和缓存访问量与算法最小值的比较。

图3:实现完全张量复用的最大有效缓冲区大小,已归一化到总张量大小。
图3:实现完全张量复用的最大有效缓冲区大小,已归一化到总张量大小。

精确数据移动评估的挑战:架构师之所以不能在早期DSE中使用更精确的数据移动估算方法,原因有二。首先,数据移动对架构内存层次结构所能利用的复用程度敏感,精确建模需要更详细的架构模型,这带来了更高的实现和运行时成本。其次,数据移动对算法的具体实现方式(即“映射”)敏感。一个张量算法可以有多种映射方式(tiling、并行化、调度、融合等),这些选择由专家程序员或优化编译器做出,以最优地利用硬件资源和算法固有的复用模式。因此,我们不能简单地使用对缓存容量敏感但只对单一映射建模的Belady算法【7, A study of replacement algorithms for a virtual-storage computer, IBM Systems journal, 1966】来确定最优流量。即使对于单一设计,进行详尽的映射空间搜索也可能耗时过长,不适用于早期DSE。

本文方法的定位:本文旨在弥合算法最小值的极度不精确性与精确模型的巨大成本和实现依赖性之间的鸿沟。我们的方法:
* 解决映射特定性:提供一个任何映射都无法超越的数据移动界限。
* 解决复用无关性:为任意给定的缓冲区容量提供一个后备存储访问界限,从而可以预测任何设计内存层次结构中所有级别的数据移动界限。
* 解决运行时成本:采用一个我们称为 Snowcat 的简单代理架构(图4),它暴露了一个极小的映射空间,并且可以被极其高效地分析或建模。

该工具可在 https://timeloop.csail.mit.edu/orojenesis 获取。

图4:Snowcat架构与真实设计的比较。
图4:Snowcat架构与真实设计的比较。

A2 方法细节

Orojenesis

术语定义

本文使用源自 TeAAL 【47, Teaal: A declarative framework for modeling sparse tensor accelerators, MICRO, 2023】的术语。
* 张量(Tensor):一个具有固定数量秩(rank)的多维数组,每个秩都有一个形状(shape)。
* 爱因斯坦求和约定(Einsum):一种简洁表示张量运算(如矩阵乘法、卷积)的符号【20, The foundation of the general theory of relativity, Annalen Phys, 1916】。例如,矩阵乘法表示为 $B_{m,n}^{M,N} = \sum_{k=0}^{K-1} A_{m,k}^{M,K} W_{k,n}^{K,N}$。
* 张量算法(Tensor Algorithm):可以表示为单个Einsum或一系列生产者-消费者级联的Einsum的计算。张量算法总是未映射的。
* 映射(Mapping):在目标架构上对张量算法进行分块(tile)、并行化(parallelize)、调度(schedule)和/或融合(fuse)的特定方式。所有合法映射的集合称为映射空间(Mapspace)
* 映射器(Mapper):一种在给定优化目标和超参数的情况下,在映射空间中寻找最优映射的算法或启发式方法。

Orojenesis 方法论

Orojenesis 是一种推导缓冲区容量与对上一级内存(后备存储)访问次数下限之间关系的方法,这个下限是给定张量算法的任何映射都无法超越的。

核心洞察与 Snowcat 架构:本工作利用的核心洞察是,任何两个连续内存级别之间的数据移动行为在极限情况下是分形的。对于内存层次结构中的任何给定级别,我们可以将其总存储容量视为一个集体的“字节池”,可实现的最大数据复用和到下一级的流量仅受此池中字节数的限制。因此,我们可以使用一个名为 Snowcat 的简单代理架构来研究可推广到复杂架构的数据移动行为。如图4所示,Snowcat 是一个单处理元件架构,只有两级存储:一个无约束的缓冲区和一个后备存储。由于只有两级存储且无需考虑并行化,其映射空间大大减小,建模复杂度也更低。例如,一个 4k x 4k x 4k GEMM 的 Snowcat 映射空间比类 Eyeriss 架构【12, Eyeriss: An energy-efficient reconfigurable accelerator for deep convolutional neural networks, IEEE journal of solid-state circuits, 2016】小7350倍。

图5:Orojenesis 流程。
图5:Orojenesis 流程。

界限推导:虽然静态编译时分析和基于优化的方法【31, 54, 55】可用于寻找数据移动界限,但它们未分析融合映射空间且界限存在差距。启发式或数据驱动的映射空间搜索方法【27, 36, 77】不保证收敛到全局最优。由于 Snowcat 架构的可管理建模成本和紧凑搜索空间,详尽搜索成为找到所有帕累托最优点(优化缓冲区大小和数据访问计数)的可行方法。

工具流程:图5展示了 Orojenesis 的整体流程。它接受以单个或链式 Einsum 形式指定的工作负载。对于给定的工作负载,流程遍历其在 Snowcat 架构上的完整无约束映射空间。对于遍历中遇到的每个映射,我们计算后备存储访问次数和每个张量的分块大小(即活跃数据足迹)。基于此信息,缓冲区的大小被精确地调整以适应映射所需的分块大小,我们称之为映射的缓冲区大小需求。在整个映射空间遍历过程中,Orojenesis 收集所有映射的缓冲区大小需求和后备存储访问次数,并不断更新不同缓冲区大小需求下实现的最佳后备存储访问次数。最终,连接该空间中的帕累托最优点,便形成了滑雪坡图。

实现细节:我们使用配置为详尽搜索的 Timeloop 映射器【57, Timeloop: A systematic approach to dnn accelerator evaluation, ISPASS, 2019】及其性能模型来实现 Orojenesis 流程。对于多 Einsum 评估,我们调整了 Timeloop 以支持融合优化。图6以GEMM为例,说明了如何使用简单的代数表达式推导分块大小和访问次数。每个张量的分块大小是其相关秩的内循环边界的乘积,总缓冲区需求是它们的和。后备存储访问次数通过将内循环分块大小与外循环迭代次数相乘计算得出。

图6:从矩阵乘法映射中推导出的缓冲区大小需求和到下一内存级别的访问量。
图6:从矩阵乘法映射中推导出的缓冲区大小需求和到下一内存级别的访问量。

Orojenesis 界限的外推:基于 Snowcat 的 Orojenesis 分析适用于多种架构分析:
1. 多级内存层次结构:如图7所示,可以在滑雪坡图的不同点上探测算法的 Orojenesis 界限,以找到内存层次结构中任意两个级别(如 L1 和 L2,L2 和 DRAM)之间的数据移动界限。注意,跨多个级别组合帕累托最优映射可能不总成立,因此多级后备存储查找产生的是一个下限,但不保证是紧凑的。
2. 并行架构:并行化可能引入数据复制,导致有效缓冲区容量减少或数据移动流量增加,从而产生一个相对于 Orojenesis 的 Snowcat 派生帕累托界限的次优点。因此,当存在并行性时,Orojenesis 返回一个较宽松但仍然正确的下限。
3. 受约束的映射空间:某些架构可能会约束合法映射的空间以简化设计。这种约束会缩小映射空间,但 Orojenesis 界限仍然保证是该空间的界限。
由于 Orojenesis 界限在这些架构属性中仍然有效,因此其界限可移植到所有可以用 Timeloop 等框架描述的张量加速器架构,无需为不同架构重新运行。

图7:16k x 1k x 1k GEMM 的多级内存访问界限,点表示已探索的映射。
图7:16k x 1k x 1k GEMM 的多级内存访问界限,点表示已探索的映射。

衍生模型

可达到运算强度(OI)模型:与源于算法固有属性的算法 OI 不同,工作负载的可达到 OI 依赖于映射空间并受硬件缓冲区容量的约束。Orojenesis 数据可用于推导如图8所示的最大可达到 OI 曲线。我们称之为 OI 平台图(OI mesa)。图中 OI 要么受缓冲区容量限制,要么受算法固有的计算与张量大小比率限制。

图8:最大可达到运算强度的 OI 平台图。
图8:最大可达到运算强度的 OI 平台图。

可达到性能模型:可达到 OI 模型可与传统的屋顶线模型【76】结合,形成一个将缓冲区容量、内存带宽和计算能力作为输入的性能模型。结果是一个新的、与映射无关的性能模型,用于指导 DSE。该模型的输出称为性能平台图(performance mesa),如图9所示。

图9:以缓冲区大小为输入的性能平台图。
图9:以缓冲区大小为输入的性能平台图。

单 Einsum 界限分析

  1. 矩阵乘法(GEMM):其 Einsum 为 $B_{m,n}^{M,N} = A_{m,k}^{M,K} W_{k,n}^{K,N}$。图10显示了不同形状 GEMM 的滑雪坡图和 OI 平台图。图10a 显示,所有 GEMM 的后备存储访问界限随缓冲区容量增加呈幂律下降,且较大的 GEMM 从增加缓冲区容量中获益更多。图10b 显示,可达到 OI 也呈幂律增加,直到达到平台顶部(峰值 OI)。峰值 OI 受限于 GEMM 的最小维度。图11显示,GEMM 的最大有效缓冲区大小约等于其最小操作数的大小。

    图10:GEMM 尺寸的影响。随着缓冲区容量的增加,较大的 GEMM 在总数据移动量上的减少更为显著。注意Y轴使用线性刻度。
    图10:GEMM 尺寸的影响。随着缓冲区容量的增加,较大的 GEMM 在总数据移动量上的减少更为显著。注意Y轴使用线性刻度。

    图11:最大有效缓冲区比率与总操作数大小的比值。
    图11:最大有效缓冲区比率与总操作数大小的比值。

  2. 卷积:其 Einsum 为 $B_{p,q,n}^{P,Q,N} = \sum_{c,r,s} A_{p+dr, q+ds, c}^{P+DR, Q+DS, C} W_{c,n,r,s}^{C,N,R,S}$。图12显示,较大的滤波器尺寸(R, S)导致更多的后备存储访问和更高的峰值 OI,并且从增加缓冲区容量中获益更多。步长(stride)和空洞(dilation)会略微增加后备存储访问。

    图12:各种卷积配置的影响。
    图12:各种卷积配置的影响。

  3. 批量矩阵乘法(BMM):广泛用于 Transformer 模型的多头注意力(MHA)【74】。其 Einsum 为 $B_{h,m,n}^{H,M,N} = A_{h,m,k}^{H,M,K} W_{h,k,n}^{H,K,N}$,H 为头数。图13显示,在总计算量固定的情况下,头数越多,后备存储访问越高,增加缓冲区容量的收益递减,且最大有效缓冲区大小和峰值 OI 均下降。这表明对于头维度较小的 BMM,增加缓冲区大小或计算资源无法提升性能,唯一的方法是增加内存带宽或头特征维度。

    图13:在总计算量固定为128 GOPs的情况下,BMM中头数的影响,维度为H=头数,M=4k,K=4k/头数,N=4k。
    图13:在总计算量固定为128 GOPs的情况下,BMM中头数的影响,维度为H=头数,M=4k,K=4k/头数,N=4k。

  4. 分组 BMM:用于多查询注意力(MQA)【67】和分组查询注意力(GQA)【4】以降低 MHA 的内存访问成本。其 Einsum 为 $B_{h,m,n}^{H,M,N} = A_{h,m,k}^{H,M,K} W_{g,k,n}^{G,K,N}$,其中 G 为组数,一个操作数的头由另一个操作数的 H 个头共享。图14显示,减少组数可以降低数据移动并提高 OI,但当缓冲区容量大于10MB时,MQA 和 GQA 带来的数据访问节省效果减弱。

    图14:在维度 H=32, M=4k, K=128, N=4k 的分组BMM中,第二个输入操作数不同组数的影响。
    图14:在维度 H=32, M=4k, K=128, N=4k 的分组BMM中,第二个输入操作数不同组数的影响。

Orojenesis 融合

当 Orojenesis 的输入是一系列 Einsum 时,由于存在融合机会,后备存储的访问界限不能简单地通过各个 Einsum 界限求和得到。

融合方法论

融合的定义:如图15所示,层融合基于两个关键标准:首先,序列中一个 Einsum 的输出必须是下一个 Einsum 的输入。其次,在内存层次结构的某个级别 m,前一个 Einsum 的输出(中间张量)应被后一个 Einsum 完全消耗,而不会溢出到更外层的内存级别。

图15:层融合定义。Einsum e的输出是e+1的输入。在内存层次的m级,e的输出被e+1完全消耗。
图15:层融合定义。Einsum e的输出是e+1的输入。在内存层次的m级,e的输出被e+1完全消耗。

非分块融合(Untiled Fusion):如果缓冲区能容纳完整的中间张量,融合不对单个 Einsum 的映射施加约束。然而,这种方式通常导致很高的缓冲区大小需求,往往是次优的。

分块融合(Tiled Fusion):为了更有效地利用缓冲区,对中间输出进行分块至关重要。这要求链中每个 Einsum 的执行顺序和粒度保持一致。我们通过强制所有映射遵循一个称为融合友好映射模板(Fusion Friendly Mapping Template, FFMT)的约束集来实现。图16展示了用于本文所有融合分析的 GEMM FFMT。

图16:GEMM融合友好映射模板(FFMT)。GEMM FFMT是四种不同约束集的并集。灰色框中最外层的循环在两个Einsum链中是可排列的。
(a) 完整(Full):只有M维度被分块,以完全消耗一个输入行并产生最终和的完整输出行。
(b) K分块(TiledK):M和K维度被分块。消耗输入行的子分区并产生部分和。
(c) N分块(TiledN):M和N维度被分块。完全消耗输入行并产生最终和输出的子分区。
(d) KN分块(TiledKN):所有维度都被分块,导致消耗输入行的子分区和产生部分和输出的子分区。
图16:GEMM融合友好映射模板(FFMT)。GEMM FFMT是四种不同约束集的并集。灰色框中最外层的循环在两个Einsum链中是可排列的。 (a) 完整(Full):只有M维度被分块,以完全消耗一个输入行并产生最终和的完整输出行。 (b) K分块(TiledK):M和K维度被分块。消耗输入行的子分区并产生部分和。 (c) N分块(TiledN):M和N维度被分块。完全消耗输入行并产生最终和输出的子分区。 (d) KN分块(TiledKN):所有维度都被分块,导致消耗输入行的子分区和产生部分和输出的子分区。

FFMT 约束:如图17所示,FFMT 映射约束因 Einsum 在融合链中的位置而异。第一个 Einsum (Einsum 0) 可以应用最宽松的 FFMT-TiledKN 模板。然而,为了避免传播部分和,后续的 Einsum 直到最后一个,都必须遵循更严格的 FFMT-Full 约束,即消耗并产生完整的 M0 行。最后一个 Einsum 可以使用 FFMT-TiledN 模板。对于两个 Einsum 的链,这是一个特殊情况,可以对两者都应用 FFMT-TiledKN,并探索不同的循环排序。

图17:为融合链中不同Einsum施加的GEMM FFMT约束。每个回环箭头表示每个内部缓冲区分块执行的次数。‘K2(0)’和‘N2(0)’表示Einsum 0的K2和N2循环边界。‘N2(E-1)’代表Einsum E-1的N2循环边界。
图17:为融合链中不同Einsum施加的GEMM FFMT约束。每个回环箭头表示每个内部缓冲区分块执行的次数。‘K2(0)’和‘N2(0)’表示Einsum 0的K2和N2循环边界。‘N2(E-1)’代表Einsum E-1的N2循环边界。

缓冲区大小需求

对于顺序融合,整个链的缓冲区大小需求由所有 Einsum 中最大的缓冲区需求决定:$BufReq = \max_e(BufReq_e)$。对于流水线融合,缓冲区需求是所有权重张量大小之和加上所有 Einsum 中最大的输入和输出分块大小之和:$BufReq = (\sum_e BufReq_{W,e}) + \max_e(BufReq_{I,e} + BufReq_{O,e})$。本文主要关注顺序融合,因为它通常更优。

后备存储访问计数

总后备存储访问计数为:$Access = Access_W + Access_{I,0} + Access_{O,E-1}$,即所有 Einsum 的总权重访问量,加上 Einsum 0 的输入访问量,以及 Einsum E-1 的输出访问量。

映射权衡

映射空间中的一个关键权衡是 M0 的大小和选择哪些层的权重被完全缓冲在片上。增加 M0 会增加缓冲区需求,但会减少整个链需要重新遍历的次数,从而减少权重重载。我们遍历这个融合映射空间以产生 Orojenesis 界限。

避免部分和传播

当前分析仅关注从生产者到消费者 Einsum 传输最终和,以避免因传输部分和而导致的重计算或额外的内存访问。

多 Einsum 界限分析

融合与非融合的比较:图18展示了融合两个 GEMM(32k x 4k x 16k 和 32k x 16k x 4k)的 Orojenesis 界限。紫色曲线代表非融合执行的基线,通过无约束地搜索每个 Einsum 的最优映射并将它们的最佳数据访问量相加得到。蓝色曲线是非分块融合,由于需要缓冲巨大的中间张量,效率低下。绿色曲线是分块融合,它在更小的缓冲区容量下就能有效减少后备存储访问。

关键洞察
* 融合并非总是最优:由于 FFMT 施加的额外层内约束,当缓冲区大小小于 10MB 时,分块融合的性能可能不如非融合映射。
* 分块融合的有效性:当缓冲区大小超过 10MB 时,分块融合开始显示其优势。当缓冲区大于 256MB 时,融合可带来高达 3.7 倍的访问次数减少。
* 缓冲区对融合的增益:从曲线的斜率可以看出,增加缓冲区容量对融合映射的益处大于非融合映射。

图18:融合32k x 4k x 16k和32k x 16k x 4k的GEMMs。
图18:融合32k x 4k x 16k和32k x 16k x 4k的GEMMs。

A7 案例研究:大语言模型

本案例研究利用 LLM 中层的融合,建立数据移动界限并指导加速器设计。目标工作负载是 GPT-3-6.7b,其构建模块(图19)在多头注意力(MHA)和前馈网络(FFN)模块中包含一系列 Einsum。

图19:LLM 构建模块架构。
图19:LLM 构建模块架构。

MHA 融合策略比较

融合 MHA 中的 BMM 对性能至关重要。我们将 FLAT【39】和 FlashAttention【15, 16】的平铺策略用 FFMT 建模并进行比较。FLAT 需要生成完整的输出行以进行 Softmax 的行式规约,而 FlashAttention 利用代数变换允许对输出行进行平铺。如图20所示,在缓冲区大小为 16MB 时,FlashAttention(绿色曲线)的后备存储访问量比 FLAT(蓝色曲线)低 6 倍以上。当缓冲区容量达到 32MB 的最大有效尺寸时,两者策略变得同样有效。

图20:MHA融合策略类似于FlashAttention [16]和FLAT [39]。
图20:MHA融合策略类似于FlashAttention [16]和FLAT [39]。

链分段的影响

图21展示了 GPT-3-6.7b 模块中一个六 Einsum 链的数据移动界限。绿色曲线代表对整个链进行融合。我们观察到,融合一个长链并不总是最优的,因为中间的 Einsum 会受到更严格的 FFMT 约束,增加了整体的缓冲区需求。为了缓解这个问题,我们探索了所有可能的分段方式,并用黄色曲线表示在不同缓冲区大小约束下的最优分段策略。结果表明,将融合链划分为较短的段可以显著减少数据访问,特别是在缓冲区较小的情况下。

图21:分段对GPT-3-6.7b中一个六Einsum链的影响。
图21:分段对GPT-3-6.7b中一个六Einsum链的影响。

最优完整 LLM 融合策略

图22显示了 GPT-3-6.7b 构建模块中所有 Einsum 顺序执行的总后备存储访问需求。分析表明,对于一个 50MB 的末级缓存,最优的融合执行相比最优的非融合映射,可以将总后备存储流量减少 2.5 倍(绝对减少 4.7 GB)。当片上缓冲区大于 320MB(最大有效缓冲区大小)时,融合的全部潜力得以释放,可实现高达 5.6 倍的 DRAM 访问减少(绝对减少 6 GB)。

图22:整个GPT-3-6.7b构建块的Orojenesis界限。
图22:整个GPT-3-6.7b构建块的Orojenesis界限。

缓冲区与计算面积配比的界限

在 DSE 中一个常见的挑战是在固定的总芯片面积下,确定缓冲区与 MAC 单元的最佳面积比。我们使用性能平台图模型来解决这个问题。以一个类似 GF100【48】的芯片规格为基准,在 40nm 工艺下,die 面积为 529 mm²,DRAM 带宽为 149 GB/s。我们调整片上缓冲区大小和 MAC 单元数量,以优化支持 GPT-3-6.7b 工作负载的硬件。如图23所示,我们绘制了芯片吞吐量与缓冲区面积占比的关系。实际可达到的性能(不透明绿色曲线)是内存受限性能(橙色线)和计算受限性能(蓝色线)的最小值。曲线显示吞吐量是缓冲区面积比的凹函数,两线交点处达到峰值性能。与非融合 LLM 相比,为融合 LLM 设计的最优硬件需要低 60% 的缓冲区面积,同时实现 2.4 倍的整体吞吐量提升。

图23:山状性能模型。垂直虚线表示在40nm技术、529 um2芯片面积的芯片中实现峰值吞吐量性能的最佳缓冲区域与总芯片面积比。
图23:山状性能模型。垂直虚线表示在40nm技术、529 um2芯片面积的芯片中实现峰值吞吐量性能的最佳缓冲区域与总芯片面积比。

A4 实验环境与结果

实验环境

  • 工作负载
    • 通用算法:多种形状的 GEMM(如 4kx4kx4k, 16kx1kx1k)、卷积、批量矩阵乘法(BMM)、分组 BMM。
    • LLM 模型:GPT-3-6.7b,特征维度(d)=4096,注意力头数(h)=32,头维度(f)=128,隐藏特征维度(c)=16384,输入序列长度(l)=32768 (批量大小16 x 序列长度2048)。
  • 硬件配置
    • 验证平台:NVIDIA GPUs (A2-2MB L2, A30-24MB L2, A100-40MB L2, H100-50MB L2);Simba 加速器【66, Simba: Scaling deep-learning inference with multi-chip-module-based architecture, MICRO, 2019】的分析模型。
    • DSE 案例研究平台:基于 GF100【48】芯片规格(40nm工艺, 529 mm² 芯片面积, 700 MHz 频率, 149 GB/s DRAM 带宽)。
  • 软件配置
    • 核心工具:Orojenesis 流程使用 Timeloop【57】的映射器和性能模型实现,并为融合优化进行了适配。
    • 验证软件:在 NVIDIA GPU 上使用 CUTLASS【72】库运行 GEMM。
    • 面积估算:使用 Accelergy【78, Accelergy: An architecture-level energy estimation methodology for accelerator designs, ICCAD, 2019】进行 MAC 和 SRAM 的面积计算。

实验结果

界限验证

  • GPU 和 Simba 验证 (图 24):我们验证了 Orojenesis 对 4k x 4k x 4k GEMM 的界限。
    • 图24a 显示,在四款 NVIDIA GPU 上运行 CUTLASS 测得的 DRAM 访问量均位于 Orojenesis 界限之上,表明界限是有效的。针对 A100 和 H100 的优化调度非常接近最优界限。
    • 图24b 显示,在 Simba 加速器模型上,不同全局缓冲区配置下的所有映射点均位于滑雪坡曲线之上,验证了其对空间加速器的有效性。
    • 图24c 进一步验证了多 Einsum 链的界限。在 Simba 上,无论是融合还是非融合执行,测得的数据访问量都始终高于相应的 Orojenesis 界限。

图24:测量的DRAM访问量与Orojenesis的对比。
图24:测量的DRAM访问量与Orojenesis的对比。

运行时比较

  • 与映射感知 DSE 的比较 (表 I):为了展示 Orojenesis 的运行时优势,我们将其与对 Simba 架构进行的完整映射感知 DSE 进行了比较。
    • 在单个 4 核 Intel i7 处理器上,评估 Simba 上的一个映射需要 3.90ms,而评估 Orojenesis 的 Snowcat 架构上的一个样本仅需 0.20ms,快了 19.5 倍。
    • 对 Simba 的 100 种不同全局缓冲区容量配置进行 DSE,共需评估 260 万个映射。而 Orojenesis 的一次详尽搜索评估了 9 万个有效映射。
    • 总体而言,Orojenesis 的运行时间比对 100 种 Simba 配置的 DSE 快 556 倍
    • 更重要的是,Orojenesis 的运行结果可移植到大量张量架构,而 DSE 的结果仅限于特定架构的特定搜索空间。对于更广泛的 DSE,Orojenesis 的速度优势会更加显著。

TABLE I: Orojenesis与Simba 100设计DSE的比较。

表1:Orojenesis 与 Simba 100 设计 DSE 的比较。
表1:Orojenesis 与 Simba 100 设计 DSE 的比较。

A5 结论

本文提出了 Orojenesis,一种计算张量算法数据移动界限的方法。Orojenesis 提供的界限比传统的算法最小值界限更紧凑,因为它考虑了缓冲区可利用的数据复用。同时,它避免了传统基于缓存研究的映射特定性以及现代张量加速器框架中难以处理的映射空间搜索问题。Orojenesis 的界限是映射无关的,为任何算法映射(包括利用生产者-消费者融合的映射)在架构层次结构中可能提取的数据移动提供了一个极限。本文进一步展示了如何使用 Orojenesis 的输出来获得广泛的架构见解,并以一个当代大语言模型的综合案例研究作为总结。

凭借其提供的紧凑界限,Orojenesis 可以帮助架构师设计出良好平衡的系统,避免因围绕算法最小流量设计时保护带不足而导致的计算资源利用率低下的陷阱。

未来的工作包括进一步收紧 Orojenesis 的界限,以考虑并行性、延迟隐藏、流水线、部分和传播以及超越平铺光栅遍历的张量遍历模式。

A6 附录

A. 工件附录

摘要:本节描述如何访问 Orojenesis 的工件(artifacts)以及如何复现论文中的关键图表。此工件可在任何装有 x86-64 CPU 的机器上运行。Orojenesis 流程已集成到 Timeloop 中。

工件清单(元信息)
* 算法:张量算法的分析性分析。
* 程序:C++, Python。
* 运行环境:x86-64 机器。需要 root 或 docker 访问权限。
* 输出:论文中的关键图表。
* 实验:各种张量的界限。
* 所需磁盘空间:20 GB。
* 准备工作流程所需时间:30 分钟。
* 完成实验所需时间:12 小时。
* 是否公开可用:是。
* 代码许可证:BSD-3-Clause 许可证。
* 存档地址https://zenodo.org/doi/10.5281/zenodo.10850531

描述
* 如何访问:建议通过存档库 (https://zenodo.org/doi/10.5281/zenodo.10850531) 访问工件,其中包含最新的更新和错误修复。

安装
1. 下载 Orojenesis 工件:在用户机器上运行提供的命令下载并解压 orojenesis.zip
2. 安装 Orojenesisorojenesis 仓库下提供了一个安装脚本 install.sh 来安装 Timeloop 和其他软件依赖。
* 原生主机:如果你在基于 Debian 的系统上拥有 sudo 权限且已安装 Python 3.10 或更高版本,建议直接执行安装脚本。
* Docker:或者,如果没有 sudo 权限,可使用 Docker 容器。Dockerfile 位于 orojenesis/docker/Dockerfile

实验工作流程
* 我们提供了 Jupyter notebooks (orojenesis_single.ipynborojenesis_multi.ipynb) 来指导生成论文中的关键图表。
* 在 orojenesis 目录下启动 Jupyter notebook,并按照终端输出中的说明在 Web 浏览器中导航到 Jupyter 界面。
* 如果无法访问 GUI,可以使用 jupyter nbconvert 命令将 notebooks 转换为 Python 脚本运行。

评估和预期结果
* 单 Einsum:运行 orojenesis_single.ipynb 会生成论文中的图1, 10, 11, 12, 13, 14 和 24b。
* 多 Einsum:运行 orojenesis_multi.ipynb 会生成图18, 20, 21, 22 和 23。

实验定制
* 要自定义输入工作负载形状和约束,请参考 Jupyter notebook orojenesis/orojenesis_example.ipynb 中的说明和示例。