橙子派的个人博客

CUDA FlashAttention：从在线 Softmax 到长序列 IO 极限

Wed, 15 Jul 2026 00:00:00 GMT

大模型把上下文从 2K 推到 8K、32K 以后，Attention 的问题不再只是“两个 GEMM 很贵”。标准实现会把 QKTQK^TQKT 写回显存，经过 mask 和 Softmax 后再读回来与 VVV 相乘。计算量仍是二次的，真正先让实现失去扩展性的往往是这个 N×NN\times NN×N 中间矩阵。

本文只讨论精确的 dense causal attention，不用稀疏近似换速度。目标是从一个可核验的 CUDA 数据流出发，解释 FlashAttention 为什么有效、如何分层到 CTA/warp/MMA、在什么 shape 下应回退，以及生产实现和教学骨架之间还差哪些工作。

1. 问题、shape 与性能指标#

设

Q,K,V∈RB×H×N×D,Q,K,V\in\mathbb{R}^{B\times H\times N\times D},Q,K,V∈RB×H×N×D,

大模型中的 scaled dot-product attention 为

O=softmax⁡(QKTD+M)V,O=\operatorname{softmax}\left(\frac{QK^T}{\sqrt D}+M\right)V,O=softmax(DQKT+M)V,

其中 causal mask Mij=0M_{ij}=0Mij=0 当 j≤ij\le ij≤i，否则为 −∞-\infty−∞。本文实测使用 B=1,H=32,D=128B=1,H=32,D=128B=1,H=32,D=128，这对应常见大模型的 head geometry；NNN 从 512 扩展到 8192。

非 causal 前向的两个矩阵乘法约需 4BHN2D4BHN^2D4BHN2D FLOP。causal kernel 若能跳过上三角 tile，实际有效计算约为

Fcausal≈2BHN2D.F_{\mathrm{causal}}\approx 2BHN^2D.Fcausal≈2BHN2D.

性能不能只看 TFLOP/s，还要同时记录：端到端 kernel 延迟、输入之外的峰值显存、长序列是否 OOM，以及 FP16 输出相对 FP32 Softmax 参考路径的误差。

2. 精度语义：低精度乘法不等于低精度 Softmax#

本文的 Q,K,VQ,K,VQ,K,V 和输出为 FP16，点积与输出 accumulator 使用 FP32。在线 Softmax 的行最大值 mmm、指数和 lll 也必须是 FP32：如果 lll 用 FP16，长行会更早丢失小概率项；如果不减行最大值，exp 会溢出。

BF16 有更大的指数范围，但尾数更短，仍应保留 FP32 的 m/l/Om/l/Om/l/O accumulator。TF32 是 FP32 GEMM 的硬件折中，不是本文 FP16 大模型路径的输入格式。FP8/INT8 可以降低 Q/K/V 流量，但需要 per-tensor 或 per-block scale，且 Softmax 状态仍应提升到 FP32。FP64 吞吐和容量代价都不适合作为大模型默认路径，只用于高精度离线校验。

3. 朴素实现为何先撞上显存#

朴素路径是三个独立阶段：

GEMM 写出 S=QKT/DS=QK^T/\sqrt DS=QKT/D；
mask 与 Softmax 读写 SSS；
GEMM 读取 P=softmax⁡(S)P=\operatorname{softmax}(S)P=softmax(S) 和 VVV，写出 OOO。

即使只把 SSS 保存成 FP16，中间矩阵也占

MS=2BHN2 bytes.M_S=2BHN^2\ \text{bytes}.MS=2BHN2 bytes.

在 B=1,H=32,N=8192B=1,H=32,N=8192B=1,H=32,N=8192 时，仅一份 SSS 就是 4096 MiB。安全 Softmax 还可能生成 mask、FP32 临时量或额外输出，因此真实峰值高于这个下界。OneFlow 对 Softmax 的拆解给出了同样的工程信号：exp、reduce 和 broadcast 分成多个 kernel 会反复读写整行，优化应先减少全局访存，而不是先抠一条指令。

从 roofline 看，两个大 GEMM 本身有较高算术强度，但 GEMM 之间的 N2N^2N2 物化流量没有产生新的模型信息。FlashAttention 不减少 dense attention 的渐近计算量，而是利用片上 SRAM 重排计算，消掉这部分 HBM 往返。

4. 在线 Softmax：不保存整行仍保持精确#

把一行 score 分成 tile。处理第 ttt 个 tile 前保存运行状态：最大值 m(t−1)m^{(t-1)}m(t−1)、指数和 l(t−1)l^{(t-1)}l(t−1)、未归一化输出向量 o(t−1)o^{(t-1)}o(t−1)。新 tile 的局部最大值和局部指数和为

m~=max⁡jsj,l~=∑jesj−m~.\tilde m=\max_j s_j,\qquad \tilde l=\sum_j e^{s_j-\tilde m}.m~=jmaxsj,l~=j∑esj−m~.

合并时先更新最大值

m(t)=max⁡(m(t−1),m~),m^{(t)}=\max\left(m^{(t-1)},\tilde m\right),m(t)=max(m(t−1),m~),

再重标定旧状态：

α=em(t−1)−m(t),β=em~−m(t),\alpha=e^{m^{(t-1)}-m^{(t)}},\qquad \beta=e^{\tilde m-m^{(t)}},α=em(t−1)−m(t),β=em~−m(t),l(t)=αl(t−1)+βl~,l^{(t)}=\alpha l^{(t-1)}+\beta\tilde l,l(t)=αl(t−1)+βl~,o(t)=αo(t−1)+∑jesj−m(t)vj.o^{(t)}=\alpha o^{(t-1)}+ \sum_j e^{s_j-m^{(t)}}v_j.o(t)=αo(t−1)+j∑esj−m(t)vj.

最后输出 O=o(T)/l(T)O=o^{(T)}/l^{(T)}O=o(T)/l(T)。这是代数重排，不是近似；只要每个 tile 都进入合并，结果与整行安全 Softmax 相同，差异来自浮点加法顺序。

FlashAttention online softmax tile flow

图 1：Q tile 常驻寄存器，K/V tile 通过双缓冲 shared memory 流过；每轮只保留 FP32 的 m/l/Om/l/Om/l/O 状态，不把 score tile 写回 HBM。

5. 从能跑到高吞吐的 CUDA 数据流#

5.1 CTA tile：让 K/V 被多个 query 复用#

一个 CTA 负责 BrB_rBr 行 query。QiQ_iQi 只加载一次并留在寄存器或 shared memory；K/V 按 BcB_cBc 个 token 分块。单个 K/V tile 从 HBM 读取一次，却服务 BrB_rBr 行 query，避免“一行 query 一个 block”重复读取整个上下文。

shared memory 预算近似为

Msmem≈2BcD s+MQ+MP,M_{\mathrm{smem}}\approx 2B_cD\,s+M_Q+M_P,Msmem≈2BcDs+MQ+MP,

其中 s=2s=2s=2 bytes 对应 FP16。增大 BcB_cBc 提高复用，却会增加 shared memory 和寄存器占用，减少每个 SM 的并发 CTA。不存在对所有 D,ND,ND,N 都最优的 tile。

5.2 warp/MMA：把乘加留给 Tensor Core#

QKTQK^TQKT 和 PVPVPV 应映射到 mma.sync 支持的 fragment，而不是逐元素 FP32 FMA。warp 先协作装载 FP16 fragment，Tensor Core 以 FP32 累加；Softmax 的 max/sum 则使用 warp shuffle 和少量 shared memory 做跨 warp 合并。

FlashAttention-2 的关键不是换一个 Softmax 公式，而是减少非矩阵乘 FLOP、在单 head 内沿 query 方向增加 CTA 并行，并调整 warp 分工以减少 shared-memory 通信。对 B×HB\times HB×H 很小而 NNN 很长的情况，这种并行度尤其重要。

5.3 K 维流水与 bank conflict#

Ampere 的 cp.async 可以把 global-to-shared copy 与当前 tile 的 MMA 重叠：stage 0 计算时预取 stage 1，随后交换 buffer。K/V shared layout 要同时满足 128-bit 合并访存和 MMA fragment 的访问方式；简单的二维行主序可能让多个 lane 命中同一 bank。常见处理是 padding 或 swizzle，而不是盲目增大 shared memory。

5.4 向量化、对齐与尾部#

head dim 为 64/128 且首地址 16-byte 对齐时，使用 uint4 或等价的 128-bit load 搬运 8 个 FP16。尾部不整齐、跨页或变长序列必须走 masked load。向量化减少指令数，但不能越界，也不能假设框架给出的任意 stride 都连续。

6. 核心 CUDA 骨架#

下面只展示在线状态更新和 tiled 主循环的边界。生产 kernel 还需要真正的 MMA fragment、swizzle layout、causal 边界和架构特化；benchmark 不使用这段教学骨架冒充工业实现。

1
struct OnlineState {
2
  float row_max;
3
  float row_sum;
4
};
5

6
__device__ __forceinline__ OnlineState merge_softmax_tile(
7
    OnlineState old_state, float tile_max, float tile_sum) {
8
  const float next_max = fmaxf(old_state.row_max, tile_max);
9
  const float old_scale = __expf(old_state.row_max - next_max);
10
  const float tile_scale = __expf(tile_max - next_max);
11
  return {
12
      next_max,
13
      old_state.row_sum * old_scale + tile_sum * tile_scale,
14
  };
15
}
16

17
template <int Br, int Bc, int D>
18
__global__ void flash_attention_forward(
19
    const half* __restrict__ q,
20
    const half* __restrict__ k,
21
    const half* __restrict__ v,
22
    half* __restrict__ out,
23
    int sequence,
24
    float scale) {
25
  extern __shared__ half smem[];
26
  half* k_stage[2] = {smem, smem + Bc * D};
27
  half* v_stage[2] = {smem + 2 * Bc * D, smem + 3 * Bc * D};
28

29
  // Q fragments and FP32 m/l/O accumulators stay live across the K/V loop.
30
  OnlineState state{ -CUDART_INF_F, 0.0F };
31
  float output_fragment[D / 32] = {};
32

33
  for (int col = 0, stage = 0; col < sequence; col += Bc, stage ^= 1) {
34
    // Production code issues aligned cp.async into the next stage here.
35
    load_kv_tile_async<Bc, D>(k, v, col, k_stage[stage], v_stage[stage]);
36
    cp_async_wait_and_sync();
37

38
    // mma_qk returns FP32 score fragments; masked tiles become -infinity.
39
    auto scores = mma_qk<Br, Bc, D>(q, k_stage[stage], scale);
40
    const float tile_max = warp_block_max(scores);
41
    const float next_max = fmaxf(state.row_max, tile_max);
42
    rescale_output(output_fragment, __expf(state.row_max - next_max));
43
    const float tile_sum = exp_scores_and_mma_pv(
44
        scores, next_max, v_stage[stage], output_fragment);
45
    state = {next_max,
46
             state.row_sum * __expf(state.row_max - next_max) + tile_sum};
47
  }
48
  store_normalized(out, output_fragment, state.row_sum);
49
}

展开收起

代码中最容易写错的是状态顺序：旧输出必须按新最大值重标定后再累加当前 PVPVPV；causal mask 必须在 tile max 之前生效；全 mask 行不能让 −∞−(−∞)-\infty-(-\infty)−∞−(−∞) 产生 NaN。

7. shape、mask 与大模型 dispatch#

工业实现至少应按以下维度分派：

D=64/128D=64/128D=64/128：使用专用 MMA/layout；其他 head dim 走通用 kernel 或库回退。
短 NNN：kernel launch 和 online bookkeeping 占比高，普通 fused SDPA 或 GEMM 可能更合适。
长 NNN、小 B×HB\times HB×H：沿 K 维做 split，多个 CTA 产生局部 (m,l,o)(m,l,o)(m,l,o)，再用同一合并公式归约。
causal：跳过完全位于上三角的 tile；对角 tile 才做逐元素 mask。
GQA/MQA：query head 映射到共享 KV head，避免物理复制 K/V；但多个 query head 的复用要和 CTA 调度一起设计。
dropout/反向：训练必须保存或重算 log-sum-exp，并保证随机数 offset 可复现，不能直接套推理 kernel。

极长序列还要考虑 32-bit 索引溢出、grid 维度、变长 batch 的有效 token 数，以及 sequence parallel/context parallel 的跨卡归约。单卡 FlashAttention 解决的是片上/显存 IO，不自动解决多卡通信。

8. 独占 GPU benchmark#

测试环境为 RTX 3060 Laptop 6 GiB、SM86、driver 581.57、PyTorch 2.11.0+cu130。输入为 causal FP16，B=1,H=32,D=128B=1,H=32,D=128B=1,H=32,D=128。每个 shape 先 warmup，再做 7 组 CUDA Event 测量，每组内重复 5 到 30 次，报告 7 组均值的中位数。

测试前 nvidia-smi 和 pmon 均无 compute process。运行时临时把允许的 power limit 从 100 W 提到 130 W，监控峰值为 100% SM、7000 MHz 显存、1965 MHz 核心、102 W、76°C；结束后恢复 100 W。监控进程只采样硬件计数，不提交 CUDA 工作。

对比项强制为 PyTorch SDPBackend.FLASH_ATTENTION 与 SDPBackend.MATH。因此下表验证的是生产 FlashAttention 数据流，而不是上一节的简化骨架。

NNN	FlashAttention	Math SDPA	加速	Flash 峰值额外显存	Math 峰值额外显存
512	0.149 ms	2.214 ms	14.82x	8.1 MiB	109.0 MiB
1024	0.447 ms	7.878 ms	17.61x	16.1 MiB	364.0 MiB
2048	1.704 ms	30.276 ms	17.77x	32.3 MiB	1312.1 MiB
4096	6.573 ms	155.493 ms	23.66x	64.5 MiB	4960.1 MiB
8192	25.854 ms	OOM	-	129.0 MiB	OOM

在 N≤4096N\le4096N≤4096 的可比区间，最大绝对误差均为 1.953125×10−31.953125\times10^{-3}1.953125×10−3，相对 L2L_2L2 误差从 2.54×10−42.54\times10^{-4}2.54×10−4 到 2.68×10−42.68\times10^{-4}2.68×10−4。这是 FP16 输出和不同归约顺序的预期量级，不应要求 bitwise 一致。

可复现入口如下：

1
from torch.nn.attention import SDPBackend, sdpa_kernel
2
import torch.nn.functional as F
3

4
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
5
    output = F.scaled_dot_product_attention(
6
        query, key, value, is_causal=True, dropout_p=0.0
7
    )

这里的 14.82x 到 23.66x 不能外推到 H100，也不能解释为“算法少算了同样倍数的 FLOP”。Math 路径承担了中间张量物化、mask 和多 kernel 流量；硬件、框架版本、shape 或非 causal 模式改变后比例会变化。

9. 局限与回退路径#

本文只验证 forward causal FP16。没有测试 backward、dropout、变长 packed sequence、RoPE 融合、FP8 scale 或多 GPU context parallel。SM86 也没有 Hopper 的 TMA、WGMMA 和更大的片上资源。

生产中优先调用框架 SDPA、FlashAttention 官方实现或 cuDNN，并让后端根据 dtype、head dim 和 mask 选择 kernel。只有自定义 bias、特殊稀疏模式或框架不覆盖的融合边界，才值得维护手写 CUDA；不支持的 shape 应明确回退，而不是静默给出错误结果。

10. 参考资料#

CUDA KV Cache：分页存储与长上下文 Decode 极致优化

Wed, 15 Jul 2026 00:00:00 GMT

自回归大模型每次只生成一个 token，却必须让新 query 访问此前所有 token 的 key/value。KV Cache 把过去 token 的 K/V 保留下来，避免每步重算整个前缀；代价是缓存容量随 batch、层数和上下文线性增长，而且 decode 每步都要重新流式读取这些历史数据。

当服务端同时处理变长请求时，单纯把缓存预留成连续的 max_seq_len 数组，会把显存浪费、碎片、拷贝和 attention kernel 寻址绑在一起。本文从 CUDA 地址计算出发，实现分页 append 与 paged GQA decode，再讨论长上下文下的 warp 划分、split-KV、量化和调度边界。

1. KV Cache 为什么成为大模型推理瓶颈#

设每层有 HkvH_{kv}Hkv 个 KV head，head dim 为 DDD，batch 为 BBB，上下文长度为 LLL，元素字节数为 sss，层数为 NlN_lNl。K 和 V 两份缓存总容量是

MKV=2BNlLHkvDs.M_{KV}=2BN_lLH_{kv}Ds.MKV=2BNlLHkvDs.

以 Nl=32,Hkv=8,D=128,L=32768N_l=32,H_{kv}=8,D=128,L=32768Nl=32,Hkv=8,D=128,L=32768、FP16、B=1B=1B=1 为例：

MKV=2×1×32×32768×8×128×2=4 GiB.M_{KV}=2\times1\times32\times32768\times8\times128\times2 =4\ \text{GiB}.MKV=2×1×32×32768×8×128×2=4 GiB.

这还没有包含模型权重、激活、allocator 保留页和 attention workspace。若把 8 个 KV head 物理展开到 32 个 query head，容量会再乘 4。GQA 的意义不只是少几个参数，它直接降低 decode 每步必须读取的缓存字节数。

单 token decode 的矩阵很“瘦”：QQQ 只有一行，Tensor Core 很难像 prefill 大 GEMM 一样饱和；随着 LLL 增大，它越来越接近显存带宽和访存延迟受限。因此要同时优化两件事：页面管理不能制造额外复制，kernel 也不能先 gather 成连续大张量再做 attention。

2. 连续缓存、分页缓存与精度语义#

连续缓存通常布局为

[B,Lmax,2,Hkv,D].[B,L_{max},2,H_{kv},D].[B,Lmax,2,Hkv,D].

它的地址计算简单，但每个请求必须预留最大长度；请求结束或增长时可能留下不可复用空洞。Paged KV Cache 把物理显存切成固定 token 数的 block：

cache[P][2][T][Hkv][D],\text{cache}[P][2][T][H_{kv}][D],cache[P][2][T][Hkv][D],

其中 PPP 是物理页数，TTT 是 page size。每个请求只保存逻辑页到物理页的 block_table。逻辑 token ttt 的地址为

pl=⌊tT⌋,o=t mod T,p=block_table[b,pl].p_l=\left\lfloor\frac{t}{T}\right\rfloor,\qquad o=t\bmod T,\qquad p=\mathrm{block\_table}[b,p_l].pl=⌊Tt⌋,o=tmodT,p=block_table[b,pl].

最后一页最多浪费 T−1T-1T−1 个 token，而不是浪费到 LmaxL_{max}Lmax。

Paged KV logical-to-physical mapping

图 1：请求看到连续的逻辑 token，物理页可以离散。append 和 decode 直接通过 page table 寻址，不生成连续 gather buffer；共享前缀可以让不同请求的逻辑页指向同一只读物理页。

本文缓存和 query 为 FP16，点积、在线 Softmax 和输出 accumulator 为 FP32。BF16 可扩大数值范围但流量不变。FP8/INT8 cache 可以近似减半流量和容量，代价是每页或每组 scale、dequant 指令和精度校准；不能只改存储类型却继续按 FP16 解释地址。FP64 不属于在线大模型 decode 的合理默认路径。

3. 朴素路径的三个隐藏成本#

3.1 按最大长度预留#

如果 100 个请求都按 32K 预留，而多数只生成几百 token，逻辑空闲空间仍占物理显存。连续 batching 能提高计算吞吐，却会放大不同生命周期带来的外部碎片。

3.2 先 gather 再 attention#

一种容易实现的 paged attention 是：根据 page table 把 K/V gather 成连续张量，再调用 GEMM/SDPA。这样每个 decode step 都额外读写整个历史缓存，流量至少又增加

2×MKV,active layer,2\times M_{KV,\mathrm{active\ layer}},2×MKV,active layer,

并产生大临时 buffer。分页的正确边界是把地址翻译融合进 attention kernel。

3.3 把 GQA 展开成 MHA#

若 Hq/Hkv=gH_q/H_{kv}=gHq/Hkv=g，repeat_interleave 会把 K/V 物理复制 ggg 份。它让普通 MHA kernel 可以运行，却消除了 GQA 的容量与带宽收益。正确映射应是

hkv=⌊hqg⌋,h_{kv}=\left\lfloor\frac{h_q}{g}\right\rfloor,hkv=⌊ghq⌋,

多个 query head 在读取时共享同一个 KV head，而不是事先复制。

4. 向量化 paged append#

每个 decode step 先把当前层的新 K/V 写进最后一页。标量 kernel 每个线程写一个 half，还可能让 K 和 V 分成两个 launch。对 D=128D=128D=128 且地址 16-byte 对齐的常见 shape，可以把 8 个 FP16 打成一个 uint4，一个 block 同时完成地址翻译和连续写入。

1
template <int HeadDim, int PageSize>
2
__global__ void append_paged_kv(
3
    const half* __restrict__ new_kv,       // [B, 2, Hkv, D]
4
    half* __restrict__ cache,              // [P, 2, T, Hkv, D]
5
    const int* __restrict__ block_table,
6
    const int* __restrict__ sequence_len,
7
    int max_pages,
8
    int num_kv_heads) {
9
  const int batch = blockIdx.x;
10
  const int kind_head = blockIdx.y;
11
  const int kind = kind_head / num_kv_heads;
12
  const int head = kind_head % num_kv_heads;
13
  const int token = sequence_len[batch] - 1;
14
  const int logical_page = token / PageSize;
15
  const int page_offset = token % PageSize;
16
  const int physical_page = block_table[batch * max_pages + logical_page];
17

18
  constexpr int kHalfPerVector = sizeof(uint4) / sizeof(half);
19
  const int vector_index = threadIdx.x;
20
  if (vector_index < HeadDim / kHalfPerVector) {
21
    const size_t source =
22
        (((batch * 2 + kind) * num_kv_heads + head) * HeadDim);
23
    const size_t target =
24
        ((((static_cast<size_t>(physical_page) * 2 + kind) * PageSize
25
            + page_offset) * num_kv_heads + head) * HeadDim);
26
    reinterpret_cast<uint4*>(cache + target)[vector_index] =
27
        reinterpret_cast<const uint4*>(new_kv + source)[vector_index];
28
  }
29
}

展开收起

这个 kernel 适用于 HeadDim % 8 == 0、输入和 cache 都满足 16-byte 对齐的路径。否则应 dispatch 到 half2 或 masked scalar 尾部，而不是越界。batch 很小时，append 只有几十 KiB，瓶颈通常是 launch latency，不应拿有效带宽低就断言 global memory 很慢；融合 K/V 和多请求才是主要收益。

5. Paged GQA decode 的 warp 数据流#

一个直接的 CUDA 映射是每个 warp 负责一个 (request, query_head)。lane 持有 D/32D/32D/32 个 query 和输出 fragment，沿 token 循环：

通过 block_table 得到 K/V page；
合并读取对应 KV head 的 K，warp reduce 得到一个 score；
用在线 Softmax 更新 m/lm/lm/l；
每个 lane 读取自己的 V fragment 并更新 FP32 输出。

1
template <int HeadDim, int PageSize>
2
__global__ void paged_decode(
3
    const half* __restrict__ q,
4
    const half* __restrict__ cache,
5
    const int* __restrict__ page_table,
6
    const int* __restrict__ lengths,
7
    half* __restrict__ out,
8
    int query_heads,
9
    int kv_heads,
10
    int max_pages) {
11
  const int lane = threadIdx.x & 31;
12
  const int warp = threadIdx.x >> 5;
13
  const int query_head = blockIdx.y * (blockDim.x / 32) + warp;
14
  const int batch = blockIdx.x;
15
  if (query_head >= query_heads) return;
16

17
  const int group_size = query_heads / kv_heads;
18
  const int kv_head = query_head / group_size;
19
  float q_fragment[HeadDim / 32];
20
  float o_fragment[HeadDim / 32] = {};
21
  load_query_fragment(q, batch, query_head, lane, q_fragment);
22

23
  float running_max = -CUDART_INF_F;
24
  float running_sum = 0.0F;
25
  for (int token = 0; token < lengths[batch]; ++token) {
26
    const int logical_page = token / PageSize;
27
    const int offset = token % PageSize;
28
    const int page = page_table[batch * max_pages + logical_page];
29
    const half* key = cache_ptr(cache, page, 0, offset, kv_head);
30
    const half* value = cache_ptr(cache, page, 1, offset, kv_head);
31

32
    float score = dot_fragment<HeadDim>(q_fragment, key, lane);
33
    score = warp_sum(score) * rsqrtf(static_cast<float>(HeadDim));
34
    const float next_max = fmaxf(running_max, score);
35
    const float old_scale = __expf(running_max - next_max);
36
    const float weight = __expf(score - next_max);
37
    #pragma unroll
38
    for (int i = 0; i < HeadDim / 32; ++i) {
39
      o_fragment[i] = o_fragment[i] * old_scale
40
                    + weight * __half2float(value[lane + 32 * i]);
41
    }
42
    running_sum = running_sum * old_scale + weight;
43
    running_max = next_max;
44
  }
45
  store_output(out, batch, query_head, lane, o_fragment, running_sum);
46
}

展开收起

这段代码表达了正确的页寻址和在线归一化，但生产实现还会让多个 warp/CTA 协作、使用 packed load、减少整数除法、缓存 page id，并对长上下文做 split-KV。

Paged GQA decode pipeline

图 2：32 个 query head 映射到 8 个 KV head；warp 沿逻辑 token 前进，每跨 page size 才更新物理页，FP32 的 m/l/Om/l/Om/l/O 状态跨页保留。

6. 从基础实现到极致优化#

6.1 page id 不要每个元素都查#

page table 查询应位于 token tile 或 page 边界，而不是每个 DDD 元素。一个 warp 广播 page id，随后 lane 连续读取 head dim。PageSize=16 时，单次表查询可以服务 16 个 token。

6.2 packed load 与布局#

对 NHD 布局 [page, token, head, dim]，同一 head 的 DDD 连续，适合 half2/128-bit load。HND 布局 [page, head, token, dim] 则让同一 head 跨 token 更连续，具体选择取决于 CTA 是并行 token 还是并行 head。不要只凭名字判断，应该从 warp 实际地址验证 transaction。

6.3 GQA 共享与 CTA 组织#

同组 query head 读取相同 K/V。若每个 warp 独立读取，L2 可能复用但 HBM 流量不稳定；更激进的实现让一个 CTA 同时处理一组 query head，把 K/V tile 放进 shared memory，再由多个 warp 消费。代价是 shared memory、寄存器和同步增加，小 batch 或短上下文可能不划算。

6.4 长上下文 split-KV#

当 batch 很小、LLL 很长时，batch * query_heads 仍可能不足以填满所有 SM。把 K 维分成 SSS 段，每段 CTA 产生 (ms,ls,os)(m_s,l_s,o_s)(ms,ls,os)，再按在线 Softmax 的合并律归约：

m=max⁡sms,m=\max_s m_s,m=smaxms,l=∑slsems−m,o=∑sosems−m,O=o/l.l=\sum_s l_s e^{m_s-m},\qquad o=\sum_s o_s e^{m_s-m},\qquad O=o/l.l=s∑lsems−m,o=s∑osems−m,O=o/l.

split 增加临时状态和第二个归约 kernel，只有并行度收益大于额外流量时才启用。应按 B,Hq,L,DB,H_q,L,DB,Hq,L,D 和 SM 数量 dispatch，而不是固定切段。

6.5 page size、碎片与 TLB 式权衡#

小 page 降低最后一页内部碎片，却增加 page table、allocator 元数据和地址切换；大 page 相反。16 token 是常见折中，不是定理。超长上下文、prefix cache 和 beam search 还要考虑只读页共享与 copy-on-write：共享前缀不能被后续 append 原地覆盖。

6.6 量化、RoPE 与 CUDA Graph#

FP8/INT8 KV cache 应让 dequant 与 attention 融合，scale 最好按 page/head 或更细粒度存储；粒度越细精度越好，scale 流量和指令越多。RoPE 通常在写 K 前完成，或在读取时按位置融合，二者影响 cache 是否可跨位置复用。CUDA Graph 要求 workspace、page table buffer 和 kernel dispatch 稳定，不能每步 cudaMalloc 或改变 host 控制流。

7. 大规模 serving 的分层边界#

CUDA kernel 只负责给定 page table 的高效读写。真正的大模型系统还需要上层 runtime：free-list/block allocator、continuous batching、请求抢占、prefix hash、跨请求共享、CPU/NVMe offload 和张量并行。把这些策略塞进一个 kernel 会失去可测试性。

合理分层是：scheduler 决定 active requests，block manager 维护逻辑页，append kernel 写当前 token，paged attention 只消费不可变的本轮 page table。OneFlow 的系统设计同样强调 Kernel 只向 CUDA stream 提交异步工作，资源生命周期和调度由上层边界管理。

8. 独占 GPU benchmark#

测试环境为 RTX 3060 Laptop 6 GiB、SM86、driver 581.57、PyTorch 2.11.0+cu130、FlashInfer 0.6.12。shape 固定为 B=8,Hq=32,Hkv=8,D=128B=8,H_q=32,H_{kv}=8,D=128B=8,Hq=32,Hkv=8,D=128、FP16、page size 16，context 从 1K 到 32K。每项 warmup 后做 7 组 CUDA Event 测量，报告组均值的中位数。

测试前后均无其他 compute PID；运行监控达到 100% SM、7000 MHz 显存、1965 MHz 核心、102 W、76°C，临时 130 W power limit 在结束后恢复到 100 W。

8.1 Paged GQA decode#

分页路径使用 FlashInfer BatchDecodeWithPagedKVCacheWrapper。对比路径把 8 个 KV head 预先展开成 32 个连续 head，再强制调用 PyTorch Flash SDPA。这个对比代表两种可运行的大模型数据路径，但 6 到 7 倍差距主要同时包含 GQA 不复制和专用 decode 调度收益，不能单独归因于 paging。

Context	FlashInfer paged GQA	Expanded contiguous Flash SDPA	延迟比	原生 KV Cache	展开后 K/V
1024	0.131 ms	0.838 ms	6.42x	32 MiB	128 MiB
4096	0.457 ms	3.283 ms	7.19x	128 MiB	512 MiB
8192	0.953 ms	6.683 ms	7.01x	256 MiB	1024 MiB
16384	1.785 ms	13.275 ms	7.44x	512 MiB	2048 MiB
32768	3.623 ms	跳过	-	1024 MiB	需要 4096 MiB

32K 对比路径在缓存之外只剩约 3914 MiB 可用，却需要 4096 MiB 展开 K/V，因此按预先设定的安全阈值跳过，没有把系统推入不受控 OOM。分页 GQA 路径完成了 1 GiB 单层缓存的 decode。

在可比 shape 上，最大绝对误差从 1.22×10−41.22\times10^{-4}1.22×10−4 降到 3.05×10−53.05\times10^{-5}3.05×10−5，相对 L2L_2L2 误差约 3.86×10−43.86\times10^{-4}3.86×10−4 到 4.00×10−44.00\times10^{-4}4.00×10−4，符合 FP16 输出和不同归约顺序。

8.2 Fused paged append#

append 对比 FlashInfer fused kernel 与两次 PyTorch advanced-index store。每次写入 8 个请求的一枚 K/V token，时间主要受 launch 和索引开销控制。

Context	Fused append	两次 indexed store	加速
1024	0.0344 ms	0.1266 ms	3.67x
4096	0.0312 ms	0.1159 ms	3.71x
8192	0.0340 ms	0.1138 ms	3.35x
16384	0.0278 ms	0.1102 ms	3.96x
32768	0.0271 ms	0.1043 ms	3.84x

append 延迟几乎不随 context 线性增长，因为每步只写新 token；context 影响的是目标页地址和后续 decode 读取量。这里有效带宽不到显存峰值并不矛盾：写入量只有 32 KiB，kernel launch 和索引固定成本占主导。

9. 局限与回退#

本文只测单 GPU、单层 cache 数据和无 RoPE/ALiBi 的 FP16 GQA decode，没有测完整模型 tokens/s。32 层 32K cache 的容量可以按公式外推，但不能把单层 kernel 延迟简单乘层数后当作端到端延迟，因为层间还有 GEMM、通信和调度。

page table 本身不能保证更快；请求长度接近且显存充足时，连续 cache 地址更简单。极小 batch、短 context 可回退到 contiguous decode；长 context 或并行度不足才启用 split-KV。量化 cache 必须以任务精度验证为前提，OOM 时的正确回退是降低 batch/context、启用 GQA/量化/offload，而不是越界复用物理页。

10. 参考资料#

CUDA LSMR 专栏（三）：递推融合与 8kv/11kv 实测

Thu, 18 Jun 2026 00:00:00 GMT

前两篇已经把 LSMR 的两个大头固定下来：Solver 不知道矩阵格式，SparseMatrix 则同时持有 CSR(A) 与 CSR(ATA^TAT)。最后一篇处理 SpMV 之间的小向量递推，并给出两组真实重建输入的可复现结果。

这一步最容易犯的错误，是看到“自写 CUDA”就只比较时间。迭代求解器首先要对齐 stop code、迭代数和残差；只有数值轨迹成立，性能数字才有意义。

1. 剩余瓶颈与性能指标#

每轮两次 SpMV 之间，LSMR 还要完成归一化、搜索方向和解更新。以长度 nnn 的两条递推为例：

hˉ←h−chhˉ,\bar h\leftarrow h-c_h\bar h,hˉ←h−chhˉ,x←x+cxhˉ.x\leftarrow x+c_x\bar h.x←x+cxhˉ.

分成两个 elementwise kernel 时，第一步读 h/hbar、写 hbar，第二步读 x/hbar、写 x，合计约 6 个长度 nnn 的 FP32 流量，即 24n24n24n 字节。融合后，每个元素读 h/hbar/x、写 hbar/x，约为 20n20n20n 字节，同时少一次 kernel launch。

节省的 4 字节乘 nnn 不是整个 LSMR 的主导项；真正价值是把确定存在的重复访问和 launch 去掉，同时保持公式边界清楚。评价仍以 solve wall time 为主，并记录迭代数、stop code、解相对差和独立计算的相对残差。

2. 递推与 FP32 误差语义#

融合必须尊重依赖：xkx_kxk 使用的是刚更新的 hˉk\bar h_khˉk，不能把旧值读进另一个无序 kernel。逐元素形式为

hˉk,i=hk,i−chhˉk−1,i,\bar h_{k,i}=h_{k,i}-c_h\bar h_{k-1,i},hˉk,i=hk,i−chhˉk−1,i,xk,i=xk−1,i+cxhˉk,i.x_{k,i}=x_{k-1,i}+c_x\bar h_{k,i}.xk,i=xk−1,i+cxhˉk,i.

两个式子对索引 iii 独立，因此可由同一个线程顺序完成。hhh 的下一步更新

hk+1,i=vk+1,i−cvhk,ih_{k+1,i}=v_{k+1,i}-c_v h_{k,i}hk+1,i=vk+1,i−cvhk,i

与前两式没有需要跨线程同步的关系，但它读写另一组状态，保留为第二个 kernel 更简单。

Fused LSMR vector recurrence

图 1：融合前后不是改变数学，而是让同一线程在寄存器中转交新的 hbar。红色路径表示被消除的一次 global-memory 往返。

FP32 下不能要求不同归约树逐 bit 相同。这里采用三层正确性标准：stop code 与迭代数应一致；解相对差应落在 FP32 迭代可解释范围；最终使用同一个 CuPy sparse matrix 重新计算 ∥Ax−b∥2/∥b∥2\|Ax-b\|_2/\|b\|_2∥Ax−b∥2/∥b∥2，两者应匹配。

3. 未融合实现的成本#

CuPy 版本把递推写成清晰的数组表达式，适合作为算法参考，但每个表达式由 Python 驱动：

1
hbar *= -coefficient
2
hbar += h
3
x += step * hbar
4
h *= -theta_over_rho
5
h += v

memory pool 可以复用分配，却不会自动把这些语句融合成一个 CUDA kernel；矩阵乘法返回临时向量后，后续加法也会增加一次遍历。对于 8kv/11kv 这类每轮 SpMV 已很快、迭代约百次的输入，host 调度和小 kernel 数量不是零成本。

纯 CUDA 库有两个可控点：矩阵递推通过 cuSPARSE 的 beta 参数原位完成；依赖明确的小向量递推手写融合。范数仍交给 cuBLAS，不重复实现工业归约。

4. 融合数据流#

从一轮的数据路径看，融合后的顺序是：

u←CSR(A)v→∥u∥2→v←CSR(AT)u→∥v∥2→(hˉ,x,h).u\xleftarrow{\mathrm{CSR}(A)}v \rightarrow \|u\|_2 \rightarrow v\xleftarrow{\mathrm{CSR}(A^T)}u \rightarrow \|v\|_2 \rightarrow (\bar h,x,h).uCSR(A)v→∥u∥2→vCSR(AT)u→∥v∥2→(hˉ,x,h).

两个 norm 会把 FP32 标量送回 host，因为停止条件需要分支；其余状态保持在设备端。LSMR 的每一轮天然依赖上一轮，不能简单把多轮并发。优化重点应放在一轮内部的数据复用，而不是制造错误的跨轮流水。

对右端项也采用一次性适配：dense b 若在 host 则复制到复用 buffer；sparse b 先 memset，再按唯一索引 scatter 成 dense u。LSMR 后续所有递推统一处理 dense 向量，不在热循环里保留 sparse-b 分支。

5. Kernel 实现与适用条件#

核心 kernel 很短：

1
__global__ void UpdateHbarAndXKernel(
2
    float* hbar, const float* h, float* x,
3
    float coefficient, float step, std::int64_t size) {
4
  const std::int64_t i =
5
      static_cast<std::int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
6
  if (i < size) {
7
    const float next_hbar = h[i] - coefficient * hbar[i];
8
    hbar[i] = next_hbar;
9
    x[i] = fmaf(step, next_hbar, x[i]);
10
  }
11
}

每个线程只处理一个连续元素，256 threads/block。这里没有 shared memory、warp shuffle 或复杂 occupancy tuning，因为算子是纯带宽型逐元素更新，寄存器状态只有少量 float；引入 shared memory 只会增加搬运。

融合适用于两式存在逐元素生产者-消费者关系、shape 相同、没有跨元素归约的场景。代价是降低单个基础算子的复用性，所以它放在 LSMR 私有 kernel 层，不暴露成通用 public API。

OneFlow Softmax 的优化也从“naive 多次 global memory”出发，用融合和寄存器缓存减少数据往返；共同原则是先列出实际访存，再决定融合，而不是看到多个 kernel 就机械合并。

6. shape、对齐与 dispatch#

本 kernel 对任意 nnn 使用尾部判断，不要求 16-byte 对齐。float4 能减少指令数，但 hbar/h/x 三个输入输出指针需要共同对齐，尾部还要单独处理；在当前 n=6177n=6177n=6177 和 233942339423394 上，SpMV 占比更大，先保留简单标量路径。

矩阵方向的 dispatch 则更重要：sparse A 走双 CSR SpMV，dense A 走 GEMV；sparse b 只在入口 scatter；多 RHS 当前逐列求解。没有把 SpMV 与 SpMM 混为一谈，也没有声称一个融合 kernel 适合所有 shape。

如果未来 profile 显示向量 kernel 成为主导，可以按对齐和 nnn 分派 float4 路径；如果多个 RHS 改为 block LSMR，则应重新设计 B/C tile 复用，而不是继续逐列调用。

7. 可调用接口与 benchmark 脚本#

C++ 调用者提供 device 输出，Context 复用所有工作区：

1
cuda_lsmr::Context context;
2
auto A = cuda_lsmr::SparseMatrix::FromCsc(context, csc_view);
3
cuda_lsmr::Solver solver(context);
4

5
cuda_lsmr::SolverOptions options;
6
options.atol = 1.0e-6F;
7
options.btol = 1.0e-6F;
8
options.max_iterations = 5000;
9

10
auto result = solver.solve(*A, dense_b, device_x, options);
11
context.synchronize();

用于对比的 pybind11 层不参与算法，只接收 CuPy 的 b.data.ptr 与 x.data.ptr。矩阵加载、SciPy 转 CSC、H2D、双格式构造和两轮 warmup 都在计时外。CuPy 显式传入 FP32 零 x0，避免其兼容路径把迭代解向量提升精度。

8. 8kv/11kv 结果、容差与局限#

测试环境为 RTX 3060 Laptop 6 GiB、CUDA 12.9.86、driver 581.57、CuPy 14.0.1。参数固定为 damp=0、atol=btol=1e-6、conlim=1e8、maxiter=5000。每组执行 7 次同步 solve，报告中位数。

数据集	shape	nnz	CUDA LSMR	CuPy	时间比	迭代数
8kv	306136×6177306136\times6177306136×6177	5,588,894	60.24 ms	127.28 ms	0.473	86 / 86
11kv	936063×23394936063\times23394936063×23394	19,747,202	167.11 ms	229.57 ms	0.728	96 / 96

FP32 benchmark comparison

图 2：同一 GPU、同一输入和停止参数下的 7 次中位数。数字只代表这两组矩阵，不用于宣称普遍快于 CuPy/cuSPARSE。

两边 stop code 均为 2，即近似满足最小二乘停止条件。解相对差为 1.84×10−71.84\times10^{-7}1.84×10−7 和 6.26×10−76.26\times10^{-7}6.26×10−7。独立重算的相对残差分别为 0.3540266 和 0.2360898，两种实现一致。残差不接近零并非失败：输入是非一致最小二乘系统，应该比较停止语义与残差是否对齐。

性能达到并超过“与 CuPy 基本持平”的目标，但边界必须保留：双格式占用更多 VRAM；当前仅 FP32/int32；无预条件器；消费级 GPU 上的结果不能外推到 A100/H100；不同 CSR 行长分布可能改变 cuSPARSE 占比。若显存不足，回退单格式；若要求逐 bit 确定性，应选择对应 cuSPARSE algorithm 并接受吞吐代价。

9. 参考资料#

CUDA LSMR 专栏（二）：用 CSR/CSC 双存储压榨双向 SpMV

Mon, 01 Jun 2026 00:00:00 GMT

上一篇把 LSMR 拆成 Solver 与 LinearOperator。这一篇只盯住 Operator 中最贵的两步：AvAvAv 和 ATuA^TuATu。核心结论很直接：既然每轮都要沿两个方向扫完整个稀疏矩阵，就分别为两个方向准备连续的“行”，而不是让一个格式勉强承担所有访问。

1. 为什么双向 SpMV 是瓶颈#

对 A∈Rm×nA\in\mathbb{R}^{m\times n}A∈Rm×n，LSMR 每轮各执行一次

ym=Avn,zn=ATum.y_m=Av_n,\qquad z_n=A^Tu_m.ym=Avn,zn=ATum.

设每个非零值和索引各 4 字节。忽略向量缓存和 offset，一次方向的矩阵读取至少为

BSpMV≥(4+4)nnz⁡=8nnz⁡ bytes.B_{\mathrm{SpMV}}\ge (4+4)\operatorname{nnz}=8\operatorname{nnz}\ \text{bytes}.BSpMV≥(4+4)nnz=8nnz bytes.

对应计算量约为 2nnz⁡2\operatorname{nnz}2nnz FLOP，理想算术强度不超过

I≤2nnz⁡8nnz⁡=0.25 FLOP/byte,I\le \frac{2\operatorname{nnz}}{8\operatorname{nnz}}=0.25\ \text{FLOP/byte},I≤8nnz2nnz=0.25 FLOP/byte,

实际还会更低。于是 roofline 上界首先由显存带宽决定。优化目标不是让每个线程多做几个 FMA，而是让 warp 读取连续的 values/indices、减少无效流量，并避免转置路径产生更差的归约或 scatter 行为。

2. CSR、CSC 与精度语义#

CSR(A) 由三个数组组成：

row_ptr 长度为 m+1m+1m+1，第 iii 行区间是 [row_ptr[i], row_ptr[i+1])；
col_ind[p] 给出第 ppp 个非零元素的列；
values[p] 给出其 FP32 数值。

因此

(Av)i=∑p=row_ptrirow_ptri+1−1valuesp vcol_indp.(Av)_i=\sum_{p=\mathrm{row\_ptr}_i}^{\mathrm{row\_ptr}_{i+1}-1} \mathrm{values}_p\,v_{\mathrm{col\_ind}_p}.(Av)i=p=row_ptri∑row_ptri+1−1valuespvcol_indp.

CSC(A) 则用 col_ptr 压缩列，row_ind 保存行号。关键观察是：CSC(A) 的内存结构恰好就是 CSR(ATA^TAT)。如果把 descriptor 的尺寸设为 n×mn\times mn×m，就能直接计算

z=ATuz=A^Tuz=ATu

而不向 cuSPARSE 请求 transpose 操作。

矩阵值、输入、输出和 SpMV compute type 都固定为 FP32。这里没有 half2 或 Tensor Core 路径，因为不规则 SpMV 的主要矛盾是索引与带宽，而且测试输入已经是 FP32；盲目把值降到 FP16 会改变最小二乘系统的条件与停止轨迹。

3. 单格式实现的瓶颈#

如果只保留 CSC(A)，ATuA^TuATu 很自然，但 AvAvAv 的行归约不再连续；只保留 CSR(A) 则相反。cuSPARSE 支持格式和 transpose 枚举，但官方文档明确指出 transpose 操作没有结果复现保证，而且具体算法的数据路径由库内部决定。

更重要的是，LSMR 不是偶尔做一次转置乘法，而是每轮固定交替两个方向。一次性多存一份矩阵，与每轮承担方向不匹配的代价，属于典型的“空间换稳定热路径”。

朴素的另一项问题是运行时转换：若在每次 A.T @ u 前生成转置格式，转换成本会直接乘以迭代数。正确边界是在 SparseMatrix 构造时完成一次转换，并让两个 descriptor 随矩阵生命周期存在。

4. 双格式数据流#

Dual sparse storage paths

图 1：输入无论是 CSC 还是 CSR，构造阶段都补齐另一种格式。求解阶段的两个方向均使用 NON_TRANSPOSE CSR SpMV。

CSR 到 CSC 使用 cusparseCsr2cscEx2。CSC 输入可以直接解释成 CSR(ATA^TAT)，再把 ATA^TAT 从 CSR 转成 CSC；所得 CSC(ATA^TAT) 的数组，正好就是 CSR(A)。这个变换只重排一次 values/indices。

双存储的字节代价可近似写成

Mdual≈2[8nnz⁡+4(m+n+2)].M_{\mathrm{dual}}\approx 2\left[8\operatorname{nnz}+4(m+n+2)\right].Mdual≈2[8nnz+4(m+n+2)].

11kv 的 nnz⁡=19,747,202\operatorname{nnz}=19{,}747{,}202nnz=19,747,202，仅双份 values/indices 已约为 301 MiB；还要为 offset、cuSPARSE workspace、LSMR 向量和内存池物理页留余量。这是明确的成本，不应只谈速度不谈容量。

5. 从基础到工业实现#

5.1 构造时转换，迭代时只换指针#

1
cusparseCsr2cscEx2(
2
    handle, rows, cols, nnz,
3
    csr_values, csr_offsets, csr_indices,
4
    csc_values, csc_offsets, csc_indices,
5
    CUDA_R_32F, CUSPARSE_ACTION_NUMERIC,
6
    CUSPARSE_INDEX_BASE_ZERO, CUSPARSE_CSR2CSC_ALG1,
7
    conversion_workspace);

构造完成后创建两个 sparse descriptor、四个 dense-vector descriptor，并查询两个方向所需 workspace 的最大值。热路径只调用 cusparseDnVecSetValues 替换指针，然后发射 SpMV。

5.2 用 beta 项吃掉 Golub-Kahan 旧向量#

前向递推的分子是

u←Av−αu.u\leftarrow Av-\alpha u.u←Av−αu.

直接设 alpha=1、beta=-alpha_k，让 cuSPARSE 在输出 u 上完成它，避免生成 Av 临时向量和一次额外 elementwise kernel。转置方向同理：

v←ATu−βv.v\leftarrow A^Tu-\beta v.v←ATu−βv.

5.3 行长分布决定负载不均#

CSR 的“连续”不等于“均衡”。若某些行只有几个非零，而另一些行有上万个非零，固定的一行一 warp 会让不同 warp 的工作量严重失衡；短行还会浪费 lane，长行则需要循环和跨 warp 归约。

本项目把通用不规则矩阵交给 CUSPARSE_SPMV_CSR_ALG1，没有虚构一个自写 kernel 普遍胜过 cuSPARSE。未来可在 Operator dispatch 中统计

ℓi=row_ptri+1−row_ptri\ell_i=\mathrm{row\_ptr}_{i+1}-\mathrm{row\_ptr}_iℓi=row_ptri+1−row_ptri

的均值、方差和长尾，再为窄分布短行选择专用 warp kernel。机制、适用条件和维护代价必须一起评估。

6. SpMV、SpMM 与 shape dispatch#

这里必须区分 SpMV 和 SpMM。单个 LSMR 右端项的 v,uv,uv,u 都是向量，核心操作是 SpMV。solve_many 虽然接受 dense RHS 矩阵，但每一列有自己的 αk,βk\alpha_k,\beta_kαk,βk、旋转和停止时刻，当前实现是逐列运行独立 LSMR，并不把它伪装成 SpMM。

真正的 block LSMR 才会让 B/C tile 在多个右端项之间复用：CSR(A) 的 values/indices 读一次，与 V∈Rn×rV\in\mathbb{R}^{n\times r}V∈Rn×r 的一小块相乘，形成 U∈Rm×rU\in\mathbb{R}^{m\times r}U∈Rm×r。当 rrr 足够大时，B tile 的连续列和输出 C tile 提高算术强度；代价是块正交化、更复杂的收敛判断和更大的寄存器/显存工作区。

当前 dispatch 因此保持诚实：r=1 走 SpMV；多列输入复用内存但逐列求解；未来只有实现块算法后才启用 SpMM。

7. 显存池与代码边界#

矩阵双存储解决稳定带宽，显存池解决稳定延迟。Context 创建私有 cudaMemPool_t，所有内部 buffer 都走显式 pool：

1
cudaMallocFromPoolAsync(
2
    reinterpret_cast<void**>(&data), bytes, pool, stream);
3

4
// Destruction remains ordered behind prior work on the same stream.
5
cudaFreeAsync(data, stream);

Stream-ordered memory pool timeline

图 2：第一次 shape 扩容向 pool 请求物理页，后续 solve 复用容量；free 也按 stream 排序，不引入 cudaFree 的全局同步语义。

这与 CuPy 默认 memory pool 的目标相同，但库使用私有 pool，避免修改进程全局默认池。DeviceBuffer::Resize 仅在新 shape 超过 capacity 时重新申请；迭代循环内没有分配。

8. 验证、边界与回退#

双路径测试使用同一个小矩阵，分别从 CSR 和 CSC 构造，比较两个方向与 dense GEMV 的结果。Compute Sanitizer 对 Solver 测试报告 0 个访问错误。真实数据则记录 pool reserved bytes：8kv 约 192 MiB，11kv 约 640 MiB；这是 pool 的物理页保留量，不等于矩阵数组的精确逻辑字节数。

双存储适用于矩阵复用多个迭代、显存充足、两个方向频率相近的场景。如果矩阵只乘一次、转置极少使用，或双份数据会逼近 OOM，应回退为单格式。若索引或 nnz⁡\operatorname{nnz}nnz 超过 int32 范围，当前 API 会拒绝输入，而不是截断。

下一篇会进入剩余热路径：怎样融合 hˉ/x/h\bar h/x/hhˉ/x/h 递推，以及为什么最终 8kv/11kv 的迭代轨迹能与 CuPy 对齐。

9. 参考资料#

CUDA LSMR 专栏（一）：从迭代公式到分层库架构

Sat, 16 May 2026 00:00:00 GMT

这组文章记录一个可调用 CUDA LSMR 库的完整实现。第一篇先不急着写 kernel，而是回答三个更根本的问题：LSMR 每轮到底搬了什么数据，哪些变化属于数值算法，哪些变化应该被隔离在矩阵后端，以及怎样让 CSC、CSR 和稠密矩阵共用同一套迭代逻辑。

1. 问题与性能指标#

给定 A∈Rm×nA\in\mathbb{R}^{m\times n}A∈Rm×n 和 b∈Rmb\in\mathbb{R}^{m}b∈Rm，LSMR 求解线性系统或最小二乘问题。带阻尼的形式是

min⁡x∥[b0]−[AλI]x∥2.\min_x \left\| \begin{bmatrix} b\\0 \end{bmatrix} - \begin{bmatrix} A\\\lambda I \end{bmatrix}x \right\|_2.xmin[b0]−[AλI]x2.

它适合 AAA 很大、很稀疏，无法显式形成 ATAA^TAATA 的场景。每轮主要工作不是复杂标量计算，而是

Avk,ATuk+1,A v_k,\qquad A^T u_{k+1},Avk,ATuk+1,

外加若干长度为 mmm 或 nnn 的向量更新与范数归约。若 CSR 的值和列索引都是 4 字节，仅顺序读取矩阵本体，一次 SpMV 至少要读取约 8nnz⁡8\operatorname{nnz}8nnz 字节；每轮前向、转置各一次，因此矩阵流量下界已经接近

Bmatrix≥16nnz⁡ bytes/iteration.B_{\mathrm{matrix}}\ge 16\operatorname{nnz}\ \text{bytes/iteration}.Bmatrix≥16nnz bytes/iteration.

这还没有计入 x/yx/yx/y 向量、row offset 和缓存未命中。每个非零元素只贡献一次乘加，算术强度很低，LSMR 在大多数非结构稀疏矩阵上首先是显存带宽问题，而不是 FP32 峰值算力问题。

因此本文采用四类指标：迭代数和 stop code 检查算法轨迹；∥Ax−b∥2/∥b∥2\|Ax-b\|_2/\|b\|_2∥Ax−b∥2/∥b∥2 检查结果；单次 solve wall time 衡量端到端迭代开销；矩阵格式与工作区字节数解释显存代价。

2. 数学定义与 FP32 语义#

LSMR 基于 Golub-Kahan 双对角化。给定归一化后的 uk,vku_k,v_kuk,vk，核心递推是

βk+1uk+1=Avk−αkuk,\beta_{k+1}u_{k+1}=Av_k-\alpha_k u_k,βk+1uk+1=Avk−αkuk,αk+1vk+1=ATuk+1−βk+1vk.\alpha_{k+1}v_{k+1}=A^Tu_{k+1}-\beta_{k+1}v_k.αk+1vk+1=ATuk+1−βk+1vk.

随后通过一组稳定的对称 Givens 旋转更新 ρ,ρˉ,ζ\rho,\bar\rho,\zetaρ,ρˉ,ζ，再更新搜索方向和解：

hˉk=hk−θˉkρkρk−1ρˉk−1hˉk−1,\bar h_k=h_k- \frac{\bar\theta_k\rho_k}{\rho_{k-1}\bar\rho_{k-1}}\bar h_{k-1},hˉk=hk−ρk−1ρˉk−1θˉkρkhˉk−1,xk=xk−1+ζkρkρˉkhˉk.x_k=x_{k-1}+\frac{\zeta_k}{\rho_k\bar\rho_k}\bar h_k.xk=xk−1+ρkρˉkζkhˉk.

Golub-Kahan data flow

图 1：一轮 LSMR 的因果链。矩阵后端只负责两个蓝色 SpMV，旋转和停止条件不应知道 A 的存储格式。

库的矩阵、向量、SpMV 累加和输出均为 FP32。范数由 cuBLAS 以 FP32 输入计算，回到 host 后，旋转与停止规则用 double 标量计算，再把向量系数转换为 FP32。这样没有把整个迭代偷偷提升为 FP64，同时降低连续旋转中 host 标量舍入的影响。FP64 不是这组重建输入的默认路径，而且消费级 Ampere 的 FP64 吞吐代价也不适合被忽略。

3. 朴素实现为什么不够成为库#

直接翻译论文或 CuPy 的 Python 循环，可以很快得到正确结果：每轮分别计算 A @ v、A.T @ u，再执行缩放、加法和范数。但它有三个工程问题。

第一，算法层直接持有具体稀疏类型时，CSC、CSR 和 dense 会产生分叉，停止规则容易被复制。第二，u *= -alpha; u += A @ v 会生成临时向量并多走几遍 global memory。第三，矩阵描述符、SpMV workspace 和迭代向量若跟着函数调用反复创建，分配和 Python 调度会进入热路径。

CuPy 的实现仍然是重要基线，因为它公开了清晰的 LSMR 递推和 SciPy 兼容语义；但“参考递推”与“复刻运行时组织”是两件事。库的目标是让数值流程保持可核验，同时把存储、调度和资源复用下沉。

4. 分层数据流#

Layered CUDA LSMR architecture

图 2：从稳定 API 到 CUDA backend 的单向依赖。上层只依赖能力，下层可以独立更换实现或 dispatch 策略。

这里借鉴 CUB 的思路，不是照搬它的 Device/Dispatch/Agent 类名，而是保留“接口、策略、执行机制分离”的关键性质：

Public API 定义尺寸、精度、所有权和 host/device 位置。
Solver 只看 rows()、cols()、AxAxAx 和 ATxA^TxATx。
Operator 根据 sparse/dense 和方向选择 cuSPARSE 或 cuBLAS。
Runtime 统一管理 stream、handle、显存池和 RAII buffer。
Kernel 层只放无法由库调用合并的小向量递推。

OneFlow 的系统设计同样把 Operator、Kernel、DeviceCtx/stream 分开。这个边界对迭代求解器很实用：算法公式是 Operator 级语义，具体 CUDA 调用和异步执行属于 Kernel/runtime 语义。

5. 从接口到执行的逐层实现#

5.1 用 LinearOperator 截断格式依赖#

1
class LinearOperator {
2
 public:
3
  virtual ~LinearOperator() = default;
4
  virtual std::int64_t rows() const noexcept = 0;
5
  virtual std::int64_t cols() const noexcept = 0;
6

7
  virtual void apply_scaled(float alpha, const float* x, float beta,
8
                            float* y) const = 0;
9
  virtual void apply_transpose_scaled(float alpha, const float* x, float beta,
10
                                      float* y) const = 0;
11
};

接口直接暴露 y←αAx+βyy\leftarrow\alpha Ax+\beta yy←αAx+βy，而不是只暴露 y=Axy=Axy=Ax。原因不是追求抽象完整，而是让 Golub-Kahan 分子

Avk−αkukAv_k-\alpha_k u_kAvk−αkuk

可以直接映射成一次 cuSPARSE/cuBLAS 调用，避免先生成 Av 再单独读写 u。

5.2 PImpl 隔离 CUDA ABI 细节#

公共头文件不包含 cuSPARSE descriptor 和内部 buffer。Context、SparseMatrix、DenseMatrix、Solver 都使用 PImpl。这样使用者看到的是稳定 C++17 接口，CUDA 12.9 的描述符字段、算法枚举和 workspace 查询留在 .cu 实现中。

5.3 Context 是资源边界#

一个 Context 拥有一个 stream、一对绑定到该 stream 的 cuBLAS/cuSPARSE handle，以及一个私有 stream-ordered memory pool。矩阵和 Solver 引用它，但不各自创建 runtime。这也明确了并发规则：需要并行 stream 时创建多个 Context，而不是让多个 host 线程竞争同一组可变描述符。

6. shape、格式与 dispatch#

当前 dispatch 有三条路径：CSC 输入构造双格式稀疏算子；CSR 输入构造同一种双格式算子；column-major dense 输入构造 cuBLAS 算子。Solver 不需要 if (is_csc)。

这与 OneFlow Softmax 文章中的 shape 分段思想一致：工业实现不应假定一个 kernel 对所有数据分布都最优。当前版本先按存储类别分派；后续如果加入按行长方差选择 CSR algorithm，或为极短行增加 warp 专用 kernel，也只修改 Operator 层。

代价同样明确：虚函数调用相对一次大 SpMV 可以忽略，但这套接口不适合在 device 内部逐元素调用；双格式存储还会接近翻倍占用矩阵显存。第二篇会单独量化这项交换。

7. 关键 CUDA 调用边界#

稀疏后端缓存矩阵和 dense-vector descriptor，只在调用前替换输入输出指针：

1
cusparseDnVecSetValues(input, const_cast<float*>(x));
2
cusparseDnVecSetValues(output, y);
3
cusparseSpMV(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha,
4
             matrix, input, &beta, output, CUDA_R_32F,
5
             CUSPARSE_SPMV_CSR_ALG1, workspace);

描述符和 workspace 的创建不在迭代循环内。这里选择 CUSPARSE_SPMV_CSR_ALG1 是当前 CUDA 12.9/FP32 实测路径，而不是承诺所有 GPU 上都最优；cuSPARSE 文档也明确说明高性能算法不保证逐 bit 可复现。

8. 验证方法、局限与回退#

框架阶段先用一个 3×23\times23×2 矩阵验证 CSC、CSR、dense 三个后端的 AxAxAx 与 ATxA^TxATx 完全一致，再用同一 Solver 测 dense b、sparse b 和多列 dense RHS。真实 benchmark 则固定 FP32、相同容差、相同最大迭代数，并同时比较 stop code、迭代数、解差与独立残差。

当前边界包括：仅实数 FP32 与 int32 索引；没有预条件器；多 RHS 是复用 workspace 后逐列求解，并非 block LSMR；当双格式超出显存预算时，应回退到单格式和 transpose SpMV，而不是冒险 OOM。

下一篇进入最关键的数据布局：为什么 CSC 输入仍要生成 CSR(A)，以及怎样让 ATuA^TuATu 也变成顺序行遍历。

9. 参考资料#

大模型FCN算子（GEMM）CUDA实现

Wed, 15 Apr 2026 00:00:00 GMT

无论是 Transformer 的 QKV 投影、MLP 的 up/gate/down projection，还是卷积经 im2col 之后的计算，本质上都在反复做同一件事：

D=αAB+βC,A∈RM×K, B∈RK×N.D=\alpha AB+\beta C,\qquad A\in\mathbb R^{M\times K},\ B\in\mathbb R^{K\times N}.D=αAB+βC,A∈RM×K, B∈RK×N.

GEMM 从全局内存到 CTA、warp 与 MMA 指令级分块的数学数据流图

图 1 把一个输出 tile 的计算拆成三个尺度：CTA 负责 BM × BN 的输出，warp 消费 CTA tile 的子块，MMA 指令只计算其中的 fragment。沿 K 维遍历时，A/B tile 被重复装载与复用，C tile 一直留在寄存器累加。

GEMM 不是“写三层循环”这么简单。真正难的是：同一套模型会同时遇到 prefill 的大矩阵、decode 的 skinny GEMM、MoE 的不规则 grouped GEMM，以及 FP4/FP8/INT8 到 FP64 的不同数值约束。为某个方形 FP16 shape 跑满，不等于实现了一个可放进框架的 FCN 算子。

本文沿用一个很实用的讲解顺序：先算清楚性能上界，再从朴素核函数逐层推出分块、访存、流水和 Tensor Core；最后把它们组织成按形状和精度 dispatch 的算子族。这里不虚构某块 GPU 的绝对跑分，所有“最快”的结论都必须在目标 GPU、目标 CUDA 和真实 shape 上重新测量。

1. 性能目标到底是什么#

单次 GEMM 的工作量为：

FLOPs=2MNK.\text{FLOPs}=2MNK.FLOPs=2MNK.

运行时间为 ttt 时，吞吐为 2MNK/t2MNK/t2MNK/t。但只报 TFLOP/s 不足以判断 kernel 好坏，因为它还受矩阵是否足够大、是否对齐、是否发生额外读写、launch 开销和数值精度影响。

一个有用的最小模型是 Roofline：

P≤min⁡(Pcompute peak, I×BWmemory),P\leq\min(P_{\text{compute peak}},\ I\times BW_{\text{memory}}),P≤min(Pcompute peak, I×BWmemory),

其中 III 是算术强度。GEMM 通过 tile 复用 A/B 的元素，通常最终受计算单元约束；而小矩阵、skinny GEMM、拆得过细的 epilogue 则可能重新落回带宽或 launch 受限。

因此，“极致优化”不是只看一项指标，而是同时做到：

A/B 的全局访存合并且对齐；
shared memory 的布局没有可避免的 bank conflict；
Tensor Core 或 CUDA Core 持续有数据可算；
寄存器、shared memory、CTA 数之间保持可接受的 occupancy；
输出的 bias、activation、residual、量化尽量在同一次写回中完成；
对不适合主路径的形状及时换 kernel，而不是硬撑。

2. 从朴素实现出发：为什么它离峰值很远#

标量 baseline 的逻辑如下：

1
template <typename T>
2
__global__ void gemm_naive(int M, int N, int K,
3
                           const T* A, const T* B, T* D) {
4
  int m = blockIdx.y * blockDim.y + threadIdx.y;
5
  int n = blockIdx.x * blockDim.x + threadIdx.x;
6
  if (m >= M || n >= N) return;
7
  float acc = 0.f;
8
  for (int k = 0; k < K; ++k)
9
    acc += float(A[m * K + k]) * float(B[k * N + n]);
10
  D[m * N + n] = T(acc);
11
}

它正确，却几乎没有复用：相邻 thread 会重复读取同一行 A 或同一列 B。即使 L1/L2 能帮上一部分忙，数据搬运仍由缓存偶然性决定，无法稳定接近峰值。优化的第一原则就是把“会被多个输出元素使用的数据”显式搬到片上。

3. 第一级：CTA tile 与层次化分块#

一个 thread block 计算 DDD 的 BM × BN tile，并沿 K 每次推进 BK。数据层级是：Global A/B → shared-memory A/B tile → warp tile → registers / MMA fragments → D tile。

在每个 K tile 中，CTA 协作将 A/B 连续搬入 shared memory；之后每个 warp 读取其中一小块，多个输出元素复用同一份数据。关键参数不是“越大越好”：

BM, BN 越大，A/B 的复用越好，但累加器和 shared memory 也更多；
BK 太小，循环与同步开销明显；太大则挤占 shared memory、降低常驻 CTA；
warp tile 决定寄存器 fragment 大小与并行度；
instruction tile 必须匹配实际架构和 dtype。

一个正确的调优流程是先枚举少量合理 tile，而不是从无穷多数字中盲猜。以常见矩阵乘数据通路为例：BM/BN 取 64、128、256 一类，BK 对 FP16/BF16 常取 32/64，warp tile 与 MMA 指令 tile 对齐；实际值必须让编译器资源报告和 profile 说话。

4. 第二级：访存布局比算术更早决定成败#

4.1 全局内存：连续、对齐、向量化#

线程应协作读取连续地址，使一个 warp 的访问能尽量合并为少量内存事务。对齐满足时，FP32 可使用 float4，FP16/BF16 可用成对或更宽的 pack；否则必须有独立的尾部路径，不能用未对齐的强制 reinterpret cast 赌运气。

向量化的价值不只是少几条指令，还在于明确表达连续访问。OneFlow 在优化 Softmax 时也把 half 打包成 half2，借助更高的指令吞吐处理成对元素；GEMM 的 A/B 装载同样应把 dtype、对齐和 layout 放入模板/dispatch 条件，而不是散落在 kernel 分支里。

4.2 shared memory：布局必须为消费端服务#

shared memory 不是“放进去就快”。若同一 warp 的线程访问同一 bank 的不同地址，会发生 bank conflict，访问被串行化。常见处理方式包括 padding、swizzle，以及按 ldmatrix/MMA 的消费顺序重排 layout。

尤其需要避免一种误区：全局内存写入连续，并不保证 shared memory 读取也连续。生产实现会分别设计 global-to-shared 和 shared-to-register 的 layout，并在 Nsight Compute 中检查 shared load/store 的冲突指标。OneFlow 对 shared buffer 调整索引顺序以避开 bank conflict 的案例，是这条原则的一个小而典型的例子。OneFlow Softmax CUDA kernel 优化

4.3 寄存器：GEMM 的真正工作台#

每个 warp 对应的输出 fragment 长时间留在寄存器，只在 K 循环结束后写回全局内存。任何把中间 A×B partial 频繁写到 global memory、再由下一 kernel 读取的设计，都会直接破坏 GEMM 的算术强度。

不过，寄存器也不是越多越好。过多 accumulator 会导致 spill 或降低 occupancy；对于 decode 的小 M，可能需要 persistent CTA 或更小 tile 来获得足够并发。

5. 第三级：让搬运和计算真正重叠#

同步式 tile loop 是“搬 A/B → 等待 → 计算 → 再搬”的串行流程。现代 GPU 的关键是多阶段流水：当 stage s 的 fragment 正在 MMA 时，stage s+1 的 A/B tile 已经在路上。

1
// stage 0: async copy A0/B0 -> shared
2
// stage 1: async copy A1/B1 -> shared, compute A0 * B0
3
// stage 2: async copy A2/B2 -> shared, compute A1 * B1
4
// ...

在支持的架构上，可用异步 copy 建立双缓冲或多缓冲；再新的数据通路可以由专门的搬运机制和 warp-specialized producer/consumer 组织。核心不在于记住某个 intrinsic 名称，而在于三个不变量：

GEMM K 维累加与多阶段流水：下一 tile 搬运和当前 tile 矩阵乘重叠

图 2 中的 C_acc 是同一输出 tile 的寄存器累加器；它只在所有 K tile 完成后才交给 epilogue。双缓冲的目的不是增加计算量，而是让 load(k+1) 与 MMA(k) 重叠。

consumer 在读取前确认对应 stage 已完成；
producer 不覆盖仍被消费的 stage；
stage 数增加带来的 latency hiding，必须大于它对 shared memory 和 occupancy 的代价。

CUTLASS 把这些策略封装为可组合的 mainloop、collective 和 epilogue，是阅读工业级 GEMM 流水线的好入口；其近代实现还覆盖了 grouped GEMM、FP8、block-scaled 类型与架构专用调度。NVIDIA CUTLASS

6. Tensor Core 不是一个 dtype 开关#

Tensor Core 路径的本质是 warp/warp-group 使用特定矩阵乘累加指令，操作 fragment 而非单个标量。要让它工作，至少要同时满足：dtype、矩阵布局、指令 tile、K 对齐和目标计算能力的约束。

典型策略是：主路径处理满足 tile 对齐的大块，边界使用 predicated load/store 或单独的 residue kernel；不要让每一次 MMA 都包在复杂的逐元素边界判断里。对于常见混合精度，输入可以是 FP16/BF16/TF32/FP8，而 accumulator 通常是 FP32；CUDA Programming Guide 明确说明 BF16 WMMA fragment 以 float 为累加类型，TF32 需要从 FP32 显式转换得到相应语义。CUDA C++ Programming Guide

7. 精度矩阵：从 INT4/FP4 到 FP64#

“支持所有精度”不应理解为一个模板把所有类型硬塞进一个指令序列，而是每种精度都有明确的输入、累加、缩放、输出与 fallback 规则。

输入 A/B	默认累加	典型用途	实现重点
INT4 / INT8	INT32 或 FP32	量化推理	per-channel/per-group scale、零点、饱和及反量化 epilogue
FP4	FP32 或框架指定精度	超低比特推理	block scale 是数据布局的一部分，不是事后乘法
FP8 E4M3/E5M2	FP32	训练/推理加速	scale、amax 管理、fast-accum 的误差边界
FP16 / BF16	FP32	训练与通用推理	MMA + FP32 accumulator；BF16 保留更大指数范围
TF32 语义	FP32	FP32 范围、可接受 reduced mantissa	仅在允许的数值合同下打开
FP32	FP32	精度基线、敏感算子	CUDA Core 或适配的 tensor 路径，严测误差
FP64	FP64	科学计算、高精度参考	走目标硬件支持的双精度路径，吞吐/布局另行调优

这里有四个容易被忽略的细节：

存储与累加分离。 FP16 输入、FP32 累加不是“可选优化”，而应是通用训练的默认数值合同。
缩放粒度属于 ABI。 INT4/FP4/FP8 的 scale 按 tensor、channel、group 还是 block 存储，会同时决定精度、访存和 kernel layout；必须在 tensor descriptor 中明确。
K 方向分割会影响数值。 split-K 改变累加顺序；高精度或可复现模式应采用固定归约树，必要时使用更高精度 partial。
低比特不等于总更快。 若 shape 太小、scale 读取占主导，或需要单独 dequant kernel，端到端速度可能反而下降；理想做法是把 dequant/scale 融合在加载和 epilogue 中。

当前 CUTLASS 的功能与变更记录展示了 FP8、NVFP4、MX 格式、block scaling 和 sparse/block-scaled GEMM 等路线；它们说明低比特 GEMM 是一套数据布局与调度问题，而不只是把 float 换成更短的类型。CUTLASS changelog

8. GEMM 不是一种形状：分段 dispatch 才是通用解#

和 OneFlow 的分段 Softmax 一样，GEMM 也不存在一个在所有形状上最优的 kernel。一个实用 dispatcher 至少把以下信息作为 key：

1
{arch, dtypeA, dtypeB, dtypeD, accumulator,
2
 layoutA, layoutB, layoutD, M, N, K,
3
 alignmentA/B/D, alpha/beta, epilogue, batch/group count,
4
 deterministic, workspace limit}

对应的主路径可这样划分：

问题特征	优先策略	原因
大 M/N/K，整齐对齐	大 CTA tile + 多阶段 Tensor Core pipeline	复用高、可计算饱和
`M` 很小、decode skinny GEMM	persistent CTA / 更小 tile / grouped GEMM	避免 CTA 不足与 launch 主导
`K` 极大，M×N tile 不足	split-K 或 Stream-K	增加并行度，随后规约
多个小而异形的 GEMM（MoE）	grouped GEMM	合并 launch，动态平衡工作
M/N/K 边界不对齐	主 kernel + residue，或 predicate tile	不让边界条件污染内层主循环
FP32 高精度或 FP64	独立高精度 kernel	不与低精度 Tensor Core 假设混用

split-K 并不是免费午餐：它增加 workspace 或原子/归约开销，只有原始网格并行度不足时才值得使用。Stream-K 的意义也正是在 tile 间重新分配 K 工作，缓解固定 tile 调度的不平衡；它应作为 benchmark 候选，而非默认开关。

9. 一个可维护的 kernel 骨架#

生产代码通常不把所有细节塞进 __global__ 函数，而是把策略变为类型和常量。下面是结构示意：

1
template <class ElementA, class ElementB, class ElementD,
2
          class Accumulator, class Tile, class Layout,
3
          class Mainloop, class Epilogue>
4
struct GemmKernel {
5
  __device__ void operator()(GemmArgs args) {
6
    // 1. 根据 CTA 坐标取得 D 的 Tile。
7
    // 2. Mainloop: 预取 A/B，管理 pipeline stage，执行 MMA。
8
    // 3. Epilogue: alpha/beta、bias、residual、activation、量化/转换。
9
    // 4. 只写回最终 D。
10
  }
11
};
12

13
GemmPlan make_plan(const GemmArgs& args) {
14
  // 静态规则先过滤不合法的 dtype/layout/alignment；
15
  // 查 shape+arch 的 autotune 缓存；缺失时仅 benchmark 少量候选。
16
}

展开收起

这种分层使主循环只关心“如何高复用地计算”，epilogue 只关心“如何生成最终输出”，dispatcher 则对形状做选择。它也让你能把框架的 bias、GELU/SwiGLU、残差、输出量化等融合进去，而不污染 MMA 内层循环。

10. 融合 epilogue：节省的往往比 MMA 更真实#

以 FCN 为例，常见路径是：

1
D = A @ W
2
Y = activation(D + bias)

若拆成两个 kernel，D 至少经历一次完整的 global write 与一次 global read。把 alpha/beta + bias + activation + cast/quantize 放在 GEMM epilogue 中，可以让 accumulator fragment 在寄存器中完成最后变换后只写一次 Y。

但融合有边界：复杂 activation、大量额外输入或分支会提高寄存器占用，可能伤害主 GEMM。判断标准永远是端到端 benchmark，而不是“kernel 数量越少越好”。

11. 如何验证“快”与“对”#

基准至少应包含 LLM 常见 shape：prefill 的大 M、decode 的 M=1..16、MLP 的大 N/K、QKV 的非方阵、MoE 的 grouped GEMM，以及难看的不对齐尾部。每项都要记录：

1
environment: GPU/SM、CUDA/驱动、时钟状态
2
problem: 矩阵 layout、M/N/K、dtype、accumulator
3
timing: warmup、重复次数、CUDA event 时间、总时间、TFLOP/s、workspace
4
correctness: 最大绝对/相对误差、随机种子、确定性模式

正确性基准要独立于被测实现：低精度输入可用 FP32 或 FP64 reference 比较，比较容忍度随 dtype 和 K 长度设置；对 FP8/FP4/INT4，额外测 scale 边界、NaN/Inf、饱和、非整齐 K 和 beta 非零。性能分析使用 Nsight Compute 观察 Tensor Core 活跃度、DRAM 吞吐、shared bank conflict、寄存器 spill、occupancy 和 warp stall 原因。

不要只对一个漂亮的 8192×8192×8192 报数。真正的大模型推理，往往被 decode、短 batch、量化 scale、融合边界和 grouped workload 决定。

12. 常见的“优化”陷阱#

用一个 tile 打天下。 大矩阵最优 tile 在 skinny GEMM 上可能只会浪费线程。
无条件加深 pipeline。 stage 多会占用更多 shared memory，可能使常驻 CTA 下降。
只测 kernel，不测 epilogue/量化/转换。 用户感知的是端到端时间。
默认 split-K。 没有并行度问题时，额外规约就是纯负担。
将 TF32/FP8 当作无损 FP32。 数值合同必须由模型和测试决定。
把未对齐指针强转成 vector 类型。 这不是优化，是未定义行为风险。
忽略 launch 开销。 对 decode 或小 GEMM，CUDA Graph、grouped/persistent 调度常比更大的 MMA tile 重要。

总结#

高性能 GEMM 的路径可以概括为：用 CTA/warp/instruction 三层 tile 建立复用，用对齐向量化和无冲突 shared layout 喂饱计算单元，用异步多阶段流水隐藏搬运延迟，用 Tensor Core 与混合精度匹配硬件能力，最后用 shape-aware dispatch 覆盖大矩阵、skinny、split-K 与 grouped workload。

对于大模型 FCN，真正可上线的版本还必须把 FP4/INT4、FP8、FP16/BF16、TF32、FP32、FP64 的数值合同和 fallback 写清楚，并把 bias/activation/quantize 融进合适的 epilogue。所谓“极致优化”，从来不是某个花哨 intrinsic，而是一套在真实输入、真实数值约束和真实硬件上持续测得最快的系统工程。

参考资料#

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

Sun, 05 Apr 2026 00:00:00 GMT

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

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

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

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

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

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

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

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

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

格式要服从结构。 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：为块结构准备的快速通道#

当非零元以 br×bcb_r\times b_cbr×bc 小块出现时，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σ 内排序，再把每 CCC 行切成切片。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 计算 Ci,n=∑k∈N(i)Ai,kBk,nC_{i,n}=\sum_{k\in\mathcal N(i)}A_{i,k}B_{k,n}Ci,n=∑k∈N(i)Ai,kBk,n。同一个 A[i,k] 可用于 N_TILE 个列，因而可以把 BBB 的连续列放入寄存器累加：

CSR SpMM 中一个稀疏值复用一段连续的 B 行，形成多个输出列累加器

图 2 的关键差别在于右端项不再是一个标量：一次读取 Ai,kA_{i,k}Ai,k 后，连续的 Bk,n:n+NtileB_{k,n:n+N_{tile}}Bk,n:n+Ntile 同时更新多个寄存器累加器，因而将 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 口径都放进设计，才是“通用”二字真正的含义。

参考资料#

使用CUDA实现并极致优化softmax

Sat, 04 Apr 2026 00:00:00 GMT

这篇文章是草稿#

这篇文章目前处于草稿状态，尚未发布。因此，它不会对普通读者可见。内容仍在进行中，可能需要进一步编辑和审查。

当文章准备发布时，您可以在 Frontmatter 中将 “draft” 字段更新为 “false”：

1
---
2
title: 草稿示例
3
published: 2024-01-11T04:40:26.381Z
4
tags: [Markdown, 博客, 演示]
5
category: 示例
6
draft: false
7
---