CUDNN 9 （4） Graph API Engines

cuDNN Graph API 支持一组图形模式（graph patterns），这些模式由大量引擎支持。这些引擎分为四个不同的类别：

• pre-compiled single operation engines，预编译单操作引擎
• generic runtime fusion engines，通用运行时融合引擎
• specialized runtime fusion engines，专用运行时融合引擎
• specialized pre-compiled fusion engines，专用预编译融合引擎

专用引擎（specialized engines），无论是使用运行时编译还是预编译，都是针对一些常见且重要的例子（例如：transformer），当前支持的模式集相当有限，但可能会不断扩充。接下来分别介绍四个类别的Engine。

Pre-Compiled Single Operation Engines

这是最基本的一类Engine，仅支持单算子的op graph，专注于提高常见深度学习单一操作的效率，比如卷积、激活函数等，也就是说不会进行算子融合。举例说明：

ConvolutionBwData

ConvolutionBwData 计算张量 dy 的卷积数据梯度（convolution data gradient）。此外，它还使用缩放因子 ɑ、ꞵ 并将此结果与先前的输出混合。

ConvolutionBwFilter

ConvolutionBwFilter 计算张量 dy 的卷积滤波器梯度（convolution filter gradient）。

ConvolutionFwd

ConvolutionFwd 计算输入 X 与滤波器数据 W 的卷积。

NormalizationBackward

NormalizationBackward 计算梯度 dX 以及尺度和偏差梯度（bias gradients） dScale 和 dBias。预编译引擎（precompiled engines）支持instance norm backward和layer norm backward，而batch norm backward由专门的运行时编译引擎（specialized runtime compiled engine）支持。前向训练过程中保存的均值和方差将作为 NormBackward 操作的输入传递。

NormalizationForward

NormalizationForward 根据输入 X 计算归一化输出 Y。此操作用于推理和训练阶段。

Generic Runtime Fusion Engines

为了使算子融合变得有趣，Graph需要支持多种操作。理想情况下，我们希望支持的模式能够灵活地覆盖不同的用例。为了实现这种通用性，cuDNN 具有运行时融合引擎（runtime fusion engines），可以在运行时根据图模式生成内核（kernel）（或多个内核）。本节概述了这些运行时融合引擎支持的模式（即具有 CUDNN_BEHAVIOR_NOTE_RUNTIME_COMPILATION 行为注释的引擎）。

Support Patterns

主要的通用模式有如下几种，其结构图如下图所示：

• ConvolutionFwd fusions: $g\; 2\; =\; (\; Y\; =\; c\; o\; n\; v\; o\; l\; u\; t\; i\; o\; n\; F\; w\; d\; (\; X\; =\; g\; 1\; (\; i\; n\; p\; u\; t\; s\; )\; ,\; W\; )\; ,\; i\; n\; p\; u\; t\; s\; )\; g\_2=(Y=convolutionFwd(X=g\_1(inputs),W),inputs)$g2=(Y=convolutionFwd(X=g1(inputs),W),inputs)
• ConvolutionBwFliter fusions: $g\; 2\; =\; (\; d\; w\; =\; c\; o\; n\; v\; o\; l\; u\; t\; i\; o\; n\; B\; w\; F\; i\; l\; t\; e\; r\; (\; d\; y\; =\; g\; 1\; (\; i\; n\; p\; u\; t\; s\; )\; ,\; X\; )\; ,\; i\; n\; p\; u\; t\; s\; )\; g\_2=(dw=convolutionBwFilter(dy=g\_1(inputs),X),inputs)$g2=(dw=convolutionBwFilter(dy=g1(inputs),X),inputs)
• ConvolutionBwData fusions: $g\; 2\; =\; (\; d\; x\; =\; c\; o\; n\; v\; o\; l\; u\; t\; i\; o\; n\; B\; w\; D\; a\; t\; a\; (\; d\; y\; ,\; W\; )\; ,\; i\; n\; p\; u\; t\; s\; )\; g\_2=(dx=convolutionBwData(dy,W),inputs)$g2=(dx=convolutionBwData(dy,W),inputs)
• MatMul fusions: $g\; 2\; =\; (\; C\; =\; M\; a\; t\; M\; u\; l\; (\; A\; =\; g\; 1\; (\; i\; n\; p\; u\; t\; s\; )\; ,\; B\; )\; ,\; i\; n\; p\; u\; t\; s\; )\; g\_2=(C=MatMul(A=g\_1(inputs),B),inputs)$g2=(C=MatMul(A=g1(inputs),B),inputs)
• Pointwise fusions: $g\; 2\; (\; i\; n\; p\; u\; t\; s\; )\; g\_2(inputs)$g2(inputs)

$g\; 1\; g\_1$g1 And $g\; 2\; g\_2$g2

其中，g1是一个有向无环图（directed acyclic graph ，DAG），能够由0个或者多个以下的operation组合而成：

• CUDNN_BACKEND_OPERATION_CONCAT_DESCRIPTOR，表示连接（Concatenation）操作。连接操作通常用于将多个张量沿着指定的轴或维度合并成一个更大的张量。
• CUDNN_BACKEND_OPERATION_SIGNAL_DESCRIPTOR，表示信号（Signal）操作。cuDNN 后端信号操作描述符指定用于更新或等待标志变量（flag variable）的操作节点。信令操作可用于在 cuDNN operation graphs之间进行通信，甚至可以与另一个 GPU 中的operation graphs进行通信。
• CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR，表示逐点（Pointwise）操作。逐点操作是一种按元素进行的操作，即对张量中的每个元素应用相同的操作，而不考虑元素之间的位置关系。逐点操作通常用于实现各种激活函数、归一化操作、元素级别的操作等。

g2同样是一个有向无环图（directed acyclic graph ，DAG），能够由0个或者多个以下的operation组合而成：

• CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR，表示逐点（Pointwise）操作，同上。
• CUDNN_BACKEND_OPERATION_RESAMPLE_FWD_DESCRIPTOR，表示前向重采样操作，指定前向重采样的操作节点，根据不同模式下重新采样的image tensor X计算output tensor Y的值。
• CUDNN_BACKEND_OPERATION_RESAMPLE_BWD_DESCRIPTOR，表示反向重采样操作，指定用于反向重采样的操作节点，根据不同的模式重新采样的输出张量梯度 dy 计算输入张量梯度 dx。
• CUDNN_BACKEND_OPERATION_GEN_STATS_DESCRIPTOR，表示将生成每个通道统计信息的操作。
• CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR，表示归约（Reduction）操作，该操作节点在一维或多维中实现输入张量 X 的归约（降维）。
• CUDNN_BACKEND_OPERATION_SIGNAL_DESCRIPTOR表示信号（Signal）操作，同上。

另外，g1和g2在某些情况下没有完全适用，存在一些限制（limitations），具体的内容参考官方文档。例如GPU compute capability过低；Signal与ruduction操作不能同时存在等。

Examples of Support Patterns

Pointwise Operations After Convolution 2

g2由一组连续的多个逐点操作（pointwise operations）组成 。

Pointwise Operations Before Matrix Multiplication

逐点操作也可以先于卷积或矩阵乘法，即g1由逐点操作组成。

Convolution Producer Node in Middle of DAG

下图中，g1为先于卷积的逐点操作的 DAG；g2是由两个逐点运算组成的DAG。注意，卷积与g2中的Pointwise:bias一同作为Pointwise:add的输入，而不是作为g2第一个节点，这是合理的。

另外，运行时融合引擎（runtime fusion engines）支持的通用模式中的每个操作都受到有关参数的一些特定约束，这些约束是对后端描述符（Backend Descriptor Types）的限制以及上一节提到的limitations的补充，具体也参考官方文档。

Specialized Runtime Fusion Engines

专门的运行时融合引擎（specialized runtime fusion engines）针对并优化当前流行的深度学习模型中常见的专门图形模式（specialized graph patterns）。这些Engines在支持的融合模式（fusion patterns）、支持的数据类型（data types）和支持的张量布局（tensor layouts）方面提供有限的灵活性。但是从长远来看，专用模式预计将更加具有通用性，因为他针对特定的模型实现的高性能的算子融合。

BnAddRelu

在类似 ResNet 的视觉模型中，批量归一化（batch normalization）后进行 ReLU activation是一种常见的模式。使用运行时编译引擎支持的 BNAddRelu 融合模式旨在优化此重复operation graph。它还支持单节点多 GPU 批量归一化，以加速多 GPU 系统中的批量归一化计算。该模式用于训练阶段的前向传播。

在单节点多 GPU 批量归一化的情况下，每个 GPU 根据其输入数据计算本地统计（local statistics）数据，并将本地统计数据写入peerTensors（peerTensors 指代在不同设备之间共享数据的张量）。每个peerTensors驻留在节点上的单独 GPU 上，用于从peers GPU读取和写入本地统计信息。接下来是全局统计计算阶段，其中每个 GPU 聚合（aggregates）来自其他peers GPU的统计数据，并计算其本地数据的批量规范输出计算的全局均值和方差。

DReluForkDBn

与 BnAddRelu 模式类似，DReluForkDBn 模式也针对类似 ResNet 的视觉网络模型。旨在用于训练阶段的反向传播。 DReluForkDBn 模式是通过运行时编译引擎支持的，该Engine通常是对 BnAddRelu 模式的补充。它还支持单节点多 GPU 批量归一化，以加速多 GPU 系统中的批量归一化反向计算。

Fused Attention fprop

多头注意力（Multi-head Attention）的前向传播，公式：
$$O\; =\; m\; a\; t\; m\; u\; l\; (\; S\; =\; g\; 4\; (\; P\; =\; m\; a\; t\; m\; u\; l\; (\; Q\; ,\; g\; 3\; (\; K\; )\; )\; ,\; V\; )\; )\; O=matmul(S=g\_4(P=matmul(Q,g\_3(K)),V))$$O=matmul(S=g4(P=matmul(Q,g3(K)),V))

这是Transformer模型中的重要模式，可以用于BERT、LLAMA等大模型。

其中，g3可以是空graph，或者是单比例操作（single scale operation）。

g4可以为空或下图 cuDNN 操作的 DAG 的组合。每个 DAG 都是可选的，如图中虚线所示：

组合必须遵循呈现它们的顺序。例如，如果要使用 padding mask 和 softmax，则 padding mask 必须出现在 softmax 之前。

这些操作在Attention中常用。在下图中，我们描述了如何为每个操作创建 DAG。

Padding Mask

Causal Mask

Softmax

Dropout

g4 能够将中间张量存储到标记为 S 的全局内存中，可用于融合multi-head attention bprop。 DAG:Softmax 和 DAG:Dropout 都具有此功能。将 S 设置为图中最后一个 DAG 的输出。

Fused Attention bprop

多头注意力（Multi-head Attention）的前向传播，公式：
$$d\; V\; =\; m\; a\; t\; m\; u\; l\; (\; g\; 5\; (\; S\; )\; ,\; d\; O\; )\; d\; S\; =\; m\; a\; t\; m\; u\; l\; (\; d\; O\; ,\; V\; T\; )\; d\; Q\; =\; m\; a\; t\; m\; u\; l\; (\; g\; 6\; (\; d\; S\; )\; ,\; K\; )\; d\; K\; =\; m\; a\; t\; m\; u\; l\; (\; Q\; ,\; g\; 7\; (\; g\; 6\; (\; d\; S\; )\; )\; )\; dV=matmul(g\_5(S),dO)\backslash \backslash dS=matmul(dO,V^T)\backslash \backslash dQ=matmul(g\_6(dS),K)\backslash \backslash dK=matmul(Q,g\_7(g\_6(dS)))$$dV=matmul(g5(S),dO)dS=matmul(dO,VT)dQ=matmul(g6(dS),K)dK=matmul(Q,g7(g6(dS)))

g5目前支持固定的DAG：
g6表示softmax和masking的反向传递，得到dP：

g 7 is the transpose of dP the output of g 6.

Fused Flash Attention fprop

cuDNN 支持Fused Flash Attention来执行 GPT、BERT 等模型中常用的scale dot product attention。该Engine支持的一般模式是 BMM-Softmax-BMM 以及许多可以选择的可选功能。

Fused Flash Attention bprop

cuDNN 支持用于Fused Flash Attention的相应反向传播graph。这可以与 fprop graph一起使用来对大语言模型 (LLM) 进行训练。

Specialized Pre-Compiled Engines

预编译的专用引擎（pre-compiled specialized engines）针对专用图形模式进行优化。由于目标明确，这些graph不需要运行时编译。

目前，该引擎支持以下模式。有些节点是可选的。可选节点由虚线轮廓表示。

ConvBNfprop

ConvBNwgrad

ConvBiasAct

ConvScaleBiasAct

DgradDreluBNBwdWeight

FP8 Fused Flash Attention

Forward Pass

Backward Pass

总结

本文介绍了四个不同类别的Engine，cuDNN针对不同的应用场景对通用\特定的模型进行了优化，使得运行模型时能够通过启发式查询或者指定的Engine去进行相应的算子融合等操作，以达到高性能的计算。