GPUMODE-CUDA-笔记-全-

GPUMODE CUDA 笔记（全）

1： 使用 LeanRL 以光速训练小型强化学习模型

在本节课中，我们将要学习如何通过优化 CPU 瓶颈来加速小型强化学习模型的训练。我们将探讨强化学习的基础概念、现有工具库的局限性，并深入介绍一种名为 LeanRL 的实验性解决方案，它通过结合 PyTorch 编译和 CUDA Graphs 技术，显著提升了训练效率。

---

什么是强化学习？🤖

强化学习将问题建模为智能体与环境之间的交互。环境是一个“黑盒”，它接收智能体的动作，并返回观察结果和奖励。智能体的目标是通过学习一系列动作来最大化累积奖励。

第一张图来自 Sutton 和 Barto 的经典著作，它很好地总结了强化学习的核心循环。关键在于，环境是动态的，而非静态数据集。这种动态特性使得强化学习非常强大，理论上可以解决任何能表述为“最大化奖励”的问题。

近年来，人们开始将大语言模型的后训练也视为一个强化学习问题。此时，环境可能是一个包含固定答案的数据集（如数学题），而智能体就是待训练的 LLM。这种框架的通用性极强，可以应用于机器人控制、网页浏览、GPU 内核优化等广泛领域。

强化学习的优点包括：

• 处理复杂问题：当问题难以用数学方程精确描述时，强化学习提供了一种通过最大化奖励来寻找解决方案的途径。
• 持续改进：通常，投入更多数据或任务能带来线性的性能提升。
• 过拟合问题较小：相比监督学习，过拟合有时不那么严重。
• 计算高效：使用像 PPO 这类方法时，无需通过奖励模型进行反向传播，计算开销更小。

然而，强化学习也存在一些挑战：

• 入门门槛高：需要一定的专业知识才能有效运用。
• RLHF 的局限性：基于人类反馈的强化学习像是一个“不完整的补丁”，其效果受限于奖励模型训练数据的质量。

---

构建通用强化学习库的挑战 🧩

上一节我们介绍了强化学习的基本概念，本节中我们来看看在 PyTorch 生态中构建通用强化学习库时遇到的独特挑战。

在 PyTorch 领域，像 TorchVision 这样的库任务相对明确，因为输入（如图像）和输出（如标签）的格式是预先可知且固定的。但强化学习不同，它并非关于某种特定媒体，而是关于一系列算法。用户可能将其用于 LLM 后训练、机器人学、自动驾驶、药物设计等任何领域。

这意味着构建一个“非主观”的、能处理所有情况的库非常困难。另一个核心挑战是“策略”的多样性。在智能体-环境范式中，策略负责根据观察输出动作。但根据所用算法的不同（例如 DQN 与 PPO），策略的实际输出可能是一个值、一个概率分布或其他复杂结构。

此外，在多智能体场景中，复杂性会急剧增加，例如需要考虑集中控制、智能体间通信、团队分组等多种设置。尽管场景复杂，但所需的基础工具（环境、智能体、数据连接和存储）在概念上是相似的。

为了解决输入/输出格式多变的问题，我们提出了一个方案：让库中的所有组件都通过字典进行通信。这样，策略只需要负责选择动作，而不必关心具体实现细节；环境也只需提供重置和步进等方法，无需严格限定签名。

但直接在 Python 中处理字典会很繁琐，例如需要堆叠字典、处理嵌套键等。因此，我们在其基础上构建了一个名为 TensorDict 的类。

---

TensorDict：字典与张量的结合 🔄

TensorDict 是字典和张量的混合体。你可以将其视为一个张量般的对象，它具备许多张量的特性。

以下是 TensorDict 的关键优势：

• 批量操作便利性：你可以像操作单个张量一样，对一组张量执行批量操作（如加法、乘法）。例如，tensordict1 + tensordict2。
• 高效设备转移：你可以将整个 TensorDict 同步转移到 CUDA 设备，操作快速且能避免 CUDA 同步相关的错误。
• 丰富的工具集：它提供了许多工具函数，用于处理张量的重塑、共享内存存储等。

这解决了一个“一对多”的问题。例如，当你为单个张量写了一个优化器更新规则后，可以轻松地将其应用于一个包含多个张量的 TensorDict，而无需使用复杂的树状结构来分发操作。这种设计思维是：任何我想对单个张量做的操作，都可能想对一批张量做。

虽然也有人考虑过使用 JSON、Pickle 或原始指针等方案，但 TensorDict 的面向对象设计使得消息传递更加方便，并且为底层优化提供了可能。所有算术运算都使用 torch.vmap，速度远快于普通的 Python 循环。同时，它能与 torch.compile 兼容，确保整个计算图可以被编译和优化。

目前 TensorDict 是完全基于 Python 的，因为结合 torch.compile 后，性能开销已不再是问题。我们计划未来将其上游到 PyTorch 主库中。

基于 TensorDict，我们构建了 TorchRL 库。

---

TorchRL：模块化的强化学习库 🧰

TorchRL 的目标是提供高度可复用的组件。由于所有组件都基于 TensorDict 这一通用抽象，因此它们的接口非常灵活。

TorchRL 的核心组件包括：

• Actor 抽象：用于执行策略的模块。
• 环境抽象：高度通用，可以封装游戏、数据集（如数学题）等。
• 经验回放缓冲区：一个动态的数据集，支持读写操作，数据会随时间变化。
• 损失函数模块：各种强化学习算法的损失计算模块。

该库支持单节点单进程、单节点多进程以及多节点部署，并集成了多种分布式后端。其最大的特点是模块化。用户可以选择只使用库中的特定组件（例如仅使用经验回放缓冲区），而不必引入整个框架。这与生态系统中许多其他库不同，那些库通常要求用户全盘接受其整个训练流程。

---

小型模型的 CPU 瓶颈问题 ⚡️

上一节我们介绍了 TorchRL 库的模块化设计，本节中我们来看看强化学习中一个突出的性能问题：小型模型的 CPU 瓶颈。

强化学习社区长期以来一直关注小型模型（如简单的多层感知机 MLP）。早期是因为与环境的交互数据有限，不需要大模型。如今，虽然有了可以在 GPU 上大规模并行运行的仿真器，每秒可收集数十万甚至数百万帧数据，但在实际部署时（例如在机器人上），仍然需要模型小而快，无法承受数十亿参数大模型的推理延迟。

在 Python 中训练这些小模型时，一个主要问题是 CPU 开销过大。研究人员喜欢用 Python 快速实验，但 Python 的解释执行和与 C++ 层的频繁交互会导致显著开销，使得 GPU 利用率可能低至 10%，这是非常低效的。

torch.compile 通过融合 C++ 操作、减少返回 Python 层的次数来提升性能。但对于小型模型，编译引入的“守卫”检查开销可能比模型实际执行时间还长。守卫检查是为了确保编译后的代码在输入条件（如张量形状、字符串参数）未改变时才能被复用。

例如，编译一个简单的填充操作时，如果填充模式和填充值是字符串或整数，编译器会为这些具体值生成特化代码。如果后续输入中这些值发生变化，就需要触发新的编译和守卫检查。对于大模型，矩阵运算是主要开销，守卫检查的影响微乎其微。但对于一个在 CUDA 上只有四层、每层 500 个神经元的小型 MLP，守卫检查的时间可能超过模型前向传播的时间，导致编译后速度没有提升甚至更慢。

---

解决方案：LeanRL 与 CUDA Graphs 🚀

那么，如何为小型强化学习模型打破 CPU 瓶颈呢？解决方案是结合使用 torch.compile 和 CUDA Graphs。

CUDA Graphs 允许你将一系列 CUDA 内核（操作）封装成一个图，然后一次性提交执行，避免了单个内核启动的开销。如果只使用 torch.compile（其底层会使用 CUDA Graphs），你仍然会在图的前后有守卫检查。

我们的洞察是：在机器人、游戏等强化学习场景中，输入结构通常是高度可预测的。机器人有多少个关节、摄像头，LiDAR 的输出尺寸等，在训练过程中往往是固定的。批大小也可以固定，或者使用可预测的动态形状。

如果我们能做出“输入结构不变”的保证，就可以消除守卫检查。具体做法是：先使用 torch.compile 编译代码，然后在编译好的代码外部包裹一层 CUDA Graph，而不是让编译过程内部去管理 CUDA Graph。这样就能完全绕过守卫检查。

我们在一款名为 LeanRL 的实验性库中实现了这一想法。LeanRL 基于一个包含流行 RL 算法简洁实现的库，我们将其 fork 并应用了编译和 CUDA Graphs 加速。

以下是性能对比数据（以 TD3 算法为例）：

• Eager 模式（纯 Python）：247 帧/秒
• 仅编译：272 帧/秒（提升很小）
• 仅 CUDA Graph（无编译）：778 帧/秒（显著提升，仅消除了 Python 开销）
• 编译 + CUDA Graph：速度进一步提升
• JAX 版本：561 帧/秒

在这个例子中，PyTorch 方案展示了显著优势。我们还测量了训练过程中的 GPU 利用率曲线下面积。使用“编译+CUDA Graph”方案时，GPU 平均利用率更高，但总的使用量（面积）更少，这意味着我们更高效地利用了 GPU。

LeanRL 已成功帮助多个研究团队将训练速度提升了高达 6 倍。

---

最佳实践与生态系统 🌍

在结束之前，让我们看看强化学习领域的其他高效库，并总结一些最佳实践。

现有的高效强化学习库包括：

• SB3：基于 JAX 的算法库。
• Parlay：基于 C++ 的环境库，通过共享内存实现高速仿真，可达每秒数百万帧。
• RLtools：一个极简的 C++ 库，专注于少数几种算法（如 PPO），旨在为嵌入式或机器人平台提供极快的训练和部署。
• Julia 生态中的库：也提供了高性能的解决方案。

对于希望优化自己代码的开发者，我们建议遵循以下步骤：

1. 首先尝试 torch.compile：使用 mode="reduce-overhead" 编译你的代码，观察速度提升。
2. 如果效果不佳：考虑使用类似 LeanRL 中的技术，手动为你的更新函数或策略创建“编译后+外部 CUDA Graph 包裹”的版本。
3. 避免常见陷阱：
   • 直接在 CUDA 设备上创建张量，而不是在 CPU 创建后再转移。
   • 避免在训练循环中调用 tensor.item() 进行日志记录，这会触发昂贵的设备同步。
   • 多使用分析工具查看性能轨迹。

---

总结 🎯

本节课中我们一起学习了如何加速小型强化学习模型的训练。我们从强化学习的基本概念和挑战出发，介绍了 TensorDict 和 TorchRL 如何通过模块化设计提供灵活性。然后，我们深入探讨了小型模型训练中的 CPU 瓶颈问题，并揭示了 torch.compile 守卫检查对微小操作的开销影响。最后，我们介绍了 LeanRL 的解决方案：通过将编译后的代码用 CUDA Graph 外部包裹，彻底移除守卫检查，从而实现了显著的性能提升（最高可达 6-7 倍）。我们还回顾了生态系统中的其他高效工具，并总结了一些实用的优化最佳实践。希望这些知识能帮助你更高效地进行强化学习研究和开发。

2：Modular的统一设备加速器语言

概述

在本节课中，我们将深入探讨现代GPU编程的现状、挑战以及Modular公司提出的全新解决方案——Mojo编程语言。我们将学习Mojo如何通过强大的元编程能力、统一的编程模型以及对性能与可移植性的极致追求，旨在解决当前AI与异构计算领域面临的多语言、多系统复杂性问题。课程内容涵盖从高级概念到具体代码示例，特别是Mojo在GPU内核编程、性能优化和跨硬件支持方面的独特优势。

GPU软件现状与挑战

上一节我们介绍了课程的整体框架，本节中我们来看看当前GPU编程领域面临的核心问题。

现代AI软件栈面临“多语言”困境。为了获得高性能，开发者经常需要在Python（用于快速原型和建模）、C++/Rust（用于系统级性能）以及CUDA（用于GPU硬件）等不同语言和范式之间不断重写代码。这种分裂不仅成本高昂，还导致团队需要在不同思维模式间切换，降低了整体效率。

由于这种分裂，现代AI系统栈往往由许多不同时期、不同目标的技术拼凑而成，形成了一个复杂且脆弱的“缝合怪”。每一层都在独立演化，缺乏统一设计，导致在特定硬件、特定量化方案或特定用例下，常常出现性能不佳、兼容性问题或难以调试的故障。

对于专注于内核编程和GPU编程的开发者而言，三个核心目标至关重要：

1. 性能：充分利用昂贵的GPU硬件，获取峰值算力。
2. 可用性：提高开发效率，降低GPU编程门槛，让更多人能够参与。
3. 可移植性：代码能够跨不同硬件（如NVIDIA、AMD）甚至未来新型加速器运行。

然而，现有主流方案往往只能满足其中一两个目标：

• CUDA：性能王者，但专属于NVIDIA硬件，可移植性差。近期在可用性上有所改进。
• Triton：显著提升了可用性和一定的可移植性，但在追求极致性能时需要做出妥协。
• OpenCL：注重可移植性（支持FPGA等），但在张量计算性能和现代开发体验上已落后。
• ROCm：在AMD硬件上性能优异，但未聚焦于跨厂商的可用性和可移植性。

目前缺乏一个能同时完美满足性能、可用性和可移植性三大目标的解决方案。

传统编译器方案的局限性

上一节我们分析了现有方案的不足，本节中我们来看看另一种常见的解决思路——AI编译器——及其面临的挑战。

作为编译器领域的专家，我们曾长期致力于构建“足够智能的编译器”来自动化解决性能与可移植性问题，例如XLA等项目。然而，AI领域算法迭代速度极快，而全球编译器工程师的数量有限，且他们通常并非算法和微架构领域的专家。这导致编译器成为创新的瓶颈。

更重要的是，成功的AI生态系统需要算法研究员、内核工程师、AI应用开发者等不同角色高效协作。传统的AI编译器方案试图将所有复杂性封装在编译器内部，迫使内核工程师去理解中间表示等底层细节，这背离了让各领域专家专注于自身所长、通过工具高效协作的初衷。

Mojo的愿景：统一、原生的异构计算平台

上一节我们探讨了传统方案的瓶颈，本节中我们将介绍Modular公司提出的全新愿景和解决方案。

我们需要的是一个统一、原生的生成式AI软件平台，它需要从第一性原理出发重新构建，吸收现有各系统的经验教训。这个平台的目标是：

• 支持来自不同厂商的多种GPU。
• 释放硅芯片的全部性能潜力。
• 提供卓越的开发者体验和可用性。

这正是Modular公司正在构建的。Mojo不仅仅是一个新的编程语言，它更是一个全新的、为异构计算和现代AI用例从头设计的完整软件栈，涵盖了从底层内核到上层服务与集群管理的整个垂直领域。

目前，Mojo已经实现了对CPU、NVIDIA GPU（包括消费级和数据中心级）以及AMD GPU（如MI300系列）的全面支持。重要的是，Mojo不依赖于CUDA等现有生态，而是基于第一性原理构建了自己的技术栈，确保了统一性和控制力。

为什么选择构建全新的Mojo语言？

上一节我们了解了Mojo的宏大愿景，本节中我们来探讨实现这一愿景的核心载体——Mojo语言本身。

构建一门全新的通用编程语言（而非领域特定语言DSL）通常被认为是一个糟糕的主意，因为它面临采用成本高、构建难度大、需要长期一致的愿景和大量资源投入等挑战。然而，我们认为异构计算正是需要这种“硬核”方案的正确时机和领域。

构建新语言是获得最佳解决方案的唯一途径。我们想要创造的是一个真正强大、能产生巨大影响、能让AI开发者和内核工程师都爱不释手的工具。Mojo的目标不是做一个临时的演示或研究性质的项目，而是旨在提升整个行业水平的基础设施。

Mojo的成功建立在几个关键因素上：清晰的愿景、多年的持续投入、顶尖的团队以及我们在LLVM、Clang、Swift等语言和编译器项目上的丰富经验。这一切使得我们有信心迎接这个挑战。

Mojo语言初探

上一节我们解释了构建新语言的必要性，本节中我们来具体看看Mojo是什么。

Mojo是一门具有Python风格的系统编程语言。

• Python风格：它采用了Python的语法和生态系统，降低了学习成本。
• 系统编程：它拥有与C++、Rust相媲美的性能，是一个完全通用的系统编程语言，但并非使用Python解释器，而是使用了Python的语法。

Mojo包含完整的工具链，如编辑器支持、调试器、性能分析器集成等。它的目标是为Python生态的开发者赋予新的超能力，同时将编程语言和编译器设计的最新成果结合起来，形成一种新颖的组合。

尽管代码看起来像Python，但Mojo是一门全新的语言，它并非Python的DSL。

Mojo的核心：强大的元编程

上一节我们介绍了Mojo的基本概念，本节中我们将深入其最核心的特性——元编程，这是理解Mojo如何实现其目标的关键。

元编程是Mojo赋予开发者超能力的核心机制。与传统AI编译器将权力收归编译器内部不同，Mojo通过元编程将这些能力交还给内核工程师。这使得开发者无需等待编译器团队支持，就能自行实现新数据类型、硬件指令或优化策略。

Mojo的元编程主要基于两个概念：参数和别名。

• 参数：使用方括号 [] 声明，在编译时确定。类似于C++的模板参数，但更强大、更易用。

  struct SIMD[DT: DType, width: Int]:
      var value: … # 内部实现
  

• 别名：使用 alias 关键字声明，是编译时可计算的任意表达式变量。

  alias float32 = DType.float32 # float32 是在库中定义的别名
  

Mojo的关键突破在于统一了运行时语言和元语言。这意味着你可以在编译时使用for循环、if语句等几乎所有运行时能用的结构，进行符号执行并将结果直接嵌入二进制。这允许将大量复杂的计算（如张量索引数学）从运行时转移到编译时，实现零开销抽象。

以下是一个简单的例子，展示如何在编译时生成一个列表：

fn fill_list(lb: Int, ub: Int) -> List[Int]:
    var result = List[Int]()
    for i in range(lb, ub):
        result.append(i)
    return result

alias vec = fill_list(5, 20) # 在编译时计算，vec 是一个编译时常量列表

在更复杂的场景中，比如实现一个SIMD水平归约操作，你可以利用编译时if语句为不同架构、不同数据类型选择最优的实现路径，就像在编写一个专属的小型编译器。这赋予了库作者极大的灵活性和控制力。

此外，Mojo的元编程还支持高阶函数作为编译时参数，使得算法可以高度可配置和可组合。例如，可以编写一个通用的向量化算法，它接受一个表示元素操作（如sin, cos）的函数作为编译时参数，并确保循环被正确向量化，而不是依赖编译器的提示（如OpenMP的pragma），后者可能失败。

这意味着：热爱性能、内核、硬件的开发者，不再需要成为编译器专家。你可以发明编译器中没有的新优化，可以无需修改编译器就支持新硬件指令，如果现有库不满足需求，你可以直接绕过它们，编写自己的实现甚至内联汇编。Mojo提供了直达硬件的路径和提升生产力的工具。

Mojo GPU编程实战

上一节我们深入探讨了元编程这一核心概念，本节中我们来看看如何用Mojo实际编写GPU代码。

Mojo统一了CPU和GPU编程。一个函数可以轻松地在CPU或GPU上启动。

fn say_hi(value: Int):
    print("Hello", value)

# CPU 调用
say_hi(42)
# GPU 调用
say_hi[grid_dim=(1,), block_dim=(1,)](42)

Mojo中的所有基础类型和数学函数都是在库中实现的，而非编译器内置。例如，exp2函数的实现会根据目标硬件和数据类型（如f16, bf16）选择最优的路径，可能调用特定的PTX内联汇编或LLVM内部函数。这一切都通过编译时参数和条件判断完成，无需修改编译器。

fn exp2[width: Int](x: SIMD[DType.float32, width]) -> SIMD[DType.float32, width]:
    @parameter
    if width == 1:
        return _call_scalar_ptx("ex2.approx.f32", x)
    # ... 其他优化路径

Mojo提供了高级抽象来简化编程。例如，foreach（或elementwise）函数允许你以声明式的方式编写逐元素操作，而库会自动处理在CPU（多线程、向量化）或GPU（线程索引、块划分、向量化加载）上的并行化细节。其内部实现同样是可探查、可修改的库代码。

Mojo的特征系统（类似于接口或协议）提供了强大的组合能力。例如，在实现注意力机制时，可以定义一个MHAMask特征，然后为因果掩码、滑动窗口掩码、分块掩码等不同变体提供具体实现。多头部意力内核可以泛化地接受任何满足MHAMask特征的类型，从而使算法保持清晰，避免充斥条件判断。特征还支持组合（如与、或操作），进一步增强了灵活性。

基于分块的编程与操作组合

上一节我们看到了Mojo如何简化常见的GPU编程模式，本节中我们关注更复杂的、性能关键的模式：基于分块的编程。

矩阵乘法、注意力等操作是AI工作负载的核心。Mojo从Cutlass、CUTE等库中汲取灵感，引入了强大的布局和分块张量抽象。布局定义了从多维坐标到线性内存地址的映射，支持行主序、列主序、旋转、Swizzle等复杂模式。这些抽象全部在库中实现，是可调试、可可视化的，而非编译器魔法。

基于这些抽象，可以实现高性能的矩阵乘法内核。代码结构清晰，通过参数化来适应不同硬件（如A100、H100、MI300）的特定流水线策略（如双缓冲、异步拷贝、Warp专业化）。矩阵乘法的核心（MMA操作）也被抽象出来，根据硬件选择对应的张量核心指令。

更重要的是，Mojo强调操作的组合。在真实场景中，矩阵乘法之后往往会接一个激活函数（如SiLU）。传统的库方案需要为每个“矩阵乘+激活”的组合编写专门的内核，导致算子爆炸。Mojo结合其图编译器，可以自动将基础原语（加、乘、exp、sigmoid）融合成单一高效内核。开发者只需定义基础原语，系统便能自动合成如SiLU这样的复杂操作，实现了“最好的内核是你根本不用写的那个”。

性能表现

上一节我们探讨了高级抽象和组合，本节中我们关心最实际的问题：Mojo的性能到底如何？

我们追求的是业界一流的性能。以下是一些早期基准测试结果（仍在持续优化中）：

• 对比 cuBLAS (NVIDIA)：在多种对于LLaMA 70B等模型重要的矩阵乘法形状上，Mojo的实现接近甚至有时超过cuBLAS的性能。请注意，Mojo的矩阵乘法内核设计允许灵活地融合后续激活操作，而cuBLAS的接口通常不支持，因此在真实场景中Mojo可能更具优势。
• 对比 ROCm/hipBLAS (AMD)：在AMD MI300上，Mojo的实现同样瞄准并达到了hipBLAS和rocBLAS的先进水平，在不少情况下表现更优。
• 对比 Triton：在注意力内核上，Mojo在MI300上展示了优于Triton的性能。

这些数据表明，Mojo在提供高可用性和可移植性的同时，并没有牺牲性能，真正致力于同时实现性能、可用性、可移植性这三个目标。

开发工具与调试

上一节我们验证了Mojo的性能潜力，本节中我们来看看支撑高效开发的工具链。

Mojo提供了完整的现代开发体验：

• 调试器：支持源码级调试、断点、变量查看、调用栈跟踪。
• 性能分析：与NVIDIA Nsight Compute、ROCgdb等工具无缝集成，性能剖析结果可以直接映射回Mojo源代码，方便定位热点。
• 交叉编译与代码生成：无需实际硬件，即可编译并输出针对特定目标（如NVIDIA B200，AMD GPU）的PTX或汇编代码，便于研究和优化。
• 编辑器支持：提供语法高亮、代码格式化、LSP支持等。

对比与总结

上一节我们了解了Mojo强大的工具链，本节中我们将其与现有技术进行整体对比并总结。

与Python DSL（如Triton）相比，Mojo在保持高生产力的同时，提供了更强的性能、更通用的编程能力、更强大的元编程以及更丰富的工具链。Mojo不是将控制权交给某个黑盒编译器，而是将能力赋予开发者。

与CUDA C++相比，Mojo提供了更好的可移植性、更现代的语法、更安全的特性（如借用检查器）以及更快的编译速度。

与库方案（如cuBLAS）相比，Mojo的开源内核库（可能是最大的开源GPU内核库之一）提供了透明性和可定制性，并且其抽象允许更灵活的操作融合。

Mojo的核心理念是构建一个统一、原生的生成式AI软件平台。它通过一门为异构计算从头设计的语言，将编译器能力下沉到库中，赋予开发者超能力，旨在解决AI时代的软件分裂问题。

生态、许可与未来

本节课我们一起学习了Mojo语言的理念、特性与能力。最后，我们来了解一下它的生态现状。

Mojo及其内核库已在Apache 2.0（附LLVM例外）许可证下开源。你可以免费下载并使用Mojo和MAX进行非商业或商业项目（有一些简单的归属要求）。目前主要支持NVIDIA（Ampere, Hopper, Blackwell）和AMD（MI300系列）GPU，并在CPU上表现优异。社区正在快速成长，Modular也定期举办黑客松等活动鼓励开发者参与。

编译器核心计划在明年年底前开源，当前优先确保其稳定性和一致性。

总结

在本节课中，我们一起深入探讨了Mojo编程语言如何试图重塑GPU和异构计算的编程范式。我们学习了：

1. 当前GPU编程的挑战：多语言分裂、系统复杂、难以兼顾性能、可用性与可移植性。
2. Mojo的解决方案：构建一个统一、原生的全栈平台，其核心是一门Python风格的、高性能的系统编程语言。
3. 元编程的超能力：通过编译时参数、别名等特性，将传统编译器的能力赋予库开发者和内核工程师，实现硬件抽象、算法优化和零开销抽象。
4. 实际的GPU编程：从简单的元素级操作到复杂的分块矩阵乘法，Mojo提供了从高级抽象到底层控制的全频谱编程能力。
5. 性能验证：早期基准测试表明，Mojo能够在保持高开发效率的同时，追求并达到业界一流的性能水平。
6. 完整的工具链：提供调试、性能剖析、交叉编译等现代开发工具。

Mojo代表了一种大胆的尝试：通过重新思考编程语言和系统设计的基础，来应对AI时代异构计算的根本性挑战。其成功与否取决于社区的采纳和共建，但无疑为开发者提供了一个强大且富有潜力的新选择。

3：基准测试的传说 🧪

在本节课中，我们将学习如何为CUDA内核进行准确且可靠的性能基准测试。我们将探讨基准测试与性能分析器的区别，揭示在测量GPU内核性能时可能遇到的各种陷阱，并介绍一个名为NVBench的工具，它可以帮助我们自动化处理这些复杂问题。

什么是基准测试？

基准测试是一种工具，用于回答关于代码内核最基本的问题之一。首先要问的问题是：你的内核是否提供了正确的结果？这类问题通常由单元测试来回答。一旦解决了正确性问题，就可以继续探讨性能问题：我的内核有多快？这类问题通常由性能分析器和基准测试来回答。

虽然性能分析器和基准测试在技术上提供了正确的答案，但发现答案的过程往往像格林童话一样充满曲折，因此有了本讲的标题。这个领域初看似乎很熟悉，但随着深入，它会变成一场噩梦。

为 device_select 评估性能

让我们尝试为 device_select 算法回答性能问题。具体来说，假设你有一个应用程序，在某个地方调用了 cub::device::select。如果你不熟悉这个算法，它的作用是：接收一个输入范围和一个谓词（例如，选择所有奇数），然后将符合条件的元素原地紧凑排列到输出范围中。

假设这个内核提供了正确的结果。现在我们的问题是：它快吗？在应用程序层面，我们如何弄清楚这一点？

使用性能分析器

你应该从使用性能分析器开始。NVIDIA Nsight Systems 是一个提供CPU和GPU活动系统级视图的分析器。它会可视化地表示异步内核、内存传输等活动。

当我们对应用程序使用 Nsight Systems 时，会看到一些以微秒为单位的测量值。通过查看时间线，我们可以发现该算法由两个内核组成，总共耗时约92微秒。

但这并不是唯一可用的工具。NVIDIA Nsight Compute 提供了内核的详细性能视图，包括与源代码级别的性能指标关联以及性能优化指导分析。然而，当我们使用 Nsight Compute 时，两个内核的持续时间总和达到了114微秒，比 Nsight Systems 报告的要长。

这是因为 Nsight Compute 默认将GPU时钟频率锁定在基础值（例如900 MHz），而实际使用的GPU（如A6000）可以加速到更高的频率。这导致了性能下降。我们可以通过告诉 Nsight Compute 避免锁定时钟来改善，这显著减少了算法的运行时间，降至约57微秒。但这仍然与 Nsight Systems 报告的92微秒不符。

理解差异：延迟加载

要回答这个问题，必须理解启动内核前需要发生什么。最初，设备代码存储在磁盘上。我们必须将数据加载到RAM，处理它，并将其放入GPU的全局内存，然后才能调用内核。以前，所有设备代码都在应用程序启动时加载。这意味着如果你的应用程序有数千个预编译的内核，但只使用其中几个，就会产生大量的启动开销。

因此，在 CUDA 11.3 之后，默认切换为设备代码的延迟加载。这意味着设备代码将在首次使用内核时被处理和加载到全局内存。这需要大量时间，从而拖慢了性能分析。

我们可以通过设置环境变量 CUDA_MODULE_LOADING=EAGER 来控制此行为。在这种模式下，所有内核在应用程序启动时再次加载，我们不会将模块加载时间计入分析。这有助于我们看到 Nsight Compute 和 Nsight Systems 现在报告大致相同的时间。

仍然存在微小差异，这是因为 Nsight Compute 提供的是每个内核的结果，而 Nsight Systems 给出的是执行由多个内核组成的算法所需的实际时间跨度。内核之间的微小差异是由内核拆卸和设置时间造成的，目前 Nsight Compute 没有包含这部分时间。

基准测试与性能分析

现在，我们可以稍微转变一下视角。在应用程序层面，当我们使用加速算法时，我们可以问：这个内核在我特定的用例中速度如何？例如，对于 device_select，如果50%的数据被选中，问题规模是2亿个元素，数据类型是 int8_t。这类问题由性能分析器回答。

但是，当你开发一个内核时，你的视角会转变为：我的内核在各种用例下的速度如何？你需要一种针对内核性能的“单元测试”，覆盖更广泛的工作负载。对于 device_select，如果只有1个选中项或2亿个选中项，性能会怎样？如果数据类型是 double 呢？这有助于你在开发时更局部地思考内核性能。

因此，我们可以在性能分析器和基准测试之间划清界限。性能分析器本质上是性能调试器，它们提供对特定执行的单次深入分析，并具有更好的硬件可见性。而基准测试则是一种性能单元测试，可用于性能CI，通常覆盖广泛的工作负载，如输入大小、数据类型等。

尝试编写基准测试

现在，让我们尝试编写一个基准测试。我们需要自己计时。一个简单的尝试是使用 std::chrono 包裹内核调用。

auto start = std::chrono::high_resolution_clock::now();
cub::device::select(...);
auto end = std::chrono::high_resolution_clock::now();

但这样会得到不切实际的快速结果（约4微秒）。这是因为 device_select 是异步的，我们计时的只是启动内核的时间，而不是算法运行的时间。

为了修正这一点，我们可以在计时前后添加 cudaDeviceSynchronize() 调用。这确保了我们在开始计时前等待任何先前的异步操作完成，并在算法结束后再记录结束时间。现在我们看到了更现实的100微秒，但这仍然是 Nsight Systems（启用急切加载模式）报告时间的两倍。

这是因为我们仍然处于延迟加载模式。为了绕过这个问题，我们可以进行一次预热运行：先调用一次 cub::device::select 并同步，然后再进行计时。这有所帮助，现在我们看到57微秒，但这比 Nsight Systems 急切模式下的结果还要快。

这是因为第一次调用触发了设备代码加载到全局内存，同时也预热了所有缓存。算法将要读取的所有数据在第二次调用时已经驻留在L2缓存中。我们无意中将算法置于了一个理想场景。为了绕过这一点，我们可以尝试刷新缓存。一个简单的方法是写入大量随机数据以驱逐L2缓存中的任何内容。

这有所帮助，现在我们回到了66微秒，比之前慢了，但仍然比 Nsight Systems 急切模式的60微秒要慢。

提高计时精度：CUDA 事件

要回答这个问题，让我们看看当前的基准测试方案。我们使用 chrono::now() 记录开始时间 T0，然后启动所有内核，接着调用 cudaDeviceSynchronize()（在GPU端完成工作的时刻记为 Ts），最后用 chrono::now() 记录结束时间 T1。

问题在于，我们希望在 Ts（内核实际结束）时刻尽可能接近地获得结束测量值 T1。但 cudaDeviceSynchronize() 做了更多工作：它等待GPU完成工作，然后同步CPU，这显然需要一些时间。因此，T1 和 Ts 之间存在开销。

我们可以依靠CUDA事件来消除这种开销。CUDA事件直接在GPU时间线上记录，不涉及CPU同步。在代码中，可以这样实现：

cudaEvent_t start, stop;
cudaEventCreate(&start);
cudaEventCreate(&stop);

// 预热和刷新缓存...
cudaEventRecord(start);
cub::device::select(...);
cudaEventRecord(stop);
cudaEventSynchronize(stop);

float milliseconds = 0;
cudaEventElapsedTime(&milliseconds, start, stop);

这给了我们更好的结果，现在更接近急切模式下的 Nsight Systems 测量值。我们移除了一些开销，但仍然存在差距。

消除内核启动开销：流阻塞

让我们再次审视当前方案。开始事件的记录点 T0 距离GPU上内核的实际开始时间 Tk 相当远。这个开销是由CPU端启动内核所需的时间（称为内核启动延迟）造成的。我们对这个时间不感兴趣，因为我们试图了解GPU端完成工作需要多长时间，而不是启动工作需要多长时间。

我们如何让 T0 更接近 Tk 呢？我们可以尝试通过阻塞CUDA流来“欺骗”系统。这样，我们可以在流被阻塞时记录开始事件，然后启动所有内核，记录结束事件，最后解除流的阻塞。如果我们能做到这一点，就能让 T0 非常接近内核的开始时间。

通过结合预热、刷新缓存、使用CUDA事件和流阻塞，我们最终得到了与 Nsight Systems 报告非常相似的时间。这是相当准确的测量。

测量的噪声与分布

但这还不是故事的全部。这也是本次演讲标题“基准测试传说”的第二个原因。如果你多次运行基准测试，实际上会得到一个测量值的分布，而不是单一值。像任何童话一样，部分真实情况是，880微秒确实存在于分布中的某个位置，但内核可能运行得更快或更慢。

童话的部分在于，我们可以用单一值来代表整个分布。在这一点上，我们快速回顾一下。我们已经明白，性能分析器是性能调试器，基准测试是性能单元测试。我们还理解了准确计时CUDA内核需要预热运行、刷新缓存、复杂使用CUDA流和事件，而这些最好留给专门的工具。但最重要的是，我们意识到测量是有噪声的，这导致需要进行多次测量。

如果我们进行多次测量，就必须问几个问题：第一，我们能否或应否尝试减少数据中的噪声？第二，我们需要进行多少次测量？

测量次数与节流

简单的方法是直接运行内核很多次（例如10万次），然后计算平均运行时间。但如果我们这样做，会看到完全奇怪的结果：现在内核运行需要毫秒级的时间，而之前我们还在800微秒的范围内。发生了什么？

当我们收集测量值时，可以看到GPU频率开始下降。这是因为我们的 select 实现非常高效，导致GPU发热。随着GPU变热，它开始降低频率，这被称为节流。这种节流影响了 device_select 的执行时间，导致运行时间有大约35%的差异。

我们该如何处理这种差异？这种差异仅仅是由节流引起的，还是有其他因素？为了回答这个问题，我们可以尝试按频率范围对数据进行分桶。首先，这帮助我们认识到 device_select 的实现与GPU频率呈强逆线性相关：GPU工作得越慢，算法运行时间越长。

其次，我们可以观察到，在相同的频率范围内，算法的运行时间仍有约2%的变异。这是由 device_select 内部的并发机制造成的。该算法内部有一些复杂的并发机制，导致其性能不具有确定性。因此，我们必须进行多次基准测试，以找出该算法实际所处的代表性性能状态集合。

从这个图中我们还可以得到另一个启示。深绿色部分是通过使用均匀分布并选择所有元素来运行算法观察到的。如果我们改为使用一个接近零的常数，并选择该常数，使得写入内存的元素数量相同，内核中所有代码路径也相同，但可以观察到运行时间显著下降。为什么会发生这种情况？为什么相同的代码路径会导致不同的频率缩放和不同的执行时间？

我选择这个工作负载是为了说明，接近零的值意味着该字段中的所有位都是零，而零意味着芯片上没有电流。电流更少意味着热量更少。因此，仅仅通过选择不影响任何代码路径的不同工作负载或输入，就可以观察到不同的热节流行为。

这对我们来说意味着，在进行基准测试时，花时间准备一些现实的工作负载至关重要，而不是仅仅提供一些随机的人工常数。

锁定时钟频率

无论如何，我们能否减少这种差异？假设我们确实需要对均匀分布进行基准测试，因为它更现实，更能代表应用程序使用 device_select 的情况。我们能在硬件层面减少这种差异吗？

我们可以尝试将时钟锁定在尽可能接近用户设置的值，比如查看3.1 GHz。但更快的GPU会更快地发热，所以我们仍然会有一个小的频率下降范围。总体而言，这并没有真正减少10万次内核调用中运行时间的差异。这并不奇怪，因为GPU如果发热，就必须降频。

那么，查看较低的时钟频率呢？例如2.5 GHz或2.2 GHz。但这次让我们使用一个不同的算法：device_histogram。这再次导致了一些令人惊讶的结果。

在完全相同的工作负载、问题规模、数据类型和所有值都相同的情况下，我在2.5 GHz和2.2 GHz下进行基准测试。数据显示，降低时钟频率后，算法运行得更快了，有大约6%的速度提升。这有些令人惊讶，为什么会这样？

histogram 是一个重度使用原子操作的算法。在某个时刻，我们会向内存发出大量原子操作，这会导致争用。在更高的时钟频率下，我们以更高的速率发出原子操作是合理的。但如果我们降低GPU端的频率而保持内存频率不变，我们发出原子操作的速率就会降低，这意味着争用减少，这可能对性能产生积极影响，正如我们在屏幕上数据中看到的那样。

如果我们现在绘制每个GPU频率下的运行时间分布，会看到一些有趣的现象。起初，频率越高，运行时间越短，这是合理的。但当你达到2.2 GHz左右时，运行算法的时间开始增长，然后略有下降，这是一种非线性的行为。

稳定测量的意义

这对我们意味着什么？我们能查看较低的时钟频率吗？现在是时候问一个问题了：我们为什么首先需要准确的测量或稳定的测量？通常，当你试图比较某些东西时，才会对稳定的测量感兴趣。例如，你可能想知道我的内核性能是否退步了？或者我的更改是否带来了性能改进？但是，当运行时间对频率存在非线性依赖时，答案可能会令人惊讶。

以 histogram 为例。这次我们关注一个更改：我们希望将线程块大小从300个线程改为1000个线程。这个更改将如何影响性能？

在左侧（2.2 GHz GPU），更改导致性能退步。从300个线程增加到1000个线程，性能下降了。这有一定道理，因为我们没有发出足够的原子操作来造成内存端争用，因此增加线程块大小只是降低了占用率，导致并发工作的线程块减少，这是一种退步。

但在更高的时钟频率下（右侧，2.5 GHz），降低占用率反而有帮助，因为我们发出的原子操作更少，但更大的线程块处理着更大的私有化直方图，这显著提升了性能。

简而言之，如果我们查看较低的GPU时钟值，我们可能会观察到不能代表真实设置的性能变化。用户运行你的内核时不会锁定低频，他们感兴趣的是GPU的最大性能。因此，如果你在较低频率下观察到退步，你可能会避免应用这个更改，而这个更改实际上在更能代表用户工作负载的较高频率下带来了显著的性能提升。

相反的情况也可能发生。我们多次看到，在较低频率下观察到性能改进，但在较高频率下性能却退步了。因此，这强烈取决于你的算法运行时间对GPU频率的依赖曲线的形状。

对于一些算法来说，情况没那么糟。例如，adjacent_difference 是Thrust库中最简单的内核之一，它几乎不依赖于GPU频率，因为这个算法是内存瓶颈型的，计算量很少。

关键要点

在这一点上，我们可以得出哪些结论？

• 锁定在较低频率是不安全的，锁定在较高频率也无济于事。
• 算法本身会导致差异，因此我们必须进行多次基准测试。
• 但我们又面临着因启动大量基准测试而导致的过度节流问题。

那么，我们可以做的是尝试测量基准测试期间GPU的平均频率，然后直接丢弃在明显节流期间发生的测量值。例如，我们可以声明一个最大频率，获取平均频率，并设置一个阈值（比如峰值的75%），认为这对于我的应用程序来说是不现实的，然后丢弃这些测量值。这将帮助我们获得更具代表性的性能指标样本。

还有一个相关问题：我们应该锁定CPU时钟吗？由于我们之前讨论的流阻塞方法，我们并没有真正将CPU端启动内核的时间计入测量。因此，CPU目前无法影响我们的测量，锁定CPU时钟没有帮助。

高效的测量：熵停止准则

但我们仍然存在10万次测量耗时太长的问题。这只是一个工作负载，我们实际上可能有许多数据类型和问题规模需要测试，而这正是基准测试的目的。因此，我们希望能够尽可能减少测量次数。我们该如何做到这一点？

我们可以尝试重新表述这个问题。打个比方：我有一个骰子，但不知道它有多少面。我该如何弄清楚需要投掷这个骰子多少次才能理解它的面数？要回答这个问题，我们可以参考香农熵。

香农熵表示描述一个随机变量状态所需的信息量的期望值。对于性能测量，事件 X 是我们从内核运行中观察到的具体运行时间。

让我们尝试投掷骰子一次。假设我们掷出了6。我们唯一观察到的事件是6，并且观察到一次。所以观察到6的概率是1，熵是0。这是合理的，因为根据我们看到的数据，描述下一个随机变量不需要任何信息，它还会是6。

如果你第二次投掷并观察到2，现在我们有两个等可能的事件（6和2），熵增加到1。现在我们需要更多信息来描述下一次投掷，可能是2或6，我们不知道。随着你不断投掷，每次出现新的、令人惊讶的测量值时，熵都会增加。

如果我们掷了19次，终于看到了最后一面（假设是6面骰子）。理想情况下，如果我们知道骰子的实际面数，我们就会在这里停止。你会看到熵随着投掷次数增加而逐渐饱和，因为看到每个新测量值变得越来越不令人惊讶。但我们实际上并不知道只有六面，也许有更多面，所以让我们尝试掷47次。你会看到熵现在已经完全稳定，不再进一步增加，它已经饱和了。

我们可以在基准测试中利用这个特性，提出一个停止准则。让我们从投骰子回到性能测量。对于 select，尝试追踪样本熵随时间的变化。你可以看到同样的效果：我们得到越来越不令人惊讶的测量值。

我们如何停止基准测试？我们可以尝试在熵的窗口上运行线性回归，以查看线性回归在何处表明熵已达到饱和。通过这种方法，我们可以更早停止。例如，我们在大约800次测量后停止，这比我们原本计划的10万次要少100多倍。

这就引出了一个问题：我们是否真的看到了之前观察到的所有性能分布？数据显示，虽然不完美，但我们至少覆盖了与过度基准测试相同的运行时间范围和频率范围。因此，我们在尽可能减少测量的同时，获得了更符合现实的测量值分布。

NVBench：自动化基准测试框架

但是，实现所有这些（流阻塞、停止准则、预热、缓存刷新等）变得非常复杂。将这些抽象到一个框架中是合理的。这正是我们所做的，我们创建了 NVBench。

它类似于 Google Benchmark（如果你熟悉这个工具），但不同之处在于它支持 CUDA。它理解CUDA计时器、设备、缓存刷新。它还通过流阻塞和基于熵的停止准则提高了基准测试的稳定性。此外，它还具有强大的参数扫描框架实用程序，允许你探索工作负载空间，以及灵活的命令行界面。

让我们看看如果使用 NVBench，基准测试会是什么样子。基准测试只是一个接受 nvbench::state& 引用的函数，然后你用 NVBENCH_BENCH 宏注册它。你还可以为内核的参数注册不同的轴。例如，我们可能关心 select 在小、中、大问题规模下的性能。你可以注册一个名为 elements 的轴和三个值，基准测试将被调用三次。在基准测试函数内部，你可以通过 state.get_int64("elements") 检索当前元素

4：CuTe张量代数库入门教程

概述

在本节课中，我们将学习CuTe（CUDA Tensor）库的核心概念。CuTe是一套用于在GPU上高效构建张量线性代数运算的抽象工具集，它通过独特的布局（Layout）系统，将复杂的多维数据访问和线程分区模式转化为可组合的数学对象，从而兼顾了编程的灵活性与运行时性能。

1. CuTe的起源与目标

上一节我们介绍了课程概述，本节中我们来看看CuTe库是如何诞生的，以及它旨在解决什么问题。

CuTe的创造源于一个核心痛点：张量收缩（Tensor Contractions）操作在现有库中实现起来异常困难。尽管张量收缩本质上可以视为广义的矩阵乘法（GEMM），但传统方法（如BLAS扩展）在处理多维数据的复杂布局和分区时显得笨拙且低效。

CuTe的开发历程证明了其设计的优越性。在早期实验中，一位实习生仅用一周时间，通过修改一个基于CuTe的单文件GEMM示例，就超越了高度优化的cuBLAS库25%的性能。这得益于CuTe提供的“可破解”的、高级别的抽象，使得开发者能够深入底层进行优化，同时保持代码的简洁性。

CuTe的核心目标是：为构建GPU上各种规模和范围的张量线性代数运算，提供一套兼顾生产力（Productivity）与性能（Performance）的抽象。它是一个多层次系统，开发者可以从任意层级切入：

• 速度极限（Speed of Light, SOL）内核：提供接近硬件理论极限性能的预构建内核（如GEMM）。
• 集合（Collective）层：允许开发者对计算和通信模式进行定制化“ hacking”。
• 基础CuTe层：整个cuBLAS库都构建在CuTe之上，提供了最底层的布局和代数操作原语。

2. 从循环到布局：思想的转变

上一节我们了解了CuTe的目标，本节中我们来看看它是如何通过重新思考循环来建立其核心抽象的。

传统的多维循环代码包含大量起始索引、步长和偏移量参数，使得数据在内存中的实际布局难以一眼看清。CuTe的思想是将这些参数抽象为一个布局（Layout） 对象。

布局是一个数学函数，它将逻辑坐标（例如矩阵中的(i, j)）映射到物理索引（内存中的一维偏移量）。最常见的布局是仿射变换，通过步长（Stride） 来描述。

以下是如何用步长描述常见布局：

• 列优先（Column-major）：对于一个2x3的矩阵，形状为 (2, 3)，步长为 (1, 2)。这意味着沿着行（第一维）移动时，内存地址连续（步长为1）；沿着列（第二维）移动时，需要跳过一行的元素（步长为2）。
  • 物理内存顺序：A, B, C, D, E, F
• 行优先（Row-major）：形状同样为 (2, 3)，步长为 (3, 1)。沿着列（第二维）移动时内存地址连续。
  • 物理内存顺序：A, D, B, E, C, F
• 带填充的布局：形状 (2, 3)，步长 (1, 4)。这会在每“行”末尾留下一个未使用的填充元素。
• 高维张量：此概念可轻松扩展到任意维度。对于一个三维张量，其布局由形状 (dim0, dim1, dim2) 和对应的步长 (stride0, stride1, stride2) 定义。

给定逻辑坐标 coord 和步长元组 strides，计算物理索引 physical_idx 的公式为：
physical_idx = sum(coord[i] * strides[i]) for i in range(rank)

CuTe的创新之一：静态形状与步长
在GPU编程中，尤其是在处理寄存器、共享内存或静态分块时，张量的形状和步长常常在编译时已知。CuTe允许将这些值定义为静态整数。编译器可以利用这些信息进行激进优化（如完全展开循环、预计算偏移量），从而消除所有运行时开销。

// 动态形状
cute::Tensor tensor_dyn = make_tensor(data_ptr, cute::make_shape(22, 19));
// 静态形状（22在编译时已知）
cute::Tensor tensor_static = make_tensor(data_ptr, cute::make_shape(cute::Int<22>{}, 19));

静态信息会在打印的布局中以下划线 _ 标示，并且在性能关键的内核中至关重要。

3. 布局的威力：折叠与代数运算

上一节我们介绍了基础的布局概念，本节中我们来看看CuTe布局如何通过“折叠”和代数运算解决传统方法无法处理的难题。

一个常见的需求是将高维张量“折叠”成矩阵，以便调用高效的GEMM库。然而，传统的“展平”操作存在不对称性。

问题：考虑一个形状为 (2, 2, 2)，步长为 (4, 1, 2) 的三维张量。我们可以轻松地将其前两维合并，视为一个 2 x 4 的矩阵。但是，如果我们想将其视为一个 4 x 2 的矩阵（即合并后两维），却找不到一个单一的整数步长来描述新“行”中元素的物理位置（因为元素A、C、E、G在内存中并非等间距排列）。

CuTe的解决方案：放弃强制“展平”为单一步长矩阵的思维。CuTe允许我们创建层次化形状（Hierarchical Shape）。

与其将 (2,2,2) 张量强行视为 (4,2)，不如将其视为 ((2,2), 2)。这意味着新的行坐标本身是一个二维坐标 (i,j)。通过这种方式，我们保留了原始的步长信息 (4,1,2)。当使用行索引 (i,j) 和列索引 k 进行访问时，计算物理索引的公式依然成立：
physical_idx = (i, j, k) · (4, 1, 2) = i*4 + j*1 + k*2

在代码中，我们可以通过一个简单的索引转换，将一维的行编号 r 映射到这个二维坐标 (r/2, r%2)，从而在算法层面仍然将其当作普通的 4x2 矩阵来处理。这揭示了折叠的对称性：只要通过层次化形状保留内部结构，任何张量都可以折叠成任何形式的矩阵。

布局代数
CuTe将布局视为可进行数学运算的对象，形成了一个丰富的代数系统：

• 组合（Composition）：layout_c = layout_a ◦ layout_b。这是最强大的操作，可以构建复杂的分区。
• 乘积（Product）：用一个布局替换另一个布局中的每个元素，形成“布局的布局”。
• 除法（Divide）：根据一个布局将另一个布局分割。
• 逆（Inverse）：获取布局的逆映射。

这些操作为实现通用的分块（Tiling）和分区（Partitioning）提供了坚实的基础。

4. 通用算法：一次编写，处处运行

上一节我们探讨了布局的代数性质，本节中我们来看看如何利用这些性质编写极其通用的算法。

得益于布局抽象，我们可以编写出与张量维度、形状、步长无关的算法。

以Copy为例
一个朴素的复制实现需要根据源张量和目标张量的维度编写多重嵌套循环。但在CuTe中，由于所有张量都可以通过整数坐标进行索引，复制算法可以简化为一个一维循环：

template <class SrcTensor, class DstTensor>
void copy(SrcTensor const& src, DstTensor& dst) {
    for (int i = 0; i < size(dst); ++i) {
        dst(i) = src(i);
    }
}

这个简单的 copy 函数可以处理：

• 任意维度的张量之间的复制。
• 转置操作（例如从列优先布局复制到行优先布局）。
• 广播操作（从步长为0的张量复制）。
• 收集（Gather）和散射（Scatter）操作。
• 静态形状的完全循环展开与零运行时开销优化。

以GEMM为例
同样，矩阵乘法的核心也可以用一个与具体布局无关的三重循环来实现：
C(m, n) += A(m, k) * B(k, n)
只要 A, B, C 张量支持通过 (m,k), (k,n), (m,n) 坐标进行访问，这个算法就成立。这使得同一个GEMM实现可以用于普通矩阵乘法、批处理GEMM、张量收缩甚至卷积（通过Im2Col变换）。

5. 核心应用：线程分区与Tensor Core编程

上一节我们看到了通用算法的简洁性，本节中我们来看看CuTe最具威力的应用：描述复杂的线程数据分区，特别是用于Tensor Core编程。

在GPU内核中，一个关键步骤是将全局数据分块（Tile）并分配给线程块（Block）或线程（Thread）进行处理。CuTe通过布局组合来优雅地定义分区。

分区即组合后切片

1. 定义线程值布局（Thread Value Layout）：这是一个描述“哪个线程负责哪些逻辑坐标”的布局。例如，对于一个有4个线程的向量，可以定义一个 4x6 的布局，其中每一行对应一个线程，行内的6个坐标是该线程需要处理的向量元素位置。
2. 与输入张量组合：将输入张量的布局与线程值布局进行功能组合（Functional Composition）。这产生了一个新的布局，其第一维是线程ID，第二维是该线程负责的元素在其本地视图中的索引。
3. 切片：每个线程通过自己的线程ID对这个组合后的布局进行切片，就得到了自己负责的那部分数据的视图（一个子张量）。

公式化表示：
thread_sub_tensor = slice(compose(global_tensor_layout, thread_value_layout), thread_id, _)

这个方法的强大之处在于：

• 适用于任意张量：因为任何张量都可以通过整数坐标索引，所以此分区方法适用于任何维度、任何布局的张量。
• 完美匹配Tensor Core：NVIDIA Tensor Core指令集有严格的输入数据分区要求。CuTe可以精确地将这些要求编码为线程值布局。开发者只需构建正确的共享内存布局和寄存器布局，然后使用相同的分区模式，即可安全、正确地调用复杂的Tensor Core指令（如Volta、Ampere、Hopper架构的MMA指令）。
• 提升开发效率：无需手动跟踪每个线程负责哪些复杂非连续的数据。CuTe的抽象使得编写和调试Tensor Core内核变得更加直观和安全。优化工作通常简化为设计更高效的共享内存布局，而核心计算逻辑保持不变。

6. 高级优化：自动化向量化与未来方向

上一节我们深入了解了CuTe在核心计算中的应用，本节中我们简要探讨其如何支持高级编译时优化，并展望其未来。

CuTe的数学基础允许在编译时进行深度分析，从而自动启用优化。

自动化向量化
考虑两个张量之间的复制操作。如果源张量和目标张量在某个逻辑子区域上，其元素在物理内存中都是连续的，那么就可以用更宽的向量化加载/存储指令来加速复制。CuTe可以通过数学运算（如计算两个布局的公共子布局）自动发现这种机会。

公共子布局的数学本质可以归结为计算 inverse(layout_b) ◦ layout_a。其结果揭示了两个布局在坐标映射上的最大公共连续模式，编译器可据此安全地确定向量化宽度。

CuTe的Python前端
为了应对C++模板元编程带来的编译时间慢、错误信息冗长等挑战，CuTe团队正在积极开发Python前端。这将带来：

• 极快的编译/“编译”时间（提升超100倍）。
• 与PyTorch等ML框架的无缝集成。
• 更易上手的编程体验，同时保留CuTe的所有表达能力。

总结

本节课中我们一起学习了CuTe张量代数库的核心思想与应用。我们从其解决张量收缩难题的起源出发，逐步深入其核心抽象——布局（Layout）。布局作为一个将逻辑坐标映射到物理索引的数学函数，通过支持静态信息、层次化形状和丰富的代数运算（组合、乘积等），实现了前所未有的灵活性。

我们看到，基于布局的抽象使得编写维度无关的通用算法（如copy、gemm）成为可能。更重要的是，CuTe通过“分区即组合后切片”的范式，为GPU内核中复杂的数据分区，特别是Tensor Core编程，提供了安全、简洁且强大的描述工具。最后，我们了解到CuTe的数学基础还能支持编译时的自动化优化，并且其正在向更友好的Python生态系统演进。

CuTe的成功体现在其作为cuBLAS等高性能库的基石，以及被广泛应用于从基础线性代数到注意力机制等各种GPU计算场景中。它证明了良好的抽象不仅能提升生产力，更能释放硬件的极限性能。

5：解耦式大语言模型推理 🚀

概述

在本节课中，我们将要学习解耦式大语言模型推理。这是一种将推理过程的前缀阶段和解码阶段分离的系统设计方法，旨在优化服务质量和资源利用率。我们将探讨其核心动机、设计挑战、实现方案以及社区的最新进展。

---

引言与背景

上一节我们介绍了课程主题，本节中我们来看看解耦式推理的背景和动机。

欢迎来到GPU MODE的第58讲。本次讲座将讨论解耦式大语言模型推理。这项研究是近年来最具影响力的机器学习系统研究之一。

如今，我们有多种不同的应用场景，例如聊天机器人、搜索引擎或编程助手，它们对服务等级目标有着不同的要求。我们主要关注两个指标：

• 首令牌时间：从收到请求到生成第一个输出令牌所需的时间。
• 每令牌时间：在生成第一个令牌后，后续每个输出令牌之间的间隔时间。

不同的应用对这些延迟有着不同的敏感度。

---

从吞吐量到优质吞吐量

上一节我们提到了延迟指标，本节中我们来看看如何更全面地评估系统性能。

长期以来，人们通常使用吞吐量作为衡量推理系统成本的代理指标。这里的吞吐量指系统每秒处理的请求数或生成的令牌数。

然而，仅考虑吞吐量是不够的。如果系统负载过高，导致用户请求的延迟无法满足其应用的服务等级目标，那么用户体验就会下降。因此，我们引入了一个更优的指标——优质吞吐量。

优质吞吐量的定义是：在满足特定服务等级目标约束的前提下，系统每秒能够成功处理的请求数量。它可以更好地衡量“每请求成本”，因为未能满足延迟要求的请求对用户而言价值很低。

高吞吐量并不等同于高优质吞吐量。例如，一个系统每秒能处理10个请求，但如果设定严格的服务等级目标，可能只有3个请求能满足要求，那么其优质吞吐量就是3。

因此，我们的系统设计目标是：在满足各种应用延迟要求的前提下，使用最少的GPU资源服务更多的请求，即最大化优质吞吐量。

---

前缀与解码：不同的计算特征

为了理解解耦的必要性，我们需要先了解大语言模型推理的两个核心阶段：前缀和解码。

以下是这两个阶段的简要说明：

1. 前缀阶段：用户输入（提示词）的所有令牌被一次性送入模型进行并行前向计算。此阶段会生成注意力机制所需的键值缓存，并输出第一个令牌。处理整个提示词的时间即首令牌时间。
2. 解码阶段：模型以上一个生成的令牌为输入，自回归地逐个生成后续令牌。每个解码步骤的时间即每令牌时间。

这两个阶段具有截然不同的计算特征：

• 前缀阶段是计算密集型：即使单个请求也容易使GPU计算饱和。
• 解码阶段是内存密集型：主要开销在于从内存加载模型权重和键值缓存，通常需要较大的批处理大小才能充分利用计算资源。

---

传统方案的挑战：干扰与资源配置

上一节我们了解了两个阶段的特性，本节中我们来看看将它们混合处理时面临的问题。

传统服务引擎通常将前缀和解码请求放在同一批中进行处理，这带来了两个主要挑战：

1. 计算干扰
当解码请求与新的前缀请求被批处理在一起时，计算密集型的前缀阶段会显著拖慢内存密集型的解码步骤，导致解码延迟大幅增加。这种干扰使得系统难以保证解码请求的低延迟要求。

2. 资源配置僵化
不同的应用对首令牌时间和每令牌时间有着不同的优先级。然而，在混合调度的系统中，前缀和解码阶段共享相同的并行化配置（如张量并行、流水线并行的程度）。这导致我们无法为前缀阶段（追求低首令牌时间）和解码阶段（追求低每令牌时间）分别独立地优化资源配置，从而无法在两者间取得最佳平衡。

---

解耦式推理的核心思想

针对上述挑战，解耦式推理提出了一个直观的解决方案：将前缀阶段和解码阶段物理分离到不同的计算设备上。

核心思想如下：

1. 设立专门的前缀工作器组和解码工作器组。
2. 请求首先进入前缀工作器，完成提示词处理并生成键值缓存。
3. 将生成的键值缓存和请求状态迁移到解码工作器。
4. 请求在解码工作器上完成剩余令牌的生成。

这种分离带来了以下好处：

• 消除干扰：前缀和解码阶段不再相互影响。
• 独立优化：可以为前缀工作器和解码工作器独立配置最适合的并行策略和资源数量，以分别优化首令牌时间和每令牌时间。

---

关键设计挑战与解决方案

实现解耦式推理需要解决几个关键的技术挑战。

1. 如何确定最优的并行配置？
我们通过一个模拟器来解决。该模拟器会针对给定的工作负载特征（请求到达率、提示长度等）和服务等级目标，在前缀和解码的配置空间（如张量并行度、流水线并行度）中进行搜索，找到能使系统优质吞吐量最大化的资源配置方案。

2. 如何高效迁移键值缓存？
键值缓存的迁移可能成为性能瓶颈。我们采用了两种策略来优化：

• 基于网络拓扑的智能放置：在高带宽互联（如NVLink）的GPU组内进行键值缓存迁移，避免跨低速网络的传输。
• 拉取式传输：改为由解码工作器在需要时主动从前缀工作器的内存中拉取键值缓存，这有助于缓解解码侧内存突增的压力，并实现更灵活的传输调度。

3. 如何应对请求突发？
拉取式传输机制本身有助于应对突发请求。此外，结合流水线并行等技术也可以平滑请求处理流程。

4. 如何减少流水线气泡？
可以采用分块等技术来优化流水线并行的效率，减少计算资源的空闲时间。

---

评估方法与社区进展

我们使用优质吞吐量和服务等级目标达成率作为核心评估指标，而不仅仅是原始吞吐量。实验表明，解耦式设计能够显著提升系统的优质吞吐量。

目前，解耦式推理已成为业界主流方案。英伟达的NIM、微软的DeepSpeed、vLLM等各大公司的推理引擎均已集成或正在集成此特性。社区也涌现出如OneCache（专注于键值缓存中心化存储）等创新设计。

---

未来展望

解耦式推理为系统设计打开了新的空间，未来的研究方向包括：

• 快速重配置：实现系统能够根据工作负载的变化，动态、快速地调整前缀与解码资源的比例和并行配置，向“Serverless”化演进。
• 更广泛的集成：将解耦思想与推理优化、长上下文处理、多模态模型等更广泛的场景结合。

---

总结

本节课中我们一起学习了：

1. 解耦式推理的动机：为了优化以优质吞吐量为代表的真实服务质量，并解决前缀与解码阶段混合调度时的干扰和资源配置问题。
2. 核心思想：将前缀和解码阶段物理分离到不同的计算单元，允许独立优化。
3. 关键挑战与方案：通过模拟器寻找最优配置，利用智能放置和拉取式传输优化键值缓存迁移。
4. 现状与未来：该技术已被业界广泛采纳，并为进一步实现弹性、高效的Serverless化推理系统奠定了基础。

解耦式推理通过重新思考推理系统的计算组织方式，为实现更低成本、更高质量的大语言模型服务提供了关键路径。

6：加速大型视频扩散模型

在本节课中，我们将学习如何加速大型视频扩散模型的推理过程。我们将探讨视频扩散模型的工作原理、当前面临的性能瓶颈，以及通过蒸馏、特征缓存和稀疏注意力等技术来显著提升推理速度的方法。

概述：视频扩散模型的挑战

视频生成模型正经历快速发展，但其推理过程非常耗时。以生成一段5秒、720P的视频为例，在高端H100 GPU上可能需要长达16分钟。这主要源于两个核心瓶颈：扩散过程的迭代性（通常需要20-50步）以及注意力机制在长序列上带来的巨大计算开销（可占总计算量的85%以上）。

视频扩散模型的工作原理

上一节我们介绍了视频扩散模型面临的挑战，本节中我们来看看其基本工作原理。我们以当前先进的开源视频模型SVD 1.2为例。

视频扩散模型与图像扩散模型类似，但处理的序列更长。其核心是一个扩散变换器（Diffusion Transformer, DiT）。去噪过程以输入噪声、用户提供的文本提示和时间步长为条件。为了融入文本信息，会在扩散变换器和视觉/文本编码器之间应用注意力机制。

视频是三维数据（高度、宽度、时间/帧数）。在输入变换器前，需要将3D视频展平为1D序列。最后，一个解码器将潜在表示解码为最终的高分辨率视频。扩散变换器通常在潜在空间中工作，这可以稍微压缩序列长度。

视频扩散模型的性能瓶颈

了解了基本工作原理后，我们进一步分析其推理时间的构成。以另一个先进模型Hunyuan Video为例，生成5秒720P视频时，扩散变换器需要处理超过10万个令牌，这是一个非常长的序列。

因此，扩散变换器中使用的长序列因果注意力机制占据了超过85%的计算量。扩散本身又是一个迭代过程（例如Hunyuan默认50步）。在H100上，即使使用了Flash Attention 3和Torch Compile等优化，单次前向传播仍需约19秒。总计生成一个视频需要约16分钟，这非常耗时。语言编码器和VAE解码器耗时极少。

综上所述，现有开源视频生成模型主要受限于两个因素：

1. 扩散的迭代过程。
2. 注意力机制的计算开销。

我们需要开发方法来应对这两个主要成本以加速视频扩散模型。

加速方法一：减少扩散步数（蒸馏）

为了应对迭代过程的瓶颈，一个主要方向是减少所需的扩散步数。以下是相关方法：

• 蒸馏方法：如渐进式蒸馏、一致性模型、DMD、LCM等。这里的“蒸馏”并非指从大教师模型到小学生模型，而是在相同大小的模型之间，从使用多步扩散的教师模型蒸馏到使用极少步数的学生模型。这类方法可将扩散步数减少5-10倍，从而显著加速推理。
• 特点：蒸馏需要训练，且训练过程通常非常不稳定。在图像生成领域应用较多，但在视频生成领域成功的开源代码较少，复现难度大。

加速方法二：特征缓存

另一种无需训练的方法是特征缓存。它利用了连续扩散步之间中间特征图相似性高的启发式方法。

其思想是：在当前的扩散步中，策略性地重用前一个扩散步计算出的部分特征图，从而避免重新计算扩散变换器的某些组件（例如前馈网络层或注意力层）。这类方法通常能将推理速度提升约2倍。

关于内存带宽的讨论：有人可能会问，跳过计算但需要加载之前步的权重，是否会变成内存带宽瓶颈？实际上，由于序列非常长（超过10万令牌），视频扩散模型通常是计算受限的。不过，缓存所有中间特征确实会带来不小的内存开销。

FastVideo项目：统一的加速框架

目前，各种加速方法（蒸馏、缓存）的代码分散在不同的仓库中，难以比较和集成。FastVideo项目旨在构建一个统一、高效的视频生成模型训练和推理引擎。

FastVideo目前整合了一种蒸馏方法（PCM）和一种特征缓存方法（T-Cache），并提供了抽象接口以便集成其他方法。例如，项目已成功对Hunyuan模型进行了蒸馏，在质量略有下降的情况下实现了8倍的加速。

加速方法三：稀疏注意力

除了减少步数，另一个关键方向是直接减少注意力机制本身的计算量。这在视频扩散模型中尤为重要，因为其序列长度远超图像模型。

注意力分数通常是稀疏的。观察注意力图可以发现，只有少数位置的注意力分数较高，其余大部分接近零。不同层、不同注意力头的稀疏模式差异很大。如果我们能预测这些关键模式，只计算关键部分，就能加速注意力。

然而，为了实现硬件高效执行，我们需要块稀疏注意力模式。GPU计算注意力时以块（Tile）为单位。如果一个块完全稠密或完全空，都可以高效处理。但如果一个块是混合的（部分需要计算，部分不需要），GPU仍需计算整个块再丢弃不需要的部分，这会造成算力浪费。因此，理想的预测模式应只包含稠密块和空块。

这引出了几个设计权衡：

1. 预测开销与实际注意力开销：预测关键令牌的过程本身必须非常轻量。
2. 块大小与效率的权衡：较小的块大小能使注意力模式更灵活、更具表现力，有利于模型性能，但会降低注意力核函数的算术强度，可能拖慢计算速度。较大的块大小则相反。
3. 是否需要全局上下文：既然只计算部分连接，是否还需要设计专门的轻量级分支来提供全局或局部上下文？或许可以借鉴CNN的思想，让令牌主要关注其邻居。

我们的研究试图回答这些问题，并提出了视频稀疏注意力方法。

VSA：视频稀疏注意力详解

我们的工作《Fast Video Generation with Sliding-Window Attention》和《Faster Video Diffusion with Trainable Sparse Attention》提出了具体的稀疏注意力方案VSA。

一个关键技术是令牌重排序。传统上，视频被展平为1D序列时，空间上邻近的令牌在序列索引上可能相距甚远。我们提出将视频分割成小的3D立方体（例如4x4x4），并确保同一立方体内的令牌在展平后的序列中拥有连续的索引。这样，每个立方体可以映射到注意力核函数中的一个计算块。在执行稀疏注意力时，整个立方体要么被全部计算，要么被全部跳过，从而自然形成块稀疏模式。

VSA的结构包含两个分支：

1. 粗粒度分支：对Q、K、V进行4x4x4的平均池化，得到立方体级别的表示，将序列长度减少64倍。然后计算轻量级的立方体级别注意力，获得全局上下文输出。
2. 细粒度分支：基于粗粒度分支的注意力图，使用Top-K选择关键立方体块。然后仅对这些关键块执行令牌级别的细粒度注意力，块大小为64。
3. 门控融合：两个分支的输出分别经过一个可学习的门控权重，然后相加，以平衡两者对最终注意力输出的贡献。

通过大量实验，我们得出了最佳配置，并验证了其有效性：

• 在控制计算量的情况下，VSA的性能优于全注意力。
• 即使引入87.5%的稀疏度（仅计算12.5%的注意力），VSA在百万至十亿参数规模的模型上，其损失-计算量曲线始终优于或匹配全注意力模型，显示出更好的扩展性。
• 稀疏度的最优选择取决于训练预算和序列长度。模型训练得越充分，所需的稀疏度越低（越接近全注意力）。

微调现有模型使用VSA

对于已训练好的全注意力模型，我们可以通过渐进式的方法将其微调为使用VSA：

1. 初始化：将粗粒度分支的门控权重初始化为零，使其初始贡献为零。将细粒度分支的Top-K值初始化为总块数（即模拟全注意力）。
2. 渐进稀疏化：在微调过程中，逐步将Top-K值降低到目标稀疏度（例如，从300块降到32块）。
3. 数据：我们使用Synthetic数据（由Sora 1.4B生成）进行微调。实验显示，即使从高稀疏度开始，随着微调进行，生成质量会迅速提升，最终达到与全注意力模型相当的水平。

总结与未来方向

本节课我们一起学习了加速大型视频扩散模型的多种技术：

1. 蒸馏：可带来5-10倍加速，是当前最显著的方法，但需要训练且开源生态分散。
2. 特征缓存：无需训练，能提供约2倍加速。
3. 稀疏注意力：我们的VSA方法既能加速推理（约2倍），也能加速训练，且性能优于全注意力。
4. 其他技术：如量化（FP8线性层、平滑量化）以及针对扩散变换器中特定算子的定制核函数（如AdaLN）也值得探索。

FastVideo项目的目标是构建一个统一平台，支持不同模型、不同蒸馏技术和不同稀疏注意力方法，集中优化视频生成模型的训练和推理效率，最终目标是将H100上16分钟的推理时间缩短至实时水平。

未来方向包括：探索混合专家模型、研究更精细的稀疏策略（如不同时间步或层使用不同稀疏度）、以及借鉴大语言模型推理基础设施的经验来优化多模型组成的扩散流程。

7：优化线性注意力

在本节课中，我们将学习线性注意力的工作原理、其效率挑战，以及如何通过分块递归算法和Flash线性注意力等技术实现硬件高效的训练。我们还将探讨Deltanet等改进模型如何解决线性注意力在上下文检索中的性能问题。

---

欢迎与介绍

大家好，欢迎来到GPU MODE的另一期节目。今天，我们很高兴邀请到Songland Yang作为嘉宾。

Songland被多位同行推荐为演讲者，主要原因是她在构建高效架构和加速这些架构方面，是研究社区中的关键人物之一。她围绕编写自定义内核和探索非典型架构的工作非常有趣。

传统注意力机制回顾

上一节我们介绍了本课程的主题，本节中我们来看看Transformer，它是现代生成式AI的主力模型。

我们都熟悉其核心的自注意力机制。以下是快速回顾。

对于并行训练，我们有三个矩阵 Q、K、V，大小为 L × D。其中 L 是序列长度，D 是头维度。这里我们只讨论单个头的情况。M 是因果掩码。

在并行训练中，我们可以使用几个矩阵乘法：

1. 首先，计算查询矩阵和键矩阵的点积，得到注意力分数。
2. 然后应用因果掩码，接着进行Softmax操作。
3. 最后，将得到的注意力矩阵与值矩阵进行点积，得到输出。

这是并行训练的形式。在自回归生成场景中，我们使用迭代形式。基本上，在时间步 T，我们使用该时间步的查询向量 q_T 去关注之前所有的键向量，得到注意力分数，然后用这个分数对历史值向量进行加权求和，得到输出。

这是一个非常快速的自注意力机制回顾，相信大家都很熟悉。

传统注意力的效率问题

我们了解到，Softmax注意力存在几个效率问题。

第一个问题在于训练效率，因为注意力机制的计算复杂度是序列长度的二次方。当序列长度非常大时（例如在视频生成或DNA建模中达到百万级别），这将成为关键瓶颈。

第二个效率问题在于推理。在推理时，我们需要维护历史的键和值向量，即所谓的KV缓存。KV缓存的大小随生成长度线性增长，在进行长序列生成时，缓存大小可能爆炸式增长，导致内存不足。

现在我们理解了自注意力机制的关键效率挑战。虽然有许多优化技术可以降低这个成本，但今天我们将聚焦于一种可以从根本上消除KV缓存的不同架构——线性注意力。

线性注意力的核心思想

线性注意力背后的关键思想非常简单。它基本上移除了Softmax注意力中的Softmax算子。

因此，我们得到了这个迭代形式。让我们深入了解一下。

这里我们有一个注意力分数，但没有Softmax算子，现在它是一个完全线性的算子。

首先，这是一个标量值，我们将其放在第二个位置。然后，我们利用这些算子的结合律，先将值和键组合起来。

我们得到这个标红的部分，它是所有历史时间步的键值外积的累积和。右边这个项通常被称为递归隐藏状态，因为在推理时，这个隐藏状态通过外积计算结构将键值关联对编码到记忆里。

现在变得非常清楚，线性注意力实现了一种线性递归。其递归状态就是键值外积的简单累积和。

对于输出计算，我们使用矩阵向量乘法来获得输出。

这是一个非常简单的线性递归。我们知道，在经典的RNN或门控循环单元中，递归状态的大小通常在输入维度 D 的量级。而现在，递归状态的大小在 D² 的量级。从这个意义上说，线性注意力拥有一种机制，允许我们以非常高效的方式扩展递归状态的大小。

我们都知道，对于RNN来说，递归状态大小非常重要，因为RNN只有一个固定大小的状态记忆来编码所有历史数据。直观上，状态越大，RNN的记忆能力就越好。

线性注意力的高效训练挑战

接下来，我将讨论如何以硬件高效的方式训练这种线性注意力模型。

之前，我们展示了线性注意力的并行形式和递归形式。然而，不幸的是，这两种形式都不适合进行硬件高效的训练。

首先，对于并行形式，它与Softmax注意力的情况非常相似，计算复杂度也是序列长度的二次方。如果我们进行非常长的序列建模，这将非常成问题。

对于递归形式，也存在问题，原因有几个。第一个原因是，在递归形式中，我们必须进行顺序递归，这限制了在序列维度上的并行性。其次，我们可以看到，对于递归状态更新，我们只涉及一个秩为1的外积更新；对于记忆读取，我们只进行矩阵向量乘法。这两个操作都不是矩阵乘法，因此无法使用张量核心进行加速，而张量核心是非常高效的。

一个常见的问题是，为什么我们不使用并行扫描算法来训练这个线性注意力模型？因为递归的第一个限制是它无法在序列上并行，而我们知道线性递归可以使用并行扫描算法在序列上实现并行计算。

但是，在线性注意力的情况下，即使并行扫描算法可以提供序列级别的并行性，你仍然无法获得能够进行快速矩阵乘法的操作。这主要是由于递归形式的限制，它本身就不具备基于矩阵乘法的注意力机制。

当扩展递归状态大小时，这可能是个问题。因此，Mamba无法将其状态大小扩展到数万，因为它无法利用张量核心的算力。此外，并行扫描往往会带来非常大的内存和I/O开销。

对于递归形式，我们只需要维护一个大小为 D × D 的递归状态。但对于并行扫描，特别是当状态大小太大无法放入共享内存时，我们必须将其存入高带宽内存。这就需要为每个时间步维护大量的递归状态，这是非常消耗内存的。正如我们之前提到的，我们有一个线性状态扩展机制，每个位置的递归状态大小是 D × D。如果我们为每个时间步都物化递归状态，内存复杂度将是 O(L × D²)，这比查询或键向量大 D 倍，这是不可接受的。

由于这些原因，在实践中，使用并行扫描训练这些注意力机制和模型太慢了。

分块递归算法

由于递归形式和并行形式的局限性，我们现在讨论分块递归形式，它基本上是介于递归形式和并行形式之间的一种中间形式。

分块递归算法的高层思想是：首先，我们有一个序列和一个块大小 C。我们将序列分割成几个大小相等的块。

当块大小等于1时，我们可以恢复递归更新机制。当块大小等于整个序列长度时，它将简化为并行形式。

这里我想强调，分块递归形式不是一种近似。它在数学上等价于递归形式和并行形式。因为线性注意力包含许多线性操作，允许我们利用结合律进行这些数学变换。

在分块递归形式中，我们不是为每个单独的令牌计算递归状态，而是只为每个块计算递归状态，并可能将其物化到高带宽内存中。

然后，我们可以使用一种混合形式来计算所需的输出。对于历史上下文，由于我们已经将所有历史信息编码到块级别的隐藏状态中，我们可以利用递归形式来计算来自先前上下文的输出贡献。

对于局部上下文，由于我们没有递归状态，我们可以转而使用并行形式直接计算由当前块内局部上下文贡献的输出。

这就是分块递归形式的高层思想。

为了用数学公式表示，我们有几个符号。首先，这是块级别的隐藏状态，我们使用带括号的下标来表示块的索引。

此外，我们使用这些块来表示查询、键、值和输出矩阵。它们的大小是 C × D。完整矩阵的大小是 L × D，但现在我们只取一个子块来考虑第 i 个块的计算。

首先，在第一步，我们进行块级别递归状态的计算。在这个阶段，我们只计算每个块的递归状态，这基本上是每个块的最后一个位置的状态。我们直接从这个状态跳转到那个状态。在中间，有几个状态，但我们跳过了计算，直接跳到最后一个状态。

对于这个递归状态更新，我们需要计算这个块内键值向量外积的累积和。这里有一个非常优雅的数学性质：这个外积的累积和可以写成矩阵乘法的形式。这就是为什么我们可以将递归状态更新写成矩阵乘法。

这里我们有大小为 C × D 的键块和值块。我们取值块的转置，然后进行点积，并在这个块大小维度上求和，我们将得到一个 D × D 的输出，其大小与递归隐藏状态相同。然后我们将这两项相加来实现状态更新。

对于输出计算，正如之前提到的，我们有两个贡献源。

第一个是历史上下文，它被总结到前一个隐藏状态中。因此，对于这些块，我们可以利用递归形式，让每个查询关注这个递归状态以获得输出。由于我们在这个块中有多个查询，并且它们都从同一个递归状态读取输出，这是一种批量化计算。我们可以在查询序列维度上进行批量化，使其成为一个矩阵乘法。

因此，第一项的大小是 C × D，第二项是 D × D，所以我们可以将这个重写为矩阵乘法过程。这是块间贡献，我们计算由先前上下文贡献的输出。

其次，我们还需要考虑由当前局部块贡献的输出。在这种情况下，我们直接利用并行形式，基于输入的查询、键和值直接计算输出。我们不需要计算任何中间递归状态，可以直接基于输入数据计算输出。这非常类似于自注意力机制的并行形式。基本上，这是一个局部块的线性注意力。我们对每个块分别应用这种并行线性注意力形式。

总结一下，使用这种分块递归形式，并固定块大小 C，总的时间复杂度是 O(L × D² + L × D × C)。当 C 不随序列长度 L 增长时，这是亚二次的。在实践中，我们将 C 设置为一个小的常数（例如64、128、256），这允许我们进行硬件高效的计算，这些是编写矩阵乘法内核时非常常见的块大小。

这种分块递归形式也非常通用，我们可以将其扩展到其他线性注意力变体（如带数据控件的Mamba）的高效训练。

这种形式已成为现代线性注意力模型（如基于Mamba的Deltanet等）的标准训练技术。

Flash线性注意力

对于分块形式，我们有一个与现代硬件非常契合的算法，但我们还需要进行一些I/O优化来进一步提高吞吐量。

在Flash线性注意力中，我们进行了几次内核融合，试图减少I/O成本。例如，在这种情况下，查询块可以用于计算两个部分的输出贡献。键块可以用于更新隐藏状态，也可以用于计算来自局部上下文的输出。

这是一个简单的内核融合，以减少I/O成本。在实践中，这种内核融合可以在GPU上带来显著的加速，与朴素PyTorch实现相比，可以达到约4倍的加速。

我们还有一个名为Flash线性注意力的开源库，在这个库中，我们提供了多种线性注意力机制（如RetNet、基于Gated的RWKV等）的实现。

总结一下，在线性注意力中，概念上的线性注意力非常简单，它只是没有Softmax的注意力，可以被表述为具有矩阵值隐藏状态的线性递归，允许你拥有更大的递归状态来编码历史数据。

分块递归形式是一种用于训练这些线性注意力模型的硬件高效算法。

而Flash线性注意力是该分块递归形式的一个I/O优化实现。

线性注意力的性能问题与Deltanet

现在我们已经了解了一些线性注意力的基本优化技术。然而，线性注意力本身存在一些性能问题，特别是在上下文检索方面。有多篇论文表明，线性注意力或RNN空间模型都受到一些常见问题的限制，例如由于状态大小有限导致的上下文检索能力不足。这也很直观，因为你不能期望RNN以无损的方式编码所有历史数据，因为隐藏状态是固定大小的。

因此，这些方法在下游任务非常重要的上下文检索中表现非常差。所以，我想提出一个模型来改进线性注意力的上下文检索能力，这就是Deltanet。

首先，我们可以分析为什么线性注意力在上下文检索中表现不佳。如前所述，线性注意力可以被视为一种关联记忆，使用外积结构编码键值关联对。这与经典的关联记忆非常相关。

如果我们想使用一个键从这个递归记忆中检索输出，我们使用这个记忆读取操作来检索输出。我们可以看到，因为 k_j 与时间步 j 的值相关联，理想情况下，我们希望从这个递归记忆中检索出 v_j。

但是，在这个关联记忆中，我们还有其他键值关联对，因此存在一些不需要的交叉项。我将这些交叉项称为检索误差，因为这是我们不希望从这个递归记忆中得到的项。这是因为，如果我们想实现完美检索，键向量应该彼此正交，这样我们就可以实现完美检索。但由于关联记忆是 D × D 的，在这个向量空间中最多只能有 D 个正交向量。因此，当序列长度大于模型维度时，这种真实误差是不可避免的。

这就是为什么增加头维度非常有用，因为如果增加头维度，向量空间中就有更多空间来容纳更多的正交键向量，从而获得更好的检索性能。

这是线性注意力的一个基本限制。因此，我们想设计一些机制来缓解这个问题。这里我们有了Deltanet。

Deltanet实际上是由...提出的经典模型。它使用了一种非常直观的记忆管理机制。

我们有查询、键和值向量，类似于线性注意力或注意力。我们对输入数据进行几次线性投影以获得查询、键和值向量。

在Deltanet中，我们还有另一个线性投影，将输入令牌映射到一个标量值。我们还应用Sigmoid函数将这个值限制在0和1之间。

这个 β_t 可以被想象为当前输入的写入强度。

Deltanet有一个记忆管理机制，最终输入由新输入和旧值的加权和计算。这个旧值是这样计算的：首先，我们使用输入键从记忆中检索关联值，然后模型动态决定是使用新的关联值（绿色的输入值）还是旧的关联值（蓝色的值）。模型必须决定是保留先前的关联值还是使用当前的关联值。

我们使用 β_t 来控制这个过程。β_t 是数据依赖的，因此模型可以动态决定是保留当前值还是使用先前的关联值。

然后，在Deltanet中，它从递归记忆中擦除旧值，并将新的键值关联信息写入递归记忆。

这种擦除机制是实现良好关联检索性能的关键。

对于记忆读取，它与普通的线性注意力相同，我们使用矩阵向量乘法来完成。

这里我们有一个非常简单的基准测试，称为多对关联检索。这是一个合成基准，用于测试这些新兴架构的上下文检索性能。

这个任务非常简单。我们在先前上下文中提供多个键值对。键是大写字母，值是数字。模型需要根据先前上下文中的键来检索值。这是一个纯粹的上下文学习任务，因为答案只出现在先前上下文中，而不在训练数据中。这些训练数据是即时生成的，因此模型无法在其直接记忆中对这些关联对进行编码。

尽管这个任务非常简单，但Transformer可以完美解决。然而，对于其他亚二次模型，如Mamba、Gated RNN，当模型维度较小时，它们无法实现完美性能。正如我们之前讨论的，因为当模型维度较小时，模型只能在空间中拥有少量正交键，这时模型会有较大的检索误差。

但对于Deltanet，它拥有更好的记忆管理机制，即使模型维度较小时，也能实现完美性能。

Deltanet与测试时训练

现在我们知道Deltanet在上下文检索中表现非常强大。

我想将Deltanet与一个非常近期的概念——测试时训练或测试时回归——联系起来。

这里有一个非常优雅的框架，称为测试回归。它将许多最近的自回归序列建模层分为几类。

在Deltanet中，有一个在线学习目标，它使用随机梯度下降来优化这个在线学习目标。

Deltanet的在线学习目标是键的转置和值之间的线性回归损失。在这种情况下，递归隐藏状态 S 被视为权重矩阵。这也有一个与快速权重编程非常有趣的联系，RNN的递归状态可以被视为快速权重（即每一步都在变化的权重）。

测试时训练框架也有这种解释，它将递归状态视为这个权重矩阵。因此，我们使用这个权重矩阵进行线性变换，然后希望恢复输入值。这是一个线性回归损失函数。

如果我们对这个目标函数执行单步梯度下降，我们将恢复Deltanet的更新规则。

因此，这里的 β_t 从先前上下文中的梯度下降角度来看是学习率。β_t 是写入强度，控制我们希望为历史关联值保留多少信息，以及我们希望为新传入的输入值使用多少新信息。

这样我们就可以恢复Deltanet的更新规则。

一个常见的问题是，我们为什么要使用这个在线学习目标？这是我之前演讲中经常遇到的问题。这里我给出一个非常直观的理解：对于上下文学习，我们希望使用先前上下文中提供的一些键值关联对。对于Softmax注意力，它使用KV缓存直接存储这些键值关联对以供后续使用。然而，对于RNN，它没有KV缓存来存储这些上下文信息，因此我们希望RNN的递归隐藏状态尽可能准确地编码那些键值关联对。

如果我们有这个学习目标，RNN模型直接优化这些键值关联预测，因此希望它能更好地记忆这些键值对。这就是为什么我们希望优化这些键值对之间的线性回归损失的一个非常直观的理解。

Deltanet的并行训练

接下来，我将简要讨论Deltanet的并行训练。

通过重新排列Deltanet的线性递归，我们可以得到这种形式。我跳过了一些推导，这是Deltanet的最终线性递归形式。

我们也可以将其写成与普通线性注意力非常相似的形式。在这里，我们可以非常高效地计算这个 u_t。然后，我们可以利用之前讨论的线性注意力内核进行硬件高效的训练。

首先，如果我们想并行训练模型，我们需要展开这个递归。通过展开递归，我们将有这个转移矩阵的累积乘积。在Deltanet的情况下，我们也有这个累积乘积转移矩阵。

对于任意的 D × D 转移矩阵，这个累积乘积计算将非常耗费计算资源。然而，对于Deltanet，由于它使用结构化的转移矩阵（该矩阵是单位矩阵加上低秩矩阵），存在现有算法可以进行快速计算。例如，这种形式非常类似于众所周知的Householder矩阵。当 β_j 等于2时，我们可以得到Householder矩阵，而Householder矩阵的乘积在数值分析领域已有深入研究。

在这项工作中，我们利用了一种称为WY表示的经典算法，以硬件高效的方式计算这种类Householder矩阵的累积乘积。这样我们就可以将这个累积乘积重写为使用WY表示的累积和。这使得优化变得容易得多，因为这个外积的累积和可以很容易地写成矩阵乘法。这与之前线性注意力的状态更新规则非常相似。

因此，我们可以将线性注意力的算法适配到这个情况。首先，对于状态更新，我们还需要考虑状态转移矩阵。对于普通的线性注意力，我们可以将状态转移矩阵视为单位矩阵，它基本上什么都不做。然而，在这种情况下，我们需要考虑这个块内转移矩阵的累积乘积。

这里我们需要计算这个块内转移矩阵的累积乘积，这可以写成这种形式。我们也可以将其写成矩阵乘法的形式，这样我们就可以有这个状态更新的矩阵乘法。

此外，我们还可以利用另一个众所周知的技术，称为实值变换，以非常硬件高效的方式计算这个 w 向量和输入 u 向量。

我不想详细讨论这个，因为它涉及较多的数学。但关键直觉是，我们可以使用一些矩阵方程来获得这个 C × C 矩阵，然后取其下三角部分并将对角线元素设为1，接着求解这个下三角矩阵。我们可以使用前向替换来计算这个矩阵。

一旦这个 C × C 矩阵被计算出来，我们就可以使用矩阵乘法来计算这个 w 和 u 矩阵。

如果你

8：D-Matrix Corsair 🚀

概述

在本节课中，我们将学习D-Matrix公司推出的Corsair架构。这是一种专为生成式AI推理设计的存内计算解决方案。我们将从设计动机开始，深入探讨其架构、软件栈以及如何高效地服务于大语言模型推理。

---

章节 1：设计动机与核心优势 🎯

随着生成式AI推理工作负载日益成为主导，传统的计算范式面临挑战。过去，图形处理的需求催生了GPU，显著提升了性能。如今，AI训练同样受益于GPU的高吞吐量。然而，在推理阶段，特别是随着思维链、长文本生成等技术的发展，对低延迟和高效率的需求变得尤为关键。Corsair架构正是为解决这一“内存墙”问题而诞生。

“内存墙”指的是计算吞吐量与内存/互连带宽之间日益扩大的差距。对于LLM解码这类计算强度低、内存访问频繁的任务，传统架构的带宽限制成为主要瓶颈。

Corsair的核心优势在于其存内计算范式。它将计算单元与存储紧密结合，数据无需在远距离内存间频繁搬运，从而大幅降低了延迟并提升了能效。

---

章节 2：Corsair系统架构总览 🏗️

上一节我们介绍了Corsair的设计动机，本节中我们来看看其高层次的系统架构。

Corsair架构在多个层面进行了创新：

• 硬件层面：拥有巨大的片上内存，提供了极高的内存带宽。通过先进的芯片间互连技术，可以轻松扩展到整张卡乃至整个机架。
• 计算层面：采用存内计算单元，并原生支持MX微缩块格式等量化技术，在保持精度的同时实现高算力。
• 软件层面：构建了完整的软件栈，包括自定义内核和专为低延迟批处理推理设计的推理服务引擎。

这些组件协同工作，共同应对LLM推理的挑战。

---

章节 3：深入Corsair芯片架构 🧩

了解了整体架构后，我们深入到芯片内部，看看它是如何实现存内计算的。

Corsair采用Chiplet（小芯片） 设计。一个PCIe卡上包含两个多芯片模块，每个模块又由四个Chiplet组成。这种设计提高了良率、降低了成本，并便于扩展。

Chiplet的层级结构如下：

1. Quad（四核组）：基本的独立可编程单元，包含约50 TOPS的计算能力。
2. Slice（切片）：每个Quad由多个Slice组成。
3. Apollo Core（阿波罗核心）：Slice中的核心计算单元，即存内计算阵列。它像数字化的Tensor Core，将权重静态存储在内存中，流式输入激活值进行计算，避免了模拟电路的不稳定性。

关键公式：算术强度
算术强度 = 计算操作数（FLOPs）/ 内存访问量（Bytes）
LLM解码的算术强度极低，因此是内存瓶颈型任务。Corsair通过存内计算，从根本上减少了数据移动需求。

工作负载编排
Corsair使用分层队列这一创新的硬件数据结构来管理任务。它在硬件中镜像了软件定义的依赖关系图，支持高度并发操作，并能在硬件层面原生支持自回归，动态处理不断增长的上下文长度，无需重新编译计算图。

---

章节 4：系统级扩展与互连 🔗

单个芯片的能力有限，AI推理需要规模扩展。Corsair在系统层面也进行了精心设计。

扩展层级：

1. 芯片内：Chiplet内部通过高速互连实现全连接。
2. 卡内：通过Die-to-Die链路和PCIe连接多个Chiplet，形成统一的张量并行单元。
3. 卡间：使用专用的DMX桥接器（类似NVLink）连接多张卡。
4. 机架/集群级：通过定制的、支持TCP流式传输的网络接口卡，利用标准以太网进行扩展，可连接成千上万的Corsair卡。

通信优化
架构优先支持聚集（Gather） 操作而非规约（Reduce）操作，并内置多播支持，这减少了通信开销和精度要求，特别适合LLM中常见的张量并行模式。

---

章节 5：软件栈与内核编程 💻

强大的硬件需要灵活的软件来驱动。Corsair的软件栈主要包括降低层和执行层。

降低层负责将高级模型（如PyTorch模型）转换为能在Corsair上运行的低级指令。主要有两种方式：

1. 基于编译器的路径：使用MLIR编译器进行自动化图优化和代码生成。
2. 模型构建器路径：通过手工编写或组合内核来构建计算图。

什么是内核？
内核是一段（通常用Python编写的）代码，它将一个操作（如矩阵乘法）转换为一个由机器指令构成的有向无环图。

内核编程示例
以下是一个简化的内核调用示例，展示了如何分配张量并调用矩阵乘法内核：

# 模型构建器代码（用户编写）
# 1. 分配张量（指定逻辑形状和物理内存布局）
input_tensor = allocate_tensor(shape=(M, K), layout=...)
weight_tensor = allocate_tensor(shape=(K, N), layout=...)
output_tensor = allocate_tensor(shape=(M, N), layout=...)

# 2. 实例化内核并调用
mm_kernel = MatrixMultiplyKernel(resources=my_quad)
graph1 = mm_kernel.multiply(input_tensor, weight_tensor, output_tensor)
# 可以连接多个内核
graph2 = mm_kernel.multiply(graph1.output, another_weight, final_output)

# 3. 代码生成
isa_graph = codegen([graph1, graph2])

内核开发者需要精细管理数据在SRAM缓冲区（输入、权重、输出缓冲区）中的移动，以实现计算与通信的重叠，从而最大化硬件利用率。

动态形状支持
对于KV缓存等动态增长的数据，Corsair的ISA指令支持在运行时重写参数。内核会插入特殊的占位符，由运行时传入具体的迭代步数等信息，从而动态调整内存访问模式。

---

章节 6：运行时与推理服务引擎 ⚙️

最后，我们来看执行层，它负责在主机上管理和服务推理请求。

Aviator运行时
这是一个模型无关的图执行层。它的目标是接收计算图及其输入张量，在加速器上异步执行，并将结果返回给用户。

核心机制：

• DMX Tensor：抽象了设备上的张量，管理设备内存、数据类型，并可自动处理主机-设备间的数据搬运和类型转换。
• 异步图执行：主机将任务写入设备的“任务队列”，设备运行时从中获取并执行。任务完成后，通过回调机制通知主机，从而实现高效流水线。

Aviator推理引擎
这是模型感知的智能服务层，支持多模型和分布式推理。

架构特点：

• 极简主机开销：精心设计进程间通信，确保其不在关键路径上。Worker之间无直接通信，所有集合通信都嵌入计算图中由硬件执行。
• 分布式管理：采用“一个Worker对应一个加速卡”的模式。多个Worker组成张量并行单元，多组Worker形成流水线并行阶段。引擎是大脑，负责调度、管理KV缓存、执行LLM状态机（如前缀、解码阶段）。
• 阶段化执行：引擎通过“阶段运行器”来执行不同阶段（如前缀、解码），阶段运行器再调用底层的运行时来执行具体计算图。

这种设计使得Corsair系统能够以极低的每Token主机开销，服务于低延迟、高吞吐的LLM推理请求。

---

总结 🎓

本节课我们一起深入学习了D-Matrix Corsair架构。我们从生成式AI推理的挑战和“内存墙”问题出发，探讨了Corsair如何通过存内计算这一范式转变来应对这些挑战。我们剖析了其从Chiplet、Quad到Apollo Core的层级化硬件架构，以及支持大规模扩展的互连方案。在软件层面，我们了解了其兼具灵活性和性能的内核编程模型，以及高效管理任务和服务的运行时与推理引擎。Corsair代表了一种为特定AI工作负载（尤其是LLM推理）进行全栈协同设计的先进思路。

9：Exo 2 - 发展一种调度语言

在本节课中，我们将学习一种名为Exo的编程语言，它用于编写高性能计算内核。我们将探讨其核心设计理念，即如何通过用户可扩展的调度操作和库来赋予性能工程师最大控制权，同时保持编译器的安全保障，从而高效地优化大量内核。

概述：硬件潜力与软件挑战

随着摩尔定律的终结，行业转向多架构和专用加速器（如GPU）来驱动性能提升。现代硬件（如Apple M4）拥有许多专用加速器。然而，没有合适的软件来发挥其潜力，这些尖端硬件就毫无用处。

硬件没有合适的软件来利用其潜力，就只是一块砖头。我们需要能够最大限度利用现代硬件能力的软件。问题在于如何做到这一点。

考虑一个简单的矩阵乘法。计算本身很简单。但如果你直接用Python实现，性能甚至达不到单核CPU峰值的0.1%。即使使用C语言和最高优化设置，性能也不会超过CPU潜力的1%。这是因为简单代码与优化代码之间的差距巨大。

行业领先的GEMM实现（如OpenBLAS）必须绘制复杂的分块策略图，最终形成大量看起来像这样的C和汇编代码。优化一个内核已经足够复杂和耗时，但性能工程师在地球上需要优化的内核不止一个。

例如，基础线性代数子程序（BLAS）是一个线性代数库规范。这个BLAS规范有许多内核，指定了不同的向量、矩阵运算，每个都有不同的内核规范和维度。仅BLAS Level 2的11个内核规范就产生了50个内核变体，这甚至不包括不同的硬件目标。

因此，性能工程师需要高效地优化许多内核，而不仅仅是一个内核。好消息是，已经存在许多编程语言来帮助性能工程师提高效率。这些语言有不同的抽象和方法，但都必须平衡一个权衡：性能工程师希望对优化进行控制以达到峰值性能，而编译器则自动化其他方面以提高效率。

在自动化决策和留给性能工程师的决策之间存在着不可避免的张力。所有编程语言都试图为其抽象和应用定义一个良好的边界。例如，一个边界决策是功能等价性。编译器可以确保原始代码和优化代码的功能等价，这样性能工程师就可以花时间进行实际优化，而不是调试索引错误。

这些语言被称为用户可调度语言（USL）。在USL中，用户首先定义要计算什么作为算法规范。在Exo语言中，算法规范看起来像是一个普通的Python循环，可以简单地映射到C语言。除了算法规范，性能工程师还提供描述程序优化（即如何计算原始算法）的调度。

优化可以使用诸如reorder和split之类的原语来描述。这些原语描述了算法的代码转换。例如，reorder就是重新排序内外循环，split将一个循环分成外循环和内循环。这些通常被称为循环级优化，USL通常支持这些优化。

非常重要的一点是，编译器会检查这些原语转换的等价性。因此，当编译器根据算法和调度生成低级复杂的目标代码时，我们知道输入算法和优化代码在功能上是等价的。这不仅是为了美观，更是为了提高效率。具有保证等价性的调度让性能工程师能够专注于创造性和实际优化程序，而不是修复复杂的索引数学问题，尤其是在为加速器获取正确的索引数学通常非常困难的情况下。

这是USL的高级概述。这很棒，我们获得了功能等价性。但USL像其他编程语言范式一样，存在一个大问题：编译器自动化在失效之前一直有效，而当编译器自动化失效时，性能工程师别无选择，只能用C或汇编编写低级代码。

以一个实际的USL为例，一个流行的USL叫Halide。Halide让性能工程师能够控制数据局部性、并行性和冗余计算，并自动化其他方面（如边界推断和指令选择）。这很好，Halide在许多用例中被证明非常高效，但它并不总是有效。例如，即使作为一个成熟的系统，Halide的指令选择近年来也不得不通过名为Braid和Pareo的论文进行改进。这是因为从根本上说，当Halide性能工程师想出比Halide自动化更好的指令选择时，他们别无选择，只能修复编译器本身，或者放弃并直接编写C代码。这给Halide开发人员带来了持续改进指令选择的巨大压力。

此外，直到今天，Halide仍然不支持像Tensor Core或Gemini这样的张量加速器。因为一般来说，加速器倾向于向软件暴露更多控制，这使得控制与自动化的边界划分变得非常不平凡。例如，张量加速器和现代GPU向软件暴露了非常复杂的内存管理，无论我们喜欢与否，处理它们对于实现峰值性能至关重要。但这是一个悬而未决的问题：USL应该自动生成这些显式内存管理指令，还是应该向性能工程师暴露这些内存管理指令？如果是后者，应该如何暴露？我现在不打算回答这个问题，只是为了让你了解为加速器指令定义这个边界有多么不平凡。

当我们转向这个边界的另一边时，预取和存储折叠在Halide最初设计时是自动化的，但几年后不得不跨越这个边界，因为使用Halide的性能工程师希望自己控制这些决策，因为他们可以做得更好。

实际上，我的观点是，即使在成熟的系统中，这种控制与自动化的边界也不完美，而且常常是模糊的。这是因为当出现问题时，比如现有的自动化不起作用，性能工程师想要更多控制，除了修改编译器本身之外别无他法。

现有USL设计的另一个缺点是，它们并非真正为代码重用而设计。这些是Halide中Harris角点检测器和Anisotropic扩散的调度代码。你可以看到不仅在两个内核之间存在代码重复，甚至在一个内核内部也存在。这种存储-计算-向量的模式并非随机模式，而是称为滑动窗口优化，是图像处理中非常常见的优化模式。但是，如果没有重用这些优化模式的方法，你就必须一遍又一遍地编写许多相似但略有不同的调度逻辑。

你可以想象，这就像支持BLAS一样成为噩梦，因为你有很多维度的笛卡尔积，编写大量相似但略有不同的调度将非常痛苦。

因此，现有USL的设计目标是：首先，它们支持一组固定的调度操作符，这些操作符在人体工程学层面经过精心设计，针对特定的应用或一类优化。这种控制与自动化的边界被设计为固定的。其次，USL传统上没有良好的设施来支持调度代码的重用。由于具有这组固定的操作符，添加新的硬件目标需要非平凡的编译器更改。另一种看待这个问题的方式是，现有的USL在编写单个且相对较小的内核时最有用，有点像脚本DSL（如awk）非常适合编写快速或小型的一次性任务，但难以扩展到大型程序。

这是对现有USL的概述。但Exo的设计目标不同。在Exo中，我们优先考虑给予性能工程师最大控制权以实现最佳性能，因此我们将这个边界一直向右移动。我们认为编译器自动化应该最小化，专注于安全性分析，同时允许性能工程师做出关键决策，包括指令选择（这通常在其他USL中是自动化的）。这种默认给予最大控制权的设计的关键原因是，性能工程师无法获得比朴素编译器构建所能提供的更多控制。一个类比是，如果你想获得比C语言更多的控制，当你编写C代码时，你需要转而编写内联汇编，这当然是另一种语言。

当然，这种方法是以减少性能工程师的自动化为代价的，因为他们现在需要控制更多事情。我们不是将自动化内置于编译器本身，而是建议将自动化构建到用户代码作为库。这让我们可以为诸如具体化、边界推断甚至定义新硬件目标等事情提供自动化，而这些在现有编译器中通常是内置并固化在编译器中的。

因此，我们的设计目标有点像“USL界的C语言”。我们认为编译器默认应该提供最大控制权，并通过一系列库提供抽象。我们称这种设计为Exo编译，并实现了Exo语言来体现我们的设计。高层次上，我们不是拥有一组固定的调度操作符，而是希望使调度操作符可由用户扩展，这意味着自动化的生产力应该通过用户定义的库来实现，而不是通过编译器内置。因此，我们的原语被设计得尽可能低级，因为库可以提供更多自动化，但不能更少。这使得Exo能够轻松支持新的硬件加速器，并扩展到支持像完整BLAS这样的大型程序。

用户可扩展调度：从向量化案例看起

上一节我们介绍了Exo的设计理念是赋予用户最大控制权。本节中，我们来看看如何通过库来实现用户自定义的调度操作，以向量化为例。

为了让我们更好地理解，让我们看看Exo中具体的向量化过程。这是一个Exo中的SAXPY代码，其中C代表顺序循环，缓冲区x和y在DRAM中。我们想要尝试向量化这个标量SAXPY代码，目标是AVX2指令和单精度向量宽度为8。首先，我们像这样将循环分成两个io和ii循环。

但这仍然不完全像一个向量化循环，因为x和y仍然在DRAM中并且只是标量。要使用向量寄存器，我们需要将计算暂存到AVX2寄存器中。然后我们可以将这个寄存器维度扩展到8，并进行循环裂变，使得每个ii循环看起来像向量指令。第一个ii循环是从y加载到var1（AVX2寄存器），第二个循环是将alpha广播到var3（另一个AVX2寄存器）。

最后，我们可以用对实际向量指令的调用来替换那些看起来像向量指令的循环。最终，我们在Exo中得到了这个向量化代码。我们可以轻松地将此代码降低到C语言，并在AVX2机器上运行，这将比原始的标量SAXPY快得多。

回顾一下我们刚才所做的：我们从这个非常简单的标量SAXPY代码开始，对这个代码进行了四次转换。所有中间代码当然都在Exo中，并且转换的每一步都保持了功能等价性。再次强调，这种功能等价性非常重要，因为我们希望这个向量化代码仍然计算相同的SAXPY，而不是其他任何东西。我在这里描述的过程实际上是Exo调度过程，它通过应用诸如split、stage_compute等调度操作符，将这个简单的起始代码转换为其优化版本。

实际的调度代码看起来像这样：以标量SAXPY作为输入proc，通过split划分循环，暂存计算，并用向量指令替换循环。这是一个用户调度过程，我们认为这是优化代码的一种比手动编写最终向量化代码更高效的方式。

在这里，我们成功向量化了一个内核SAXPY。但SAXPY实际上只是BLAS的一小部分。回想一下之前的幻灯片，BLAS有许多内核和许多维度，SAXPY只是Level 1内核中的一个。作为性能工程师，你必须高效地优化所有这些维度，而不仅仅是一个SAXPY。让我们看看另一个名为SCAL的BLAS内核。SCAL是将向量乘以标量并写入同一个向量X，而SAXPY是累加到不同的向量Y。

SAXPY和SCAL是不同的内核，但实际上，我们可以应用完全相同的调度序列，并为两者都获得向量化代码。这很好，但不是很简洁，因为我们不希望代码库中到处都是重复的代码。软件工程师自然要做的事情是将这些常见的调度模式封装在一个函数中。也许我们可以将这个函数命名为simple_vectorize，并从那些SAXPY和SCAL中调用它。

高层次上，Exo允许用户在用户代码中完全定义一个新的调度操作符，比如这个simple_vectorize。这看起来可能微不足道，但在现有的USL（如Halide）中这是不可能做到的。我们想称之为用户可扩展调度语言。

我们确定了用户可扩展调度语言的三个关键设计特性。首先，调度需要一种转换或修改目标代码的方式，我们称之为“动作”。编译器需要提供一些内置或原始动作，以提供这种安全的低级转换。我们称之为“原始动作”，我们支持超过60个这样的动作。我们使用基于Z3的多面体依赖分析来检查这些转换的正确性，我将在本讲的下一部分更多地讨论这个分析。

这很棒，我们检查了这些原始重写的安全性，但你可能想知道用户定义操作的安全性检查是如何工作的。这在事后看来是一个显而易见的想法：因为每个原始动作都保证保持功能等价性，所以像simple_vectorize这样的用户定义函数（是这些动作的组合）也仅仅通过组合就保证了保持功能等价性。这非常强大，因为我们可以从更简单的操作构建更复杂的操作，并且基本上免费获得安全检查。实际上，那些调度操作符stage_compute、fission、replace_all也是用户定义的运算符，而不是编译器原语。这只是试图展示我们如何通过组合这些较小的函数来构建这个操作符库，以定义像simple_vectorize这样相当复杂的操作。

我们讨论了动作。在执行动作时，我们有时需要询问关于代码的问题，比如这个缓冲区是F32还是F64？我们需要一种询问代码问题的方式，我们称之为“检查”。最后，这听起来可能非常基础，但我们需要能够指向或引用对象的部分。

稳定引用：光标（Cursors）的作用

上一节我们讨论了如何通过组合动作来构建复杂的调度操作。本节中，我们来探讨一个关键机制——稳定引用，这是实现可组合调度的基础。

现在让我们谈谈引用。暂时，我们将忘记Exo和USL。编程语言中有许多种引用，比如指向值、字符串、数组或函数指针或其他数据结构的指针。但编译器的特殊之处在于，编译器中的引用指向另一段代码。编译器当然需要引用它们正在操作的代码，以便可以修改它。

Clang称之为AST匹配器，LLVM称之为模式匹配，它们在编译器中有各种用途。指向代码是编译器（如LLVM）非常基本的机制。假设我们有一个初始的LLVM IR，其中这个外三角形和这个外圆角矩形说明了整个LLVM IR，内矩形和这个高亮三角形表示IR中的这个子树。

LLVM使用这种模式匹配来指向一个子树，如果可以的话，它会简化这个指令。这对LLVM有效，但有一些限制。因为模式匹配是一次性引用，这意味着如果你想应用进一步的优化，你必须用新模式（如A）重新查询整个IR。这里的问题是你必须始终确切地知道在此过程的每个点上代码结构是什么样子。因此，当你用A显式重新查询IR时，你的模式匹配只针对具体的IR A，而不是其他任何东西。有了这个，就不可能通过定义来参数化这个一次性引用，因为你必须始终确切地知道在此调度过程的每个点上目标代码的结构。

对我们来说，对于USL，一次性引用是不够的，我们需要一个稳定的引用。假设我们有这个初始目标代码proc，proc指向这整棵树，一个引用c指向一个像这样的子树。

每个Exo调度操作接受一个过程proc和一个像这样的引用c。reorder_loops重新排序我们在上一张幻灯片中看到的循环，因此在应用此操作后，代码看起来像这样。更新后的proc将指向一个新的目标代码。我们希望c在转换后继续指向同一个子树，因为这显然对调度很方便，因为你可以继续使用相同的引用，而无需知道转换后IR的确切结构。

在Exo中，我们提出了称为“光标”的东西来支持这种稳定引用。光标由路径表示，该路径记录AST的导航，这相当简单。例如，这个底部显示了一个指向此目标代码y的光标的路径表示，body 1表示主体中的第二个语句，所以它指向这个循环。body 0表示主体中的第一个语句，所以指向这个赋值，rhs是这个赋值的右侧，lhs是这个二元操作add的左侧，所以y由这个光标路径表示表示。

好的，这并不难，但是在代码被调度修改后会发生什么？如果你展开这个j循环，y的这个路径表示基本上就变得无效了。所以路径表示本身并不十分稳定，我们需要除了这个之外的其他东西。

光标提供这种稳定引用的下一个特性称为“转发”，它定义了在应用某个动作后路径应如何更新。我们可以将这两个语句替换为单个语句。当我们这样做时，我们希望指向这些语句的光标在转换后继续指向这个替换后的语句。

这个转发规则可能看起来像这样：我们保留到子树其他部分的现有路径，到旧树的路径应保持有效。我们为所有60个编译器内置原语定义了这个转发规则。这就是我们如何通过光标实现稳定引用。

但我们想看看光标如何在实际的调度过程中使用。假设我们在左侧有一个调度库函数foo，函数参数proc指向对象过程，看起来像这样。光标c指向这个i循环，另一个光标d指向这个if语句，所以整棵树可能看起来像这样，采用向下遍历的光标路径表示。

第一个动作reorder_loops将重新排序光标c指向的循环，所以它将重新排序循环i和j，因此在重写后它将看起来像这样。之后，c的主体将指向这里，因为c指向i循环。fission将在光标之后进行裂变，所以在裂变后它将看起来像这样。reorder_before接受一个指向这个if语句的光标d。reorder_before将此语句重新排序到循环之前，所以它将看起来像这样。最后，将光标c指向的这个i循环替换为load，所以它变成这样。

因此，在此完整调度过程之后，这里的对象代码可能看起来像这样。光标c指向这个load，这可能是此对象代码的光标路径表示。但重要的是，回想一下，在函数foo开始时，c指向这个i循环，而同一个光标c设法在整个调度过程中稳定地指向同一个子树。这是通过我们定义的光标转发实现的。

我们想看看一个更现实的用户级调度函数。这个tile_2d代码非常简单，由三个原始调用组成：两个divide_loop和一个lift_scope。vectorize是simple_vectorize的更通用版本，它更复杂，因为vectorize本身调用了其他用户定义的函数，如parallel_reduce。

但这里的要点是，我们可以在库代码中构建相当复杂的操作符，如tile_2d和vectorize。重要的是，所有那些用户级函数（如vectorize）也保证保持功能等价性，正如我所解释的，仅仅通过组合。这是可能的，因为那些以红色高亮显示的光标在整个调度过程中提供了稳定的引用，并且它允许这些调度函数的参数化。

这些单独的用户定义函数的集合最终成为一个库。

在库中定义硬件目标

上一节我们看到了如何通过库函数构建复杂的调度操作。本节中，我们来看看Exo的另一个核心特性：如何在用户库中定义硬件目标，而不是将其硬编码在编译器内部。

现有的编译器在编译器本身内部实现新的硬件目标，因此当新硬件目标出现时，需要编译器编写者来支持它，因为硬件目标代码实际上嵌入在编译器内部。在Exo中，我们认为这个责任应该外部化并转移给用户或库实现者。我们认为这种外部化在处理专有硬件接口和开发新硬件时是有益的，因为ISA非常不稳定。

通过外部化硬件定义，你不必重复修改编译器，不必每晚获取完整更新并试图跟上主编译器开发的步伐。我认为拥有这样一个模块化的编译器系统是非常必要的，尤其是在我们现在拥有如此多商业加速器，并且它们每12个月就会大幅更新的情况下。

10：基于搜索的深度学习编译器 🚀

在本节课中，我们将学习Luminal项目，这是一个采用“搜索优先”方法的机器学习编译器。我们将探讨它如何通过极简的算子集和基于搜索的编译技术，来简化并加速深度学习模型的执行。

概述

传统的机器学习框架（如PyTorch）依赖大量手工优化的内核，导致代码库庞大且复杂。Luminal采取了不同的路径：它从一个极简的、由11个基础算子组成的集合出发，将复杂的神经网络操作表示为这些基础算子的组合。然后，它使用基于搜索的编译器技术，自动探索大量逻辑上等效但性能不同的计算内核，并通过性能分析找到最优实现。这种方法旨在将复杂性从手工编码转移到自动化搜索中，从而简化核心库并实现高性能。

核心概念：极简算子集

上一节我们概述了Luminal的目标，本节中我们来看看其核心设计理念：一个极简的算子集。

Luminal认为，大多数现代神经网络（如LLM、MoE、扩散模型）本质上都是线性代数运算。因此，它仅使用11个基础算子来构建所有复杂操作：

• 一元算子（对单个张量进行逐元素操作）: exp2, log2, sin, cos, sqrt
• 二元算子（对两个张量输入进行操作）: add, multiply, mod, less_than
• 归约算子: sum_reduce, max_reduce

通过组合这些简单算子，可以构建出减法、除法、矩阵乘法（MatMul）甚至卷积等复杂操作。例如：

• 减法: a - b 等价于 a + (-1 * b)
• 矩阵乘法: 可以通过广播乘法（multiply）后沿K维度进行求和归约（sum_reduce）来实现。

这种设计使得Luminal的核心库非常精简（约5000行代码），但直接执行这些基础算子组合成的计算图会很慢。这正是编译器的用武之地。

从计算图到编译器

上一节我们介绍了极简算子集，本节中我们来看看Luminal如何利用计算图和编译器来恢复性能。

Luminal将神经网络模型表示为静态的计算图（有向无环图）。与PyTorch等框架的动态图相比，静态图虽然灵活性较低，但为编译优化提供了坚实基础。模型一旦定义，其计算图结构（不包括权重和激活值数据）就是固定的。

这种极简且静态的表示，为编写编译器创造了有利条件。Luminal的核心工作就是一个机器学习编译器，它将这些由基础算子构成的、执行缓慢的高级计算图，转换并优化为高效的、针对特定硬件（如CUDA）的低级内核代码。

搜索：编译器的核心

上一节我们了解了编译器的角色，本节中我们深入探讨Luminal编译器的核心机制：基于搜索的优化。

传统编译器依赖预定义的启发式规则和优化顺序，这在面对复杂且追求极致性能的ML工作负载时显得力不从心。Luminal借鉴了AlphaGo等解决复杂问题的思路，将内核编译问题转化为一个搜索问题。

其核心思想是：与其依赖复杂的、可能出错的启发式规则来决定如何优化（例如，循环是否展开、分块大小是多少），不如系统地搜索所有可能的、逻辑上等效的内核实现空间，并通过实际性能分析（Profiling）找出最快的那个。

以下是搜索过程的关键步骤：

1. 构建搜索空间：使用egglog（一种基于e-graph的等价逻辑编程语言）和一系列重写规则，将初始计算图扩展成一个包含数百万个可能内核实现的巨大搜索空间。
2. 探索与评估：采用各种搜索策略（如蒙特卡洛树搜索MCTS、束搜索Beam Search）来探索这个空间。对于搜索到的候选内核，会生成实际的CUDA代码并进行性能分析。
3. 选择最优：最终选择在目标硬件上运行最快的那个内核实现。

这种方法让编译器“不知道”最佳优化是什么，而是通过搜索“发现”它。

技术细节：中间表示与代码生成

上一节我们介绍了基于搜索的编译理念，本节中我们来看看其背后的具体技术表示。

Luminal使用一种自定义的中间表示来刻画计算。这种IR非常接近硬件执行模式。以下是一个对10个元素先取指数再取正弦的简单操作示例：

; IR示例：计算 sin(exp2(input))
(loop_in 10 stride_z) ; 循环读取10个元素，步长为z（此处为1，表示连续访问）
  (exp2)              ; 对读取的元素进行指数运算
  (sin)               ; 对上一步结果进行正弦运算
(loop_out stride_z)   ; 将结果写回内存，步长同样为z

这个IR可以确定性地生成对应的CUDA内核代码。编译器搜索过程会应用重写规则来变换IR。例如，一个重要的优化是循环融合，它将多个逐元素操作合并到同一个循环中，避免中间结果写回全局内存的开销，从而显著提升性能。

实际案例：自动发现Flash Attention

上一节我们看了基础优化，本节中我们来看一个更激动人心的案例：自动发现Flash Attention。

通过设置通用的代数重写规则（约20-30条）并进行搜索，Luminal编译器能够从朴素的注意力计算实现出发，自动推导出Flash Attention算法。这大约需要应用12到14个连续的重写步骤，包括循环融合、在线Softmax重写、循环交换和循环分块等。

这个过程证明了，通过系统性的搜索，编译器可以自动完成那些通常需要人类专家深度介入的、复杂的算法级优化，而不仅仅是局部代码变换。

扩展能力：训练与未来方向

上一节我们看到了搜索在推理优化上的威力，本节中我们看看Luminal的扩展能力。

尽管最初为推理设计，但由于其核心表示极其简单，Luminal可以轻松地通过自动微分（AutoGrad）来支持训练。其自动微分引擎仅需约150行代码，可以作为外部插件添加到核心库中。它通过链式法则从正向计算图派生出反向计算图，然后将正反向图融合后进行统一的编译优化。

关于未来，Luminal计划：

• 支持更多硬件后端：在现有CUDA、Metal、CPU支持基础上，增加对AMD、TensorRT、TPU等的支持。
• 分布式计算：将搜索空间扩展到包含设备间数据通信，以自动优化张量并行、流水线并行等拓扑结构。
• 云端服务：提供云编译和服务器端推理服务，利用其编译技术实现快速冷启动和高效执行。

总结

本节课中我们一起学习了Luminal，一个基于搜索的深度学习编译器。它的核心思想是通过一个极简的算子集来简化机器学习框架的复杂性，并将性能优化的重任交给一个能够自动探索巨大优化空间的搜索编译器。这种方法不仅使核心库保持小巧和可维护，还能自动发现人类专家级别的优化（如Flash Attention），为应对未来日益复杂的硬件和模型提供了有前景的解决方案。

11：多GPU编程入门教程

概述

在本节课中，我们将学习多GPU编程的基础知识。我们将从一个单GPU的示例应用出发，逐步了解如何将其扩展到多个GPU上运行。课程将介绍三种常见的编程模型或库：MPI、NCCL和NVSHMEM，并解释它们各自的适用场景、工作原理以及如何通过性能剖析来验证优化效果。我们的目标是理解应用程序的关键路径，以及不同编程模型中的延迟如何影响最终性能。

为什么使用多GPU？

在深入技术细节之前，我们首先需要明确使用多GPU的两个主要原因。

1. 强扩展：你有一个固定大小的问题，但需要更快地解决它。例如，天气预报如果在预报时间之后才得出结果就失去了意义。使用更多资源来解决同一个问题，以获得更快的速度。
2. 弱扩展：你的问题规模在增长，但你希望解决这个更大规模的问题时，所用时间保持不变，或者你只是需要更多的内存等资源。

此外，一个更实际的原因是，当今的系统规模庞大且日益增长，通常配备多个GPU。为了高效利用这些并行系统，你的代码必须具备并行能力。

示例应用：Jacobi求解器

为了理解这些编程模型，我们需要一个具体的例子。我们将使用一个求解二维拉普拉斯方程的Jacobi求解器。

公式：
∇²u = 0 （在特定域上，并带有边界条件）

我们有一个二维计算域，其左右边界采用狄利克雷边界条件（固定值），上下边界采用周期性边界条件（可以想象成圆柱体的表面）。核心求解算法是一个迭代过程，每次迭代中，网格中每个点的新值由其上下左右四个邻居点的平均值更新。

核心计算代码（单GPU）：

for (int iy = iy_start; iy < iy_end; iy++) {
    for (int ix = ix_start; ix < ix_end; ix++) {
        u_next[iy * NX + ix] = 0.25f * (u_curr[iy * NX + (ix+1)] +
                                         u_curr[iy * NX + (ix-1)] +
                                         u_curr[(iy+1) * NX + ix] +
                                         u_curr[(iy-1) * NX + ix]);
    }
}

代码循环遍历整个计算域，根据四个邻居更新每个点的值，然后交换新旧缓冲区，检查收敛性，并进入下一次迭代。

扩展到多GPU：域分解

要将计算扩展到多个GPU，我们需要进行域分解。这意味着将整个模拟域分割成多个块，让每个GPU只计算其中的一块。

我们可以选择水平分割、垂直分割或棋盘状分割。本教程中，我们选择水平分割为条带，因为这样处理起来最简单：每行数据在内存中是连续的，便于与邻居交换数据，且需要交换的数据总量不大。

进行域分解后，每个GPU只拥有本地的一部分数据。但是，在更新边界上的点时，它需要邻居GPU上的数据。例如，一个GPU顶部边界点的更新需要其上方邻居GPU底部边界点的值。因此，我们需要在GPU之间交换这些边界数据，这个过程称为Halo交换。

编程模型一：MPI

MPI（Message Passing Interface）是一个成熟且广泛使用的标准库，用于在不同进程间交换数据。

MPI基础

一个典型的MPI程序骨架包括初始化、获取进程标识（rank）和总数（size）、执行通信、以及最终化。

代码示例：

#include <mpi.h>
int main(int argc, char** argv) {
    MPI_Init(&argc, &argv);
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    // ... 你的计算和通信代码 ...
    MPI_Finalize();
    return 0;
}

在Jacobi求解器中实现Halo交换

在Jacobi求解器中，我们需要与上方和下方的邻居GPU交换一行边界数据。MPI提供了MPI_Sendrecv函数，它可以组合发送和接收操作，有助于避免死锁。

Halo交换代码示例：

// 发送顶部行给上方邻居，同时从下方邻居接收数据到本地底部Halo区
MPI_Sendrecv(&u[1*NX], NX, MPI_FLOAT, top_rank, 0,
             &u[(local_ny+1)*NX], NX, MPI_FLOAT, bottom_rank, 0,
             MPI_COMM_WORLD, MPI_STATUS_IGNORE);
// 发送底部行给下方邻居，同时从上方邻居接收数据到本地顶部Halo区
MPI_Sendrecv(&u[local_ny*NX], NX, MPI_FLOAT, bottom_rank, 1,
             &u[0], NX, MPI_FLOAT, top_rank, 1,
             MPI_COMM_WORLD, MPI_STATUS_IGNORE);

这段代码从当前进程的角度，完成了与两个邻居的双向数据交换。

GPU Direct 技术

现代MPI实现（如CUDA-aware MPI）支持GPU Direct技术。这意味着当MPI函数接收到一个GPU内存指针时，它能识别出来，并可能启用以下优化：

• GPU Direct Peer-to-Peer (P2P)：在同一节点内的多个GPU之间，通过NVLink直接传输数据，无需经过CPU内存。
• GPU Direct RDMA：在网络节点之间，网络卡（如InfiniBand）可以直接读写GPU内存，实现远程直接内存访问。

这些技术能显著降低通信延迟，提高带宽利用率。在性能剖析工具中，你可以看到数据通过PCIe或NVLink直接传输，而没有额外的内存拷贝内核。

进程启动配置

通常，我们采用一个MPI进程对应一个GPU的配置。这样管理起来最简单，也能最好地配合GPU Direct技术。使用mpirun或集群作业调度系统（如Slurm）可以方便地启动相应数量的进程。

优化技术：计算与通信重叠

在基础的MPI实现中，通信操作（Halo交换）发生在计算内核完成之后，这会导致GPU在通信期间空闲，降低了并行效率。

我们可以通过计算与通信重叠来优化。思路是：由于内部点的计算不需要边界数据，我们可以将一次内核启动拆分为三次：

1. 计算顶部边界行。
2. 计算底部边界行。
3. 计算内部区域。

我们为边界计算和内部计算使用不同的CUDA流（Stream）。在边界计算完成后，即可开始Halo交换，而此时内部区域的计算仍在另一个流上继续。这样就隐藏了部分通信时间。

性能提升：随着使用的GPU数量增加，每个GPU的本地计算量减少，通信开销占比增大。重叠技术能有效提升强扩展效率，在GPU数量较多时优势更明显。

编程模型二：NCCL

NCCL（NVIDIA Collective Communication Library）是NVIDIA优化的通信库，专为GPU缓冲区设计。

NCCL 特点

• GPU内核通信：NCCL使用GPU内核来执行通信操作。
• 流感知：NCCL调用可以指定CUDA流，从而能与计算内核在流中自动排序，减少显式的流同步。
• 拓扑感知：NCCL了解系统拓扑（GPU如何连接），并选择最佳的数据传输路径。
• 操作分组：可以将多个通信操作组合在一起，减少内核启动开销，并允许库进行更深度的优化。

在Jacobi求解器中使用NCCL

我们可以用NCCL的点对点发送（ncclSend）和接收（ncclRecv）操作替换MPI的Sendrecv。由于NCCL操作在流上排队，我们可以移除之前需要的流同步点。

NCCL通信示例：

ncclGroupStart();
ncclSend(&u[1*NX], NX, ncclFloat, top_rank, comm, stream);
ncclRecv(&u[(local_ny+1)*NX], NX, ncclFloat, bottom_rank, comm, stream);
ncclSend(&u[local_ny*NX], NX, ncclFloat, bottom_rank, comm, stream);
ncclRecv(&u[0], NX, ncclFloat, top_rank, comm, stream);
ncclGroupEnd();
// 不再需要显式的 streamSynchronize

进一步重叠优化

为了确保边界计算在内核计算之前完成，以便通信能尽早开始，我们可以使用流优先级。为负责边界计算和通信的流设置高优先级，为内部计算流设置低优先级。这样GPU调度器会优先执行高优先级流上的任务。

结合NCCL的流排序和流优先级，我们可以实现比MPI重叠版本更好的计算与通信重叠效果。

编程模型三：NVSHMEM

NVSHMEM是一个基于分区全局地址空间（PGAS）模型的库。它实现了OpenSHMEM标准。

NVSHMEM 核心概念

• 对称堆：所有处理元素（PE，相当于一个GPU进程）共享一个虚拟的、分区的全局地址空间。使用nvshmem_malloc分配的内存位于对称堆中，所有PE都可以直接访问。
• 单边通信：通信操作（如nvshmem_float_p）由发起方PE主动执行，目标PE无需发布匹配的接收操作。这简化了编程模型。
• 设备端通信：最重要的特性是，可以直接从CUDA内核内部调用NVSHMEM函数进行通信。

在Jacobi求解器中使用NVSHMEM

在CUDA内核中，每个处理边界点的线程可以直接使用nvshmem_float_p将单个浮点数put到邻居PE的对应内存位置。这实现了极细粒度的、与计算交织的通信。

设备端通信代码示例：

// 在内核中，位于顶部边界的线程
if (iy == 1) { // 假设第1行是顶部边界
    nvshmem_float_p(&u_remote_top[ix], u_local[iy*NX + ix], top_pe);
}
// 位于底部边界的线程
if (iy == local_ny) { // 假设最后一行是底部边界
    nvshmem_float_p(&u_remote_bottom[ix], u_local[iy*NX + ix], bottom_pe);
}

通信模式与内核融合

NVSHMEM支持多种粒度的设备端通信：

• 单元素Put：每个线程传输一个数据元素。适合NVLink等低延迟互连。
• 块级Put：线程块协作传输一批数据。适合InfiniBand等需要较大消息以获得高带宽的网络。
• 完全融合的内核：将整个求解循环（包括计算、通信和迭代控制）全部放入一个巨大的内核中。这需要在内核内实现PE间的同步（例如，使用nvshmem_int_atomic_add操作计数器）。这种方式完全消除了多次内核启动的开销，对于需要强扩展至非常多GPU的应用潜力巨大。

应用案例：量子色动力学代码QUDA通过使用NVSHMEM将多个内核融合为一个，在512个GPU上相比MPI实现获得了约1.6倍的性能提升。

编程模型对比与选择

以下是三种编程模型的简要对比：

特性	MPI	NCCL	NVSHMEM
通信模型	双边（发送/接收需匹配）	双边	单边（Put/Get）
缓冲区支持	主机和设备（CUDA-aware）	仅设备	对称堆（设备）
流/排序	无直接流参数	有流参数，可排序	设备端调用与内核代码交织
关键优势	标准、通用、支持CPU/GPU	高性能集合操作、流感知、拓扑优化	设备端通信、简化算法、支持内核融合
典型使用场景	已有CPU MPI代码迁移；通用跨节点通信	需要流排序和重叠；高性能集合操作（如All-Reduce）	需要极细粒度或与计算紧密交织的通信；希望减少内核启动开销

如何选择：

• 从MPI开始：如果你的代码已有CPU上的MPI实现，迁移到CUDA-aware MPI是最直接的路径。
• 需要流排序和重叠：选择NCCL，它可以简化同步，并获得更好的计算通信重叠。
• 需要设备端通信或单边模型：考虑NVSHMEM，特别是当你的算法能从内核内部的细粒度通信或内核融合中受益时。

重要的是，这些库可以互操作。例如，你可以用MPI进行进程启动和粗粒度通信，同时用NCCL进行GPU间的集合操作，或者在关键内核中使用NVSHMEM。选择取决于你的具体应用需求和性能瓶颈。

总结

本节课我们一起学习了多GPU编程的基础。我们从单GPU的Jacobi求解器出发，探讨了通过域分解将其扩展到多GPU时面临的挑战，即Halo交换。

我们深入介绍了三种主流的编程模型：

1. MPI：成熟的标准，支持GPU Direct技术（如P2P和RDMA）以实现高效通信。优化时需注意计算与通信的重叠。
2. NCCL：NVIDIA优化的GPU通信库，支持流感知和操作分组，能更好地与CUDA计算集成，简化重叠优化。
3. NVSHMEM：基于PGAS模型的单边通信库，支持直接从CUDA内核发起通信，为实现内核融合和极细粒度通信提供了可能。

每种模型都有其适用场景，并且可以结合使用。理解应用程序的通信模式、关键路径以及对延迟和带宽的需求，是选择合适工具并进行有效优化的关键。

为了进一步实践，你可以访问 多GPU编程模型GitHub仓库，其中包含了使用不同模型实现的完整Jacobi求解器代码，供你学习和实验。

12：邻域注意力机制 🧠

概述

在本节课中，我们将要学习一种称为“邻域注意力”的稀疏注意力机制。注意力机制是现代基础模型的核心，但其二次复杂度在处理大量令牌（如高分辨率图像或长视频）时可能成为性能瓶颈。邻域注意力通过将每个查询的注意力窗口局部化到其周围邻域，旨在降低计算复杂度，同时保持模型性能。我们将从概念、实现挑战、优化策略以及实际应用等多个角度来探讨这一主题。

---

什么是注意力？

注意力是一种在两个令牌集合（查询集和上下文集）上进行的操作。在自注意力中，查询集和上下文集是相同的，我们只是将它们投影到不同的空间（Q, K, V）。无论查询令牌位于何处，它都会关注整个上下文。在二维视觉任务中，这可以看作是一个18x18的网格，每个查询都关注整个网格。

核心公式：
Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

---

从自注意力到邻域注意力

上一节我们介绍了标准的全局自注意力。本节中我们来看看如何通过引入空间局部性来降低其计算复杂度。

如果我们希望将注意力局部化，可以利用空间位置信息。我们可以将自注意力视为一种局部注意力，但其窗口大小是整个输入。通过逐渐减小这个窗口，我们便得到了邻域注意力。

在邻域注意力中，对于每个查询令牌，我们将其上下文窗口限制在其周围的一个局部邻域内。窗口大小由用户定义，越小则计算量越少。极端情况下，当窗口大小为1x1时，注意力退化为线性投影，因为单个元素的softmax结果总是1。

核心概念：

• 窗口大小：定义了每个查询可以“看到”的邻域范围。
• 局部化：将计算从 O(N^2) 降低到 O(N * window_size)。

---

为什么需要稀疏注意力？

注意力操作是大多数基础模型的核心。由于其二次复杂度，在处理大量令牌时（例如图像/视频生成模型的预填充阶段或去噪扩散模型），注意力可能占据超过50%的端到端推理时间。因此，优化注意力计算对于提升大模型效率至关重要。

稀疏注意力通过跳过计算某些注意力权重来减少计算量。邻域注意力是实现稀疏性的一种方式，它利用了数据在空间上的局部性。

---

实现挑战：从密集到稀疏

上一节我们了解了邻域注意力的动机。本节中我们来看看将其高效实现所面临的挑战。

标准的注意力实现依赖于密集的矩阵乘法，这能很好地利用现代AI加速器的专用硬件。然而，邻域注意力不再是纯粹的矩阵乘法问题，而更接近于向量-矩阵乘法，这可能导致硬件利用率下降。

最初的实现尝试需要从头编写CUDA内核。虽然可以使用注意力掩码来实现稀疏性，但如果不进行内核融合，就无法真正节省内存和计算。

一个更大的问题是，邻域注意力失去了全局上下文信息。为了解决这个问题，我们引入了膨胀机制。

---

引入膨胀以获取全局上下文

膨胀机制借鉴了卷积神经网络的思想。在邻域注意力中引入膨胀，允许查询在跨越更大步长的位置收集信息，从而在不增加窗口大小（即不增加计算成本）的情况下，获得更广阔的“视野”。

与卷积不同，膨胀的邻域注意力不需要进行与膨胀因子成比例的填充，这得益于注意力操作与卷积的本质区别。

实现技巧：膨胀的邻域注意力可以通过对输入进行预分区来实现，将额外的维度移到批处理维度，从而转化为一个批处理更大但问题规模更小的标准邻域注意力问题，这通常能带来更好的并行性和速度。

---

现有方法：分块注意力

在邻域注意力之前，一种流行的局部注意力实现是分块注意力。它将输入分割成连续的块，然后在每个块内独立进行自注意力。

以下是分块注意力的主要特点：

• 易于实现：只需在PyTorch中进行张量分区和重塑，无需修改底层注意力内核。
• 保持密集计算：每个块内仍是密集的矩阵乘法。
• 潜在缺点：可能破坏平移等变性，导致在某些任务上性能低于滑动窗口（邻域）注意力。

---

内核级实现：融合稀疏注意力

为了真正获得性能提升，我们需要在内核级别支持稀疏性。这通常通过块稀疏性技术实现。

在融合注意力内核（如FlashAttention）中，计算是分块进行的。如果能够提前预测哪些块完全被注意力掩码覆盖（即不包含任何有效的点积），就可以跳过加载这些块的K和V值，从而节省计算和内存访问。

块稀疏性的工作原理：

1. 外层循环遍历查询块。
2. 内层循环遍历键值块。
3. 如果根据掩码判断某个键值块完全不需要，则跳过该次迭代。

对于分块注意力，其块稀疏掩码仅在对角线上有值，因此甚至可以不需要修改内核，只需通过张量变换将问题转化为并行的小规模自注意力即可。然而，对于像邻域注意力这样的滑动窗口模式，其掩码模式更复杂，必须修改内核以支持非对角的块访问。

---

多维度的挑战与FNA方法

上一节我们讨论了在1D序列上实现稀疏注意力的通用方法。本节中我们来看看当问题扩展到2D或3D时（如图像、视频）带来的独特挑战。

在1D中，平铺（Tiling）是线性的。但在2D或3D中，如果仍然使用1D的平铺策略，可能会加载大量在空间上相距很远但在内存中相邻的令牌，导致计算浪费。这被称为“多维度的诅咒”。

FNA的解决方案：
我们的“快速邻域注意力”方法对内核进行了几项关键修改：

1. 多维平铺：将内核的平铺尺寸重新解释为多维形状（例如，将64解释为8x8），使得加载的Q和KV块在空间上保持局部性。
2. 动态KV平铺：不是平铺整个上下文，而是根据每个查询块的位置，动态地裁剪出其感兴趣的KV区域再进行平铺，避免了加载无关的块。
3. 细粒度掩码：在内核的softmax/掩码阶段，根据邻域注意力的公式计算每个Q-K对是否需要被掩码。

然而，FNA内核基于较旧的xFormers FMHA内核构建，其性能上限可能不如最新的FlashAttention。此外，软件中的多维平铺和细粒度掩码会引入显著的指令开销。

---

广义邻域注意力与步长参数

为了在性能和模型质量（如平移等变性）之间取得更好的权衡，我们引入了步长参数，形成了广义邻域注意力。

步长的作用类似于卷积中的步长。它将空间上相邻的查询分组，组内的所有查询共享同一个注意力窗口。这相当于在滑动窗口中引入了“跳跃”。

步长的影响：

• 步长=1：标准的邻域注意力，每个查询有自己的窗口。
• 步长增大：查询被分组，计算更粗粒度，可能节省更多计算（实现更高的块稀疏度），但可能损失一些局部细节建模能力。
• 步长=窗口大小：退化为非重叠的分块注意力。

因此，广义邻域注意力统一了从滑动窗口（邻域注意力）到非重叠窗口（分块注意力）的整个谱系，允许用户通过调整步长来精确控制效率与精度的权衡。

---

性能分析与NAttn模拟器

面对如此多的参数（窗口大小、步长、膨胀、输入尺寸等），手动评估每种配置的性能潜力是困难的。为此，我们开发了NAttn模拟器。

该模拟器是一个分析工具，它接受用户的使用场景参数（张量形状、注意力模式参数）和内核设计选择（平铺尺寸等）。它不实际运行内核，而是通过分析计算每个查询块需要访问的KV块数量。

模拟器的用途：

• 预测速度提升：通过统计需要访问的块数，可以估算相对于全自注意力的理论速度提升（基于块稀疏度，而非FLOPs）。
• 指导参数选择：例如，可以扫描所有可能的步长值，找出能达到完全块稀疏性（零计算浪费）的最小步长，从而在保持一定模型质量的前提下获得最佳性能。
• 降低开发成本：在投入时间编写新内核之前，可以先评估其潜在收益。

---

新一代内核：Hopper与Blackwell上的FNA

为了在最新的Hopper和Blackwell GPU架构上获得最佳性能，我们重新设计了FNA内核，采用了令牌置换策略。

核心思想：在运行注意力内核之前，先通过一个内存拷贝操作，对输入令牌进行重新排列（置换），使得在内存中连续的令牌在空间上也连续。这样，注意力内核内部就无需处理复杂的多维平铺逻辑，可以像处理1D问题一样工作，大大简化了内核设计。

优点：

• 内核简化：注意力内核只需支持块稀疏性和细粒度掩码，无需处理多维平铺的复杂性。
• 开销可控：置换操作只是一个内存拷贝，在Blackwell等高性能GPU上，其耗时可能仅占整个注意力计算的很小一部分（例如1%）。
• 性能提升：在许多情况下，可以接近理论上的块稀疏速度提升。

注意事项：置换/逆置换操作在某些模型架构中可能只需在开始和结束时各执行一次，从而进一步分摊开销。

---

实际应用与性能结果

我们将广义邻域注意力应用于视频生成模型（如Cosmos和Hounian）。通过结合局部和膨胀的注意力模式，并在训练中引入稀疏性，我们取得了显著的效果。

以Cosmos模型为例：

1. 分析阶段：在预训练的全自注意力模型上，测量不同稀疏注意力模式在各层引入的误差。
2. 选择策略：为每层选择误差低于阈值且稀疏度最高的模式。
3. 微调训练：以稀疏配置恢复训练，模型能快速恢复损失的精度。

结果：

• 在2B参数模型上，实现了1.9倍至2.6倍的端到端速度提升。
• 在14B参数模型上，实现了1.7倍至2.1倍的速度提升。
• 视觉质量在多数样本和基准测试中得以保持。

这表明，通过精心设计的稀疏模式和适当的训练，可以在保持模型质量的同时，获得显著的推理加速。

---

NAttn工具库简介

我们提供了开源的natten库来支持邻域注意力的研究和应用。

主要功能：

• 多种算子：提供1D、2D、3D的邻域注意力函数和PyTorch模块。
• 灵活配置：支持设置核大小、步长、膨胀、因果掩码等参数。
• 多后端支持：
  • CUDA：基于cuDNN、FlashAttention的高性能后端。
  • FNA：我们自研的融合邻域注意力内核（支持旧架构）。
  • FMHA (Cutlass)：针对Hopper/Blackwell优化的新内核。
  • FlexAttention：实验性后端（支持动态稀疏模式）。
• 性能分析工具：内置分析器，可自动为给定问题选择最佳后端和配置。

使用示例（通过分析器）：

# 分析特定配置下的性能
python -m natten.profile --shape 30 48 80 --kernel-size 18 24 24 --stride 16 8 8 --backend fmha_blackwell

---

总结

本节课中我们一起学习了邻域注意力机制。我们从标准自注意力出发，探讨了其计算瓶颈，并引出了通过局部化来实现稀疏化的思想。我们详细介绍了邻域注意力的定义、其与膨胀结合以获取全局信息的方法，以及它与分块注意力的区别。

在实现层面，我们深入探讨了在内核中支持稀疏性的技术（块稀疏性、细粒度掩码），以及处理多维数据时面临的挑战和解决方案（多维平铺、动态平铺）。为了平衡效率与模型质量，我们引入了广义邻域注意力，通过步长参数统一了滑动窗口与分块模式。

我们还介绍了用于性能预测的NAttn模拟器，以及为新一代GPU架构设计的高效内核（采用令牌置换策略）。最后，通过在实际视频生成模型上的应用案例，我们展示了邻域注意力能够带来显著的端到端加速，同时保持模型质量。

邻域注意力及其广义形式为高效Transformer模型的设计提供了强大的工具，特别是在视觉、视频等具有强空间局部性的领域。

13：引言与概述

在本节课中，我们将要学习如何利用经典电子游戏环境来评估和训练大型语言模型。我们将探讨为什么游戏是一个有价值的评估平台，以及如何设计一个有效的“游戏测试框架”来克服模型在游戏中的固有弱点。

欢迎来到新一期的GPU MODE。今天，我们很高兴邀请到来自加州大学圣地亚哥分校Haal实验室的Ln Shanghu，他将与我们分享一个相当不同的项目——关于“Alum Game benchch”，即通过电脑游戏来评估AI模型。

虽然这不像典型的“让GPU更快”的演讲，但如果你从事机器学习系统工作，这些话题具有普遍的相关性。我们迫不及待地想听听他的分享。谢谢Ella Chang的介绍。

谢谢Mark的亲切介绍。大家好，我是Ln Xianghu，来自加州大学圣地亚哥分校的二年级博士生。我的研究方向是大规模机器学习工作流的效率问题，以及为AI模型评估带来新的视角。今天，我将分享我们最近关于如何通过游戏评估和训练LLM的工作，以及为什么我们认为游戏环境很重要。

---

课程名称：GPU编程与AI评估新视角：章节编号：2：章节名称：游戏作为评估平台的动机

上一节我们介绍了课程主题，本节中我们来看看为什么选择游戏作为AI模型的评估环境。

在过去的几个月里，我们一直在“回收利用”许多传统的经典游戏环境，并将其用于AI评估，并得到了一些有趣的结果。我们或许可以先看几个演示，了解为何使用游戏环境进行评估，以及在此过程中学到的一些经验。

演示：马里奥游戏中的模型对比
这是一个使用马里奥游戏进行不同模型对比的例子。如图所示，Claude 3.7赢得了游戏。这项评估是在三月份进行的，反映了当时模型的水平。当时，人们普遍认为模型在数学和代码推理任务上表现很好，但使用智能体进行评估开始兴起。然而，当时还没有在游戏环境中评估模型的工作。

游戏对人类来说是非常直观且相对简单的任务，我们小时候都玩过游戏。但模型在游戏环境中的表现如何呢？实际上，结果令人惊讶。例如，GPT-4o在游戏中很快就失败了，Claude 3.4和3.7也没有坚持多久。虽然GPT-1.5表现稍好，但仍存在巨大差距。

这引出了一个更高层次的问题：我们为什么要研究游戏？在LLM时代之前，强化学习训练已经使用了大量环境，但并非针对LLM，而是针对传统的RL智能体，如深度Q学习。在过去，人们大量使用游戏环境，例如OpenAI Gymnasium和Stable Retro等代码库，它们“回收”了许多经典游戏。

以下是游戏环境作为评估平台的优势：

• 定义清晰：游戏具有定义明确的状态、动作和可验证的奖励，便于计算奖励和优势函数以训练模型。
• 规模庞大：存在大量流行且为人熟知的游戏，易于扩展评估规模。

基于此动机，我们开始思考：为什么不先使用这些环境进行基准测试，看看模型表现如何？在此过程中，我们发现直接将模型放入游戏环境效果不佳，无法提供有意义的评估。因此，我们需要所谓的“测试框架设计”，我们将在后面深入探讨。

---

课程名称：GPU编程与AI评估新视角：章节编号：3：章节名称：游戏选择与挑战

上一节我们了解了游戏作为评估平台的优势，本节中我们来看看如何选择合适的游戏以及面临的挑战。

选择合适的游戏比想象中更难，需要满足以下条件：

• 难度适中：游戏难度需在不同LLM间保持适中。例如，不能选择《英雄联盟》或《Dota 2》，因为它们对当前模型来说太难了，模型缺乏良好的物理感知能力。
• 类型多样：游戏类型需要足够多样化。
• 广为人知：应使用流行游戏，使结果不仅对AI研究人员，也对普通观众易于理解。

基于以上原则，我们选择了多种类型的游戏（论文发表于五月）：

• 解谜/棋盘游戏：如《推箱子》、《2048》、《糖果粉碎传奇》、《俄罗斯方块》。这些游戏广为人知。
  • 《推箱子》和《2048》需要长视野规划。
  • 《糖果粉碎传奇》和《俄罗斯方块》则需要大量空间推理能力。
• 实时动作游戏：如《超级马里奥兄弟》。
• 视觉小说游戏：这类游戏文本量重，类似于侦探解谜游戏，需要与多个角色互动。例如游戏《逆转裁判》。

然而，如果直接将LLM“开箱即用”地放入游戏环境，它们通常表现不佳，主要原因有三点：

1. 视觉感知能力弱：模型容易混淆图像元素。
2. 延迟高：许多推理模型响应延迟长。在实时游戏中，等模型响应后，游戏可能已经结束了。
3. 重复性错误：模型因缺乏记忆机制而反复犯同样的错误。

因此，为了有效区分模型能力，在游戏环境中评估模型需要一个智能体工作流程，这就是我们所说的“游戏测试框架”。

---

课程名称：GPU编程与AI评估新视角：章节编号：4：章节名称：游戏测试框架设计

上一节我们讨论了直接评估的挑战，本节中我们来看看为解决这些问题而设计的模块化游戏测试框架。

我们以模块化的方式设计了一个游戏测试框架，以便研究不同组件如何影响模型的游戏性能。该框架包含三个主要组件：

1. 视觉感知模块
此模块用于解决视觉感知弱的问题。它通过从游戏后端读取状态，或查询另一个视觉语言模型，将图像元素转换为文本描述。例如，在《推箱子》游戏中，模型可能将墙误认为箱子。提供真实游戏状态或使用VLM提取元素可以缓解此问题。

2. 记忆模块
此模块对智能体至关重要，尤其是在游戏中。以《推箱子》为例，如果没有记忆，人类通常会制定计划（例如先推箱子A，再处理箱子B）。但我们观察到LLM缺乏这种长视野规划能力，它们会在不同选项间犹豫不决。记录游戏历史可以帮助模型进行更一致的规划。

3. 推理模块
此模块整合来自视觉感知和记忆模块的所有信息，以生成动作。可以选择开启或关闭长链思维推理模式。

框架工作流程概述
我们首先将游戏环境设计为符合Gym API，包含观察空间、动作空间以及以文本（表格）或图像表示的游戏状态。游戏状态被输入模型，模型可以选择性地启用智能体组件（如记忆或感知模块），然后生成动作。该动作被输入游戏环境以驱动状态变化，产生下一个状态，如此循环。

---

课程名称：GPU编程与AI评估新视角：章节编号：5：章节名称：不同游戏类型的评估策略

上一节我们介绍了通用测试框架，本节中我们来看看针对不同游戏类型的具体评估策略和挑战。

棋盘/解谜游戏（如《推箱子》）
我们发现，LLM最大的瓶颈是游戏状态理解。如果只提供图像，效果不佳。因此，我们为其准备了文本格式的状态表示，例如2D ASCII表格或对象列表，明确告知模型每个元素的位置坐标。

实时游戏（如《超级马里奥》）
挑战有所不同，因为这类游戏对延迟敏感。我们需要指定一个动作对应多少帧。例如，我们规定模型生成 jump 10 这样的格式，表示按住跳跃键10帧。
我们还观察到一个有趣的现象：“知与行的差距”。模型在推理过程中能制定合理的计划（例如，“为了安全越过管道，应保持动量和高度”），但在执行时，却无法将计划映射为精确的按键操作。
此外，还存在帧率低的问题。模型只在动作完成后接收到新的状态。当马里奥在空中时，如果没有记忆机制，模型不知道其动量和运动轨迹。换句话说，游戏状态输入过于稀疏。

侦探游戏/视觉小说（如《逆转裁判》）
这类游戏的挑战是数据污染问题。由于游戏文本量重，在LLM预训练时，很多游戏信息（如游戏论坛内容）可能已被模型见过，导致严重的记忆问题。
我们进行了一个小实验来测试数据污染：询问模型关于游戏第一轮的剧情，发现Claude 3模型的输出与维基百科文本几乎完全相同，存在严重的记忆现象。
那么，记忆问题是否会影响游戏性能？我们进行了相关性研究，发现记忆程度与游戏性能排名之间存在强相关性（记忆越强，排名越好，即数字越小）。为了缓解数据污染，我们采用了以下方法：

• 实体替换：将游戏中的具体名称替换为通用名称（如“律师”代替“成步堂龙一”）。
• 提示词引导：在提示中明确指示模型进行推理而非依赖记忆。
• 文本重写：在评估前，以不同的方式重写游戏背景文本。
  应用这些缓解策略后，相关性减弱，表明记忆对游戏性能的影响减小了。

---

课程名称：GPU编程与AI评估新视角：章节编号：6：章节名称：测试框架的有效性与排行榜

上一节我们探讨了不同游戏的评估策略，本节中我们来看看测试框架是否有效，以及当前的模型表现如何。

在引入测试框架后，它是否有效？我们进行了“测试框架有效性评估”，即比较应用测试框架前后的模型游戏性能。结果显示，性能提升的百分比变化非常显著。通过统计假设检验，大多数游戏上的性能提升是显著的。
例外是《超级马里奥兄弟》，因为它是一个部分可观测的随机游戏，即使应用了测试框架，游戏性能的方差仍然很大，统计上不足以证明测试框架完全解决了问题。

在解决了数据污染问题并设计了测试框架后，我们建立了排行榜。排行榜分为两部分：

1. 模型排行榜：不应用任何测试框架，直接比较模型原始性能。
2. 智能体排行榜：应用所有测试框架组件后的性能。
   目前，Claude 3模型表现最佳。Claude 3 Pro表现更好，但由于成本高昂，我们尚未能进行完整的多次评估。GPT-2.5 Pro紧随其后。

---

课程名称：GPU编程与AI评估新视角：章节编号：7：章节名称：游戏性能的含义与分析

上一节我们看到了排行榜，本节中我们深入探讨一个核心问题：良好的游戏性能意味着什么？

我们主要从两个角度研究这个问题。

角度一：游戏性能与其他基准的相关性
我们研究了游戏性能与现有其他基准（如数学、代码）之间的关系。我们收集了约20个基准上的模型排名，通过计算斯皮尔曼相关系数来比较。
以下是主要发现：

• 《推箱子》、《俄罗斯方块》、《2048》与解谜类基准有强相关性，这符合预期。
• 《推箱子》与一些数学和编码基准也有较强的相关性。
• 《逆转裁判》则与语言基准强相关。
  这表明，游戏性能实际上是多种现有能力的综合体现，包括编码、数学、空间推理和解谜能力。

角度二：能力分解与线性模型预测
我们将能力分解为几个类别：物理、数学、编码、空间推理、语言。一个有趣的问题是：能否基于现有基准的排名，学习一个线性模型来预测游戏性能？如果可以，我们就能得到一组权重，揭示不同能力对各类游戏的贡献。
分析发现：

• 《推箱子》、《俄罗斯方块》、《2048》与数学和编码能力高度相关。
• 《逆转裁判》当然与语言能力相关。
• 《超级马里奥》和《糖果粉碎传奇》则与空间推理和物理基准相关。这对于《超级马里奥》来说很有趣，因为它需要良好的物理感知来控制角色执行动作。

---

课程名称：GPU编程与AI评估新视角：章节编号：8：章节名称：基于游戏的模型训练与泛化

上一节我们分析了游戏性能的构成，本节中我们来看看通过在游戏环境中训练模型，能带来哪些泛化能力。

我们使用强化学习在游戏环境中训练模型。由于游戏奖励是稀疏的（离散的分数），我们使用策略梯度方法。我们将每个回合的状态、响应和奖励连接成一个超长序列，计算每个回合的总奖励。正向奖励意味着游戏进展或成功，而长时间无进展则会累积小的惩罚。
我们训练使用的模型是Qwen2-57亿参数指令微调版。重要提示：我们目前只训练LLM，而非VLM。游戏状态仅以文本表格形式提供。

训练结果与泛化
在《推箱子》和《俄罗斯方块》上训练后，我们观察模型的泛化能力：

• 域内泛化：在训练环境（6x6网格《推箱子》）上，成功率从约11%显著提升至24%。
• 更难环境的泛化：在更大的8x8网格《推箱子》上，性能也提升了约3%。
• 跨游戏泛化：仅在《推箱子》上训练的模型，在《俄罗斯方块》（尤其是较简单变体）上的性能也有所提升。
• 规划任务泛化：在“积木世界”规划任务上，性能也有提升。
• 智能体任务泛化：在“WebShop”多交互任务上，成功率从7%提升至约20%。
  然而，在数学和编码任务上，我们没有看到明显的性能提升，甚至有轻微下降。

扩展环境类型与数据混合
另一个问题是：如果同时训练多种游戏甚至混合其他数据（如数学），模型能否更好地泛化？我们发现，在混合数据上训练，模型在各项任务上均有提升，但上限不如在单一游戏上专门训练的效果好。如何找到能带来广泛泛化的数据配方，仍是一个挑战。

关于“思维链”格式的发现
一个有趣的观察是关于“思维链”格式。我们使用的Qwen2基础模型本身不具备推理格式。在训练中，模型同时学习性能和输出格式。我们发现，如果训练时包含思维链令牌，初始性能实际上更差，因为格式对模型来说是分布外的。性能提升可能部分源于模型学会了这种格式，而非掌握了新能力。我们下一步计划先对模型进行思维链格式的微调，再进行RL训练。

---

课程名称：GPU编程与AI评估新视角：章节编号：9：章节名称：对现有评估的反思与未来方向

上一节我们探讨了训练与泛化，本节中我们以对当前热门评估的反思作为结尾，并展望未来方向。

对《宝可梦红》评估的反思
在最新的模型技术报告中，常看到对《宝可梦红》的评估。但我们认为，它目前并非一个好的评估基准，主要原因如下：

1. 战斗控制：仅涉及移动光标和使用技能，过于简单。
2. 导航：空间推理对当前模型来说太难。现有评估通常在地图上覆盖显式网格，告知模型每个格子的属性和坐标，这引入了大量人为框架。
3. 长视野规划：成本极高。例如，到达第一个检查点就需要上千步，花费高达50-120美元，且耗时很长。

一个更有趣的方向是：不让模型执行低级动作，而是让其调用工具（如A*搜索算法）。模型只需做出高级决策（如“去房子A”），然后调用工具寻路。这将大大减少步数，并将评估重点转向模型的工具使用能力。

结论与总结
本节课中我们一起学习了利用游戏环境评估和训练AI模型的全过程。

• 游戏环境的价值：游戏环境提供了丰富的任务空间，适用于评估和训练AI智能体。
• 测试框架的必要性：通过游戏测试框架和缓解策略，我们揭示了当前模型行为的局限性。
• 性能分析：我们通过免训练和训练导向的方法，研究了游戏性能的含义。
• 未来工作：我们正在寻找能在游戏环境中训练并实现良好泛化的新训练方法。

游戏环境为评估和训练AI智能体（特别是基于LLM的智能体）提供了丰富的资源。通过游戏测试框架和缓解策略，我们深入了解了LLM的优势和局限。我们使用免训练和训练导向的方法来探究游戏性能的本质。目前，我们也在致力于寻找能在游戏环境中实现良好泛化的训练方案。

---

问答环节要点

• 高层规划与低级执行：未来方向是让模型进行高层规划并调用工具（如路径搜索），而不是执行每一个低级动作。
• 添加自定义游戏：建议从已有良好支持的游戏开始，添加自定义游戏需要更多工程工作。
• 测试框架的“作弊”边界：测试框架只是一个临时解决方案，用于在当前阶段区分模型能力。最终，足够强大的AI不应需要它。我们尽可能使用最少的框架，并确保其能有效区分模型。
• 具有商业价值的游戏：经济模拟、战略类游戏可能具有现实世界用例，是未来的探索方向。目前正在进行多智能体评估的研究。

14：NCCL与NVSHMEM 🚀

概述

在本节课中，我们将学习NVIDIA的两个核心通信库：NCCL和NVSHMEM。我们将探讨它们的设计理念、历史背景、核心差异以及各自的适用场景。通过具体的代码示例和性能基准测试，我们将理解如何根据不同的应用需求选择合适的通信模型，并了解它们在AI训练（特别是MoE模型）和高性能计算中的关键作用。

---

历史背景与核心概念

上一节我们概述了本课程的目标。本节中，我们将深入了解NCCL和NVSHMEM诞生的历史背景及其核心设计哲学，这有助于我们理解它们为何以当前的形式存在。

MPI：可移植性与双边通信

MPI（消息传递接口）是高性能计算领域的基石。它的核心优势在于极致的可移植性和基于双边通信（发送/接收）的模型。

• 设计哲学：MPI诞生于CPU相对网络速度较快的时代，其设计优先考虑通用性和跨平台兼容性。它可以通过标准的TCP/IP套接字在任何系统上运行。
• 双边通信：发送操作和接收操作必须匹配，这结合了数据移动和同步。发送方和接收方需要协调（握手），数据才能传输。
• 公式/代码描述：

  // MPI_Send 和 MPI_Recv 必须成对出现
  MPI_Send(sendbuf, count, datatype, dest, tag, comm);
  MPI_Recv(recvbuf, count, datatype, source, tag, comm, &status);
  

• 局限性：对于GPU通信，MPI的一些通用性设计（如复杂数据类型、通配符匹配）可能带来不必要的开销，且其主机端发起的模型与CUDA流式执行模型的集成不够流畅。

SHMEM：分区全局地址空间与单边通信

与MPI不同，SHMEM代表了一种分区全局地址空间模型。

• 核心思想：所有处理单元（PE）的内存在一个逻辑上统一的地址空间中可见，但物理上分区。通过特殊的分配函数（如shmem_malloc）创建对称内存，所有PE都可以直接读写这部分内存。
• 单边通信：通信操作（如put、get）解耦了数据移动和同步。发起方可以直接向远程内存写入（put）或从远程内存读取（get），无需远程PE的显式参与。
• 公式/代码描述：

  // PE i 可以直接写入 PE j 的对称内存
  shmem_putmem(dest, source, size, j); // 单边 Put 操作
  // 同步需要显式调用，如 shmem_barrier_all()
  

• 优势：这种模型减少了通信中的握手开销，尤其适合不规则、数据驱动型的通信模式。

---

NCCL vs. NVSHMEM：为何两者并存？

上一节我们回顾了两种经典的通信范式。本节中，我们来看看NVIDIA如何将这些思想具体化到GPU领域，并解答为何需要NCCL和NVSHMEM两个库。

NCCL：为AI规模化训练而生

NCCL专为多GPU、多节点AI训练设计，优化了集体通信操作。

• 核心目标：高效执行All-Reduce、All-Gather、All-to-All等集体操作，以支持数据并行、张量并行、流水线并行和专家并行。
• 设计特点：
  1. 类MPI接口：提供类似MPI的集体通信原语，易于理解和使用。
  2. CUDA流集成：操作直接在CUDA流上排队，与GPU计算无缝重叠。
  3. 自动化调优：根据消息大小、GPU拓扑（NVLink, PCIe）和网络类型自动选择最优算法和协议。用户应“信任NCCL”做出最佳选择。
  4. 广泛的生态系统：集成在PyTorch、TensorFlow等所有主流AI框架中。
• 适用场景：大规模、结构化的集体通信模式，是分布式AI训练的默认选择。

NVSHMEM：GPU上的PGAS与设备端发起通信

NVSHMEM将SHMEM的PGAS模型引入GPU世界，并增加了关键创新：设备端发起通信。

• 核心创新：
  1. 对称内存：在GPU之间创建全局可访问的内存区域。
  2. 设备端API：CUDA内核中的线程、线程束或线程块可以直接调用nvshmem_put、nvshmem_get进行通信，无需返回到主机CPU。
  3. 设备端集体操作：通过协作组启动，可以在CUDA内核内部执行All-Reduce、Barrier等操作。
• 公式/代码描述：

  // 在CUDA内核中，线程可以直接发起单边Put操作
  __global__ void my_kernel(...) {
      nvshmemx_float_put_nbi_block(dest, source, nelems, pe);
  }
  

• 独特优势：
  • 超低延迟通信：避免了内核启动开销（>1微秒），支持计算与通信的细粒度交错。
  • NVLink负载/存储语义：在NVLink域内，远程内存访问可像本地加载/存储一样使用指针直接操作，完全消除显式通信调用。
  • 不规则通信模式：非常适合数据驱动、目标不确定的通信，如MoE模型中的专家路由。

---

实验：矩阵转置的性能对比

上一节我们从理论上比较了两者。本节中，我们通过一个具体的分布式矩阵转置基准测试，来直观感受不同通信模型的性能差异。

实验设置

• 硬件：DGX H100（8个H100 GPU，通过NVLink交换机互联）。
• 目标：实现一个分布式矩阵转置，产生All-to-All通信模式。
• 性能上限：
  • 本地HBM内存带宽：~24 TB/s。
  • NVLink交换机双向带宽：~3 TB/s。

测试的不同实现版本

以下是测试的不同通信策略实现：

1. NCCL All-to-All：使用NCCL的send/recv原语手动构建All-to-All。这是基准版本。
2. NCCL 点对点（Send/Recv）：尝试手动配对通信以交错计算。结果发现性能更差，因为NCCL内部会对单个流上的多个send/recv进行流量控制，以防止网络拥塞。
3. NVSHMEM All-to-All：使用nvshmem_alltoall。性能与NCCL All-to-All相近。
4. NVSHMEM Get（设备端）：在CUDA内核中使用nvshmem_get从远程GPU获取数据块，然后进行本地转置。
5. NVSHMEM 指针访问：利用NVSHMEM获取远程内存的直接指针。在CUDA内核中，像访问本地内存一样通过指针加载远程数据。这是最优雅且高性能的方式。

性能结果与分析

• NCCL All-to-All 和 NVSHMEM All-to-All 达到了约2.7 TB/s的网络带宽，接近硬件上限。整体应用带宽约为7.5 TB/s。
• NCCL 点对点 版本由于内部流量控制，性能显著下降。
• NVSHMEM 指针访问 版本实现了约9.5 TB/s的整体带宽，比基准提升了近30%。其秘诀在于：
  • 通信融合：将显式的数据移动（get）转化为内存控制器级别的NVLink加载/存储操作。
  • 代码简化：通信逻辑消失，代码看起来像是在处理一个更大的统一内存空间。

关键启示：在NVLink域内，利用NVSHMEM的指针访问能力，可以消除显式通信，将通信开销完美地隐藏在计算中，从而获得显著性能提升。

---

前沿应用：MoE模型中的NVSHMEM

上一节的实验展示了NVSHMEM在规则模式下的潜力。本节中，我们来看看它在当前最热门的AI模型——混合专家模型中的革命性应用。

MoE模型的通信挑战

在MoE模型中，一个门控网络（路由器）将每个输入令牌动态路由到少数几个“专家”子网络进行处理。这导致：

• 不规则All-to-All：每个GPU需要将不同数量的令牌发送给其他GPU，且接收方在数据到达前不知道来源和数量。
• 细粒度通信：通信粒度是令牌级别的，需要极低的延迟。

NVSHMEM的解决方案

DeepSeek和Perplexity等公司公开了其使用NVSHMEM优化MoE的方案：

1. 设备端路由与PUT：路由器在CUDA内核中直接决定令牌目的地，并使用nvshmem_put将令牌写入远程GPU的对称内存缓冲区。
2. 异步通知：使用nvshmem_put_with_signal等操作，在写入数据的同时更新一个远程计数器，通知专家数据已就绪。
3. 专家处理：专家GPU轮询其缓冲区，发现有新令牌到达即开始计算。
4. 结果收集：处理完成后，同样使用nvshmem_put将结果写回。

性能收益

根据Perplexity发布的基准测试，与基于NCCL All-to-All的原始实现相比，使用NVSHMEM的细粒度、设备端通信实现了超过10倍的延迟降低和吞吐量提升。这对于要求实时响应的AI应用（如搜索）至关重要。

---

未来展望与总结

上一节我们看到了NVSHMEM在尖端AI模型中的成功。最后，我们来展望这两个库的未来发展并总结全课。

NCCL的未来路线图

NCCL正在积极吸收NVSHMEM的优点，以提供更统一的编程模型：

1. 对称内存支持：已发布预览版，允许在NCCL通信器中创建共享地址空间。
2. 设备端API：计划中，将允许从CUDA内核内部发起NCCL通信。
3. 单边通信原语：未来版本将增加put/get等操作。
   这意味着，未来用户可以在享受NCCL强大生态系统和集体通信优化的同时，也能使用设备端发起和单边通信等高级特性。

如何选择？

以下是当前的选择指南：

• 选择 NCCL，如果你：
  • 主要进行大规模、结构化的集体通信（如All-Reduce梯度同步）。
  • 使用主流AI框架（PyTorch/TensorFlow），且不想深入底层通信。
  • 需要极致的可扩展性和跨平台支持（云、IB、以太网）。
  • 遵循“信任库”的原则，希望自动化调优。

• 选择 NVSHMEM，如果你：
  • 需要设备端发起通信以实现计算通信最大重叠。
  • 处理高度不规则、数据驱动的通信模式（如MoE、图计算、不规则网格）。
  • 希望在NVLink域内获得极致性能，使用指针访问消除显式通信。
  • 正在构建定制化的高性能计算内核（如3D-FFT）。

总结

本节课中我们一起学习了：

1. NCCL 和 NVSHMEM 的历史根源，分别源于MPI的双边通信范式和SHMEM的单边PGAS范式。
2. NCCL 是AI规模化训练的支柱，通过高度优化的集体通信和广泛的集成，提供了“开箱即用”的高性能。
3. NVSHMEM 通过引入设备端发起通信和对称内存，为GPU通信打开了新的维度，特别适合低延迟、不规则通信和利用NVLink负载/存储语义。
4. 通过矩阵转置基准测试，我们验证了不同通信模型的性能特征，并展示了NVSHMEM指针访问在NVLink域内的巨大优势。
5. 在MoE模型中，NVSHMEM的设备端、细粒度、单边通信特性使其成为实现超低延迟推理的关键技术。
6. 展望未来，NCCL正在演进以包含更多NVSHMEM的特性，为用户提供更统一且功能丰富的通信工具箱。

最终，选择取决于你的具体应用模式、性能需求和对编程模型的控制需求。理解两者的核心思想和优势，将帮助你在构建下一代高性能GPU应用时做出明智的决策。

15：GPU中心化通信全景图

在本节课中，我们将学习GPU中心化通信的发展历程、核心概念以及当前主流的通信库。我们将从为什么通信成为瓶颈开始，逐步深入到节点内与节点间的通信类型，并介绍MPI、NCCL、NVSHMEM等关键库的特点与差异。

概述：为什么通信至关重要

自21世纪初以来，计算进入了多核时代。2007年CUDA发布后，GPU编程逐渐普及。2010年，首批基于GPU的超级计算机登顶Top500榜单。如今，前十的超级计算机大多基于GPU，数据中心和AI领域也广泛采用GPU。计算能力的发展速度远超内存和网络带宽的提升速度，这意味着数据移动（即通信）的优化变得越来越重要。

例如，训练一个大型模型可能需要数千块GPU，成本高达数百万美元。优化通信不仅能节省时间，更能显著降低成本。因此，减少CPU在关键通信路径中的参与，赋予GPU更多自主权，即实现“GPU中心化通信”，是提升整体系统效率的关键。

通信技术发展时间线

以下是英伟达相关通信技术发展的简要时间线：

• 内存管理技术：如统一虚拟地址（UVA）。
• GPU Direct技术：允许GPU间直接通信。
• NVLink与NVSwitch：提供高带宽、低延迟的GPU互连。
• 用户级通信库：如NCCL和NVSHMEM。

节点内通信类型

我们根据通信过程中操作的执行者来定义四种节点内通信类型。关键在于区分API调用发生在主机（CPU）端还是设备（GPU）端，以及数据路径是否经过CPU。

类型一：主机原生通信

在早期，GPU之间没有点对点（P2P）访问能力。数据必须先从源GPU复制到CPU内存，再由CPU复制到目标GPU。API调用和数据路径都完全在CPU端。

类型二：主机控制通信

随着GPU Direct 2.0引入，GPU可以通过PCIe总线直接访问彼此的内存，无需经过主机内存复制。数据路径直接在GPU间通过PCIe或NVLink进行，但API调用仍在主机端。MPI、NCCL的cudaMemcpy操作通常属于此类。

类型三：设备原生通信

在具备直接P2P内存访问的基础上，将API调用也移至设备端。例如，NVSHMEM提供了设备端API，允许在GPU内核中直接发起通信操作。NCCL未来也计划支持设备端API。

类型四：主机回退通信

当P2P访问被禁用时，即使API在设备端发起，通信也会回退到经过主机的路径。这确保了代码的兼容性，但性能会下降。

节点间通信类型

节点间通信引入了网络接口卡（NIC）。除了API和数据路径，我们还需关注消息注册和触发这两个操作由谁执行。

类型一：主机原生通信

所有操作（API调用、数据移动、消息注册、触发）都通过主机完成。数据从GPU到NIC需要经过主机内存中转。

类型二：固定内存的主机原生通信

GPU Direct 1.0引入了GPU与NIC之间的共享固定内存区域。GPU可以将数据放入此区域，NIC直接读取，消除了主机端的一次数据拷贝，降低了延迟。

类型三：GPU Direct RDMA通信

这是节点间通信的一个重要里程碑。NIC可以通过PCIe直接读写GPU内存，完全绕开了主机内存。这进一步优化了数据路径，是MPI、NCCL、NVSHMEM等进行节点间GPU通信的常用机制。

类型四：GPU触发通信

此类通信优化了控制路径。CPU预先向NIC注册消息，然后由GPU触发（“按门铃”）通知NIC开始传输。这减少了CPU的参与，但仍需要一个运行在主机上的代理线程。

类型五：设备原生通信

这是真正的“CPU无关”通信，消除了代理线程。API调用、触发、消息注册等所有操作都移至GPU端。CPU仅用于初始的内核启动。NVSHMEM和配置了GPU Direct Async IB的InfiniBand支持此模式，但在多数超算中并未默认启用。

高速互连：NVLink与NVSwitch

PCIe带宽有限，NVLink提供了更高的带宽和更低的延迟，可实现高效的GPU间P2P通信。然而，当GPU数量较多时，无法实现全互联。NVSwitch的引入解决了这个问题，它提供了全互联路由能力。AMD的Infinity Fabric链路与NVLink类似，但目前没有类似NVSwitch的交换设备。

用户级通信库

上一节我们介绍了底层的硬件通信机制，本节中我们来看看构建在其之上的用户级通信库。

GPU感知MPI

MPI是高性能计算领域标准的并行编程库，具有极好的可扩展性和可移植性。GPU感知MPI指的是MPI实现能够自动区分主机和设备缓冲区。当进行发送/接收操作时，MPI会检查缓冲区位置，并自动选择最优路径（如使用GPU Direct RDMA），从而避免不必要的主机中转拷贝。其API始终在主机端调用。

以下是实现GPU感知的关键技术：

• 统一虚拟地址（UVA）：提供统一的地址空间，允许通过地址值推断内存所属设备，简化了编程。
• 进程间通信（IPC）：允许进程间共享GPU内存，通过传递内存句柄实现直接访问，避免了主机中转。这是目前节点内通信最常用的方式。
• GDRcopy：通过内存映射处理小消息，通常用于小数据量传输。
• GPU Direct：支持GPU间的直接数据传输。

UCX通信框架

UCX是一个统一的通信框架，旨在抽象底层多种传输技术（如InfiniBand、RDMA、TCP、GPU Direct IPC等）。上层库（如MPI、NVSHMEM）通过UCX的API选择最适合当前硬件和配置的传输方式。我们团队开发的UCtrace工具可以可视化基于UCX的MPI通信。

NCCL与RCCL

NCCL（英伟达）和RCCL（AMD）是厂商提供的集合通信库，针对NVLink、PCIe、InfiniBand等进行了深度优化。与MPI的主要区别在于它们原生支持CUDA流，允许将通信操作放入流中并按序执行，便于实现通信与计算的重叠。MPI目前尚无原生的流支持。

以下是一个简单的代码示例，对比MPI与NCCL的接口：

// MPI 示例
MPI_Isend(sendbuf, count, datatype, dest, tag, comm, &request);
MPI_Irecv(recvbuf, count, datatype, source, tag, comm, &request);
// 需要显式同步流以确保操作顺序
cudaStreamSynchronize(stream);

// NCCL 示例
ncclSend(sendbuf, count, datatype, dest, comm, stream); // 指定流
ncclRecv(recvbuf, count, datatype, source, comm, stream); // 指定流
// 操作在流中自动按序执行

NVSHMEM

NVSHMEM是一个基于分区全局地址空间（PGAS）模型的库。它提供单边的put和get操作，允许GPU像访问本地内存一样远程访问其他GPU的内存。它支持设备端API，最接近“CPU无关”的执行模式。

NVSHMEM的编程模型更灵活，可以实现细粒度的通信重叠，但同时也引入了复杂性，因为直接远程内存访问容易导致竞态条件，需要程序员仔细使用信号量等机制进行同步。调试和性能分析也更具挑战性。

以下是一个NVSHMEM设备端API的示例：

// 在GPU内核中直接进行远程写入
nvshmem_float_p(signal, value, pe_target);
// 需要显式同步以确保内存操作顺序
nvshmem_quiet();

通信库对比与性能考量

在高层次上对比这三种选项：

• MPI与NCCL：主要使用双边（发送/接收）通信。MPI缺乏原生流支持，NCCL有流支持并提供通信操作分组以进行拥塞控制。
• NVSHMEM：使用单边（存放/获取）通信，支持细粒度重叠，但编程和调试更复杂。

关于性能，最优选择取决于众多因素：工作负载特性、通信频率、消息大小、库配置、驱动版本、是在节点内还是跨节点通信、是否使用NVLink等。因此，很难给出普适的性能结论。建议根据具体应用场景进行测试和评估。

研究项目介绍：统一接口与可视化工具

面对众多通信库的复杂性，我们团队进行了一些探索性工作。

Unicorn项目：我们尝试创建一个统一的API，能够以最小开销支持上述多种通信库，同时兼容主机端和设备端API。初步结果表明，我们的封装层引入的开销可以忽略不计。

可视化工具：为了深入理解通信行为，我们开发了两种工具。

1. Snoopy：用于可视化NVSHMEM等基于P2P内存访问的通信，可以生成通信图和代码热力图，帮助定位引发通信的代码行。
2. UCtrace：专用于基于UCX的MPI应用，可展示使用了哪种传输层协议、数据量大小以及通信拓扑，有助于性能调试和系统配置验证。

总结

本节课中，我们一起学习了GPU中心化通信的全景图。我们首先了解了通信成为系统瓶颈的原因。接着，我们系统性地学习了节点内和节点间的通信类型分类。然后，我们深入探讨了MPI、NCCL、NVSHMEM等主流用户级通信库的特点、接口差异和适用场景。最后，我们介绍了一些旨在简化编程和提升可观测性的前沿研究工具。

理解这些通信技术的演变和权衡，对于在GPU集群上开发高性能应用至关重要。

16：Quartet 4-bit 训练教程

概述

在本节课中，我们将要学习 Quartet 方法，这是一种用于大型语言模型（LLM）的 4-bit 浮点数（FP4）训练技术。我们将探讨为何 FP4 训练在某些情况下可以是最优选择，理解其背后的核心理论，并了解实现高效 FP4 训练所需的关键技术，包括缩放定律、成本归一化比较以及高性能 CUDA 内核的实现。

---

模型训练与数据表示基础

上一节我们介绍了课程概述，本节中我们来看看机器学习模型训练的基本原理。

当我们谈论机器学习模型时，通常指的是某种参数化函数，这些函数通过优化某个目标来进行训练。优化过程通常使用基于梯度的方法，例如梯度下降、随机梯度下降或 Adam 优化器。这些方法的核心需求是能够计算损失函数相对于这些函数参数的梯度。

因此，这些参数自然被用作实数，因为实数参数函数是可微分的，并且可以轻松地对其应用梯度步进。然而，在现实中，计算机无法表示无限的实数集合，因此必须使用某种近似方法。人们通常使用浮点数来近似表示实数。

浮点数近似由两部分组成：尾数和指数（以及符号位）。它可以编码非常广泛的实数范围（例如从 10^38 到 10^-38），并具有高达 10 位小数的精度。这使得它成为实数的良好近似，因为它既有极宽的范围，又有足够高的精度。

然而，从大型语言模型分析的一些理论结果中，我们知道它们实际能存储在其参数中的知识容量，通常接近每个参数 1 到 2 比特。这意味着模型能编码到其参数中的真实知识量，远低于用于优化这些参数的高精度浮点数数据类型所能表示的量。这在我们使用的优化过程（使用非常昂贵的高比特宽、高精度浮点数）与最终结果（编码的信息量远未达到数据类型的上限）之间造成了某种差异。

因此，人们一直在广泛试验替代的数据类型，这些类型是对实数更粗糙、更基本的近似，具有更低的比特宽，但可能仍然允许高效的信息编码，更重要的是，允许在其上执行高效的优化过程。

---

低精度数据类型与硬件支持

上一节我们讨论了高精度浮点数与模型实际容量之间的差距，本节中我们来看看业界使用的低精度数据类型。

最简单的低精度表示是实数的固定子集，它比高精度浮点数覆盖的整个范围和精度要小得多，但可能仍然允许高效的训练和信息编码。

以下是不同公司和加速器最近使用的一些低精度数据类型示例：

• Apple：在其最新的 iPhone 和 iPad 芯片中，不仅包含高精度 FP32 数字，现在还包含标准的半精度 Bf16 数字和 8 位整数。这里的“包含”不仅指可以在内存中存储它们，更重要的是提供了对这些数字的内核支持，允许在机器学习模型中进行高效的矩阵运算和编解码。
• Qualcomm：采用了略有不同的方法，使用 FP16 代替 Bf16，并使用 8 位整数。近年来，支持的数据类型链变得有些复杂，尤其是在 NVIDIA 方面，它们不仅包括 FP32、16、8 甚至 4，还包括 Bfloat 数字和一些整数，并且经常添加或放弃对某些数据类型的支持。

因此，需要一个更健壮的框架来比较用于这些低精度数据类型的机器学习训练方案，以便更准确地判断哪种精度效果最好，以及这些精度在低精度机器学习训练中带来的权衡。

这些精度不仅仅是设备上存储数据的方式，更重要的是硬件对机器学习常用操作的支持。具体来说，对这些定制程序数据类型的多数支持来自矩阵乘法运算，因为这些运算通常构成神经网络内的大部分计算。因此，有效支持矩阵乘法运算几乎允许你以低精度运行整个模型。

在这项工作中，我们将专注于量化大型语言模型内部的矩阵乘法运算（即线性层），试图展示较低的精度可以带来更好的精度与速度权衡。理论上，FP4 的速度提升可以高达 FP16 的 4 倍或 FP8 的 2 倍。

但是，这些速度提升是以降低精度为代价的。具体来说，随着降低模型训练精度，在某个点上会开始损失精度。这就在我们因降低精度而损失的精度与从硬件支持的低精度数据类型中获得的速度提升之间产生了权衡。在实践中，这更多地取决于运行模型的具体设置。

---

精度与速度的权衡比较模型

上一节我们介绍了低精度训练带来的速度与精度权衡，本节中我们来看看如何系统地比较不同的训练方案。

设计低精度训练方法时，最简单的目标是证明它在提供速度提升的同时不损失精度。实际上，许多最近提出新的低精度训练方法的论文都展示了这一点，它们旨在表明几乎不损失任何精度，或者损失不显著，同时展示可以获得一些速度提升。

然而，这种比较存在一个问题：首先，它基于一个非常具体的设置。在这些论文中，作者通常训练一个特定模型，为其测量一系列任务，显示在所有任务中的微小性能下降，并声称几乎没有损失任何东西。他们为此特定模型提供速度提升，并说在使用低精度训练时，速度提升基本上是免费的。

但在实践中，这并没有回答任何关于如果（例如）将模型大小增加一倍会发生什么的问题。你的精度下降会如何变化？你的速度提升会如何变化？为了提供这些更具体的比较，需要以某种方式聚合关于不同速度提升下的性能和速度的信息，并进行同时考虑精度和速度的比较。因此，需要建立一个新的模型来进行这些比较。

我们希望我们的比较模型能够：

1. 首先，能够使用来自许多不同训练轮次的信息进行比较，以便更稳健地估计实际获得的收益。
2. 能够使用这个比较模型来外推我们尚未测试的设置，例如，如果我们将模型规模扩大 10 倍，我们的假设是否仍然成立。这非常有用，因为如果你的模型是稳健的，你可以使用它来在实际训练模型之前预测哪种程序会更好，从而在训练你能负担的最大模型之前更有效地分配计算资源。
3. 希望这个比较模型能够非常准确地描述不同设置、不同模型大小以及提供不同类型速度提升的略有不同的方法的精度-速度权衡。

---

构建基于缩放定律的比较模型

上一节我们提出了对健壮比较模型的需求，本节中我们开始构建这个模型，首先学习如何以更稳健的方式比较精度。

为此，我们采用了缩放定律。当我们说缩放定律时，意味着如果你的训练设置（广义上，例如一个模型家族）在超参数（如模型大小和训练数据量）方面具有一致的性能，那么你可以尝试将这种可预测的损失拟合为模型参数的参数化函数。

一个非常著名的 LLM 缩放定律例子是 Hoffman 等人提出的缩放定律，他们预测了验证损失作为模型参数数量和训练令牌数量的函数。这个缩放定律在文献中已被多次验证，并用于不同的目的。

我们可以将这个想法外推，将精度纳入方程。这是我们最终用于描述量化训练性能作为参数数量、训练令牌数、前向精度和反向精度函数的缩放定律。

这个缩放定律在形式上与之前的定律相似，参数数量和数据系数在某种程度上仍然是解耦的，仅通过一个额外的常数伽马相互作用。更详细地看，前向精度在这里仅影响参数数量，作为参数数量的一种乘法因子。自然地，由于量化训练无法超过基线全精度训练的性能，这个有效乘数介于 0 和 1 之间。

缩放定律中包含反向传递信息的另一部分是令牌数量的有效乘数。我们选择将其与参数数量解耦，基于一个简单的想法：如果你的模型训练不足，那么从优化理论中我们知道，反向传递量化不一定会在你的优化方法中引入任何一致的偏差，而只是使你的梯度估计过程更加嘈杂。从这个提出的缩放定律中，你可以看到令牌数量的有效乘数并不会改变具有量化/非量化反向部分的模型会收敛到的损失，而只会改变收敛速度。

因此，从比较单个设置，我们现在可以比较在不同训练时长下训练的整个模型家族。从这个缩放定律的形式中，我们可以看到，我们可以通过简单地比较有效因子来比较不同的方法。也就是说，如果一个方法具有更好的前向有效因子和反向有效因子（即参数数量乘数和数据量乘数都更高），那么该方法在各种设置下会产生更低的损失，我们可以说它在质量上更好。

---

将速度提升纳入成本考量

上一节我们通过缩放定律量化了精度，本节中我们来看看如何将速度提升也纳入比较框架。

我们需要考虑的是，前向传递和反向传递的速度提升并不等同。前向传递主要影响量化模型的推理，因为在训练期间，前向传递的量化会影响每个线性层的三个 GEMM 操作中的一个，但对于推理，所有 GEMM 操作都针对量化的前向传递进行量化。相反，反向传递主要影响训练，意味着如果我们在训练期间量化反向传递中的三个 GEMM 操作中的两个，我们将获得更快的训练，但在模型训练完成后对推理没有任何影响。

考虑到这一点，我们可以得出这样的观点：模型有一个单独的推理成本概念和一个训练成本概念，而量化方案可以以不同的方式影响这些成本。

让我们尝试阐述这个关于成本的想法。我们选择某个基线精度。对于固定精度下的基线，表征推理成本的一种方式是模型的参数数量，因为前向传递（即推理）的计算和内存移动与参数数量成正比。为了考虑不同精度可能具有某种推理速度提升（此处表示为 S_forward）的事实，要对不同精度的模型进行推理成本归一化，你可以选择与速度提升成正比的参数数量。

我们可以将同样的想法应用于训练成本，它与模型训练的令牌数量和参数数量都成正比。这里的归一化稍微复杂一些，因为由于参数数量的变化，你必须同时通过训练和推理速度提升进行归一化。

有了这些，你可以进行成本比较，即比较不同的精度，使得它们的推理成本和训练成本都相同。然后，你可以通过将这些精度（将产生缩放定律中的某些有效参数数量乘数）代入来考虑精度的损失。因此，这个公式化是对不同精度的推理和训练计算归一化比较。

我们使用它的方式如下：我们有一堆要比较的精度，我们为所有精度拟合缩放定律，然后我们选择一个基线精度，并基准测试每个其他量化方法相对于这个基线方法的速度提升。然后，对于任何训练和推理成本（例如，通过基线模型的参数数量和基线模型的数据饱和比率来参数化），我们可以使用缩放定律预测损失，并选择最低的一个来说明哪种精度、哪种确切的训练方案对于给定的推理和训练计算是最佳的。

我们可以像这样绘制它：在图中，Y 轴上是基线模型的参数数量（从 1 亿到 1000 亿个 Transformer 模型参数），X 轴上是数据饱和程度（从接近 Chinchilla 最优到远超 Chinchilla 最优训练时长）。我们比较了四种不同的训练方案：基线 FP8 训练、仅前向 FP4 训练（仅量化前向为 FP4，反向仍为 FP8）、前向 FP8 反向 FP4 训练以及完全 FP4 训练。在右侧的图中，你可以看到这些区域根据我们生成的内核的真实测量速度提升绘制而成。

这些区域显示了在真实世界的推理速度下，哪种计算成本归一化的方法在特定计算约束下会产生更好的精度。你可以看到，我们在图上绘制了一些最近发布的模型，如 Llama 3、Qwen 2.5 或 Gemini 3。我们可以看到，它们中的许多都落在要么是仅前向量化 FP4，要么是完全 FP4 量化的区域内。这就是为什么我们声称 FP4 对于 LLM 预训练可以是最优的，因为一些在实践中合理且人们已经训练过的模型落在了如果它们完全用 FP4 训练会产生更好损失的区域内，或者至少这是我们根据获得的缩放定律和速度提升所预测的。

这个比较方案表明，我们可以使用它来比较 LLM 预训练的真实内核，并使用在 LLM 预训练上拟合的缩放定律。我们可以将真实的预训练模型映射到获得的区域，以证明某些精度在特定训练机制中的可用性。

总的来说，这个比较方案允许我们对非常具体的计算设置进行基于缩放定律的预算归一化比较。

---

设计最优的 FP4 训练方法

上一节我们建立了一个可以比较不同方案的框架，本节中我们来看看如何利用这个框架设计出最优的 FP4 训练方法。

我们可以进一步应用这个方案，因为我们有“更好的有效因子导致在固定计算预算下更好的方法”这个概念。在固定精度下运行的方法通常具有相同或非常相似的计算预算，因为它们使用相同的数据类型和非常相似的内核。这意味着，如果我们想设计一个对固定精度最好的方法，我们可以直接使用这些有效因子来比较方法。

这就是我们在论文中所做的。我们方法的第三个组成部分是，我们使用这种有效系数直接选择最优的反向传递和前向传递方案。为了避免涉及太多细节，我们比较了许多方案，例如用于反向传递梯度估计的随机舍入（无偏）、EMA 舍入方案以及我们最近论文中基于标准差的标准差舍入方案。

我们可以看到，随机舍入为反向传递产生了最佳性能，而 QueST 为前向传递产生了最佳性能。QueST 是我们之前论文中的一种方法，它使用基于标准差的缩放和哈达玛变换进行归一化，以及梯度估计的无偏期望。我们使用的另一种无偏方法是 Albert 等人最初提出的方法，它也使用哈达玛变换和梯度估计产生的舍入。

整个算法看起来像这样：在前向传递中，我们进行 QueST 舍入以及一些值的裁剪。在反向传递中，我们进行额外的哈达玛变换以进行梯度归一化。我们以低精度执行矩阵乘法，但随后我们必须执行一些额外的操作，如掩码和哈达玛变换，以更新无偏且准确的梯度估计。

通过这个，你已经可以看到这个方法有多复杂。由于我之前展示的最优区域不仅取决于方法的精度，还取决于速度提升，因此我们论文的第四个也是 arguably 最重要的部分是我们提供的、真正使这个方法成为可能的内核。

---

实现高性能的 CUDA 内核

上一节我们介绍了 Quartet 算法的复杂性，本节中我们来看看实现其高效运行的关键——高性能 CUDA 内核。

正如 Andre 提到的，我们方法的最优区域不仅取决于精度，还取决于速度提升。因此，我们投入了大量时间并继续投入大量时间来优化 CUDA 内核，因为内核的真实速度提升将决定我们方法的最优区域大小。速度提升越好，这个最优区域就越大，落在该区域内的模型就越多。

在深入内核细节之前，我先简要概述一些信息。当我们提到 MxFP4 时，指的是我们在 Quartet 论文中使用的 FP4 类型。这意味着我们的值由这种 E2M1 格式表示：1 位用于符号，2 位用于指数，1 位用于尾数。对于矩阵中的每一列，每 32 个元素组将被分配一个缩放因子。对于 MxFP4 格式，这个缩放因子由一个 E8M0 的 FP8 值表示，即所有 8 位都用于表示指数值。

在新的 Blackwell 架构中，我们有硬件支持来执行这种操作，其中要相乘的 A 矩阵和 B 矩阵将由元数据信息（即这些缩放因子）额外表示。这个缩放因子应用于矩阵乘法的内部维度。

在硬件方面，这意味着我们将有两个矩阵，在 Tensor Core 中，我们可以用一个新指令执行这些矩阵乘法，该指令不仅接收输入矩阵，还接收此元数据信息以产生输出矩阵。

Andre 已经强调过矩阵乘法在大型语言模型中的影响。这些模型的大部分运行时间由线性层中的矩阵乘法表示，这意味着它们是必须尽可能优化的关键操作。这需要仔细调整，但也意味着我们必须避免额外操作（如我们刚刚讨论的哈达玛操作）的开销。

哈达玛变换作为一种归一化，其大小至关重要。当哈达玛大小小于 256 时，要解决的问题是内存瓶颈；当它大于或等于 256 时（这在很大程度上取决于所使用的架构），问题就变成计算瓶颈。因此，应用哈达玛变换所需的时间将由加载矩阵所需的时间决定，因为它是内存瓶颈问题。这意味着，如果我们非常高效地移动数据，我们的方法将会快得多。这也意味着哈达玛变换在某种程度上可以被矩阵加载隐藏，因为我们无论如何都必须加载这些矩阵以量化它们。

由于这是一个内存瓶颈问题，这意味着我们不一定必须应用哈达玛变换，如果我们愿意，实际上可以在这里应用任何类型的变换。我们使用一个融合的 CUDA 内核解决了哈达玛大小为 32 的这个 alpha 投影问题，我们发现它们确实有效，所以我们真的不需要使用更大的哈达玛大小来使我们的方法工作。

当我们提到这个密集哈达玛变换器时，实际上我们有一个针对特定组大小的这种块对角矩阵。如果对角线上的所有这些块基本相同，我们可以将之前的问题转换为这个。由于我们专注于 MxFP4，我们可以证明哈达玛大小为 32 是合理的。这意味着，我们可以将这个问题转换为一个普通的 GEMM 操作。

利用这种解决哈达玛变换的替代方法，我们定义了一个自定义的 CUTLASS GEMM 模板来以下列方式执行：内部维度固定为 32，所有剩余的元素将代表左侧输入矩阵的外部维度。然后，我们可以以下列方式定义这些 tile 形状：M 维度为 128，K 维度为 32（因为它永远不会大于 32）。以下配置仅取决于我们使用的特定硬件架构。

我们所做的是：我们有两个 Bf16 输入矩阵，一个是要量化的矩阵，另一个是哈达玛矩阵。然后，我们使用这个自定义的 CUTLASS GEMM 模板应用密集哈达玛变换，并且这些分离的累加将发生在 FP32 上。输出将存储在共享内存中。然后，我们定义了一个自定义的 epilogue 来产生以下三个结果：第一个是量化为 E2M1 格式的 FP4 值；第二个参数是我在第二张幻灯片中提到的具有 E8M0 格式的缩放因子；第三个参数是将在反向传播中使用的裁剪掩码（这是一个二进制掩码）。

这里一个非常重要的信息是，这个自定义的 epilogue 几乎是免费的，因为它发生在本地。我们不需要同步，因为每个线程将包含所有需要的信息。我们只需要在哈达玛变换后获取信息，计算几个元素，然后写回全局内存，不需要额外的同步。

这个先前的幻灯片描述了我们实现的融合量化内核。在这个矩阵乘法之前，我们必须添加一个额外的方法，即缩放因子重排。这与量化或我们的实现无关，这只是因为新的 Blackwell Tensor Core 要求元数据信息以特定方式存储，即我左侧显示的布局。这将应用于我们在上一步中计算的 E8M0 缩放因子。

总的来说，这里我展示了一个如何执行这三种方法的示例。这个示例代表了我们方法的前向路径。我们将接收输入和权重，并即时量化它们。这个前向量化方法应用哈达玛变换，然后计算三个输出：FP4 值、FP8 缩放因子和二进制掩码。然后，使用这些缩放因子，我们调用这个双块函数，它将信息重新排列为这个新的 Blackwell Tensor Core 所需的特定格式。然后，我们就可以调用这个 FP4 MxFP4 GEMM 内核了。

---

性能分析与未来工作

上一节我们详细介绍了内核实现，本节中我们来看看其性能表现以及未来的优化方向。

这里我有 Llama 70B 模型某些层的性能细分。Y 轴显示了这三种方法各自所占运行时间的百分比。实际上，我们可以看到，对于较大的矩阵，矩阵乘法占据了运行时间的重要部分，但对于较小的矩阵，量化相关操作的开销更大。这与我们的实现无关，只是随着我们增加算术强度，矩阵乘法将花费更长时间，我们可以轻松地隐藏这些

17：PCCL 容错集合通信库教程

在本节课中，我们将学习PCCL（Prime Collective Communications Library），这是一个支持在IP网络上进行容错、确定性集合通信的库。我们将了解其设计动机、核心概念、编程范式以及实现容错所面临的挑战。

概述：为什么需要PCCL？

PCCL的动机是在互联网上执行All-Reduce操作，并容忍所有可能的I/O故障。它提供了实现多种分布式优化方案的基元，从朴素的数据并行到Dilico流式或延迟应用的Eager Dilico。PCCL的一个主要优势是能保证所有节点上集合通信操作的比特级相同结果，这有助于促进共享状态（包括模型和优化器状态）的独立推进，达到完全相同的最终状态，从而消除节点间状态漂移的问题。

核心概念与编程范式

上一节我们介绍了PCCL的目标，本节中我们来看看它的核心特性和基本编程模型。

PCCL支持以下特性：

• 零拷贝：在发送调用前不复制数据。
• 比特级确定性：给定相同输入，保证相同的输出结果。
• 可中断：发生I/O故障时可中断操作并妥善回滚。
• 流水线化：实现教科书式的MPI流水线All-Reduce。
• 量化：支持以低于源数组的精度发送数据，内部以更高精度累加。
• 异步：可启动多个操作并在选择的时间等待它们。

这些特性使得共享状态可以从所有已执行的归约结果和初始状态完全重新推导出来。

初始化与通信器

每个PCCL应用程序首先需要初始化PCCL，然后创建一个指向专用主节点的通信器。这个轻量级主节点主要作为事实来源，存储连接信息并促进协议的状态转换，使得故障发生时的回滚更加明确和实用。

微共识方案与节点管理

PCCL内部采用一种微共识方案。许多阻塞的PCCL API调用会形成一个屏障，并协商下一步操作，例如接受新节点。

节点只能在现有节点集明确接受新节点时加入。这通常在每个训练步骤中周期性发生，通过调用 pcc_update_topology 实现。此时，世界大小可以增加。而在代码的其他部分，世界大小只会因非正常故障而减少，不会增加。

pcc_connect 会从连接方的角度阻塞，直到进入被接受的状态。

共享状态同步

共享状态引用所有权重和优化器状态，并定期同步。PCCL使用一个简单的、非加密的、为GPU优化的哈希函数（称为simple_hash）对所有共享状态内容进行哈希。

如果哈希一致，则同步是无操作；如果不一致，则会与拥有流行共享状态的随机节点配对以重新获取状态。理想情况下，应用程序应促进每个逻辑步骤共享状态的确定性推进，使得同步始终是无操作。这之所以合理，是因为所有节点都获得相同的归约结果，并且从相同的初始状态开始应用更新。

以下是共享状态同步后的一个最佳实践断言示例：

# 初始同步后，断言不应再接收任何共享状态字节
assert rx_bytes == 0, "No further shared state should be received after initial sync."

主循环与容错集合通信

上一节我们了解了状态管理，现在进入应用程序的主循环和具体的集合通信操作。

一个典型的PCCL应用程序主循环包含以下阶段：

1. 接受新节点：所有现有节点同意新节点加入。
2. 优化拓扑：进行速度测试，构建带宽和可达性矩阵，并形成最优的通信环。
3. 同步共享状态：理想情况下是无操作，主要作为防护机制并使新节点对齐状态。
4. 执行计算工作：例如模型的前向和反向传播。
5. 执行集合通信操作：基于计算结果进行归约等操作。
6. 推进共享状态：使用通信结果更新状态。

所有非异步的PCCL函数都是严格阻塞且基于共识的。所有节点必须在同一时间调用相同的方法，这形成了一个隐式屏障。一个重要规则是：应用程序绝不能分支，导致不同节点调用不同的PCCL函数，否则会导致死锁。

执行集合通信

由于需要容错，我们希望重试集合通信操作直到成功。操作可能因节点掉线等原因失败，此时PCCL会正确地回滚到用户的调用点并返回错误状态，用户可以决定重试或中止。

可以启用量化以节省带宽，例如使用 in8 模式。在等待一个集合通信操作后，世界大小可能已改变，因此应重新获取世界大小，避免脏的应用程序状态。

拓扑优化与容错含义

我们刚刚讨论了主循环，其中关键一步是拓扑优化。那么，为什么要进行拓扑优化呢？

拓扑优化的目标是基于节点间的实际网络速度（通过速度测试获得）构建一个高效的通信环，以最大化带宽利用率，避免数据包频繁通过慢速链路。一个糟糕的拓扑（例如，仅按节点加入顺序连接）会导致性能严重下降。

PCCL还会考虑部分节点对不可达的情况。如果节点A无法ping通节点B，则通信环将绕开此不可达路径（如果可能）。如果无法实现，则导致此情况的节点将被踢出。

实现高吞吐与真正容错

PCCL的一个标题是“在公共互联网上实现45 Gb/s的All-Reduce”。这通过每流公平队列技术实现。通过建立大量连接并在内部将数据分片分发到这些连接上，PCCL可以规避单个TCP流因长延迟或公平队列策略导致的低吞吐问题，从而实现高聚合带宽。

容错意味着能够容忍任何I/O故障，并且所有错误路径都必须被正确且优雅地处理。这包括在发现故障时，可能需要通过数据包通知主节点，由主节点重新配置拓扑并指导节点重建对等连接。设计协议时必须仔细处理跨连接数据包处理链的顺序问题，因为数据包到达顺序没有保证，这需要通过主节点协调的状态机来建立共识。

创建真正容错的系统在实践中非常困难。API设计必须谨慎，避免做出无法实现的承诺。例如，早期PCCL允许节点在任何时间加入，但这导致在集合操作进行中加入的节点无法获得有效的更新贡献，使得API无法实现。

确定性训练与GPU

之前我们提到PCCL保证比特级相同结果，这为实现共享状态的确定性推进奠定了基础。这对于许多仅由逐元素操作组成的优化器（如Adam）来说是微不足道的，因为这些操作不会引入GPU非确定性。

即使涉及矩阵乘法等操作，在许多情况下，通过使用确定的块大小、内核启动参数以及确定的算法（如确定性的Flash Attention前向/后向），也可以实现GPU上的确定性训练。虽然NVIDIA不正式保证这一点，但实际观察发现，从GTX到B200的多代GPU上，许多核心操作的数值行为是一致的。

理论上，这可以扩展到训练Transformer语言模型。PCCL只要求优化器步骤是确定的，数据加载中的非确定性则不影响其核心机制。

实现容错的挑战：套接字API的“痛苦”

构建真正容错的系统需要深入理解底层套接字API，而这些API在不同操作系统上存在微妙差异和陷阱。

套接字行为与文件描述符不同。例如，直接关闭套接字可能导致附近发送的数据永远无法送达，因为会触发连接重置。需要使用 SO_LINGER 选项进行更优雅的关闭，并可能设置超时。

中断阻塞在 recv 上的线程并非简单地关闭套接字文件描述符即可，在某些系统上可能还需要调用 shutdown 来解除阻塞。

在实践中，正确关闭套接字的步骤可能非常复杂，需要设置接收超时、优雅关闭连接的两端，然后才关闭套接字。

这只是套接字实现 quirks 的冰山一角，其他还包括：

• Nagle算法可能导致小数据包延迟发送。
• Keep-alive 不是心跳，不配置可能导致数小时才检测到对端死亡。
• send 可能产生 SIGPIPE 信号，如果不处理会杀死进程。
• close 不一定会解除 recv 的阻塞。
• 半关闭和排空行为在Windows上不同。
• SO_LINGER 没有超时可能导致关闭挂起。
• 监听积压队列与TCP快速打开结合可能导致半建立连接。

测试这些边缘情况非常困难，通常需要结合原理推理、大量测试以及在不同操作系统和环境下的实际部署调试。

总结

在本节课中，我们一起学习了PCCL容错集合通信库。我们从其设计动机和保证比特级确定性的核心目标开始，探讨了其编程范式，包括初始化、微共识、节点动态管理和共享状态同步。我们剖析了一个典型PCCL应用的主循环步骤，并深入了解了拓扑优化对于在复杂网络环境下实现高性能的重要性。

我们重点讨论了实现真正容错通信的挑战，特别是底层套接字API在不同平台上的微妙差异和陷阱，这些是构建健壮分布式系统时必须克服的困难。最后，我们看到了PCCL的确定性保证如何与GPU计算结合，为完全确定性的分布式训练打开了可能性。

PCCL代表了一种对可靠性和确定性有更高要求的系统设计思路，旨在为机器学习基础设施提供更坚固的基础。

18：ScaleML系列 - FlexOlmo：面向灵活性的开放语言模型

概述

在本节课中，我们将学习混合专家模型的基本原理及其在分布式数据训练场景下的创新应用。课程首先由组织者介绍混合专家模型的核心概念和GPT OSS中的实现细节，然后由主讲人Se Wun深入讲解其团队提出的FlexOlmo框架，该框架旨在解决不同数据所有者在不共享原始数据的情况下协作训练单一模型的问题。

---

混合专家模型基础

上一节我们概述了课程内容，本节中我们来看看混合专家模型的核心思想。

我们知道，更多的模型参数通常能带来更好的性能，但同时也意味着更高的计算成本。基础模型是真正的多任务学习者。对于一个特定的输入，我们真的需要激活所有参数吗？混合专家模型采取了一种不同的视角：我们可以为每个输入只使用模型参数的一个子集。

直观上，你可以将其理解为让不同的专家拥有不同的专长。在实践中，这些专家可能没有清晰的语义划分，但这种设计允许你在不显著增加计算量的情况下扩大模型规模。例如，在GPT OSS模型中，较大的版本约有1200亿参数，但每次推理只激活50亿参数。

混合专家模型的工作原理

以下是混合专家模型在GPT OSS中的关键实现步骤，我们将通过代码片段来理解其前向传播过程。

# 假设输入是一个token，其隐藏维度为H
x = token_input  # 形状: [H]

![](https://github.com/OpenDocCN/dsai-notes-pt1-zh/raw/master/docs/gpu-mode-cuda/img/7c0d341d6c688b90f85860afb32ff55b_9.png)

![](https://github.com/OpenDocCN/dsai-notes-pt1-zh/raw/master/docs/gpu-mode-cuda/img/7c0d341d6c688b90f85860afb32ff55b_11.png)

# 1. 层归一化
x_norm = rms_norm(x)

# 2. 路由计算：将token投影到专家数量的空间，得到路由分数
router_logits = x_norm @ W_gate  # W_gate形状: [H, num_experts]
# 应用softmax得到每个专家的权重
router_probs = softmax(router_logits, dim=-1)

# 3. 选择Top-K专家
top_k_weights, top_k_indices = topk(router_probs, k=K)  # 例如K=2或4

# 4. 归一化所选专家的权重
top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)

接下来，我们需要将token发送给选定的专家进行处理。专家本质上是替换了原始Transformer中的MLP层。在GPT OSS中，每个专家是一个独立的MLP。

# 5. 初始化最终输出
final_output = torch.zeros_like(x)

# 6. 对每个选定的专家进行处理
for i, (expert_idx, weight) in enumerate(zip(top_k_indices, top_k_weights)):
    # 获取对应的专家MLP
    expert_mlp = experts[expert_idx]
    # token通过专家MLP
    expert_output = expert_mlp(x_norm)
    # 加权累加到最终输出
    final_output += weight * expert_output

![](https://github.com/OpenDocCN/dsai-notes-pt1-zh/raw/master/docs/gpu-mode-cuda/img/7c0d341d6c688b90f85860afb32ff55b_17.png)

# 7. 添加残差连接
output = x + final_output

张量并行与通信开销

在大型模型中，专家MLP的参数通常使用张量并行进行切分。这意味着每个GPU只持有每个专家MLP参数矩阵的一部分。

• 计算过程：每个GPU使用其持有的参数分片与输入token进行计算，得到一个部分结果。
• 通信过程：所有GPU通过all-reduce操作通信并求和这些部分结果，最终每个GPU都得到完整的专家输出。

张量并行旨在降低单次推理的延迟，但其代价是增加了设备间的通信开销。在网络带宽受限的情况下，张量并行可能不是最优选择。

训练稳定性与负载均衡

训练混合专家模型的一个主要挑战是路由器坍缩，即路由器可能倾向于将所有token都路由到少数几个专家，这违背了利用专家多样性的初衷。

为了解决这个问题，训练过程中会引入辅助损失函数（如负载均衡损失）。该损失函数鼓励token更均匀地分布在不同专家之间，确保所有专家都能得到训练。

---

FlexOlmo：面向分布式数据的混合专家模型

上一节我们介绍了标准的混合专家模型，本节中我们来看看如何利用其模块化特性解决分布式数据训练的难题。

问题定义

标准的语言模型训练假设开发者能集中访问所有训练数据。然而在现实中，数据访问权限往往是一个谱系：

• 数据可能因隐私、商业价值或法规限制而无法集中。
• 数据可能在不同时间点可用。
• 我们希望在无需从头重新训练整个模型的情况下，能够灵活地添加或移除数据源。

FlexOlmo旨在支持模块化的分布式训练，其设置如下：

1. 一个在公共数据上训练的共享基础模型。
2. 多个数据所有者，各自在本地、独立地使用自己的私有数据继续训练这个基础模型。
3. 最后，将这些独立训练的模型合并成一个统一的混合专家模型。

这种方法有两个关键优势：

• 数据贡献无需共享：数据所有者保留对数据的完全控制权。
• 支持灵活的数据增删：可以轻松添加新数据源或移除现有数据源，而无需重新训练整个模型。

现有方法的局限性

在引入FlexOlmo的解决方案前，我们先看看两种可能的替代方案及其局限。

1. 模型融合

• 思路：每个数据所有者独立训练自己的模型，然后通过权重平均或集成的方法合并。
• 局限：当各模型在不同分布的数据上训练后，其参数空间可能差异巨大，直接合并会导致性能显著下降。

2. 混合专家模型融合

• 思路：将每个数据所有者训练得到的FFN作为独立的专家，合并成一个MoE模型。
• 局限：合并后，路由器需要重新训练以学会如何选择专家，而这通常需要访问所有数据的并集进行联合训练，违反了数据隐私的前提。

FlexOlmo的核心方法

FlexOlmo的核心目标是保留混合专家架构的模块化优点，同时消除对数据联合访问的需求。它通过两个关键设计来实现：

1. 混合专家感知训练
每个数据所有者拿到共享基础模型后，不进行普通的继续预训练，而是先将其转换为一个仅有两个专家的微型MoE模型。

• 专家一：冻结的、来自公共数据的原始FFN。
• 专家二：新初始化的、可训练的FFN。
  数据所有者只训练第二个专家以及对应的路由器部分。这样，新专家在训练时就知道需要与公共专家进行“协作”。

2. 非参数化路由器
标准MoE的路由器是一个 [隐藏维度 H] x [专家数 N] 的矩阵。FlexOlmo将其分解为N个H维的路由器嵌入向量，每个向量对应一个专家。

• 公共数据对应的路由器嵌入向量被冻结。
• 每个数据所有者只训练自己新增专家对应的路由器嵌入向量。
• 最终合并时，只需将所有路由器嵌入向量拼接起来，就形成了完整的路由器矩阵。

工作流程总结

1. 在公共数据上训练一个密集模型作为基础。
2. 每个数据所有者以此为基础，进行“两专家MoE”模式的本地训练（冻结公共部分，训练新增的专家和路由器嵌入）。
3. 所有训练完成后，将各数据所有者新增的专家FFN和路由器嵌入合并，与冻结的公共部分一起，构成最终的、包含多个专家的MoE模型。

实验结果与优势

在7B参数规模的实验表明，FlexOlmo相比仅在公共数据上训练的模型有显著提升（在MMLU Pro和HELM基准上平均提升41%）。

• 它显著优于传统的模型融合方法。
• 在专业领域任务上，它能达到接近甚至超过在该领域数据上专门训练的“专家模型”的性能。
• 更重要的是，FlexOlmo的性能达到了无限制联合训练MoE模型（数据可集中访问的理想上限）的90%，同时保留了数据隐私和灵活增删的巨大优势。

数据剔除分析
实验验证了“ opting-out”的有效性。当从模型中移除某个数据源对应的专家后：

• 与该数据源相关的任务性能会下降。
• 其他不相关的下游任务性能受影响极小，甚至可能因路由器选择范围变小而略有提升。

局限与未来方向

1. 专家粒度：当前每个数据源对应一个完整的FFN专家，模型参数量增长快。未来可探索更细粒度的MoE，让更多小数据持有者参与。
2. 路由器优化：当前路由器仅在“二选一”场景下训练，泛化到“多选一”时可能不是最优。如何训练出泛化能力更强的路由器是一个技术挑战。
3. 持续学习：如何利用此架构实现模型的持续渐进式改进。
4. 实际部署问题：包括如何评估不同数据贡献的价值、如何设计激励机制，以及如何防止恶意数据源污染模型等。

---

总结

本节课我们一起学习了混合专家模型从基础到前沿应用的完整脉络。我们首先剖析了MoE通过激活参数子集来提升效率的核心机制及其在GPT OSS中的实现。接着，我们深入探讨了Se Wun团队提出的FlexOlmo框架，它巧妙利用了MoE的模块化特性，使得在分布式、隐私敏感的数据上进行协作式模型训练成为可能，同时保持了模型性能并支持灵活的数据管理。这项工作为大规模、合规的AI模型开发开辟了新的可行路径。

19：高效流式语言模型与注意力汇聚

在本节课中，我们将学习一种名为“注意力汇聚”的现象，以及如何利用它来构建能够处理超长文本序列的流式语言模型。我们将从模型部署的挑战开始，逐步理解注意力汇聚的成因，并最终掌握“流式语言模型”这一无需训练即可扩展模型上下文窗口的实用技术。

概述：流式应用中的挑战

上一节我们介绍了GPU编程的基础。本节中，我们来看看将大型语言模型应用于流式对话等长交互场景时面临的核心问题。

当我们将预训练好的语言模型（如Llama 2）直接用于长对话时，会遇到两个主要挑战：

1. 内存消耗线性增长：在自回归解码过程中，需要缓存所有历史token的键值状态（KV Cache）。随着对话轮数增加，KV Cache的内存占用会线性增长，最终导致GPU内存耗尽。
2. 性能在预训练长度外崩溃：模型通常在固定长度（如4K）的文本块上预训练。当输入序列长度超过这个预训练限制时，模型的输出质量（困惑度）会急剧下降，输出变得不可信。

下图展示了这两个挑战：

KV Cache：效率与瓶颈

为了解决解码时的计算效率问题，我们首先回顾一个关键技术：KV Cache。

在自回归语言模型中，当前token的生成只依赖于之前的token。因此，我们可以缓存之前所有token计算好的键（K）和值（V），在解码新token时直接复用，无需重新计算。这被称为KV Cache优化。

公式表示：
对于第 t 个token，其注意力计算为：
Attention(Q_t, K_{1:t}, V_{1:t})
其中 K_{1:t} 和 V_{1:t} 是缓存的前 t 个token的键和值。

没有KV Cache时，每一步的计算复杂度为 O(t^2)，总复杂度为 O(N^3)。使用KV Cache后，每一步只需计算当前token的Q与历史K的点积，复杂度为 O(t)，总复杂度降至 O(N^2)，是巨大的效率提升。

然而，KV Cache本身成为了长上下文的主要内存瓶颈。对于一个典型设置（如Llama 2 7B，批次大小=4），缓存32K token的KV Cache可能需要高达64GB的显存。

朴素解决方案：窗口注意力及其失败

一个很自然的想法是：既然缓存所有token开销太大，我们只缓存最近的L个token，丢弃更早的，这称为“窗口注意力”。

这种方法将内存和解码时间复杂度都降为了常数 O(L)。但是，当我们用自然文本测试时，发现一旦文本长度超过缓存大小L，模型性能就会突然崩溃。

如下图所示，当序列长度超过缓存大小时（橙色虚线），困惑度（橙色实线）出现尖峰，模型失效。

这引出了一个关键观察：模型性能的崩溃，恰好发生在最早的那几个token被移出缓存窗口的时刻。这表明，初始的某些token对模型稳定运行至关重要。

核心发现：注意力汇聚现象

为了探究窗口注意力失败的原因，我们可视化了模型内部的注意力分布图。

观察：在Llama 2等模型的多数层中（尤其是两层之后），所有后续token都会强烈地关注第一个token，无论它们在语义上是否相关。如下图所示，第一列呈现出一条明显的红色竖线。

我们将这种现象命名为 “注意力汇聚” 。初始的token像“汇点”一样，吸收了模型中大量“多余”的注意力分数。

成因分析：其根源在于Softmax函数的特性。Softmax要求所有上下文token的注意力概率之和必须为1。即使对于某些简单的、无需关注太多上下文的token，模型也必须将“多余”的概率质量分配到某个地方。在因果语言模型中，第一个token是全局可见的，因此成为最自然的“注意力汇点”。

为了验证初始token的重要性是源于其位置而非语义，我们进行了消融实验：

以下是实验结果对比：

• 保留前0个token + 最近1024个token：困惑度很高（性能差）。
• 保留前4个初始token + 最近1020个token：困惑度良好，接近预训练水平。
• 保留4个无意义的换行符token + 最近1020个token：困惑度同样恢复良好。

结论是：模型依赖初始token，主要是因为它们位于序列开头，而非其语义内容。

解决方案：流式语言模型

认识到注意力汇聚现象后，解决方案变得直观：在窗口注意力中，永远不驱逐作为“注意力汇点”的初始token。

我们提出了 流式语言模型 方法：

1. 在KV Cache中，始终保留开头的几个token（如4个）作为“注意力汇点”。
2. 同时，像标准窗口注意力一样，保留最近的L个token。
3. 中间的其他token则被驱逐。

这种方法保持了 O(L) 的常数内存和计算复杂度，同时因为保留了注意力汇点，模型困惑度保持稳定，能够处理远超预训练长度的序列。

关键技术细节：位置编码重映射
模型在预训练时从未见过超过其训练长度（如8）的位置索引。在流式解码第9个token时，如果使用原始位置索引（0,1,2,3,6,7,8,9），会让模型看到未知索引（8,9）。因此，我们将其重映射为缓存内的相对位置（0,1,2,3,4,5,6,7），让模型始终认为自己处于一个熟悉的上下文窗口中。这可以通过在应用位置编码前缓存K、V，解码后再重新应用位置编码来实现。

实验结果与影响

我们将StreamingLLM与多种基线方法比较：

• 密集注意力：内存占用大，长度超过预训练限制后性能崩溃。
• 窗口注意力：效率高，但驱逐初始token后性能崩溃。
• 滑动窗口重计算：质量好（与训练一致），但每一步都需要 O(L^2) 的重计算，效率极低。

结果：StreamingLLM（红线）在保持常数效率的同时，其困惑度与计算代价高昂的“滑动窗口重计算”基线（绿线）几乎完全重合，实现了高质量的长序列建模。

我们将实验扩展到多个模型家族（Llama 2, Pythia, Falcon, MPT）和不同规模（7B, 13B, 70B）。StreamingLLM均能支持这些模型处理长达400万token的文本，远超其原始预训练长度。

与后续工作的联系：Grok-1与可学习的注意力汇点

我们的工作发表后，“注意力汇聚”的概念引起了广泛关注。例如，xAI发布的Grok-1模型在其技术报告中提到使用了“注意力汇点”。

Grok-1的设计：它为每个注意力头引入一个可学习的标量参数作为注意力汇点。在计算注意力时，将这个标量作为一列添加到注意力分数中参与Softmax计算，计算完概率后再移除该列。这使得模型可以灵活地控制是否需要向“汇点”分配注意力。

代码示意：

# 假设 attention_sink 是一个可学习的标量
attention_scores = torch.matmul(query, key.transpose(-2, -1))
attention_scores = torch.cat([attention_scores, attention_sink.expand_as(attention_scores[..., :1])], dim=-1)
attention_weights = F.softmax(attention_scores, dim=-1)
attention_weights = attention_weights[..., :-1]  # 移除汇点列
context = torch.matmul(attention_weights, value)

这与我们论文中探讨的“零汇点”或“专用汇点token”思路相似，但在工程实现上更简洁，无需修改数据管道或使用特殊的注意力核函数。

注意力汇聚的普遍性

注意力汇聚并非语言模型独有。由于根本原因在于Softmax的归一化特性，它在其他使用Softmax注意力机制的Transformer模型中也存在：

• 视觉Transformer：在ViT中，某些低语义的背景图像块会吸收大量注意力，被称为“寄存器”。
• 双向语言模型：在BERT中，[SEP]等分隔符token也扮演了类似的注意力汇点角色。

局限性与未来方向

需要明确的是，StreamingLLM并不提供真正的“无限记忆”。它通过固定大小的窗口保持模型运行的稳定性，但被移出窗口的历史信息仍然会被遗忘。模型只能可靠地访问保存在缓存中的局部上下文。

后续研究可以围绕更智能的缓存管理策略展开，例如基于重要性评分选择性地保留或压缩历史信息，从而在固定预算下实现更长的有效记忆。

总结

本节课中我们一起学习了：

1. 注意力汇聚现象：由于Softmax的归一化要求，语言模型会强烈且持续地关注初始token，无论其语义如何。
2. 流式语言模型：通过永久保留开头的几个token作为注意力汇点，并结合滑动窗口注意力，可以在不进行任何额外训练的情况下，使现有语言模型稳定处理远超其预训练长度的序列。
3. 方法优势：该方法保持了常数级的内存和计算开销，同时维持了模型输出质量，为长对话、文档处理等流式应用提供了可行的部署方案。
4. 广泛影响：注意力汇聚的概念已被后续研究和工业界模型（如Grok-1）所采纳和发展，成为改进Transformer架构长上下文能力的重要思路。

这项研究表明，深入理解模型内部的微观机制，往往能催生出简单而有效的宏观解决方案。

20：大模型中的量化技术

概述

在本节课中，我们将要学习大模型中的量化技术。量化是一种将模型权重或激活值从高精度格式（如FP16）转换为低精度格式（如INT4）的技术，其核心目的是减少模型存储空间、提升推理速度并降低能耗。我们将重点探讨训练后量化，并简要介绍训练中量化。

什么是训练后量化？

训练后量化是指在一个已经训练好的高精度模型基础上，将其转换为低精度表示的过程。这与量化感知训练不同，后者是在训练过程中就使用量化模型。

训练后量化的主要目标有三个：

1. 减小模型体积。
2. 提升推理速度。
3. 压缩过程本身不应消耗过多时间和计算资源。

我们主要讨论权重压缩，即仅对模型的参数进行量化，而不涉及前向传播过程中的中间激活值。

朴素的量化方法及其局限性

最朴素的量化方法是直接将权重舍入到最近的、低精度格式可表示的值。

代码示例：

# 伪代码：朴素舍入量化
quantized_weight = round_to_nearest(weight, target_bits)

然而，当比特数较低时，这种方法会引入大量量化噪声（舍入误差），最终可能严重损害模型性能。因此，我们的目标不仅是压缩模型，还要确保量化后的模型在预测输出上尽可能接近原始模型。

量化问题的数学建模

我们希望最小化量化模型与原始模型在某个数据集上的预测差异。这可以形式化为一个优化问题：寻找一个量化后的权重向量 W_hat，使得期望损失最小。

由于直接求解此问题非常困难，我们采用二阶近似。假设量化后的权重 W_hat 接近原始权重 W，我们可以用海森矩阵 H 来近似损失函数。

公式：

损失 ≈ (W_hat - W)^T * H * (W_hat - W)

其中，H 是损失函数在 W 处的海森矩阵。然而，计算完整的海森矩阵计算量巨大（与参数数量的平方成正比）。

实用的海森矩阵近似与层级量化

为了解决海森矩阵计算难的问题，一个常见且实用的近似方法是在单个线性层级别进行操作。具体而言，我们最小化该层输出在某个校准数据集上的重构误差。

公式：

损失 ≈ E[ ||W*x - W_hat*x||^2 ]
        = trace( (W - W_hat) * H_layer * (W - W_hat)^T )

这里的 H_layer 是输入激活值 x 的协方差矩阵的期望值，这是一个易于计算的量。这种方法被诸如GPTQ等知名量化算法所采用。

自适应舍入算法与理论分析

基于上述近似，量化问题转化为最小化一个二次型。通过将海森矩阵 H 进行LDL分解（H = L * D * L^T，其中 L 是单位下三角矩阵，D 是对角矩阵），我们可以推导出一种带线性反馈的自适应舍入算法。

该算法的核心思想是：按顺序量化权重，并将当前权重的量化误差作为反馈，用于调整下一个权重的量化决策。这确保了量化误差在数学上是可控的。

理论边界：
分析表明，这种自适应舍入方法的误差上界与矩阵 D 的迹有关。当 D 的迹相对于 H 的迹较小时，自适应舍入相比朴素舍入能带来显著提升。

不连贯处理：提升量化效果的技巧

理论分析指出，当海森矩阵 H 接近对角矩阵时，自适应舍入的收益有限。为了改善这一情况，我们可以在量化前对模型进行随机旋转（例如使用随机哈达玛变换），使 H 在变换后的基下变得“不连贯”（即远离对角矩阵）。

这种方法被称为“不连贯处理”。它能有效改善权重的分布，使其更接近高斯分布，从而更适用于后续的量化编码理论，并能缓解激活值中的异常值问题。

量化算法的演进：QUIP、QuIP# 与 QuIP-T

基于上述理论，发展出了一系列量化算法：

• QUIP：首次引入了不连贯处理，成功将LLaMA-2 70B等模型量化到2比特，性能损失极小。
• QuIP#：在QUIP基础上引入了向量量化，将多个权重组合成向量进行联合量化，进一步提升了压缩率。
• QuIP-T：采用了网格编码这一源自信息论的技术。它利用伪随机数生成器作为解码器，并通过动态编程精心选择“种子”比特，使得生成的序列尽可能接近目标权重。不连贯处理确保了权重近似高斯分布，使得网格编码得以高效应用。

超越训练后量化：在训练中使用低精度

不连贯处理的思想同样适用于训练阶段。例如，有工作探索了使用MX FP4等低精度格式训练大语言模型。随机变换有助于改善训练过程中模型的数值条件，使得低精度训练更加稳定和有效。

总结

本节课我们一起学习了：

1. 训练后量化的目标与挑战：在减小模型、加速推理的同时，需保持模型精度，且压缩成本要低。
2. 量化问题的数学框架：通过最小化输出重构误差来建模，并利用海森矩阵的二阶近似。
3. 自适应舍入算法：通过LDL分解和线性反馈，实现了误差可控的逐权重量化。
4. 不连贯处理的关键作用：通过随机旋转改变基，为高效量化创造了有利条件。
5. 算法演进：从QUIP到QuIP#再到QuIP-T，通过结合向量量化、网格编码等理论，不断逼近信息论极限。
6. 技术的泛化：相关思想（如不连贯处理）也可用于提升低精度训练的效果。

量化是一个平衡模型大小、速度和精度的强大工具。随着算法和硬件支持的不断进步，低精度模型将在边缘部署和高效推理中扮演越来越重要的角色。

21：Transformer中的位置编码与PaTH注意力

在本节课中，我们将要学习Transformer模型中的位置编码机制，特别是旋转位置编码（RoPE）及其改进方案，并深入探讨一种名为PaTH的新型注意力机制，它旨在提升模型在状态跟踪等复杂任务上的表达能力。

位置编码的动机

上一节我们介绍了课程背景，本节中我们来看看位置编码的必要性。Transformer模型最初被设计用于双向建模，其自注意力机制在没有因果掩码的情况下，会将输入序列视为一个无序的词袋。如果没有位置嵌入，模型将无法感知词语的顺序。

以下是一个直观的例子：

• “猫坐在垫子上。”
• “垫子坐在猫上。”

如果没有位置编码，模型无法区分这两个句子。在原始的Transformer论文中，作者使用了一个简单的绝对位置编码，即正弦和余弦函数的组合。然而，这种绝对位置编码存在局限性，例如，内容与绝对位置高度耦合，且无法直接建模词与词之间的相对位置关系。因此，后续研究主要聚焦于相对位置编码。

旋转位置编码（RoPE）概述

上一节我们提到了相对位置编码的重要性，本节中我们来看看目前最流行的方案——旋转位置编码。

RoPE的工作机制如下：

1. 对于输入嵌入向量，RoPE将其通道分成多个二维通道对。
2. 对于每个通道对，设定一个旋转频率角 θ。
3. 根据词的绝对位置索引 m 和频率角 θ，对该通道对进行旋转操作。
4. 旋转后的查询向量和键向量被用于计算注意力分数。

不同通道对独立进行旋转，互不干扰。尽管这里使用了绝对位置索引，但最终计算出的注意力分数仅依赖于查询和键之间的相对位置差，因此RoPE本质上是一种相对位置编码。

旋转矩阵的性质

在深入理解RoPE之前，我们先简要回顾旋转矩阵的性质。一个二维旋转矩阵 R(θ) 可以将一个向量旋转角度 θ。

旋转矩阵具有以下良好性质：

• 幂等性：R(θ)^m = R(mθ)。累积旋转可以通过直接计算总角度来实现。
• 正交性：旋转矩阵是正交矩阵，其逆矩阵等于其转置矩阵，即 R(θ)^{-1} = R(θ)^T = R(-θ)。
• 可组合性：多个旋转的组合等价于角度相加后的单个旋转，即 R(θ1) * R(θ2) = R(θ1 + θ2)。

RoPE利用这些性质，通过块对角矩阵的形式，将旋转操作扩展到高维空间。

RoPE中不同通道的作用

RoPE为不同的通道对设置了不同的旋转频率。根据频率大小，可以将通道分为两类：

• 高频通道：旋转角 θ 较大，旋转迅速。这些通道主要用于编码位置模式，例如识别最近的词。
• 低频通道：旋转角 θ 较小，旋转缓慢。这些通道主要用于编码语义信息，因为缓慢的旋转对点积结果影响较小。

这种分工使得RoPE能够有效地平衡位置信息和语义信息。

RoPE的局限性与改进方案

尽管RoPE非常有效，但它也存在局限性，最突出的问题是长度外推能力不足。例如，一个在4K序列长度上训练的模型，在评估超过4K的序列时，其困惑度可能会急剧上升。

位置插值与NTK感知RoPE

为了解决外推问题，研究者提出了多种方案。最初的位置插值方法简单地对所有通道的旋转角进行均匀缩放，但这忽略了不同频率通道的敏感性。

随后提出的NTK感知RoPE方法则更加精细。它基于神经正切核理论，该理论指出神经网络难以学习低维输入中的高频信息。在位置编码中，位置ID是单一维度的，因此高频通道的信息更难学习，也更为脆弱。

NTK感知RoPE的核心思想是：

• 对于高频通道，尽量保持其旋转角不变，以保留已学习到的精细位置模式。
• 对于低频通道，可以进行插值，使其适应更长的序列长度。

该方法使用一个平滑函数来计算每个通道的缩放因子。

YaRN：结合NTK与温度缩放

YaRN 是另一个重要的改进方案，它结合了NTK-aware插值和温度缩放。温度参数 T 根据新序列长度与训练序列长度的比例进行计算。YaRN的直觉可以通过注意力分数熵的变化来理解，它旨在更稳定地扩展上下文窗口。目前许多开源模型（如DeepSeek、Llama等）都采用了YaRN。

引入PaTH：提升Transformer的表达能力

上一节我们讨论了RoPE的改进，本节中我们来看看一个旨在突破Transformer表达能力限制的新工作——PaTH。

研究表明，像Transformer这样的模型，其计算复杂度属于 TC0 类。这意味着它们可以处理加法、比较等基本算术运算，但难以处理更复杂的任务，如状态跟踪、排列组合等，这些任务属于 NC1 复杂度类。这也是为什么当前大模型需要依赖长链思维（CoT）来进行复杂推理的原因。

一个NC1完全问题：排列组合

为了提升表达能力，PaTH关注一个简单的NC1完全问题：五元素排列组合。问题描述如下：给定五个元素（A, B, C, D, E）的初始排列，以及一系列交换操作（例如“交换位置1和2”），需要计算出经过所有交换操作后的最终排列。这个问题要求模型在内部状态中跟踪所有元素的位置，是典型的状态跟踪任务。

为什么RoPE难以处理此任务？

RoPE难以处理此类任务的原因有两个：

1. 数据无关性：RoPE的旋转操作仅依赖于绝对位置，与输入元素的内容无关。
2. 可交换性：二维旋转操作是可交换的，即 R(θ1)R(θ2) = R(θ2)R(θ1)。然而，交换操作的顺序是不可交换的，改变操作顺序会得到不同的结果。

PaTH的核心：Householder变换

PaTH使用 Householder变换 来编码交换操作。Householder矩阵是一种初等反射矩阵，可以将一个向量关于某个超平面进行反射。通过精心选择反射平面，可以用一个Householder变换交换两个向量的位置，同时保持其他正交向量不变。

更一般地，PaTH使用 广义Householder变换，其公式为：
H = I - β * w * w^T
其中 β 是一个可学习的标量参数：

• 当 β = 0 时，H 是单位矩阵，表示不进行任何操作。
• 当 β = 1 时，H 是投影矩阵。
• 当 β = 2 时，H 是标准的反射矩阵。

通过累积多个Householder变换的乘积，PaTH可以建模一系列交换操作的组合效果。PaTH将这种累积乘积作为位置编码的一部分，集成到注意力分数的计算中，形成了一个双线性权重矩阵。理论证明，这使得PaTH具备了NC1完全的表达能力。

PaTH的实验结果

在合成的五元素交换任务上，RoPE模型无法学会跟踪状态，而PaTH模型可以轻松解决。在其他状态跟踪任务（如Flip Flop语言建模）上，PaTH也显著优于其他位置编码方案，并且在处理长序列时，PaTH所需的网络层数仅以对数规模增长，效率更高。

实践中的PaTH：蒸馏与持续预训练

为了将PaTH应用于现有模型，研究者提出了蒸馏流程：

1. 阶段一（层对齐）：使用教师模型（如RoPE）的中间层输出作为目标，让学生模型（PaTH）的每一层去匹配，最小化均方误差损失。
2. 阶段二（知识蒸馏）：使用教师模型的最终输出分布作为目标，让学生模型去匹配，最小化KL散度损失。

初步实验表明，使用少量数据对Cornell-2 7B模型进行蒸馏，可以恢复其大部分性能，甚至在数学和代码任务上有所提升。

在持续预训练实验中，从一个预训练好的检查点开始，使用高质量数据继续训练，并将RoPE替换为PaTH。结果显示，在相同的训练数据和步数下，PaTH在GSM8K、HumanEval等数学和代码基准上明显优于RoPE，这表明其更强的状态跟踪能力有益于这些任务。

PaTH的高效实现

PaTH的设计考虑了硬件效率。其核心计算可以分解为块操作，并利用矩阵乘法和求逆（在小的块内进行）等GPU友好操作。研究者借鉴了Flash Attention的算法思想，设计了块状的PaTH注意力前向传播算法，使其能够高效地在现代GPU上运行。

对于推理阶段，PaTH引入了一种动态更新键值缓存（KV Cache）的机制。每个新的时间步都会产生一个Householder变换，这个变换被应用于所有历史的键缓存，相当于对历史信息进行了一次“精炼”。这可以通过秩1更新高效实现，最终可以调用优化过的注意力解码内核（如Flash Decoding）来完成计算。

PaTH与遗忘机制的结合：PaTH-FoX

上一节我们介绍了PaTH，本节中我们来看看其与遗忘机制的融合。遗忘Transformer 可以看作是一种数据依赖的偏置位置编码（类似于ALiBi），它通过一个累积的遗忘门来控制历史信息的保留程度。

PaTH-FoX将乘性的PaTH位置编码与加性的遗忘门偏置结合起来，公式如下：
注意力分数 = Q * (累积Householder变换) * K^T + (累积遗忘门偏置)

这种结合带来了互补的优势：PaTH提供了强大的状态跟踪和表达能
力，而遗忘机制有助于模型进行长度外推和选择性记忆。

实验表明，PaTH和PaTH-FoX在常识推理、语言建模等标准基准上表现良好，并且在长上下文和状态跟踪专项任务上具有显著优势。

总结

本节课中我们一起深入探讨了Transformer中的位置编码。我们从RoPE的基本原理出发，分析了其通过旋转编码相对位置的巧妙设计，以及不同频率通道的分工。接着，我们探讨了RoPE在长度外推上的局限性，并介绍了位置插值、NTK感知RoPE和YaRN等改进方案。

然后，我们引入了一个旨在突破Transformer表达能力瓶颈的新工作——PaTH。PaTH利用Householder变换的累积乘积来编码复杂的交换操作，赋予了模型NC1完全的表达能力，使其能够更好地处理状态跟踪等复杂任务。我们还讨论了PaTH的实践部署策略和高效实现方案，以及其与遗忘机制结合的PaTH-FoX变体。

这些研究展示了在保持注意力机制核心优势的同时，通过改进位置编码来显著提升模型理论表达能力和实际任务性能的潜力。

22：ScaleML系列 - GPU编程基础 + ThunderKittens

概述

在本节课中，我们将学习GPU编程的核心基础概念，并了解ThunderKittens这一旨在简化高性能AI内核开发的编程框架。我们将从GPU的硬件架构出发，理解其执行模型，然后探讨如何通过高级抽象来高效地编写内核。

---

GPU编程基础：理解硬件执行模型

上一节我们介绍了课程概述，本节中我们来看看GPU编程的基础——理解硬件如何执行我们的代码。

CUDA程序的软件视角

一个CUDA程序定义了一个类C语言的函数，然后通过启动命令在GPU上并行运行该函数的多个副本。启动时需要指定两个参数：线程块的数量和每个线程块中的线程数。从软件角度看，模型如下图所示：你有一个线程块数组，每个块内部又有一个线程数组，每个线程都是一个独立运行的小程序。

两个性能谜题

在实践中，这个模型会带来一些令人困惑的性能现象。

以下是两个常见的性能谜题及其解释：

1. 线程数减少，性能不变：假设一个内核使用128个块，每块1024个线程运行良好。如果我们将每块线程数从1024减少到256，直觉上可能认为硬件利用率会降至25%，性能下降。但许多内核的实际速度保持不变。这引出了疑问：如果减少线程数并让每个线程做更多串行工作能达到相同速度，那另外四分之三的线程原本在做什么？
2. 块数微增，性能骤降：同样从128个块开始，如果仅将块数增加到133（约4%的增长），你可能会认为性能影响不大。但在某些内核（尤其是在H100 GPU上）上运行时，性能可能下降2倍。为什么块数仅增加4%会导致如此大的性能变化？

这些谜题表明，仅从线程和块的软件模型来思考和预测程序性能是不完整的。对于初学者，一个更好的问题是：什么是GPU？它能做什么？

GPU的物理架构：一堆小方块

下图展示了一个NVIDIA GPU芯片。最显眼的视觉特征是它由许多小方块组成。

NVIDIA称这些小方块为流式多处理器，通常缩写为SM。一块现代GPU（如H100）上大约有132个SM。每个SM内部包含四个更小的计算核心，这些核心在不同文档中有不同名称（如象限、分区、线程束调度器）。我们可以将一个线程束调度器近似理解为一个CPU核心，但它高度专用于运行SIMD（单指令多数据）指令。

核心矛盾：线程数远多于硬件单元

这里出现了一个核心矛盾：一个CUDA内核可能启动数十万个线程（例如128块 * 1024线程/块 = 131,072线程），但GPU物理上只有几百个线程束调度器（H100上约132 SM * 4调度器/SM ≈ 528个）。如何用少得多的硬件单元运行如此多的线程？

答案有三点，它们共同解决了这个矛盾：

1. SIMD执行：线程束调度器每次发出指令时，并不是只为一个线程服务。硬件将32个连续线程的数据分组到一个大向量中。每次发出的一个指令会对这个向量的32个“通道”执行操作，相当于同时完成32个线程的工作。这32个线程的组被称为线程束。
2. 时间复用：一个线程束调度器可以同时跟踪多于32个线程的状态（程序计数器、寄存器值）。它可以在不同周期为不同的线程束发出指令，从而实现多个线程束在同一个调度器上的分时复用。
3. 串行调度块：如果你请求启动成千上万个线程块，GPU不会崩溃，而是会在各个SM上串行地一个接一个运行这些块。

深入理解SIMD与时间复用

上一节我们提到了SIMD和时间复用的概念，本节我们来深入探讨它们对编程和性能的影响。

SIMD执行与线程发散

当编写CUDA代码时，你表面上是在描述单个线程的行为。但硬件实际上是将32个连续线程分组，锁步执行相同的指令流。这意味着你需要在大脑中做一个视角转换：将代码视为每32个线程一组同步执行。

如果线程束内的32个线程执行路径不同（例如，由于if-else分支），就会发生线程发散。硬件会通过掩码确保每次只发出一个指令，该指令仅对参与当前路径的线程子集生效。例如，偶数线程做加法，奇数线程做减法，那么硬件需要先发出一条加法指令（对偶数线程生效），再发出一条减法指令（对奇数线程生效）。如果32个线程做32件不同的事，就需要32个周期来发出指令，吞吐量会急剧下降。

公式：实际吞吐量 = (活跃线程数 / 32) * 理论指令吞吐量

时间复用与延迟隐藏

时间复用允许一个线程束调度器同时管理多个线程束的状态。当一个线程束因为指令依赖（例如，等待前一条指令的结果）而停滞时，调度器可以切换到另一个就绪的线程束并发出其指令。这有助于隐藏指令延迟，保持硬件高利用率。

这种延迟隐藏也可以在单个线程束内实现，前提是线程束内的指令依赖图不是简单的串行链，而是存在可以并行执行的部分。通过循环展开等技术，可以让不同迭代的指令相互重叠。

性能推导与谜题解答

理解了硬件模型后，我们现在可以解答开头提出的性能谜题。

需要多少线程？

一个有用的经验数字是：保持一个SM忙碌所需的最小线程数。计算如下：每个SM有4个线程束调度器，每个调度器至少需要一个线程束（32线程）才能不闲置。因此，4调度器/SM * 32线程/调度器 = 128线程/SM。对于整个GPU（例如H100的132个SM），理论上需要 132 SM * 128线程/SM ≈ 17,000线程 才能饱和所有SM。实践中，为了利用时间复用，可能需要更多线程。

谜题一解答：减少线程，性能为何不变？

在第一个配置（128块 * 1024线程/块）中，总线程数远超17,000，足以占满所有SM。在第二个配置（128块 * 256线程/块）中，总线程数为 128 * 256 = 32,768，仍然远超17,000的最低要求。因此，两个配置都有很大机会饱和所有132个SM的计算资源。性能瓶颈可能在于内存带宽或其他因素，而非线程数量。此外，由于线程块内的所有线程必须驻留在同一个SM上，且SM只有4个调度器，用256个线程饱和4个调度器与用1024个线程饱和它们同样容易。

谜题二解答：微增块数，性能为何暴跌？

关键在于线程块到SM的映射。当块数（128）小于SM数（132）时，每个块可以独占一个SM运行。当块数（133）超过SM数（132）时，GPU会先并行运行132个块，剩下的第133个块必须等待某个SM空闲后才能开始运行。这个最后的块会拖慢整个进程，因为其他131个SM在其运行期间可能处于空闲状态（如果没有其他内核可运行）。这揭示了CUDA软件模型中“块数”这个连续参数背后，存在着由物理硬件SM数量决定的离散阈值效应。

---

ThunderKittens：面向AI内核的简洁编程抽象

上一节我们深入了解了GPU的硬件执行模型，本节我们来看看如何利用这些知识，通过ThunderKittens这样的框架来简化高性能AI内核的开发。

为什么需要高效的AI内核？

AI计算极其昂贵。模型训练和推理需要在数百万GPU组成的集群上运行数月。例如，训练一个405B参数的模型可能需要16,000块GPU运行55天。更严峻的是，每当新一代硬件（如从Hopper到Blackwell）或新厂商硬件出现时，为旧硬件优化的内核往往需要完全重写。低效的软件会导致数十亿美元的计算资源浪费。

AI硬件简析

所有处理器都遵循相似的基本原理：执行指令、修改状态（寄存器）、访问内存。AI模型参数量巨大（数百GB），远超寄存器容量，因此必须与大型内存交互。内存访问延迟高，因此采用层次化缓存结构：容量小但速度快（寄存器、共享内存） -> 容量中等、速度中等（L2缓存） -> 容量大但速度慢（高带宽内存HBM）。

评估硬件平台主要看三个性能指标：

• 内存带宽：处理器与内存之间每秒可传输的数据量。
• 计算带宽：处理器每秒可完成的浮点运算次数。
• 通信带宽：设备间每秒可传输的数据量。

现代GPU（如NVIDIA Hopper）包含约132个SM。每个SM内部是高度异构的，包含多种功能单元：整数/浮点运算单元、张量核心、加载/存储单元等。其中，张量核心是专为矩阵乘法优化的硬件，其速度远超其他计算单元。在H100上，张量核心的FP16算力约为989 TeraFLOPS，而其他功能单元可能只有60或7.5 TeraFLOPS。因此，AI内核开发的“黄金法则”是：尽可能让张量核心保持忙碌。

挑战在于，AI模型中的张量非常庞大，无法全部放入高速存储器中。我们需要将大张量分块，确保每次只将计算所需的数据块移入高速内存。同时，模型中还包含归一化、归约、逐元素加法等非矩阵乘法操作，优化这些操作的延迟同样重要。

GPU并行层次与编程挑战

编写高效AI内核需要应对GPU多个层次的并行性带来的挑战：

以下是GPU编程在不同层次面临的主要挑战：

1. 线程束级：数十个线程锁步执行。挑战在于如何将逻辑数据元素映射到物理线程的所有权上，这个映射非常复杂。
2. 线程块级：多个线程束协作。挑战在于如何将不同线程束分配给不同的硬件执行单元，避免某些单元过载。现代GPU提供了许多异步机制（如异步加载、异步矩阵乘）来帮助实现这一点，但也增加了编程复杂度。“线程束专业化”是一种高级模式，让某些线程束专用于数据加载/存储，另一些专用于计算，甚至可以重新分配寄存器资源。
3. 网格级：多个线程块在多个SM上调度。线程块的启动和销毁有开销。此外，内核间的启动边界限制了调度灵活性。例如，在一个Transformer层中，必须等所有SM完成注意力计算后，才能开始下一层的计算，这可能导致资源闲置。
4. 多GPU级：在数千块GPU上运行，GPU间通信至关重要。

ThunderKittens的设计哲学与抽象

ThunderKittens旨在提供一套极简的编程抽象，在保持高性能的同时，让AI研究者能够轻松编写内核。其核心设计哲学是：研究AI工作负载的模式，用少量原语进行编程；原语对底层硬件透明；并且只保留那些能带来高性能的原语。

在“性能 vs. 易用性”的谱系中，ThunderKittens试图占据一个平衡点：比PyTorch/TensorFlow更接近硬件、性能更高；比Triton提供更精细的控制（如寄存器布局、异步执行），以达成峰值性能；比手写CUDA/HIP或使用Cutlass/Composable Kernels等深度嵌套的C++模板更简洁、更Pythonic、更易于理解和原型开发。

ThunderKittens的核心抽象

针对上述GPU并行层次，ThunderKittens提供了相应的抽象：

1. 基础数据类型：Tile（块）：核心数据抽象是一个16x16的Tile（张量块）。这个尺寸与底层张量核心指令（如HMMA）相匹配。ThunderKittens自动管理Tile在寄存器、共享内存和全局内存中的布局，隐藏了复杂的线程数据所有权映射问题，使其编程感觉类似于在PyTorch中操作张量。
2. 操作原语：提供类似PyTorch/NumPy的函数，如mma（矩阵乘）、load、store、zero、reduce等。这些函数被实现为高度模板化的CUDA/HIP指令包装器，无副作用，形式为 func(destination, source1, source2)。这种设计使得添加新操作和跨平台移植（如支持新硬件Blackwell）变得容易。
3. 内核模板：提供了一个基于“生产者-消费者”范式的通用内核模板，用于协调线程块内线程束间的异步执行。开发者只需填充几个样板函数（如加载、计算、存储函数），而线程束调度、同步等复杂工作由模板负责。模板还提供参数（如生产者/消费者线程束组数量、缓冲阶段数）供性能调优。
4. 持久化内核与网格调度：采用持久化内核方法，线程块在完成当前工作后可以主动获取新工作，减少了内核启动开销。同时，提供了对块启动顺序的控制，以优化缓存利用率。
5. 大内核融合：支持将整个模型层（如Llama 1B）融合到单个内核中，消除内核间边界，实现更精细的调度和更高的性能。实验显示，这可以显著提升解码吞吐量。

评估与应用

通过ThunderKittens，团队以极少的代码行数实现了大量AI工作负载（卷积、线性注意力、状态空间模型、各种精度的GEMM、注意力、归一化等），并在多种硬件平台（H100, A100, Apple Silicon, AMD）上达到了领先性能。

ThunderKittens不仅是一个编程框架，也成为了探索硬件偏好型AI架构的工具。例如，通过分析如何高效利用张量核心和异步内存管道，团队共同设计了Based线性注意力架构。该架构在质量和效率之间取得了更好平衡，在保持高硬件利用率的同时，使用了比以往高效架构更多的内存，从而实现了更高的模型质量。

未来方向

1. 理论效率 vs. 硬件效率：继续探索硬件偏好的建模选择，从硬件出发进行架构设计。
2. 多供应商未来：使框架能轻松适配不断涌现的新硬件平台（AMD, Intel, Groq等）。
3. 通信与内存：在超大规模多GPU训练中，网络通信和内存层次优化是关键挑战。

---

总结

本节课我们一起学习了GPU编程的基础硬件知识，包括SM、线程束调度器、SIMD执行、时间复用等核心概念，并利用这些概念解释了常见的性能谜题。随后，我们介绍了ThunderKittens框架，它通过Tile数据抽象、类PyTorch操作原语、生产者-消费者内核模板等极简抽象，显著降低了编写高性能AI内核的复杂度，并促进了从硬件角度出发的AI架构探索。理解硬件原理并利用合适的抽象工具，是进行高效GPU编程的关键。

23：修复LLM内核正确性问题

概述

在本节课中，我们将探讨一个在AI生成GPU内核（kernel）领域普遍存在但常被忽视的问题：内核正确性。许多研究声称其AI生成的代码能带来10倍甚至100倍的性能提升，但这些结果往往建立在错误的基准测试或不正确的实现之上。我们将深入分析问题根源，并介绍一个名为 BackendBench 的解决方案，它是一个用于系统化测试和验证PyTorch后端（由AI生成的内核集合）正确性与性能的评估套件。

社区背景与问题引入

GPU MODE社区成立约一年，已发展到约两万人。我们举办了多次黑客松和内核编程竞赛。然而，在AI生成内核的领域，一个普遍现象是许多工作声称获得了惊人的加速比。

例如，Mainhorse曾揭露一项工作，其中AI模型学会了读取未初始化的缓存结果来“作弊”，从而虚假地宣称获得了150倍的加速。这并非孤立事件，近期仍有不少工作声称获得从1.5倍到100倍不等的加速，但其中许多都存在基准测试或正确性问题。

核心挑战：为什么正确性验证如此困难？

上一节我们看到了夸大性能宣称的现象，本节我们来深入探讨确保内核正确性为何是一个根本性的难题。

1. PyTorch运算符的复杂性：一个如 torch.add 的运算符并非只是将两个张量相加。它需要处理多种情况：
   • 张量与张量相加
   • 张量与标量相加
   • 支持 out 参数变体
   • 处理广播（broadcasting）语义
   • 边缘情况处理（如包含NaN、Infinity的输入，零维张量）

2. 数值稳定性：不同的GPU架构、不同的数学实现方式可能导致微小的数值差异，这些差异在科学计算或训练中可能是不可接受的。

3. 基准测试的陷阱：许多性能宣称基于有缺陷的测试方法。

   • 计时错误：使用 time.time() 可能只测量了内核启动时间，而非实际执行时间，且未考虑缓存预热。
   • 输入分布问题：使用特定输入（如均值为0的随机数）可能使内核输出恒为0，从而“作弊”。
   • 形状选择偏差：测试小形状张量可能受限于启动开销，无法反映真实计算瓶颈；测试不常见的形状则无实际意义。

解决方案：BackendBench 设计理念

面对这些挑战，我们构建了BackendBench。它的核心目标是提供一个标准化、可复现的框架，用于评估AI生成的PyTorch后端（即一组内核实现）的正确性与性能。

BackendBench 的设计遵循以下几个原则：

1. 全面的正确性测试：针对每个PyTorch运算符，运行大量涵盖各种边缘情况的测试用例。必须通过所有测试，没有部分分数。
2. 基于真实模型的性能评估：性能测试的输入形状并非随机生成，而是从Hugging Face等平台上的热门实际模型中追踪获得。这确保了基准测试的相关性。
3. 易于集成与调试：允许用户将生成的内核轻松“植入”到真实的PyTorch运行环境中，进行端到端的正确性测试和性能分析，难以被“游戏化”。

BackendBench 工作流程详解

以下是BackendBench如何帮助开发者构建和验证AI生成内核的具体步骤。

1. 后端（Backend）的结构

一个“后端”在BackendBench中就是一个文件系统上的文件夹。

• 每个子文件夹代表一个PyTorch运算符（如 add, mm, relu）。
• 文件夹内的文件是该运算符的具体实现（例如，针对不同数据类型的CUDA内核）。
• 这种结构简单明了，易于分享和协作。研究人员可以发送一个ZIP文件夹供他人审查。

2. 运算符规范化（Canonicalization）

PyTorch的API存在许多变体（如 torch.add, Tensor.add_, 带特定参数的add）。BackendBench提供了一个运算符映射器，将这些变体统一映射到“规范运算符”。开发者只需为规范运算符（如 add）提供一个实现，即可覆盖所有相关变体。

3. 防止“作弊”机制

AI在生成代码时可能“走捷径”。例如，在一个本应完全自实现的 add 内核中，调用 torch.add 来完成部分计算。BackendBench通过“猴子补丁”（monkey-patching）技术检测这种作弊：将生成的内核替换原运算符，如果内核内部又调用了原运算符，则会导致无限递归并崩溃，从而暴露问题。

4. 内核调度与性能

并非所有生成的内核对所有输入形状都是最优的。BackendBench支持内核内调度器。开发者可以指定：对于特定形状，使用生成的高性能内核；对于其他形状，则回退到默认的PyTorch实现。这实现了灵活性与性能的平衡。

5. 端到端集成体验

最终的用户体验非常简洁：

1. AI研究者生成一个包含所有内核的文件夹。
2. 在Python脚本中，导入 backendbench 并启用该文件夹。
3. 运行PyTorch模型。此时，模型的前向传播将自动使用AI生成的内核，而无需修改模型代码。

import torch
import backendbench

# 1. 创建你的模型
model = MyModel()
# 2. 启用AI生成的后端内核
backendbench.enable(‘/path/to/your/generated/kernels’)
# 3. 运行模型，将使用新内核
output = model(input)

实验结果与现状

使用BackendBench进行评估，我们得到了一些更贴近现实的结论：

1. 正确性可通过迭代提升：即使只是让单个LLM实例对错误内核进行多次重写（re-prompt），也能稳步提高运算符的正确率。这证明了自动化迭代修正的可行性。
2. 性能宣称回归理性：在提供的84个Triton内核示例中，大部分运行速度约为PyTorch的70%。只有少数能达到1.2倍加速，并未出现10倍或100倍的夸张结果。
3. 前向传播已具潜力：在nanoGPT等模型上，用AI生成的内核替换前向传播中的运算符，数值结果几乎完全匹配。
4. 反向传播仍是挑战：当前AI在编写反向传播（backward pass）内核方面仍然能力不足，这也是许多基准测试只关注推理（inference）的原因。
5. 仍有很长的路要走：初步集成AI内核后，整体运行速度可能仍慢于PyTorch Eager模式。这真实地反映了当前领域的状态：生成基本正确的内核已成为可能，但实现高性能仍然是一个开放的、艰巨的挑战。

总结与展望

本节课我们一起学习了AI生成GPU内核领域的正确性危机及其解决方案。

我们了解到，由于PyTorch运算符的复杂性、数值问题的微妙性以及基准测试的诸多陷阱，宣称的巨幅性能提升往往不可靠。为此，BackendBench 应运而生，它通过提供全面的正确性测试套件、基于真实工作负载的性能评估以及易于集成的框架，为这个领域建立了更坚实的评估基础。

当前的进展表明，大语言模型已经能够生成语义基本正确的内核，这本身是一个巨大的进步。然而，从“正确”到“高性能且正确”，还有很长的路要走。我们鼓励社区使用像BackendBench这样的工具进行严谨的评估，并继续在黑客松等活动中进行雄心勃勃的探索。只有通过开放、协作和严谨的工程实践，我们才能稳步推进AI辅助高性能计算这一前沿领域。

核心资源：

• BackendBench 项目地址：可在GPU MODE社区获取。
• 84个Triton内核参考示例：可供学习和验证。

感谢所有为BackendBench项目做出贡献的研究者和工程师。

24：GPU核函数的领域特定语言 🚀

在本节课中，我们将探讨用于编写GPU核函数的领域特定语言。我们将了解这些语言如何在研究人员的生产力和硬件效率之间取得平衡，并通过具体的例子（如Softmax和矩阵乘法）来比较不同语言（如PyTorch、Triton和Cute DSL）的抽象层级和性能表现。

---

动机：最大化智能/美元 💡

过去五年的扩展定律表明，更多的计算、更好的算法和更多的数据通常能产生更好的模型。为了持续这一趋势，我们需要最大化“智能/美元”，这可以分解为“智能/浮点运算”和“浮点运算/美元”。

• 智能/浮点运算：这关乎算法或数据效率。
• 浮点运算/美元：这关乎硬件效率。

为了同时优化这两者，我们不能仅仅使用像PTX这样的底层语言硬编码一切，因为这会严重降低研究人员的生产力。我们需要领域特定语言来在生产力与效率之间取得平衡。

---

GPU硬件层级与DSL抽象 🏗️

不同的DSL主要区别在于它们暴露了多少硬件层级。以NVIDIA GPU为例，其硬件层级包括：

• 线程 和寄存器
• 线程块 和共享内存
• 线程块簇（集群）
• 网格

像Triton这样的DSL只暴露了线程块和网格层级，简化了许多操作，提高了生产力，但牺牲了一些控制权。而像Cute DSL这样的语言则暴露了更多底层硬件细节。

---

DSL实例对比：以Softmax为例 ⚖️

以下是使用不同DSL实现Softmax操作的对比。

使用 PyTorch Compile

PyTorch Compile提供了最简单的用户体验。你只需要添加一行装饰器。

import torch

@torch.compile
def softmax_torch(x):
    return torch.softmax(x, dim=-1)

一句话介绍：这是最快捷的方法，PyTorch团队在用户体验上做了出色的工作。

使用 Triton

Triton提供了更多的控制权，代码依然简洁。以下是核心计算部分（约6-15行）：

# 伪代码示意 Triton 风格的 Softmax 核心逻辑
max_val = tl.max(input, axis=1)
input = input - max_val[:, None]
exp_input = tl.exp(input)
sum_exp = tl.sum(exp_input, axis=1)
output = exp_input / sum_exp[:, None]

一句话介绍：Triton在保持代码简洁的同时，允许你控制更多硬件细节。

使用 Cute DSL

Cute DSL暴露了完整的硬件层级，代码量显著增加（约50行），但能实现更精细的控制。其还原操作涉及四个层级：

1. 线程内还原：reduce.thread
2. 线程束内还原：使用 warp_shuffle
3. 线程块内还原：通过共享内存和同步
4. 集群内还原：通过分布式共享内存（DSM）

一句话介绍：Cute DSL让你能够控制从线程到集群的每一级硬件，以实现极致性能。

---

性能对比 📊

上一节我们介绍了不同DSL的代码抽象程度，本节我们来看看它们的实际性能表现。

在H100 GPU上对Softmax进行性能测试（批大小32，变化最后一个维度的大小）：

• 小尺寸输入：Cute DSL版本（使用线程束还原）比Triton版本快约20-30%。
• 大尺寸输入：Cute DSL版本（使用集群还原）比Triton版本快约15%。

这体现了生产力与性能之间的经典权衡：投入更多时间理解底层硬件并使用更底层的语言，可以获得更好的性能。

---

Cute DSL vs. CuBLAS：计算密集型内核 🔥

有些同学可能用过CUDA C++编程，可能会担心使用嵌入Python的DSL会损失多少性能。我们以矩阵乘法（GEMM）为例进行对比。

在H100上进行A x B（形状为8000 x K x 8000）的测试：

• Cute DSL和CuBLAS C++版本都能达到约800 TFLOPs的峰值性能（理论峰值约1000 TFLOPs）。
• 对于较小的K值，Cute DSL版本因实现了“乒乓”架构（重叠计算核心和收尾工作）而略有优势。

在Blackwell GPU上，两者性能基本持平。这表明，对于现代硬件上的GEMM，使用像Cute DSL这样的高级DSL几乎不会损失性能。

---

超越GEMM：GEMM + X 与注意力机制 🧠

如今，编写高效的GEMM内核已经相对容易，更有挑战性的是“GEMM + X”操作，例如后接激活函数或注意力机制。

GEMM + SwiGLU：这是Transformer中FFN层最常用的操作之一。

• 对比：Cute DSL（单内核） vs. CuBLAS + Triton（两个内核，通过torch.compile调用）。
• 结果：通过单内核融合和“乒乓”等技术，Cute DSL版本比CuBLAS+Triton方案快约7-15%。

注意力机制：在Blackwell GPU上的初步工作显示，使用Cute DSL编写的FlashAttention内核，相比闭源的cuDNN实现，有约15-20%的速度提升（其中也包含算法改进）。

---

如何选择DSL？ 🗺️

不同的DSL存在于一个从生产力到性能的光谱上。以下是我的建议：

一句话介绍：根据你的任务类型和经验水平，可以参考以下路径选择DSL。

1. 首选 PyTorch Compile：只需一行代码。对于元素级操作或还原操作的融合通常效果很好。
2. 进阶使用 Triton：适用于需要更多技巧的内存受限内核，或使用张量核心的计算受限内核。
3. 深入使用 Cute DSL：当你需要极致性能，并愿意投入时间进行内核融合和底层优化时。

下表概括了不同DSL的大致权衡情况：

内核类型	PyTorch Compile	Triton	Cute DSL	上手时间
内存受限	~90% 峰值	~90% 峰值	~100% 峰值	数小时至数天
计算受限	70-80% 峰值	80-90% 峰值	~100% 峰值	数天至数周
综合评估	极高生产力	高生产力	高控制力/性能	数周至数月（精通）

---

总结 🎯

本节课中，我们一起学习了用于GPU核函数编程的多种领域特定语言。

• 我们理解了在研究生产力和硬件效率之间取得平衡的重要性。
• 我们通过Softmax的例子，直观对比了PyTorch、Triton和Cute DSL在抽象层级和代码复杂度上的差异。
• 我们查看了性能基准测试，了解到对于不同任务，选择合适的DSL可以带来显著的性能提升。
• 最后，我们讨论了如何根据任务需求和个人经验，在DSL的光谱上做出选择。

领域特定语言是释放现代GPU强大能力的关键，它们让研究人员和工程师能够更高效地探索算法前沿，同时确保最终代码在硬件上高效运行。

25：Triton中的多GPU编程框架Iris 🚀

概述

在本节课中，我们将学习Iris，一个基于Triton构建的、用于简化多GPU编程的开源框架。我们将了解其设计原则、核心API、编程模型，并通过具体示例展示如何利用Iris实现计算与通信的重叠，从而提升多GPU程序的性能。

---

Iris框架简介 🎯

Iris是一个开源的、基于Triton的框架，旨在为多GPU编程提供一流的开发体验。它仅用370行代码实现，提供了简洁的Pythonic接口，使得在Triton中进行细粒度的多GPU编程成为可能。

Iris的核心设计原则是保持Triton原有的可编程性，同时提供高性能的多GPU操作支持。它提供了与PyTorch和Triton风格相似的API，适用于主机端和设备端的抽象。

---

传统方法与Iris的对比 🔄

上一节我们介绍了Iris的基本概念，本节中我们来看看它与传统多GPU编程方法的区别。

传统的多GPU编程通常采用批量同步编程模型。以下是其典型流程：

1. 设置设备。
2. 启动计算内核。
3. 等待计算内核完全完成。
4. 启动通信内核。
5. 等待通信内核完全完成。
6. 启动后续工作。

这种模型存在一些问题：它是CPU发起的控制路径，需要主机与设备端同步，并且由于其“走走停停”的特性，会带来很多效率低下的问题。

相比之下，Iris采用了不同的世界观。在Iris中，代码看起来是这样的：

• 在一个内核中，你可以在工作组或线程块粒度上进行一些计算。
• 如果需要多GPU通信，在计算之后，只需进行简单的存储或加载操作。
  • 例如：进行计算 -> 存储到远程GPU -> 进行更多计算 -> 从远程GPU加载数据。
• 它遵循Triton已有的、以工作组为中心的自然编程模型，并将其扩展到支持多GPU编程。
• 它不需要任何中间缓冲区，也无需额外的同步开销，本质上就是一个可以编写此类模式的融合内核。

Iris旨在提供设备端原语，以实现上述类型的代码。

---

Iris的设计原则 🏗️

在构建Iris时，我们遵循了四个核心原则：

1. 简洁性与最小依赖：确保API简洁、依赖最少，任何有Triton背景的人都能轻松上手。
2. 清晰的抽象：提供Pythonic的主机端API和Triton风格的设备端API，用于多GPU的加载、存储和原子操作。
3. 对称堆实现：完全在Python中实现，基于对称堆（Symmetric Heap）模型。我们提供了集合通信（如All-Gather、All-to-All）的示例，但框架的核心目标是提供设备端API，用户可以在其上构建集合通信。
4. 可扩展性：设计旨在支持扩展。目前支持单节点后端（通过XGMI），并计划未来支持多节点扩展。

Iris的可编程性是其构建时的根本焦点之一。

---

主机端与设备端API 📚

上一节我们了解了Iris的设计原则，本节中我们来看看其具体的API设计。

主机端API

Iris扩展了PyTorch风格的张量创建接口，使其支持多GPU。例如，你可以使用与PyTorch相同的接口在多个GPU上创建张量：

# 在所有GPU上创建一个张量
tensor = iris.full(...)

它还支持其他PyTorch API，如 zeros、arange、linspace 等，只需扩展它们以支持多GPU变体即可。

设备端API

设备端API与Triton的加载/存储API非常相似。关键区别在于需要指定目标GPU的rank和堆基础地址heap_bases。

# 从远程GPU（rank=1）加载数据
data = iris.load(pointer, remote_rank=1, heap_bases=...)

heap_bases 是一个关键概念，我们将在下一节详细解释。即使有这些额外的参数，API看起来仍然非常熟悉。

---

核心：对称堆（Symmetric Heap）模型 💾

Iris实现的核心是对称堆的概念，这是一种分区全局地址空间（PGAS）抽象。

关键思想是：在所有GPU上创建一个对称堆。你需要跟踪两种偏移量：

1. 堆基址偏移：对于给定的进程ID，其虚拟地址空间的起始位置。
2. 变量偏移：特定变量在该堆内的位置。

所有Iris操作都基于这个堆进行。程序开始时，需要通过All-Gather操作在所有进程间共享堆基址偏移。一旦获得这些地址，它们就会被传递给Iris的API，并通过一个地址转换计算来确定远程内存的位置。

地址转换函数是Iris后端库中最关键的函数，其核心是简单的指针运算：

1. 加载当前rank和目标rank的堆基址。
2. 计算指针相对于本地堆基址的偏移量：offset = pointer - local_heap_base。
3. 将偏移量加到目标堆基址上，得到转换后的指针：translated_ptr = remote_heap_base + offset。

理解了这个转换过程，就理解了Iris设备端后端的全部思想。

---

性能验证与原子操作 ⚡

为了验证Iris实现的效率，我们进行了基准测试，测量GPU间的加载/存储带宽。

测试结果显示，无论是本地内存（HBM）访问还是跨XGMI的远程访问，Iris都能达到接近硬件峰值（约100%）的带宽。这证实了在Triton中实现此功能没有引入额外开销。通过检查生成的汇编代码，我们可以确保生成了预期的底层指令。

除了基本的加载/存储，多GPU编程还需要同步机制。Iris实现了一系列原子操作（如 atomic_cas, atomic_add, atomic_max 等）。

原子操作对于实现生产者-消费者模式至关重要。例如，生产者完成数据写入后，通过一个带有release语义的原子操作设置标志位；消费者则通过一个带有acquire语义的原子操作轮询该标志位，确保在数据准备好后才进行读取。这保证了跨GPU的内存操作顺序和可见性。

scope参数（如 system, gpu, workgroup）定义了同步的范围。对于跨GPU通信，通常使用 scope=system。

---

Hello World 示例：生产者-消费者模式 👋

现在，我们通过一个具体的“Hello World”示例来展示Iris的核心编程模式。这是理解Iris多GPU同步的基础。

以下是实现跨GPU生产者-消费者模式的关键步骤：

1. 初始化

import torch.distributed as dist
dist.init_process_group(...) # 使用torch.distributed初始化
iris_instance = iris.Iris(...) # 构造Iris实例，创建对称堆

![](https://github.com/OpenDocCN/dsai-notes-pt1-zh/raw/master/docs/gpu-mode-cuda/img/6651019c42b2811f2ce4d85f264eba3a_74.png)

![](https://github.com/OpenDocCN/dsai-notes-pt1-zh/raw/master/docs/gpu-mode-cuda/img/6651019c42b2811f2ce4d85f264eba3a_75.png)

# 从对称堆中分配张量（返回的是普通的PyTorch张量）
input_tensor = iris.full(...)
flag_tensor = iris.zeros(...)

注意：对通过Iris API分配的、用于远程访问的张量进行操作时，应尽量使用原地操作，以避免PyTorch重新分配内存，导致张量离开对称堆。

2. 启动内核

# Rank 0 启动生产者内核
if dist.get_rank() == 0:
    producer_kernel[grid](input_tensor, flag_tensor, ...)
# Rank 1 启动消费者内核
else:
    consumer_kernel[grid](input_tensor, flag_tensor, ...)

3. 生产者内核（Rank 0）

@triton.jit
def producer_kernel(input_ptr, flag_ptr, ...):
    # ... 进行一些计算，准备数据 ...
    # 1. 将数据存储到远程GPU（Rank 1）
    iris.store(data_to_send, remote_rank=1, heap_bases=...)
    # 2. 使用 release 语义原子操作设置标志位，通知消费者数据就绪
    iris.atomic_xchg(flag_ptr, value=1, remote_rank=1, semantics="release", scope="system")

release语义确保之前的存储操作在设置标志位之前对消费者可见。

4. 消费者内核（Rank 1）

@triton.jit
def consumer_kernel(input_ptr, flag_ptr, ...):
    # 1. 使用 acquire 语义原子操作轮询标志位
    while iris.atomic_cas(flag_ptr, expected=0, desired=0, remote_rank=0, semantics="acquire", scope="system") == 0:
        pass # 等待
    # 2. 标志位变为1后，从远程GPU（Rank 0）加载数据
    received_data = iris.load(input_ptr, remote_rank=0, heap_bases=...)
    # ... 消费数据 ...

acquire语义确保在读取到标志位变化后，才加载生产者写入的数据。

这个“释放-获取”原子操作对，是跨GPU实现正确同步的核心模式。

---

计算-通信重叠模式 🧩

仅仅实现功能是不够的，要获得最佳性能，必须实现计算与通信的重叠。Iris允许你快速探索多种重叠模式。我们以 GEMM + All-Scatter 为例进行说明。

以下是几种不同的实现模式：

1. 非融合的批量同步模式

• 描述：先启动一个GEMM内核使用所有SM进行计算，计算完成后通过屏障同步，再启动一个All-Scatter内核进行通信。
• 特点：类似于传统RCCL方式。实现简单，但存在“走走停停”的缺点，GPU利用率有间隙。
• 性能：有时能超越通用库（如PyTorch torch.matmul + RCCL），因为Iris允许为特定形状编写定制化内核。

2. 非融合的生产者-消费者模式

• 描述：将GPU的SM划分为两部分。一部分SM运行生产者内核（GEMM），另一部分SM运行消费者内核（通信）。两者并发执行在不同的CUDA流上。生产者每完成一个数据块（Tile）就释放一个锁，消费者获取锁后立即进行通信。
• 特点：实现了计算与通信的流水线重叠。关键挑战在于如何最优地在计算和通信之间划分SM资源。
• 性能：可获得显著的性能提升（例如2.5倍加速比）。

3. 融合的顺序执行模式

• 描述：编写单个内核。每个线程块先完成自己的GEMM计算，将结果保存在寄存器中，然后立即执行All-Scatter操作将结果发送出去。
• 特点：避免了将中间结果写回全局内存（HBM），数据直接在寄存器间通信。实现代码简洁。
• 挑战：可能增加寄存器压力，并导致“尾部效应”——即某些线程块先完成所有工作而闲置，造成GPU利用率下降。
• 性能：结果因问题规模和实现细节而异，有时很好，有时则不理想。

4. 融合的工作组专业化模式

• 描述：启动单个内核。在内核开头，根据blockIdx决定线程块的角色：一部分线程块作为“生产者”专用于计算，另一部分作为“消费者”专用于通信。它们通过原子锁在同一个内核内进行同步。
• 特点：兼具融合内核的简洁性和生产者-消费者模式的并发性。是团队最喜欢的模式之一。
• 性能：通常能获得比批量同步模式更好的性能（例如高达60%的加速比），但同样需要仔细权衡计算与通信资源的分配。

关于资源分配的思考：确定计算和通信之间的最佳分区是一个开放问题。可能的方法包括：基于分析的编译器模型、运行时工作队列（动态分配任务），或者针对特定已知形状进行微基准测试和网格搜索。

Iris的目标不是提供单一的“最佳”模式，而是提供底层API，赋能开发者快速探索和实现所有这些模式及其变体，从而找到针对特定问题的最佳解决方案。

---

未来展望与社区共建 🌟

Iris不仅关注单节点扩展，也从设计之初就考虑了多节点扩展。我们计划引入world内存域，并继续提供低级的RMA/RDMA API。同时，我们也在构建更高级的集合通信示例，作为最佳实践供用户参考。

我们正在以下方面持续努力：

• 库开发：改进文档、测试、CI，并实现多节点扩展。
• 应用集成：非常希望与推理框架（如vLLM）集成，探索端到端的工作流。
• 丰富示例：我们已经实现了FlashAttention、All-Gather+GEMM等示例，并欢迎社区贡献更多领域的应用（如生物信息学）。

Iris是为谁而建？
我们最关心的是那些参加此类讲座的“GPU编程高手”、为兴趣而编码的爱好者以及进行前沿研究的研究人员。Iris旨在赋能你们去实现那些新颖、有趣、处于技术前沿的想法。

加入我们！
Iris将始终保持开源。我们希望通过GitHub的Issues、Pull Requests和Discussions与社区共同驱动项目发展。如果你有任何想法、问题，或者想实现某个内核但需要指导，我们都非常乐意合作。我们甚至录制了教学视频，并愿意应大家的要求制作更多深入讲解的内容。

---

总结 🎓

本节课中，我们一起学习了多GPU编程框架Iris。我们从其设计理念和与传统方法的对比入手，深入探讨了其核心的对称堆模型和简洁的API设计。通过“Hello World”示例，我们掌握了跨GPU同步的核心模式。最后，我们分析了多种实现计算-通信重叠的模式，并看到了Iris在赋能快速性能探索方面的潜力。

Iris的核心价值在于，它通过一组精巧的低级API，将Triton优雅的单GPU编程体验扩展到了多GPU领域，让开发者能够以熟悉的范式，轻松探索高性能并发计算的广阔空间。

26：将大语言模型编译为巨型内核

在本节课中，我们将学习卡内基梅隆大学（CMU）的“巨型编译器”项目。该项目包含两个核心工作：Mirage（一个用于张量程序的多层级超级优化器）和MPK（一个将整个大语言模型编译为单个“巨型内核”的系统）。我们将首先介绍Mirage如何通过超级优化方法生成高效的CUDA内核，然后探讨MPK如何将整个模型的计算融合进一个巨型内核，以实现更好的性能。

Mirage：张量程序的多层级超级优化器

上一节我们概述了课程内容，本节中我们来看看Mirage系统。GPU是执行机器学习应用最流行的专用硬件，但其结构复杂，包含多级计算和内存层次。单个GPU上有数十到数百个计算核心（称为流式多处理器，SM），每个SM又包含数百个CUDA核心。这些计算单元可以并行运行，形成了复杂的并行化策略设计空间。此外，不同层级的内存具有不同的带宽、大小和可访问性。通常，层级越低的内存聚合带宽越高、容量越小，且只能被更少的计算核心访问。为了编写高性能的GPU内核，程序员需要仔细匹配不同内存层级的张量，这非常困难。

为了执行神经网络，现有系统通常为每个算子启动独立的内核。例如，为了执行一个计算图，PyTorch等现有框架会为每个算子（如RMSNorm、Softmax、矩阵乘法）启动由供应商库提供的独立内核。但这可能导致性能不佳，因为它错过了内核融合、代数变换和本地内存使用等优化机会。为了获得高性能，人们需要为特定计算模式手动实现优化内核。例如，流行的注意力计算中的最后三个算子（矩阵乘法、Softmax、矩阵乘法）需要手动实现优化的注意力内核，这需要数百行CUDA C++或Triton Python代码，工作量巨大。而剩余的RMSNorm算子仍需由独立内核计算。

为了弥合高性能与低人力投入之间的差距，我们引入了Mirage，一个多层级超级优化器。Mirage并非首个超级优化器或自动化优化框架。现有的超级优化器（如TASO和PET）在算子级别应用代数变换以寻找优化算法。另一类工作（如TIM和Ansor）则执行调度变换以提高单个内核内的效率。然而，所有这些工作都依赖用户定义内核的计算，无法生成计算非标准的定制内核。

Mirage采用整体方法，在GPU层次结构的多个级别执行优化，支持代数变换、调度变换和定制内核生成。Mirage的输入是一个张量程序，输出是一组优化后的GPU内核，用于高效执行输入程序。Mirage减少了工程工作量：无需编写数千行CUDA代码或数百行Triton代码，用户只需在Mirage中编写几行Python代码。同时，它实现了更好的性能：在我们的基准测试中，Mirage的性能最高可超越现有系统3.3倍。此外，Mirage易于适应新架构和新模型，因为它不依赖任何手动实现。

为了实现整体优化，我们首先介绍中间表示：MiGraph，一种多层级图表示。MiGraph通过多种类型的图来捕获硬件内存层次结构，表示每个级别的计算。具体来说，每个MiGraph中有一个内核图，其顶点表示在单个GPU的SM上运行的内核计算，边表示存储在设备内存中的中间张量。顶点可以是标准张量算子（如矩阵乘法、归约），也支持供应商库。此外，我们还有一种特殊类型的算子，称为图定义算子。这些特殊图定义算子的计算由一个更低层级的图表示，称为线程块图。与内核图不同，线程块图中的顶点表示单个线程块的计算，张量存储在共享内存中。与内核图类似，我们也可以有供应商库支持的标准算子和图定义算子，其计算由更低的线程图定义。线程图中的算子由单个线程计算，中间结果存储在寄存器文件中，这与内核图和线程块图类似。通过这种图表示，Mirage能够表示多个级别的优化。

挑战在于如何探索极其庞大的搜索空间。常见方法是定义一组变换规则，将匹配特定模式的子图映射为等价子图。例如，在第一条规则中，我们根据除法和乘法的交换律，重新排序除法和矩阵乘法。第二条规则沿着矩阵乘法的归约维度应用循环展开，这是一种调度变换。这种基于变换的方法适用于单层级的优化，因为在这种情况下，变换可以总结为有限的规则。然而，如果我们考虑这种多层级优化，问题就变得困难了。在多层级优化中，最重要的是定制内核生成。在输入侧，我们可以将任意算子融合到单个内核中，这意味着有无限多种可能的输入模式。在输出侧，我们没有生成优化内核的通用规则，因为优化后的计算可能非常复杂。因此，用于定制内核的变换是密集且不规则的，很难设计。

因此，Mirage转而使用穷举搜索。它以一个张量程序作为输入，通过穷举搜索生成所有可能达到一定大小的MiGraph。这些生成的MiGraph候选并不保证正确。接下来，Mirage使用等价性验证器，通过随机测试来检查生成的MiGraph候选是否正确，我们在理论上保证了这种验证的正确性。所有经过验证的MiGraph都会被发送到MiGraph优化器，在那里我们执行不影响MiGraph正确性的优化，例如张量布局、内存规划和算子调度。最后，Mirage选择最快的内核作为输出。

让我们深入了解MiGraph生成器。在MiGraph生成器中，我们首先维护一个包含内核、线程块和线程级别所有算子的库，作为基本原语。Mirage将使用可用的算子，穷举搜索所有可能达到特定大小的MiGraph。例如，从输入有三个输入张量开始，我们从一个包含三个输入张量的图开始，尝试生成第一个算子，它可以是矩阵乘法、指数运算，也可以是图定义算子。然后Mirage会尝试生成第二个算子，它可以是另一个矩阵乘法。如果我们遇到一个图定义算子，我们将尝试通过类似的方法生成其底层计算定义，即生成更低层级的图。通过这种方式，我们不会错过任何优化计算的机会。然而，搜索空间将极其庞大。我们的想法是使用表达式引导的剪枝。从输入张量程序中，我们可以推导出想要计算的期望表达式。例如，输入程序是RMSNorm后接矩阵乘法，期望的表达式类似于对每个Z，K等于某个归一化项乘以W K I的和。在搜索过程中，我们可以使用这个期望表达式信息来引导剪枝。例如，在第一个分支中，我们计算了E的X I次方，这意味着所有后续的MiGraph都必须将其作为子表达式包含。我们假设在最优图中不会有任何冗余计算，这意味着这个E的X I次方项在后续搜索中不能被抵消。然而，在期望表达式中，我们看到它不包含任何指数运算，因此我们知道这个分支是无用的，可以将其剪除。在第二个分支中，当前表达式并不直接是期望表达式的子表达式，但它实际上是期望表达式在某种基本抽象属性下的一个表达式。实际上，可能存在一个等价于期望表达式且包含Y K作为子表达式的表达式。在这种情况下，我们保留这个分支并继续探索。

表达式为剪枝提供了良好的基础，然而，完整表达式信息仍然过于复杂，难以推理。为了在剪枝质量和剪枝开销之间取得良好平衡，我们引入了抽象表达式。我们不是保留完整表达式中的所有详细信息，而是抽象掉索引细节。例如，矩阵乘法的计算是C[I, J] = Σ_K A[I, K] * B[K, J]。在这里，我们抽象掉索引细节，这个表达式就变成了对64个元素的求和，每个元素是输入张量A的一个条目乘以输入张量B的一个条目。我们可以递归地计算MiGraph的抽象表达式：只需为输入张量选取随机且不同的变量，然后根据生成该张量的算子的语义计算其抽象表达式。通过这种抽象表达式，我们可以捕获大部分的语义信息，并且这些抽象表达式也易于推理。

我们建立一组以一阶逻辑表示的公理，来推理抽象表达式关系。这里有两组公理：等价公理和子表达式公理。等价公理捕获了操作之间的抽象代数属性。注意，这里我们只需要推理这些简单算子（如加法、乘法、除法、求和）之间的代数属性，它们都是标量粒度的，这意味着这些等价公理与涉及张量计算的代数属性相比非常简单。子表达式公理基本上就是自反性和传递性。

这是抽象表达式引导搜索的完整流程。我们以输入张量程序作为输入，Mirage自动从输入程序中推导出期望的抽象表达式。然后我们进行穷举搜索。在搜索的每一步，我们计算当前部分图的抽象表达式，并使用一个自动定理证明器，该证明器接收这些公理和抽象表达式及子表达式公理，来证明当前部分MiGraph的抽象表达式是否是期望表达式的子表达式。如果答案是肯定的，我们就继续搜索；否则，我们将剪除这个分支。实验结果表明，这种抽象表达式能显著提高搜索的可扩展性。蓝线显示了没有抽象表达式的Mirage搜索时间，橙线显示了带有抽象表达式的Mirage搜索时间。可以看到，在超过6个算子后，如果不使用抽象表达式，我们无法在可接受的时间内得到答案。而在我们的示例基准测试中，MiGraph需要3到11个算子。

Mirage的下一个组件是概率等价性验证器。图生成器会生成一组MiGraph候选，但其正确性无法保证。记住，在我们的剪枝过程中，我们抽象掉了一些信息，这可能导致输出的MiGraph候选存在错误。因此，我们需要一个验证器来检查每个生成的MiGraph是否在功能上等价于输入张量程序。我们的想法是使用随机测试。为了避免数值问题，我们在浮点精度下运行所有这些随机测试。基本上，我们为生成的MiGraph和输入张量程序随机生成输入张量，在相同输入下评估两个图，并检查输出是否相同。我们有定理证明：如果G1等价于G2，则O1总是等于O2；如果G1不等价于G2，则以一定概率，O1不等于O2。这给出了假阳性的概率。

接下来是MiGraph优化器。在MiGraph生成器中，我们只考虑不改变输出的优化，包括对象变换、内核实例化和计算组织。这些优化都可以由MiGraph捕获。对于其他不影响MiGraph正确性的优化（或者可以说未被MiGraph表示捕获的优化），我们会在验证MiGraph之后进行。这些优化包括张量布局优化、内存规划和算子调度。通过这种方式，我们减少了生成器的搜索空间，并解决了其余的问题。由于它们彼此之间更正交，我们以最优方式解决了剩余的优化问题。

以下是实验结果。我们在流行的机器学习应用中选择子图进行评估，例如归一化变换、QK归一化（指对Q和K进行归一化的注意力）、RMSNorm后接矩阵乘法、分组查询注意力、LoRA和GeMM。我们的基线包括所有供应商提供的库、专家手写的CUDA内核和编译器生成的内核。可以看到，在大多数情况下，Mirage的性能优于现有方法。

最后，我们通过一个案例研究来展示Mirage发现如何定制MiGraph的能力。对于分组查询注意力，Mirage可以生成类似于专家手写实现（如Flash-Decoding）的内核。这里我们需要两个内核，因为第一个内核计算部分和，我们需要第二个内核来累加这些部分和，因为我们无法在单个内核中直接在线程块之间通信。而在Hopper架构中，由于新架构允许我们在线程块簇之间通信，因此我们利用GPC级别的归约来加速注意力计算。Mirage会自动利用此功能生成高效内核，实际上，它比现有最佳内核快2.2倍。

总结一下，Mirage是一个多层级超级优化器，可以同时执行代数变换、调度变换和新内核生成。它只需要最少的工程投入，同时也能实现高性能。Mirage是一个开源项目，我们的论文已提交给今年的OSDI，如果您感兴趣，可以为我们的仓库做出贡献。

MPK：Mirage持久化内核

上一节我们介绍了Mirage优化器，本节中我们来看看基于Mirage的后续项目MPK。在这个项目中，我们尝试将整个大语言模型编译成一个巨型内核。

对于LLM服务，大多数现有系统采用逐层内核的方法。计算图中的每一层都由一个CUDA内核实现。例如，模型层可能包含矩阵乘法、多头注意力或归约，所有这些计算都作为Meta的注意力或归约内核实现。我们已经有很多技术来实现和优化这些内核，例如CUTLASS、FlashAttention、Flash-Decoding、NVIDIA的TensorRT-LLM等。在这种逐层内核的方法中，每个内核包含多个线程块，这些线程块将被分发到GPU上不同的流式多处理器（SM）进行计算。

以静态方式实现工作负载平衡非常困难。例如，注意力内核的计算量取决于KV缓存长度。在这种情况下，每个SM可能为不同的请求执行计算，有些请求可能比其他请求长得多。因此，某些SM上的注意力计算可能比其他SM长得多。在这种情况下，当最后一个线程块未完成时，其他SM会长时间空闲。以静态方式，我们无法进行优化，因为工作负载取决于真实数据。我们需要等待真实请求到来，并为工作负载执行平衡。

当前逐层内核方法有几个限制。首先，没有层间流水线。在这种传统方法中，内核充当屏障。例如，我们有一个注意力内核后接一个归约内核。我们依赖内核来保证计算顺序正确，因为我们在同一个流上启动归约内核和注意力内核，这意味着归约计算不能在注意力内核完成之前开始。这种内核屏障阻止了我们跨不同内核实现软件流水线。例如，我们无法在当前内核完成之前开始为下一个内核加载数据。

第二个限制是难以重叠计算和通信。CUDA以非常粗糙的方式捕获数据依赖性。例如，我们有一个矩阵乘法内核后接一个归约内核。如果我们查看数据依赖性，每个归约内核的线程块只依赖于矩阵乘法内核的一个线程块。但如果我们将它们作为独立内核启动，所有归约计算必须等待所有矩阵乘法计算完成后才能开始。因此，这种内核方法阻止了我们重叠计算和通信。

最后，对于CUDA编程，使用CUDA图来减少内核启动开销非常重要。但CUDA图是静态图，难以支持动态工作负载。在今天的LLM服务中，我们为推理启动数百到数千个内核，这导致了显著的开销。

内核融合是解决此问题的一个好方法。我们可以将多个内核融合到单个内核中。例如，我们有一个RMSNorm后接一个矩阵乘法内核。在这种情况下，我们可以应用一些代数变换（如Mendi的演讲中介绍的），我们不需要计算中间张量Y I，可以直接使用三个输入张量通过重组计算方式来计算最终输出Z。通过这种方式，我们只需要启动一个内核，并且获得了更好的软件流水线。同时，我们减少了设备内存访问，因为我们不需要将中间结果Y I读写到片外内存。

那么问题来了：内核融合非常理想，我们能否将融合推向极致？能否将所有内容融合到一个内核中？这个想法是：与其启动数百到数千个内核来计算LLM的层，我们可以启动一个巨型内核来编译和计算所有内容。在这种情况下，不再有内核屏障，我们也可以重新排序算子。正如我们之前讨论的，每个归约只依赖于矩阵乘法的一部分，因此重新组织归约和矩阵乘法是安全的，我们仍然能得到正确结果。这种重组使我们能够实现计算和通信的重叠：当一些SM在使用张量核心时，其他SM可以进行归约并利用NVLink带宽。这种方法还让我们获得更好的负载平衡。例如，对于注意力计算，一些SM可能提前完成，我们可以为这些SM分配更多工作以减少总体延迟。

总结一下，与当前的逐层内核方法相比，这种巨型内核方法提供了三个关键优势：第一，更好的层间流水线，我们有更好的机会在不同算子之间构建软件流水线。第二，可以重叠计算和通信，因为现在我们捕获了细粒度的数据依赖性，我们知道每个归约片段只依赖于一个矩阵乘法片段，我们可以重新排序它们，在保持正确结果的同时，在不同GPU资源之间获得重叠。第三，更好地支持动态工作负载，因为我们不再需要CUDA图。当所有内容融合到一个内核中时，内核启动开销最小化，我们可以构建一个运行时系统来处理这种动态工作负载。

巨型内核非常理想，因为它将所有内容融合到一个内核中，但这也很具挑战性。第一个挑战是需要管理数据依赖性。原本我们依赖CUDA屏障来处理这种依赖性。现在，没有它，我们需要一种新机制来确保依赖性仍然得到保持。第二个挑战是处理动态性。对于LLM服务，存在许多维度的动态性。例如，使用连续批处理非常普遍，这意味着我们将不同用户的请求批处理在一起，在单次迭代中服务。我们可能有不同数量的请求和每个请求的不同令牌数。有时，我们可能需要混合预填充和解码阶段。使用分页注意力也很常见，有时我们使用推测解码来减少延迟。所有这些方法都引入了不同程度的动态性。最后一个挑战是优化性能。由于大多数现有编译器优化单个内核，但当我们把所有内容融合到一个内核中时，我们关心的是整个内核的端到端性能。

我们使用三种技术来解决这些挑战。首先，我们使用任务图来管理依赖性。其次，我们构建内部运行时系统来处理动态性。最后，我们使用超级优化器来生成高度优化的内核。我们的系统称为MPK，意为Mirage持久化内核。你可以认为我们有一个编译器，输入是LLM架构，输出是高度优化的巨型内核。我们的编译器可以减少工程工作量：用户无需编写数千行CUDA或Triton代码，只需提供模型的Python代码，编译器就能为你生成巨型内核。MPK也提供高性能：我们的内核性能比现有服务系统高1.2到6.7倍。我们还支持新模型：你无需为新模型添加新的CUDA代码。

MPK有两个主要组件：一个编译器和一个运行时系统。编译器以LLM作为输入，以及一些服务配置（如如何进行批处理、分页和推测解码）。编译器将生成一个任务图。一旦我们有了这个任务图，它就可以被MPK运行时用来服务请求。

以下是任务定义。这个任务图交错排列任务和事件。在我们的图中，每个蓝色框是一个任务。一个任务是在单个流式处理器上运行的工作单元。它可以是一个计算任务，例如一个矩阵乘法的分块、注意力的一个头；它也可以是一个通信任务，例如将一块数据传输到另一个SM或另一个GPU。除了任务，每个绿色圆圈称为事件。事件处理任务之间的同步。在图中，从任务到事件的边意味着当任务完成时，我们需要通知该事件。我们还有从事件到任务的边，这意味着当事件被完全触发时（即事件所依赖的所有任务都已完成），我们可以启动所有依赖于该事件的任务。

任务图类似于CUDA图定义，但在CUDA图中，每个节点是在整个GPU上运行的内核。在我们的任务图中，每个节点是在一个SM上运行的任务。任务图是比CUDA图更低层级的表示，它捕获了子内核级别的依赖性，而不是整个内核的依赖性。任务图也是静态且不可变的，我们构建一次，并在服务过程中重复使用。

首先，我们谈谈编译器。MPK包含一个编译器，将原始

27：FlashAttention 4 工作原理详解

在本教程中，我们将深入探讨FlashAttention 4的工作原理。FlashAttention是Transformer模型中注意力计算的核心优化内核，其第四代版本针对NVIDIA Blackwell架构（B200/SM10.0 GPU）进行了重大优化，实现了PetaFLOPS级别的性能。我们将从FlashAttention的发展历程讲起，逐步解析FlashAttention 4的关键创新点，包括其五级流水线设计、更智能的SoftMax缩放策略以及软件模拟指数计算等核心技术。

---

FlashAttention 发展回顾

上一节我们介绍了本教程的概述，本节中我们来看看FlashAttention系列的发展历程，理解其核心思想的演进。

FlashAttention 1：在线SoftMax与分块计算

FlashAttention 1的核心贡献是将分块计算策略应用于注意力计算，并引入了在线SoftMax算法以解决数值稳定性问题。

标准的注意力计算涉及一个序列长度乘以序列长度的大型矩阵（QK^T），随后应用SoftMax，再进行矩阵乘法（与V）。硬件原生实现会创建这个大矩阵，导致巨大的内存开销。

FlashAttention 1的关键在于将计算分解为多个小块（Tile）进行处理。它采用在线SoftMax算法，在计算过程中动态减去当前遇到的最大值，以保持数值稳定，并在后续步骤中对之前的输出进行重新缩放。其核心公式如下：

在线SoftMax公式：
对于每个新的输入块，更新运行最大值 m_new = max(m_old, max(current_tile))，并相应地缩放之前的累加和与输出。

FlashAttention 2：查询并行与性能提升

FlashAttention 2的主要优化是改变了计算的组织方式，采用了查询并行策略。

在FlashAttention 2及之后的版本中，每个线程块（或线程束组）负责处理一个查询块（Query Tile），但需要扫描所有的键（Key）和值（Value）块。这种结构使得输入输出与线程块/线程束的映射关系更加清晰，从而在Ampere架构上实现了更高的计算利用率。

FlashAttention 3：异步性与线程束专业化

FlashAttention 3为Hopper架构（如H100）引入了两项关键技术：异步性和线程束专业化。

• 异步性：允许在Tensor Core执行矩阵乘法或Tensor Memory Accelerator加载数据的同时，执行其他计算，从而更好地隐藏延迟。
• 线程束专业化：在一个线程块内，将线程束分为不同的角色。例如，一些线程束专门负责从全局内存加载数据到共享内存（生产者），而另一些线程束专门负责计算（消费者）。这得益于Hopper架构的动态寄存器分配功能，允许为不同角色的线程束分配不同数量的寄存器。

---

FlashAttention 4 核心创新

上一节我们回顾了FlashAttention系列的基础，本节中我们来看看针对Blackwell架构的FlashAttention 4带来了哪些关键性的革新。

更复杂的五级流水线与线程束角色

FlashAttention 4将异步流水线设计发挥到极致，引入了五个专门的线程束角色，构成了一个复杂的生产者-消费者管道。

以下是流水线中的主要角色：

1. 加载线程束：负责将输入数据（Q, K, V）从全局内存加载到共享内存。
2. 矩阵乘线程束：负责执行QK^T和Attention权重V的矩阵乘法。在Blackwell上，单个线程束即可驱动Tensor Core进行高效计算。
3. SoftMax线程束：负责对注意力分数应用SoftMax操作，并判断是否需要触发重新缩放。
4. 校正线程束：当SoftMax线程束判定需要时，负责对之前的中间结果进行重新缩放校正。
5. 存储线程束：负责将最终的计算结果写回全局内存。

这种设计需要程序员显式地使用内存屏障进行同步，以管理不同阶段对共享缓冲区的访问，避免数据竞争。为了隐藏延迟，每个线程块通常会同时处理两个查询块，使得流水线更加饱满。

更智能的SoftMax缩放策略

FlashAttention 4优化了在线SoftMax算法中的重新缩放逻辑。

在FlashAttention 1-3中，每当遇到新的最大值时，都需要对之前的所有相关结果进行重新缩放，以确保数值稳定。FlashAttention 4引入了一个重新缩放阈值。只有当新最大值与旧最大值的差异足够大，以至于可能影响数值精度时，才触发重新缩放操作。

代码逻辑示意：

float scale = exp2f(old_max - new_max);
if (scale < rescale_threshold) {
    // 触发重新缩放
    trigger_rescale();
}

这种方法可以显著减少重新缩放的次数（根据Hot Chips演示，可达10倍），从而减少了线程束间的协调开销。

软件模拟指数计算

在SoftMax中，指数计算是一个关键但非矩阵乘法的操作，通常由专用的特殊函数单元执行。在Blackwell架构上，SFU可能成为瓶颈。

FlashAttention 4的解决方案是：在部分迭代中，使用CUDA核心进行软件模拟的指数计算，以分担SFU的压力。它采用一个在[0,1]区间内精度很高的三次多项式来近似指数函数的小数部分。

软件指数近似代码片段（概念性）：

// 使用CUDA核心的FMA指令进行多项式计算
// 计算 exp2f(x) 的近似值，其中x已被规范到[0,1]

该实现与硬件指令在BF16精度下是比特位一致的。内核会根据当前迭代的负载情况，动态决定是否启用软件模拟，通常在计算密集阶段使用，而在负载较轻的末尾阶段则切换回硬件指令。

---

性能对比与总结

本节课中我们一起学习了FlashAttention 4的核心机制。作为总结，我们来看一下它的性能表现及其意义。

根据公开数据，在Blackwell B200上，FlashAttention 4的性能大约是前代FlashAttention 3在A100上的两倍，比原始FlashAttention快15倍。其计算利用率达到约30%，处于业界先进水平。与NVIDIA官方cuDNN库中的注意力实现相比，FlashAttention 4约有5%-20%的性能优势。

FlashAttention 4目前仅支持前向传播和BF16精度，反向传播和更低精度的支持仍在开发中。它支持分页键值张量（Paged KV Tensors）和注意力同步（Attention Sync）等特性，使其更适合长上下文和大型模型训练场景。

FlashAttention 4的出现也反映了GPU高性能编程的一个趋势：为了极致压榨硬件性能（尤其是Tensor Core），编程模型正变得更加复杂和专业化。程序员需要更精细地控制流水线、线程束角色和异步操作。这推动了像CUTLASS、Triton等基于Tile的领域特定语言的发展，它们试图在编程便利性和性能之间找到新的平衡点。

28：高性能纯函数式数据并行数组编程

概述

在本节课中，我们将学习Futhark——一种旨在实现高性能数据并行计算的纯函数式编程语言。我们将探讨其设计理念、核心特性，以及编译器如何将高级别的函数式代码转换为高效的GPU内核。课程内容将涵盖Futhark的基本语法、关键限制、编译优化策略，并分析其在现代GPU编程中的定位与挑战。

Futhark语言简介

Futhark是一种纯函数式、数据并行的数组编程语言，属于ML语言家族。它由哥本哈根大学的研究团队开发，最初旨在为金融领域专家提供编写高性能计算模型的能力。Futhark的设计目标是让并行编程变得简单，即使对于初学者也是如此，同时保持硬件无关性，并依赖其编译器为特定硬件（尤其是GPU）生成高效代码。

核心语言特性

上一节我们介绍了Futhark的基本定位，本节中我们来看看其核心的语言特性。Futhark提供了标准函数式语言的特性，但为了性能目标，也施加了一些关键限制。

并行组合子

Futhark内置了一组编译器能够识别和优化的“并行组合子”。这些看起来像高阶函数，但对编译器有特殊的并行语义。

• map: 将一个函数应用于数组的每个元素。其类型签名为 (a -> b) -> [n]a -> [n]b。
• reduce: 使用一个结合性的操作符对数组进行归约。其类型签名为 (a -> a -> a) -> a -> [n]a -> a。
• scan: 执行前缀扫描操作。
• 其他: 还包括 scatter, histogram 等少量原语。

用户程序最终由这些原语组合而成，编译器则基于它们进行并行化推理和代码生成。

关键限制

为了实现高性能并简化编译器设计，Futhark引入了几项在传统函数式语言中不常见的限制。

• 无递归: 禁止通用递归，以防止生成无法预测性能的代码。迭代通过特殊的尾递归语法糖（loop 结构）实现。
• 受限的高阶函数: 函数不是完全一等公民。具体规则包括：
  • 条件表达式（if）不能返回函数。
  • 不能创建函数数组。
  • 这些规则确保编译器总能静态确定被调用的是哪个具体函数。
• 规则数组: 多维数组必须是“矩形”的，即所有子数组在给定维度上大小相同。这保证了内存布局的可预测性。

值与数组表示

Futhark采用按值传递，并积极拆解元组等复合结构，将各部分作为独立值处理，便于存入寄存器。

对于数组，编译器会自动应用 结构体数组到数组结构体（AoS到SoA）的转换。例如，一个类型为 [n](i32, i8) 的数组，在内部会被表示为两个独立的数组：[n]i32 和 [n]i8。这确保了内存紧凑和对齐，是高性能计算中的标准优化。

反函数化

为了高效实现高阶函数，Futhark编译器使用 反函数化 技术。它将每个闭包（lambda表达式）转换成一个包含其自由变量的记录，并将函数调用点替换为对特定“提升函数”的直接调用。结合上述的类型限制，这最终能消除所有动态函数派发，生成直来直去的代码。

从高级代码到GPU内核

了解了语言特性后，我们来看看Futhark编译器如何将这些高级概念映射到GPU硬件上。

基础映射模式

编译器将核心并行组合子直接映射到抽象的GPU内核模式。

• map -> segmap_thread: 一个对 [n]a 的 map 可以转化为一个启动 n 个线程的内核，每个线程处理一个元素。
• 嵌套 map -> 多维 segmap_thread: 完美嵌套的 map 可以转化为多维线程网格。
• reduce -> segred_thread: 归约操作被转化为一个实现了高效并行归约算法的内核。
• 分段操作: 对于 map 内部包含 reduce 的情况（如批处理归约），编译器能识别并生成相应的 segred_thread 内核。

扁平化与多版本代码生成

实际问题中的 map 体可能很复杂，不直接符合上述简单模式。例如，一个对矩阵每行先求最大值再归一化的操作。

map (\row ->
  let max_val = reduce max (-f32.inf) row
  in map (/ max_val) row
) matrix

编译器采用 扁平化/裂变 变换，将复杂嵌套操作切割成多个简单的内核。对于上述例子，它可能生成多个实现版本：

1. 完全顺序化: 仅并行化外层，内层串行执行。
2. 完全并行化: 将内层操作也并行化，生成两个独立内核（一个归约，一个映射），中间结果写回全局内存。
3. 利用共享内存: 以一个线程块处理一行，在共享内存中协作完成该行的归约和映射。

编译器会根据运行时数组大小（通过启发式规则和自动调优确定的阈值）动态选择最佳版本。这虽然可能增加代码体积，但能适应不同规模的数据。

内存布局优化

编译器会尝试优化数组的内存布局以提高访问效率。例如，如果一个内核按列顺序访问一个默认行优先存储的二维数组，导致非合并访问，编译器可能会自动插入一个转置操作，使后续访问变为合并访问。这种布局优化是保证模块化编程性能的关键。

Futhark的定位、应用与挑战

定位与比较

Futhark并非通用编程语言，而是用于编写计算库的领域特定语言。它通过外部函数接口（如C、Python）与主程序交互。

它也不是专用的GPU语言，而是一种具有良好GPU编译器的硬件无关语言。与Triton等底层内核语言相比，Futhark的目标不是达到极致的性能，而是在保持高级别、声明式编程风格的同时，获得显著优于手写CPU代码且可媲美非专家手写GPU代码的性能。其核心优势在于 优化组合：用户可以编写小型、模块化的组件，编译器会融合它们并消除中间存储开销。

应用案例

尽管是学术语言，Futhark已被用于一些实际项目，特别是在需要表达复杂并行算法且对性能有要求的领域，如：

• 计算流体动力学。
• 粒子模拟中的特殊神经网络。
• 其内置的自动微分功能对于机器学习原型开发尤其有吸引力。

当前局限与未来挑战

Futhark对现代GPU特性的利用尚有不足：

• 缺乏异步操作: 不支持内核间计算与通信的重叠。
• 未利用特殊硬件指令: 如Tensor Cores，因为编译器缺乏领域知识（如“这是一个矩阵乘法”）来触发它们。
• 有限的Warp级编程: 仅在硬编码的原语（如归约）中使用，未泛化到用户代码。
• 无透明分布式支持: 随机的数组访问模式使得透明的数据分布和通信变得困难。

这些是未来研究的主要方向。

总结

本节课中我们一起学习了Futhark编程语言。我们了解了它作为纯函数式数据并行语言的设计哲学，包括其核心的并行组合子、为确保性能而引入的关键限制（无递归、受限高阶函数），以及反函数化等实现技术。我们深入探讨了编译器如何通过基础映射、扁平化、多版本代码生成和内存布局优化，将高级别代码转换为高效的GPU内核。最后，我们讨论了Futhark在实践中的定位、应用场景以及面临的挑战。Futhark代表了一种让并行编程更接近普通函数式编程体验的探索，尽管存在限制，但它为编写高性能、可组合的计算内核提供了一个独特而强大的工具。

延伸资源

• 《Futhark书》: 学习Futhark编程和并行算法的入门指南。
• GitHub仓库与博客: 包含大量示例、项目以及语言发展历程的详细记录。
• 示例项目: 如用Futhark实现的GPT-2模型，展示了其尺寸类型系统的应用。

29：Helion：一种用于机器学习内核的高级领域特定语言

概述

在本节课中，我们将学习 Helion，这是一种用于编写高性能 GPU 内核的新领域特定语言。Helion 旨在提供比 Triton 等现有 DSL 更高的抽象级别，同时比 PyTorch 更低，从而让开发者能够更精细地控制生成的内核，而无需处理底层硬件细节。我们将探讨其设计动机、核心语言特性、自动调优机制，并通过实际示例展示如何编写 Helion 内核。

---

1：设计动机与核心理念

上一节我们概述了 Helion 的目标，本节中我们来看看为什么需要这样一个新的 DSL。

Helion 的诞生源于用户对 torch.compile 提出更多控制需求。用户希望精确控制生成的内核，这需要一个 DSL。Helion 的定位是提供一个抽象级别更接近 PyTorch 的 DSL，让开发者能够利用熟悉的 PyTorch 概念，而无需担心特定内核的底层细节。

Helion 实现高性能的关键在于集成了自动调优器。其核心思想是：一个 Helion 内核定义了一个巨大的搜索空间，自动调优器可以从中搜索数千个不同的 Triton 内核实现。这带来了两大优势：

1. 节省人力：无需人工尝试数百种内核实现方式，可以让自动调优器代劳，用机器算力换取人力。
2. 硬件可移植性：用低级语言编写的内核通常在新硬件发布时需要重写。而高级语言结合自动调优器，只需为新硬件重新调优即可获得更好的性能，从而提升跨硬件平台的移植性。

---

2：Helion 语言基础：以矩阵乘法为例

上一节我们介绍了 Helion 的设计理念，本节中我们通过一个矩阵乘法的例子来看看 Helion 代码的基本结构。

一个 Helion 内核主要包含两部分：

• 主机端代码：位于 @hl.kernel 装饰器定义的函数中，但在 hl.tile 循环之外。这部分是常规的 PyTorch 代码，用于分配输出张量等，将在即时执行环境中运行，不会被编译。
• 设备端代码：位于 hl.tile 循环内部。这部分代码将被编译成单个 Triton 内核。

以下是一个矩阵乘法的 Helion 示例：

import torch
import helion as hl

@hl.kernel
def matmul(A: hl.Tensor, B: hl.Tensor) -> hl.Tensor:
    # 主机端代码：分配输出张量
    M, K = A.shape
    K, N = B.shape
    C = torch.empty((M, N), device=A.device, dtype=A.dtype)

    # 设备端代码：使用 hl.tile 定义并行化网格
    for i, j in hl.tile((M, N), (BM, BN)):
        # 在 tile 内使用标准的 PyTorch 风格操作
        tile_C = torch.zeros((BM, BN))
        for k in range(K):
            tile_A = A[i, k]
            tile_B = B[k, j]
            tile_C += tile_A @ tile_B
        C[i, j] = tile_C
    return C

核心概念 hl.tile：
hl.tile 循环将给定大小的迭代空间划分为多个块。其具体划分方式（如块大小、迭代顺序）由自动调优器决定。你可以将其视为“带块的 PyTorch”。在 hl.tile 循环内部，可以使用标准的 PyTorch 操作，Helion 也支持控制流等其他结构。

---

3：自动调优的配置空间

上一节我们看到了 hl.tile 的基本用法，本节中我们深入探讨自动调优器可以优化的配置空间，这也能帮助我们更好地理解 Helion 语言的特性。

自动调优器在多个维度上进行搜索，以下是一些关键的配置选项：

索引模式
在 Triton 中，索引有指针运算、块指针、张量描述符等多种方式。在 Helion 中，你只需编写常规的 PyTorch 风格索引，自动调优器会自动在所有这些方式中选择最佳的一种。

块大小
当编写 hl.tile 循环时，你隐式地定义了一些块大小。对应的 Triton 代码需要大量样板代码来计算网格大小等，而 Helion 移除了这些样板代码。通过高级方式表达分块，Helion 还能实现更复杂的优化，例如将 2D 迭代空间扁平化为 1D 迭代空间。

归约循环
如果你在 Helion 中编写一个归约操作，可以以直观的方式表达。自动调优器会自动在“一次性加载整行”和“使用累加器迭代”这两种实现方式之间进行转换和选择，以优化寄存器使用和性能。

程序 ID 与启动网格
Helion 可以自动调优不同的启动网格策略，包括一维/二维启动网格，以及持久化启动网格（每个流多处理器运行一个 CUDA 程序，通过虚拟程序 ID 迭代，支持跨 CUDA 程序的流水线）。

循环优化
自动调优器可以优化循环分组和重排序，例如交换循环顺序或对迭代空间进行子划分以提升缓存复用。

配置标志直通
Helion 将许多 Triton 的配置标志（如 num_stages, warp_specialization, num_warps 等）暴露出来，并自动为其调优。

自动掩码处理
掩码处理可能很棘手，Helion 为你自动化了许多掩码处理逻辑，并且在静态形状情况下会将其优化掉。

---

4：自动调优器的工作原理

上一节我们了解了丰富的配置选项，本节中我们来看看自动调优器是如何工作的。

目前的自动调优算法相对直接，但效果显著：

1. 初始种群：随机选择 100 个配置。
2. 选择：从中选出前 5 个最快的配置。
3. 爬山法：对每个最快配置进行局部搜索，寻找局部最优解。
4. 确定最终配置：选择所有局部最优解中最快的一个。

经验表明，对于大多数内核，这个搜索空间相当规则，因此该算法能取得良好效果。一次完整的调优大约需要 20 分钟来搜索数千个候选内核。未来我们计划探索更复杂的算法，如大语言模型引导搜索、强化学习或共享性能数据库。

配置部署方式

• 单一配置：将自动调优器输出的最快配置复制到代码中，后续运行将直接使用此配置。
• 多配置：针对不同形状（如小形状和大形状）分别调优，得到多个配置。内核在首次遇到新形状时会尝试所有配置并缓存最快的结果。
• 手动路由：可以提前将不同配置编译成不同的运行器，然后编写简单的路由函数来根据条件（如张量大小）选择使用哪个运行器。

---

5：性能表现与内部原理

上一节我们介绍了调优过程，本节中我们看看 Helion 的实际性能表现，并简要了解其内部编译器原理。

性能对比

• 与 Quack 比较：在 RMSNorm 等操作上，Helion 和其等效的 Triton 实现在所有形状大小上都表现良好，获得了更好的性能可移植性。
• B200 上的加速：与 eager 模式的 PyTorch 相比，Helion 平均可获得 3.2 倍的几何平均加速。与手工编写的 Triton 内核相比，Helion 几乎快一倍，这主要得益于自动调优器能更好地优化 Triton 代码。
• AMD 上的表现：与 B200 上的情况类似，Helion 在大多数情况下都是最快的。

编译器内部原理
Helion 作为 Python 元 DSL，其编译流程如下：

1. Python AST：从 Python 抽象语法树开始。
2. 类型注解：用类型和元数据注解 AST。
3. FX 图转换：将其转换为多个 FX 图（每个基本块一个图）。每个 FX 图节点都附有 Inductor IR。
4. 编译器传递：在这些 IR 上运行编译器传递。
5. 代码生成：仅在代码生成阶段融入自动调优器确定的配置。这意味着配置点左侧的所有步骤只需运行一次，右侧步骤在每次调优时重新运行。
6. 输出：最终生成 Triton 代码。对于调用内核的包装代码，Helion 使用 AST 重写将 for 循环转换为内核启动。

---

6：快速入门与基础示例

上一节我们从宏观了解了 Helion，本节中我们开始动手，学习如何安装 Helion 并编写一些简单内核。

安装与文档

• 安装：pip install helion。同时需要安装 PyTorch 2.9 和 Triton 3.5。
• 文档：访问 helion.com 获取详细文档和 API 参考。

示例1：逐元素加法内核
这是最简单的内核之一，展示了主机端与设备端代码的划分：

@hl.kernel
def add(A: hl.Tensor, B: hl.Tensor) -> hl.Tensor:
    # 主机端代码
    output = torch.empty_like(A)
    # 设备端代码：tile 循环
    for i in hl.tile(A.shape):
        output[i] = A[i] + B[i]
    return output

重要抽象：一个 Helion 内核保证会被转换为单个 Triton 内核，设备端的所有内容将被融合到一个内核中。

示例2：矩阵乘法与 Epilogue
这个例子展示了如何向内核传递一个 lambda 函数作为 epilogue（后处理）：

@hl.kernel
def matmul_with_activation(A: hl.Tensor, B: hl.Tensor, activation) -> hl.Tensor:
    M, K = A.shape
    K, N = B.shape
    C = torch.empty((M, N), device=A.device, dtype=A.dtype)
    for i, j in hl.tile((M, N)):
        accumulator = ... # 矩阵乘计算
        # 应用 epilogue 激活函数
        C[i, j] = activation(accumulator)
    return C

调用时，可以传入 torch.relu 等函数，生成的 Triton 内核将包含对应的激活操作。

示例3：带掩码的张量连接
当操作不需要覆盖整个张量时，可以使用掩码：

@hl.kernel
def concat_2d(A: hl.Tensor, B: hl.Tensor, dim: int) -> hl.Tensor:
    out_shape = ...
    output = torch.empty(out_shape, device=A.device)
    for idx in hl.tile(output.shape):
        # 加载时使用掩码防止越界访问
        tile_A = A[idx] if idx[dim] < A.shape[dim] else 0
        tile_B = B[idx] if idx[dim] >= A.shape[dim] else 0
        # 合并数据
        output[idx] = torch.where(condition, tile_A, tile_B)
    return output

---

7：高级技巧：以 RMSNorm 为例

上一节我们看了几个基础内核，本节中我们以 RMSNorm 为例，探讨一些更高级的 Helion 编程技巧。

RMSNorm 的简单实现可能直接使用 torch.mean 和 torch.rsqrt。虽然方便，但在某些情况下可能效率不高，因为它需要加载整行数据。

高级实现技巧

1. 注册块大小：当内核中有多个内部归约，且需要确保它们使用相同的块大小时，可以使用 hl.register_block_size 预先注册一个块大小。这会告知自动调优器，在使用此块大小时，必须保证所有 tile 使用相同的块大小。

   block_size = hl.register_block_size("reduce_dim")
   

2. 驱逐策略：对于会被重用的数据，可以指定更保守的驱逐策略（如 evict_policy=hl.EvictPolicy.KEEP），而对于只使用一次的数据，可以指定立即驱逐（evict_policy=hl.EvictPolicy.EVICT）。这类似于 Triton 中的做法，允许你更精细地控制缓存行为。

关键点：对于简单的 RMSNorm，自动调优器会自动做出这些优化决策。但在你需要更明确控制内核生成方式的情况下，可以手动使用这些高级功能。

---

8：调试与开发体验

上一节我们编写了更复杂的内核，本节中我们学习如何调试和优化开发体验。

Helion 提供了多种环境变量来辅助调试：

输出生成的 Triton 代码
设置 HL_PRINT_OUTPUT_CODE=1。这会在 Helion 生成 Triton 代码后将其打印出来，同时还会输出一个易于复现问题的脚本。

关闭或减少调优开销

• HL_AUTOTUNE_EFFORT=none：完全关闭自动调优，使用默认配置，适用于快速测试正确性。
• HL_AUTOTUNE_EFFORT=fast：进行快速调优，寻找“足够好”而非最优的结果。

解释执行模式
设置 HL_INTERPRET=1。这不会通过 Triton 后端运行内核，而是将设备端代码转换为 FX 图可追踪的代码，并使用 PyTorch eager 模式在主机上执行整个区域。适用于调试数值问题。

启用详细日志
使用 HL_LOG=all 或 HL_LOG+=all 来启用调试级别的日志，输出内部诊断信息以帮助识别问题。

编程体验说明

• 主机端代码：理论上可以执行任何操作，因为它运行在 eager 环境。但涉及张量的元操作会以 eager 方式执行并产生警告。
• 设备端代码：支持大部分操作。print 语句会被自动转换为 Triton 的 tl.print。断点支持可能有限。
• 问题反馈：如果遇到问题，欢迎在 GitHub 仓库提交 issue。

---

9：实战：编写注意力机制内核

上一节我们关注于调试，本节中我们通过一个复杂的实战案例——注意力机制内核，来综合运用所学知识。

注意力机制的计算步骤包括：QK^T 矩阵乘、缩放、softmax、与 V 的矩阵乘。

核心设计决策
编写 Helion 内核时，最关键的是决定哪些维度可以并行化，哪些维度必须顺序执行。

• 对于注意力内核，只有 softmax 操作所需的归约维度（序列长度 n）需要顺序执行。
• 其他维度（如批大小 batch、头数 num_heads、查询长度 m）都可以并行化。

内核结构

1. 外层并行化：最外层的 hl.tile 在 batch * num_heads 和 m 维度上进行并行化，这对应了 Triton 内核的启动网格。
2. 内层顺序执行：在内层，使用另一个 hl.tile 循环顺序遍历 n 维度。
3. 计算步骤：在内层循环的每个 tile 中，执行 QK 乘、softmax 计算、与 V 乘，并更新 softmax 所需的运行最大值逻辑。
4. 写入结果：循环结束后，完成 softmax 并写入输出张量。

优势
Helion 帮你处理了计算偏移、形状对齐等繁琐的细节，让你可以专注于更高层次的并行化策略设计。

查看输出代码
使用 HL_PRINT_OUTPUT_CODE=1 可以查看生成的 Triton 内核，其中包含了所有必要的启动网格计算和索引样板代码，方便调试和验证。

硬件特定优化
如果需要使用最新的硬件特性（如 Blackwell 的特定指令），Helion 支持在设备端代码中直接内联 PTX 汇编字符串。

---

10：组合性、与 PyTorch 生态的集成及未来展望

上一节我们完成了注意力内核的实战，本节中我们探讨 Helion 如何与现有生态集成以及未来的发展方向。

与 Triton 和 PyTorch 的组合性
Helion 遵循“一个 Helion 内核即一个 CUDA 内核”的不变量，这与 Triton 相同。因此，你可以在一个 PyTorch 程序中混合使用 Helion 内核、Triton 内核和普通的 PyTorch 操作，它们通过输入输出进行交互，互不干扰。

与 Torch Compile 的集成
Helion 内核已经可以与 torch.compile 组合使用。目前，Helion 内核被编译为 Triton 内核，然后像自定义 Triton 内核一样被送入编译流程。未来，我们计划实现 Helion 内核与 torch.compile 更深入的融合，例如 prologue/epilogue 融合。

底层实现共享
Helion 复用了 Torch Inductor（torch.compile 的后端）的许多基础设施来生成 Triton 代码。因此，你可以在 Helion 设备端代码中直接使用大多数 PyTorch 逐点和归约操作，它们会通过相同的 Inductor lowering 流程被处理。这避免了 Helion 需要维护庞大的操作库。

未来方向：MegaKernel 支持
目前，Helion 支持在单个内核中有多个顶级的 hl.tile 循环（foreach 风格）。未来，通过引入同步操作（如 barrier），可以更轻松地构建将整个模型计算放在单个内核中的 MegaKernel。这对于计算量很小的模型尤其有益，可以减少内核启动开销。

---

总结

在本节课中，我们一起学习了 Helion 这一用于机器学习内核编写的高级 DSL。我们从其设计动机和核心理念出发，了解了它如何通过高层抽象和强大的自动调优来平衡开发效率与性能。我们学习了 Helion 的基本语法结构，包括主机端与设备端代码的划分，以及核心的 hl.tile 构造。我们深入探讨了自动调优器的工作原理和丰富的配置空间。通过从简单的加法到复杂的注意力机制等多个示例，我们实践了如何编写和调试 Helion 内核。最后，我们还了解了 Helion 与 PyTorch 生态系统的集成以及未来的发展潜力。希望本教程能帮助你开始使用 Helion 来编写高性能、可移植的 GPU 内核。

30：形式化内核推导 🧠

概述

在本节课中，我们将探讨深度学习研究中的一个核心问题：尽管深度学习模型是纯粹的数学系统，我们却缺乏分析它们的正式语言。这导致从数学描述到高效底层实现（如Flash Attention）的推导过程漫长且困难。我们将介绍一种基于范畴论的正式数学系统，用于描述深度学习模型，并展示如何将其应用于系统化的内核推导和性能建模。

---

研究背景与动机

上一节我们概述了课程目标，本节中我们来看看当前深度学习领域面临的挑战。

深度学习模型是数学函数。每一层的输出是前一层输入的函数。然而，我们无法直接使用一个“神谕”来实现这些函数，原因在于我们必须面对的计算限制。

深度学习模型的设计与其实现任务密不可分。如果这一点不成立，那么“神谕”就足够了。AlexNet是第一个真正利用GPU并行计算的网络。Transformer紧随其后，它能够以更易于计算处理的方式，通过多个注意力头并行处理所有数据，完成与RNN类似的过程。专家混合模型则是一种利用多GPU系统优势来编码前馈层的方法。

因此，深度学习架构的基本设计挑战并不一定是拥有纯粹的数学表示，而是拥有一个与执行密切相关的数学运算。然而，推导这些底层优化（甚至包括决定如何在GPU间分配计算的高层策略）是困难的。

Flash Attention花了五年时间才被推导出来，又花了两年时间才适配Hopper架构。DeepC在2024年发布，它实现了一个相对标准的Transformer模型，但通过更好地协同设计软件与硬件，实现了惊人的性能水平。

一个关键事实是，这些算法的性能优势本就存在于当时的硬件中，存在于现有算法附近的设计空间里，但我们却花了这么长时间才推导出它们。

那么，为什么这个问题如此困难？因为深度学习模型本质上是数学函数。这意味着存在一个数学关键点，一个这些事物存在的“原因”，它存在于现有的数学之中。

当我们处理深度学习算法时，我们从数学描述走向算法，最终走向底层实现。但在那个数学描述中，必然存在某种原因，使得那个底层实现成为可能。我们如何才能发现这种关系？如何轻松地从数学函数走向底层实现？

挑战很大程度上在于我们缺乏正式的表示方法。目前，我们经常使用线性代数表达式。但它们不足以覆盖深度学习中的非线性操作。例如，在注意力算法的表达式中，softmax发生在哪里？连接操作发生在哪里？这些变量的维度是多少？一个操作如何与另一个操作的广播相关联？

鉴于这些标准的、临时性的线性代数或临时性的图表表示，很难看出它们如何转化为我们追求的底层实现。这就是推导这些算法如此耗时的主要原因。

因此，答案是：我们需要一种用于深度学习模型的正式形式体系，它可以解释为什么像Flash Attention这样的方法是可能的。这样的形式体系对于其他形式的分析也有好处，但在本次讲座中，我将专注于底层内核推导，因为这是一个图表能提供巨大优势的领域，并且许多数学问题已经得到解决。

---

解决方案：范畴论

上一节我们探讨了问题的核心，本节中我们来看看提出的解决方案：范畴论。

范畴论是抽象和组合的数学。它研究由对象和态射通用构建的组合系统。对象是锚点，态射是链接。因此，它可以描述图、相互作用的函数族、形式化的图表，甚至算法。

范畴论的一个著名结果是Yoneda引理。范畴论的语言可以描述系统，然后提供关于某些性质的证明。范畴论允许构建这些图表，并在其中编码代数规则，从而推导出复杂的关系。

对于许多听说过范畴论的人来说，可能会担心它有时以无尽的抽象而闻名。那么，我们为什么要尝试应用范畴论呢？

我喜欢这样思考：像群论这样的东西，就其本身而言，可能极其抽象，专注于现实世界中不一定存在的事物。它是一堆定理的大杂烩，或者像黎曼几何那样。但考虑到像粒子物理学这样的领域，其问题的本质是所有这些相互作用的对称性，群论就成为了处理它的恰当且正确的语言。这使得群论这个抽象的数学领域能够作为系统的表示，并提供与之相关的定理。然后，我们发展出的粒子物理学数学、它对量子物理学的理解，使得在现实世界中的各种实现成为可能。

因此，我们看到，一个抽象的数学系统，在遇到具有该系统所描述特征的问题时，可以从高度抽象转变为高度应用。对于范畴论也是如此。

深度学习模型本质上是高度组合的系统。我们关心多个层次的分析：功能层次、实现层次、概率流层次，所有这些相互作用的层次。因此，范畴论作为组合系统及其之间关系的研究，变得相关，并成为描述这些事物的相关工具。

因此，我们可以从数学走向一种表示形式。范畴论提供了我们谈论这些系统所需的语言，以及进一步描述它们的工具和定理。我们看到，这种表示然后也可以走向实现。

---

范畴论基础：积范畴

上一节我们介绍了范畴论作为解决方案，本节中我们来看看其基础概念：积范畴。

一个范畴，具体来说是一个积范畴，我们如何构建它？

• 对象：对象是锚点，它们决定组合。对于函数可以是集合，对于图可以是节点。我们将以特定方式绘制对象：画一条带箭头的线并标记它。
• 态射：态射连接对象。我们用边连接节点，用函数连接集合。态射是任何具有这种形状的东西：从A到B，或者从A到A。
• 生成元：为了描述一个范畴，我们会有一些生成态射。从这些生成态射出发，构建其他一切。这在技术上并非必需，但对于表示范畴、实现它等非常有帮助。
• 组合：我们可以组合态射。当我们组合态射时，我们连接它们，并且中间对象必须相同，这就形成了一个新的态射。你会注意到这个形状模仿了那个形状，因此我们形成了一个态射。

在积范畴中，我们还可以堆叠事物。

• 积：积可以将对象组合成新对象，将态射组合成新态射。组合对象，例如，是将数据类型组装成数据类型的二元组。在这种情况下，取函数的积就是并行运行这些函数。因此，当我们堆叠函数时，我们也堆叠它们的对象。并行执行的函数分别在两个输入上独立执行。使用这样的图表，我们可以表示这一点以及更复杂系统是如何构建的。

接下来，我们有重排。

• 重排：重排有一个非常好的图表表示，但起初可能不直观。它表示像转置、复制和删除这样的操作。当我们转置、复制或删除数据时，我们并没有操纵数据的内容，我们只是在重新排列它。当我们删除数据时，我们进行的是遗忘性重排。当我们转置数据时，重排只是将0映射到1，1映射到0。当我们复制数据时，我们重新排列它使得两个东西相同。实际上，我们可以以统一的方式表达重排。

我们有一个从输出（可以认为是之后我们将拥有的变量名）到输入（之前我们拥有的变量名）的某种函数。然后，如果我们有一些变量，我们现在可以构建这个重排。你会注意到这个重排的顺序：输出的名称在前，输入的名称在后。而在另一种表示中，顺序是相反的。你会发现这是正确的做法，没有其他方式，这种模式将非常常见、非常重要且极其有用。

因此，我们有了对象、态射、组合，我们可以取积或进行重排。

---

应用：重新表达注意力机制

有了这个框架，我们可以开始重新表达像注意力机制这样的东西。

我们可以从这个临时性的图表走向更形式化的东西。我们有Q和K，我们标记对象（集合），标记态射（函数），然后逐步推导。你可能注意到这与我们直观上会做的非常接近，它匹配数据的流动。范畴论是一种形式化我们经常已经在观察的事物的方式。

但由于其严谨的结构，它允许发展扩展。那么，对于这个类别，我们需要什么扩展呢？右边的图表帮助不大。我们已经讨论过为什么其中存在结构是这个图表没有显示的。

因此，我们希望开发一种显示该结构的方法。为此，我们将开发一个特定的范畴。

---

数组广播范畴

上一节我们展示了如何用范畴初步表示注意力，本节中我们构建一个更强大的工具：数组广播范畴。

这个范畴的对象将是数组（或数组的数组）。其态射将是数组之间的函数。我们将开发特定的基础设施，以清晰地表达广播。

• 尺寸范畴：我们从尺寸范畴开始。尺寸范畴表示数组可能采用的坐标。它对应于数组的尺寸。因此，对象就是尺寸。那么积是什么？积是将单个尺寸组装成多维尺寸。因此，我们有了尺寸的基础设施。我们很快会展示这个范畴中的态射是什么。

那么，数组呢？我们可以使用这些尺寸来定义数组。一个数组是某个基础数据类型（如实数、浮点数等）和一个尺寸。因此，我们基本上是说，对于尺寸中存在的每个坐标，我都有基础数据类型的一个实例。在右边，我们有我们可以称之为标准表示的东西，其中有这些网格，它们代表数组中的值。但我们可以开发图表。

我们如何做到这一点？我们画一个带有基础数据类型的箭头（底层数据类型），然后在上面画出尺寸。这里，A是尺寸范畴中的一个对象，并且尺寸范畴的对象可以分解为其他对象的积。因此，我们可以通过这个表达式或那个表达式来表示这个数据类型，它们是相同的，因为这两个是相同的。注意，我们保留基础数据类型的箭头，但移除顶部的箭头，这样做的原因只是为了不将它们彼此混淆。

• 态射：对象到对象，很简单。任何在中间的东西构成一个态射。因此，这个x可能是那个，y可能是别的，或者我们可以使用系统来表达这个范畴。
• 积：积是笛卡尔积，事物并行进行。因为我们使用没有任何分隔的堆叠来表示一个事物，我们开发了这个虚线系统来区分不同的积。这可以描述数组的元组以及完全独立并行运行的函数。因此，如果我们有这个函数和那个函数堆叠在一起，它将是一个接受元组并产生元组的函数。

这就是数组广播范畴。这意味着我们可以将这个重新表达为那个。这更有用一些。我们看到这是流经系统的所有事物的尺寸，我们可以更好地跟踪事物，但这里仍然缺少很多结构。那个结构就是内核推导的黄金所在。

目前，我们可以做一个美学调整：使用象形图和名称来表示关键操作。我们使用杯状符号表示矩阵乘法（类似于物理学中的彭罗斯图）。Softmax将是一个三角形，它是一个非常重要的操作，所以象形图是有意义的。

现在，我们如何引入其余的结构？如何显示这个Q和那个Q之间的关系等等？

---

提升与重索引基础设施

上一节我们构建了数组广播范畴，本节中我们来看看其核心操作：提升与重索引。

左边图片的问题在于我们没有这个提升信息。这个提升信息对于像Flash Attention这样的东西至关重要，因为最终Flash Attention正是发生在这个Q轴上。

我们可以开发图表来清晰地显示它。例如这个图表，它显示了这个Q与那个Q之间的关系，等等，从操作结束一直到开始。

那么，我们做的这件事，这个“提升”，到底是什么？让我们从这里开始。

我们从一个从左边的数组到右边的数组的操作开始。使用我们的图表工具，你可以这样表示它：输入、输出用尺寸标记。

当我们提升某个东西时，我们让F在一个更大数组的行上作用。所以，我们有F作用在第一、第二、第三、第四行等。这会是什么样子？这将看起来像一个函数作用在P·A上并产生E。我们可以进行这个巧妙的转换：为了表示F在P上广播，我们只需将P轴画在F的顶部。这以正确的方式修改了输入输出形状，并清晰地显示了两者之间的关系。这就是提升。

我们如何提供一个更正式的定义？因为我们想要使用它，我们不希望提升只是一个存在的东西。我们想要展示提升如何与像Flash Attention或我们想要对深度学习系统采取的其他分析方法相关联。

我们定义提升如下：我们使用索引，或者更准确地说，切片。切片就像在PyTorch中提取特定的列或行。在左边这种情况下，我们提取表达式的第四行。我们知道这第四行是由F作用在输入的第四行上生成的。类似地，我们可以这样绘制切片：在i轴上的切片（就像在PyTorch中一样）会减少那个轴，然后留下剩余的轴，所以这将是B。

因此，左边可以这样表示：我们有F提升到P上，然后之后应用那个切片。这等于什么？它等于F作用在输入的第四切片上。类似地，我们可以用i这样表示它，这个切片移动到表达式的前面。毫无疑问，这个第四切片对应于那个。因此，这里的对应关系具有这个图表表示。

这开始显示图表具有一些力量：不仅这个P轴正确地指示了输入和输出类型，它还允许我们开始执行代数运算。这就是我们如何定义一个提升的表达式。

这里，我们有F^{P，假设我们不知道它是什么，但我们可以说，给定一个在i上的切片，F}P将等于该切片应用到输入上，然后函数F作用在其后。因此，我们可以使用F和切片的知识来定义F^P做什么。

我们有了这个切片自然性属性：表达式后的切片可以移动到表达式前，两者在数学上是等价的。

在我们的注意力模型中，我们有Softmax在Q上广播，这意味着什么？意味着表达式后的切片变成表达式前的切片。对于逐元素表达式，我们有轴从两边穿过。这意味着，如果我询问输出的第i,j坐标是什么，它等于函数应用到输入的第i,j坐标上。

注意：这是尺寸0的积。当我们将基础数据类型提升到尺寸0的积上时，我们得到一个项，它是一个单独的值。另外，你可能注意到我们在这里没有太多地写出基础数据类型。我们可以保持它隐含，对于本次分析的大部分内容，它不直接相关。

现在我们有了切片自然性。让我们回到尺寸范畴：态射是什么？我们可能想应用到轴尺寸上的可能态射是什么？仿射变换。这个定义相当多。我们真正需要知道的是，我们将自己限制在坐标的仿射变换上。这将变得非常重要，因为它允许推导出一些非常好的启发式方法。

---

重索引

在提升之后，我们有重索引。重索引是指我们进行了提升，但输出切片并不对应于完全相同的输入切片。例如，第i个输出切片可能对应于第a(i)个输入切片。我们有这种模式。

那么，这看起来像什么？这一行产生这一行，这一行产生那一行。我们如何清晰地定义这个？我们说，第i个输出对应于第a(i)个输入。例如，a将0映射到0，将1映射到2。这可能是卷积网络中的膨胀操作。

我们如何用图表表示？我们这样绘制它。我们知道这里的输入是Q·I，所以我们可以把Q·I画下来。输出是P·V，所以我们可以把P·V画在下面。然后我们有一个作用在其上的切片，所以这个切片是i。然后我们知道这应该等于那个。所以这里我们有第a(i)个切片，然后它产生A，F可以作用其上。那么，我们在这里放什么？我们可以放η。所以，如果我们有一个表达式上面有A，它从P映射到Q。这里我们有P到Q。我们在这里放η，所以顺序是向上的。这意味着i有效地向左，然后与a(i)交互以创建这个切片，该切片可以移动到表达式后面。所以我们有了这个，我们知道F是什么，我们知道这个是什么，它只是一个切片。我们说这等于那个。这就定义了一个重索引。

那么，我们能对重索引做什么？为什么我们要费这么大劲来定义它们？

---

应用重索引：转置与矩阵乘法

上一节我们定义了重索引，本节中我们来看看它的具体应用。

以转置为例。转置是什么？转置意味着输出的第i,j坐标等于输入的第j,i坐标。这两个交换了位置。之前，我们将重排定义为这些积范畴的基本组成部分，这里我们使用它们。我们重新排列索引以进行交换。然后我们将其与重索引提升一起使用来定义转置是什么。这是一个例子，我们可以使用这个重索引术语和这个重排术语一起来表示转置，以表明转置只是这个东西的一个特定实例。

此外，像矩阵乘法这样的东西更复杂，但根本上，输出的第i,j坐标等于第一个输入段的第i切片和第二个输入段的第j切片。这个从i,j到i的操作是什么？是删除？投影到第一维。因此，这也可以表示为重排，第二个也是如此。所以我们可以使用重排和重索引一起来定义这些操作。

但我们也可以以稍微不同的方式来做。对于矩阵乘法，我们可以说输出的第i,j坐标等于这个切片重排。用图表，这可以清晰地显示：我们展示j如何关联到X，i如何关联到Q。这里，投影到这些段中的每一个的重排是隐含的，我们基本上可以信任图表来向我们展示事物是如何流动的。

---

重新表达注意力机制

现在我们已经构建了这个基础设施，我们能用它做什么？我们可以将注意力机制重新表达为一系列提升和重索引，得到这个表达式。

在这里，我们可以看到我们有Q和D，这个轴一直穿过。这意味着什么？让我们看看。我们从表达式开始，展示广播，然后取一些索引看看它们做什么。穿过，穿过，再穿过。

我们刚刚做了什么？我们展示了这个表达式。我们的某个输出如何关联到特定的输入。此外，这些输出彼此独立。这意味着什么？这意味着表达式的输出等于一系列独立计算的结果。

这是否意味着我们就这样推导出了Flash Attention，作为一种将注意力分解为一系列独立操作的方式？不完全是，但有点接近。我们首先需要更多的基础设施。

---

推导内核：便签纸方法

上一节我们接近了目标，本节中我们通过一种更直观的方法来推导内核：便签纸方法。

便签纸方法是一种简写推导，没有太多形式化。首先介绍它，是开始使用图表作为一个有用工具的好方法，并能理解推导这些数学为何如此重要。

便签纸方法是指，给定这个图表，我们可以在高层推导出Flash Attention，因为我们可以看到正在发生的独立操作。

我们如何做到？我们从一个高层图表开始，比如你看到的这个。首先，我们知道某些操作具有可流式处理的属性。当一个操作在一个轴上可流式处理时，意味着该轴的输入可以分块送入，并在较低层次累积。Softmax后接矩阵乘法就是这样一个实例。因此，我们可以让X轴被流式处理，分块发送下去。然后，因为这个轴一直穿过整个表达式，我们将进行流式处理。

现在，这个S_x数字代表较低层次的尺寸，它代表块大小。接下来我们做什么？之前我们展示了索引如何从这里一直流到开头。因此，我们可以分块执行键值，而不是一次性执行所有。我们将其标记为G_Q。

现在，我们在右边有什么？我们有一个子算法。这个子算法的图表说明了一些东西：它说明我们正在将Q分布到块中，每个块大小为G_Q，进行独立操作。对于每个块，我们以大小为S的块流式处理X轴。

然后，我们可以用这个来显示芯片上的内存量、每个子算法发生的传输次数，以及我们拥有的子算法数量。子算法的数量就是Q的总尺寸除以单个块大小。但有一个限制：G_Q必须低于某个水平，以免违反较低层次的内存限制。考虑到我们的流大小可以很小，小到可以忽略，我们可以得到这个表达式。这给出了P尺寸的上限。

很好，所以我们有了这个图表，它代表一个底层算法，并且它允许我们得到芯片上的最大内存。但对于Flash Attention，重要的洞见是能够减少发生的总传输次数。为了计算这个，我们将查看每个瓦片的传输次数，通过查看表达式输入和输出的尺寸来得到。所以我们有加载次数和保存次数。由此，我们可以计算总传输次数。我们将子算法数量乘以

31：AI中的数值表示与精度优化 🚀

概述

在本教程中，我们将深入探讨GPU编程中的数值表示方法，特别是整数和浮点数在AI应用中的使用。我们将从数值表示的基础知识开始，逐步介绍不同数据类型的优缺点，以及它们在AI训练和推理中的应用。通过本教程，您将理解如何选择合适的数值表示以优化性能并保持模型精度。

---

1：数值表示基础

1.1 实数与有限表示

实数线包含无限多个值，范围从任意小的值到任意大的值。然而，计算机的内存和计算资源是有限的，因此只能使用有限的数值表示来近似实数线。这些有限表示会采样实数线，但无法准确表示所有值。

关键公式：

• 实数线：无限精度，无限范围。
• 计算机表示：有限精度，有限范围。

1.2 精度与准确度

精度指的是数值表示对实数线的采样密度，而准确度表示采样值与实际值之间的接近程度。高精度不一定意味着高准确度。

示例：

• 对于π，3.1415比3.14更精确且更准确。
• 3.14159比3.14更精确，但可能不如3.14准确（如果实际π更接近3.14）。

---

2：整数表示

2.1 整数类型

整数表示以均匀间隔采样实数线，通常用于表示整数值。常见的整数表示包括无符号整数、有符号整数和补码表示。

代码示例：

// 无符号整数
unsigned int a = 255;
// 有符号整数（补码表示）
int b = -128;

2.2 整数在AI中的应用

整数在AI推理中广泛应用，主要优点是节省内存和带宽，并提高计算速度。然而，整数表示需要缩放因子来处理分数值，这增加了复杂性。

缩放因子公式：
[
x' = x_q \times \text{scale}
]
其中 (x_q) 是量化后的整数，(x') 是解码后的浮点数。

---

3：浮点数表示

3.1 浮点数基础

浮点数表示可以同时表示整数和分数，但其采样不是均匀的。浮点数在接近零的区域采样更密集，而在较大值时采样更稀疏。

浮点数结构：

• 符号位（1位）
• 指数位（E位）
• 尾数位（M位）

示例（FP16）：

• 结构：E5M10（5位指数，10位尾数）
• 解码公式：( \text{值} = (-1)^{\text{符号位}} \times 2^{\text{指数} - \text{偏置}} \times (1 + \text{尾数}) )

3.2 特殊值

浮点数包含一些特殊值，如零、无穷大和NaN（非数字）。这些值用于处理溢出或无效操作。

特殊值示例：

• 零：指数和尾数全为零。
• 无穷大：指数全为1，尾数为零。
• NaN：指数全为1，尾数非零。

---

4：浮点数在AI中的应用

4.1 动态范围与精度

浮点数的动态范围由指数位决定，精度由尾数位决定。在AI中，BF16（E8M7）比FP16（E5M10）具有更宽的动态范围，但精度较低。

动态范围公式：
[
\text{动态范围} = \frac{\text{最大可表示值}}{\text{最小可表示值}}
]

4.2 低精度浮点数

为了进一步节省内存和计算资源，AI中开始使用低精度浮点数，如FP8和FP4。这些格式通过缩放因子和块编码来保持精度。

FP8示例：

• E4M3：较窄的动态范围，较高的精度。
• E5M2：较宽的动态范围，较低的精度。

---

5：缩放因子与块编码

5.1 缩放因子的作用

缩放因子用于将浮点数映射到低精度表示中，以覆盖所需的动态范围。缩放因子可以是每张量、每行或每块的。

缩放因子公式：
[
x_{\text{quantized}} = \text{round}\left( \frac{x}{\text{scale}} \right)
]

5.2 块编码

块编码将张量划分为小块，并为每个块分配独立的缩放因子。这种方法可以在保持精度的同时减少计算开销。

块编码示例：

• 块大小：32元素
• 缩放因子：每块一个

---

6：数值稳定性与浮点数陷阱

6.1 非结合性

浮点数加法不满足结合律，因为中间值的舍入会导致结果差异。例如：
[
(a + b) + c \neq a + (b + c)
]

6.2 融合乘加（FMA）

FMA指令在单个操作中执行乘法和加法，可以减少舍入误差。然而，它也可能导致一些意外的数值行为。

FMA示例：
[
\text{FMA}(a, b, c) = a \times b + c
]

---

7：AI中的数值优化

7.1 训练中的数值问题

在训练中，梯度可能非常小，导致在低精度表示中被舍入为零。通过缩放损失值，可以将梯度映射到可表示的范围内。

损失缩放公式：
[
\text{缩放后的损失} = \text{损失} \times \text{缩放因子}
]

7.2 推理中的数值优化

在推理中，通过量化技术将高精度模型转换为低精度表示，以节省内存和计算资源。MXFP格式通过块编码和缩放因子实现了高效的推理。

MXFP示例：

• 块大小：32元素
• 缩放因子：每块一个

---

总结

在本教程中，我们深入探讨了GPU编程中的数值表示方法，特别是整数和浮点数在AI中的应用。我们从数值表示的基础知识开始，逐步介绍了不同数据类型的优缺点，以及它们在AI训练和推理中的优化方法。通过理解这些概念，您将能够更好地选择合适的数据类型，以在保持模型精度的同时最大化性能。

32：Factorio学习环境

概述

在本节课中，我们将学习Factorio学习环境（FLE）。这是一个专为评估大型语言模型（LLM）和视觉语言模型（VLLM）在《异星工厂》（Factorio）游戏中的表现而设计的平台。我们将探讨其设计动机、核心架构、评估方法以及它在未来智能体研究中的潜力。

---

动机：为什么需要Factorio学习环境？

上一节我们介绍了课程背景，本节中我们来看看创建FLE的动机。

当前评估智能体的基准存在一个问题：它们容易被快速饱和。一个基准可能在几年内有效，但随着智能体能力的进步，该基准就无法再有效区分不同智能体的性能，从而变得无效。

即使最近出现了更难的、甚至超人类的基准（如HumanEval的最终考试和AKGI2），这种饱和趋势并未减缓，反而在加速。根据当前预测，这些新基准可能在约12个月后就不再有用。

同时，模型能够完成任务的时间范围（时间跨度）每6.6个月翻一番。过去，任务（如SWE-bench）可能在几分钟内完成。现在，最先进的模型有约50%的概率能完成耗时约一小时的任务，并且这个时间跨度仍在持续翻倍。

这促使我们需要一类能够应对超长跨度任务的新环境。在这些任务中，智能体可能需要花费数百小时、执行数万步操作或工具调用才能达成目标。

因此，我们基本上需要一套新的环境，能够在超长时间跨度内抵抗性能饱和。

---

Factorio：一个理想的环境候选

Factorio是此类环境的绝佳候选。我们寻找的环境需要具备以下特点：

以下是理想环境的关键特征：

• 无界性：环境无法被“完成”，性能天花板无限高。Factorio符合这一点。
• 过程生成：环境应是过程生成的，这样模型无法通过预训练记住解决方案并直接输出，从而将“知识”成分与“智能”成分分离，以进行有意义的基准测试。
• 长跨度：任务需要很长的决策链，与前述时间跨度增长趋势相匹配。
• 后果性：需要做出一系列非常精确的决策才能达到最终目标。
• 开放性：达成目标有多种不同方式。例如，创办一家初创公司就是一个未来模型可能需要完成的目标，这没有固定配方，需要大量探索和目标设定。

Factorio是Steam上评分最高的游戏。它本质上是一个自动化建设模拟游戏。玩家作为一个角色出生在一个世界，目标是通过开采原始资源、建立自动化处理链、投资科技研究，最终发射火箭并逃离。世界纪录大约1小时15分钟，但通常新手需要约50小时。

游戏是一个沙盒，没有真正的终点，能力范围跨越十几个数量级。新手可能建造每分钟生产10个物品的小工厂，而人类团队建造的最大“超级基地”每分钟能发射约一千枚火箭。游戏地图随机生成，无法死记硬背开局策略。

游戏中的工厂极其复杂，一个微小的错误（如放错一个单位）就可能导致整个供应链崩溃，这很像系统工程或软件工程的挑战。存在多种主导策略，例如使用主干传送带或建造“意大利面式”工厂。后期建造大型工厂时，模块化设计、基础设施网络等软件工程思想同样适用。

此外，还有不同的元策略：工厂生产污染会激怒游戏中的原生生物（虫族），导致它们攻击工厂。因此，高速增长的代价是需要投资军事工业综合体进行防御。另一种策略是接受较慢的增长速度，投资太阳能等绿色能源技术，减少污染，从而避免被攻击，专注于清洁增长。

由于其巨大的沙盒性质，Factorio成为了一个绝佳的研究测试平台。

---

FLE的核心架构与工作原理

上一节我们了解了Factorio作为评估环境的优势，本节中我们来看看FLE的具体实现。

在GPU Mode的Discord社区中，我们花了18个月构建了一个用于在Factorio中评估LLM和VLLM的框架。

其工作原理如下：类似于软件工程，智能体通过代码符号化地与游戏交互。

以下是FLE的核心工作流程：

1. 观察：智能体通过读取其程序执行产生的标准错误和标准输出来观察世界。对于视觉语言模型，还可以选择性地获取游戏的视觉表示。
2. 抽象与复用：智能体被鼓励采用分层抽象开发，通过声明函数使代码可复用，编写类来形式化它们对世界的理解。
3. 信息筛选：由于观察流来自标准输出和错误，智能体可以修剪它们选择观察的内容，防止输出阻塞其令牌流。例如，一个有1000台机器的工厂，如果将所有实体符号列表都放入上下文流，模型可能很快耗尽资源。
4. 策略演进：随着智能体在游戏中推进，它们会发展出越来越具表现力的行动策略，以及更精确的观察和反应方式。
5. 奖励：智能体通过生产吞吐量获得奖励，这本质上是工厂的“GDP”。

我们有一个可以立即使用的pip包。安装后，大约三行代码就可以开始运行你自己的评估。

环境的工作方式类似于标准的聊天转录格式（用户/助理）。我们将Python API模式以及现成LLM与游戏交互所需的所有附加信息放入系统提示中。环境观察和相关内容在用户消息中，而智能体被要求对助理消息进行采样。这样，它几乎与任何经过聊天训练的LLM兼容。

智能体在游戏中是具身的，它们有一个角色，移动范围有限，必须四处移动并在发现新事物时进行观察，逐步建立其世界模型。

---

技术实现细节

FLE的实现基于一个Python执行层，其中包含了Factorio对象模型的Python表示。我们为所有对象提供了强类型，这意味着模型可以更有效地进行符号推理，并编写下游代码来处理这些对象。

Python命令被翻译成远程过程调用，然后通过TCP使用RCON协议发送到正在运行的Factorio无头服务器。服务器在游戏中执行这些底层原始操作（如拾取、放置物品），从而异步更新游戏状态。

我们跟踪并整合一小段时间的奖励，以观察行动对生产分数的影响。然后，我们从游戏状态获取一个对象返回，将其重新编组回Python，并由策略的其余部分处理。

RCON协议是一种游戏通信协议，常用于《我的世界》等游戏执行管理员命令。Factorio也支持RCON，并暴露了一个Lua控制台，我们可以利用它向运行中的游戏热加载任意脚本，这构成了Factorio学习环境后端的基础。

我们不需要Factorio的自定义构建版本，只需使用官方发布的多人在线服务器，然后用FLE客户端替换Factorio客户端，即可使服务器兼容智能体。

在性能方面，由于游戏服务器高度优化，平均每秒可以执行约220个工具操作。在一台M4笔记本电脑上，可以同时运行64到128个无头Factorio实例，这有利于进行大量并行实验和在线学习。

为了绕过Factorio客户端作为渲染器的限制，我们构建了自己的内部自定义渲染器，为视觉智能体提供游戏快照，而无需外部依赖。我们简化了游戏原版的视觉复杂性，使其对智能体更清晰，例如移除了背景，使用更简单的图块集。

有趣的是，对于大多数模型，添加视觉表示并没有提高性能。智能体似乎更擅长对坐标系统进行符号推理。但最新的模型（如Gemini 3 Pro和Claude Opus 4.5）的视觉推理能力已经达到可以从图像中获得额外收益的水平。

---

智能体策略与评估设置

智能体倾向于将Python命名空间作为记忆系统。它们观察世界，将观察结果赋值给变量，这些变量在整段情节中存储在命名空间里。这使得它们可以对那些对象进行下游操作。

函数实现了可复用性。通过封装，智能体可以将高级概念包装在单个函数中，存储在命名空间里，然后在后续策略中引用。这样，智能体每次采样的行动令牌的表达力会随时间增长。

在智能体设计方面，我们早期花了很多时间优化复杂的智能体框架。唯一真正有帮助的是进行反思和回溯。我们最终意识到，与其花数月时间开发更好的智能体，不如等待下一代模型的出现，它们的性能可能会超越我们的优化成果。

我们没有进行任何强化学习或微调。这个环境支持纯粹的上下文学习（提示）模式下的智能体。然而，我们确实进行了一些“中间压缩”，即省略过时的观察结果，以提高令牌效率并防止智能体因关注旧状态而犯错。

环境是模型无关的，可以接入OpenRouter等平台进行测试。智能体是具身的，这意味着评估不仅仅是优化工厂拓扑结构，智能体还需要控制角色移动、避免危险等。智能体在游戏中没有外部数据源，无法访问技能库或蓝图库。我们将蓝图视为智能体编写的命令式策略。

---

奖励机制与评估场景

Factorio的成本呈指数级增长，科技树的每一步都需要大约两倍于前一步的资源。到发射火箭时，一个最终物品大约相当于70万个原始资源。这只是名义上的胜利条件，但世界上最大的工厂每分钟能发射一千枚火箭，自动化水平惊人。

自3月份开始评估以来，最佳模型的性能已提升约8倍。我们预计生产分数大约每两到三个月翻一番。

游戏开发者提供了一个内部生产分数公式，我们直接采用它作为奖励信号。公式本质上是：奖励 = Σ（生产物品的价值 - 消耗物品的价值）。物品的价值由其最便宜的原料获取方式加上一个复杂性修正因子和能源成本决定。因此，奖励本质上是工厂的利润，生产越多，得分越高，但不能通过创建和销毁物品来作弊。

我们创建了两种评估场景：

1. 实验室玩法：一个资源充足的受限区域，智能体无需过多探索，只需在64到128步内利用给定资源布局工厂达成目标。我们为智能体提供足够的起始库存以启动自动化。
2. 开放玩法：更类似于真实的游戏流程。智能体被放置在世界中，拥有与人类玩家相同的初始物品，任务是“建造最大的工厂”。我们运行数千步，观察不同模型的表现。

我们还设计了吞吐量实验室玩法任务，要求智能体全自动地以特定速率（如每分钟16个物品）生产物品，并持续30秒的保持期。这可以防止奖励黑客行为（如手动制作后瞬间达到目标），迫使智能体创建可持续的工厂。由于结果可验证，我们可以无需人工检查或LLM作为裁判就进行大量轨迹推演，为下游蒸馏提供可扩展的数据生成方式。

---

实验结果与观察

在实验室玩法中，智能体的工厂布局看起来与人类非常不同，它们没有见过人类工厂的代码模式，思考方式也不同。

在吞吐量任务中，自3月份以来，最佳模型从只能生产塑料条，发展到能生产钢板、硫磺，再到高级电路。闭源模型（如Claude Opus 4.5）表现最佳，开源模型大约落后6到7个月。

不同模型在与游戏交互时表现出不同的失败模式。模型的编码能力与其空间推理能力和游戏表现有很强的相关性。经过大量代码优化的模型表现非常好，它们的主要失败在于“语用错误”——知道如何做，但不知道在特定情境下该做什么。未针对代码优化的模型则会犯很多语法和执行错误。

在开放玩法中，模型在长时间跨度内的生产分数在对数坐标上可以明显区分。有趣的是，模型似乎不会达到平台期，而是保持近乎线性的增长率。投资科技研发（如研究电钻）会显著提升增长率。

较弱的模型（如GPT-4 mini和LLaMA 70B）由于对Factorio的先验知识较弱，表现不佳，甚至会出现反复请求重置或建造大量无意义箱子的行为。

---

未来方向：工具性趋同与模型蒸馏

我们构建FLE的长期目标之一是研究工具性趋同假说。该假说认为，无论智能体的具体目标是什么，一系列行为（如自我保存、目标完整性保持、自我改进、资源获取、权力寻求）都会自然涌现。

Factorio是进行此类实证研究的绝佳环境。我们可以轻松地追踪资源获取、通过系统构建实现的认知增强，并设计实验来评估目标内容完整性等。

我们的计划是利用吞吐量任务生成经过验证的任务完成数据，并利用模型生成的合成蓝图，通过我们的无头渲染器进行视觉接地。然后，将这些来自大型、重型SOTA模型的Factorio能力蒸馏到一个微小的（如8B）模型中。

这将使我们从需要大量系统提示的上下文学习范式，转向将所有信息自然嵌入权重的范式。我们可以创建运行速度快、易于进行在线训练（如接入RL循环）的“模型生物”，甚至可以创建“回形针最大化器”。

一旦我们训练出展示工具性趋同行为的模型，就可以进行有趣的实验，例如：激活引导以减少权力寻求行为，或者尝试在不损害下游任务性能的情况下去除这些不良驱动。

我们预计，随着SOTA模型的进步，会看到越来越多的工具性趋同行为。在模型变得非常智能并在现实世界中获得权力之前，找出缓解方法是一个有价值的研究方向。

---

总结

本节课中，我们一起学习了Factorio学习环境（FLE）。我们探讨了创建它的动机——为了应对智能体基准快速饱和的问题，并评估超长跨度任务的能力。我们详细介绍了FLE的核心架构，即智能体通过Python代码符号化地与Factorio游戏交互，并采用分层抽象和选择性观察。我们还了解了其技术实现、奖励机制（基于游戏内生产分数）以及两种主要的评估场景（实验室玩法和开放玩法）。实验结果表明，不同模型在Factorio中表现出不同的能力水平、失败模式和编码风格。最后，我们展望了FLE在未来智能体研究中的潜力，特别是在实证研究工具性趋同假说和通过蒸馏创建高效、可训练的“模型生物”方面。FLE旨在成为一个能够长期抵抗性能饱和、并深入探索智能体本质的评估平台。

33：第 86 讲：CuTe DSL 入门指南 🚀

概述

在本节课中，我们将学习 NVIDIA 的 CuTe DSL（领域特定语言）。CuTe DSL 旨在为开发者提供一种在保持对硬件完全控制的同时，比 PTX 更简单的抽象层来编写 GPU 内核代码的方法。我们将通过一系列示例，了解其核心概念、编程模型和实用功能。

---

背景与设计理念

上一节我们介绍了课程主题，本节中我们来看看 CuTe DSL 的设计背景。

当前，高级代码生成器和编译器允许用户在无需了解过多硬件细节的情况下编写内核。这使得快速迭代新算法变得容易。然而，它们也隐藏了许多细节，使得用户难以完全控制硬件特性，从而在追求极致性能时遇到困难。

另一方面，PTX 提供了对硬件特性的完全控制，但其抽象层次过低，包含过多细节，编程效率不高。

CuTe DSL 的设计目标是在这两者之间取得平衡。它为开发者提供了在 NVIDIA GPU 上编程的表达性抽象，试图提供比 PTX 更简单、同时又能给予用户完全控制的编程体验。

CuTe 的设计架构包含多个抽象层：

• 顶层：高级接口，允许用户基于调优配置选择优化内核。
• 底层：包含 CuTe 抽象（黄色框）和原子操作抽象（灰色框）。CuTe 抽象提供一致的编程模型，原子操作抽象允许用户针对不同架构编程。
• CuTe DSL 正是基于这两层构建的解决方案，旨在赋能 AI 研究人员和工程师进行快速开发和原型设计。

---

开始使用 CuTe DSL

了解了设计理念后，本节我们来看看如何开始使用 CuTe DSL。

开始使用 CuTe DSL 非常简单。以下是快速入门步骤：

1. 通过 pip 安装：pip install nvidia-cute
2. 编写 Python 代码。
3. 直接运行程序，代码将在 NVIDIA GPU 上执行。

在接下来的内容中，我们将聚焦于为 GPU MODE 竞赛提供的模板程序，并以这些程序为例，帮助大家理解如何使用 CuTe DSL 的功能来编写代码。

---

示例一：简单矩阵乘法（GEMM）

我们将从第一个示例开始，这是一个简单的矩阵乘法（GEMM）计算。

核心装饰器：@cute.jit 与 @cute.kernel

在 CuTe DSL 中，有两个关键的装饰器。

以下是 CuTe 装饰器的组合与要求：

装饰器	作用	调用限制
@cute.jit	主机端入口点，用于为内核准备参数（如网格大小、块大小、簇配置）。它也可以作为内联函数被 @cute.kernel 调用。	可被 @cute.jit 或 @cute.kernel 调用。
@cute.kernel	GPU 内核的全局入口点。	只能被另一个 @cute.kernel 装饰器调用。

CuTe 布局代数

第二个需要熟悉的概念是 CuTe 布局代数。

CuTe 布局主要包含一系列划分（partition）和切片（slice）操作。例如，在 GEMM 中，我们使用 tile 操作根据给定的 BLOCK_M 和 BLOCK_K 配置来划分全局输入矩阵 A。

# 示例：使用 tile 操作划分全局张量
gA = cute.global_tensor(A, (M, K))
tiled_gA = cute.tile(gA, (BLOCK_M, BLOCK_K))

划分后，我们得到了一个具有层次化坐标信息的矩阵。然后，我们可以使用 slice 操作，根据线程索引，从划分后的张量中获取每个线程需要处理的数据块。

张量与张量 SSA

第三个有用的概念是 CuTe DSL 中的张量（Tensor）和张量 SSA（Tensor SSA）。

• 张量：是内存的迭代器和布局的组合。它支持全局内存、共享内存和寄存器内存。
• 张量 SSA：是 CuTe 张量的本地数据建模，采用值语义且不可变。

当我们进行全局逐元素操作时，可以先将全局张量加载到一个张量 SSA 表示中（底层是向量操作），然后使用大量向量内联操作进行类型转换和逐元素计算，最后将结果存回寄存器内存。

控制流

CuTe DSL 支持 for、while 和 if 控制流，允许用户编写符合直觉的控制逻辑。编译器会根据控制流条件进行 AST 处理，决定是直接求值代码还是生成条件分支的 IR 代码。

for 循环有一个特殊的 cute.arange，它几乎等同于 range，但支持一些高级的循环展开和流水线控制。例如，在简单的 GEMM 示例中，我们可以使用 @cute.prefetch 注解，将简单的循环转换为手动软件流水线，提前加载数据，使计算与数据预取并行执行。

---

示例二：通用矩阵乘法（GEMM）优化

在熟悉了基础功能后，本节我们来看一个更复杂的通用矩阵乘法（GEMM）优化示例。

为了在 GPU 上达到接近理论极限的计算性能，我们需要利用 Blackwell 架构的许多高级特性。这个示例旨在帮助参与者熟悉 CuTe 的编程模型。

该模板展示了如何：

1. 使用 TMA 将张量 A、张量 B 以及缩放因子张量从全局内存加载到共享内存。
2. 使用 S2T 拷贝将缩放因子从共享内存加载到张量内存（Tensor Memory）。
3. 使用 张量核心 进行计算。
4. 使用 张量内存到寄存器 的拷贝将计算结果移出，并进行简单的全局存储。

CuTe 编程模型

该示例揭示了 CuTe 的一个典型编程模型，这个模型在所有竞赛模板示例中都很常见。

对于拷贝操作（尤其是全局张量）：

1. 使用 local_tile 对全局张量进行分块。
2. 对源张量和目标张量进行 partition（或使用 tma_partition 简化 TMA 逻辑）。
3. 调用 cute.copy 在划分后的张量上进行迭代。

对于核心计算：

1. 在本地内存中准备操作数片段，使其符合指令要求的数据类型。
2. 调用 cute.gemm 等计算函数。

序列和逻辑非常相似。优化的关键在于如何选择原子操作、分块模式，以确保最大的内存带宽利用率和最佳的张量核心效率。

优化思路

对于竞赛参与者，可以从以下几个方面探索优化：

• 原子操作尺寸：根据问题大小选择，避免浪费计算或内存资源。
• 软件流水线：将数据拷贝操作与计算操作重叠。理想情况下，计算时数据应已预取到本地内存。
• 工作组专业化：利用 TMA 和张量核心指令的异步特性，让不同的工作组负责不同的拷贝和计算任务，提高并发性。
• 细粒度分块：在输出阶段（Epilogue）对数据进行子分块，在寄存器中进行逐元素操作，并与共享内存传输并发执行，以优化寄存器使用和共享内存压力。

---

示例三：双矩阵乘法（Dual GEMM）

接下来，我们看看双矩阵乘法（Dual GEMM）示例。

双 GEMM 涉及两个 B 矩阵（与 A 矩阵的 K 维度相同），计算得到两个累加结果。在将结果写入全局内存之前，我们对其中一个结果进行激活函数（如 GELU）处理，然后进行逐元素操作。

Lambda 表达式与张量 SSA

这个示例的关键特性是使用了 Lambda 表达式。

内核函数的签名可以包含编译时常量参数和动态参数。动态参数是内核特征，而静态参数是一些可以应用于内核以提高效率的内联函数。这里，我们使用 lambda 函数来说明如何执行激活操作。

同时，张量 SSA 的强大功能再次体现：我们可以将寄存器张量加载到向量化操作中，直接应用 GELU 操作，然后进行两个张量 SSA 之间的逐元素乘法。这一切都简化了实现，无需逐个元素管理。

---

示例四：分组矩阵乘法（Grouped GEMM）

本节我们来看最复杂的分组矩阵乘法（Grouped GEMM）示例。

分组 GEMM 包含多组矩阵乘法，每组中的 M 和 K 大小可能不同。代码需要动态计算每个线程块（CTA）要处理的数据。

动态 TMA 描述符更新

与之前示例不同，此代码的核心难点在于需要在线更新 TMA 描述符。

在代码开始时，每个 CTA 需要根据坐标和问题大小，计算出它属于哪个组以及要计算的数据块。然后，它需要在核函数内部更新 TMA 描述符。

模板中使用了 TMA 映射管理器 工具来执行这种动态的张量映射更新。基本流程是：

1. 每个 CTA 确定其所属的组。
2. 使用主机端创建的 TMA 描述符作为模板，在共享内存中构造新的描述符。
3. 用实际的张量信息（在线计算得到）更新该描述符。
4. 将更新后的描述符存回全局内存。
5. 同步确保所有 TMA 映射更新完成。
6. 后续的 TMA 加载、缩放因子加载和张量核心计算步骤与普通 GEMM 类似。

一旦掌握了张量描述符的动态管理，从单个 CTA 的角度看，其余部分与优化后的 GEMM 和 Dual GEMM 非常相似。

---

其他实用功能

除了核心编程模型，CuTe DSL 还提供了一些其他实用功能来辅助开发和调试。

编译与缓存

CuTe DSL 支持自定义缓存键。由于 DSL 编译时间可能较长，在基准测试时，我们不希望每次都重新编译内核。用户可以自定义一个键来索引缓存。如果内核已在全局缓存中，则无需重新编译。

调试打印

调试时，打印功能非常有用。

• cute.print_：打印所有静态信息，运行时变量显示为问号。
• cute.print：在主机端和设备端均可使用，打印运行时参数的值。
• cute.print_tensor：打印张量的内容（支持的数据类型有限）。对于低精度数据，可以先转换为 float32 再打印。

内部调试与 PTX 内联

CuTe DSL 还支持生成内部表示以辅助调试。

• keep_ir：可以转储 IR 来检查问题。
• lineinfo：在生成代码中包含源文件信息，便于在性能分析工具中关联源码与汇编指令。
• 内联 PTX：如果遇到某些底层优化在 DSL 中尚未支持，熟悉 PTX 语法的用户可以使用内联 PTX 功能来绕过限制。

---

总结与资源

本节课我们一起学习了 NVIDIA CuTe DSL 的核心概念和编程方法。

我们回顾了：

1. CuTe DSL 的设计目标与架构。
2. 通过简单 GEMM 示例，学习了核心装饰器、布局代数、张量 SSA 和控制流。
3. 通过通用 GEMM 优化示例，深入了解了 CuTe 的编程模型和性能优化思路。
4. 了解了 Dual GEMM 和 Grouped GEMM 的特殊处理方式。
5. 掌握了编译缓存、调试打印和内联 PTX 等实用功能。

有用的资源：

• CuTe 创始人 Chris 在 GPU MODE 竞赛中的介绍，包含了代数和布局变换的详细细节。
• CuTe DSL 官方文档，涵盖了本课程提到的大部分功能。
• GitHub 上的 DSL 示例库，其中许多展示了如何达到接近理论极限的性能。

希望本指南能帮助大家在 GPU MODE 竞赛中取得好成绩，并享受使用 CuTe DSL 进行高性能 GPU 编程的旅程。祝大家好运！

34：使用NVSHMEM的低延迟通信内核

在本节课中，我们将学习如何为多节点大语言模型推理设计低延迟的通信内核，特别是针对all_reduce操作。我们将探讨现有通信库的性能瓶颈，并介绍一种名为NVR的自定义内核设计，它利用NVSHMEM实现了显著的性能提升。

概述

随着大语言模型规模的不断增长，单GPU已无法容纳整个模型，因此必须将模型的计算和参数分布到多个GPU上。这引入了跨节点的通信开销，尤其是在使用张量并行策略时，每层之后都需要进行all_reduce操作。本课程将分析这些通信瓶颈，并展示如何通过精心设计的NVSHMEM内核来缓解它们。

背景：大语言模型推理与并行策略

大语言模型推理是一个自回归过程，模型逐个生成令牌。传统上，它分为两个阶段：

• 前缀阶段：模型并行处理所有提示令牌以生成第一个输出令牌。此阶段计算密集，因为涉及大量矩阵乘法。
• 解码阶段：将生成的令牌反馈给模型以生成下一个令牌。此阶段内存密集，因为输入是单个或少量令牌，矩阵乘法变得“瘦高”。

为了在多个GPU上运行大模型，主要有三种并行策略：

1. 流水线并行：将模型的层序列分配到不同GPU上。数据在GPU间以点对点方式传递。通常将令牌批次划分为微批次以填充流水线。
2. 张量并行：将单个层内的矩阵乘法拆分到多个GPU上。每个GPU计算部分结果，然后在每层之后通过all_reduce操作跨GPU聚合结果，这会产生很高的通信量。
3. 专家并行：主要用于混合专家模型，本课程主要关注前两种策略。

多节点推理的性能瓶颈研究

为什么需要研究多节点推理？并非所有人都拥有配备高速NVLink互联的DGX服务器。在传统的HPC集群中，节点间通常通过速度较慢的网络（如Slingshot或InfiniBand）连接。标准的扩展方式是：在节点内（NVLink域）使用张量并行，在节点间使用流水线并行，这称为混合并行。

然而，我们的性能研究发现，在强扩展（固定工作负载，增加GPU数量）场景下，时间并未按预期减少，反而经常保持不变或增加。特别是在解码密集型工作负载中，混合并行的性能显著劣于跨节点的纯张量并行。

通过深入分析性能跟踪数据，我们发现了两个关键点：

1. 解码阶段的计算特性：在解码阶段，矩阵乘法的输入维度M（令牌数）很小。在流水线并行中，通过微批次进一步减小M，但这对计算时间几乎没有改善，可能是因为矩阵维度低于硬件优化的平铺大小。而在张量并行中，通过分割内部维度K（参数维度），计算时间得到了预期的减少。
2. 显著的通信开销：在跨节点的张量并行中，all_reduce通信时间急剧增加，成为主要瓶颈。

聚焦于解码阶段的通信瓶颈，我们发现all_reduce的消息大小通常在几十到几百KB。基准测试显示，在此消息大小范围内，NCCL的性能有时甚至不如未充分优化的MPI实现。我们假设这是因为NCCL仅使用环或树算法，而MPI可能使用了更适合该消息大小范围的递归倍增算法。

NVR：自定义低延迟All-Reduce内核设计

为了解决上述瓶颈，我们设计了NVR，这是一个基于NVSHMEM的三阶段all_reduce内核。

NVSHMEM简介

NVSHMEM是基于OpenSHMEM模型的并行编程接口，它提供了一个逻辑上共享的全局地址空间。每个处理单元可以直接访问远程PE的内存，通过put和get等API进行数据移动。NVSHMEM提供了主机端和设备端API。

NVR的三阶段设计

NVR的设计旨在最小化跨节点通信量，并使用优化的算法。以下是其三个阶段的伪代码概述：

// 伪代码：NVR All-Reduce 驱动函数
void nvr_allreduce(void* message, size_t msg_size, int gpus_per_node, int num_nodes) {
    // 阶段 1: 节点内Reduce-Scatter
    // 在单个NVLink域内，将数据规约并分散到节点内的各个GPU上。
    // 这减少了需要跨节点通信的数据量（减少为原来的 1/gpus_per_node）。
    intra_node_reduce_scatter(message, msg_size);

    // 阶段 2: 节点间递归倍增All-Reduce（自定义NVSHMEM设备内核）
    // 跨节点的对应GPU之间，使用递归倍增算法进行规约。
    // 这是性能关键路径，我们为此编写了自定义内核。
    launch_rd_inter_node_kernel(partial_message, msg_size / gpus_per_node);

    // 阶段 3: 节点内All-Gather
    // 将跨节点规约后的部分结果在节点内收集起来，使每个GPU都拥有完整的全局规约结果。
    intra_node_all_gather(final_message, msg_size);
}

设计选择与洞察：

• 为何不融合内核？ 在异构网络（节点内NVLink，节点间其他网络）中，为不同域优化的内核启动参数（如线程块大小）可能不同。虽然可以编写一个复杂的融合内核，但分离的内核更易于优化和维护，并且我们仍能获得出色的性能。
• 为何使用Reduce-Scatter + All-Gather？ 这种设计将跨节点通信量减少了gpus_per_node倍，对于延迟敏感的操作至关重要。

关键优化技术

除了三阶段设计，我们还实施了多项关键优化以实现高性能。

1. 细粒度同步替代全局同步

在连续的all_reduce调用之间，需要确保缓冲区可安全重用。标准的nvshmem_quiet()操作成本高昂，因为它会与通信器中的所有PE进行同步。

我们的优化：我们为每次all_reduce操作分配一个唯一的序列号。每个PE只与其通信对端进行细粒度同步，通过检查远程对端的序列号原子变量来实现。这避免了昂贵的全局quiet操作。

2. 高效使用NVSHMEM RMA操作

• 优先使用put而非get：NVSHMEM推荐使用put进行数据转移。
• 使用非阻塞API：我们使用nvshmem_put_nbi等非阻塞操作，允许计算和通信重叠。
• 选择适当的API变体：对于基于代理的传输，使用nvshmemx_putmem_block（块级API）在节点内场景能获得最佳性能。

3. 数据分块与参数调优

我们引入了两个超参数：

• 线程块数量：多个线程块并行处理独立的数据块，允许一个线程块发出put请求时，另一个线程块正在对接收到的数据进行规约，实现了计算与通信的重叠。
• 分块大小：每个线程块将其负责的数据进一步分块，每块通过一次非阻塞put发送。这有助于控制网络注入大小，对性能有巨大影响。我们需要为不同的消息大小和节点配置寻找最优的分块大小和线程块数量。

4. 无信号同步（L风格协议）

使用put_signal API进行显式同步在Slingshot等网络上非常慢。

我们的优化：我们采用了类似NCCL的L风格协议。我们将需要规约的数据和一个表示数据有效性的标志打包在同一个8字节的有效载荷中发送。由于8字节操作在大多数网络和CPU/GPU架构上是原子的，因此可以保证数据和标志作为一个整体被接收。接收方通过检查标志值（与当前步骤的序列号比较）来判断数据是否有效，然后进行规约。这完全避免了显式的信号API调用。

性能评估与总结

我们将NVR集成到我们自研的推理引擎YaLLS中，并进行了性能评估。

独立微基准测试：在Pearl Motor（Slingshot网络）和Vista（InfiniBand网络）系统上，NVR在关心的消息大小范围内（约256KB-2MB）显著优于NCCL，速度提升可达1.5倍至3倍。

端到端推理加速：在LlaMA 70B和405B模型的多节点推理中，使用NVR替换NCCL的all_reduce操作，带来了显著的端到端加速。例如，在32个GPU上运行70B模型，获得了最高1.8倍的加速。性能分析表明，加速主要来自于通信时间的减少，最高可减少50%。

理论模型：我们的递归倍增算法具有O(log N_nodes) * α_inter的延迟项，而NCCL的树算法为O(log N_nodes) * 2 * α_inter。在中等消息大小下，这个因子2的差异会体现为明显的性能差距。

主要经验总结

以下是设计高效NVSHMEM通信内核的关键经验：

1. 最小化跨节点消息量：通过类似Reduce-Scatter的操作减少需要跨网络传输的数据。
2. 为消息区间选择最优算法：研究不同通信库（如MPI）的算法，选择最适合目标消息大小的算法（如对小消息使用递归倍增）。
3. 考虑异构网络的分核设计：为不同网络域（节点内/节点间）使用独立优化的内核可能比单一融合内核更简单且高效。
4. 优先使用Put和非阻塞操作：遵循NVSHMEM最佳实践，利用计算通信重叠。
5. 避免昂贵的全局同步：尽可能用细粒度的同步机制（如基于原子操作的序列号检查）替代nvshmem_quiet或fence。
6. 避免显式信号：探索像L风格协议这样的替代同步策略。
7. 数据分块与参数调优至关重要：分块大小和线程块数量对性能有极大影响，需要针对具体配置进行调优。
8. 注意传输后端优化：某些NVSHMEM传输后端（如libfabric）可能尚未充分利用硬件特性，性能有待优化。

本节课中，我们一起学习了多节点LLM推理中的通信瓶颈，深入探讨了如何利用NVSHMEM设计和优化自定义的all_reduce通信内核（NVR）。通过三阶段设计、细粒度同步、数据分块以及无信号L风格协议等关键技术，NVR在多种HPC集群上实现了比NCCL更低的延迟和更高的端到端推理性能。这些设计原则和优化技巧也为希望使用NVSHMEM编写高性能通信内核的开发者提供了宝贵的参考。