SparseDitto:突破350倍格式悬崖,用LLM闭环定制GPU稀疏内核
SparseDitto: Customizing GPU Kernels for Different Sparsity Patterns with LLM-Based Agentic System

在 GPU 高性能计算与深度学习底层系统中,稀疏矩阵计算始终是一个充满“性能陷阱”的领域。与计算模式规整、内存访问连续的密集矩阵乘法(GEMM)不同,稀疏矩阵向量乘(SpMV)、稀疏矩阵稠密矩阵乘(SpMM)以及稀疏矩阵相乘(SpGEMM)的实际执行效率,不仅取决于矩阵的维度,更完全被矩阵中非零元的空间分布形态所左右。同一个算子在同一块 GPU 上执行相同的矩阵,仅仅因为采用了不同的存储格式与线程调度策略,性能差距就可以达到令人瞠目的程度——例如在英伟达官方数学库 cuSPARSE 中,针对同一个矩阵执行 SpMM,压缩稀疏行(CSR)格式与分块 ELL(Blocked-ELL)格式之间的耗时差距最高可达 350 倍。
ArXiv URL:https://arxiv.org/abs/2608.05033v1
针对这种高度敏感的算子特性,系统研发人员通常有两种选择:要么依赖 cuSPARSE 等通用库,但这往往意味着在面对长尾分布矩阵时必须承受巨大的性能折损;要么针对特定场景手写专用算子(如面向 Tensor Core 的 DTC-SpMM,或面向自适应分块的 CB-SpMV),但这些系统往往固化了特定的结构假设,一旦矩阵分布或底层微架构发生变化,性能就会急剧劣化。来自美国东北大学与明尼苏达大学的研究团队提出了全新解决方案 SparseDitto。该系统将大语言模型(LLM)的开放代码生成能力与基于真实硬件测量的闭环反馈相结合,为任意给定的稀疏矩阵、算子和目标 GPU 动态定制最优的底层 CUDA 内核。
评测显示,在覆盖 60 个代表性矩阵与多种算子的全量测试中,SparseDitto 在 NVIDIA RTX PRO 6000 上相比 cuSPARSE 实现了 2.68 倍的几何平均加速,峰值加速比达到 146.61 倍;在 NVIDIA H200 上则取得了 2.79 倍的几何平均加速,最高达 78.5 倍。此外,在全图图卷积网络(GCN)的端到端训练中,由 SparseDitto 生成的定制 SpMM 算子将单个 Epoch 耗时缩减至 cuSPARSE 的三分之一以下。这项研究的核心价值不仅在于刷新了稀疏算子的性能天花板,更在于证实了基于结构特征引导的 LLM 智能体系统完全有能力跳出人类专家的固定模板库,在严苛的物理硬件约束下自主探索出更优的系统级代码。
稀疏算子的非对称鸿沟与专门系统的局限
要理解 SparseDitto 的设计动机,必须先看清稀疏矩阵在 GPU 硬件上的运行困局。对于一个稀疏矩阵而言,行长度的方差直接决定了线程块(CTA)与线程束(Warp)之间的负载均衡;列索引的离散程度决定了访问稠密矩阵或向量时的显存合并与缓存命中率;非零元聚集成块的紧密程度决定了是否值得启用分块格式以利用 Tensor Core;而在 SpGEMM 这种双稀疏输入算子中,中间乘积的产生速度与碰撞概率更是直接统治着片上累加器的排队与同步延迟。
稀疏优化的代价往往存在显著的非对称性:采用适配当前稀疏模式的表示与调度方案,通常可以获得数倍的稳健加速;然而,一旦选择了不匹配的策略,系统将付出高达数十甚至数百倍的灾难性代价。以 Blocked-ELL 格式为例,这种格式将非零元强行对齐到固定大小的稠密小块中,以便在硬件层面走通高吞吐的张量计算路径。如果目标矩阵本身天然具备稠密子块,这种设计能带来极致的计算能效;但如果目标矩阵具有极高的离散度与不规则的行长,为了满足分块对齐,系统不得不填充海量的零元(Padding),导致绝大多数显存带宽和张量计算资源都在搬运与计算完全无用的假数据。
近年来学术界提出了多种专用的稀疏系统与编译器,如面向 SpMV 的 CB-SpMV、面向 SpMM 的 DTC-SpMM、面向 SpGEMM 的 HSMU-SpGEMM,以及基于 TVM 生态的稀疏编译器 SparseTIR。这些系统在特定输入与特定硬件上表现出色,但各自的架构假设非常狭窄。DTC-SpMM 假设矩阵能够被有效压缩为 Tensor Core 所需的小块;HSMU-SpGEMM 依赖共享内存哈希表来进行稀疏累加,一旦中间产物数量超出片上共享内存容量,性能同样会骤降。因此,任何固定的内核实现,实际上都在代码层面同时冻结了两个假设:它所假定的稀疏模式,以及它在编写时所针对的目标 GPU 微架构特性。
统一设计框架与整体架构
为了打破这种固化假设,SparseDitto 首先将稀疏 GPU 内核的构建过程抽象为一个统一的设计空间 $\Pi = \langle R, S, \theta_H \rangle$。其中,$R = \langle L, I \rangle$ 代表数据表示层,包含具体的非零元排布格式 $L$(如 CSR、COO、SELL、BSR 等)与稀疏元数据组织方式 $I$;$S = \langle P, D \rangle$ 代表执行调度层,定义了工作负载的切分粒度 $P$ 与线程级的数据流映射 $D$;$\theta_H$ 则表示与目标硬件强绑定的执行配置,涵盖网格与线程块尺寸(Grid/Block Dimensions)、寄存器分配上限以及共享内存配额等。

上图完整展示了 SparseDitto 的分层协作管线。系统并非盲目地让语言模型从零编写 CUDA 代码,而是构建了一套“结构分析—能量排序—架构规划—代码合成—实测闭环”的严密工程链条:
-
结构化特征提取:在不运行任何 GPU 计算内核的前提下,通过轻量级扫描快速提取输入矩阵的 36 维结构特征;
-
基于能量模型的策略初筛:利用可解释的加性模型对已知的优化策略进行打分,快速锚定搜索的大致方向;
-
架构感知的层次化规划:将策略候选与目标硬件的具体参数(如 SM 数量、共享内存上限、寄存器约束)结合,实例化为具体的候选规范;
-
代码生成与硬件闭环调优:由代码生成智能体与验证智能体在目标 GPU 上进行编译、数值精度校验与性能 Profiling,通过多轮迭代收敛出最优实现。
从结构感知到可解释策略排序
自动化内核生成面临的首要挑战是搜索空间的组合爆炸。未受约束的 LLM 即使能生成语法正确的 CUDA 代码,也极大概率会选择错误的并行粒度或内存布局。SparseDitto 解决这一问题的抓手是 36 个物理特征指标,涵盖四类核心信息:18 个矩阵固有模式特征(如非零元密度分布、行长变异系数、带宽与对角聚集度)、8 个由表示格式引入的开销特征(如 ELL 或 BSR 格式下的填充率估计)、10 个算子诱导的特征(如 SpGEMM 中基于压缩估计得到的中间产物分布与累加器压力),以及目标 GPU 的硬件配置。
在获取这 36 个特征后,系统通过一个轻量级的可解释加性模型(Additive Energy Model)对专家经验库中的候选策略进行先验排序。候选策略集合涵盖了学术界和工业界经过验证的主流方案,例如 SpMV 中的自适应二维分块与不同切片高度的 SELL 格式,SpMM 中的 DTC 块状张量核心调度与列窗口打包技术,以及 SpGEMM 中的片上共享内存哈希与混合累加方案。
与黑盒分类器不同,该模型将每种策略的偏好能量建模为各个特征变换后的线性叠加:
\[E_o(s \mid x) = \beta_{o,s} + \sum_j h_j(x_j)^{\mathsf{T}} w_{2,o,j,s}\]这种设计的精妙之处在于可解释性与低计算开销。每个特征在最终评分中的贡献都是显式解耦的。分析表明,模型并非机械依赖单一特征,而是综合衡量非零元规模、行长离散度与算子需求。例如在 SpMM 中,稠密维度 $K$ 的大小会显著改变策略走向;而在 SpGEMM 中,中间结果上限与行合并冲突率则直接决定了累加器的选型。更重要的是,这个能量模型输出的不是最终的固定代码,而是一组排名前列的设计路线建议,为后续规划器留出了充足的自由探索空间。
架构感知的规划器与闭环生成
得到策略排序后,如何将其落地为符合具体 GPU 硬件约束的代码设计?这由层次化架构感知规划器(Hierarchical Architecture-Aware Planner)完成。

如上图所示,规划器接收矩阵结构特征、能量模型的策略排名以及目标 GPU 的物理规格(如 SM 数量、Warp 调度限制、共享内存与寄存器文件大小)。规划器的职责是对底层资源进行显式建模与边界截断。它会精确推算当前配置所需的元数据显存占用 $M_{\mathrm{meta}}$、中间计算工作空间大小 $W$ 以及动态分配的共享内存 $S_t$,确保候选方案满足严格的硬件预算:
\[B_t \le T_{\max}, \quad r_t B_t \le R_{\max}, \quad S_t \le S_{\max}\]这一机制从根本上避免了 LLM 常见的“幻觉参数”——例如在共享内存中开辟超出物理硬件规格的缓冲区,或者因单个线程占用寄存器过多导致 Warp 占用率极低甚至编译失败。
规划器会同时派生出多条相互独立的搜索分支。在初始阶段,算力预算被均匀分配给各个分支,以保证策略的多样性;随着搜索推进,实测性能更优且通过数值验证的分支将获得更多的迭代配额。
在执行层面,代码生成智能体负责将具体的规划草案翻译为高质量的 CUDA 源码,并配齐对应格式的元数据预处理脚本。验证智能体则在物理 GPU 上编译执行,对比参考实现(如 cuSPARSE)输出的数值结果。针对 SpGEMM 这种输出非零元结构本身就动态可变的算子,验证模块不仅校验非零元的值,还会在列索引排序合并后严格比对输出的行指针、列索引与非零元总量,确保绝对与相对误差均在可信阈值内。通过不断回传真实的编译错误信息、硬件 Profiling 指标(如 Warp 活跃度、缓存命中率)给 LLM,系统以测量驱动的闭环完成多轮微调。
硬件实测:击败 cuSPARSE 与专用系统的根源
研究团队在 SuiteSparse 矩阵集的 60 个代表性稀疏矩阵上进行了全面测试,覆盖了网络图、有限元网格、电路仿真等多种截然不同的拓扑结构。在 RTX PRO 6000 与 H200 两个完全不同的微架构平台上,SparseDitto 在 94% 的任务中均优于 cuSPARSE。在 RTX PRO 6000 上,SpMV、SpMM(不同列宽 $K \in {8, 32, 128, 256}$)与 SpGEMM 的整体几何平均加速比为 2.68 倍;在配备 HBM3e 的顶级加速卡 H200 上,这一数字进一步提升至 2.79 倍。
为了揭示性能跃升的深层硬件机制,研究人员使用 NVIDIA Nsight Compute 深入剖析了内核在底层显存层次结构中的行为。以面向 Tensor Core 的 DTC-SpMM 为例,Nsight 分析给出了极具反常识的结论:DTC-SpMM 的缓存命中率并不低,且从设备显存搬运的总字节数比 SparseDitto 略少,然而在实际执行时间上,SparseDitto 却比 DTC-SpMM 快了整整 3.99 倍。
产生这一现象的根本原因在于有效计算密度与硬件管线占空比。DTC-SpMM 强行采用 $16 \times 8$ 的稠密小块,但在测试集的大多数矩阵上,小块内的实际非零元占用率中位数仅有 10.1%。这意味着系统搬运的数据中近 90% 都是为了迎合格式而填充的无效零。尽管缓存命中率高,但发射出的指令几乎都在做无用功。探测表明,DTC-SpMM 的 Tensor Core 流水线利用率仅达到理论峰值的 1.5%,且由于寄存器和共享内存占用过大,驻留 Warp 数量只有 SparseDitto 的三分之一,完全无法掩盖访存延迟。而 SparseDitto 则根据矩阵分布,自适应地选择了更精简的切片布局或线程级映射,保持了极高的数据实际利用率和极低的空转开销。
在算子表现的分化上,SpGEMM 展现了最宽广的性能分布区间,峰值加速比达到了 146.61 倍。这是因为与单纯由访存带宽受限的 SpMV/SpMM 相比,双稀疏相乘所产生的中间产物乘积数量在不同矩阵间存在数量级的剧烈波动,通用库的固定哈希或合并策略极易发生严重的哈希碰撞或频繁的全局内存回写,而定制内核通过行分箱(Row Binning)与自适应片上工作区彻底消解了这一瓶颈。
突破预设经验:自主涌现的全新算子设计
SparseDitto 最引人注目的能力在于:它并非仅仅是在人类已知的格式与超参数之间做网格搜索,而是具备“跳出预设模板库”的代码合成能力。在超过三分之一的获胜测试用例中,最终跑出最高性能的内核所采用的设计策略,甚至根本没有出现在能量模型的候选词表 $V_o$ 中。
一个典型的案例来自流体动力学模拟矩阵 Shen/e40r0100。该矩阵拥有 17,281 行,平均每行仅有 32 个非零元,密度为 $1.9 \times 10^{-3}$,在计算 $AA$ 矩阵相乘时,平均每行会派生出 967 个中间项。如果严格遵循能量模型的首选建议,系统会采用经典的混合累加策略,该分支在实测中相比 cuSPARSE 取得了 1.55 倍的加速。
然而,架构感知规划器派生出的另一个自由搜索分支,却根据特征分析挖掘出了该矩阵独特的局部拓扑特性:该矩阵呈现出强烈的对角优势,中位数行带宽仅为 481 列,这意味着任意一行计算出的中间产物实际上高度集中在特定的列区间内,而绝非离散地随机撒向整个 17,281 列的全局空间。基于这一观察,智能体自主构造了一个专家库中从未定义过的累加器架构:
-
它为每个输出行分配一个 CTA,并在共享内存中开辟一段长度为 4,096 列的稠密值条带(Value Stripe),辅以代际标记(Generation Tag)以消除反复重置内存的开销;
-
在预处理阶段,系统将输入矩阵按照列区间进行分段打包,使得每个 CTA 仅需精准流式读取该条带所涵盖的列分段;
-
所有的局部合并直接在私有稠密条带中无锁就地累加,全程不调用任何全局原子操作(Atomic Operation)。
这个完全由智能体自主设计并落地的无锁条带内核,最终在物理硬件上跑出了 8.47 倍的惊人加速,将固定策略分支远远甩在身后。这一事实有力地证明了,当赋予 LLM 严密的结构特征感知与真实的底层微架构反馈时,智能体完全能够根据工作负载的独特物理属性,发明出超越人类通用库与既有规则模板的特定领域代码。
在端到端应用评估中,将 SparseDitto 生成的 SpMM 内核直接替换进经典的图神经网络模型——基于 Reddit 数据集的两层全批次图卷积网络(GCN)进行训练,在特征维度 $K=8$ 时将训练 Epoch 耗时从 24.6 毫秒大幅压缩至 7.3 毫秒,实现了 3.39 倍的端到端提速;即便在 $K=256$ 的计算密集场景下,依然稳健保持了 1.14 倍的端到端加速优势。
总结与展望
在 GPU 体系结构向着极度异构与专用计算单元(如 Tensor Core)演进的今天,底层软件栈正面临着严峻的碎片化危机:硬件微架构迭代愈发迅速,而稀疏模式的多样性又使得一套代码打天下的传统库开发模式难以为继。
SparseDitto 的成功实践展示了一条全新的系统软件演进路径。它表明,解决复杂高性能计算问题的有效范式,并非单纯依赖更大参数量的模型进行端到端黑盒猜想,而是将领域专家的结构化先验转化为严密的特征表达与资源边界约束,再让 LLM 智能体在明确的物理约束空间中施展自由代码生成与重构能力,最终依托真实物理硬件的执行测量来指导收敛。这种闭环机制既消解了传统编译器的保守性与开发周期长的问题,又彻底驯服了大语言模型在底层代码生成中的“幻觉”隐患,为未来智能化的算子自动调优与系统底层软件开发提供了极具启发性的范式参考。