This paper is the first to assess the state of existing sparse matrix
multiplication algorithms on GPU for the butterfly structure, a promising form
of sparsity. This is achieved through a comprehensive benchmark that can be
easily modified to add a new implementation. The goal is to provide a simple
tool for users to select the optimal implementation based on their settings.
Using this benchmark, we find that existing implementations spend up to 50% of
their total runtime on memory rewriting operations. We show that these memory
operations can be optimized by introducing a new CUDA kernel that minimizes the
transfers between the different levels of GPU memory, achieving a median
speed-up factor of x1.4 while also reducing energy consumption (median of
x0.85). We also demonstrate the broader significance of our results by showing
how the new kernel can speed up the inference of neural networks.

通过一项全面的基准测试，本研究评估了目前 GPU 上用于蝶形稀疏矩阵乘法算法的现状，旨在为用户提供一个简单的工具来选择最佳算法实现。结果显示现有实现在内存重写操作上耗费了总运行时间的高达 50％。同时，引入一种新的 CUDA 核心可以优化这些内存操作，最多可以使计算速度提高 1.4 倍，降低能耗 0.85 倍。我们还展示了新核心的广泛意义，以其加速神经网络的推断。

提升推理速度：蝶式稀疏矩阵乘法的高效 GPU 内存管理

Make Inference Faster: Efficient GPU Memory Management for Butterfly  Sparse Matrix Multiplication

We provide an optimized implementation of the forward pass of
FlashAttention-2, a popular memory-aware scaled dot-product attention
algorithm, as a custom fused CUDA kernel targeting NVIDIA Hopper architecture
and written using the open-source CUTLASS library. In doing so, we explain the
challenges and techniques involved in fusing online-softmax with back-to-back
GEMM kernels, utilizing the Hopper-specific Tensor Memory Accelerator (TMA) and
Warpgroup Matrix-Multiply-Accumulate (WGMMA) instructions, defining and
transforming CUTLASS Layouts and Tensors, overlapping copy and GEMM operations,
and choosing optimal tile sizes for the Q, K and V attention matrices while
balancing the register pressure and shared memory utilization. In head-to-head
benchmarks on a single H100 PCIe GPU for some common choices of
hyperparameters, we observe 20-50% higher FLOPs/s over a version of
FlashAttention-2 optimized for last-generation NVIDIA Ampere architecture.

我们提供了 FlashAttention-2 前向传递的优化实现，使用了自定义融合的 CUDA 内核，以适应 NVIDIA Hopper 架构，并使用开源的 CUTLASS 库编写。在此过程中，我们解释了将在线 softmax 与连续的 GEMM 内核融合在一起的挑战和技术，利用 Hopper 特定的 Tensor Memory Accelerator（TMA）和 Warpgroup Matrix-Multiply-Accumulate（WGMMA）指令，定义和转换 CUTLASS 布局和张量，重叠复制和 GEMM 操作，并选择 Q、K 和 V 注意力矩阵的最优瓦片大小，同时平衡寄存器压力和共享内存利用率。在单个 H100 PCIe GPU 上进行的对比性测试中，针对某些常见的超参数选择，我们观察到与针对上一代 NVIDIA Ampere 架构进行优化的 FlashAttention-2 版本相比，FLOPs/s 高出 20-50%。