随着加速计算不断提升 AI 和科学计算各个领域的应用程序性能，人们对 GPU 优化技术的兴趣也越来越浓厚，以确保应用程序获得尽可能好的性能。作为应用程序开发者，有很多方法可以对软件堆栈上下进行 GPU 编程。在本文中，我们将介绍堆栈的一些不同级别，并深入探讨最低级别：手写并行线程执行 (PTX) 代码。

加速计算软件堆栈

现在，您无需编写特定于 GPU 的代码，即可使用 GPU 完成大量工作。库开发者和软件工程师已经为您完成了底层工作。例如，您可以在堆栈中进行高级工作，使用蓝图构建完整的 AI 工作流。或者，您可以在 PyTorch 等框架中开发应用程序，您可以在其中指定模型、适当的 GPU 代码和库来自动执行程序。

您还可以使用全套 NVIDIA CUDA-X 库开发应用程序，其中包括量子计算、数据处理、物理 AI、基因测序、边缘计算、药物研发等领域特定的库。如果这些特定于域的库不包含您所需的所有功能，您可以使用 OpenACC 等编译器指令对 GPU 进行编程，也可以使用 libcu++ 对 C++ stdpar 算法等库和 C++ 标准库进行编程。

在上述所有情况下，您不是在编写特定于 GPU 的代码，而是依赖于由专家工程师精心设计、实施和优化的库或编译器指令。

但是，在某些情况下，您可能必须实现自己的 GPU 代码，因为当前不存在用于满足您所需功能的库。然后，您可以进一步向下移动堆栈，并直接使用高级语言 (例如 C++、Fortran 和 Python) 编写 CUDA GPU 代码。

最后，在极少数情况下，开发者可能会选择更深入地使用 PTX 直接编写代码中对性能极为敏感的部分。与大多数性能优化技术一样，您期望的控制越多，在堆栈中提取性能所需的越低。应谨慎考虑这一权衡：除了增加开发和调试的复杂性之外，手写低级代码带来的性能提升可能无法移植到其他 GPU 架构。

正如我们在之前的博文中所展示的，PTX 是 GPU 的汇编语言。直接编写 PTX 是一种非常先进的优化技术，对于大多数开发者来说并不是必需的，因此应将其视为万不得已的工具。不过，在某些情况下，通过编写 PTX 实现的精细控制可直接提高特定应用程序的性能。这些情况通常发生在应用程序对性能非常敏感的部分，其中性能提升的每一部分都具有显著的优势。所有可用的 PTX 指令都在 PTX ISA 文档中。

在这篇博客文章中，我们将深入探讨一个示例，其中使用手写 PTX 来提高某些 AI 模型实现中出现的重要算法的性能。

编写 PTX

在进入示例之前，我们将列出一些在应用中包含手写 PTX 代码的方法。换句话说，原则上是如何做到这一点。以下示例展示了真实场景并显示了性能变化。

内联 PTX

在代码中包含 PTX 的一种标准方法是使用内联 PTX。这是我们将在下面展示的方法，有关语法和语义的详细信息请参阅文档。这与在 CPU 上编写组装代码非常相似。

cuda::ptx 命名空间

在代码中包含 PTX 的另一个选项是使用 libcu++，其中包含命名空间 cuda::ptx，可提供直接映射到 PTX 指令的函数。这有助于在 C++ 应用程序中轻松使用特定的 PTX 指令。有关 cuda::ptx 命名空间的更多信息，请参阅 cuda::ptx 命名空间文档。

CUTLASS 示例

为了说明如何手动编写 PTX 代码，我们将使用线性代数的特定示例。一般来说，如果您的运算可以表示为线性代数运算 (例如 GEMM) ，则推荐使用 NVIDIA CUBLAS 在 GPU 上运行。CUBLAS 已经针对许多矩阵的大小和形状进行了高度优化，并且具有多个数值精度可供选择。

有时，CUBLAS 中的功能无法完全表达您想要执行的操作，或者您希望在 GEMM 之前或之后直接执行计算。有时，您可以通过将其他运算与 GEMM 运算融合来提高性能，而不是先调用一些函数，然后调用 CUBLAS 和更多函数。这有很多好处，因为融合内核可能会实现更多优化，例如更高效地使用数据。

这正是 NVIDIA CUTLASS 库的用武之地。CUTLASS 包含一系列 CUDA C++ 模板抽象，用于在 CUDA 内的各个级别和规模上实现高性能矩阵乘法 (GEMM) 和相关计算。由于 CUTLASS 支持对 GEMM 和类似 GEMM 的操作进行更多的控制和自定义，因此与 CUBLAS 相比，CUTLASS 涉及的开发者代码略多一些。

CUTLASS 包含大量手写 PTX，因为它在每个 GPU 架构上都以最佳性能进行设计。这使得 CUTLASS 成为说明手写 PTX 实际应用的现成示例。

GEMM 以及 top_k 和 softmax

我们将演示的特定运算是 GEMM 与 top_k 和 softmax 算法的融合。这是运行混合专家神经网络时的常见操作。我们将重点介绍 NVIDIA Hopper 架构。由于这是一种常用的操作，因此 CUTLASS 已经为此提供了一个带有一些内联 PTX 的特殊内核，因此可以直接演示 CUTLASS 如何将手写 PTX 融入其高性能 GPU 代码。

在这篇博文中，我们使用：

• 版本CUTLASS 的 3.9.2
• NVIDIA GH200 GPU
• 驱动程序版本 570.140
• CUDA 工具包版本 12.8

按照 CUTLASS 网站上的构建说明，我们使用构建选项 -DCUTLASS_NVCC_ARCHS = 90a 进行 cmake，以确保启用 Hopper 架构的完整功能集。CUTLASS 资源库中有许多示例展示了最新架构上的各种功能。完成 cmake 后，我们导航至构建目录 (例如，build/examples/ 61_hopper_gemm_with_topk_and_softmax) ，以便构建和运行示例代码。

我们执行 make 来构建代码，可执行文件已构建完毕并可随时运行。该应用程序接受一些不同的选项作为输入，包括矩阵大小 m、n 和 k、容错 epsilon，以及为生成 GFlop/s 基准测试数据而运行的迭代次数。

选择 m = 1024、n = 8 (默认) 、k = 4096、迭代 = 1000000 以及 epsilon 为 1e-4，即可获得以下输出。在此基准测试中，类似于 LLM 执行，m 是 token 的数量，n 是专家的数量，k 是专家的嵌入维度，top_k 的值为 2 (在测试代码中进行硬编码) 。

$ ./61_hopper_gemm_with_topk_and_softmax --m=1024 --k=4096  --iterations=1000000 --eps=1e-4
Disposition: Passed    Relative error: 1.52478e-05
Problem Size: 1024x8x4096x1
Avg runtime: 0.011765 ms
GFLOPS: 5704.11

在此基准测试示例中，性能为 5704 GFlop/s。我们将令牌数量 ( m 个参数) 最多改为 16384，并生成以下性能表。

表 1。基准测试代码的性能，包括将内联 PTX 用于 top_k 和 softmax 函数

删除内联 PTX

此基准测试示例将 GEMM 与 top_k 和 softmax 融合在一起，并调用在 top_ k 函数中使用内联 PTX 的函数，前提是 k 的值为 2 或 4。(注意：`k` 不同于上述矩阵维度 k。)

在某些条件下，它还使用 inline PTX 作为 softmax 函数。此外，当不符合特定条件时，top_k 和 softmax 都有使用 CUDA C++ 编写的后备例程。您可以直接更改 top_k 和 softmax 的内部函数，以注释掉 PTX 函数的调用，并运行备用 CUDA C++ 代码。在本例中，这将使我们能够量化手写 PTX 的价值。

为从此示例中删除内联 PTX，我们编辑了 cutlass/include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp 文件，以注释掉内联 PTX 函数的使用。在此文件的顶部附近，您会看到一些使用内联 PTX 编写的函数，这些函数以“top_2”和“top_4”开头。例如，以下是您将遇到的第一个 PTX 功能。

CUTLASS_DEVICE
Array<float, 2> top_2_reduce_scalar(Array<float, 2> a, float scalar) {
  Array<float, 2> out;
  asm volatile(
      "{\n"
      "  .reg .f32 mx;\n"
      "  .reg .pred p;\n"
      "  max.f32 mx, %3, %4;\n"
      "  setp.gtu.f32 p, %2, %4;\n"
      "  selp.f32 %1, mx, %2, p;\n"
      "  selp.f32 %0, %2, %4, p;\n"
      "}\n" : "=f"(out[0]), "=f"(out[1]) : "f"(a[0]), "f"(a[1]), "f"(scalar));
  return out;
}

无需了解此代码的每个细节。主要目的是展示一个简短的内联 PTX 功能的外观示例。

在源代码下方，您还会看到 softmax 函数。这些是我们将省略的 PTX 函数，以查看性能变化。

在同一源文件中，您可以找到作为 if 语句的一部分调用这些函数的位置。我们只需注释掉调用内联函数的 if 语句，然后保留语句的 else 部分。这将省略对内联 PTX 函数的调用，转而执行使用 CUDA C++ 编写的代码。

例如，有一个名为 add_element_to_desc_sorted_array 的函数，如果 k = 2 或 k = 4，则会分别调用 top_2 或 top_4 PTX 函数，或者调用该算法的 C++ 实现。此函数的代码如下：

void add_element_to_desc_sorted_array(cutlass::Array<Element, N>& a, Element b) {
  if constexpr (N == 2 && is_same_v<Element, float>) {
    a = top_2_reduce_scalar(a, b);
  }
  else if constexpr (N == 4 && is_same_v<Element, float>) {
    a = top_4_reduce_scalar(a, b);
  }
  else {
    // slower generic path with branching, slower, and can cause register spill
    CUTLASS_PRAGMA_UNROLL
    for (int k = 0; k < N; ++k) {
      if (a[k] < b) {
        // Shift down
        CUTLASS_PRAGMA_UNROLL
        for (int l = N - 1; l > k; --l) {
          a[l] = a[l-1];
        }
        a[k] = b;
        break;
      }
    }
  }
}

为确定手写 PTX 函数的效果，我们将对这些 PTX 函数的调用注释掉，并仅允许代码执行 C++ 版本的算法，如下所示：

void add_element_to_desc_sorted_array(cutlass::Array<Element, N>& a, Element b) {
/* BEGIN COMMENT
  if constexpr (N == 2 && is_same_v<Element, float>) {
    a = top_2_reduce_scalar(a, b);
  }
  else if constexpr (N == 4 && is_same_v<Element, float>) {
    a = top_4_reduce_scalar(a, b);
  }
  else {
 END COMMENT */
    // slower generic path with branching, slower, and can cause register spill
    CUTLASS_PRAGMA_UNROLL
    for (int k = 0; k < N; ++k) {
      if (a[k] < b) {
        // Shift down
        CUTLASS_PRAGMA_UNROLL
        for (int l = N - 1; l > k; --l) {
          a[l] = a[l-1];
        }
        a[k] = b;
        break;
      }
    }
//}  COMMENT THE END OF THE ELSE
}

我们对函数 merge_desc_sorted_arrays 和 masked_softmax 进行类似更改，以删除 if/else 语句，从而从本示例中删除手写 PTX 函数 $ ./61_hopper_gemm_with_topk_and_softmax --m=1024 --k=4096 --iterations=1000000 --eps=1e-4 Disposition: Passed Relative error: 1.52478e-05 Problem Size: 1024x8x4096x1 Avg runtime: 0.011765 ms GFLOPS: 5704.110、tg_ 11、tg_ 12、tg_ 13 和 tg_ 14。

以下是性能结果。

表 2。对于 top_k 和 tg_ 16 函数，基准测试代码在无内联 PTX 的情况下的性能

将这些结果与表 1 中的结果进行比较后，您会发现当使用手写 PTX 而非 CUDA C++ 代码时，性能将从 7% 提高到 14%。带回家所传达的信息并不是此处具体显示的绝对性能提升，而是在某些精心选择的情况下，手写 PTX 可以带来性能提升。应仔细分析性能和可移植性权衡，以确定在您的应用中包含手写 PTX 的可行性。

这是一个经过高度优化的示例代码，我们之所以选择它，是因为它具有由 NVIDIA CUTLASS 工程师编写的手写 PTX，显示出显著的性能提升。

此示例强化了以下指导：在绝大多数情况下，开发者应将 PTX 的手写工作交给 CUTLASS、CUBLAS 和其他 GPU 库的开发者，并在这些库的基础上进行构建。

总结

在本文中，我们展示了一个示例，展示了 CUTLASS 如何使用手写 PTX 来提高某些 AI 模型中使用的特定融合 GEMM 操作的性能。我们不想让人觉得每个开发者都应该编写 PTX。绝大多数开发者不需要这样做。手动编写 PTX 应该是万不得已的工具。

尽管如此，手写 PTX 是一项可供所有开发者使用的技术。这是一种先进的专业技术，如果使用得当，可以成为高级 GPU 程序员工具箱中的另一种工具。

这是 CUDA 平台的一大优势，开发者可以在任何适合自己的级别使用 NVIDIA 堆栈，从应用程序级别一直到编写组合代码 (PTX) ，以及介于两者之间的任何级别。

致谢

感谢以下 NVIDIA 贡献者：Ali Hassani