SparseDitto:用LLM Agent自动生成匹配稀疏模式的GPU Kernel

📅 2026/8/27 6:08:57
SparseDitto:用LLM Agent自动生成匹配稀疏模式的GPU Kernel
之前在做模型稀疏化加速时我遇到过一种很尴尬的情况模型剪枝之后计算量确实降了但直接跑在 GPU 上推理速度反而没有明显提升甚至偶发变慢。调查下来发现问题往往出在稀疏模式与 GPU Kernel 的匹配度上——非结构化稀疏和 2:4 结构化稀疏在访存方式、线程分配、指令选择上的差异远比“把零跳过”复杂得多。如果每个稀疏模式都靠人工写一套 CUDA 或 Triton Kernel工作量非常大而且难以维护。SparseDitto 这个方向正是针对“不同稀疏模式需要不同 GPU Kernel”的工程痛点提出的它用 LLM 驱动的 Agentic System 自动分析稀疏模式、定制内核、编译评测闭环。本文围绕这个概念从稀疏模式与并行计算的底层关系讲起再拆解 SparseDitto 的系统设计思路最后给出一套原理性演示代码和工程落地建议。适合正在做模型推理加速、GPU Kernel 开发和自动化编译优化的开发者阅读。1. 背景稀疏计算与 GPU Kernel 的“定制困境”1.1 稀疏性看似“省算力”实则“费工程”深度学习模型中的权重矩阵在经过剪枝Pruning之后会产生大量零元素。直观来看零元素不参与有效计算跳过它们就能减少浮点运算量也就是常说的 FLOPs 下降。但 GPU 是一个高度并行的 SIMT 架构所有线程按照 warp 调度执行指令。如果矩阵中的零元素分布不均匀直接跳过零会导致两种严重后果线程束发散同一个 warp 内不同线程计算量不同部分线程在忙部分线程空转最终执行时间以最慢的线程为准。访存模式不规则非结构化稀疏矩阵需要使用 CSR、CSC 这类压缩存储格式每个元素的坐标需要额外保存访存不再是规则的连续访问shared memory 和 cache 的利用率都会下降。所以稀疏矩阵要真正跑得快不只是“省算力”还要设计一个符合硬件特性的 Kernel这本身就是一项专门的工程能力。1.2 为什么不同稀疏模式需要不同的 GPU Kernel稀疏模式Sparsity Pattern决定了矩阵中零元素的位置规律。常见的有非结构化稀疏零元素完全随机分布压缩率高但难以高效映射到 GPU 硬件。结构化稀疏按照行、列、块、通道等规律剪枝例如 NVIDIA 的 2:4 结构化稀疏每 4 个元素保留 2 个。Block 稀疏以固定大小的块为单位进行稀疏化比如 32x32 的块要么保留要么丢弃。Fine-grained细粒度与 Coarse-grained粗粒度稀疏粒度不同对应的计算密度也不同。很关键的一点在于GPU Kernel 的优化策略高度依赖这些模式。Block 稀疏可以用 BSRBlock Compressed Sparse Row格式按块加载数据减少索引开销2:4 结构化稀疏可以使用硬件提供的稀疏矩阵指令非结构化稀疏则往往需要特殊的数据重排和 lookup table。这意味着一种 Kernel 不可能适配所有稀疏模式。传统做法是为每种模式投入大量人工开发这也是很多团队做模型稀疏化时“理论加速比很高、实际收益却不高”的核心原因。2. SparseDitto 是什么用 LLM 智能体系统做内核定制2.1 从人工调优到自动化定制SparseDitto 这个名字很有意思。Sparse 对应稀疏Ditto 可以理解为“复制后再适配”的含义。整体来看它要解决的问题是当你拿到一个新的稀疏模式时能否自动生成一个匹配该模式的高性能 GPU Kernel而不是靠人工从头写。这个思路比传统的 AutoTVM、Ansor 更进一步。传统自动调优系统主要在一个预设的代码模板空间里搜索参数比如 tile 大小、unroll 因子而 SparseDitto 这类 LLM-Based Agentic System 是在“代码结构”层面做生成和搜索Agent 可以增加循环变换、改变数据布局、重写访存逻辑甚至组合多个 Kernel。2.2 LLM-Based Agentic System 解决什么问题LLM-Based Agentic System 是指以大型语言模型为核心决策组件由多个 Agent 协作完成复杂任务的系统。在 SparseDitto 的场景里Agent 系统承担的角色包括理解稀疏模式的统计特征和存储格式。根据硬件类型如 A100、H100、消费级显卡选择 Kernel 模板。生成候选 Kernel 代码。编译、运行、采集性能数据。分析性能瓶颈迭代修改代码。这里的关键不是“让 LLM 写一段能跑的代码”而是让 Agent 像工程师一样形成一个“分析—生成—评测—再修改”的闭环。这也是 Agentic System 和普通代码生成模型的最大区别系统具备反馈和迭代能力而不是一次性输出结果。3. 稀疏模式与并行计算的底层关系3.1 常见稀疏模式的数据结构差异理解 SparseDitto 之前有必要先梳理几种常见稀疏模式在底层数据结构上的差异。CSRCompressed Sparse Row常用于非结构化稀疏核心是三个数组values、column indices、row pointers。读取第 i 行的非零元素需要通过 row_ptr[i] 和 row_ptr[i1] 定位区间。优点是存储紧凑缺点是访存时 column index 需要额外读取且非零元素的位置不规则。BSRBlock Compressed Sparse Row是 CSR 的块扩展版把矩阵分成固定大小的块只保存非零块。它更适合计算密集的矩阵乘法因为块内数据连续可以高效加载到 shared memory 或寄存器。Bitmask 方式则用位图标记哪些位置是非零适合规则稀疏模式比如 2:4 稀疏。NVIDIA 的 cuSPARSE 和 Tensor Core 对这类模式有专门的指令支持。这些格式差异直接决定了 Kernel 的循环结构和访存指令不能混用。3.2 稀疏模式对访存、负载均衡和指令选择的影响从 GPU 编程视角看稀疏模式的影响体现在三个层面。第一个层面是访存。非结构化稀疏的数据往往需要 gather 操作即根据索引取数这种随机访问很难预取容易造成 cache missBlock 稀疏的数据则更适合用连续内存加载甚至可以使用向量化 load 指令一次读取多个元素。第二个层面是负载均衡。GPU Kernel 通常把矩阵分给不同的 thread block 处理。如果每个 block 分配到的非零元素数量差异过大会出现计算资源闲置。SparseDitto 这类系统需要分析稀疏模式的空间分布选择合理的分区策略比如按行分区、按非零数量分区或者按块分区。第三个层面是指令选择。现代 GPU 提供了很多专用指令。Tensor Core 上的稀疏矩阵乘法指令要求特定的稀疏格式普通 CUDA Core 的 FMA 指令则更通用。对 Kernel 生成系统来说能否根据稀疏模式识别出“当前模式适合走 Tensor Core 还是 CUDA Core 路径”决定了最终的性能上限。4. 系统设计与核心模块拆解基于 LLM-Based Agentic System 来做 GPU Kernel 定制SparseDitto 一类系统通常包含以下核心模块。4.1 稀疏模式分析模块这个模块负责接收用户输入可能是模型权重文件、稀疏矩阵文件或者剪枝配置然后输出一个结构化的“稀疏模式描述”。具体做什么呢比如读取一个剪枝后的权重矩阵统计非零率、非零元素分布是否符合 2:4 规律、是否可以划分为固定大小的 block、行间非零数量的方差等。这些统计结果会成为后续内核生成模块的约束条件。实际工程中这个模块可以输出类似这样的元数据{ matrix_shape: [4096, 4096], non_zero_ratio: 0.25, pattern_type: structured_2_4, block_size: [2, 4], row_balance_score: 0.98, suggested_format: bitmask, hardware: cuda:0 }这段描述的作用是告诉 Agent 系统“当前稀疏模式是什么适合用什么数据结构硬件环境是什么”从而缩小代码搜索空间。4.2 内核生成与搜索模块这是系统的核心。拿到稀疏模式描述后多个 Agent 会协作生成候选 Kernel。第一个 Agent 负责选择 Kernel 模板比如基础矩阵乘法模板、稀疏感知模板、融合偏置和激活函数的模板。第二个 Agent 负责填充模板中的关键参数比如 tile 大小、block 大小、unroll 因子、是否使用向量化访存。第三个 Agent 负责修改循环结构比如把全局内存访问改成 shared memory 分块加载。与传统模板搜索不同LLM Agent 的优势在于可以跨模板重组代码而不是只在固定参数空间里做枚举。比如检测到稀疏模式是 Block 稀疏Agent 可以主动把外层循环从“遍历行”改成“遍历非零块”并为每个 block 生成独立的计算谓词这种方式已经超出了传统自动化调优的搜索范围。4.3 评测与反馈闭环生成的 Kernel 不可能一次就达到最优性能因此评测与反馈闭环非常重要。系统会完成以下步骤编译候选 Kernel。在目标 GPU 上运行采集耗时、显存占用等指标。与 baseline 或上一次迭代的 Kernel 对比。将性能数据反馈给 LLM Agent。Agent 基于这些反馈决定是否修改代码。比如当发现访存受限时可能会尝试更小的 tile 或不同数据布局当发现计算受限时可能会增加 ILP指令级并行或切换成 Tensor Core 路径。这个闭环本质上把“人工性能调优经验”封装成了可执行的 Agent 策略也是 SparseDitto 区别于一次性代码生成工具的关键设计。5. 关键技术流程演示原理性实现需要提前说明以下不是 SparseDitto 项目的官方代码而是为了理解系统原理设计的一套教学演示。核心思路是通用的你可以根据实际项目需求替换为生产级实现。5.1 环境准备本文示例以 Python Triton 为主因为 Triton 写 Kernel 比 CUDA 更简洁且便于演示“由 Agent 生成 Kernel”的思路。建议环境Python 3.9 以上。PyTorch 2.x演示矩阵乘法时用到。Triton通过pip install triton安装。一个 NVIDIA GPU建议显存 8GB 以上。版本需要根据实际环境调整重点演示设计思路。如果暂时没有 GPU也可以把代码中的 Kernel 执行部分替换成单元测试先跑通 Agent 逻辑再在 GPU 上验证。5.2 稀疏模式描述与元数据先定义一个稀疏模式描述器它分析权重矩阵并输出元数据。这是整个 Agent 系统输入的起点。# 文件路径sparse_ditto_demo/sparsity_analyzer.py import numpy as np def analyze_sparsity(matrix: np.ndarray): 分析稀疏矩阵的基本模式返回描述字典。 注意这里只统计通用特征实际项目还会结合硬件特性做更多分析。 total matrix.size non_zero np.count_nonzero(matrix) non_zero_ratio non_zero / total rows, cols matrix.shape # 简单检查 2:4 结构每 4 个连续元素中是否有 2 个非零 structured_2_4 True if cols % 4 0: reshaped matrix.reshape(-1, 4) for i in range(reshaped.shape[0]): if np.count_nonzero(reshaped[i]) ! 2: structured_2_4 False break else: structured_2_4 False # 行间非零数方差用于判断负载均衡难度 row_nnz np.count_nonzero(matrix, axis1) row_balance_score 1.0 - float(np.std(row_nnz) / (np.mean(row_nnz) 1e-6)) return { matrix_shape: [rows, cols], non_zero_ratio: round(float(non_zero_ratio), 4), pattern_type: structured_2_4 if structured_2_4 else unstructured, row_balance_score: round(row_balance_score, 4), } if __name__ __main__: mock_weight np.zeros((16, 16), dtypenp.float32) for i in range(0, 16, 2): mock_weight[i, 0:2] 1.0 mock_weight[i, 2:4] 0.0 print(analyze_sparsity(mock_weight))运行后输出类似{ matrix_shape: [16, 16], non_zero_ratio: 0.125, pattern_type: unstructured, row_balance_score: 0.0 }这个输出的意义在于Agent 拿到pattern_type和row_balance_score之后就能决定下一步走什么优化路径。比如structured_2_4优先走 Tensor Core 稀疏推理路径unstructured则考虑用查表法或负载均衡重排。5.3 基于 LLM 的 Agent 编排逻辑下面用 Python 模拟 LLM Agent 系统的调度逻辑。实际项目中这里可以接入 OpenAI API、本地 LLM 或者微调后的代码模型本文用规则函数代替 LLM 输出方便你理解整个链路。# 文件路径sparse_ditto_demo/agent_scheduler.py from sparsity_analyzer import analyze_sparsity import numpy as np class LLMAgent: 模拟 LLM Agent 的决策接口。 真实场景中这里的 decide_kernel_strategy 会调用大模型生成代码 本文直接用规则返回策略方便演示系统流转流程。 def decide_kernel_strategy(self, meta: dict) - dict: if meta[pattern_type] structured_2_4: return { kernel_template: tensor_core_2_4, data_format: bitmask, tile_size: [128, 128], suggested_unroll: 4, } if meta[non_zero_ratio] 0.1: return { kernel_template: block_sparse_bsr, data_format: bsr, block_size: [32, 32], use_shared_memory: True, } return { kernel_template: csr_custom, data_format: csr, tile_size: [64, 64], load_balance: row_balanced, } def build_agent_task(matrix: np.ndarray): meta analyze_sparsity(matrix) agent LLMAgent() strategy agent.decide_kernel_strategy(meta) return meta, strategy if __name__ __main__: matrix np.random.rand(256, 256).astype(np.float32) matrix[matrix 0.9] 0.0 meta, strategy build_agent_task(matrix) print(稀疏模式描述, meta) print(内核生成策略, strategy)这里面有两个要点。第一analyze_sparsity的结果不只是给人看的它充当了 Agent 的“观测状态”。第二decide_kernel_strategy返回的kernel_template、data_format、tile_size等字段会直接驱动下一步的内核代码拼装。如果换成真实 LLM这里就是大模型生成 JSON 或者代码的位置可以用 structured output 约束格式。5.4 生成一个 Triton Kernel 示例当 Agent 决定使用 block 稀疏策略时可以生成类似下面的 Triton Kernel。这个 Kernel 的思路是按固定 block 扫描矩阵只对包含非零数据的块执行乘加运算。# 文件路径sparse_ditto_demo/kernels/block_sparse_matmul.py import torch import triton import triton.language as tl triton.jit def block_sparse_matmul_kernel( a_ptr, b_ptr, c_ptr, rows, cols, k_dim, block_k: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, ): 简化的 block 稀疏矩阵乘法演示 Kernel。 实际场景中Agent 会根据稀疏模式调整 BLOCK_M / BLOCK_N / block_k 并在循环内跳过全零块。这里保留完整循环便于理解结构。 pid_m tl.program_id(0) pid_n tl.program_id(1) offs_m pid_m * BLOCK_M tl.arange(0, BLOCK_M) offs_n pid_n * BLOCK_N tl.arange(0, BLOCK_N) offs_k tl.arange(0, block_k) a_ptrs a_ptr offs_m[:, None] * k_dim offs_k[None, :] b_ptrs b_ptr offs_k[:, None] * cols offs_n[None, :] acc tl.zeros((BLOCK_M, BLOCK_N), dtypetl.float32) for k_start in range(0, k_dim, block_k): a tl.load(a_ptrs, maskoffs_k[None, :] k_dim - k_start, other0.0) b tl.load(b_ptrs, maskoffs_k[:, None] k_dim - k_start, other0.0) acc tl.dot(a, b) a_ptrs block_k b_ptrs block_k * cols offs_cm pid_m * BLOCK_M tl.arange(0, BLOCK_M) offs_cn pid_n * BLOCK_N tl.arange(0, BLOCK_N) c_ptrs c_ptr offs_cm[:, None] * cols offs_cn[None, :] tl.store(c_ptrs, acc) def block_sparse_matmul(a: torch.Tensor, b: torch.Tensor) - torch.Tensor: assert a.shape[1] b.shape[0] rows, k_dim a.shape cols b.shape[1] c torch.empty((rows, cols), devicea.device, dtypetorch.float32) BLOCK_M 16 BLOCK_N 16 block_k 16 grid (rows // BLOCK_M, cols // BLOCK_N) block_sparse_matmul_kernel[grid]( a, b, c, rows, cols, k_dim, block_kblock_k, BLOCK_MBLOCK_M, BLOCK_NBLOCK_N, ) return c if __name__ __main__: A torch.randn(64, 32, devicecuda, dtypetorch.float32) B torch.randn(32, 64, devicecuda, dtypetorch.float32) C block_sparse_matmul(A, B) print(输出形状, C.shape) print(与 torch 结果是否接近, torch.allclose(C, A B, atol1e-2))这是一个最简版本的 Triton 矩阵乘法 Kernel用来展示“Agent 生成的代码大概长什么样”。实际 SparseDitto 类系统会在这个内核基础上增加稀疏块跳过逻辑比如把非零块索引提前构建好循环时直接按索引加载而不是逐个遍历。这里故意保留完整遍历是希望代码结构更容易理解也不会因为过度精简在某些 Triton 版本上无法运行。5.5 运行与验证思路把上面几个文件按目录放好sparse_ditto_demo/ ├── sparsity_analyzer.py ├── agent_scheduler.py └── kernels/ └── block_sparse_matmul.py分别运行python sparse_ditto_demo/agent_scheduler.py python sparse_ditto_demo/kernels/block_sparse_matmul.py第一个脚本输出稀疏模式描述与内核生成策略验证 Agent 链路。第二个脚本在 GPU 上执行一个 Triton Kernel验证生成的代码可以运行。这里要强调真实的 SparseDitto 系统会复杂得多包括 kernel 编译错误重试、benchmark 数据采集、多轮 prompt 迭代等。但核心链路就是这个闭环——观察稀疏模式、决策内核策略、生成代码、运行验证。6. 工程落地中的常见问题与排查思路在实际构建类似系统时有几类高频问题值得提前关注。问题现象常见原因解决思路LLM 生成的 Kernel 编译失败代码与硬件架构不匹配或使用了不存在的 APIAgent 增加编译错误反馈循环解析报错后重写代码同时做好模板约束生成的 Kernel 性能反而不如 baseline调度策略选择错误比如非结构化稀疏走了 dense 路径完善稀疏模式分析模块增加硬件性能计数器的数据反馈Agent 迭代次数过多优化时间过长反馈信号设计不合理策略空间过大先用规则缩小候选模板范围再做细粒度代码生成Triton Kernel 在部分 GPU 上报错Triton 版本与 CUDA 版本不兼容固定 Triton 版本使用容器化环境复现记录每轮生成的运行环境稀疏模式分析耗时过长统计逻辑使用了 Python 循环而非向量化操作用 PyTorch 或 numpy 的批量运算改写统计逻辑必要时做采样估算这些问题的预防思路其实是一致的不要让 LLM 完全自由发挥要给 Agent 设置好边界比如模板库约束、编译错误模板解析、性能对比基线。Agentic System 的强项是自动化和迭代效率但依然需要工程框架兜底。7. 最佳实践与工程建议7.1 明确什么场景适合 LLM-Based Kernel 定制并不是所有矩阵计算都需要 SparseDitto 这类系统。模型已经非常规则、稀疏模式固定且已知的场景人工写一个 Kernel 可能更快。LLM-Based Agentic System 的价值在于探索复杂或非固定的稀疏模式或者团队没有专职 CUDA 工程师时用自动化能力替代部分人工开发。建议按下面三个标准判断是否值得引入稀疏模式是否足够多样人工维护成本是否过高。是否反复出现新的剪枝策略或新的稀疏格式。是否有明确的性能反馈闭环可以自动化评估。7.2 安全与可复现性建议使用 LLM 自动生成 GPU Kernel 时要格外关注代码安全性。LLM 生成的代码可能包含未定义行为、内存越界或死循环。在实际生产环境使用前必须做三件事静态审查、沙箱编译、小规模数据验证。可以对 Agent 系统增加一个权限边界生成的 Kernel 只有在测试环境下通过验证后才允许进入正式推理链路。可复现性方面建议记录每次生成实验的完整上下文包括稀疏模式元数据、Agent 配置、LLM 模型版本、生成的 Kernel 源码、编译器版本、GPU 型号、性能数据。这样后续某次优化异常时可以回溯定位也能用于微调 Agent 的决策策略。7.3 与现有推理框架的结合方式SparseDitto 这类系统的产出是一个或一组高性能 GPU Kernel。工程落地时可以把这些 Kernel wrap 成 ONNX Runtime 的自定义 OP或者集成进 TensorRT 插件、PyTorch 自定义 autograd Function。比较好的实践是采用“先离线生成再在线加载”的架构Agent 在开发环境完成内核搜索和编译缓存生产环境直接加载编译后的二进制或 Triton JIT 缓存避免推理延迟被 LLM 调用影响。另外建议在推理框架内增加“性能回归看板”。因为稀疏模式可能会随模型版本变化每次模型更新后都触发一次内核性能对比如果新模型上某些 Kernel 性能下降超过阈值就自动重新触发 Agent 优化流程。8. 总结与下一步学习方向围绕 SparseDitto 做一轮拆解之后可以清楚看到这类系统的本质把 GPU Kernel 定制从一个“依赖资深工程师经验”的任务变成一个“可观测、可迭代、可自动执行”的工程流程。稀疏模式分析负责定义问题LLM Agent 负责生成方案性能评测闭环负责验证方案三者缺一不可。如果接下来想深入这个方向建议按这个路径学习先掌握 Triton 和 CUDA 的基础 Kernel 编写理解 tile、shared memory、warp 等概念。熟悉常见稀疏格式的底层存储结构尤其是 CSR、BSR、bitmask 之间的转换成本。研究传统自动调优工具比如 TVM Ansor、Triton Autotune理解它们的能力边界。再尝试搭建一个最小 Agentic 闭环稀疏模式分析、模板选择、代码生成、编译评测、反馈迭代。最后引入 LLM 作为代码生成与策略决策组件并设计合理的约束和评测机制。这中间最值得投入时间的不是让 Agent 生成更多代码而是把评测反馈设计得足够细致。只有反馈信号足够准确Agent 才能像工程师一样持续改进 Kernel 性能。这也是 SparseDitto 这类系统从“论文概念”走向“工程可用”的关键一步。