Coruscant: Co-Designing GPU Kernel and Sparse Tensor Core to Advocate Unstructured Sparsity in Efficient LLM Inference

文章标题: Coruscant: Co-Designing GPU Kernel and Sparse Tensor Core to Advocate Unstructured Sparsity in Efficient LLM Inference
作者/机构: Donghyeon Joo (University of Maryland, College Park), Helya Hosseini (University of Maryland, College Park), Ramyad Hadidi (d-Matrix), Bahar Asgari (University of Maryland, College Park)

A1 主要贡献

本文旨在解决在现代硬件上利用非结构化稀疏性进行高效大语言模型（LLM）推理的难题。尽管非结构化剪枝在较高稀疏度下能更好地保持模型精度，但由于硬件支持不足，LLM剪枝通常被限制在硬件友好的结构化模式上。这导致剪枝技术的潜力无法被完全发挥。

核心问题: LLM推理，尤其是解码（decode）阶段，主要由权重投影（GeMM）操作主导，这些操作是内存密集型的（如图1所示）。现有的稀疏格式在30%至70%的稀疏度范围内（这是LLM剪枝的典型范围）压缩效率低下，而现有的SpMM（稀疏矩阵-密集矩阵乘法）GPU核（Kernel）也无法在该稀疏度范围内提供有效的计算加速。这迫使研究人员采用NVIDIA支持的2:4半结构化稀疏格式，但这种格式将稀疏度固定在50%且会牺牲模型精度。

研究目标: 本文的目标是充分利用剪枝后LLM中的非结构化稀疏性，以减少数据移动、加速计算并降低权重矩阵在全局内存中的占用，从而为KV缓存释放更多空间，支持更长的上下文和更大的批处理大小。为了实现这一目标，本文致力于解决开发一种既能提高压缩率又能实现高效计算的稀疏格式所面临的挑战。

创新与贡献:
为了应对上述挑战，本文提出了Coruscant，一个协同设计GPU核与稀疏张量核的方案，以支持LLM推理中的非结构化稀疏性。主要贡献如下：
* Coruscant GPU SpMM核: 开发了一个专门的SpMM核，该核直接操作于一种新提出的基于位图的稀疏格式——Coruscant稀疏格式。该格式能在30%至70%的广泛稀疏度范围内，以比现有格式更高的压缩率压缩LLM权重。Coruscant GPU核通过减少内存带宽压力，性能优于先前的核。
* Coruscant稀疏张量核: 设计了一种稀疏张量核微架构，与SpMM核协同工作。它能够直接在压缩的稀疏格式上进行计算，无需解压，从而通过微小的架构改动实现进一步的加速。
* 显著加速效果: 在RTX 6000 Ada GPU上，与cuBLAS相比，Coruscant GPU核在50%/60%/70%的稀疏度下分别实现了高达1.50倍/1.72倍/2.02倍的加速；与Flash-LLM相比，在30%/40%/50%的稀疏度下分别实现了高达1.47倍/1.49倍/1.48倍的加速。在Llama 2 7B的端到端生成任务中，Coruscant核在商用GPU上使token吞吐量增加了26%（达到135 tokens/s），而使用Coruscant稀疏张量核则增加了40%（达到206 tokens/s）。


(a) Prefill和Decode阶段的延迟比较 (b) Prefill和Decode阶段的Roofline比较
图1：Llama 2-7B和13B的执行剖析，输入/输出长度分别为256和128。(a) Decode阶段主导了总运行时间，其中权重投影(GeMM)是主要瓶颈。(b) 矩阵乘法的Roofline图：灰色代表prefill，红色代表decode。

A3 背景知识与动机

2.1 LLM中的非结构化稀疏性

非结构化剪枝的优势与2:4模式的局限性。一些先前的研究从根本上在LLM权重中产生了非结构化稀疏模式。然而，由于现代GPU架构的限制，强制采用2:4结构化剪枝模式已成为一种普遍的解决方案，以便利用当前稀疏张量核架构带来的计算优势，尽管这会降低模型精度并将稀疏度固定在50%。著名的剪枝方法包括SparseGPT【16，Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. In Proceedings of the 40th International Conference on Machine Learning, 2023】，它通过迭代多个剪枝步骤将权重值转换为稀疏矩阵。在决定剪枝哪些值时，Wanda【49，Mingjie Sun, Zhuang Liu, Anna Bair, and J. Zico Kolter. A Simple and Effective Pruning Approach for Large Language Models. In Proceedings of the 12th International Conference on Learning Representations, 2024】使用激活-权重乘积，而SAMSP【46，Hang Shao, Bei Liu, and Yanmin Qian. One-Shot Sensitivity-Aware Mixed Sparsity Pruning for Large Language Models. In ICASSP 2024 - 2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2024】使用基于Hessian迹的敏感度作为元素级重要性分数。此外，DSnoT【67，Yuxin Zhang, Lirui Zhao, Mingbao Lin, Yunyun Sun, Yiwu Yao, Xingjia Han, Jared Tanner, Shiwei Liu, and Rongrong Ji. Dynamic Sparse No Training: Training-Free Fine-Tuning for Sparse LLMs. In Proceedings of the 12th International Conference on Learning Representations, 2024】与其他剪枝方法协同工作，以减少模型精度下降。表1比较了在有无2:4稀疏模式限制的情况下应用这些剪枝技术的模型困惑度。如表所示，施加2:4稀疏模式限制比非结构化稀疏性更能降低模型精度。注意，SAMSP【46】和DSnoT【67】未报告2:4稀疏模式下的性能指标。尽管非结构化剪枝有保持精度的潜力，但这些方法必须采用2:4稀疏格式才能实现高效部署。

Table 1: 50%稀疏度下非结构化与2:4稀疏模式的性能比较。

† WikiText-2上的困惑度【34】。结果来自【72】。‡ 越低越好。

确定理想的目标稀疏度范围。在不考虑硬件限制的情况下，我们首先进行了一项研究，以确定理想的稀疏度目标范围。图2展示了在WikiText2【34，Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In arXiv preprint arXiv:1609.07843, 2016】数据集上，不同剪枝方法在非结构化稀疏性下不同稀疏度的困惑度。在多种剪枝技术和不同规模的模型中，我们观察到困惑度（即精度下降）在30%稀疏度时开始显著增加，并在70%稀疏度之后呈指数级加速。对于低于30%的稀疏度，由于稀疏性不足，剪枝的效果减弱，导致的困惑度影响与30%时相似。此外，尽管困惑度较高，70%的稀疏度仍然具有宝贵的潜力，例如可以像He等人【19，Shwai He, Guoheng Sun, Zheyu Shen, and Ang Li. What Matters in Transformers? Not All Attention Is Needed. In arXiv preprint arXiv:2406.15786, 2024】探索的那样，在不同层应用不同的稀疏度。基于这些观察，我们的目标是加速30%到70%稀疏度范围内的SpMM，这超出了传统稀疏核和加速器【5, 8, 9, 20, 43, 62】通常处理的目标稀疏度范围。

非结构化稀疏性的压缩效率挑战。尽管非结构化稀疏性允许我们扩大稀疏度范围，但它面临着另一个压缩效率的挑战。压缩剪枝后LLM权重的非结构化稀疏模式的主要挑战是，现有的稀疏格式无法在30%到70%的稀疏度范围内高效地压缩稀疏矩阵。这一限制源于表示非零位置的方法。诸如压缩稀疏行（CSR）、压缩稀疏列（CSC）和Flash-LLM的Tiled-CSL【58，Haojun Xia, et al. Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity. Proc. VLDB Endow. 17, 2, 2023】等流行格式，都使用一个指针来引用一个非零组的起始索引（CSR是行，CSC是列，Tiled-CSL是tile）。然后，每个组内的非零位置表示为索引。虽然索引的确切位宽影响压缩率，但这通常在我们的目标稀疏度范围内导致巨大的开销。对于坐标列表（COO）格式，情况更糟，它用索引表示非零元素的两个维度。像ELL + COO这样的混合方案在理想条件下（即每行包含相同数量的非零元素，允许每个条目只有一个索引）实现了略优的压缩。然而，为每个非零元素存储一个索引的开销导致稀疏表示在50%稀疏度时超过了密集矩阵的内存占用。


图2：在非结构化稀疏性下，不同稀疏度和方法的LLM困惑度（越低越好）比较。本文目标范围为30%至70%。缺失值表示先前工作中没有可用结果。


图3：稀疏格式之间的压缩率比较。与基于位图的表示（Coruscant）相比，使用索引值（CSR, COO, Tiled-CSL）表示非零位置在我们目标稀疏度区域会产生高压缩率（相对于原始大小的百分比）。红色区域表示稀疏格式的大小超过了密集矩阵。

现有稀疏格式的压缩率分析。图3比较了应用于一个大小为9216 × 9216的float16矩阵的不同稀疏格式的压缩率。我们假设每个索引和偏移指针为16位。图中显示，在50%的稀疏度下，CSR和Tiled-CSL的压缩率超过100%，使其比原始矩阵更大。这是因为，为了替代被移除的零，每个非零元素的索引被存储起来。偏移指针也引入了少量额外开销。COO格式的表现更差，在70%的稀疏度下压缩率仍为90%。基于这些观察，本文提出使用一种基于位图的稀疏格式，其细节将在第3.1节中解释。

2.2 用于LLM的GPU核

现有GPU核方案。GPU上的SpMM可以利用各种核实现，从密集计算到稀疏优化方法。CuBLAS【41，NVIDIA. cuBLAS API Reference Guide, 2024】是Nvidia的密集矩阵乘法核，作为我们密集矩阵操作的基线。CuSPARSELt【42，NVIDIA. cuSPARSELt: A High-Performance CUDA Library for Sparse Matrix-Matrix Multiplication, 2024】同样由Nvidia提供，支持2:4稀疏矩阵乘法，并作为展示具有减少数据传输和为特定稀疏格式调整的硬件支持的GPU核性能的代理。Flash-LLM【58，Haojun Xia, et al. Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity. Proc. VLDB Endow. 17, 2, 2023】是为LLM推理设计的SpMM核。它引入了Tiled-CSL格式，该格式标注了每个tile内的非零值及其位置索引，并采用了一种“加载为稀疏，计算为密集”的方案。

Flash-LLM的流水线与局限性。图4展示了Flash-LLM核的流水线，它重叠了分块矩阵乘法的内存和计算操作。流水线从一个前导的iteration-1开始，其中矩阵A的一个压缩tile被写入GPU处理器的寄存器（gmem2reg）。矩阵A tile的共享内存被初始化为零（rst_smem），然后矩阵A tile被解压缩到共享内存中（extract）。同时，一个密集矩阵B的tile被直接加载到GPU处理器的共享内存中（ld_dense）。在下一个迭代中，当第二组矩阵tile经历上述过程时，前一迭代的第一组矩阵tile在张量核内进行计算（smem2tc）。这个过程持续进行，直到所有tile组被处理完毕。然而，Flash-LLM存在局限性。尽管Flash-LLM可以从GPU全局内存和GPU处理器之间减少的数据传输中受益，因为矩阵A是以压缩形式传输的，但计算仍然以密集格式进行。此外，Flash-LLM仅在稀疏度为70%及以上时才有效，此时Tiled-CSL能够将原始矩阵压缩到更小的大小。


图4：Flash-LLM [58]的加载和计算流水线。每次迭代都是内存阶段（写）和计算阶段（读和计算）的流水线执行，内存阶段将数据写入GPU处理器内存，计算阶段张量核对加载的数据进行操作。压缩数据到共享内存的解压（extract）也在内存阶段进行。

2.3 稀疏张量核

张量核与2:4稀疏支持。张量核是Nvidia Volta【39，NVIDIA. NVIDIA TESLA V100 GPU ARCHITECTURE, 2017】代GPU引入的专用矩阵乘法单元。如图5所示，张量核位于GPU处理器，即流式多处理器（SM）中，以加速小块矩阵的乘法。张量核的编程接口以矩阵乘加（MMA）操作的形式暴露，这些操作对$M \times K$和$K \times N$矩阵进行运算，并累加到$M \times N$的输出矩阵上，具体维度因GPU代次和指令而异。Ampere【40，NVIDIA. NVIDIA A100 Tensor Core GPU Architecture, 2020】代引入了稀疏张量核，增加了直接计算2:4稀疏格式的硬件能力。如图6所示，利用一个2位索引来表示稀疏矩阵操作数A中四个连续元素内的非零位置，一个选择器单元从密集矩阵操作数B的四个操作数中选择两个B元素，与两个非零A操作数相乘。这有效地将先前MMA操作的K维度加倍，以实现理论上的2倍加速（未来称为计算跳过）。在LLM推理的内存密集型SpMM操作背景下，更重要的是减少数据传输以及直接在压缩的矩阵tile A上操作而无需解压的能力。


图5：每个GPU处理器内的张量核，执行$M \times K$和$K \times N$的矩阵乘法。矩阵C和D用于累加多个MMA指令的输出来计算更大维度的矩阵乘法。


图6：稀疏张量核的2:4选择器。2位索引（粉色）表示非零位置，用于提取密集向量元素（紫色）。压缩的非零元素（绿色）与提取的密集元素相乘，并与先前的输出累加。

先前稀疏张量核研究的局限性。为了实现硬件支持的解压缩，先前的工作提出了对稀疏张量核的微架构增强，以支持非结构化稀疏性。DSTC【55，Yang Wang, et al. Dual-Side Sparse Tensor Core. In 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA), 2021】和RM-STC【21，Guyue Huang, et al. RM-STC: Row-Merge Dataflow Inspired GPU Sparse Tensor Core for Energy-Efficient Sparse Acceleration. In 56th Annual IEEE/ACM International Symposium on Microarchitecture (MICRO), 2023】扩展了稀疏张量核的能力，以支持两个矩阵操作数中的任意稀疏性。然而，这两种设计都针对计算密集型的大型方阵矩阵乘法进行了优化（RM-STC为$M = N = K = 2048$，DSTC为$M = N = K = 4096$），在这种情况下，跳过计算会带来显著的性能提升。为了利用稀疏性，这两种架构都将tile压缩成基于位图的稀疏格式，并选择性地对非零元素执行MMA操作。然而，这需要在更广泛的输出矩阵区域上分散部分结果，这需要额外的累加逻辑和中间缓冲区，从而使硬件复杂化并增加了面积开销。

LLM推理负载的独特性。相比之下，LLM推理工作负载呈现出根本不同的计算模式。在这里，代表批处理大小的N维度远小于M和K维度，导致了形状高度不平衡的“瘦”SpMM。这些工作负载是内存密集型的，性能主要由内存流量而非计算吞吐量决定。在这种情况下，最大化稀疏矩阵操作数的压缩对于减少内存带宽至关重要，而在计算跳过架构中常见的跨宽区域输出累加的需求可能会引入不必要的开销。我们将在第5.3.3节中进一步评估这三种设计的性能影响和开销。

A2 方法细节

Coruscant由三个关键部分组成：Coruscant稀疏格式，它能最大限度地压缩30%至70%的非结构化稀疏性；Coruscant GPU核，它利用这种稀疏格式在商用GPU上加速SpMM；以及Coruscant稀疏张量核，它直接在这种稀疏格式上操作，通过消除解压开销来提升GPU核的性能。


图7：Coruscant格式——tile偏移量中的每个元素都是一个指向tile起始位置的指针，该位置由相应tile的非零（NNZ）元素定义。每个大小为64的列tile由其非零元素和一个64位位图表示。

3.1 Coruscant稀疏格式

格式设计与优势。基于图3所示的压缩率比较，先前支持非结构化稀疏性的稀疏格式——CSR、COO、Flash-LLM的Tiled CSL和HYB——必须存储一个或多个与非零元素数量成正比大小的显式坐标数组。在中等稀疏度（30-70%）下，每个非零元素的索引字段占据了总内存占用的主导地位，甚至超过了密集矩阵。在Chou等人【11，Stephen Chou, Fredrik Kjolstad, and Saman Amarasinghe. Format Abstraction for Sparse Tensor Algebra Compilers. In Proceedings of the ACM on Programming Languages, Vol. 2, 2018】的基于层次的稀疏张量抽象中，我们的格式（图7）可以理解为两个维度级别的组合：用于tile划分的压缩级别（level-0）和一个由64位位图支持的自定义单例级别（level-1），用于编码每个tile内的坐标存在性。与存储显式坐标数组的传统单例级别不同，我们的level-1支持使用warp同步位掩码操作进行内存高效的坐标值迭代，同时保持有序、唯一和紧凑的属性。这种方法显著提高了压缩率，在70%稀疏度下达到36.25%，在50%稀疏度下达到56.25%，而先前的稀疏格式在这些稀疏度下无法有效压缩非结构化稀疏矩阵。在level-0中，我们将矩阵逻辑上划分为列式tile。通过按列布局tile，访问连续tile的线程将以轮询方式接触不同的共享内存bank，从而消除了解压过程中的bank冲突。我们使用64的tile大小，这是基于每个线程可用的GPU寄存器数量以及Coruscant GPU核的公式化设计，具体细节在第3.2节中解释。

3.2 Coruscant GPU核

利用稀疏格式加速SpMM。为了最大限度地减少GPU全局内存和处理器内的数据传输，Coruscant GPU核利用了我们的稀疏格式。通过减少数据传输和高效解压位图，该核加速了LLM中内存密集型的SpMM操作。

Coruscant SpMM公式化。一个压缩的稀疏矩阵tile被加载到GPU处理器的寄存器中，然后解压成共享内存中的一个密集矩阵tile。每个张量核使用这个密集tile来执行密集的MMA操作。我们采用了Flash-LLM的流水线（图4），其中寄存器和共享内存充当乒乓缓冲，有效地将“加载为压缩”和“计算为密集”的迭代过程流水线化。在此机制的基础上，Coruscant集成了基于位图的稀疏格式，并引入了一种高效的解压算法。

核的公式化细节。图8展示了Coruscant核的公式化。每个线程块由四个warp组成，每个warp映射到GPU处理器的子分区（SMSP）。因此，一个$256 \times 64$的矩阵操作数A的tile被进一步分解为分配给每个32线程warp的$64 \times 64$的tile。在一个warp内，每个线程被分配给两个长度为64的列。在“加载为压缩”阶段，每个线程将两个压缩的tile和相应的位图加载到其寄存器中，然后解压到共享内存中。在下一节中，我们将解释tile形状背后的原因，即避免共享内存bank冲突。每个warp通过迭代K维度来计算输出操作数的一个tile。在每次迭代中，一个矩阵A的tile和一个矩阵B的tile被用来执行分块矩阵乘法，直到输出tile被完全计算出来。


图8：Coruscant GPU核的SpMM公式化。矩阵（蓝色）被分解为分配给处理器的tile（紫色）。每个处理器迭代K维度（红色）。在一个处理器内，每个warp被分配给tile（绿色），每个warp线程被分配给两个列tile（黄色）。

Coruscant稀疏格式的解压。解压成完整矩阵是必要的，因为当前的张量核设计只支持完整矩阵tile或2:4压缩格式。当压缩列tile的非零元素和相应的位图位于每个线程的寄存器中时，每个线程按照算法1将指定的列解压到共享内存中。每个tile的解压过程会迭代压缩的非零向量，一旦迭代次数达到该tile中的非零元素数量，迭代就终止（第3行）。采用固定长度的tile大小64的迭代（第1行），而不是迭代动态数量的非零元素或通过位图进行循环，这对于向CUDA编译器提供静态的非零向量访问索引至关重要。这会促使编译器将非零元素放置在线程的寄存器中。如果非零向量的访问是动态的，编译器会将压缩的非零向量放置在本地内存的栈帧中，尽管有充足的寄存器，这会产生不必要且代价高昂的数据传输。我们还发现，使用前导零计数（CLZ）指令来推导每个非零元素的位置比迭代位图的每一位要快，后者会产生额外的内循环。然后，相应的1位从位图中移除（第6-7行）。计算出非零元素在线程块的$256 \times 64$ tile内的位置（第7-8行），该位置用于将非零元素放置在共享内存中的正确位置。

避免共享内存Bank冲突。Coruscant使用列式分块来避免共享内存bank冲突，这是解压到共享内存期间核性能的关键因素。GPU处理器的共享内存由32个bank组成，每个bank宽4字节。图9比较了不同分块方案（行式分块、矩形分块和Coruscant的列式分块）所产生的bank冲突。当一个线程的tile在行维度上扩展时，例如行式分块或矩形分块，就会发生bank冲突，即两个线程试图将其非零元素放置到同一个共享内存bank中。这会导致写操作串行化并产生额外的warp停顿。通过为每个线程分配两列，我们确保每个线程写入一个特定的共享内存bank，从而防止因共享内存bank冲突而产生的额外停顿。

其他优化。最后，为了减少对内存系统的请求或波前数量，每个tile的非零元素被填充到8的倍数，这样每个线程都可以通过单个LDG.128指令得到服务。为了最大化核的占用率，我们采用K维度切分，将K维度的迭代分开，并将每一部分分配给不同的线程块。这确保了计算在GPU处理器之间更均匀地分布。我们的实验结果（详见第5节）证明了上述减少全局内存和GPU处理器之间数据传输在内存密集型SpMM操作中的影响，因为Coruscant GPU核在70%稀疏度下比Flash-LLM实现了高达1.48倍的加速。此外，我们观察到Coruscant GPU核在低至30%的稀疏度下比密集基线cuBLAS和稀疏基线Flash-LLM【58】都快，这为LLM剪枝开辟了各种设计机会。


图9：(a) 行式分块和 (b) 矩形分块在多个线程同时解压到同一个bank时可能会遇到共享内存bank冲突，而在 (c) 列式分块中，每个线程写入不同的bank，从而完全避免了bank冲突的可能性。


图10：Coruscant核中解压所占的百分比。

3.3 Coruscant稀疏张量核

设计动机与目标。Coruscant稀疏张量核在Coruscant GPU核的基础上运行，以释放非结构化稀疏性的全部潜力，并利用Coruscant稀疏格式实现最大压缩。我们首先介绍稀疏张量核的设计动机，即消除解压到共享内存的开销，包括共享内存停顿和位图解码的计算指令。然后，我们介绍了Coruscant稀疏张量核的整体架构，以及我们为warp级MMA和张量核内部MMA操作引入的数据流变化。最后，我们介绍了位图解码器，这是Coruscant稀疏张量核的一个关键组件，它能直接操作Coruscant稀疏格式。Coruscant GPU核减少了从全局内存到每个处理器的数据传输，加速了内存密集型SpMM。然而，由于当前的GPU和张量核不原生支持我们的格式，核必须首先使用算法1将每个tile解压成共享内存中的密集矩阵，然后才能在张量核上执行。如图10所示，解压开销高达Coruscant核总执行时间的36%，这表明还有改进的空间。

架构改进。根据我们的观察，我们建议通过对稀疏张量核进行架构更改，来移除额外的解压指令以及从寄存器到共享内存的数据传输。具体来说，我们在张量核内部增加一个小单元，直接解码位于寄存器文件中的压缩非零向量。我们的方法对现有的稀疏张量核架构引入了最小的改动，并允许稀疏张量核在两种模式下运行：原始的密集模式和Coruscant稀疏模式。

性能提升原理。图11重申了Coruscant核和张量核修改的两阶段影响如何提高性能。由于核减少了全局内存到处理器内存的传输，它提高了内存密集型SpMM的操作强度——计算量相同，操作数字节数减少。通过提出的架构修改以在硬件中进行解压，图10的解压开销被移除，从而减少了核延迟并提高了吞吐量。


图11：Coruscant核和STC修改对roofline分析的影响。(a) 应用位图压缩通过减少数据移动来提高操作强度。(b) 将解压卸载到硬件消除了显式的解压指令和共享内存传输，减少了核延迟并提高了性能。


图12：Coruscant稀疏张量核概览。位图解码器（绿色）选择与该周期处理的密集向量元素相对应的所有线程的非零元素。

Coruscant微架构概览。图12提供了Coruscant稀疏张量核架构的概览。我们采用了Raihan等人【44，Md Aamir Raihan, Negar Goli, and Tor M. Aamodt. Modeling Deep Learning Accelerator Enabled GPUs. In 2019 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS), 2019】的Volta【39】张量核微架构模型，该模型每个周期能够计算128个MAC。在每个32线程的warp中，八个线程使用一个octet，四个线程使用一个线程组。在一个线程组内，每个线程编程四个乘法器和一个归约加法器树，以每周期执行4个MAC（一个lane）。自Volta以来，各种研究已经对不同代的GPU及其张量核进行了基准测试和逆向工程【1, 33, 51】，揭示了张量核的吞吐量每代翻倍。这促使我们在Volta【39】张量核的微架构基础上工作，然后扩展到每一代的吞吐量。位图解码器被放置在MAC单元的前面，以选择和流式传输相关操作数，如图12所示。我们将在接下来的小节和图14中更详细地解释。增加了一条旁路路径，允许张量核在原始的密集矩阵乘法模式下运行，从而支持密集和稀疏的warp级MMA指令。


图13：(a) warp级MMA的数据流。矩阵A的大小增加以容纳每个线程的压缩tile，而矩阵B的大小减少以保持每个周期的MAC数量不变。(b) 带有位图解码器的张量核内部数据流。


图14：位图解码器，一个交错到张量核流水线中的组合逻辑，用于每周期生成4个寄存器索引。

Coruscant稀疏张量核数据流。由于Coruscant稀疏格式将非结构化稀疏模式局部化到64元素tile，这要求Coruscant张量核采用不同的数据流来直接操作稀疏格式的tile。图13a说明了warp级MMA公式的变化。虽然Coruscant GPU核利用商业支持的HMMA指令，其输入维度为M=16, N=8, K=16，执行一个（$16 \times 16$）×（$16 \times 8$）的矩阵乘法，但我们现在遵循64元素的非结构化稀疏窗口。因此，一个MMA指令现在变为M=32, N=1, K=64，以执行一个（$32 \times 64$）×（$64 \times 1$）的矩阵乘法，这与原始HMMA指令的总MAC数量保持一致。由于无需解压，从而避免了共享内存bank冲突，我们重新制定了Coruscant稀疏格式，以压缩每个（$64 \times 1$）的行tile。因此，每个warp服务于一个MMA指令的（$64 \times 32$）矩阵A tile，每个线程提供一个大小为（$64 \times 1$）的压缩行tile。在稀疏张量核内部，我们保留了每个线程4个MAC的分配方案。图13b详细说明了Coruscant张量核内部的数据流。在每个周期，每个warp的密集向量B以四为一组进行处理。在每个周期中，所有线程tile的位图充当有效信号，以决定每个密集向量元素是否用于乘法。有效的密集向量位置与相应的非零元素相乘。每个线程累加矩阵C的一个元素，该元素被写回到寄存器文件。

位图解码器架构。图14展示了我们之前在图12中提出的位图解码器，它直接获取非零元素和相应的密集向量元素。位图解码器按warp操作，为每个周期的所有线程提供操作数。GPU SM中的每个线程在其寄存器中维护一个64位位图（❶）和一个最多64个FP16值的关联数组（❷）；在每个周期，一个线程级lane检测器（❸）解析位图的下四位，通过一串加法器（❹）计算相应非零FP16元素的索引，然后发出派生的寄存器地址，供每个线程读取相应的非零元素，同时将未使用的槽位填充为零，以每周期向张量核lane发送固定的4个元素（❺）。同时，密集向量指针增量器（❻）跟踪对密集矩阵操作数（❼）的共享访问索引，并且每周期四个连续的操作数元素被广播到所有线程lane。我们验证了Coruscant张量核的寄存器访问带宽不超过Volta张量核的带宽，Raihan等人【44】分析其为每个warp每周期384字节。位图解码器使用静态获取四个2字节元素的方式处理密集向量。对于压缩的非零元素，最大操作数读取是每个线程四个元素，导致每个warp 128个操作数。此外，还读取32个4位位图值和32个用于累加的操作数（每个线程一个）。在最大情况下，这需要读取164个2字节操作数和16字节的位图数据，总计每周期344字节，这远在Volta张量核的寄存器访问带宽之内。

A4 实验环境

  • 硬件配置：

    • GPU: RTX 6000 Ada，配备48GB VRAM。

  • 软件配置：

    • 代码库：使用Huggingface Transformer库【14】进行端到端评估，并移植了Coruscant核及稀疏张量核模拟。
    • 性能分析工具：使用Nsight Systems测量核时间，使用Nsight Compute收集GPU执行统计数据。
    • 硬件模拟与合成：稀疏张量核的行为通过在Coruscant核内移除位图解码和共享内存访问来模拟，类似于VectorSparse【71】的方法。位图解码器的RTL实现使用Design Compiler在45nm技术下合成，并使用Stillmaker等人【48】的公式缩放到7nm技术。

  • 模型与数据集：

    • 端到端评估模型：Llama 2 7B和13B【52】，使用Wanda【49】方法剪枝到30%、50%和70%的目标稀疏度。
    • 核性能评估模型：使用OPT 30B和66B【66】的权重矩阵，尺寸分别为(7168 × 7168)和(9216 × 9216)。

  • 实验任务与参数：

    • Llama 2 7B：输入序列32个token，生成512个token。
    • Llama 2 13B：输入序列64个token，生成1280个token。
    • 所有执行均使用FlashAttention【13】以最小化中间激活占用的内存。

  • 评估指标：

    • 端到端性能：每秒生成的token数（tokens per second）。

A4 实验结果

5.1 端到端LLM评估

  • 性能提升：如图15所示，与使用cuBLAS的密集模型相比，Coruscant在Llama 2 7B和13B上显著提升了token生成吞吐量。

    • 对于Llama 2 7B，Coruscant核在批大小为8/16/32时，吞吐量分别提升高达103/135/114 tokens/s。
    • 对于Llama 2 13B，Coruscant核在批大小为8/16时，吞吐量提升高达48/47 tokens/s。
    • Coruscant稀疏张量核（STC）在此基础上进一步提升性能，7B模型平均增加50 tokens/s，13B模型平均增加20 tokens/s。

  • 内存占用优势：如图15b所示，Coruscant的权重压缩带来了显著的内存优势。密集模型在处理Llama 2 13B、批大小为32的任务时因内存不足（OOM）而失败，而Coruscant通过压缩权重释放了GPU内存，成功执行了该任务，并将吞吐量提升了76 tokens/s（Coruscant核）和98 tokens/s（Coruscant STC）。


图15：端到端推理性能比较。标记缺失（Llama 2 13B, bs32）表示因内存溢出（OOM）错误而无法执行。

5.2 核性能与分析

  • 核延迟比较 (图16)：

    • Coruscant核在30%-70%的目标稀疏度范围内全面优于基线cuBLAS、Flash-LLM、SparTA和Sputnik。
    • Flash-LLM在低于50%稀疏度时性能不如cuBLAS，但在70%稀疏度时开始展现优势。
    • Coruscant核的性能优势随着稀疏度的增加而扩大，但在更高稀疏度下与Flash-LLM的差距缩小，因为后者的解压开销更小。Coruscant在LLM剪枝的目标稀疏区间内达到了最佳的性能平衡。

  • 内存占用减少 (图17)：

    • 与cuBLAS相比，Coruscant在30%至70%的稀疏度下，内存占用分别减少了14%至55%。
    • 与Flash-LLM相比，Coruscant在30%-70%的所有稀疏度下都表现出更低的内存占用，这不仅加速了内存密集型SpMM，还为KV缓存等其他张量释放了宝贵的GPU显存。

  • 内存开销分析 (图18)：

    • Coruscant的高压缩率显著降低了全局内存加载指令相关的平均warp停顿周期，优于cuBLAS和Flash-LLM。
    • 内存利用率数据显示，Coruscant通过压缩减少了数据传输，有效缓解了GPU全局内存的压力。


图16：不同稀疏度下的核延迟比较。


图17：N=8时核内存占用比较。


图18：内存统计数据比较。

5.3 稀疏张量核评估

  • 与Coruscant核的性能对比 (图19)：

    • Coruscant STC通过消除解压和共享内存写入，使张量核利用率平均提高了2.5倍。
    • 共享内存停顿周期减少到Coruscant软件核的平均5%。
    • 综合来看，Coruscant STC相比软件核实现了平均1.3倍的加速。

  • 面积与功耗 (表2)：Coruscant STC引入的位图解码器硬件开销极小。即使在最新的Hopper架构上，增加的面积也仅占GPU总面积的0.018%。

  • 与其他稀疏张量核设计的比较 (表3, 图20)：

    • 与DSTC和RM-STC相比，Coruscant STC的硬件复杂度和面积开销极低（分别减少了96%和95.4%）。
    • Coruscant的压缩率更高，因为它采用更大的tile尺寸（64），减少了非零指针的存储开销。在内存密集型的LLM推理中，Coruscant以压缩为中心的设计更为有效。

  • 与半结构化稀疏性的比较 (图21, 22, 表4)：

    • Coruscant软件核因解压开销而慢于硬件支持的2:4稀疏核。
    • Coruscant STC的性能与2:4核相当（50%稀疏度）甚至超越（70%稀疏度），证明了压缩带来的好处可以弥补计算跳过的不足。
    • 最关键的是，Coruscant支持的非结构化稀疏性在模型精度（困惑度和TriviaQA得分）上远优于半结构化稀疏性，后者在追求硬件效率的同时严重损害了模型质量。

  • 对通信开销的影响 (图23)：权重压缩不仅减少了片上内存传输，也显著降低了模型卸载（offloading）场景下从CPU/SSD到GPU的数据传输延迟，这对于多设备部署具有重要意义。

  • 可扩展性与通用性：
    • 稀疏范围 (图24a)：在<30%的稀疏度下，Coruscant核性能不如cuBLAS；在>70%时，Flash-LLM更优，但此时模型精度已严重下降。
    • 批处理大小 (图24b)：随着批处理大小（N）增加，计算密集度提高，Coruscant核的优势减弱。但由于长上下文推理限制了批大小，Coruscant的优势在实际场景中依然存在。Coruscant STC即使在N=256时仍有优势。
    • 任务通用性 (图25)：Coruscant在预填充（prefill）阶段比cuBLAS慢。对于生成长度短的任务（如摘要），端到端性能可能受影响。但对于生成长度长的任务（如对话、推理），解码阶段的持续加速能够摊销预填充的开销，实现总体加速。


图19：Coruscant核与Coruscant稀疏张量核的性能比较。

Table 2: Coruscant-STC的面积和功耗分解。

Table 3: 稀疏张量核的面积比较。可用数据来自各自的论文。


图20：在OPT30B操作数上，不同稀疏度下STC设计的性能比较。


图22：非结构化和半结构化稀疏性的性能比较。

Table 4: 不同稀疏度下半结构化和非结构化稀疏性的TriviaQA分数。


图21：与cuSPARSELt的比较。我们的方法使STC能够处理完全非结构化的稀疏性（30-70%），其性能水平可与NVIDIA的2:4半结构化硬件相媲美。


图23：卸载的OPT-66B GeMM的延迟比较。位图压缩减少了通过PCIe的传输时间。


图24：(a) 额外稀疏范围和 (b) 更大N维度下的延迟。红线表示cuBLAS。


图25：Llama 2 7B在不同生成长度下，prefill/decode阶段与cuBLAS的延迟差异。红色标签表示端到端延迟慢于密集推理。

A5 结论

本文为解决非结构化稀疏性的计算效率挑战，提出了一个由三部分组成的解决方案Coruscant。
首先，我们引入了Coruscant稀疏格式，这是一种基于位图的格式，在LLM剪枝的典型稀疏度范围（30%-70%）内，其压缩能力优于现有稀疏格式。
其次，我们开发了Coruscant SpMM核，它利用高度压缩的Coruscant稀疏格式来减少GPU全局内存与处理器之间的数据传输。与基线实现（cuBLAS）相比，Coruscant核实现了高达2.02倍的加速。
第三，基于Coruscant核的深入分析，我们发现通过对张量核进行微小的架构改动可以实现额外的加速。通过增加一个开销极小的位图解码器，在Coruscant稀疏张量核上运行的Coruscant GPU核相较于基线实现了高达2.75倍的加速。
将Coruscant应用于Llama 2模型，我们观察到token生成吞吐量显著增加：仅使用软件核就额外增加了135 tokens/s，而借助Coruscant稀疏张量核则额外增加了206 tokens/s，最高可达847 tokens/s的生成吞吐量。