通用稀疏矩阵乘法CUDA算子实现

稠密矩阵乘法的优化，核心是在有限的片上存储里尽可能多地复用数据；稀疏矩阵乘法却刚好相反：数据位置是不规则的、每一行工作量不相同，算术强度往往又低得可怜。把一个稠密 GEMM 的 tile 套到 CSR 上，通常只会得到一个更复杂、更慢的 kernel。

本文把稀疏矩阵乘法拆为两个最常见的算子：

y=Ax\quad\text{(SpMV)},\qquad C=AB\quad\text{(SpMM)}.

CSR 格式的单行 SpMV：行偏移确定非零段，warp lane 分担 gather 与 shuffle 归约

图 1 展示 CSR SpMV 的基本数学映射：row_offsets[i:i+1] 给出第 i 行的非零元区间；warp 的各 lane 并行计算部分和，最后只在 warp 内做 shuffle 归约并写出一个 $y_i$ 。这也是短行不需要 shared memory 的原因。

其中 $A\in\mathbb R^{M\times K}$ 为稀疏矩阵， $x\in\mathbb R^K$ ， $B\in\mathbb R^{K\times N}$ 。目标不是制造一个“万能 kernel”，而是搭建一个按稀疏格式、行长度分布、右端项宽度和精度自动选路的算子族。这才是工业框架中通用稀疏算子的正确打开方式。

先说结论：稀疏优化是在和不规则性谈判#

假设 CSR 有 nnz 个非零元。一次 FP32 SpMV 至少要读取 value、column index、向量元素，并写回一个结果。即便不计缓存未命中，浮点运算也只有大约 2 * nnz 次；相对于数据移动，它天生接近带宽受限。

SpMM 的情况更好。每读取一次 $A_{ik}$ ，可以把 $B_{k,:}$ 的一段复用到多个输出列， $N$ 越大，计算与访存的比值越高。因此下面三条是全文的主线：

格式要服从结构。 COO 方便构造，CSR 适合行方向 SpMV/SpMM，BSR/Blocked-ELL 才有机会把结构化块交给 Tensor Core；不能把格式转换的代价藏起来。
线程映射要服从行长度。 一 warp 一行、一个 CTA 一行、一个 CTA 多行、长行分段，适用的行长度区间完全不同。
调度要服从问题本身。 平均 nnz/row 不够，方差、P99 行长、列索引局部性、N、dtype、对齐和 batch 都是 dispatch key。

这也是 OneFlow 优化 Softmax 时最值得借鉴的思想：不是执着于一种实现，而是按形状分段，利用寄存器、shared memory、warp 原语和占用率约束选择真正可启动、可跑满的方案。不同的是，稀疏场景的“形状”还多了稀疏模式本身。

1. 数据格式不是文件格式，而是执行计划#

1.1 CSR：默认主路径#

CSR 使用三段数组：

1
row_offsets: [M + 1]        第 i 行位于 [row_offsets[i], row_offsets[i+1])
2
col_indices: [nnz]
3
values:      [nnz]

它的好处是同一行的 value 与 index 连续，天然适合输出按行写回；缺点是行长可能极不均衡。绝大多数“普通稀疏”输入先进入 CSR 主路径，基准实现也应先从它开始。

1.2 COO：输入友好，执行不应死守 COO#

COO 的 (row, col, value) 三元组适合图边列表、动态 append 和分块拼接，但随机写回 y[row] 往往需要原子操作。对于静态矩阵，预处理成按 (row, col) 排序的 CSR；对于在线构造的 COO，采用 warp-aggregated atomic 或分段归约，并把转换成本单独计入 benchmark。

1.3 BSR / Blocked-ELL：为块结构准备的快速通道#

当非零元以 $b_r\times b_c$ 小块出现时，BSR 把一个索引摊薄到一个稠密小块上。对于固定块、固定宽度模型层，它不仅降低 index 开销，还能让块内计算变得规整。Blocked-ELL 则在每个 block-row 中填充等数量的块，适合块分布较均匀的结构化稀疏。

请注意：通用非结构化 CSR 不能直接获得 Tensor Core 加速。 需要 BSR、2:4 structured sparsity 或其它满足硬件约束的布局；否则所谓“Tensor Core 稀疏加速”通常只是营销话术。cuSPARSE 的通用接口也把 COO、CSR、Blocked-ELL 明确区分为不同格式与算法路径。NVIDIA cuSPARSE 格式文档

1.4 SELL-C- $\sigma$ ：用可控填充换取规则访存#

对行长差异很大但行可重排的矩阵，可在长度相近的窗口 $\sigma$ 内排序，再把每 $C$ 行切成切片。SELL 牺牲少量 padding，换得同一 warp 跨行访问的连续性。它很适合短行图数据或有限元网格；但若行重排会破坏外部语义或 padding 比例太高，就应回退到 CSR。

2. SpMV：把一行的归约做对#

2.1 朴素版本为什么不够#

一个 thread 负责一整行，逻辑最简单：

1
for (int p = row_offsets[row]; p < row_offsets[row + 1]; ++p)
2
  sum += values[p] * x[col_indices[p]];
3
y[row] = sum;

但行长超过几十时，一个 thread 串行读取，延迟无法隐藏；行长很短时，又会浪费一个 block 中的大量线程。优化的第一步不是上 shared memory，而是给一行匹配合适的并行粒度。

2.2 短行：一个 warp 处理一行#

当 nnz_per_row 大致不超过 32 或 64，最自然的映射是一个 warp 一行。lane t 处理 p = begin + t, begin + t + 32, ...，最后用 shuffle 做树形归约：

1
template <typename Value, typename Acc>
2
__device__ __forceinline__ Acc warp_sum(Acc v) {
3
  #pragma unroll
4
  for (int delta = 16; delta > 0; delta >>= 1)
5
    v += __shfl_down_sync(0xffffffff, v, delta);
6
  return v;
7
}
8

9
template <typename Value, typename Acc, int WARPS_PER_CTA>
10
__global__ void csr_spmv_warp(const int* __restrict__ row_offsets,
11
                              const int* __restrict__ col_indices,
12
                              const Value* __restrict__ values,
13
                              const Value* __restrict__ x,
14
                              Value* __restrict__ y, int m) {
15
  const int lane = threadIdx.x & 31;
16
  const int warp = threadIdx.x >> 5;
17
  const int row = (blockIdx.x * WARPS_PER_CTA) + warp;
18
  if (row >= m) return;
19

20
  Acc sum = Acc(0);
21
  for (int p = row_offsets[row] + lane; p < row_offsets[row + 1]; p += 32)
22
    sum += static_cast<Acc>(values[p]) * static_cast<Acc>(x[col_indices[p]]);
23
  sum = warp_sum<Value, Acc>(sum);
24
  if (lane == 0) y[row] = static_cast<Value>(sum);
25
}

这里没有 __syncthreads()，也没有 shared memory。warp 内 shuffle 比 block 归约更轻，且多 warp CTA 可以同时服务多行，是短行 SpMV 的第一选择。

2.3 中长行：一个 CTA 一行，但不要盲目增大 block#

当行长达到几百或几千，一个 warp 不够。可以用 128/256/512 threads 的 CTA 处理一行：每个 warp 局部归约，warp 0 再归约 warp sums。block size 不能写死：它同时影响每个 SM 的常驻 CTA 数、寄存器压力和同步点的空闲时间。

实践中应在初始化阶段对候选 128, 256, 512 查询 occupancy，结合行长区间选择；这和 OneFlow 在 block reduce 中用可启动 CTA 数约束 block size 的思路一致。对于很长的一行，CTA 数太少也会让整个 GPU 吃不饱。

2.4 超长行：分段，再归并#

一条 10 万非零元的行不能只分配一个 CTA。把它切成 SEGMENT_NNZ 的 segment，每个 CTA 输出一个 partial sum，第二个小 kernel 再归约；若只允许一次 launch，则对同一行做 atomicAdd，但需评估原子竞争。典型选择是：

极少数长行：partial buffer + 第二阶段归约，性能更稳定；
多数行都超长：直接按 tile 设计 CTA-per-segment，避免“行”成为唯一调度单位；
结果对严格可复现敏感：不要使用浮点 atomic 作为默认路径，固定二叉归约顺序。

2.5 索引访问的现实#

values[p] 和 col_indices[p] 连续，容易合并；x[col_indices[p]] 却是 gather。可做的优化包括行内列排序、对热点向量使用只读缓存语义、预取下一批 index/value、以及把索引压缩为 32-bit（维度允许时）。但不要把 __ldg 当魔法：随机列、工作集远大于 L2 时，主要瓶颈仍然是 DRAM 延迟。

3. SpMM：右端项宽度改变了一切#

SpMM 计算 $C_{i,n}=\sum_{k\in\mathcal N(i)}A_{i,k}B_{k,n}$ 。同一个 A[i,k] 可用于 N_TILE 个列，因而可以把 $B$ 的连续列放入寄存器累加：

图 2 的关键差别在于右端项不再是一个标量：一次读取 $A_{i,k}$ 后，连续的 $B_{k,n:n+N_{tile}}$ 同时更新多个寄存器累加器，因而将 index 与稀疏 value 的读取开销摊薄到多个输出列。

1
template <typename Value, typename Acc, int N_TILE>
2
__global__ void csr_spmm_row_tile(/* CSR A, dense B, dense C */) {
3
  // 一个 CTA 可处理一行或一组短行；每个 warp 持有若干输出列。
4
  // 对每个 A(i,k)，协作读取连续的 B(k, n:n+N_TILE)，
5
  // 以 Acc 累加到寄存器 fragments，最后一次性写 C。
6
}

真正的实现不能只写这个空壳，至少还需要下面四层策略。

3.1 `N` 很小：沿行并行，避免寄存器爆炸#

N=1 就是 SpMV；N=2/4/8 时每线程持有的 accumulator 很少，应优先沿 nnz 方向协作，把 CTA 的线程给长行。若仍强行让一个 thread 存很多 C 的列，寄存器压力会降低 occupancy，得不偿失。

3.2 `N` 中等：一 warp 一行 × 一列 tile#

当 N=16..128，每个 warp 处理一行和一个连续 N_TILE。B[k, n:n+N_TILE] 是连续的，适合 half2、float4 或 int4 风格的向量化 load；累加 fragment 放寄存器。每个非零元的 index 开销被多个输出列分摊，这是 CSR SpMM 最常见的甜点区。

3.3 `N` 很大：二维 CTA 与 CTA-per-row-group#

N 很大时，单个 CTA 负责一行会产生太多输出列。把 N 分块，并让 CTA 负责 (row_group, n_tile)；短行合并成 row group，长行独占多个 CTA。此处必须保证 C 的写回是唯一所有者，或明确安排 reduction 阶段，避免无意义的原子累加。

3.4 块稀疏：才值得谈 Tensor Core#

对 BSR，外层遍历非零 block，内层是小型稠密 GEMM。此时可以：

将 A block、B tile 异步搬到 shared memory；
以 ldmatrix/fragment 组织矩阵片段；
使用适配架构的 WMMA/MMA/WGMMA 指令；
用双缓冲或多阶段流水重叠搬运和计算。

如果稀疏模式满足硬件支持的 2:4 结构化规则，也可进入专用稀疏 Tensor Core 路径。CUTLASS 已提供 dense、block-sparse、grouped GEMM 等可组合的参考实现；它是研究 tile 布局与流水线的好教材，但上线时仍要在目标 GPU 和实际形状上 profile。CUTLASS 文档

4. 一个真正通用的 dispatch 设计#

把所有情况塞进模板参数会造成编译爆炸；反过来，运行时每一步动态判断也会拖慢热路径。一个平衡的设计是：首次见到矩阵时提取轻量 profile，构造稳定的 dispatch key。

1
key = {
2
  format, dtype, index_width, M, K, N,
3
  avg_nnz, p95_nnz, max_nnz, row_length_cv,
4
  sorted_columns, block_shape, alignment, arch
5
}

建议的规则表如下。

条件	首选实现	退化/备注
CSR，短行，`N=1`	warp-per-row SpMV	多 warp/CTA，shuffle 归约
CSR，长行，`N=1`	CTA-per-row 或 segmented SpMV	极长行走 partial reduction
CSR，`N=8..128`	row × N-tile SpMM	B 连续列寄存器化
CSR，短行占多数且方差大	SELL 或 CTA-per-row-group	先检查 padding 成本
BSR / Blocked-ELL	block sparse SpMM	可选择 Tensor Core 路径
COO，动态输入	COO segmented / warp atomic	后台合并为 CSR
`nnz/(M*K)` 接近稠密	稠密 GEMM	转换后直接走 GEMM 常更快

首次 profile 可以缓存到矩阵 descriptor；动态稀疏矩阵则按版本号失效缓存。对线上模型，再加一个小型 autotune：只在候选的 2~4 个 kernel 中用 CUDA event 测量，缓存最快配置，绝不在每次调用都 benchmark。

5. 从低精度到高精度：存储精度和累加精度必须分开#

A/B 存储	建议累加	适用场景	要点
INT8	INT32 / FP32	量化推理	scale 可按 tensor、channel 或 block；先定义溢出策略
FP8 / INT4	FP32（或按误差预算）	低比特推理、块稀疏	scale 元数据必须随 block 布局一起设计
FP16 / BF16	FP32	训练与常规推理	`half2`/BF16 向量化；不要用 FP16 累加替代 FP32 默认值
TF32 输入语义	FP32	FP32 范围、可放宽 mantissa	仅在架构与数值约束满足时启用
FP32	FP32	基准与高精度主路径	行内规约顺序影响 bitwise 一致性
FP64	FP64	科学计算	通常更受带宽和较低双精度吞吐约束

对 SpMV，低精度不一定自动更快：index 和随机 x gather 仍是固定成本，value 减半却不等于端到端时间减半。对块稀疏 SpMM，低精度更有价值，因为它能减轻 A/B tile 搬运并匹配矩阵乘指令。所有低精度路径都应以更高精度 reference 校验，并报告相对误差、最大误差与是否允许非确定性。

6. 框架级融合与内存管理#

稀疏算子常常位于 SpMM + bias + activation、消息传递聚合、embedding 聚合或 attention 的投影附近。若中间张量会立即被消费，融合 epilogue 可以少一次 C 的全局内存往返。但融合的边界要克制：过多分支会增加寄存器占用，导致主 kernel 降速。

工程上建议把算子拆为：

执行链为：descriptor/profile → format conversion/cache → dispatch → main kernel → optional reduction → optional fused epilogue。

所有临时缓冲区通过框架 allocator 或 stream-ordered memory pool 管理；不要在每次 SpMV/SpMM 调用中 cudaMalloc。同样，不要为了一次调用把 CSR 无条件转换成 BSR——格式转换要么被长期复用，要么本身就是亏损交易。

7. 性能与正确性要一起测#

不要只报“GFlop/s”。稀疏计算里名义 FLOPs 容易掩盖填充、格式转换和原子冲突。推荐至少输出：

1
matrix: M、K、nnz、density、avg/p95/max nnz per row、格式
2
operator: SpMV/SpMM、N、dtype、index width、是否排序、是否转换
3
performance: 端到端时间、kernel 时间、转换摊销时间、有效带宽、有效 GFLOP/s
4
correctness: 相对 FP64/FP32 reference 的 atol/rtol、最大误差、确定性
5
environment: GPU、SM、CUDA、驱动、时钟策略、warmup/重复次数

有效带宽可用“实际读写字节 / kernel 时间”估算，但请明确它是否包含 index、输出、partial buffer 和格式转换。用 Nsight Compute 观察 dram__throughput、warp 执行效率、寄存器数、occupancy、L2 命中和长行 kernel 的负载均衡；这比猜测某个 intrinsic 是否“更快”可靠得多。

8. 常见误区#

“CSR + Tensor Core”：CSR 本身不提供规则的矩阵片段；需要 BSR 或受支持的结构化稀疏。
“行越长 block 越大”：大 block 可能减少常驻 CTA，并增加同步停顿；应以 profile 和 occupancy 选取。
“FP16 一定两倍快”：SpMV 常受 index 和随机 gather 限制，低精度收益可能很小。
“COO atomic 总能接受”：高冲突行会让 atomic 成为串行点；分段归约更可控。
“格式转换免费”：一次性矩阵转换很可能比计算本身还贵。
“只用平均行长做选择”：同样平均值的均匀矩阵与幂律图，最优 kernel 可以完全不同。

总结#

稀疏矩阵乘法没有一个能统治所有输入的 kernel。可维护、高性能的实现应该把问题拆成三件事：用 descriptor 识别结构，用 format 和线程映射承接结构，用 dispatch/autotune 选择实测最快的实现。

对 SpMV，核心是让每一行获得恰当的 warp/CTA/segment 粒度；对 SpMM，核心是借助右端项宽度提高复用，并在块结构存在时才进入 Tensor Core 路径。把格式转换、精度、确定性和 benchmark 口径都放进设计，才是“通用”二字真正的含义。

音乐

音乐

先说结论：稀疏优化是在和不规则性谈判#

1. 数据格式不是文件格式，而是执行计划#

1.1 CSR：默认主路径#

1.2 COO：输入友好，执行不应死守 COO#

1.3 BSR / Blocked-ELL：为块结构准备的快速通道#

1.4 SELL-C- $\sigma$ ：用可控填充换取规则访存#

2. SpMV：把一行的归约做对#

2.1 朴素版本为什么不够#

2.2 短行：一个 warp 处理一行#

2.3 中长行：一个 CTA 一行，但不要盲目增大 block#

2.4 超长行：分段，再归并#

2.5 索引访问的现实#

3. SpMM：右端项宽度改变了一切#

3.1 `N` 很小：沿行并行，避免寄存器爆炸#

3.2 `N` 中等：一 warp 一行 × 一列 tile#

3.3 `N` 很大：二维 CTA 与 CTA-per-row-group#

3.4 块稀疏：才值得谈 Tensor Core#

4. 一个真正通用的 dispatch 设计#

5. 从低精度到高精度：存储精度和累加精度必须分开#

6. 框架级融合与内存管理#

7. 性能与正确性要一起测#

8. 常见误区#

总结#

参考资料#

支持与分享

音乐

文章目录

音乐

音乐

先说结论：稀疏优化是在和不规则性谈判#

1. 数据格式不是文件格式，而是执行计划#

1.1 CSR：默认主路径#

1.2 COO：输入友好，执行不应死守 COO#

1.3 BSR / Blocked-ELL：为块结构准备的快速通道#

1.4 SELL-C-σ\sigmaσ：用可控填充换取规则访存#

2. SpMV：把一行的归约做对#

2.1 朴素版本为什么不够#

2.2 短行：一个 warp 处理一行#

2.3 中长行：一个 CTA 一行，但不要盲目增大 block#

2.4 超长行：分段，再归并#

2.5 索引访问的现实#

3. SpMM：右端项宽度改变了一切#

3.1 N 很小：沿行并行，避免寄存器爆炸#

3.2 N 中等：一 warp 一行 × 一列 tile#

3.3 N 很大：二维 CTA 与 CTA-per-row-group#

3.4 块稀疏：才值得谈 Tensor Core#

4. 一个真正通用的 dispatch 设计#

5. 从低精度到高精度：存储精度和累加精度必须分开#

6. 框架级融合与内存管理#

7. 性能与正确性要一起测#

8. 常见误区#

总结#

参考资料#

支持与分享

音乐

文章目录

1.4 SELL-C- $\sigma$ ：用可控填充换取规则访存#

3.1 `N` 很小：沿行并行，避免寄存器爆炸#

3.2 `N` 中等：一 warp 一行 × 一列 tile#

3.3 `N` 很大：二维 CTA 与 CTA-per-row-group#