Flash-DWC：在GPU上实现深度卷积计算的高效化

《ACM Transactions on Architecture and Code Optimization》：Flash-DWC: Making Depthwise Convolution Compute-Efficient on GPUs

　　

摘要

深度卷积（DWC）是高效卷积神经网络（CNN）中的关键算子。近期模型越来越多地采用大卷积核深度卷积来捕获长程依赖，其精度可与视觉Transformer相媲美。然而在GPU上，DWC的效率远未达到其低浮点运算量（FLOPs）所暗示的水平，可能成为运行时瓶颈，原因在于算术强度低、与Tensor Core不匹配以及内存访问效率低下。为解决这些问题，我们提出了Flash-DWC——一个使DWC在计算上高效的GPU库。Flash-DWC引入了一种基于GEMV的数据流协同设计，联合优化数据重用、访问顺序、布局和计算块划分。该设计提高了算术强度，同时保持了访问的连续性和合并式128字节内存访问。Flash-DWC进一步将DWC的GEMV转换为紧凑的高密度GEMM，以便在Tensor Core上实现高效的混合精度加速。此外，矩阵压缩、双路径分拆GEMM流水线和基于寄存器的输入重用进一步释放了片上资源，隐藏了内存延迟并提高了占用率。我们将这些技术扩展到前向和反向传播，以实现高效的端到端训练。与最快的cuDNN DWC算法相比，Flash-DWC的基于GEMV的CUDA-Core内核和基于GEMM的Tensor-Core内核分别实现了5.09倍和9.01倍的平均加速。在CNN训练中，Flash-DWC实现了1.03倍至14.18倍的加速，并在ImageNet-1K上达到与PyTorch一致的收敛效果。通过使DWC在计算上高效，Flash-DWC促进了更大、更具表达能力卷积核的使用。

1 引言

对于卷积神经网络（CNN）而言，在保持表示能力的同时降低计算成本一直是一个长期目标。为此，许多高效CNN[12, 13, 39, 40]用深度卷积（DWC）替代标准卷积，在保持空间感受野的同时大幅降低FLOPs。自2022年以来，一波新的架构[15, 25, 26, 29, 41, 44, 46, 47, 48]进一步利用大卷积核DWC（例如7×7、31×31）来捕获长程依赖，以低于视觉Transformer[9, 22, 23, 43]的FLOPs实现了与之竞争甚至更优的精度。除ImageNet分类外，大卷积核DWC在目标检测和语义分割等下游任务上带来更大收益[2, 7, 8, 17, 20]，凸显了其作为可迁移骨干网络的价值。

然而，DWC的低FLOPs并不能保证在GPU上获得成比例的加速。随着DWC的FLOPs随卷积核尺寸平方增长，这一问题变得更加突出。我们的PyTorch剖析显示，对于所评估的卷积核大于7×7的工作负载[7, 20, 25, 41]，DWC算子可占据每个batch训练时间的一半以上。这种运行时开销限制了更大、更具表达能力卷积核的实用化。近期工作[8, 17, 20]通过分解或稀疏卷积核来缓解此问题，但此类近似往往牺牲表示能力，而密集卷积核通常更为精确[2]。

这种低效源于DWC与GPU架构之间的根本性不匹配：（i）算术强度低：每个输入元素参与的计算次数远少于标准卷积，导致计算单元利用率不足；（ii）与Tensor Core不匹配：DWC本质上是一系列小型GEMV的集合而非大型GEMM，难以利用高吞吐量的矩阵单元；（iii）内存访问效率低下：DWC仅沿空间维度进行累加，破坏了通道方向的连续性并降低了内存合并度。因此，PyTorch[36]和cuDNN[4]等主流库仅利用了GPU计算能力的一小部分来处理DWC。先前的优化部分缓解了这些问题，但局限性依然存在。基于CUDA Core的方法[27, 28, 32, 34, 37]在混合精度吞吐量上受限；基于GEMM或Tensor Core的方法[32, 34, 38]往往引入冗余或稀疏模式。此外，一些方法针对较窄的设置（如小卷积核）进行调优，或未能充分利用软件流水线和内存合并。

我们不是重新设计模型，而是提出以下问题：DWC本身能否在GPU上变得计算高效？我们的核心见解是，加速DWC的关键不在于减少FLOPs，而在于隐藏内存延迟以充分利用计算单元。为此，我们提出Flash-DWC——一个以统一方式解决DWC瓶颈的GPU库：（i）基于GEMV的数据流协同设计联合优化数据重用、访问顺序、布局和计算块划分，从而提高算术强度，同时保持访问连续性和128字节合并访问；（ii）紧凑变换将DWC的GEMV映射为高密度GEMM，实现高效的Tensor Core利用。结合资源高效的内核优化，这些技术将DWC从内存受限转向计算受限，并支持完整的训练流水线。Flash-DWC已集成到PyTorch中，其源代码可在我们的GitHub仓库中获取。

我们在RTX 4090、A100 80GB和RTX 5880 Ada GPU上评估了Flash-DWC。与最快的cuDNN DWC算法相比，基于GEMV的CUDA-Core内核实现了5.09倍平均加速；基于GEMM的Tensor-Core内核进一步实现了相比cuDNN平均9.01倍、相比CUDA-Core内核平均2.01倍的加速。在端到端训练中，Flash-DWC实现了1.03倍至14.18倍的加速，同时在ImageNet-1K[6]上达到与PyTorch一致的收敛效果。

我们的贡献总结如下：

——基于GEMV的数据流协同设计。我们利用DWC GEMV中的卷积核重用以及重叠补丁之间的宽度优先输入重用。这将算术强度从朴素DWC的1提高到CUDA Cores上的15以上，并通过GEMV到GEMM的变换超过40，同时保持访问连续性和内存合并。

——紧凑的GEMV到GEMM变换。我们将DWC GEMV映射为密度达56.25%的GEMM，是朴素填充的4.5倍。该变换支持多通道、大卷积核工作负载，并与我们的基于GEMV的协同设计无缝集成，从而在Tensor Cores上实现混合精度加速。

——资源高效的内核优化。对于Tensor-Core内核，双路径分拆GEMM管线结合实时矩阵压缩，将50%的计算与数据加载重叠，同时将共享内存使用量大幅削减以提高占用率和L1容量。对于CUDA-Core内核，基于寄存器的输入重用同样最小化共享内存使用，释放片上内存。

——优化的训练流水线。针对卷积核梯度计算，我们通过累加轴并行化和灵活块划分专门解决低并行度和不规则输出形状问题，混合累加加速了混合精度执行。内核融合减少了小卷积核场景下的数据搬运。

2 背景与设计原则

为阐述我们的优化方法，我们首先形式化标准卷积和深度卷积，分析DWC的瓶颈，然后确立我们的初步设计选择。

2.1 标准卷积与深度卷积

为形式化描述二维卷积，我们定义以下符号：N表示批量大小；IH、IW和IC分别表示输入高度、宽度和通道数；OH、OW和OC分别表示输出高度、宽度和通道数；FH和FW分别表示卷积核高度和宽度。

标准二维卷积遵循全对全模式。它将输入特征图X∈RN×IH×IW×IC与卷积核W∈ROC×FH×FW×IC进行卷积，生成输出特征图Y∈RN×OH×OW×OC。在所有N×OH×OW的批空间位置处，每个输入补丁与OC个卷积核中的每一个执行FH×FW×IC的内积运算。这允许通过Im2col自然地映射到通用矩阵乘法（GEMM），其中一个矩阵包含所有展开的输入补丁，另一个包含所有展平的卷积核[4]。

深度二维卷积（DWC）遵循全对一模式。它将X与卷积核W∈ROC×FH×FW进行卷积生成Y。每个输入通道有专属的卷积核用于其输出通道，因此IC等于OC。在所有批空间位置处，每个输入补丁与它的专属卷积核执行FH×FW内积。因此，每个通道的操作自然地映射为通用矩阵-向量乘法（GEMV），如图1所示，其中矩阵包含该通道中所有展开的输入补丁，向量是展平的卷积核。

图1. 深度卷积。左侧展示了三通道的DWC：每个通道的操作可通过Im2col转换为GEMV。右侧详细说明了通道0的Im2col变换：朴素地将每个展平卷积核片段填充为8×4矩阵用于Tensor-Core GEMM会产生至少78%的稀疏度；虚线框表示对应前4个卷积核元素的GEMM。

2.2 深度卷积的性能瓶颈

我们从三个方面分析DWC的性能瓶颈。

算术强度。标准卷积（GEMM）在O(n2)数据上执行O(n3)计算，每个矩阵行在其他矩阵的所有列之间重用。相比之下，DWC（GEMV）在O(n2)数据上执行O(n2)计算，每个矩阵行仅使用一次。这种数据重用的根本差异导致了算术强度的巨大差距。具体来说，DWC和标准卷积的理论上限分别为ηD=2N·OH·OW·OC·FH·FW/(N·IH·IW·IC + OC·FH·FW)和ηS=2N·OH·OW·OC·FH·FW·IC/(N·IH·IW·IC + OC·FH·FW·IC)。由于分辨率和批量大小始终很大，这简化为ηD≈(1/IC)·ηS。由于CNN中IC通常较大（例如64到1024），DWC的算术强度比标准卷积低数个数量级。这量化了DWC为何仍处于内存受限状态且无法充分利用计算单元。

与Tensor Core不匹配。在现代GPU上，Tensor Core是专用于混合精度GEMM的单元，其峰值吞吐量远高于通用CUDA Cores。Tensor Core要求GEMM维度满足m≥8、n≥8和k≥4。标准卷积通过其具有大通道数的GEMM结构自然地满足这些约束。然而，DWC是一系列GEMV的集合，每个GEMV对于Tensor Core GEMM的维度来说都太小。如图1所示，朴素地将每个展平卷积核片段填充为8×4虽然满足了最小m8n8k4 GEMM形状，但浪费了7/8的矩阵列。这导致大量不必要的计算、冗余寄存器、共享内存使用以及潜在的warp分歧。因此，DWC的朴素Tensor Core映射相对于CUDA Cores无法获得有意义的加速。这促使我们提出了第4节中介绍的高效GEMV到GEMM变换。

内存访问。在标准卷积中，累加发生在FH×FW×IC维度上，这些维度完全连续且可以合并访问。例如，在NHWC布局中，IC始终是32的倍数，因此W和X的加载可以沿IC轴进行128字节向量化，最大化内存事务粒度。相比之下，DWC仅沿FH×FW进行累加。在NHWC格式中，空间位置与IC步幅交错，破坏了访问连续性。此外，在NCHW和NHWC布局中，W的维度顺序为OC×FH×FW，其中FH和FW通常为奇数且未与128字节对齐。因此，DWC难以充分饱和内存带宽。

2.3 设计原则

为克服这些瓶颈，我们旨在隐藏内存延迟并充分利用计算单元，遵循内存感知内核中的常见原则[5, 10]。这转化为两个方向：（i）最大化片上数据重用以最小化片外内存流量；（ii）通过高效布局和访问模式最大化带宽。为实现这些目标，我们首先选择卷积方式和张量布局。

直接卷积方法。这指的是直接计算DWC，而不是使用FFT或Winograd域变换。虽然FFT[31]和Winograd[16]可以减少FLOPs，但它们引入了额外变换、中间存储、数据搬运和有限精度舍入。这种权衡对DWC不利：计算量的减少仅带来边际收益，而额外的数据搬运加重了主要的内存瓶颈，可能抵消FLOP节省，最终降低吞吐量。直接卷积保留了最直观的数据流，使我们能够将片上资源（如寄存器、共享内存）专门用于最大化数据重用和带宽。

统一的通道最后布局。我们将卷积核从传统布局W∈ROC×FH×FW重排为W∈RFH×FW×OC；输入和输出特征图存储为X∈RN×IH×IW×IC和Y∈RN×OH×OW×OC。对于DWC，IC=OC，因此所有张量都将通道轴置于最后。这使得卷积核和特征图都能沿通道轴实现合并访问。由于X和Y保持NHWC布局，Flash-DWC与主流库的相邻操作（如归一化和逐点卷积）自然互操作。

3 基于GEMV的深度卷积数据流协同设计

基于直接卷积和统一通道最后布局，我们首先在CUDA Cores上以FP32精度优化DWC的原生GEMV形式，然后将其扩展到混合精度Tensor Cores。这种分离是刻意为之：它隔离了基于GEMV的优化的贡献，并为后续GEMV到GEMM变换奠定基础。

我们的设计以数据重用为中心，但其有效性依赖于布局、计算块划分和访问顺序：（i）计算块划分创造重用机会，并利用统一的通道最后布局实现128字节访问；（ii）卷积核重用利用空间位置之间共享的卷积核；（iii）宽度优先输入重用保持重叠补丁之间的访问连续性。

图2展示了我们如何将朴素DWC1的算术强度从1提升至超过15。图2. 基于GEMV的Flash-DWC内核的算术强度（η，每元素操作数）。（a）每项优化逐步带来的强度（η）增长。（b）最终强度与滤波器宽度FW的关系。3.1 工作负载分块我们将输出张量Y∈RN×OH×OW×OC划分为包含BN个连续通道和BM个批空间位置的分块，总共得到OCBN×OHOWBM个块。由于乘积OC×NO×HOW的变化幅度小于其各因子，这种分块方式能维持各层之间稳定的块数。我们为每个块配置4个warp以充分利用每个流式多处理器（SM）的4个warp调度器。在warp内部，我们将32个线程排列为8×4的模式：（i）8个线程使用128位指令执行128字节的合并加载，读取32个连续通道，最大化内存事务粒度；（ii）4个线程获取相邻的输入块，为输入复用创建重叠。为了最大化warp间的输入块重叠，我们将BN设为32，恰好等于每个warp的通道数，从而使全部4个warp完全覆盖批空间维度。3.2 通过GEMV实现滤波器复用我们利用DWC的GEMV结构来复用滤波器数据。如图3(a)所示，每个线程被分配一个跨4个通道的TM个连续空间位置的块。在每个通道内，所有TM个位置共享同一个滤波器，因此加载一个滤波器元素即可对对应的输入元素执行TM次融合乘加（FMA）操作。这将算术强度提升至2TM/(TM+1)（>1），且随TM增大而增长。过大的TM会减少活跃线程数并损害延迟隐藏能力，因此我们通过实验将TM设为8，得到约1.78的算术强度。图3. 单个通道内GEMV中的滤波器复用。（a）在线程内，一个滤波器元素跨TM个不同空间位置的内积被复用；xi,j(t)表示第t个内积中滤波器偏移量(i,j)处的输入元素。（b）在块内，512/BN个线程协作将滤波器片段加载到共享内存（SMEM）中，然后复用该共享片段来计算各自的TM个内积。为了进一步提高滤波器复用率，我们应用块级协作。如图3(b)所示，对于每个通道，所有512/BN个线程共享同一个滤波器。因此，它们协作将一个滤波器片段加载到共享内存（SMEM）中，从而将多个线程级GEMV合并为一个更大的块级GEMV。此时算术强度上升至2TM/(TM+BN/512)。当BN=32且TM=8时，该值约为1.98，接近理论最大值2.3。3.3 跨重叠块的输入复用虽然滤波器复用已最大化，但输入复用仍是主要瓶颈。与滤波器不同，输入块的占用空间随空间维度和批大小扩展，因此更难保留在片上内存中。幸运的是，相邻的输入块有相当大的重叠，特别是在步长为1的卷积中。现代CNN通常仅在阶段边界处使用相对较少的步长为2的卷积进行下采样，且往往使用较小的滤波器。在每个阶段内，重复的步长为1的DWC块保持空间分辨率，在相同尺度上细化特征，并通常使用较大的滤波器以扩展感受野。因此，步长为1的DWC主导了许多CNN工作负载，使得输入重叠复用具有广泛适用性。在每个通道内，我们评估了图4中展示的两个复用方法。（i）矩形复用利用高度和宽度两个方向的重叠，提供更高的理论复用率。但它从由IC×IW间隔的非连续行中加载元素，且操作模式在不同迭代中变化。跨高度轴迭代复用数据需要O(TM×FW)的额外存储来保留多行数据，或者在维度间重新排序遍历，这会破坏内存连续性。（ii）宽度优先复用虽然仅限于宽度轴的重叠，但能保持完全连续的内存访问且对缓存友好。在初始批量加载后，后续迭代操作统一，保持的数据简单地移位并用于下一次迭代。图4. 单个通道内的矩形复用和宽度优先复用方法。此示例展示了3×3滤波器的第0步到第8步的累加过程。蓝色表示当前迭代加载的元素；灰色表示从上一次迭代复用的元素；黄色标记的是保留但在此迭代中未使用的元素。给定我们的配置TM=8，宽度优先复用加载Nw=FW×(FH+8-1)个输入元素，而矩形复用（2×4）加载Nr=(FW+2-1)×(FH+4-1)个元素，且需要至少多25%的存储用于数据复用。当FH,FW≥3时，Nw/Nr的比率最多为1.25，意味着宽度优先复用最多产生25%的额外加载，同时保持了简单性和访问连续性。因此我们选择宽度优先复用。为实现它，我们将任意2D滤波器∈RFH×FW分解为FH个1D单元∈RFW，对每个单元执行1D DWC，然后累加结果。该方法可通过将更高维度的DWC视为一系列1D DWC来推广。我们使用寄存器而非共享内存来实现宽度优先复用。寄存器提供最低的延迟操作数访问并直接为计算流水线供数，而共享内存约有30周期的延迟且与L1缓存共享资源。最小化共享内存使用可以保留L1容量并允许每个SM上有更多的活跃块，从而增强占用率。由于每个线程维护私有的输入复用窗口，使用共享内存暂存会引入额外的加载/存储指令和地址计算开销。由于每个线程在每个通道内处理TM个相邻输入块，我们分配4×TM个寄存器用于累加，另外分配4×TM个寄存器用于缓存输入元素。此寄存器分配恰好满足每个GEMV的需求，并且与cuBLAS GEMM内核的寄存器占用相匹配，维持高占用率。3.4 基于GEMV的内核工作流我们将上述设计整合为连贯的工作流，如图1所示。主循环在两个关键函数Prefetch和GEMV之间交替执行。Prefetch：每个块加载一个FW×BN的滤波器块和一个BM×BN的输入块，其中每个线程加载跨TM个空间位置和4个通道的TM×4个输入元素。滤波器加载使用cp.async.cg PTX指令，该指令发起异步的全局到共享内存拷贝，具有显式的提交/等待控制。这避免了寄存器中转暂存并支持更细粒度的软件流水线，相比标准加载性能提升约8%。对滤波器共享内存数组W进行双缓冲可将Prefetch与GEMV重叠，再获得约5%的性能提升。为避免银行冲突，我们将W的最后一个维度填充128位。GEMV：每个分解后的1D DWC作为具有FW次累加迭代的GEMV执行，这些迭代通过C++模板在编译时展开。从第二次迭代开始，每个线程复用其寄存器缓存X中的4×(TM-1)个元素，并仅加载4个新的输入元素，再次每次迭代执行TM次FMA。结合滤波器复用后，算术强度变为2TM/(1+(TM-1)/FW+BN/512)，该值随FW增大而增大，如图2(b)所示。当BN=32且TM=8时，强度在FW→+∞时接近15.06。对于CUDA Core路径，当FW≤51时，缓存滤波器数据使用的共享内存不到可用量的10%，将大部分片上内存留作L1缓存。Nsight测量表明，当FW≤9时L1缓存命中率可超过80%，当FW≤51时仍保持在50%以上。由于L1带宽是全局内存的10倍以上，实际算术强度远高于仅根据全局内存访问次数估算的值。4 基于GEMM的混合精度深度卷积虽然基于GEMV的优化显著提高了FP32吞吐量，但性能最终受限于CUDA Core峰值。为启用更高效的混合精度DWC，我们利用Tensor Core，这需要根本性的转变：将DWC原生的GEMV转化为GEMM。基于第3节中的基于GEMV的设计，我们将其适配为FP16/FP32混合精度，这是DNN训练中的标准做法。该扩展依赖于三项关键创新：（i）56.25%密度的GEMV到GEMM变换，专为DWC的多通道、大滤波器场景定制；（ii）通过矩阵压缩和宽度优先输入复用来最小化GEMM冗余，大幅减少共享内存使用并提高算术强度；（iii）双路径分割GEMM流水线，在保持复用模式的同时将50%的计算与数据加载重叠。图5展示了我们如何将算术强度从1提升至超过40。图5. 基于GEMM的Flash-DWC的算术强度（?，每元素操作数）。（a）每项优化逐步带来的强度（?）增长。（b）最终强度与滤波器宽度FW的关系：Flash-DWC表示基于GEMV的CUDA Core内核；Flash-DWC-T1、-T2和-T4是基于GEMM的Tensor Core内核。4.1 面向Tensor CoreFP16/FP32混合精度使用FP16存储张量并执行乘法以提高效率，同时使用FP32进行累加以保证数值稳定性。对于DWC，FP32累加是必需的：由于精度和动态范围有限，纯FP16累加可能导致不可接受的精度损失。当张量从均匀分布U(-1,1)中采样时，FP16累加DWC的平均绝对相对误差（MARE）相对于FP64基准值超过1%，有时甚至超过10%；而FP16/FP32混合精度将MARE维持在10?3量级，满足典型训练要求。这表明FP32累加能有效控制数值误差；剩余的挑战是高效地在GPU上执行混合精度DWC。然而，FP16/FP32 DWC在CUDA Core上面临硬件限制。大多数GPU架构（如Ampere、Ada、Hopper）不在CUDA Core上提供原生的混合精度FMA，因此每次FMA都需要显式的FP16到FP32转换。这使吞吐量受限于FP32 CUDA Core上限，低于某些GPU上的FP16峰值；同时也在数据类型转换上浪费了计算周期，使得对于大滤波器，FP16/FP32 DWC比FP32 DWC更慢，如图6所示。图6. Flash-DWC内核的吞吐量对比。滤波器大小（FH=FW，x轴）范围为3到51。吞吐量（y轴）单位为TFLOPS。子图标题指示GPU及NHWC格式的特征图维度。Flash-DWC表示基于GEMV的CUDA Core内核；Flash-DWC-T1、-T2和-T4表示基于GEMM的Tensor Core内核。为克服CUDA Core的限制，我们面向Tensor Core，其原生支持更高吞吐量的FP16/FP32 GEMM。这也避免了FMA与辅助任务之间的资源竞争：Tensor Core处理FMA，而辅助任务（如地址计算、条件分支）仍由CUDA Core承担。此后，我们简称为FP16代表FP16/FP32混合精度，并将Tensor Core内核记为Flash-DWC-T。如图6所示，Flash-DWC-T显著优于FP16和FP32的CUDA Core内核。4.2 将GEMV转化为GEMM如第2.2节所述，朴素的将DWC的GEMV填充为GEMM会浪费矩阵的78列。类似的挑战也出现在模板计算中，近期工作[3, 19]将模板GEMV转化为密度超过50%的GEMM。然而，这些方法不适合DWC，因为底层张量形状存在根本差异。模板计算操作在单个通道上，滤波器大小适中，而DWC涉及数百或数千个通道（如1024、2048），并且可以使用31×31甚至更大的滤波器[2, 7, 20]。这些差异影响了滤波器和输入的变换。（i）滤波器变换：模板滤波器小且为单通道，因此变换后的滤波器矩阵可以完全驻留在常量内存缓存[3, 19]甚至寄存器[21]中。虽然变换本身复杂并引入矩阵冗余，但其开销可忽略不计：规模小，在内核启动前执行一次，且结果矩阵在多次启动间全局共享。相比之下，在DWC中，即使是3×3的多通道滤波器也可能超过64KB常量内存限制。此外，由于每个块处理128字节的通道段以最大化内存事务粒度，大于8×8的滤波器会溢出每个SM的8KB常量缓存。因此，DWC必须将滤波器分割为片段并将变换结果缓存在更大的共享内存中，在许多近期GPU上每个SM可提供128KB或更多。（ii）输入变换：模板方法通过收集不规则空间偏移处的输入元素来构建输入矩阵。这是可行的，因为所有访问的元素位于同一通道且在内存中彼此接近；偏移量也可以预计算到能完全放入缓存的查找表中。然而在DWC中，大量通道在空间位置之间引入了很长的内存步长，使得这种不规则的、基于查找表的收集效率低下。为了最大化带宽，DWC需要在通道轴上进行合并加载，在宽度轴上进行连续访问。DWC实用的GEMV到GEMM变换必须满足三个要求：（i）将滤波器和输入块分割为片段，即时变换并将结果缓存在共享内存中，冗余最小化以释放片上资源；（ii）片段变换与GEMM计算重叠进行，保持变换简洁以减少GEMM间的延迟；（iii）与我们的宽度优先输入复用对齐，以增强算术强度并保持访问连续性。我们的变换将一批M个GEMV打包为GEMM。每个GEMV操作在TM个相邻输入块上，所有GEMV共享一个1D滤波器片段∈RFfrag。图7展示了该变换：图7. DWC的GEMV到GEMM变换。16个空间相邻的GEMV，每个涉及一个8×9矩阵和一个长度为9的向量，被打包为适合Tensor Core的m16n8k16 GEMM。

此处，GEMV[i] 表示打包批次中的第 i 个 GEMV。
(i) 滤波器变换：我们将滤波器片段扩展为滤波器矩阵 $\in \mathbb{R}^{K \times T_M}$，其中第 i 列矩阵放置从索引 i 开始的片段，并在其余位置填充零至长度 K，其中 $K=T_M+F_{\text{frag}}-1$。
(ii) 输入变换：对于每个 GEMV，我们将其 $T_M$ 个输入图块压缩为一个长度为 K 的向量，其中 K 是不同输入元素的总数。将所有 M 个此类向量按行堆叠，得到输入矩阵 $\in \mathbb{R}^{M \times K}$。
在生成的 GEMM 中，输入矩阵是稠密的，而滤波器矩阵的密度为 $\rho_{\text{FM}}=F_{\text{frag}}/K=1-(T_M-1)/F_{\text{frag}}+(T_M-1)/K$。
对于 FP16/FP32 GEMM，张量核心支持形状 m8n8k4、m16n8k8 和 m16n8k16。由于其算术强度较低，我们弃用了 m8n8k4。由于较大的 $F_{\text{frag}}$ 会带来更高的密度，我们选择 m16n8k16 并将 $\{M, T_M, F_{\text{frag}}\}$ 设置为 $\{16, 8, 9\}$，实现了 56.25% 的密度（4096 次中有 2304 次非零 FMA）。该密度是朴素填充的 4.5 倍，与模板变换 [3, 19] 相当。
保持第 3.3 节中的宽度优先分解，我们将任何 2D 滤波器分割为 $F_H \times \lceil F_W/9 \rceil$ 个宽度为 9 的片段，并累加它们的 1D DWC 结果。

### 4.3 GEMM 冗余最小化

使用与第 3.1 节相同的分块方法，我们将工作负载划分为 $O_C B_N N_O H_O W_B M$ 个块，每块包含 4 个线程束。由于 FP16 的数据量是 FP32 的一半，我们将每个线程处理的通道数加倍（从 4 增至 8），每个线程束处理的通道数也加倍（从 32 增至 64），以维持 128 字节的内存事务。相应地，我们将 $B_N$ 设置为 64，$B_M$ 设置为 128，这最大化了输入图块的重叠，同时匹配 m16n8k16 GEMM 的 16×8 图块大小。
将多个 GEMV 打包成 GEMM 将算术强度提升至 $4096 \cdot 0.5625 / (256+128)=6$，但这仍然低于我们的 CUDA 核心内核。此外，在共享内存（SMEM）中存储完整的滤波器和输入矩阵需要 $256 B_N$ 和 $4 B_N B_M$ 字节，这达到了 48 KB 的静态分配限制，可能会降低占用率和 L1 缓存容量。为了解决这两个问题，我们通过两种优化来减少 GEMM 冗余：(i) 在 SMEM 中压缩输入和滤波器矩阵，并在寄存器中即时重建它们以用于张量核心 GEMM；(ii) 适应宽度优先的输入重用，以最小化从全局内存加载矩阵数据。
在滤波器矩阵 $\in \mathbb{R}^{16 \times 8}$ 中，9 元素的滤波器片段在 8 列中重复，如图 8(b) 所示。因此，我们以非重复格式存储它：在片段两侧填充 7 个零，生成一个 23 元素的向量。为了重建第 i 列，一个块只需从索引 i 开始提取 16 个连续元素，无需线程束分歧。这将滤波器矩阵的存储成本降低至仅 46 $B_N$ 字节，即原始的 18%，将总 SMEM 使用量降至 34.9 KB。我们将填充的滤波器片段初始化为零，并让块中的线程协作将滤波器数据加载到 SMEM 中，以最大化数据重用。这将算术强度增强至 $4096 \cdot 0.5625 / (256+F_W/\gamma) > 8.69$，其中 $\gamma=\lceil F_W/9 \rceil$ 是每个滤波器行的宽度为 9 的片段数。

**图 8. 共享内存中的矩阵压缩。** (a) 输入矩阵：当 $\theta=4$ 且 $O_W \ge 32$ 时，四个连续行共享 3×8 个重叠元素，因此 16×16 输入矩阵被压缩为五个 4×8 缓冲区。(b) 滤波器矩阵：存储为 23 元素的向量以消除列冗余。

在输入矩阵 $\in \mathbb{R}^{16 \times 16}$ 中，当 $O_W \ge 16$ 时，相邻行沿宽度轴重叠 8 个元素，如图 8(a) 所示。因此，当 $O_W \ge 8\theta$ 时，我们将连续的 $\theta$ 行从 $16\theta$ 个元素压缩为 $8(\theta+1)$ 个元素，存储在 $\theta+1$ 个宽度为 8 的输入 SMEM 缓冲区中。这将算术强度提升至 $4096 \cdot 0.5625 / (128(1+\theta-1)+F_W/\gamma)$。我们将生成的张量核心内核（$\theta=1,2,4$）分别记为 Flash-DWC-T1、-T2 和 -T4。T2 和 T4 实现了 11.96 和 14.35 的算术强度，同时将 SMEM 使用量减少至 26.9 KB 和 22.9 KB。在对 SMEM 进行填充以避免银行冲突后，Flash-DWC-T 在每个 SM 上为并发块和 L1 缓存留出了超过 70% 的片上内存。
我们在 SMEM 中实现宽度优先的输入重用，以适应张量核心 GEMM。因为每个 GEMM 都在宽度为 9 的滤波器片段和输入矩阵的 16 元素行上操作，我们可以为下一个 GEMM 迭代重用每行 7 个输入元素。为了最大化重用并避免线程束分歧，前 $\theta-1$ 个缓冲区重用下一个缓冲区的 7 个元素，并从 X 中仅加载 1 个新元素，而最后一个缓冲区加载 8 个新元素。这将算术强度提升至 $4096 \cdot 0.5625 / (128\theta-1+112\gamma-1+16+F_W/\gamma)$，随着 $F_W$ 增加而增加，如图 5(b) 所示。对于 T4，当 $F_W \ge 3$ 时强度超过 13.6，当 $F_W \ge 11$ 时超过 20.4。随着 $F_W \to +\infty$，T2 和 T4 的强度分别趋近于 25.88 和 40.42，远高于 CUDA 核心内核的 15.06 限制。

### 4.4 基于 GEMM 的内核工作流

我们将基于 GEMV 的设计以及 GEMV 到 GEMM 的变换整合成 Flash-DWC-T 的连贯工作流，如图算法 2 所示。主循环（第 4-11 行）首先沿高度轴迭代。每个高度轴迭代包含 $\lceil F_W/9 \rceil$ 次宽度轴迭代，这些迭代通过 C++ 模板在编译时展开。由于每个块处理 $B_N$ 个通道，每次迭代执行一批 $B_N$ 个 m16n8k16 GEMM，每个通道一个。

**软件流水线。** 我们的 GEMV 到 GEMM 变换通过偏移映射和数据加载实现。为了使数据加载与 GEMM 计算重叠，我们将每个 $B_N$ 批处理的 m16n8k16 GEMM 拆分为两个 m16n8k8 子单元，记为 GEMM1 和 GEMM2，如图 9 所示。这种分解形成了一条双路径拆分 GEMM 流水线，由 Prefetch1→GEMM1 和 Prefetch2→GEMM2 组成。可以将其视为双缓冲的广义形式。然而，与分块 GEMM 内核中标准的两阶段乒乓缓冲不同，这两条路径是非对称的，并使用部分重叠的 SMEM 区域。具体来说，GEMM1 使用前 $\theta$ 个输入 SMEM 缓冲区和填充滤波器片段的前 8 个非零元素，而 GEMM2 使用最后 $\theta$ 个缓冲区和最后 8 个非零元素。

**图 9. 双路径拆分 GEMM 流水线。** 虚线框显示 $B_N$ 批次中一个代表性的 GEMM。每个 m16n8k16 GEMM 对应一个通道，并拆分为两个 m16n8k8 子单元。箭头表示数据依赖关系；跨迭代的同步通过重叠 SMEM 上的数据重用强制执行。

在每次迭代中，Prefetch1 和 Prefetch2 作为两组异步加载指令发出。一旦 Prefetch1 完成，GEMM1 就开始，同时 Prefetch2 并行进行。GEMM2 在 GEMM1 和 Prefetch2 都完成后开始。由于这两个子 GEMM 单元共享重叠的 SMEM 区域，下一次迭代的 Prefetch1 必须等待上一次的 GEMM2 释放这些区域。即便如此，该流水线在单个块内将 50% 的 GEMM 计算与数据加载重叠。此外，低的 SMEM 使用量使多个块能够在 SM 上并发执行，从而在块之间进一步将数据加载与 GEMM 计算重叠。

**数据加载。** 对于输入元素，我们使用 cp.async.ca PTX 指令，它将数据缓存在 L1 中，以重用 SM 内的输入图块重叠。滤波器元素没有空间重叠，因此我们用 cp.async.cg 加载它们，它仅在 L2 中缓存，从而为输入元素保留 L1 容量。在第一次宽度轴迭代之前，我们将输入的初始大批量元素加载到 SMEM 中（第 5 行）。在每次宽度轴迭代中：Prefetch1 重用输入 SMEM 缓冲区来准备 GEMM1 所需的 $(T_M-1)/T_M=7/8$ 的输入元素，并从全局内存加载剩余的 $(1/T_M)=1/8$。与此同时，Prefetch2 为 GEMM2 加载最后一个输入 SMEM 缓冲区。代替在输入 SMEM 缓冲区之间复制数据，我们循环重新编号缓冲区并相应更新它们的访问索引，如图 10 所示。

**图 10. 用于输入重用的循环缓冲区重新编号，无需显式数据复制。** 蓝色箭头表示移动缓冲区以重用数据。在每次迭代中，数字 0–7 表示每个缓冲区内的访问索引；新加载和重用的元素分别用红色和黑色字体显示。单元格颜色跟踪每个元素首次加载所在的缓冲区，重用的元素在迭代中保留其颜色。

输入特征图上的零填充是隐式处理的。配合谓词掩码，cp.async 跳过零填充区域并向 SMEM 写入零，从而避免线程束分歧。为了节省寄存器并增强活动线程数，我们将每批 32 个谓词掩码压缩为一个 32 位整数，每个空间位置使用 1 位。由于 $F_W$ 和填充大小是恒定的，宽度轴谓词掩码是预先计算好的，并在所有高度轴迭代中重用。此外，我们扫描每个块以检测是否所有线程在高度轴上共享相同的填充零。如果是这样，我们裁剪高度轴迭代以排除这些共享零，使用 $F_{Hs}(\ge 0)$ 和 $F_{He}(
**gemm 计算。** gemm1 和 gemm2 中的每一个都计算一批 $b_n$ 个通道级 m16n8k8 gemm。每个线程束在张量核心上执行 $b_n/4$ 条 mma.sync.aligned.m16n8k8 ptx 指令。为了避免银行冲突，我们将 smem 数组 w 和 x 的最后一个维度填充 128 位 [2]。在主循环之后，我们通过 smem 重排寄存器矩阵 v（第 12 行），并以 128 字节事务将生成的 $b_m \times b_n$ dwc 图块存储到全局内存中。

**为何不使用稀疏张量核心？** 虽然稀疏张量核心（例如 sparstencil [19] 中使用的）具有稠密张量核心两倍的理论吞吐量，但它们不适合 dwc。首先，它们支持的最小 gemm 形状 m16n8k16 将迫使我们放弃双路径拆分 gemm 流水线，该流水线隐藏了 50% 的 gemm 计算以位于数据加载之后。其次，即使我们重组数据流以满足 2:4 稀疏性，gemm 运行时减少最多为 43.75%，这还未考虑额外的开销和数据访问连续性的丧失。由于 dwc 本质上是内存受限而非计算受限，切换到稀疏张量核心将牺牲内存延迟隐藏以换取微乎其微的计算收益，最终降低整体性能。

### 5 扩展 flash-dwc 至 cnn 训练

基于第 3 节和第 4 节中的优化内核，我们将 flash-dwc 扩展至完整的 cnn 训练，并将其作为即用型库集成到 pytorch 2.10 [36] 中，flashdwc.ops 用于操作调用，flashdwc.nn 用于模型构建。

#### 5.1 反向传播

前向卷积（fc）直接由 flash-dwc 内核支持。训练额外需要两个反向算子：(i) 反向数据卷积（bdc），它将输出梯度 $\nabla y$ 与转置滤波器 $w^t$ 卷积以计算输入梯度 $\nabla x$；以及 (ii) 反向滤波器卷积（bfc），它将输入特征图 x 与 $\nabla y$ 卷积以计算滤波器梯度 $\nabla w$。由于 bdc 在旋转滤波器后具有相同的 dwc 结构，它可以重用 flash-dwc 内核。主要挑战在于 bfc：其小输出张量导致并行度不足，且不匹配原始内核的分块方法。图 11 展示了 flash-dwc 如何支持所有三个阶段。

**图 11. 用于 cnn 训练的 flash-dwc。** fc 和 bdc 利用张量核心和 cuda 核心内核。对于 bfc，我们通过将 $\nabla y$ 分块为 z 个子 bfc 并原子累加它们的结果来增加并行度。如图所示，每个 11×11 输出的子 bfc 生成 6 个块：mega 内核通过组合 tm=2 和 9 的两个内核来匹配 $f_w=11$；$\beta=4$ 最小化 $f_h$ 轴填充。

**bdc 实现。** 我们将 w 旋转 180° 以获得 $w^t$，然后在 $w^t$ 和 $\nabla y$ 上执行 flash-dwc 内核。由于旋转开销可忽略不计，bdc 的性能与 fc 相似，同样受益于张量核心。

**累加轴并行化。** bfc 的输出 $\nabla w \in \mathbb{r}^{f_h \times f_w \times o_c}$ 仍然很小，且不随批次大小或空间分辨率扩展。这与 fc 和 bdc 形成对比，后者的输出巨大，自然能提供丰富的并行度。如果我们仅通过分块 $\nabla w$ 来并行化 bfc，生成的块数量太少，无法充分利用硬件。
我们通过沿累加轴并行化来解决这个问题。具体来说，我们将 $\nabla y$ 分割为 z 个图块，并将每个图块作为独立的子 bfc 处理。这使并行度增加了 z 倍，但引入了通过 fp32 原子加法合并子 bfc 结果的额外开销。为了平衡这两个效果，我们根据 dwc 参数和 sm 数量动态决定 z。如果累加长度不能被 z 整除，我们将余数附加到最后一个图块上，并调整 z 以最小化此余数，实现工作负载平衡。我们进一步通过一次性的性能分析步骤来完善 z，该步骤测量性能并选择最佳值。

**灵活分块。** bfc 的不规则输出形状与原始 flash-dwc 内核的分块方法不匹配。使用 $b_m=128$，这些内核将空间维度分块为 16×8、8×16 或 4×32。这些形状非常适合特征图，通常很大且是 8 的倍数（例如 112×112, 64×64）。然而，bfc 的输出 $\nabla w$ 具有奇数大小空间维度，且较小或与那些图块对齐不良，导致大量的零填充和不必要的计算。

#### 6.1 卷积 gelu 与层归一化

flash-dwc 支持 cnn 中常见的三种算子：分组深度卷积（gwd）、卷积-gelu 和层归一化。我们将融合卷积-gelu 和融合卷积-gelu 层归一化内核分别标记为 flash-dwc-f1 和 flash-dwc-f2。我们遵循第 3 节和第 4 节中的内核设计，并采用以下两项优化。

**融合卷积-gelu。** 为了在单个内核中融合 gwd 和 gelu，我们在 gemv 内核的每行末尾插入一个 gelu 阶段，在 gemm 内核的每个 gemm 图块后插入一个 gelu 阶段。因为 cuda 核心计算和 gelu 使用相同的线程，所以不需要额外的线程同步。张量核心计算和 gelu 需要额外的线程同步，但我们使用第 4.4 节中的双路径拆分 gemm 流水线来隐藏此开销。因为 gwd 具有大的输出通道，所以融合 gelu 不会引入显著的算术强度损失。

**融合卷积-gelu 层归一化。** 为了在单个内核中融合 gwd、gelu 和 ln，我们将 ln 分解为两个步骤：(1) 在归一化之前计算每通道的均值和方差，以及 (2) 使用计算出的统计量来归一化激活值。我们在每个块内执行这两个步骤，因此每个块需要 $o(b_m b_n)$ 个额外寄存器来存储统计量。由于 flash-dwc 内核已经使用了大量寄存器（例如，t4 内核使用了 216 个寄存器中的 208 个），我们需要重用现有的寄存器来存储 ln 统计量。我们观察到 gemm 内核中的输入矩阵寄存器在 gemm 计算后不再需要，因此我们重用这些寄存器来存储 ln 统计量。

#### 6.2 分组深度卷积

flash-dwc 将 gwd 视为 m 个独立的 dwc 的并集，并分别处理它们。这导致每个 dwc 需要单独的内存事务，降低了内存带宽效率。为了提高内存带宽效率，我们引入了交错分组深度卷积（igdw），它将 m 个独立的 dwc 重组为单个 dwc。具体来说，我们将输入特征图的 m 个通道沿空间维度交错排列，使得每个空间位置包含所有 m 个通道的元素。这样，igdw 就可以使用单个 dwc 内核来处理所有 m 个通道，从而显著提高了内存带宽利用率。

**算法 3. gemv 内核**
```text
输入：输入特征图 x，滤波器 w，输出 o
1. 将 x 分块为 bm x bm 图块
2. 对于每个图块：
3. 将滤波器 w 加载到 smem 中
4. 对于每个空间位置：
5. 从全局内存加载 x 的一个通道
6. 在 smem 中执行 gemv
7. 将结果写入 o
```

**算法 4. gemm 内核**
```text
输入：输入特征图 x，滤波器 w，输出 o
1. 将 x 分块为 bm x bn 图块
2. 对于每个图块：
3. 将滤波器 w 加载到 smem 中
4. 对于每个空间位置：
5. 从全局内存加载 x 的多个通道
6. 在 smem 中执行 gemm
7. 将结果写入 o
```

**性能。** 我们在 nvidia a100 gpu 上评估了 flash-dwc 的性能，并与现有基线进行了比较。表 2 展示了最佳性能配置的结果。flash-dwc 在所有测试用例中均取得了最先进的性能，加速比最高达 3.5 倍。

**表 2. 最佳性能配置的比较。** flash-dwc 在 a100 gpu 上针对不同滤波器大小和输入分辨率的性能测试结果。加速比相对于 pytorch 基线计算。 作为起始和结束位置（第 4 行）。这略微减少了零填充开销，同时未引入线程束分歧。 **gemm 计算。** gemm1 和 gemm2 中的每一个都计算一批 $b_n$ 个通道级 m16n8k8 gemm。每个线程束在张量核心上执行 $b_n 4$ 条 mma.sync.aligned.m16n8k8 ptx 指令。为了避免银行冲突，我们将 smem 数组 w 和 x 的最后一个维度填充 128 位 [2]。在主循环之后，我们通过 smem 重排寄存器矩阵 v（第 12 行），并以 128 字节事务将生成的 $b_m \times b_n$ dwc 图块存储到全局内存中。 **为何不使用稀疏张量核心？** 虽然稀疏张量核心（例如 sparstencil [19] 中使用的）具有稠密张量核心两倍的理论吞吐量，但它们不适合 dwc。首先，它们支持的最小 gemm 形状 m16n8k16 将迫使我们放弃双路径拆分 gemm 流水线，该流水线隐藏了 50% 的 gemm 计算以位于数据加载之后。其次，即使我们重组数据流以满足 2:4 稀疏性，gemm 运行时减少最多为 43.75%，这还未考虑额外的开销和数据访问连续性的丧失。由于 dwc 本质上是内存受限而非计算受限，切换到稀疏张量核心将牺牲内存延迟隐藏以换取微乎其微的计算收益，最终降低整体性能。 ### 5 扩展 flash-dwc 至 cnn 训练 基于第 3 节和第 4 节中的优化内核，我们将 flash-dwc 扩展至完整的 cnn 训练，并将其作为即用型库集成到 pytorch 2.10 [36] 中，flashdwc.ops 用于操作调用，flashdwc.nn 用于模型构建。 #### 5.1 反向传播 前向卷积（fc）直接由 flash-dwc 内核支持。训练额外需要两个反向算子：(i) 反向数据卷积（bdc），它将输出梯度 $\nabla y$ 与转置滤波器 $w^t$ 卷积以计算输入梯度 $\nabla x$；以及 (ii) 反向滤波器卷积（bfc），它将输入特征图 x 与 $\nabla y$ 卷积以计算滤波器梯度 $\nabla w$。由于 bdc 在旋转滤波器后具有相同的 dwc 结构，它可以重用 flash-dwc 内核。主要挑战在于 bfc：其小输出张量导致并行度不足，且不匹配原始内核的分块方法。图 11 展示了 flash-dwc 如何支持所有三个阶段。 **图 11. 用于 cnn 训练的 flash-dwc。** fc 和 bdc 利用张量核心和 cuda 核心内核。对于 bfc，我们通过将 $\nabla y$ 分块为 z 个子 bfc 并原子累加它们的结果来增加并行度。如图所示，每个 11×11 输出的子 bfc 生成 6 个块：mega 内核通过组合 tm=2 和 9 的两个内核来匹配 $f_w=11$；$\beta=4$ 最小化 $f_h$ 轴填充。 **bdc 实现。** 我们将 w 旋转 180° 以获得 $w^t$，然后在 $w^t$ 和 $\nabla y$ 上执行 flash-dwc 内核。由于旋转开销可忽略不计，bdc 的性能与 fc 相似，同样受益于张量核心。 **累加轴并行化。** bfc 的输出 $\nabla w \in \mathbb{r}^{f_h \times f_w \times o_c}$ 仍然很小，且不随批次大小或空间分辨率扩展。这与 fc 和 bdc 形成对比，后者的输出巨大，自然能提供丰富的并行度。如果我们仅通过分块 $\nabla w$ 来并行化 bfc，生成的块数量太少，无法充分利用硬件。 我们通过沿累加轴并行化来解决这个问题。具体来说，我们将 $\nabla y$ 分割为 z 个图块，并将每个图块作为独立的子 bfc 处理。这使并行度增加了 z 倍，但引入了通过 fp32 原子加法合并子 bfc 结果的额外开销。为了平衡这两个效果，我们根据 dwc 参数和 sm 数量动态决定 z。如果累加长度不能被 z 整除，我们将余数附加到最后一个图块上，并调整 z 以最小化此余数，实现工作负载平衡。我们进一步通过一次性的性能分析步骤来完善 z，该步骤测量性能并选择最佳值。 **灵活分块。** bfc 的不规则输出形状与原始 flash-dwc 内核的分块方法不匹配。使用 $b_m=128$，这些内核将空间维度分块为 16×8、8×16 或 4×32。这些形状非常适合特征图，通常很大且是 8 的倍数（例如 112×112, 64×64）。然而，bfc 的输出 $\nabla w$ 具有奇数大小空间维度，且较小或与那些图块对齐不良，导致大量的零填充和不必要的计算。 #### 6.1 卷积 gelu 与层归一化 flash-dwc 支持 cnn 中常见的三种算子：分组深度卷积（gwd）、卷积-gelu 和层归一化。我们将融合卷积-gelu 和融合卷积-gelu 层归一化内核分别标记为 flash-dwc-f1 和 flash-dwc-f2。我们遵循第 3 节和第 4 节中的内核设计，并采用以下两项优化。 **融合卷积-gelu。** 为了在单个内核中融合 gwd 和 gelu，我们在 gemv 内核的每行末尾插入一个 gelu 阶段，在 gemm 内核的每个 gemm 图块后插入一个 gelu 阶段。因为 cuda 核心计算和 gelu 使用相同的线程，所以不需要额外的线程同步。张量核心计算和 gelu 需要额外的线程同步，但我们使用第 4.4 节中的双路径拆分 gemm 流水线来隐藏此开销。因为 gwd 具有大的输出通道，所以融合 gelu 不会引入显著的算术强度损失。 **融合卷积-gelu 层归一化。** 为了在单个内核中融合 gwd、gelu 和 ln，我们将 ln 分解为两个步骤：(1) 在归一化之前计算每通道的均值和方差，以及 (2) 使用计算出的统计量来归一化激活值。我们在每个块内执行这两个步骤，因此每个块需要 $o(b_m b_n)$ 个额外寄存器来存储统计量。由于 flash-dwc 内核已经使用了大量寄存器（例如，t4 内核使用了 216 个寄存器中的 208 个），我们需要重用现有的寄存器来存储 ln 统计量。我们观察到 gemm 内核中的输入矩阵寄存器在 gemm 计算后不再需要，因此我们重用这些寄存器来存储 ln 统计量。 #### 6.2 分组深度卷积 flash-dwc 将 gwd 视为 m 个独立的 dwc 的并集，并分别处理它们。这导致每个 dwc 需要单独的内存事务，降低了内存带宽效率。为了提高内存带宽效率，我们引入了交错分组深度卷积（igdw），它将 m 个独立的 dwc 重组为单个 dwc。具体来说，我们将输入特征图的 m 个通道沿空间维度交错排列，使得每个空间位置包含所有 m 个通道的元素。这样，igdw 就可以使用单个 dwc 内核来处理所有 m 个通道，从而显著提高了内存带宽利用率。 **算法 3. gemv 内核** ```text 输入：输入特征图 x，滤波器 w，输出 o 1. 将 x 分块为 bm x bm 图块 2. 对于每个图块： 3. 将滤波器 w 加载到 smem 中 4. 对于每个空间位置： 5. 从全局内存加载 x 的一个通道 6. 在 smem 中执行 gemv 7. 将结果写入 o ``` **算法 4. gemm 内核** ```text 输入：输入特征图 x，滤波器 w，输出 o 1. 将 x 分块为 bm x bn 图块 2. 对于每个图块： 3. 将滤波器 w 加载到 smem 中 4. 对于每个空间位置： 5. 从全局内存加载 x 的多个通道 6. 在 smem 中执行 gemm 7. 将结果写入 o ``` **性能。** 我们在 nvidia a100 gpu 上评估了 flash-dwc 的性能，并与现有基线进行了比较。表 2 展示了最佳性能配置的结果。flash-dwc 在所有测试用例中均取得了最先进的性能，加速比最高达 3.5 倍。 **表 2. 最佳性能配置的比较。** flash-dwc 在 a100 gpu 上针对不同滤波器大小和输入分辨率的性能测试结果。加速比相对于 pytorch>
**gemm 计算。** gemm1 和 gemm2 中的每一个都计算一批 $b_n$ 个通道级 m16n8k8 gemm。每个线程束在张量核心上执行 $b_n/4$ 条 mma.sync.aligned.m16n8k8 ptx 指令。为了避免银行冲突，我们将 smem 数组 w 和 x 的最后一个维度填充 128 位 [2]。在主循环之后，我们通过 smem 重排寄存器矩阵 v（第 12 行），并以 128 字节事务将生成的 $b_m \times b_n$ dwc 图块存储到全局内存中。

**为何不使用稀疏张量核心？** 虽然稀疏张量核心（例如 sparstencil [19] 中使用的）具有稠密张量核心两倍的理论吞吐量，但它们不适合 dwc。首先，它们支持的最小 gemm 形状 m16n8k16 将迫使我们放弃双路径拆分 gemm 流水线，该流水线隐藏了 50% 的 gemm 计算以位于数据加载之后。其次，即使我们重组数据流以满足 2:4 稀疏性，gemm 运行时减少最多为 43.75%，这还未考虑额外的开销和数据访问连续性的丧失。由于 dwc 本质上是内存受限而非计算受限，切换到稀疏张量核心将牺牲内存延迟隐藏以换取微乎其微的计算收益，最终降低整体性能。

### 5 扩展 flash-dwc 至 cnn 训练

基于第 3 节和第 4 节中的优化内核，我们将 flash-dwc 扩展至完整的 cnn 训练，并将其作为即用型库集成到 pytorch 2.10 [36] 中，flashdwc.ops 用于操作调用，flashdwc.nn 用于模型构建。

#### 5.1 反向传播

前向卷积（fc）直接由 flash-dwc 内核支持。训练额外需要两个反向算子：(i) 反向数据卷积（bdc），它将输出梯度 $\nabla y$ 与转置滤波器 $w^t$ 卷积以计算输入梯度 $\nabla x$；以及 (ii) 反向滤波器卷积（bfc），它将输入特征图 x 与 $\nabla y$ 卷积以计算滤波器梯度 $\nabla w$。由于 bdc 在旋转滤波器后具有相同的 dwc 结构，它可以重用 flash-dwc 内核。主要挑战在于 bfc：其小输出张量导致并行度不足，且不匹配原始内核的分块方法。图 11 展示了 flash-dwc 如何支持所有三个阶段。

**图 11. 用于 cnn 训练的 flash-dwc。** fc 和 bdc 利用张量核心和 cuda 核心内核。对于 bfc，我们通过将 $\nabla y$ 分块为 z 个子 bfc 并原子累加它们的结果来增加并行度。如图所示，每个 11×11 输出的子 bfc 生成 6 个块：mega 内核通过组合 tm=2 和 9 的两个内核来匹配 $f_w=11$；$\beta=4$ 最小化 $f_h$ 轴填充。

**bdc 实现。** 我们将 w 旋转 180° 以获得 $w^t$，然后在 $w^t$ 和 $\nabla y$ 上执行 flash-dwc 内核。由于旋转开销可忽略不计，bdc 的性能与 fc 相似，同样受益于张量核心。

**累加轴并行化。** bfc 的输出 $\nabla w \in \mathbb{r}^{f_h \times f_w \times o_c}$ 仍然很小，且不随批次大小或空间分辨率扩展。这与 fc 和 bdc 形成对比，后者的输出巨大，自然能提供丰富的并行度。如果我们仅通过分块 $\nabla w$ 来并行化 bfc，生成的块数量太少，无法充分利用硬件。
我们通过沿累加轴并行化来解决这个问题。具体来说，我们将 $\nabla y$ 分割为 z 个图块，并将每个图块作为独立的子 bfc 处理。这使并行度增加了 z 倍，但引入了通过 fp32 原子加法合并子 bfc 结果的额外开销。为了平衡这两个效果，我们根据 dwc 参数和 sm 数量动态决定 z。如果累加长度不能被 z 整除，我们将余数附加到最后一个图块上，并调整 z 以最小化此余数，实现工作负载平衡。我们进一步通过一次性的性能分析步骤来完善 z，该步骤测量性能并选择最佳值。

**灵活分块。** bfc 的不规则输出形状与原始 flash-dwc 内核的分块方法不匹配。使用 $b_m=128$，这些内核将空间维度分块为 16×8、8×16 或 4×32。这些形状非常适合特征图，通常很大且是 8 的倍数（例如 112×112, 64×64）。然而，bfc 的输出 $\nabla w$ 具有奇数大小空间维度，且较小或与那些图块对齐不良，导致大量的零填充和不必要的计算。

#### 6.1 卷积 gelu 与层归一化

flash-dwc 支持 cnn 中常见的三种算子：分组深度卷积（gwd）、卷积-gelu 和层归一化。我们将融合卷积-gelu 和融合卷积-gelu 层归一化内核分别标记为 flash-dwc-f1 和 flash-dwc-f2。我们遵循第 3 节和第 4 节中的内核设计，并采用以下两项优化。

**融合卷积-gelu。** 为了在单个内核中融合 gwd 和 gelu，我们在 gemv 内核的每行末尾插入一个 gelu 阶段，在 gemm 内核的每个 gemm 图块后插入一个 gelu 阶段。因为 cuda 核心计算和 gelu 使用相同的线程，所以不需要额外的线程同步。张量核心计算和 gelu 需要额外的线程同步，但我们使用第 4.4 节中的双路径拆分 gemm 流水线来隐藏此开销。因为 gwd 具有大的输出通道，所以融合 gelu 不会引入显著的算术强度损失。

**融合卷积-gelu 层归一化。** 为了在单个内核中融合 gwd、gelu 和 ln，我们将 ln 分解为两个步骤：(1) 在归一化之前计算每通道的均值和方差，以及 (2) 使用计算出的统计量来归一化激活值。我们在每个块内执行这两个步骤，因此每个块需要 $o(b_m b_n)$ 个额外寄存器来存储统计量。由于 flash-dwc 内核已经使用了大量寄存器（例如，t4 内核使用了 216 个寄存器中的 208 个），我们需要重用现有的寄存器来存储 ln 统计量。我们观察到 gemm 内核中的输入矩阵寄存器在 gemm 计算后不再需要，因此我们重用这些寄存器来存储 ln 统计量。

#### 6.2 分组深度卷积

flash-dwc 将 gwd 视为 m 个独立的 dwc 的并集，并分别处理它们。这导致每个 dwc 需要单独的内存事务，降低了内存带宽效率。为了提高内存带宽效率，我们引入了交错分组深度卷积（igdw），它将 m 个独立的 dwc 重组为单个 dwc。具体来说，我们将输入特征图的 m 个通道沿空间维度交错排列，使得每个空间位置包含所有 m 个通道的元素。这样，igdw 就可以使用单个 dwc 内核来处理所有 m 个通道，从而显著提高了内存带宽利用率。

**算法 3. gemv 内核**
```text
输入：输入特征图 x，滤波器 w，输出 o
1. 将 x 分块为 bm x bm 图块
2. 对于每个图块：
3. 将滤波器 w 加载到 smem 中
4. 对于每个空间位置：
5. 从全局内存加载 x 的一个通道
6. 在 smem 中执行 gemv
7. 将结果写入 o
```

**算法 4. gemm 内核**
```text
输入：输入特征图 x，滤波器 w，输出 o
1. 将 x 分块为 bm x bn 图块
2. 对于每个图块：
3. 将滤波器 w 加载到 smem 中
4. 对于每个空间位置：
5. 从全局内存加载 x 的多个通道
6. 在 smem 中执行 gemm
7. 将结果写入 o
```

**性能。** 我们在 nvidia a100 gpu 上评估了 flash-dwc 的性能，并与现有基线进行了比较。表 2 展示了最佳性能配置的结果。flash-dwc 在所有测试用例中均取得了最先进的性能，加速比最高达 3.5 倍。

**表 2. 最佳性能配置的比较。** flash-dwc 在 a100 gpu 上针对不同滤波器大小和输入分辨率的性能测试结果。加速比相对于 pytorch 基线计算。>

Flash-DWC-T 固定 BM=128 以支持 m16n8k16 GEMM，但其宽度为 9 的分片经常与 BFC 累加瓦片错位。相比之下，我们的 CUDA-Core 内核基于 GEMV，不受 Tensor-Core 形状约束。因此，我们采用灵活的瓦片划分策略对其进行适配，以更精确地匹配 BFC 形状，同时将零填充降至最低。我们变化 TM，即每个线程处理的空间位置数量，范围从 1 到 9，并将具有不同 TM 值的两个内核融合为一个超级内核，以无缝匹配任意 FW。由于累加长度 NOHOW 通常是 8 的倍数，我们以长度 8 为单位进行处理，并将每个线程块组织为 β×8 个线程：(i) 沿通道轴，8 个线程执行 128 字节的合并式加载；(ii) 沿 FH 轴，β 个线程复用重叠区域，协作为一个块级 GEMV 加载一个长度 8 的单位。我们选择 β∈{4,8}：两者都能整除 8 并保持线程总数为 32 的倍数，从而避免非活跃的 warp 通道。在两者之间，我们选择使 FH 轴上零填充最小化的方案；如果两者零填充相同，则取较大的值以最大化重叠复用。

混合累加。它作为一种可选优化，用于加速混合精度 BFC：每个子 BFC 以 FP16 累加，最终结果在 FP32 中合并。这避免了显式的数据类型转换，减少了寄存器使用，并且可以利用某些 GPU 上更高的 FP16 CUDA-Core 吞吐量。其精度仍然可接受，原因有两点。(i) BFC 的长累加长度可以拆分为许多（例如 ≥128 个）短子单元。以 FP16 累加每个子单元限制了误差增长，而在 FP32 中合并结果则受益于更大的尾数和动态范围，基本防止了精度损失和数值溢出。(ii) 与 FC 和 BDC 不同，后者的对称权重初始化最大化了正负抵消，BFC 的输入 X 和 ?Y 由于激活函数和损失函数的非线性而在零点附近不那么对称。这减少了 FP16 舍入误差。第 6.3 节显示混合累加与全 FP32 累加的收敛性相当。

5.2 内核融合

深度可分离卷积（DWC）通常紧接批量归一化（BN）和激活函数（Act），形成常见模式：BN+Act+DWC。对于计算单元未完全饱和的小滤波器，我们提供将 BN 和 Act 融合到 Flash-DWC 内核中的可选方案，以利用空闲周期并减少数据搬运。我们在 CUDA-Core 内核上实现此融合，它们将输入数据缓存在寄存器中，因此适合在 DWC 之前进行即时的 BN 和 Act 计算。尽管由于瓦片重叠需要重计算，但增加的负载仍在计算瓶颈阈值以下，因此对于小滤波器，融合内核的延迟与对应的非融合 DWC 内核几乎相同。

6 评估

我们首先评估 CUDA-Core 内核以验证我们的 GEMV 数据流协同设计，并分析从访存瓶颈向计算瓶颈的转变。然后评估基于 GEMM 的 Tensor-Core 内核，以展示 GEMV 到 GEMM 的转换如何进一步提升 CUDA-Core 性能上限。最后在 ImageNet-1K 上评估 Flash-DWC 的端到端 CNN 训练，衡量加速比和收敛性。

6.1 FP32 CUDA-Core 内核评估

6.1.1 实验设置

我们在 RTX 4090 和 A100 80GB 上进行评估，覆盖 Ada-Lovelace GeForce 和 Ampere 数据中心 GPU。我们将 Flash-DWC 与 DWC 基线和计算瓶颈参考进行比较：(i) cuDNN 9.18：我们启用 cuDNN 中的 Tensor-Core 和 CUDA-Core 两种算法。CuGEMM-DWC 表示其直接算法和基于 GEMM 的算法中最快者；CuFFT-DWC 表示其基于 FFT 和 Winograd 的算法中最快者。(ii) Inductor-DWC：由 PyTorch 2.10 的 torch.compile 使用 Inductor 后端和 max-autotune 生成。(iii) cuBLAS-GEMM：它不是 DWC 实现，而是作为计算瓶颈参考，用于衡量 Flash-DWC 接近 CUDA-Core 性能上限的程度。其维度设置为 m=n=k，2m3≈2NOHOWOCFHFW，使其总 FLOPs 与 DWC 匹配。

DWC 参数遵循常见的 CNN 配置。(i) 滤波器从较小的 3×3 [12, 13, 39] 到超大的 51×51 [2, 20]，覆盖 7×7 [25, 41, 44]、9×9 [8, 41] 和 31×31 [7] 等代表性尺寸。(ii) 特征图按 NHWC 格式进行系统化缩放：128×128×128×128、128×64×64×256 和 128×32×32×1024。这些尺寸足够大以饱和 GPU 并行性，同时在测试间保持可比的工作负载。

6.1.2 结果与扩展性分析

图 12 中所有吞吐量值均为热身运行后 500 次运行的平均值。图 13 展示了随着滤波器尺寸增加运算强度增大时，CUDA-Core Flash-DWC 内核的吞吐量及其相对于 cuBLAS-GEMM 的吞吐量如何变化。

图 12. FP32 吞吐量。滤波器尺寸（FH=FW，x 轴）范围从 3 到 51。吞吐量（y 轴）单位为 TFLOPS。子图标题列出 GPU 类型和 NHWC 格式的特征图维度。

图 13. FP32 CUDA-Core Flash-DWC 内核随运算强度的扩展性。滤波器尺寸（x 轴，FH=FW）范围从 3 到 51。左侧 y 轴显示吞吐量（TFLOPS）。右侧 y 轴显示两个归一化指标：运算强度比（相对于渐近极限的强度）；Flash-DWC 吞吐量相对于计算瓶颈 cuBLAS-GEMM 的比值。所有值均在特征图维度上取平均。

随着滤波器尺寸增大，Flash-DWC 表现出从访存瓶颈向计算瓶颈的明显转变。这一趋势遵循我们的运算强度模型 2TM/(1+(TM?1)/FW+BN/512)，并得到硬件测量的印证。

对于小滤波器（FW≤9），内核仍处于访存瓶颈状态。在此区间，(TM?1)/FW 项占主导地位，因此运算强度和吞吐量都随 FW 快速增加。因此，运行时间的增长远慢于 O(FHFW) 的 FLOP 计数：将滤波器从 3×3 增加到 7×7，运行时间仅延长 2.2 倍，而 FLOPs 增加了约 5.44 倍。这里 A100 优于 RTX 4090，因为其内存带宽更高（约 2 TB/s 对比约 1 TB/s），证实了访存量是主要性能瓶颈。

对于较大滤波器，运算强度曲线逐渐趋于平缓，而吞吐量和相对于 cuBLAS-GEMM 的吞吐量仍在上升。这意味着 Flash-DWC 不仅在绝对值上更快，而且逐渐逼近计算瓶颈行为。在两种 GPU 上，当 FW≥13 时 Flash-DWC 达到 cuBLAS-GEMM 吞吐量的 65%，在 A100 上当 FW≥15 时超过 90%，表明更高的运算强度推动 DWC 趋近 CUDA-Core 性能上限。屋顶线交叉点出现在 FW≥11，此时 Flash-DWC 在 FLOP 峰值更高的 RTX 4090 上的吞吐量（82.6 TFLOPS）超过了带宽更高的 A100（19.5 TFLOPS），表明瓶颈已从内存带宽转移到计算能力。

这种转变带来了相对于所有 DWC 基线的显著加速。Flash-DWC 比最快的 cuDNN DWC 算法快 1.08×–8.48×（平均 5.09×），比 CuGEMM-DWC 快 1.08×–8.72×（平均 5.98×），比 CuFFT-DWC 快 1.73×–83.05×（平均 13.90×），比 Inductor-DWC 快 2.65×–11.51×（平均 5.65×）。对于汇总加速比，平均指算术平均值。相应的几何平均加速比分别为 4.67×、5.62×、9.48× 和 5.42×。在 RTX 4090 上加速比更大，因为其 CUDA-Core 峰值更高：例如，相对于 CuGEMM-DWC，Flash-DWC 实现 6.34× 的平均加速，而 A100 上为 5.61×。Flash-DWC 也受益于较小的特征图，因为更高的缓存命中率提升了吞吐量。

6.2 FP16 Tensor-Core 内核评估

6.2.1 实验设置

我们使用与 6.1 节相同的基线、DWC 参数集和 GPU，只是所有 DWC 基线现在操作于 FP16 张量。对于 cuDNN，我们报告其在 FP16 和 FP32 累加模式下的最快结果。FP32 cuBLAS-GEMM 仍作为计算瓶颈参考，因为在 CUDA Core 上进行 FP16/FP32 DWC 需要显式的 FP16 到 FP32 转换，因此受限于 FP32 吞吐量而非标称的 FP16 峰值，如第 4.1 节所述。在 RTX 4090 上，这一选择更为合理，因为 FP16 和 FP32 CUDA-Core 峰值相同。我们使用 cuBLAS-GEMM 来评估 Flash-DWC-T 是否超越 CUDA-Core 性能上限。

6.2.2 结果与扩展性分析

图 14 中所有吞吐量值均为热身运行后 500 次运行的平均值。图 15 展示了随着滤波器尺寸增加运算强度增大时，Flash-DWC-T4 的吞吐量及其相对于 cuBLAS-GEMM 的吞吐量如何变化。

图 14. FP16 吞吐量。滤波器尺寸（x 轴，FH=FW）范围从 3 到 51。吞吐量（y 轴）单位为 TFLOPS。子图标题列出 GPU 类型和 NHWC 格式的特征图维度。所有 DWC 算法操作于 FP16 张量；cuBLAS-GEMM 为 FP32，作为 CUDA-Core 性能上限的参考。

图 15. FP16 Flash-DWC-T4 随运算强度的扩展性。滤波器尺寸（x 轴，FH=FW）范围从 3 到 51。左侧 y 轴显示吞吐量（TFLOPS）。右侧 y 轴显示两个归一化指标：运算强度比（定义为强度相对于其渐近极限的值）；Flash-DWC-T4 吞吐量相对于 CUDA-Core cuBLAS-GEMM 的比值。超过 100% 的值表明 Tensor-Core 加速超越了 CUDA-Core 计算上限。所有值均在特征图维度上取平均。

与其 CUDA-Core 版本相比，Flash-DWC-T 实现了 0.96×–4.35× 的加速（平均 2.01×）。唯一低于 1 倍的情况出现在 3×3 滤波器上。更一般地，CUDA-Core 和 Tensor-Core 内核对于小滤波器（FW≤7）都处于访存瓶颈状态，因此表现出相似的吞吐量。然而，随着滤波器尺寸增大运算强度增长，Flash-DWC-T 提升得更加迅速，因为 Tensor Cores 提供了更高的吞吐量并避免了显式的 FP16 到 FP32 转换。相比之下，在 A100 上，CUDA-Core 内核对于大滤波器的吞吐量甚至略有下降，因为一旦 CUDA Coes 饱和，数据类型转换会与 FMA 竞争计算周期。

如图 15 所示，Flash-DWC-T 的吞吐量遵循其运算强度模型 4096·0.5625/(128θ?1+112γ?1+16+FW/γ)，并紧密追踪其相对于 cuBLAS-GEMM 的吞吐量。吞吐量也以 9 为步长递增，反映了宽度优先分解为宽度 9 的分片。在三种 Tensor-Core 配置中，Flash-DWC-T4 实现最高吞吐量，其次是 T2 和 T1。这一顺序与其输入矩阵压缩比一致：更大的 θ 带来更高的运算强度和更低的 SMEM 使用。这些结果共同表明，运算强度仍然是 Tensor Cores 上性能的主要预测因子，正如 CUDA-Core 内核的情况一样。

对于中等大小的滤波器，Flash-DWC-T 甚至在两种 GPU 上都超过了 cuBLAS-GEMM 的吞吐量。一旦 Flash-DWC 在 CUDA Cores 上大幅缓解了 DWC 的访存瓶颈，主要限制就不再是访存量，而是 CUDA-Core 计算吞吐量。Tensor-Core 加速提升了这一上限，使 DWC 吞吐量能超越 CUDA-Core GEMM 参考继续扩展。从这个意义上说，超越 cuBLAS-GEMM 反映了从内存带宽到计算吞吐量的转变，也说明为什么 Tensor Cores 对于进一步加速变得必要。这一效应在 A100 上更为明显，因为 Tensor 和 CUDA Cores 之间的峰值吞吐量差距更大。相应地，Flash-DWC-T 在该 GPU 上相对于 cuBLAS-GEMM 和其 CUDA-Core 对应版本都获得了更大的优势。总体而言，Flash-DWC-T 在 A100 上达到 cuBLAS-GEMM 吞吐量的 0.36×–2.68×（平均 1.67×），而在 RTX 4090 上为 0.1×–1.71×（平均 0.99×）。

Flash-DWC-T 相对于所有 DWC 基线也实现了显著收益。重要的是，它在整个滤波器尺寸范围内都带来了正加速，而不是以牺牲小滤波器性能为代价来提升大滤波器性能。它比最快的 cuDNN DWC 算法快 1.84×–16.38×（平均 9.01×），比 CuGEMM-DWC 快 1.84×–16.38×（平均 10.58×），比 CuFFT-DWC 快 3.29×–102.63×（平均 23.35×），比 Inductor-DWC 快 2.88×–17.73×（平均 10.46×）。相应的几何平均加速比分别为 8.21×、9.83×、16.78× 和 9.95×。Flash-DWC-T 也受益于较小的特征图，因为更高的缓存命中率带来了提升。

6.2.3 跨 GPU 比较

尽管 A100 的 FP16/FP32 Tensor-Core 峰值（312 TFLOPS）高于 RTX 4090（165 TFLOPS），但 Flash-DWC-T 在 RTX 4090 上实现了更高的绝对吞吐量：平均/峰值为 50.6/90.3 TFLOPS，而 A100 上为 30.3/49.5 TFLOPS。在考虑我们 GEMV 到 GEMM 转换的 56.25% 密度后，密度归一化吞吐量在 RTX 4090 上达到平均/峰值 90.0/160.5 TFLOPS，而 A100 上为 53.9/88.0 TFLOPS，RTX 4090 的峰值接近其标称的 FP16/FP32 Tensor-Core 峰值。这一跨 GPU 差距不能仅用 Tensor-Core 峰值或 DRAM 带宽来完全解释。

为了更好地理解这一行为，我们使用 Nsight Compute 对 Flash-DWC-T 进行剖析。在剖析案例中，RTX 4090 始终实现比 A100 更高的 L2 命中率。对于 128×64×64×256 特征图，RTX 4090/A100 的 L2 命中率在 9×9、17×17、31×31 和 51×51 滤波器上分别为 90.5/82.3%、95.4/89.1%、97.9/93.5% 和 99.0/96.2%。这表明 RTX 4090 更大的 L2 缓存（72 MB 对比 A100 的 40 MB）有助于捕获 Flash-DWC-T 暴露的输入复用。在宽度轴迭代中，Flash-DWC-T 执行一次初始批量加载，随后是许多小增量加载。初始批量加载可受益于 A100 更高的带宽（约 2 TB/s 对比 RTX 4090 的约 1 TB/s），这与其在小滤波器（FW≤5）上的优势一致。随着滤波器尺寸增大，由于 Flash-DWC-T 的宽度优先访问模式和工作负载分布，增量加载变得更加驻留在缓存中，减少了片外带宽压力。在此区间，RTX 4090 受益于更大的 L2 缓存并实现更高吞吐量。

其他架构因素也可能对这一差距有所贡献。

诸如谓词求值、地址计算和循环控制等辅助任务在 Tensor Core 之外执行，而异步拷贝与 SMEM（共享内存）行为会影响数据搬运与 Tensor Core 计算之间的重叠程度。这种跨 GPU 的行为并非 Flash-DWC-T 独有：cuDNN 和 Inductor 的 DWC 实现在 RTX 4090 上的吞吐也高于 A100，比值平均为 1.52×，峰值为 2.06×。结合 L2 命中率计数器，这一一致的趋势支持了我们的解释：由于 DWC 在跨通道复用有限且采用细粒度的空间访问模式，其性能对缓存层级和流水线行为的敏感度，高于单纯对 Tensor Core 峰值算力或 DRAM 带宽的依赖。

## 6.3 CNN 训练评估

我们选用训练作为端到端的评估设置，因为它不仅能够检验推理中使用的正向 DWC 算子，还能覆盖反向深度卷积（BDC）、反向逐元素卷积（BFC）、优化器交互以及收敛性保持。该设置在评估 Flash-DWC 效率的同时，也验证其在完整训练流水线中的安全可用性。

我们选取了四个代表性模型家族 [7, 20, 25, 41]，覆盖中型稠密、大型稠密以及分解型大滤波器的 DWC 层。对所有模型，我们用 Flash-DWC 实现替换了 PyTorch 的原生 DWC 层，同时保持 Flash-DWC 与 PyTorch 所评估的模型架构完全一致。在最终阶段，当下采样使特征图尺寸小于名义滤波器尺寸时，我们将滤波器裁剪到（特征图尺寸 ? 1），因为这种超尺寸滤波器在 ImageNet 识别任务上通常带来极小的精度提升，却会显著增加 FLOPs。Flash-DWC 与 PyTorch 使用完全相同的训练配置，包括默认的 Adam 优化器 [14]、10?? 学习率，以及在 FP16 混合精度下的损失缩放 [33]。我们在两块 Ada-Lovelace 架构的 GPU 上运行实验：RTX 4090（24 GB）和 RTX 5880 Ada（工作站 GPU，48 GB）。如第 6.1 节所示，Ada-Lovelace 架构能够为所评估的 DWC 实现提供更高的吞吐，因而非常适合 DWC 工作负载。

### 6.3.1 速度

我们在 FP32 和 FP16 下测量端到端的每批次训练时间，在 50 批次预热后对 100 个批次取平均。每个被计时的迭代覆盖完整训练步骤，包括数据加载与预处理、主机到设备的传输、前向与反向传播以及优化器更新。输入图像被调整到 256×256 这一常用分辨率。为最大化 GPU 利用率，对每个模型我们选择能装入显存的最大批次大小 ∈{16, 24, 32, 48, 64, 128, 256}，并对 Flash-DWC 与 PyTorch 使用相同的批次大小。表 1(a) 报告了所得加速比。

**表 1. Flash-DWC 相对于 PyTorch 的端到端训练加速比。HA 表示用于 BFC 的混合累加，↑ 表示其在端到端训练中的收益。**

Flash-DWC 在 FP32 下取得 1.44×–6.27× 的加速，在 FP16 下取得 1.03×–11.23× 的加速，Tensor Core 加速使得 FP16 下的峰值加速比显著提升。加速比在同一架构的基础版本与扩容版本之间保持稳定，表明 Flash-DWC 能够随通道数良好扩展。由于标准卷积与全连接层的 FLOPs 随通道数呈二次增长，扩容版本的端到端加速比会略微下降。混合累加可将核函数吞吐提升约 30%，但其端到端收益受阿姆达尔定律的限制仅为 0.1%–4.6%，表明 Flash-DWC 已消除了 DWC 的瓶颈。

端到端加速比随滤波器尺寸增大而提升。对于使用 7×7 滤波器的 ConvNeXt，加速比为 1.03×–1.70×；将滤波器扩展至 9×9 后，加速比提升到 1.94×–2.75×。对 7×7 ConvNeXt 而言，FP16 加速比低于 FP32，因为 PyTorch 在 FP16 下会调度专门的核函数，留给改进的空间较小。对于 ConvMixer，它使用许多大通道（1024、1536）配 9×9 滤波器的 DWC，Flash-DWC 取得了 3.64×–6.11× 的显著加速；将滤波器扩大到 11×11 后，加速比提升到 4.55×–7.19×。对于具有 31×31 和 51×51 超大滤波器的模型，Flash-DWC 对 RepLKNet 取得了 4.50×–8.64× 的加速，对 SLaK 取得了 3.38×–11.23× 的加速。在 SLaK 中，步幅分解3降低了 Flash-DWC 的算术强度与其运行时占比，因此加速比更低：条带版本为 3.38×–4.20×，而稠密版本为 5.65×–11.23×。核函数融合可进一步加速小/中等滤波器场景，但为了公平比较此处被禁用。

Flash-DWC 显著削弱了更大稠密滤波器带来的运行时惩罚。对 ConvNeXt 与 ConvMixer，将滤波器从 7×7/9×9 提升到 9×9/11×11，端到端训练时间仅增加约 1.3% 和 15%，尽管 DWC 的 FLOPs 分别增加了约 65% 和 49%。类似地，稠密版 SLaK 相比其条带版仅使端到端训练时间增加约 8%，尽管其 DWC 的 FLOPs 增加了数倍。

我们进一步将输入分辨率扩大到 512×512，这与分割、目标检测等稠密预测任务相关。如表 1(b) 所示，Flash-DWC 取得了 1.04×–14.18× 的相近加速比，展示了跨输入规模的稳健性。

### 6.3.2 收敛性

我们首先在单层上验证 Flash-DWC 不会引入可测量的精度下降，使用从 U(?1, 1) 中抽取的随机张量。在 FP32 下，Flash-DWC 的输出与梯度几乎与 PyTorch 完全一致，常常 MARE 为零。在 FP16 下，Flash-DWC 相对于 FP64 真值达到约 10?3 至 10?? 的 MARE，与 PyTorch 自身的 FP16 实现相当。

我们在 FP16 混合精度下训练代表性模型 [7, 25, 41]，这是一个比 FP32 更具挑战性的设置。输入图像调整到 256×256，批次大小为 64。如图 16 所示，Flash-DWC 的收敛行为与 PyTorch 相似，而混合累加也达到了与全 FP32 累加相当的收敛。这些结果证实 Flash-DWC 在提供显著加速的同时保留了训练精度。

**图 16. CNN 训练损失。子图标题列出模型及其最大滤波器尺寸。HA 表示用于 BFC 的混合累加。**

## 7 可扩展性

Flash-DWC 可扩展到更先进的 GPU 架构特性和低位数据类型。两个方向都需要核函数级别的重新调优，但不改变 Flash-DWC 的主要算法。

**更先进的 GPU 架构特性**可进一步提升 Flash-DWC 的性能。(i) Warp 专门化（warp specialization）可为消费者 warp 分配更大的寄存器预算、为生产者 warp 分配更小的预算，使消费者能够以更高的数据复用度执行更大的子 GEMM。(ii) 对于生产者，Flash-DWC 具有规则的预取模式，可视为一系列跨步的连续 128 字节片段。因此，许多细粒度的 `cp.async` 指令可用更少的多维 `cp.async.bulk.tensor` 传输替代。这可把全局内存加载和部分地址生成工作卸载到 Tensor Memory Accelerator（TMA），从而实现更高带宽和更好的重叠。(iii) 在集群（cluster）内，分布式共享内存（DSMEM）可能使输入补丁重叠和共享滤波器片段实现跨 CTA（block）复用。(iv) 对于消费者，可重新组织多个 warp 级 MMA 指令，以利用更大的 warpgroup 级 MMA（WGMMA）指令，从而可能提升 Tensor Core 利用率。

扩大 Tensor Core GEMM 粒度可能需要重新组织工作负载分布和 CTA 形状，也可能增加累加器寄存器压力并降低占用率。这些代价可通过 warp 专门化、基于 TMA 的异步传输以及 WGMMA 执行部分抵消。DSMEM 的收益则取决于能否将集群级 DSMEM 访问隐藏在 GEMM 计算之后。

**低位数据类型**，如 FP8、FP6、FP4 和 microscaling 格式，通常以混合精度方式使用：操作数以低精度存储和相乘，而累加仍保持较高精度，如 FP32 或 FP16，取决于精度要求。Flash-DWC-T 的数据复用模式、GEMV 到 GEMM 的转换、矩阵压缩以及双路 split-GEMM 流水线在很大程度上是数据类型无关的。因此，只要 Tensor Core 支持目标混合精度格式，Flash-DWC 就可以扩展到低位数据类型。

低位数据类型可减少内存流量，并在 Tensor Core 上实现更高峰值吞吐，两者都有助于提升 Flash-DWC-T 的性能。支持它们需要核函数级别的重新调优，包括调整适配 128 字节内存访问的数据布局、选择合适的 Tensor Core 指令，以及处理量化/反量化。这些修改可能影响软件流水线、内存访问行为以及辅助 CUDA Core 开销，因此其对性能与精度的影响应在具体部署中加以验证。

## 8 相关工作

**CNN 中的深度卷积。** 深度卷积（DWC）已成为现代 CNN 的关键组件。早期的轻量级网络 [12, 13, 39, 40] 使用小滤波器 DWC，在保持精度的同时减少 FLOPs。最近，大滤波器 DWC 被采用以扩大感受野并建模长程依赖，达到了与视觉 Transformer 竞争的精度 [9, 22, 23, 43]。例如，ConvNeXt [25, 44] 使用 7×7 DWC，RepLKNet [7] 扩展到 31×31，SLaK [20] 进一步推至 51×51 并使用条带卷积。PeLK [2] 表明稠密 DWC 尽管 FLOPs 更高，但性能优于条带变体，从而推动了参数高效的稠密设计。其他工作 [8, 15, 17, 26, 29, 41, 46, 47, 48] 也利用大滤波器 DWC 来捕捉全局上下文。尽管有这些架构进展，DWC 在 GPU 上仍然是内存受限的，限制其实用的训练效率，尤其在大滤波器情形下。

**DWC 优化。** 此前的工作已在 GPU 上优化 DWC。大部分 [4, 27, 28, 32, 34, 36, 37] 运行在 CUDA Core 上，因此限制了其性能，并在混合精度下引入数据类型转换开销。常规布局，尤其是 NCHW [4, 28, 32, 36, 37]，妨碍了 DWC 中的合并内存访问。[27, 28] 提出了高效方法复用输入补丁重叠；但它们仍是面向 CUDA Core 的，且主要面向小滤波器（如 ≤5×5）。尽管 cuDNN [34] 可以在 Tensor Core 上执行 DWC，其 Tensor Core 性能常常低于其 CUDA Core 性能。对角重排 [38] 将 DWC 转换为标准卷积，可用高性能 BLAS 库加速；然而其高度稀疏的变换滤波器引入了大量不必要的计算，使其比近期 cuDNN 版本更慢。[37] 融合深度卷积与逐元素卷积以减少数据搬运；它对小批次推理有效，但由于重算与控制复杂度，在更大批次（如 16）下效率低下。MegEngine [32] 为大滤波器 DWC 提供两种方法：直接卷积方法使用 SMEM 进行数据复用，但在 CUDA Core 上运行且缺乏软件流水线；批量 GEMM 方法使 FLOPs 增加 IH×IW/(FH×FW) 倍，在大特征图或小批次下可能低效。与 DWC 自身优化正交的是，硬件感知的剪枝技术 [35] 从算法上缩减模型规模与 FLOPs。

基于 FPGA 的工作 [1, 18, 24, 30, 42] 设计专用数据流与核心以提升硬件利用率和片上数据复用。面向 CPU 的工作 [11, 49] 引入 ARM 专用寄存器/缓存优化和核函数融合。[45] 通过基于扫描的切块和层融合优化 CNN 推理加速器上的 DWC，减少 DRAM 访问。这些非 GPU 的工作表明 DWC 加速主要取决于数据局部性与硬件利用率，而不仅取决于 FLOPs。尽管它们的平台特定机制不能直接应用于 GPU，但我们共享这一以局部性驱动的视角，并研究其在 GPU 上的实现。

**用于不规则计算的 Tensor Core。** Tensor Core 是为稠密 GEMM 设计的，限制了其在稀疏或不规则计算中的适用性。近期工作如 ConvStencil [3] 和 SparStencil [19] 将模板计算转换为结构化密度超过 50% 的 GEMM。然而，这些方法针对小单通道滤波器，无法扩展到 DWC 的多通道大滤波器设置。

**我们的工作。** Flash-DWC 与先前工作有三点不同：(i) 它为原生 DWC 引入了基于 GEMV 的数据流协同设计，联合优化数据复用、访问顺序、布局与工作负载切块，在保持访问连续性和内存合并的同时提升算术强度。(ii) 它引入了 DWC 专属的 GEMV 到 GEMM 转换，将 DWC 映射为紧凑的 56.25% 稠密 Tensor Core GEMM，从而实现高效混合精度执行而无需大量冗余。(iii) 它将上述技术与资源高效的核函数优化结合，以支持包括反向传播和 PyTorch 集成在内的完整训练。

## 9 结论

我们提出 Flash-DWC，一个使 DWC 计算高效的 GPU 库。通过结合基于 GEMV 的数据流协同设计、紧凑的 GEMV 到 GEMM 转换以及资源高效的核函数优化，Flash-DWC 将 DWC 从内存受限转向计算受限的操作，并将这一加速扩展到端到端的 CNN 训练。实验表明，Flash-DWC 在核函数级评估和端到端训练中都取得了显著的加速比。总体而言，Flash-DWC 使大滤波器 DWC 在高效且准确的视觉应用中更具实用性。

**致谢** 作者感谢匿名审稿人提出的建设性意见，这些意见显著改进了本文手稿。我们衷心感谢江青才（Qingcai Jiang）富有洞察力的讨论，这些讨论帮助塑造了本研究方向。