TVM Relax：通过跨层次抽象实现动态 shape 的 LLM 高效部署

动态形状计算已成为现代机器学习工作负载中的关键部分，尤其是在新兴的大语言模型中。这些模型的成功推动了它们在各种后端环境中的普遍部署需求。本文介绍了Relax，一种用于优化端到端动态机器学习工作负载的编译器抽象。Relax 引入了:
一种跨层次的抽象，将计算图、循环级张量程序和外部库调用封装在一个统一的表示中。
还引入了一等符号形状注解（first-class symbolic shape annotations），用于全局跟踪动态形状计算，从而实现跨层次的动态形状感知优化。

我们基于这种方法构建了一个端到端的编译框架，用于优化动态形状模型。

实验结果表明，Relax 在各种 GPU 上实现了与最先进的系统相当的性能，并且能够将新兴模型部署到更广泛的环境中，包括手机、嵌入式设备和浏览器。

标题：Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
组织：CMU、OpenAI、SJTU、NVIDIA、ByteDance、UW、UIUC、Hyperbolic
论文：https://arxiv.org/abs/2311.02103
代码：https://github.com/apache/tvm...

1. Introduction

机器学习（ML）应用如今在日常生活和更广泛的经济领域中无处不在。随着 GPT-4 和开源大语言模型（LLMs）的出现，构建更强大的现代 AI 系统（用于处理文本、图像、音频等）的机会已经显现。这些模型的成功也导致了对它们在各种后端环境中的部署需求不断增长，包括服务器、个人电脑、车辆和移动设备。机器学习框架负责将这些模型部署到不同的后端。ML 编译器旨在通过将模型计算转换为通用的程序抽象、执行优化，并在各种平台上生成高性能代码，来缩小模型与后端之间的差距。

然而，为了支持不同硬件后端中的模型张量操作，需要大量的工程工作，尤其是因为大多数模型都使用动态张量形状。动态形状包含可能依赖于程序值的变量，这使得执行关键优化（如静态内存规划）变得更加困难。例如，语言模型必须处理可变大小的输入消息、KV 缓存上下文长度、词汇表大小以及其他形状动态性的来源。

典型的 ML 编译器包括多个抽象层次（或中间表示）以及单次降级（single-shot lowering）。

高层计算图描述模型使用数据流图的高级张量操作（例如矩阵乘法、重塑），以促进全局计算重写。
张量程序描述张量操作的低级循环和索引缓冲区访问，以在此级别上启用细粒度的内核优化（例如循环分块）。
算子库（算子库）允许将张量操作卸载到供应商优化的例程中。

大多数端到端 ML 编译器使用计算图作为高级表示，并将张量程序和算子库视为通常不透明的外部函数。动态形状通常在每个函数内处理。Relay 和 IREE 在计算图中通过“未知”注解处理动态形状，但不跟踪动态形状之间的关系。

PyTorch 编译器启用了即时（JIT）图跟踪，并在每个跟踪函数内处理动态形状跟踪。JIT 方法消除了跨函数边界进行形状跟踪的需求，但也限制了其在受限环境（如移动设备和 WebGPU）中的应用。DietCode、CoRA 和 SparseTIR 专注于在每个张量程序内优化动态形状。Halide 在张量程序中跟踪动态形状，并提供从张量程序调用外部库函数的原语。尽管在每个层次和函数内处理动态形状方面取得了进展，但跨这些层次和函数的优化仍存在挑战。

首先，随着我们针对一系列新兴平台，我们必须启用提前编译（AOT），这需要跨函数进行完整的程序优化。此外，用户定义的算子（如自定义量化解码）需要计算图优化能够感知外部函数。最后，单次降级使分析或转换张量程序变得更加困难，然后使用结果来指导高层图优化。

动态形状跟踪贯穿所有这些挑战中的每一个增量转换，因为丢失形状关系信息可能会显著阻碍程序中的算子和函数之间进行优化的能力。为了解决这些挑战，我们引入了 Relax，这是一种针对新兴端到端动态机器学习模型的全面 AOT 编译器程序抽象。Relax 使我们能够封装计算图、循环级张量程序和外部库函数的抽象，这在本文中被称为跨层次抽象。

我们还引入了一等符号形状，用于跟踪和表示动态形状维度之间的关系。Relax 使用变量表示符号形状维度，并在可能的情况下静态地跟踪动态形状，同时在需要时提供动态回退。跨层次抽象与一等符号形状允许跨这些抽象层次进行分析和优化，同时在优化过程中保留 IR 中的符号形状信息。我们引入了一系列优化，以实现动态形状感知的算子融合、张量程序工作区提升、内存优化、图卸载和张量算子优化。最后，我们在这些元素之上构建了一个端到端的编译框架。本文的贡献如下：

提出了一种跨层次抽象的设计，以实现 ML 框架中传统抽象层次之间的优化和分析。
展示了一种程序抽象，它使用一等符号形状方法，全局跟踪张量算子、子图函数调用以及张量程序和外部库函数的调用之间的动态形状关系，从而实现全程序符号形状跟踪和跨抽象层次的动态形状感知优化。
构建了一个 AOT 端到端编译框架，以支持将新兴模型部署到各种硬件后端，包括许多现有框架支持不足的新兴后端。

实验结果表明，Relax 能够将新兴 LLMs 编译和优化到一系列新兴设备和环境中，包括手机、嵌入式设备和通过 WebGPU 的浏览器。此外，Relax 在性能上与经过大量优化的特定平台解决方案相当。Relax 已被纳入一个主要的开源项目，并支持新兴机器学习模型的通用部署。

2. Overview

本节描述了我们的方法的关键见解，并概述了本文的内容。上面的图 1 总结了我们的整体方法，重点介绍了两个关键设计，这些设计使得编译器能够跨所有抽象层次对动态机器学习模型进行优化。

首先，我们观察到 ML 编译器通常需要经过多个抽象层次才能将机器学习模型带到目标平台。典型的层次包括计算图、张量程序和外部库。传统上，ML 编译器专注于每个单独抽象层次内的优化，并进行单向的单次降级（single-shot lowering）从一个层次到下一个层次。Relax 将计算图、张量程序和库统一到一个单一的跨层次抽象中，允许通过不同的方法逐步优化或部分降低计算的某些部分，同时考虑所有抽象层次的分析。这一设计使我们能够使用不同的方法逐步优化或部分降低计算的某些部分，同时考虑所有抽象层次的分析。

其次，我们观察到，尽管新兴的机器学习工作负载涉及动态形状计算，但我们可以通过考虑形状之间的关系来进行大量的静态分析和优化。因此，我们引入了注解（Annotation），这些注解可以通过符号变量和符号形状计算来跟踪中间计算的形状。我们的方法通过全局跟踪这些动态形状计算，跨子图函数调用和张量程序及外部库函数的调用，来表示整个程序中的动态形状，并启用动态形状感知优化。我们引入了 Relax 的抽象设计（见第 3 节 Relax Abstraction），并讨论了由我们的设计启用的一组具体的跨层次编译器优化（见第 4 节跨层次算法与抽象）。

关键设计

Relax 的关键设计包括以下几个方面：

跨层次抽象（Cross-level abstraction）：将计算图、张量程序和外部库统一到一个单一的抽象中，允许不同层次之间的交互和优化。
一等符号形状（First-class symbolic shapes）：通过符号变量和符号形状计算来跟踪动态形状，使得编译器能够在优化过程中保留形状信息。
逐步优化（Incremental optimization）：允许在不同的抽象层次上逐步进行优化，而不是一次性完成所有优化。

通过这些设计，Relax 能够在动态形状感知的情况下，对整个程序进行全局优化，从而提高性能和内存使用效率。

3. Relax Abstraction

本节介绍 Relax 的整体抽象设计。我们首先介绍语言构造，然后讨论 Relax 中的一等符号形状和跨层次抽象。

3.1 语言构造

Relax 是一种命令式编译器抽象，具有头等函数（first-class functions），专注于在高层（大多数 ML 编译器所指的“图层”）操作张量。Relax 程序对整个张量应用高级算子，并可以在函数之间传递张量（或张量元组），或者调用低级张量程序或外部库函数以对张量进行循环级操作。本节描述 Relax 的三个主要元素：结构注解、数据流块以及同一抽象层次内和跨抽象层次的函数调用。这些构造通过在编译器优化的联合转换中提供符号形状指导，将整个系统结合在一起，同时在这些转换过程中保留程序中的符号形状信息。

注解（Annotations）

Relax 中的每个值都与一个注解相关联，该注解传达结构信息，类似于静态类型。下面的表 1 总结了 Relax 中的不同注解及其使用示例和解释。我们设计注解语法以嵌入 Python AST（抽象语法树），并在函数签名中将符号表达式引用为字符串（例如，“n*4”），当符号变量尚未声明时。

为了丰富注解的形状表达能力并确保形状注解的覆盖范围，我们重用了循环级张量程序的表达式系统，以便形状注解支持张量程序支持的所有整数表达式，并且编译器符号表达式分析（例如表达式相等性证明）可以利用公共表达式。注解在编译时指示值的整体类型（例如张量、元组）以及关于值的其他信息，例如张量的形状和数据类型。注解是全程序符号形状跟踪和跨抽象层次动态形状感知优化的基础。

数据流块（Dataflow Blocks）

Relax 中的数据流块划定了一个无副作用的程序（子）区域，没有控制流，即一系列纯操作的直线序列，以简化程序转换。例如，在对 Relax 数据流块执行死代码消除时，可以安全地移除未使用的算子，而无需考虑是否可能通过移除具有副作用的算子来影响程序的可见行为。

函数调用（Function Calls）

Relax 纳入了函数调用，这些调用可以在同一抽象层次内（即允许一个图层函数调用另一个图层函数）或跨抽象层次，即允许图层函数调用外部张量程序函数和外部库函数。

调用外部循环级张量程序和外部库函数是跨层次抽象的基础元素，这在第 3.3 跨层次抽象中进行了详细探讨。我们使用 TensorIR 作为循环级张量程序抽象，但相同的原则也可用于其他循环级抽象。

3.2 一等符号形状抽象（First-Class Symbolic Shape Abstraction）

张量的形状在 ML 框架的上下文中是非常有用的信息，尤其是对于内存规划。然而，张量形状通常在编译时未知，必须动态处理。一种推理动态形状维度的方法是引入一个任意（或未知）值来表示动态维度，如 ONNX、Relay 和某些 MLIR 方言中所做的。不幸的是，这种方法未能保留可能有用的信息，例如形状维度之间的关系或约束（例如，如果一个张量的维度是，另一个可能是）。

这些信息对于编译器优化非常有价值，而将维度标记为任意值则完全抹去了这些信息。相反，我们引入了一等符号形状注解（如图 3 所示），以更好地推理和优化动态形状模型。符号形状注解使用符号表达式（由整数变量和常数组成）描述每个形状维度。

因此：

对于完全静态的模型，Relax 形状注解涵盖了现有的基于静态形状的注解。
对于具有混合动态和静态维度的模型，Relax 不仅可以符号化地表达形状维度，还可以跟踪维度之间的符号关系。

这些符号形状关系帮助我们应用更多的动态形状感知优化。例如，我们将知道图 3 中 flatten 算子后的总元素数是，并且与输入相同，这表明可能的缓冲区重用。

符号形状注解与未知动态形状注解的比较

符号形状注解能够全面跟踪形状关系，并促进高级动态形状感知优化，而未知动态形状注解则无法做到这一点。

在编译时并不总能跟踪形状关系。例如，图 3 中的 unique 算子，其输出张量的形状取决于输入的运行时值，我们无法推断出数据依赖算子的输出形状。对于这种情况，我们提供了粗粒度的注解（例如表 1 中的 Shape(ndim=2)，表示形状有两维但两者都未知）。为了处理这种情况，我们引入了一个特殊的构造match_cast，它为值断言一个符号形状注解，允许引入新的符号变量。

编译器为每个match_cast插入运行时检查，如果违反了约束条件，则抛出错误。在图 3 的特定例子中，尽管编译器在 unique 算子后不知道 lv2 的形状，但可以使用match_cast假设它具有形状(𝑚,)（如其别名 lv3 所示）。match_cast可以由前端和编译器通插入，以建议更具体的符号形状，并作为开发人员在程序中指示形状信息的有价值工具。

3.3 跨层次抽象（Cross-Level Abstraction）

本节描述 Relax 中启用跨层次抽象的构造。我们的主要目标是设计能够自然表示和优化计算图、外部张量程序和库之间交互的原语。为此，我们必须协调每个抽象层次的不同特性。具体来说：

计算图抽象倾向于使用返回新张量的纯算子。这允许我们通过有向无环图组织计算，并在不担心副作用的情况下有效地进行图重写。
另一方面，大多数低级计算的张量程序和库采用目标传递风格Destination-Passing Style， DPS 接口，即接受计算结果张量作为输入并直接修改它们，而不是分配和返回新张量。

知识补充：Destination-Passing Style, DPS
Destination-Passing Style (DPS) 是一种函数调用方式，它允许函数直接将结果写入预先分配好的目标位置，而不是函数内部自己创建新的变量来返回结果。这种方式可以提高内存管理的效率，因为它避免了在函数内部频繁地分配和释放内存。
通俗理解
你有一个盒子，想把一些东西装进去。通常做法是，你先有一个空盒子，然后把东西装进去。但在 DPS 中，不是先有空盒子，而是别人已经给你准备了一个盒子，你只需要把东西放进去。这样做的好处是，盒子已经存在了，你不需要再去创建新的盒子，考虑大小合适，因为别人帮你准备好了。
例子
假设我们有一个简单的加法函数，它接收两个数字并返回它们的和。在常规的方式中，函数会创建一个新的变量来存储结果并返回它。但在 DPS 中，我们会预先分配一个变量来存储结果，然后把这个变量传递给函数，函数直接在这个变量上进行操作。

# 常规方式
def add(a, b):
    return a + b

result = add(3, 4)
print(result)  # 输出：7

在常规方式中，add 函数创建了一个新的变量 result 来存储结果。下面是 DPS 方式：

# DPS方式
defadd_dps(a, b, result):
    result[0] = a + b

result = [0]  # 预先分配一个存储结果的变量
add_dps(3, 4, result)
print(result[0])  # 输出：7

在 DPS 方式中，我们预先分配了一个变量 result，然后将其传递给 add_dps 函数。函数直接在这个变量上进行操作，将结果写入其中。这样做的好处是，我们可以在函数外部控制内存的分配和管理，而不是让函数内部去处理这些事情。

为此，我们在 Relax 中引入了一个要求，即显式地将输入和输出内存在低级张量程序中传递，以符合 DPS。DPS 将内存管理问题从低级代码中抽象出来，从而简化代码生成。我们引入了两个外部函数调用原语（如图 4 所示）以桥接抽象层次。

首先，我们引入了 call_tir，它允许从图层直接调用张量程序。我们设计了 call_tir 的语义（如下图图 5 所示），以直接映射到低级函数的 DPS 调用。

这种方法允许我们在图层转换期间为 call_tirs 分配高级语义，并在后续阶段将它们降低为带有内存管理的 DPS 调用。值得注意的是，call_tir 还接受输出形状的注解以及可能的其他符号表达式作为参数，以将形状信息从图层传递到循环级张量程序。

这种形状信息对于循环级程序的优化至关重要，例如算子融合。通过从图层向张量程序传递符号形状信息，我们可以允许张量程序生成代码，这些代码专门针对大多数静态维度，并仅在必要时使用动态维度（如前面图 4 中的 n）。

其次，我们引入了 call_dps_library 原语，以允许从图层直接调用外部算子库。在实践中，它引入了极大的原型设计灵活性，因为可以从 Relax 程序轻松调用外部例程。call_dps_library 的约定与 call_tir 类似，只是被调用者是库函数的名称。这些函数由注册表提供，并链接到最终的可运行模块。值得注意的是，我们将图、张量程序和库紧密结合在一起，所有这些在 ML 编译器中都是关键的，同时我们利用并补充了编译器中用于 GPU 源代码生成的其他较低层。

跨层次抽象的好处

上图图 6 总结了跨层次抽象启用的常见优化模式：

部分降低（Partial lowering）：而不是一次性做出所有降低决策，一个 pass 可以对部分计算做出分派决策或循环级优化。例如，当我们希望将融合算子的降低决策替换为不同的库时，我们可以将程序传递到后续 pass 以处理其他算子。
分析反馈（Analysis feedback）：我们可以分析张量程序的循环模式，并自动注解它们的算子属性。通常，编译器开发人员需要为系统中的每个高级算子手动注解属性。通过采用基于分析的属性，可以大大减少高级算子注解的工程成本。
跨层次转换（Cross-level transforms）：有时，只有在某些低级优化之后才能发现优化机会。例如，张量程序分析可能决定张量程序需要一个临时工作区。在这种情况下，我们可以联合转换张量程序和图层，以在图层插入工作区分配，从而允许工作区也作为全局内存规划的一部分。

虽然每种优化模式本身都很有用，但真正的优势在于将它们结合起来。例如，我们可以部分降低到库，然后使用其他技术优化剩余组件。我们将在第 4 节中详细讨论跨层次优化。

4. 跨层次算法与优化

本节描述了一组利用所提出的跨层次抽象在端到端编译框架中实现的具体算法和优化。

4.1 形状注解推导（Shape Annotation Deduction）

注解中的符号形状是优化 pass 的重要信息源。为了最大化利用这些信息，Relax 在模型构建期间以及编译器 pass 之间自动跟踪和推导中间值的符号形状注解，允许 pass 推导形状之间的等式和关系，并启用额外的优化。同时，这也增加了推导效率的需求，因为推导在每个 pass 中运行。
每个张量算子都有一个注册的形状推导规则，该规则根据输入的形状注解和值（如 reshape 的情况）返回输出注解。我们采用前向推导方法，基于输入推导表达式的注解。对于外部函数调用原语 call_tir 和 call_dps_library，输出注解是其参数的一部分（见前文图 4），将直接用于推导。

前向推导的优点是简单和局部性，有效地避免了在处理过程中跨全局上下文的同步复杂性。前向推导还利用了 match_cast 的显式信息。通过前向推导，可以在编译器通期间高效地局部推导出任何新变量的形状注解。此外，为了尽可能提供形状信息，我们还认识到传播跨子图函数调用的全局形状关系的重要性，以适应优化 pass （如融合）的中间结果，其中子图函数本身包含动态输入和输出。下图图 7 展示了子图函数调用的符号形状推导。我们的系统能够根据函数 subfn 的符号关系在调用者中正确传播形状。

关键设计原则

函数边界的隔离符号关系：Relax 中的函数签名提供了参数和返回值的注解，允许仅通过函数签名进行函数调用的注解推断。这允许函数作为具有 Callable 注解的一等值使用。例如，我们可以通过仅查看其签名来推断 subfn 调用的返回形状。
前向符号推导：Relax 基于符号形状关系执行前向形状推导。作为安全网，当无法推断出更具体的信息时（例如对于数据依赖的算子），返回粗粒度注解，这允许符号推导在常见情况下成功，但也支持一般情况。
支持参数注解中的符号表达式：除了符号变量，我们还支持函数参数注解中的通用算术表达式，这对于简化子图函数的跨算子融合等转换至关重要。

4.2 跨层次动态形状感知算子融合（Cross-Level Dynamic Shape–Aware Operator Fusion）

算子融合有助于将多个算子组合在一起，减少算子的整体内存加载成本。下图图 9 展示了 Relax 中算子融合的一般流程。Relax 程序可以包含标准（例如矩阵乘法）和自定义算子的张量程序，这些自定义算子可能没有对应的图层算子，例如用循环编写的量化解码。

候选的模式种类包括 ElementWise、Broadcast、Injective、Reduction、OutputEwiseFusible 或 Opaque（作为回退）。这些模式种类描述了张量程序的数学属性，然后这些信息被 FuseOps pass（算法 2）用于通过基于模式匹配的图分区将张量程序调用分组到子图函数中（一个示例模式是将 ElementWise 张量程序融合到 OutputEwiseFusible 的后面，例如矩阵乘法+ReLU）。

最后，我们应用 FuseTensorIR，这是一种跨层次转换，通过合并每个子图函数中调用的张量程序来联合更新张量程序和图层调用站点。与静态形状融合不同，我们确保上述所有步骤通过跟踪符号变量并生成额外的符号参数来支持参数注解中的符号表达式（如第 4.1 节中所述）。

分析反馈的重要性

分析反馈显著减少了手动工作量，因为我们可以用一个轻量级的 pass 自动化张量程序的模式分析，而传统的单次降级抽象需要大量且不灵活的手动算子注解。

4.3 动态形状感知内存规划（Dynamic Shape–Aware Memory Planning）

内存是现代 ML 应用中的重要资源。大多数 ML 编译器可以通过比较静态形状张量的大小并在提前分配一组固定的内存块来规划内存重用，以减少运行时内存分配。通常，编译器无法对编译时未知形状采取相同的方法，必须依赖运行时内存分配器。

然而，借助符号形状抽象，我们可以分析和比较动态张量形状并相应地规划它们的重用。下图图 10 和算法 3 展示了我们如何应用动态形状感知的内存规划。

我们首先将 call_tir 和 call_dps_library 的外部函数调用降低到显式的内存分配和 DPS 调用，以便我们可以暴露这些分配以进行规划。为了支持动态形状，算法 3 中的 RequestReuseWithSymShape 利用符号表达式分析来证明两个符号表达式是否相等，然后适当地从池中选择可重用的存储。

我们还可以在已知符号值的上界时（例如，用户注解的 LLMs 中的上下文长度）静态分配足够的内存，这允许提前创建静态内存分配计划，即使存在动态形状也是如此。这种可预测的内存使用估计对于在内存受限的后端上部署动态 ML 模型至关重要。此外，静态规划可以提高模型性能，通过启用 CUDA Graph 卸载（第 4.5 节），这依赖于静态内存分配。

4.4 跨层次张量程序工作区提升（Cross-Level Tensor Program Workspace Lifting）

除了图层的内存分配外，有时张量程序可能还需要为中间工作区进行全局内存分配。例如，StreamK 调度的矩阵乘法将矩阵乘法分解为两个约简阶段。

第一阶段：将部分累积结果写入中间全局内存缓冲区。
第二阶段：消费并累积这些结果。

如图 11 所示，我们可以通过分析反馈检测张量程序中的此类全局内存分配，并联合重写张量程序和图层调用站点，将分配提升到图层。重要的是，提升的分配可以由第 4.3 节中的内存规划进行规划，从而进一步增加整体内存重用。只有在 Relax 的跨层次抽象中保留所有跨层次转换中的形状关系时，这种优化才有可能。

4.5 CUDA Graph 卸载（CUDA Graph Offloading）

CUDA Graph 是一种减少 GPU 驱动程序级别内核启动开销的优化，从而提高整体系统性能。它通过捕获多个 GPU 内核启动并将它们作为一个组重放，而不是单独启动每个内核。为了重放捕获的内核启动，GPU 驱动程序要求所有由内核访问的全局内存都是常量大小并提前静态分配。这对于将 CUDA Graph 应用于动态形状 ML 模型提出了重大挑战。

借助静态内存规划，Relax 可以提前静态分配所有内存，即使对于动态形状的张量也是如此。我们构建了一个 pass ，分析计算图并将满足 CUDA Graph 条件的子图提升为子图函数。该通插入处理 CUDA Graph 捕获或重放的运行时内置函数。在运行时，子图函数的首次运行触发 CUDA Graph 捕获；后续运行自动重放捕获的 CUDA Graph。通过这些，我们将通常仅适用于静态模型的 CUDA Graph 扩展到广泛的动态工作负载。

4.6 通过部分降低进行张量算子优化（Tensor Operator Optimizations via Partial Lowering）

现代 ML 框架通过两种方法优化张量算子：将计算卸载到平台特定的算子库，或者利用编译器张量程序优化和代码生成。大多数现有的 ML 编译器在图层和较低层之间做出这些决策，使得很难组合不同的降低方法。

例如，如果我们想引入一个新的算子库，我们需要仔细检查现有的降低策略并相应地更新。Relax 通过部分降低（下图图 12）应用张量算子优化。我们在 Relax 中注册了一组“（子图模式，库函数）”对，并构建了基于模式匹配和重写的 pass，这些 pass 检测图层中的特定模式（例如带尾随处理的矩阵乘法），并将检测到的区域部分降低到外部库调用。

Relax 还允许用户注册模式以实现可定制性。

此外，基于 TensorIR 调度转换，我们构建了一组基于分析的动态形状感知调度规则，以通过最小化内存加载来优化张量程序。我们还可以包含 pass，以应用 Ansor 风格的自动调整，用于基于分析的调度规则无法处理的稀有张量程序（例如复杂的卷积）。重要的是，所有这些转换可以组合在一起并协同工作。Relax 实现了快速开发，只需要一个相对简单的部分降低 pass 来定制算子的优化。

4.7 优化与降低管道（Optimization and Lowering Pipeline）

Relax 在跨层次抽象上使用固定顺序的管道（没有固定点）来优化、降低并最终将端到端模型构建为可运行模块。图 13 展示了一个示例管道。

我们优先部分库降低（第 4.6 节），以利用目标平台上的外部库函数。接下来，我们遍历整个程序，为所有高级算子调用生成张量程序，并将算子调用降低到相应张量程序的 call_tir。然后，我们可以应用算子融合、张量程序工作区提升、内存规划和 CUDA Graph 卸载。

值得注意的是，某些张量程序优化（例如工作区提升）需要在图优化之前应用，这需要 Relax 的跨层次抽象设计。在最后阶段是将模型构建为可运行模块。在图层，一个基本任务是将符号变量与程序输入张量中的具体形状值关联，并在运行时计算符号表达式。我们创建一个整数宿主张量，用于存储程序中所有符号表达式的运行时值。在转换开始时，我们填充程序输入张量中的符号变量值。

然后，我们生成加载符号表达式、计算符号表达式并将结果存储到相应位置的张量程序。最后，当需要将张量形状作为一等值时，我们插入函数调用以构造形状元组。经过此转换后，我们擦除所有注解，留下一个主要由低级函数调用组成的程序。

然后，这些调用将被翻译为一系列虚拟机指令，每个指令都是对生成的或内置函数的调用。对于优化的低级张量程序，我们直接生成相应的 GPU 代码。我们将图层虚拟机指令和 GPU 代码打包在一起，形成一个单一的整体端到端模块，该模块可以在编译目标平台上运行。

5. 实验评估

我们在 Apache TVM 之上实现了 Relax。本节的评估旨在回答以下问题：

Relax 在大型语言模型（LLM）推理性能上是否能与现有框架竞争（第 5.1 节）？
提出的抽象和优化对性能和内存使用的影响是什么（第 5.2 节）？
Relax 是否能够支持这些新兴的 LLM 在更广泛的平台（第 5.3 节）？
Relax 在更广泛的模型集上的表现如何（第 5.4 节）？

5.1 大型语言模型推理评估

本节评估 Relax 在 NVIDIA GPU 和新兴的 AMD 及 Apple GPU 上对端到端 LLM 的性能。我们评估了 LLM 生成解码阶段的每 token 延迟，以包含序列长度和批量大小的动态性。我们的评估在 Llama3-8B、Gemma1.1-7B 和 Qwen2-7B 上进行，使用 float16 权重和激活，在 NVIDIA RTX 4090、AMD Radeon 7900 XTX 和 Apple M2 Ultra 上进行。

我们与以下基线框架进行比较：HuggingFace Transformers（v4.41.2）与 PyTorch（v2.3.1） Eager 模式和编译模式，vLLM（v0.5.0.post1），以及手工优化的 LLM 推理系统 llama.cpp（172c825）。如果可用，基线启用 FlashAttention。我们使用类似 PyTorch 的 nn.Module 接口构建 Relax IR。我们测量生成 32 个 token 的解码时间，并计算每序列的每 token 延迟。重要的是，Relax 仅需编译一次模型，即可处理任意批量大小和序列长度。

图 14 到 16 显示，Relax 在不同平台上始终提供具有竞争力的性能。在这种情况下，符号形状分析使 Relax 能够生成仅在批量大小和序列长度维度上动态的张量程序。跨层次抽象允许无缝组合图和张量程序优化，适用于任何 GPU 后端，消除了为每个单独后端手动编写内核的需求。更重要的是，跨层次抽象使我们能够在批量大小为 1 时使用编译器优化的矩阵-向量乘法张量程序，同时能够对其他批量大小应用部分库降低以利用算子库。

这种灵活性显著提高了特殊情况的性能。值得注意的是，虽然 HuggingFace Transformers 的 PyTorch 编译模式支持动态序列长度，但它仍然需要静态 KV 缓存，这取决于显著的模型定义更改，并且仅对少数模型可用。此外，并非所有基线都很好地支持所有平台。手工优化的 llama.cpp 在 Apple GPU 上表现出色，但在 NVIDIA GPU 上的表现较差。而 PyTorch 编译模式和 vLLM 缺乏对 Apple GPU 的支持。相比之下，Relax 能够以具有竞争力的性能支持所有平台。

5.2 可组合优化的效果

Relax 的抽象允许灵活组合优化，例如 CUDA Graph 卸载、算子融合、部分库降低和代码生成。我们使用 float16 的 Llama3-8B 在 NVIDIA RTX 4090 上评估这些优化。图 17 显示了消融研究的结果。部分库降低贡献最大，在大批量大小时性能提升高达 27%，因为它将重型矩阵乘法（约占所有算子的三分之一）降低到 cuBLAS 库内核。

算子融合通过融合约五分之一的算子（如 RMSNorm 和逐元素加法）来减少启动的内核和 GPU 全局内存访问。CUDA Graph 卸载总体上通过减少 GPU 驱动程序级别的内核启动开销带来了 1-2%的性能提升。所有这些可组合优化共同提高了整个系统的性能。

对内存使用的影响：我们通过测量在 float16 Llama3-8B 预填充阶段和解码阶段的总分配激活内存量来研究静态内存规划的内存减少效果。预填充阶段处理长度为 128、256、512、1024 的连续输入，解码阶段处理大小为 1、16、32 和 64 的连续批次。内存规划根据序列长度和批量大小的上限进行规划。当禁用内存规划时，我们使用运行时内存池来回收未使用的内存。如表 2 所示，静态内存规划在连续预填充阶段减少了 22%的激活内存，在解码阶段减少了 40%。

通过静态内存规划，我们始终在所有输入长度和批量大小之间重用内存，即使它们随时间变化。相比之下，没有内存规划的系统会在输入形状变化时反复分配动态大小的内存，这在实际应用中是不可预测的。这可能导致更高的内存使用，除非所有内存都静态规划。

此外，通过提前分配所有内存，我们能够为新兴模型启用 CUDA Graph，从而获得进一步的性能提升。重要的是，所有这些优化都依赖于所提出的抽象。例如，在没有形状信息的最佳情况下，我们无法运行静态分析，如静态内存规划和静态图捕获以进行优化，这可能导致额外的内存和延迟开销。

5.3 在更多新兴平台上的评估

本节评估将新兴模型部署到现有解决方案支持较少的一系列新兴平台。我们在以下设备上评估单序列 LLM 推理性能：搭载 Apple A16 的 iPhone 14 Pro、搭载 Qualcomm Snapdragon 8 Gen 2 的 Samsung S23、搭载 ARM Mali GPU 的 Orange Pi 5、搭载 AMD APU 的 Steam Deck、NVIDIA Jetson Orin 开发套件，以及在搭载 Apple M3 Max 笔记本上的 WebGPU。我们对大多数情况使用 4 位量化的 Llama3-8B，而对移动电话设备使用 Llama2-7B 以适应 VRAM 限制。

如表 3 所示，Relax 在移动设备上提供超过 5 个 token /秒的吞吐量，在 Orange Pi 5 上为 2.3 个 token /秒。此外，Relax 是第一个在这些平台（除了 NVIDIA Jetson Orin）上启用 GPU 加速 LLM 推理的解决方案。没有内存规划提前分配所有所需内存并保持在预算内，这些模型甚至无法在某些环境中运行，因为内存限制。

我们进一步将 Relax 与 Samsung S23 上的 llama.cpp 进行比较。Relax 实现了高达 55%的更高吞吐量（图 18）。值得注意的是，llama.cpp 仅使用 CPU，因为缺乏针对 Android GPU 的内核，而 Relax 通过编译自动生成优化的 GPU 代码，从而不仅在 Android 平台上部署新兴模型，而且更广泛地在更多新兴平台上部署。

5.4 在额外模型集上的评估

我们还在额外的模型集上研究 Relax。Whisper 是一个实现为编码器-解码器 Transformer 的自动语音识别（ASR）模型。我们评估了使用 Whisper-large-v3 转录 30 秒语音的时间，并将 Relax 与 HuggingFace Transformers、WhisperX、Faster Whisper 和 whisper.cpp 进行比较。

如图 19 所示，Relax 在 NVIDIA 4090 上带来了 14%的速度提升，并在 Apple GPU 上具有竞争力。LLaVA 是一个集成了预训练的 CLIP 视觉编码器和 LLM Vicuna 的大型多模态模型，用于通用的视觉和语言理解。我们评估了使用基线为 HuggingFace Transformers、vLLM 和 llama.cpp 的图像生成 32 个 token 的时间。结果如图 20 所示，Relax 在两个平台上都高效地支持视觉编码器以及 LLM 的预填充和解码阶段。

6. 相关工作

硬件厂商优化的库（如 cuDNN、CUTLASS、MKL-DNN 和 MIOpen）通常被 ML 框架用于支持各种硬件后端的张量操作。这些库是平台特定的，并且需要大量的工程开发成本来覆盖不断增长的算子、数据格式和布局的需求。Relax 通过允许这些库与经过动态形状感知优化的循环级代码一起使用来补充这些库。依赖这些库的框架可以利用 Relax 在库或生成代码之间进行选择。
随着大型语言模型的出现，还出现了一系列针对这些特定工作负载优化的框架。这些框架通常依赖于每个特定后端的手动优化。它们可以利用 Relax 来减少支持更广泛工作负载和新兴后端的工作量。

在循环级代码转换和优化方面也有很多工作。

Triton 和 Graphene 是用于优化 GPU 上张量化程序的抽象。
DietCode、CoRA 和 SparseTIR 专注于具有形状动态性和不规则性的张量程序优化。
Mosaic 是一个结合库分派和稀疏张量程序优化的稀疏编译器。
Cortex 启用了递归计算的张量程序优化。

我们在实现中使用 TensorIR 作为跨层次设计中的张量程序抽象，但我们可以与其他抽象结合以支持更广泛的张量程序优化。

ML 编译器旨在表示和优化端到端模型计算。高级计算通常被表示为计算图风格的方言。

TVM 的 Relay 和 MLIR 方言将动态维度表示为未知，并且不跟踪动态形状关系。
IREE 提供了基于 MLIR 的端到端编译。
Nimble 利用运行时分桶来支持动态算子。
DISC 启用了形状作为一等值，但不跟踪符号形状。
TorchInductor 将原生符号形状支持带入 PyTorch 编译器，专注于为从 TorchDynamo 派生的 TorchFX 图生成内核。
PyTorch 编译器为跟踪的子图存储全局符号变量表，并与其 JIT 中心的设计协同工作，避免跨函数的符号形状跟踪。

Relax 通过抽象和全局跟踪跨函数的符号形状来补充 PyTorch 编译器，从而实现 AOT 编译和更广泛的新兴后端部署。因此，Relax 可以用作 PyTorch、JAX 等框架的后端，以将模型部署到更多新兴后端。

Axon 是一种函数式语言，它在类型中考虑形状推导，并应用约束求解器来确定形状关系；与 Relax 不同，它在无法静态推导形状时不会描述动态回退。（注意，Relax 仍然可以应用类似的约束求解方法，尽管它会增加额外的编译时间成本。）

Halide 通过 Func::define_extern() 在张量程序中支持外部函数调用。Relax 将这种机制扩展到图和张量程序级别，将外部库与这些级别桥接在一起。此外，大多数现有的 ML 编译器遵循多级单次降级方法，而 Relax 通过跨层次抽象实现了跨函数的全局符号形状跟踪。Relax 关于支持动态形状和跨层次优化的见解可以用来改进这些 ML 编译器框架。

7. 结论

我们介绍了 Relax，这是一种用于新兴平台的端到端动态机器学习的抽象。跨层次抽象和一等符号形状使得动态形状模型的可组合优化成为可能，并使我们能够构建一个 AOT 端到端整体框架，将新兴模型部署到各种新兴后端。

Relax 在各个平台上的性能与最先进的系统相当，包括在 NVIDIA GPU 上将 LLM 解码 token 延迟减少 27%。我们希望这项工作将鼓励对动态形状感知程序抽象的进一步研究，并突出 ML 编译器的新可能性。

END

作者：TVM Relex
来源：NeuralTalk

推荐阅读

欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。