LADDER: Enabling Efficient Low-Precision Deep Learning Computing through Hardware-aware Tensor Transformation

文章标题：LADDER: 通过硬件感知的张量变换实现高效的低精度深度学习计算
作者/机构：Lei Wang†⋄∗, Lingxiao Ma⋄, Shijie Cao⋄, Quanlu Zhang⋄, Jilong Xue⋄, Yining Shi‡⋄∗, Ningxin Zheng⋄, Ziming Miao⋄, Fan Yang⋄, Ting Cao⋄, Yuqing Yang⋄, Mao Yang⋄ (中国科学院大学†, 北京大学‡, 微软研究院⋄)

A1 主要贡献

核心问题：
随着深度学习模型对性能需求的不断增长，利用低精度计算已成为一种趋势。然而，现有的硬件和软件对新兴的自定义低精度数据类型的支持不足且效率低下。硬件（如GPU）只能集成有限的几种标准数据类型计算单元，难以跟上算法需求的快速演进。软件方面，优化细粒度的低比特数据访问以适应粗粒度的内存系统（如NVIDIA GPU的4字节共享内存库宽）非常复杂，容易导致带宽浪费，需要大量非平凡的优化工作。这种软硬件支持的不足和低效严重阻碍了模型和加速器的创新。

研究目标：
本文旨在弥合不断演进的自定义数据类型与当前硬件支持的固定精度格式之间的鸿沟。目标是在不修改现有硬件的条件下，通过一种新颖的编译器设计，系统性地支持通用低比特精度自定义数据类型，并在现代加速器上实现高效的深度学习计算。

创新点：
1. 分离存储与计算的通用方法：本文观察到，硬件的内存系统可以存储任意数据类型，而计算单元支持的标准数据类型可以无损地表示大多数自定义数据类型。基于此，LADDER提出了一种通用方法：以自定义低精度数据类型存储和传输张量，通过类型转换以硬件支持的标准数据类型进行计算。这能够有效利用低比特数据类型在节省内存流量和占用方面的优势。
2. tType通用类型系统和扩展张量表达式：为了方便地实现快速演进的自定义数据类型（如块状数据类型MXFP），LADDER引入了一个名为tType的通用类型系统。tType是一个分块（tile-wise）的数据类型，通过显式指定类型宽度、元素形状和类型转换函数来定义所有常见的自定义类型。基于此，LADDER扩展了现有的张量表达式，使其能够原生支持为每个张量标注tType，从而系统地将带有自定义数据类型的DNN计算转换为标准的计算流水线。
3. 张量调度原语：为了优化涉及自定义数据存储、访问和类型转换的计算流水线，LADDER引入了一套新的张量调度原语，包括slice（切片）、map（映射）、pad（填充）和convert（转换）。这些原语能够将默认的计算流水线转换为经过优化的、性能更佳的等价形式，以适应不同的内存布局和硬件特性。
4. 硬件感知的优化策略：为了在复杂的变换空间中找到最优解，LADDER采用了一种分层的、硬件感知的优化策略。该策略将DNN计算建模为分块级别的数据流图，然后使用一种粒度感知的调度策略进行优化：较低层内存提供首选的数据访问粒度作为提示（hint），上层根据此提示决定最优的计算粒度，并通过张量变换进行对齐。
5. 系统实现与开源：LADDER基于TVM、Roller和Welder实现，并且已经开源。其核心DNN操作编译功能已作为BitBLAS库发布，可集成到现有的DNN和LLM框架中，为深度学习生态系统提供高效的低精度计算能力。

A3 背景知识/关键Observation/设计原则

2.1 深度学习中的精度要求

• 低比特数值精度：深度学习模型正从传统的FP32精度转向更低的精度以提升效率。例如，自动混合精度（AMP）训练中使用了FP16/BF16【索引36，Mixed precision training，2018，International Conference on Learning Representations】，而Transformer Engine【索引37，Fp8 formats for deep learning，2022，arXiv】和MS-AMP【索引39，Fp8-lm: Training fp8 large language models，2023，arXiv】等系统甚至开始使用FP8。在推理阶段，模型量化更为激进，通常降至8位或4位【索引14，Llm. int8 (): 8-bit matrix multiplication for transformers at scale，2022，arXiv；索引22，Gptq: Accurate post-training quantization for generative pre-trained transformers，2022，arXiv；索引48，Smoothquant: Accurate and efficient post-training quantization for large language models，2023，International Conference on Machine Learning】，甚至有研究探索2位乃至1位【索引12，Quip: 2-bit quantization of large language models with guarantees，2023，arXiv；索引46，Bitnet: Scaling 1-bit transformers for large language models，2023，arXiv】。图1展示了这种从高精度向低比特格式转变的趋势。
  
• 组级别的精度缩放：为了提升低精度模型的准确性和鲁棒性，一种常见方法是使用缩放因子。相比于传统的张量级或通道级缩放，组级别（group-wise）缩放因其更细的粒度能更好地捕捉子张量或组的分布，从而提升性能。例如，在训练后量化（PTQ）中，通常偏好使用大小为128或64的组，并用FP16进行缩放【索引22，Gptq: Accurate post-training quantization for generative pre-trained transformers，2022，arXiv】。在OCP-MXFP中，一个8位的共享缩放因子被应用于一个32元素的组【索引41，Microscaling data formats for deep learning，2023，arXiv】。
• 混合精度操作：由于不同张量对低比特量化的敏感度不同，混合精度操作应运而生。例如，混合精度训练结合了FP32、FP16和FP8等不同精度的张量。在LLM量化中，权重可以被量化到更低的位数，而激活值由于量化挑战更大，则需要更高的位数。这导致了如W4A16（4位权重，16位激活）、W2A16、W1A8等混合精度操作【索引12，Quip: 2-bit quantization of large language models with guarantees，2023，arXiv；索引22，Gptq: Accurate post-training quantization for generative pre-trained transformers，2022，arXiv；索引46，Bitnet: Scaling 1-bit transformers for large language models，2023，arXiv】。

2.2 GPU中不充分的精度支持

• 硬件演进滞后于算法需求：硬件加速器（如GPU）在不断适应深度学习的数据类型需求。从NVIDIA的Fermi架构支持FP32/FP64，到Pascal架构引入FP16，再到Turing架构的INT4/INT8，Ampere架构的BF16，以及最新Hopper架构的FP8，GPU的多功能性在不断增强。然而，硬件的发展通常滞后于算法或模型的需求。当遇到不支持的数据类型时，必须通过在更高精度的支持类型上进行转换或模拟，这可能导致显著的性能问题和效率低下。
  

2.3 低精度计算的低效性

• 优化挑战与性能瓶颈：低精度计算因其细粒度的数据访问和特殊的硬件单元（如TensorCore）而尤其难以优化。如表1所示，在NVIDIA V100、A100和AMD MI250上对标准矩阵乘法进行的测试表明：首先，低精度计算的硬件利用率普遍较低，平均不到60%，即使是主流的FP16，平均利用率也仅在60%左右。其次，一些硬件支持的精度并未得到软件的良好支持，例如大多数深度学习编译器不支持在A100和MI250上进行INT8计算。最后，硬件难以快速响应新的精度需求，例如FP8仅在下一代NVIDIA Hopper架构中可用，而混合精度计算（如F16×NF4）在所有最新的GPU上都不被支持。

2.4 我们的洞见

• 通过张量变换解决计算流水线瓶颈：本文以FP16×INT8的混合精度矩阵乘法为例阐述了核心洞见，如图2所示。DNN操作通常实现为一个计算流水线，从多层内存层次结构中加载小数据块（tile）到顶层核心进行计算。每一层内存都有其偏好的最小访问粒度，例如L1层的8字节事务长度。
• 对齐问题与解决方案：当数据块在内存中的存储与硬件的访问粒度或指令形状（如ldmatrix.2x2.f16加载2x2块）不对齐时，会导致带宽利用率低下（如图2左侧所示）。此外，由于缺乏FP16×INT8的计算指令，即使数据成功加载到寄存器也无法计算。
• 核心思想：分离存储、计算并进行变换：我们观察到，通过根据数据类型宽度、内存事务长度和指令形状来变换张量布局，可以规避对齐问题。例如，图2右侧展示了将每个2x2的数据块在L1层中连续存储，使得上层的加载指令能充分利用带宽。同时，由于计算指令只支持FP16，我们可以在数据从L2加载到L1的过程中，将INT8张量转换为FP16。
• 最终效果：通过这种方式，从L2到L1的数据加载因低比特数据类型而流量更低；从L1到L0的数据加载通过事务对齐充分利用了内存带宽；最终的计算通过类型转换在硬件计算单元上得到加速。这个例子表明，即使硬件不支持某种自定义数据类型，通过精心设计的张量布局和数据类型变换，仍然可以调度和优化其DNN计算。
  

A2 方法细节

3. LADDER 设计

系统概述：针对第2节中的观察，LADDER被设计为一个将数据类型视为一等公民的DNN编译器，它引入了张量变换来支持自定义数据类型上的高效DNN计算。图3展示了其系统架构。

核心抽象 tTile：LADDER的核心是TypedTile（tTile）抽象，它在基于分块（tile）的张量抽象上增加了数据类型信息（即tType，见§3.1）。算法设计者可以使用常用数据类型（如FP16）或定义自定义数据类型（如MXFP8, NF4）作为tType，并在此数据类型上定义DNN计算。随后，LADDER将输入的DNN模型转换为一个基于tTile的数据流图（tTile-graph），其中算子被定义为基于tTile的计算任务（tTile-operator）（见§3.1）。

硬件抽象 tTile-device：此外，LADDER将硬件加速器抽象为一个多层层次结构，每一层的需求都由一个tTile（tTile-device，见§3.1）来表示。tTile-device显式地描述了每一层的要求，例如支持的数据类型、事务大小等。通过将tTile-graph中的tTile与tTile-device对齐，以tTile-graph表示的DNN计算便可以在硬件加速器上执行。

调度机制与策略：给定初始的tTile-graph和硬件规格，LADDER会将DNN模型编译成在加速器上的高效执行计划。为了在tTile-device上调度tTile-graph并满足硬件层次结构的要求，LADDER将其调度机制与策略分离。在机制方面，LADDER提出了四种tTile变换原语：slice、map、pad和convert，它们能够将一个tTile变换为另一个等价的tTile（见§3.2）。

优化空间与策略：调度器会将初始的tTile-graph调度为一个能够精细控制tTile配置、变换以及在硬件层次结构上布局的tTile-graph。tTile抽象扩大了DNN计算的调度空间，并在内存占用效率和延迟效率之间开启了新的权衡。在策略方面，LADDER基于观察采用启发式方法，并提供了一种硬件感知的、分层的策略来优化延迟效率（见§3.3）。

代码生成：最后，编译好的、由tTile-graph表示的计划将被生成为给定硬件加速器的可执行代码。

3.1 tTile 抽象

• tType：根据第2节的观察，DNN计算中的数据类型通常在元素级或块级粒度上定义。为了表达这些数据类型，LADDER引入了tType的概念（图4(a)）。具体来说，tType表示一种由一组同质元素构成的数据类型。这些元素的布局是一个n维数组shape。每个元素共享相同的类型，并使用nElemBits位来存储。这组元素还共享相同的元数据。如第2节所述，数据类型通常可以被某些更高位的数据类型无损表示。c_tTypes表示一个tType可以通过c_func函数无损地转换为另一个tType。无论是现有的常用数据类型还是新的自定义数据类型，都可以用tType表示。例如，FP16类型可以表示为shape=[1]、nElemBits=16的tType。元素级的NF4类型可以表示为shape=[1]、nElemBits=4且元数据中包含共享值映射的tType。NF4类型可以无损地表示为FP16，因此在c_tTypes中可能有一个<FP16, NF4_to_FP16_func>条目。块级的OCP-MXFP8类型可以表示为shape=[32]、nElemBits=8且元数据中包含共享缩放因子的tType。

• tTile：基于代表数据类型的tType，LADDER提出了tTile来表示特定数据类型的张量在细粒度分块上的形式。具体来说，如图4(b)所示，一个tTile被定义为一组具有相同数据类型dtype和n维数组布局shape的同质元素。一个tTile中的元素共享一个元数据。此外，tTile中的元素按行主序存储。

• tTile-Operator：一个DNN算子（例如MatMul）通常被实现为一组独立且同质的任务，每个任务处理输入张量的一个分块并输出输出张的一个分块。通过tTile抽象，一个特定数据类型的张量在细粒度的分块级别上被表示。因此，LADDER可以利用tTile来将自定义数据类型的DNN算子表示为一组独立且同质的细粒度任务，即tTile-operator。具体来说，如图4(c)所示，一个tTile-operator显式地表示了对形状为shape的元素进行的张量计算任务。get_input_tTiles()和get_output_tTiles()返回此计算任务的输入和输出tTile。compute()则为输入和输出tTile执行在张量表达式expr中定义的计算。

• tType注解的张量表达式：tTile-operator的计算被定义为一个基于索引的lambda表达式expr【索引13，Tvm: end-to-end optimization stack for deep learning，2018，arXiv；索引40，Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，PLDI；索引57，Roller: Fast and efficient tensor compilation for deep learning，2022，OSDI】。然而，现有张量编译器【索引13，Tvm: end-to-end optimization stack for deep learning，2018，arXiv；索引21，Tensorir: An abstraction for automatic tensorized program optimization，2023，ASPLOS；索引52，Ansor: Generating high-performance tensor programs for deep learning，2020，OSDI；索引56，Amos: enabling automatic mapping for tensor computations on spatial accelerators with hardware abstraction，2022，ISCA；索引57，Roller: Fast and efficient tensor compilation for deep learning，2022，OSDI】中的张量表达式主要关注描述索引和形状，不能灵活地指示计算过程中的数据类型。例如，它无法表达一个FP16张量乘以一个FP16张量，并以FP32作为累加类型。为了支持混合数据类型上的计算表达，需要在计算过程中表达数据类型。因此，LADDER在张量表达式中引入了tType注解，以明确指示计算过程中的数据类型，包括输入、输出和中间数据，从而表示混合数据类型上的计算。例如，一个FP16类型的张量A[M,K]乘以一个NF4类型的张量B[N,K]，以FP32作为累加类型，并输出一个FP16类型的张量C[M,N]，可以如图5(a)所示进行表达。通过带有tType注解的张量表达式，给定形状后，LADDER可以推断出相应的输入和输出tTile。

• tTile-Graph与硬件抽象：通过基于tTile的细粒度表示，DNN算子可以将一个DNN模型表示为一个细粒度的tTile-graph，其中每个节点是一个tTile-operator，每条边代表两个tTile-operator之间的依赖关系。现代硬件加速器通常具有硬件层次结构，包括内存层（如DRAM、寄存器）和计算单元。硬件层次结构中的每一层都有其数据访问偏好。具体来说，内存层通常要求通过事务（transaction）进行访问，其中事务是某个粒度下的一段连续数据或一个特定形状的数据。例如，NVIDIA GPU的共享内存要求事务是32个4字节的bank。计算单元通常也要求处理特定粒度下特定形状的数据。例如，NVIDIA GPU中的hfma2指令处理的粒度是两个FP16值。这些要求可以被描述为tTile。

• tTile-device：因此，LADDER将硬件加速器抽象为一个由多个层级构成的层次结构，每个层级都由tTile描述，即tTile-device。每个层级要么是内存层，要么是计算单元，其要求（表示为某个粒度上的形状）被描述为一个tTile，而粒度则被描述为tType。图5(b)展示了NVIDIA A100 GPU使用FP16张量核心的tTile-device。FP16张量核心的MMA指令要求分别处理[16,16]和[8,16]粒度的两个输入。这可以表示为dtype=FP16、shape=[16,16]的tTile。FP16张量核心的数据加载指令要求加载[16,2]的数据，粒度为half8（即8个FP16值），可以表示为dtype=16B、shape=[16,2]的tTile。此外，充分利用共享内存的要求可以表示为dtype=4B、shape=[32]的tTile。全局内存的32字节事务要求可以表示为dtype=1B、shape=[32]的tTile。

3.2 tTile 变换

• 变换目的：tTile显式地描述了细粒度的张量存储和硬件层次结构的要求。在tTile-graph中表示的DNN计算应与tTile-device对齐以实现高效执行。幸运的是，根据我们在§2中的观察，流水线中的张量存储和访问可以被变换为逻辑上等价的格式，而每种格式在硬件层次结构中都有不同的性能影响。因此，LADDER提出了tTile变换机制，以支持将一个tTile的布局或tType变换为等价的tTile。
• 计算流水线阶段：具体来说，LADDER将一个tTile-operator的计算流水线扩展为硬件层次结构上的三个阶段：Transform-Load（变换-加载）、Compute（计算）和Transform-Store（变换-存储）。Transform-Load将tTile从较低的内存层加载到较高的内存层，并进行tTile变换。Compute在计算单元上执行tTile-operator的计算任务。Transform-Store将tTile从较高的内存层存储到较低的内存层，并进行tTile变换。

• 变换原语：LADDER提供了四种原语来将一个tTile变换为等价的tTile，如图6所示。
  • Slice：slice原语从tTile_input的地址index处切片出一组形状为shape的元素，并以out_shape的新tTile形式返回。slice原语通常用于表示数据分块（tiling）。
  • Map：map原语修改tTile中元素的布局。给定map_func，map原语将每个元素的地址映射到期望的地址。例如，在图5(d)中，从L2内存层到L1内存层的TransformLoad_L1B利用map原语和map_func修改了元素的地址。
  • Pad：pad原语使用pad_value在pad_shape给定的每个边界上对tTile_input进行填充。pad_shape的长度是tTile_input的shape长度的2倍，分别描述了每个维度的左边界和右边界。
  • Convert：convert原语将tTile_input的tType转换为给定的new_tType。给定的new_tType应在tTile_input的tType的c_tTypes中。convert将对tTile_input中的每个元素调用给定new_tType对应的c_func，并返回期望的new_tType的tTile。例如，在图5(d)中，TransformLoad_L1B使用convert原语将tType从NF4转换为FP16，以满足核心的FP16 tType要求。
• 变换的应用：通过上述四种原语，一个tTile可以通过slice和pad改变形状，通过map修改元素布局，或通过convert转换tType，从而变换为另一个等价的tTile。这使得tTile-operator的tTile能够变换以与tTile-device对齐，从而在硬件层次结构中被高效处理。
• 变换示例：图5展示了一个示例，一个FP16张量A[32,63]乘以一个NF4张量B[32,63]，以FP32作为累加，输出一个FP16张量C[32,32]（图5(a)），在一个四层tTile-device（从L2到核心，图5(b)）上执行。具体来说，图5(c)展示了执行的伪代码。A和B的tTile被变换并以FP16类型从L2加载到L1。然后，tTile通过ldmatrix加载到L0，并由mma指令处理，该指令在L0中以FP32累加中间结果。最后，L0中的C的tTile被变换并以FP16存储到L2。图5(d)(e)展示了NF4张量B为与tTile-device对齐而进行的详细变换，张量A的变换与此类似。具体来说，mma和ldmatrix指令要求L1中的数据类型为FP16。每一层也有其事务要求，如图5(b)所示。因此，TransformLoad_L1B对[16,63]进行切片并填充到[16,64]，这与L2的事务要求对齐。然后，TransformLoad_L1B将其转换为FP16并映射到另一个元素布局，以与L1和L0的事务要求对齐。我们得到了L1中的FP16 L1_B[16,64]。接着，TransformLoad_L0B利用ldmatrix对L1_B进行切片，并在L0上得到FP16 L0_B[16,16]，这与L1、L0和mma核心的要求对齐。

3.3 硬件感知的 tTile-Graph 调度

• 调度任务：给定表示为tTile-graph的DNN计算，要将其调度到tTile-device上，我们可以将每个tTile-operator的tTile计算流水线（即Transform-Load、Compute和Transform-Store）映射到tTile-device上。具体来说，我们可以将每个tTile-operator划分为多个tTile以适应每个内存层的容量，调度tTile变换以使tTile与硬件层的要求对齐，并协调算子间的tTile配置和变换以进行整体优化。最终，整个tTile-graph被调度为一个数据流水线，其中一个tTile-operator节点的tTile在硬件层次结构中上下移动，并通过边传递给后继的tTile-operator节点。
• 调度空间与挑战：tTile-graph的调度空间变得非常大，因为tTile在DNN计算调度中开启了另一个维度（即张量变换）。此外，tTile变换引入了内存占用效率和延迟效率之间的新权衡，这给调度带来了更多的复杂性和挑战。以在NVIDIA GPU上计算FP16张量和NF4张量的MatMul为例，由于硬件支持限制，需要将NF4类型转换为FP16。此转换必须在从L1到L0的Transform-Load之前完成，因此可以安排在L2或L1进行。当转换在L2上进行时，会占用更多L2和L1的内存，但在后续从L2到L1再到L0的tTile移动中不会占用计算单元。当转换在L1上进行时，会节省L2的内存和L2的内存带宽，但会占用计算单元进行类型转换。当算子受计算单元限制时，前一个选项可以实现更低的延迟但内存占用更高。当算子受内存IO限制时，后一个选项在延迟和内存占用上都能实现更好的性能。此外，由于转换只需在从L1到L0的Transform-Load之前完成，此转换可以融合到前一个算子中执行，以实现更好的端到端性能。
• 调度策略：面对如此大的调度空间，LADDER提供了一种面向延迟的策略，旨在最小化端到端延迟。具体来说，LADDER提出了一种基于硬件感知的逐层调度策略：较低层的内存提供首选的数据访问粒度作为提示（hint），表示为一个tTile，而上层通过与这个tTile表示的粒度对齐来决定最优的计算粒度，并进行变换。为了减少巨大的调度空间并在合理的时间内调度出合适的计划，LADDER采用了基于我们观察的启发式方法。
• 调度算法：算法1描述了基于提示的逐层调度策略。它接收一个表示为tTile-graph的DNN模型g和表示为tTile-device的硬件规格D，并返回调度后的tTile-graph gret。

Data: g: tTile-graph; D: tTile-device   
Result: gret : scheduled tTile-graph   
1 Function GetDeviceHint(g, D):   
2 D = SelectDeviceConfig(g, D);   
3 HintShape = None, HintGranularity = None;   
4 for layer ∈ D.layers do   
5 HintGranularity = LCM(HintGranularity, layer.tTile.type);   
6 for layer ∈ D.layers do   
7 layer.tTile = convert(layer.tTile, HintGranularity);   
8 HintShape = LCM(HintShape, layer.tTile.shape);   
9 for layer ∈ D.layers do   
10 layer.tTile.shape = HintShape;   
11 return D;   
12 Function ScheduleTransform(op,D,lid):   
13 tTileh = op.tTile[lid -1];   
14 tTilel = op.tTile[lid ];   
15 ScheduleSlice(tTilel , tTileh);   
16 if LCM(tTilel .shape, tTileh.shape) != tTilel .shape then   
17 SchedulePad(tTilel , tTileh, D);   
18 if tTilel .type != tTileh.type then   
19 ScheduleConvert(tTilel , tTileh, D);   
20 if nBits(tTileh.shape[-1]) != nBits(D.layers[lid ].shape[-1]) then   
21 ScheduleMap(tTilel , tTileh, D);   
22 return op.transform[lid -1];   
23 Function ScheduleConnectedGraph(g, D):   
24 D = GetDeviceHint(g, D);   
25 for lid in length(D.layers) do   
26 for op ∈ g[lid ] do   
27 op.tTile[lid ] = ScheduleTiling(op,D,lid );   
28 if lid > 0 then   
29 op.transform[lid ] = ScheduleTransform(op,D,lid );   
30 g = ProfileAndSelect(g);   
31 return g;   
32 Function Schedule(g,D):   
33 g = ExtractConnectedGraph(g, D);   
34 for gconn ∈ g do   
35 gconn = ScheduleConnectedGraph(gconn, D);   
36 return g;

最初，该策略将图调度为子图（第33行）。每个子图表示一个计算流水线，它将`tTile`从最低内存层加载到核心，然后将结果存储回最低内存层。一个子图可以是一个`tTile-operator`或一组可以融合的`tTile-operator`。`ExtractConnectedGraph`可以利用现有的DNN编译器工作【索引13，Tvm: end-to-end optimization stack for deep learning，2018，arXiv；索引43，Welder: Scheduling deep learning memory access via tile-graph，2023，OSDI】。对于一个子图，它首先从硬件中推断出提示。具体来说，它首先选择合适的硬件配置（例如，计算核心）（第2行），偏好硬件支持的比特位最接近的`tType`。因为位数更多的数值类型通常需要更多的晶体管来实现硬件指令，并且通常性能较低。例如，在NVIDIA A100 GPU中，NF4类型可以转换为FP16或FP32进行处理，LADDER将选择FP16核心（312 TFlops）而不是FP32（19.5 TFlops）。然后，它通过比特对齐为每个硬件层找到对齐的粒度和形状，并配置提示（第1-11行）。以NVIDIA A100为例（图5(b)），`HintGranularity`是`ldmatrix`要求的16B，`HintShape`是`[4,8]`，其中内维是128B，与全局内存的32B事务和共享内存的128B事务对齐。然后，策略从顶层（即核心）到底层（即DRAM）逐层调度此子图（第25-29行）。在每一层，策略首先通过`ScheduleTiling`调度`tTile-operator`的分块并带有提示（第27行），然后调度`tTile`变换（第29行）。如果`ScheduleTiling`（第27行）调度的算子分块是`[4,8]`和16B的倍数，后续的`ScheduleTransform`可以使此调度与`tTile-device`对齐。此外，`ScheduleTiling`可以利用现有的张量编译器【索引13，Tvm: end-to-end optimization stack for deep learning，2018，arXiv；索引52，Ansor: Generating high-performance tensor programs for deep learning，2020，OSDI；索引57，Roller: Fast and efficient tensor compilation for deep learning，2022，OSDI】。在`ScheduleTransform`中，策略将检查形状和类型与`tTile-device`的对齐情况，并调度相应的变换来对齐`tTile`（第12-22行）。调度后可能会有一些候选方案，将对其进行性能评测并返回最佳方案（第30行）。

• ScheduleMap：在调度map变换中的map_func是非平凡的。LADDER提出了一种推断map_func的方法，即将tTile中的元素按行主序映射到所需的事务大小。图5(e)展示了一个例子：在16B的粒度下，为了将L0中的shape[16,2]映射到L1中所需的shape[8]，元素按行主序被展平，得到shape[4,8]。map也可以支持其他的map_func。这种调度策略不保证最优。然而，如§5所示，这种调度策略已经可以超越最先进的技术，并能够在GPU上实现高效的低精度DNN计算。我们也希望，所提出的调度机制带来的这个优化空间可以被未来关于更先进调度策略的研究进一步探索。

4. 实现

代码与依赖：LADDER的实现包含约5000行代码，包括Python和C++，基于开源DNN编译器：TVM【索引13，Tvm: end-to-end optimization stack for deep learning，2018，arXiv】、Welder【索引43，Welder: Scheduling deep learning memory access via tile-graph，2023，OSDI】和Roller【索引57，Roller: Fast and efficient tensor compilation for deep learning，2022，OSDI】。LADDER修改了TVM以实现核函数调度和生成核函数代码，同时利用Roller来推断高效的tTile配置。Welder是目前最先进的能够全面优化DNN模型的DNN编译器，被用于端到端的图优化。

工作流程：LADDER的输入是一个PyTorch程序。对于PyTorch内置的数据类型，LADDER不需要对DNN模型程序做任何修改。此外，对于PyTorch不支持的新数据类型，LADDER通过自定义算子扩展了PyTorch，以表达用户定义数据类型上的张量表达式。给定PyTorch程序，LADDER将其导出为ONNX图。LADDER也扩展了ONNX以表示新数据类型上的计算，其中tType注解的张量表达式保存在ONNX图节点的属性中。有了导出的ONNX图和目标硬件加速器的基于tTile的规范文件，LADD-ER会自动将ONNX图转换为tTile-graph并执行调度。然后，LADDER为目标硬件加速器生成设备代码。

硬件支持：我们为NVIDIA GPU和AMD GPU实现了LADDER，因为它们是DNN最流行的加速器。在本节的其余部分，我们将详细描述在NVIDIA GPU上的LADDER实现，并简要描述在AMD GPU上的实现。此外，如果新的硬件指令（例如最新Hopper GPU中的FP8张量核心）和其他硬件加速器（例如Graphcore IPU）符合基于tTile的硬件抽象并提供在硬件层次结构上加载和存储数据的编程接口，LADDER也可以移植到它们上面。

4.1 LADDER on NVIDIA CUDA GPUs

4.1.1 tType 和 tTile

• 数据类型支持：LADDER实现了常见数据类型的tType，例如FP32、FP16、INT8、FP8、MXFP、INT4、NF4、INT1。
• 存储方式：GPU是单指令多线程（SIMT）架构，偏好一组线程在不同数据上处理相同指令。因此，LADDER分别存储tTile中的元素和每个元数据。图7展示了在NVIDIA GPU上一个形状为[32,32]的MXFP8 tTile的存储方式。元素存储在一个数组中，而共享的缩放因子存储在另一个数组中。为了访问一个tTile，连续的线程处理连续的元素，从而实现合并访问。
  
• 非2的幂比特宽度处理：需要注意的是，可能存在一些数据类型的nElemBits不是2的n次方，例如3比特【索引22，Gptq: Accurate post-training quantization for generative pre-trained transformers，2022，arXiv】。为了支持这些数据类型，LADDER以4B的粒度进行存储，这是由于GPU的规格所致，例如10个3比特的值可以存储在一个4B（32位）的粒度中。

4.1.2 使用 PTX 指令优化代码生成

• PTX的应用：NVIDIA不提供用于编程的汇编指令。取而代之的是，NVIDIA引入了并行线程执行（PTX）作为NVIDIA GPU的低级虚拟机，其中PTX虚拟机上的指令集架构（ISA）可以被视为NVIDIA GPU的指令级API【索引8，PTX ISA，https://docs.nvidia.com/cuda/parallel-thread-execution/index.html】 。CUDA C++代码首先被编译成PTX代码，然后再编译成机器码执行。
• 具体优化：CUDA为某些单元同时提供了C++ API和PTX API。例如，张量核心提供了WMMA C++ API和MMA PTX API，其中一个WMMA API会被nvcc编译器编译成一组MMA指令。MMA PTX API比WMMA C++ API具有更大的灵活性和更好的性能。LADDER在张量核心上使用MMA PTX API进行代码生成，并使用cp.async指令来支持Ampere GPU上的新异步内存复制功能【索引3，NVIDIA A100 Tensor Core GPU Architecture】。此外，我们观察到将低比特整数（如INT4）转换为浮点数（如FP16）可能会引入显著的开销。LADDER使用LOP3指令【索引8，PTX ISA，https://docs.nvidia.com/cuda/parallel-thread-execution/index.html】实现了低于4比特的整数转换。我们修改了TVM中的代码生成模块以实现这些优化 。

4.2 LADDER on AMD ROCm GPUs

• 硬件相似性与抽象：AMD GPU与NVIDIA GPU相似，也具有一个硬件层次结构，包括所有CU共享的全局内存、每个CU中的本地数据存储（类似于共享内存）、寄存器和核心。因此，与NVIDIA GPU类似，一个AMD GPU可以被抽象为一个具有不同tTile配置的四层tTile-device。
• 实现细节：ROCm为AMD GPU提供了HIP编程模型【索引1，AMD ROCm Platform，https://github.com/RadeonOpenCompute/ROCm】，其功能与CUDA类似并支持大多数CUDA语句。我们在TVM中实现了一个新的用于HIP的代码生成后端，以支 持AMD ROCm GPU。此外，我们使用MFMA（Matrix Fused-Multiply Add）ISA级API来利用矩阵核心（相当于NVIDIA的张量核心）。

A4 实验环境

• 硬件配置：

  • NVIDIA GPUs: Tesla V100 (16GB), A100 (80GB), RTX A6000 (48GB), RTX 4090 (用于部分算子评测)。
  • AMD GPUs: AMD Instinct MI250 (128GB)。
  • 操作系统: Ubuntu 20.04。

• 软件配置：

  • 驱动与工具包: CUDA 12.1, ROCm 5.7.0。
  • 实现与依赖: LADDER基于PyTorch，并利用TVM, Welder, Roller进行编译和优化。
  • 基线系统:
    • NVIDIA: Welder, PyTorch-Inductor, ONNXRuntime, TensorRT, AMOS, TensorIR, vLLM。
    • AMD: Welder, PyTorch-Inductor, ONNXRuntime, TensorIR。
    • 算子级: cuBLAS, CUTLASS, vLLM, cuDNN, AMOS, TensorIR。

• 模型与数据集：

  • 大型语言模型 (LLM): LLAMA-70B, BLOOM-176B。
  • 计算机视觉模型: ResNet-50, ShuffleNet-V2, ViT-Base。
  • 音频模型: Transducer Conformer-l。
  • 评测数据集: WikiText-2 (用于评估LLM的困惑度PPL)。

• 数据类型配置:

  • 实验覆盖了多种权重(W)和激活(A)的数据类型组合（表示为WtypeAtype），均来自前沿研究文献，确保了模型质量。
  • LLM: WFP16A16, WINT4A16, WNF4A16, WFP8A8, WMXFP8A8, WINT1A8。
  • 其他模型: WFP16A16, WFP8A8, WMXFP8A8, WINT1A4, WINT4A4, WINT8A4等。

• 评测场景:

  • LLM: 测试了不同批大小(BS)和序列长度(SEQ)的组合，如(BS=1, SEQ=1)代表在线推理的decoding阶段，(BS=32, SEQ=1)代表离线推理，(BS=1, SEQ=4096)代表pre-fill阶段。
  • 其他模型: 批大小为1和128，分别代表在线和离线推理场景。

A4 实验结果

5.2 NVIDIA GPU 评估

• 端到端推理延迟 (图8, 9, 10)

  • 实验内容: 在A100, V100, RTX A6000上评测了LLAMA, BLOOM, ResNet等多个模型的端到端推理延迟。
  • 实验结果:
    • 在标准的WFP16A16配置下，LADDER相比于Welder在A100上对LLAMA、BLOOM、ResNet等模型分别取得了1.0×、1.2×、2.0×等加速比，这得益于其对不规则形状的更优核函数调度。
    • 在LLM广泛使用的WINT4A16配置下，LADDER相比vLLM平均取得了2.3×的加速。
    • 对于其他系统不支持的自定义数据类型（如WINT1AINT8），LADDER展示了巨大优势，在BLOOM-176B单层上相比Welder实现了高达10×（A100）、13.3×（V100）和14.6×（A6000）的加速。
  • 分析结论: LADDER不仅在标准精度上表现出色，而且能够高效支持自定义低精度数据类型，在不同GPU上均取得显著的性能提升。
    
    
    

• 内存使用 (图11)

  • 实验内容: 在A100上测量了LLM推理在不同数据类型配置下的内存占用。
  • 实验结果: 内存使用量随着比特宽度的降低几乎呈线性下降。在WINT1AINT8配置下，相比WFP16A16，LLAMA模型的内存占用最多减少了74%，BLOOM模型最多减少了85%。
  • 分析结论: 使用低精度数据类型是缓解LLM内存压力的关键策略，LADDER能够有效地实现这一点。
    

• 编译时间 (表2)

  • 实验内容: 对比了LADDER与AMOS, TensorIR, Welder在ResNet和ShuffleNet上的端到端编译时间。
  • 实验结果: LADDER的编译时间比AMOS快两个数量级，比TensorIR快一个数量级，但略高于Welder。
  • 分析结论: LADDER通过扩展调度空间来支持低精度，这带来了编译开销，但其启发式策略将开销控制在合理范围内，远优于其他需要大量搜索的编译器。
    

• 算子基准测试 (图12, 13)

  • 实验内容: 在A100和RTX 4090上对LLAMA和ResNet中的核心MatMul和Conv2d算子进行性能评测。
  • 实验结果: 在A100上，相比WFP16A16，WINT4A16平均加速1.8×，WINT1AINT8平均加速4.5×。在RTX 4090上，LADDER在硬件支持的WFP8A8上性能优于cuBLAS，与CUTLASS相当，并在自定义类型上展现出比A100上更大的加速潜力。
  • 分析结论: LADDER生成的核函数性能优越，能够充分利用硬件特性，并在支持新硬件（如4090的FP8 Tensor Core）和自定义类型方面表现出色。
    
    

• 优化分解 (图14)

  • 实验内容: 展示了LADDER应用于LLAMA-70B核函数的各项优化（Tile-aware变换、PTX级优化、整体调度）所带来的逐步性能提升。
  • 实验结果: Tile-aware变换带来2.0×加速，PTX级优化进一步带来高达1.7×加速，整体调度策略最终带来高达2.5×的加速。
  • 分析结论: LADDER的各项优化组件都对最终性能有显著贡献，协同作用下实现了卓越的计算效率和适应性。
    

• 比特宽度扩展性 (图15)

  • 实验内容: 评估了权重和激活的比特宽度从高到低变化时，性能的相应变化。
  • 实验结果: 随着比特宽度降低，加速比相应提升。在内存受限的解码场景（SEQ=1），降低权重比特宽度能持续带来加速。在计算受限的编码场景（SEQ=4096），性能受限于混合精度中的高精度计算部分，加速比变化不大。
  • 分析结论: LADDER能灵活支持任意比特宽度，并展示了不同场景下精度与性能的权衡关系。
    

• 低精度LLM的效率与准确性 (图16)

  • 实验内容: 评估了多种SOTA低精度方法（如GPTQ, BitNet）在LLAMA2模型上的准确性（PPL）和由LADDER实现的推理延迟。
  • 实验结果: WFP8、WNF4、WINT4等方法在几乎不损失PPL的情况下，分别带来1.6×、1.7×、2.5×的平均加速。新兴的BitNet-b1.58 (WINT2A8)在LLAMA2-70B配置上实现了4.6×的加速，并展现了兼顾准确性和效率的潜力。实验还表明，量化后的大模型（如LLAMA2-13B-WINT4）在效率和准确性上均可超越未量化的小模型（LLAMA2-7B-WFP16）。
  • 分析结论: LADDER为低精度模型研究提供了可靠的性能反馈，验证了低精度计算在提升模型效率和保持准确性方面的巨大价值。
    

5.3 AMD GPU 评估 (图17)

• 实验内容: 在AMD Instinct MI250上对比了LADDER与Welder, PyTorch-Inductor, ONNXRuntime的端到端性能。
• 实验结果:
  • 在WFP16A16配置下，LADDER相比Welder在LLAMA、BLOOM、ResNet等模型上平均取得了1.5×到10.5×的加速，尤其在ShuffleNet上因更好的算子融合能力达到14.1×加速。
  • 在WINT4A16配置下，LADDER在LLAMA和BLOOM上相比Welder分别取得了高达3.8×和4.5×的加速。
• 分析结论: LADDER的设计是跨平台的，在AMD GPU上也表现出强大的性能优势，能够高效生成利用矩阵核心的核函数并实现更多融合机会。
  

A7 补充细节

6. 讨论

• 当前实现的局限性：LADDER当前的实现主要集中在模型推理上，本节将讨论其一些局限性和未来的工作方向。
• 多GPU服务：对于像BLOOM-176B和LLAMA2-70B这样无法装入单个GPU的大型模型，需要多GPU部署。多GPU支持与LADDER是互补的。LADDER专注于在单个硬件加速器上支持低精度计算，而多GPU框架【索引28，Efficient memory management for large language model serving with pagedattention，2023，SOSP；索引30，AlpaServe: Statistical multiplexing with model parallelism for deep learning serving，2023，OSDI；索引32，nnScaler: Constraint-guided parallelization plan generation for deep learning training，2024，OSDI；索引53，Alpa: Automating inter-and intra-operator parallelism for distributed deep learning，2022，OSDI】则专注于模型划分和跨多GPU的并行计算调度。LADDER可以与多GPU框架协作，为多GPU上的低精度模型启用并行计算，即由多GPU框架划分模型并调度分区后的计算到LADDER在单个设备上执行。将LADDER与多GPU框架的集成留作我们未来的工作。
• 低精度训练：LADDER的设计并不局限于推理。低精度模型的训练和推理都需要系统和硬件的低精度支持，并且训练中的反向计算与前向计算类似。低精度模型训练可以从以下两方面获益：1）利用更高效的低精度计算单元，例如A100上支持的WINT8AINT8张量核心吞吐量是WFP16AFP16的2倍，而WINT4AINT4则是4倍；2）低精度模型表示带来的较小内存占用可以支持更大的批处理大小，从而可能提高硬件利用率。我们将低精度训练留作未来的工作。

A5 结论

本文介绍了LADDER，这是首个旨在优化通用低精度计算在GPU等加速器上的深度学习编译器。LADDER通过暴露一个通用的类型系统（tType）和扩展的张量表达式，使用户能够轻松实现和表达深度学习中的新数据类型。它引入了一套新的张量调度原语，以促进计算流水线中张量存储、访问和类型转换等方面的优化。LADDER的分层、硬件感知的优化策略能够驾驭复杂的变换空间，展示了其系统性支持各种低比特精度自定义数据类型的能力。这在无需硬件修改的情况下，提升了现代加速器上DNN计算的性能。这项创新不仅赋能模型设计者探索数据类型优化，也为硬件供应商提供了一个灵活的解决方案来扩展对多样化精度格式的支持。