我们提供的 NVIDIA CUDA 深度神经网络库(cuDNN) 是一个专门为深度学习应用而设计的 GPU 加速库，旨在以先进的性能加速深度学习基元。
cuDNN 与 PyTorch、TensorFlow 和 XLA (加速线性代数) 等热门深度学习框架集成。这些框架抽象化了直接 GPU 编程的复杂性，使您能够专注于设计和训练模型，而无需担心底层硬件。cuDNN 可作为底层性能引擎，确保以更高的效率执行这些框架上的操作。
最近，扩展的点积注意力 (SDPA) 已成为大型语言模型 (LLM) 等重要工作负载中的性能关键基元。cuDNN 增加了对该基元的支持，并一直在使用闪存注意力和其他优化来提高其性能释放，同时扩展功能支持面，以支持一系列注意力用例。
在 NVIDIA H200 Tensor Core GPU 上，cuDNN 可以在 FP8 中实现高达 1.2 PFLOPS 的性能。作为端到端示例，我们的团队在为 Lama2 70B LoRA 微调启用 cuDNN FP8 SDPA 后，测得了 1.15 倍的速度提升。该实验使用了 NVIDIA NeMo 和 NVIDIA Transformer 引擎 在 8-GPU H200 节点上运行。
| 软件堆栈 | e2e 加速 |
| 禁用 cuDNN 的 NeMo | 1x (基准) |
| 在 BF16 中使用 cuDNN SDPA 的 NeMo | 1.11 倍 |
| 在 FP8 中使用 cuDNN SDPA 的 NeMo | 1.15 倍 |
在本文中，我将详细介绍 cuDNN SDPA 可实现的性能，介绍使用方法，并简要总结 cuDNN 9 中的其他一些值得注意的新功能。
经扩展的点积产品注意力性能
NVIDIA 正式通过在 NVIDIA Omniverse 中开源融合的 Multihead Attention (fMHA) 内核来支持注意力机制。APEX 库将注意力算法融合到单个内核中。Tri Dao 的创新工作以此内核为起点，实现了闪存注意力的形式，带来了巨大的性能改进和功能增强。欲了解更多信息，请参阅 Dao-AILab/flash-attention 在 GitHub 上的项目页面，并查看 FlashAttention-2：通过更好的并行性和作业分区提高注意力 相关论文。
然后， NVIDIA 通过更快、更灵活的实现，在融合注意力方面取得了进展。此实现现在是 NVIDIA 的最新技术成果，现在已经集成到 NVIDIA Hopper 架构 的 GPU 中。
XLA 目前提供了 cuDNN SDPA 的路径。您可以通过 JAX SDPA API 或依靠 XLA 编译器将 JAX/PyTorch 中的自定义实现降低到 cuDNN SDPA。
目前，PyTorch 渴望模式 SDPA 不使用 cuDNN，但基于 cuDNN 的实现正在进行中。有关更多信息，请参阅 PyTorch 的相关 Pull Request：Fprop 和 Bprop。
cuDNN SDPA 实施封装了以下内容：
- 深入了解基础 NVIDIA 硬件架构
- 实现从非闪存到闪存注意力 v2 的所有先进的 SDPA 算法，以及介于两者之间的所有算法
- 启发式算法，可根据问题大小和目标 GPU 架构自动设置性能旋钮 (例如图块大小)
这可在 NVIDIA GPU 上实现出色的可用性能。图 1 和图 2 显示，在各种用例中，cuDNN 9 BF16 的速度比现有最佳 GPU 快 2 倍，而 PyTorch 即时实现 中 cuDNN FP8 的实现速度提升高达 3 倍。
性能提升可实现更长的序列长度以及更短的模型预训练和微调时间。与其他公开基准（例如 flash-attention/benchmarks）相比，本文仅报告 GPU 时间，不包括主机用度。
SDPA 作为 cuDNN 图形
cuDNN 中的 SDPA 可指定为张量运算的 cuDNN 图形。对于给定的图形，cuDNN 库拥有一组可以执行该图形的引擎。
虽然一些图形可能没有合适的引擎，但其目的是为在 GPU 上原子化执行的任何 cuDNN 图形至少提供一个引擎，通常在一个融合内核中，但有时在一小部分协作内核中。您可以将其视为整个框架图形的“子图”。
SDPA (各种形式的 SDPA) 是此类子图的理想示例。支持这些模式的引擎旨在尽可能灵活，而不会影响显著的性能，而是使用现有的最佳算法 (如 Flash Attention)。这意味着您可以更改注意力计算，但仍然可以在 GPU 上高效运行。本节旨在解释目前提供的灵活性和支持。
图 4 显示了前向传播 (fprop) 用例的 cuDNN 图形。这是神经网络中注意力机制 SDPA 的前向计算所涉及的运算序列。
因果关系掩膜逻辑降低到单个生成索引、逐点和选择运算。softmax 逻辑降低到六个 cuDNN 运算 (图 4)。这种 cuDNN 运算的粒度使您能够灵活地表达自定义模型。图 5 显示了可能的组合。
图 5 显示，在头部维度更大 (256) 且无因遮罩的情况下，cuDNN FP8 正向闪光注意力可实现高达 1.2 PFLOPS 的运算性能。
您还可以使用 BMM1 和 Softmax 之间的任意逐点运算构建自定义图形。这种灵活性支持尚未发现的新变体。
当创意研究人员调整典型的注意力模式时，他们就不太容易回到优化程度较低的框架实现，从而在性能上达到峰值。图 2 适用于完整的训练运行。反向道具注意力图也具有类似的灵活性。
SDPA 使用演示
有多个 API 切入点可用于创建和运行 cuDNN 图形：
- 前端 API (包括 C++和 Python 变体)
- 后端 API (仅支持 C 语言)
上一节中描述的所有 cuDNN 图形概念都适用于这两个级别的 API.但是，cuDNN 团队建议您使用 Python 或 C++与 cuDNN 前端 API 进行交互，除非您需要 C 接口。前端 API 明显更简洁，并增加了一些便利。
例如，前端 API 扩展了操作节点的概念，使节点能够封装多个操作和它们之间的数据流。换言之，它们是方便节点，可以抽象出常见图形模式 (例如 SDPA) 的细节。节点仍然可以灵活配置众所周知的变体。此 SDPA 使用演练从最简单的情况开始，即在 Python 中创建的 SDPA 节点。
我们在 前端库中的 SDPA Python 示例 中演示了配置选项和基本使用流程：
- 初始化一个具有适当数据类型的
cudnn.pygraph对象。 - 创建一个包含张量维度、布局和数据类型的
Tensor对象。 - 创建缩放的点积闪光注意力节点并提供所需的配置。
- 构建图形并提供设备指针。
- 执行图形。
# 1. cuDNN graph
graph = cudnn.pygraph(
io_data_type=cudnn.data_type.BFLOAT16,
intermediate_data_type=cudnn.data_type.FLOAT,
compute_data_type=cudnn.data_type.FLOAT,
)
# 2. tensor attributes
q = make_tensor_attr(graph, q_gpu, "q")
k = make_tensor_attr(graph, k_gpu, "k")
v = make_tensor_attr(graph, v_gpu, "v")
attn_scale = make_tensor_attr(graph, attn_scale_cpu, "attn_scale", is_pass_by_value=True)
# 3. sdpa node with the configurations
o, stats = graph.scaled_dot_product_flash_attention(
name="scaled_dot_product_flash_attention",
q=q,
k=k,
v=v,
is_inference=false,
attn_scale=attn_scal,
bias=None,
use_alibi_mask=false,
use_padding_mask=false,
seq_len_q=None,
seq_len_kv=None,
use_causal_mask=true,
dropout=None,
)
# 4. build the graph and provide the device pointers
graph.build()
variant_pack = {
q: q_gpu,
k: k_gpu,
v: v_gpu,
o: o_gpu
}
# 5. execute the graph
graph.execute()
如果您想获得比 SDPA 节点默认情况下提供的更高的灵活性，则构建底层 SDPA 图形的代码是开源的，因此可以进行自定义。
例如，以下代码示例可自定义 SDPA 节点内的扩展节点。有关更多信息，请参阅 scaled_dot_product_flash_attention.h 文件。
// Optional scale
if (options.inputs.Attn_scale) {
// Lower options to scale options
auto attn_scale_output = std::make_shared();
attn_scale_output->set_is_virtual(true);
Pointwise_attributes scale_attributes;
scale_attributes.set_name("attn_scale");
scale_attributes.set_mode(PointwiseMode_t::MUL);
scale_attributes.inputs.IN_0 = last_output;
scale_attributes.inputs.IN_1 = options.inputs.Attn_scale;
last_output = scale_attributes.outputs.OUT_0 = attn_scale_output;
auto scale_node = std::make_unique(std::move(scale_attributes), context);
sub_nodes.emplace_back(std::move(scale_node));
}
如果您想详细了解如何加速您的自定义 Transformer，请参阅 NVIDIA cuDNN 的官方文档。
cuDNN 9 的其他显著特性
除了 SDPA 改进之外，cuDNN 9 还引入了其他几项重要改进：
- 混合输入精度支持矩阵和卷积
- 改进错误报告
- 硬件前向兼容性
- 简化安装
混合输入精度支持矩阵和卷积
要求输入操作数的数据类型为同一类型 (FP16、FP32) 的矩阵乘法 (matmul) 和卷积 API 不适合 AWQ (激活感知权重量化) 等情况，其中激活在 FP16 中，权重可能在 INT8 中。假设您想以更高的精度计算，则必须投射数据。如果不是在线的，这将需要额外的内存和转换成本。
cuDNN 现在支持混合输入精度矩阵和卷积，其中 A 和 B 操作数可以是不同的数据类型，并提供用于性能和内存优化的在线融合类型转换。您可以选择将操作数转换为其他类型。cuDNN 在优化的内核中处理所需的转换。
图 6 显示了 cuDNN 混合输入精度矩阵与未融合工作流程之间的加速。灰色条表示输入 A 和 B 分别采用 FP16 和 INT8 精度的情况，其中 A 转换为 INT8，然后使用 INT32 累加进行 INT8xINT8 矩阵乘法运算。绿色条表示 A 从 INT8 上转换为 FP16，然后使用 FP32 累加进行 FP16xFP16 矩阵乘法运算的情况。
改进错误报告
日志记录在软件开发中至关重要，对于使用 cuDNN 的深度学习框架等复杂系统尤其如此。过去，cuDNN 的一个常见痛点是难以调试错误和警告。我们一直在不断改进错误报告以解决这一问题。
cuDNN 9 增加了以下功能：
- 更具体的错误代码
- 进行分类以帮助组织数量增加的错误码
- 符合日志记录规范的嵌入式日志记录级别
cuDNNGetLastErrorString是一个新的函数，用于以编程方式检索最新的错误消息。
欲了解更多信息，请参阅 错误报告和 API 日志记录 部分内容，详见开发者指南。
硬件前向兼容性
在版本 9.0.0 之前，cuDNN 库在库发布之日支持最新的公开可用 GPU 架构。例如，cuDNN 8.9.x 通过 NVIDIA Hopper (即计算能力 9.0) 提供支持。不支持在未来的 GPU 架构上运行 8.9.x。
但是，cuDNN 9 对 API 的一大子集具有硬件前向兼容性。这意味着，仅将此 API 子集与 cuDNN v9 结合使用的程序将在未来的 GPU 上正常运行，并且这些程序的用户不会被强制升级其 cuDNN 安装以使用未来的 GPU。
在计算能力大于 9.0 的 GPU 上运行时，前向兼容性不是纠错，而是意味着库可以找到功能等效的实现，并使用 PTX JIT 针对新架构进行定位。
欲了解更多关于支持限制和最佳实践的信息，请参阅 硬件前向兼容性 部分内容，在开发者指南中。
简化安装
在 Python 环境中，除了库之外，您现在还可以使用 pip 安装新的 Python 前端：
# cuDNN library for CUDA 12
pip install nvidia-cudnn-cu12
# cuDNN frontend
pip install git+https://github.com/NVIDIA/cudnn-frontend.git
cuDNN 9 还简化了 RPM 和 Debian 元软件包的安装流程。例如，以下命令可在 Ubuntu 22 上安装 cuDNN：
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get install -y cudnn
这个新流程简化了安装 Debian 软件包时的密钥环设置，并使用 cudnn 元软件包，以提高安装效率。
如果您需要了解更多信息和完整说明，请参阅 cuDNN 安装指南，以获取详细的安装步骤和配置说明。
后续步骤
如果您有反馈、问题或意见，可以在 cuDNN NVIDIA 开发者论坛(适用于 cuDNN 库主题) 或 NVIDIA/cudnn-frontend 在 GitHub 上发布 (适用于前端主题)。另外，您可以 下载 cuDNN 开始使用。
请持续关注，了解未来更多新的 cuDNN 功能。本文中概述的进展是一个重要的里程碑，但未来还会有更多进展。
随着 AI 不断将行业推向软硬件集成的极限， NVIDIA 不断优化性能并改善用户体验，以便 cuDNN 能够在深度学习框架和图形编译器中得到更有效、更广泛的使用。
致谢
NVIDIA cuDNN 团队与公司内的许多其他团队紧密合作，为这篇博文提供了技术内容。
