矩阵、注意力与 GPU：拆解 DeepGEMM 与 FlashMLA

LLM 推理的瓶颈：矩阵乘

LLM 推理的时间几乎都花在矩阵乘上。一层 Transformer 里，Q/K/V/O 四个投影是四个 GEMM，MLP 的两个线性层是两个 GEMM，注意力本身还是两个 GEMM。一个 token 走完一层，绝大多数 FLOPs 都在做 $C = A \times B$。

前端其实不陌生：CSS matrix3d 是 4×4，canvas 的 transform 是仿射矩阵，图像卷积是滑动窗口的点积。区别只在于规模——GPU 上的 GEMM 是 $4096 \times 4096 \times 4096$ 起步，而且要在 0.1 毫秒量级完成。

这篇文章想解决三个问题：

1. 一个矩阵乘在 GPU 上到底是怎么被拆开、搬运、计算的；
2. 注意力为什么会退化成矩阵乘，MLA 又用什么代数技巧把 KV cache 压小；
3. FlashMLA 为什么能在 H800 上跑到 3000 GB/s、580 TFLOPS，它的调度到底在调度什么。

数学只用到点积、softmax 和结合律。代码会给出 JS 版本，方便把概念跑起来。

---

一、矩阵乘法：三重循环，两个瓶颈

先统一记号。DeepGEMM 的约定是

$$D = C + A B$$

（Python 里写作 A @ B。）其中 $A$ 是 $M \times K$，$B$ 是 $K \times N$，$D$ 是 $M \times N$。默认 NT 布局：A 行主序（normal），B 以转置形式存放（transposed），也就是计算 $A \cdot B^\top$。这是推理里最常见的形状——权重按 [N, K] 存，激活按 [M, K] 存，两者做内积。

朴素实现长这样：

function gemmNaive(A, B, C, M, N, K) {
  for (let i = 0; i < M; i++) {
    for (let j = 0; j < N; j++) {
      let acc = 0;
      for (let k = 0; k < K; k++) {
        acc += A[i * K + k] * B[k * N + j];
      }
      C[i * N + j] += acc;
    }
  }
}

三层循环，$2MNK$ 次浮点运算。以 $M=N=K=4096$ 为例，算力需求是 137 GFLOP。H800 的 BF16 峰值约 989 TFLOPS，理论上 0.14 ms 就能算完。

但朴素实现的访存是灾难：每算一个输出元素，就要读 $2K$ 个输入。总共读 $2MNK$ 次浮点数，约 549 GB。按 HBM 3.35 TB/s 算，需要 164 ms。

算 0.14 ms，搬 164 ms。差了一千倍。

这就是 GEMM 优化的全部动机：不是让乘法更快，是让同一个数被用更多次。

分块：把复用做进循环顺序

把输出切成 $T \times T$ 的小块，外层按块遍历，内层在块内累加：

function gemmTiled(A, B, C, M, N, K, T = 32) {
  for (let i0 = 0; i0 < M; i0 += T) {
    for (let j0 = 0; j0 < N; j0 += T) {
      for (let i = i0; i < Math.min(i0 + T, M); i++) {
        for (let j = j0; j < Math.min(j0 + T, N); j++) {
          let acc = 0;
          for (let k = 0; k < K; k++) {
            acc += A[i * K + k] * B[k * N + j];
          }
          C[i * N + j] += acc;
        }
      }
    }
  }
}

循环顺序变了，效果完全不同：现在处理一个 $T \times T$ 输出块时，A 的 $T$ 行、B 的 $T$ 列会被反复使用，缓存命中率大幅提升。这叫 cache blocking，GPU 上的「tiling」是同一件事，只是"缓存"换成了片上 shared memory。

判断一个 kernel 是算力瓶颈还是带宽瓶颈，用算术强度（arithmetic intensity，FLOP/Byte）：每从 HBM 读 1 字节，能做多少次浮点运算。

场景	算术强度	H800 机器平衡点	结论
朴素 GEMM	~2	~295 (BF16)	带宽受限，慢 1000 倍
128×128 分块	~32	295	仍差一个数量级
分块 + 寄存器复用 + 大 N	100~500	295	进入算力受限区

H800 的机器平衡点怎么算：$989 \times 10^{12} / 3.35 \times 10^{12} \approx 295$ FLOP/Byte。任何算术强度低于这个值的 kernel，Tensor Core 都会饿着。

DeepGEMM 和 FlashMLA 的所有设计，本质都是在提高算术强度，或者在接受带宽受限的前提下把带宽打满。

图 1 · 一个输出 tile 同时复用 A 的 BLOCK_M 行与 B 的 BLOCK_N 列；算术强度决定它被带宽还是被算力卡住。

---

二、Hopper 上的一次 GEMM 长什么样

硬件层次

一张 H800 SXM5：

层级	容量 / 带宽	作用
HBM3	80 GB，3.35 TB/s	权重、激活、KV cache
L2	50 MB 左右	跨 SM 复用
Shared Memory	228 KB / SM	手写管理的片上暂存
寄存器	256 KB / SM（65536 个 32-bit）	累加器、线程私有
SM	132 个	每个 SM 4 个 warp scheduler

一次 GEMM 的循环体，就是把数据从 HBM 搬到 shared memory，再从 shared memory 喂给 Tensor Core，算完写回。搬运和计算必须重叠，否则每一级都在等上一级。

三个关键指令

WGMMA（Warpgroup Matrix Multiply-Accumulate）：Hopper 的 Tensor Core 指令，一个 warpgroup 是 128 个线程，一条 WGMMA 可以算出 $64 \times N$ 的一块输出，累加器驻留在寄存器里。操作数 A 可以直接来自 shared memory，不需要先读进寄存器——这是 Hopper 相对 Ampere 的大变化。

TMA（Tensor Memory Accelerator）：一个专门的拷贝引擎。CPU 侧构造一个 tensor map 描述多维布局，GPU 侧一条指令就能把一块 2D/3D 数据从 HBM 搬进 shared memory，不占寄存器，不占线程。它还能做 multicast：集群里多个 CTA 共享同一份权重，只读一次 HBM。

mbarrier：异步屏障。TMA 拷贝完成、WGMMA 读完 shared memory，都通过 barrier 通知，线程不需要自旋等待。

warp specialization：把线程分工写死

DeepGEMM 的 SM90 FP8 kernel 是这套东西的教科书实现。入口在 deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh:39：

sm90_fp8_gemm_1d1d_impl(__nv_fp8_e4m3* gmem_a_ptr, __nv_fp8_e4m3* gmem_b_ptr,
                        int* grouped_layout,
                        cute::TmaDescriptor* tensor_map_buffer,
                        uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
                        const __grid_constant__ cute::TmaDescriptor tensor_map_a_base,
                        ...)

第 50 行断言了线程分工：

DG_STATIC_ASSERT(kNumTMAThreads == 128 and kNumMathThreads % 128 == 0, "Invalid Threads");
DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling");

一个 128 线程的 TMA warpgroup 负责搬运（实际只有一条线程真正发指令），其余线程按 128 一组组成 math warpgroup 发 WGMMA。两类线程走完全不同的分支：同一个文件 171 行的 TMA warp 调 warpgroup_reg_dealloc<24>()，释放寄存器；250 行 math warpgroup 调 warpgroup_reg_alloc<240>()，把寄存器抢过来给累加器用。

同一个 SM 上，寄存器总数固定在 65536 个。给搬运线程少一点、给计算线程多一点，是硬约束下的零和分配（sm90_fp8_gemm_1d1d.cuh:153-154 定义了 24 和 240 两个数值；未完全展开时分支取 40 / 232，正常路径走 24 / 240。新手容易以为只是调参，其实是 occupancy 的直接来源）。

主循环的骨架：

// TMA warp：每个 stage 等 consumer 释放，然后一次性发 4 条 TMA
empty_barriers[stage_idx]->wait(phase ^ 1);
tma::copy<BLOCK_M, BLOCK_K, 0>(&tensor_map_sfa, &full_barrier, smem_sfa[stage_idx], ...);
tma::copy<BLOCK_N, BLOCK_K, 0>(&tensor_map_sfb, &full_barrier, smem_sfb[stage_idx], ...);
tma::copy<BLOCK_K, BLOCK_M, kSwizzleAMode>(tensor_map_a_ptr, &full_barrier, smem_a[stage_idx], ...);
tma::copy<BLOCK_K, BLOCK_N, kSwizzleBMode>(tensor_map_b_ptr, &full_barrier, smem_b[stage_idx], ...);
full_barrier.arrive_and_expect_tx(...);

// Math warpgroup：等数据到齐，发 WGMMA
full_barriers[stage_idx]->wait(phase);
#pragma unroll
for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++k) {
    WGMMA::wgmma(desc_a, desc_b, accum, k);
}

多 stage 环形缓冲 + 双向 barrier，形成一条流水线：TMA 在搬第 $i+2$ 块时，Tensor Core 在算第 $i$ 块，CUDA Core 在处理第 $i-1$ 块的缩放。

另外两个细节值得注意：

• 持久化调度：scheduler/gemm.cuh:185 的 get_next_block() 让 CTA 处理完一个 tile 后继续取下一个，而不是退出重开。省掉了 block 启动开销，也让流水线不会被打断。
• PDL：cudaGridDependencySynchronize()（第 157 行）配合 Programmatic Dependent Launch，让下一个 kernel 在当前 kernel 结束前就开始预热，隐藏 launch 延迟。

图 2 · TMA warp 只负责搬，math warpgroup 只负责算；多 stage + mbarrier 把 HBM、SMEM、Tensor Core 串成流水线。

---

三、FP8：省一半带宽，多一个 scale

为什么要量化

BF16 每个数 2 字节，FP8 每个数 1 字节。对带宽受限的 GEMM，直接翻倍；对算力受限的 GEMM，Tensor Core 的 FP8 吞吐是 BF16 的两倍（H800：1979 vs 989 TFLOPS）。

FP8 有两种格式：e4m3（4 位指数、3 位尾数，最大 448）和 e5m2（5 位指数、2 位尾数，范围大但精度低）。矩阵乘用 e4m3，因为动态范围够用，精度更重要。

问题是 e4m3 只有 3 位尾数，相对误差约 $2^{-4}$。整层共用一个 scale 会把小值全部压成 0，所以必须分块量化。

DeepSeek-V3 的训练配方：

• 激活：每 token、每 128 个通道一个 scale（$1 \times 128$）；
• 权重：每 $128 \times 128$ 的块一个 scale。

sm90_fp8_gemm_1d1d.cuh:51 的 BLOCK_K == 128 断言，就是在硬编码这个 K 方向粒度。仓库里 1d1d 对应 A、B 都按行给 scale（shared memory 里 SFA 是 BLOCK_M 个 float、SFB 是 BLOCK_N 个 float），1d2d 对应 B 用二维块 scale，也就是 V3 权重那套 $128 \times 128$ 配方。

两级累加：Tensor Core 算，CUDA Core 修

FP8 的乘加在 Tensor Core 里完成，累加器是 FP32。但每个 K 块的数据有各自的 scale，直接累加是错的。

正确做法是：每个 K 块单独算出一段部分和，乘上 $\text{scale}_A \times \text{scale}_B$，再加进最终累加器。这就是 DeepGEMM 说的 two-level accumulation（promotion）。

代码在 sm90_fp8_gemm_1d1d.cuh:302 到 325 行：

#pragma unroll
for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++k) {
    WGMMA::wgmma(desc_a, desc_b, accum, k);   // 只累加 FP8 乘积
}
...
// Promote with scales
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++i) {
    final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0];
    final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1];
    ...
}

为什么不能每算一条 WGMMA 就做一次 promotion？因为 promotion 走 CUDA Core，是标量/向量乘加，吞吐远低于 Tensor Core。做太频繁，CUDA Core 变成瓶颈，Tensor Core 等它。

早期版本的 DeepGEMM README 里提到，他们让 promotion 和 MMA 在指令级交错（FFMA interleaving），并调 SASS 的 yield / reuse 位；2025 年 7 月之后，NVCC 12.9 能自动做这个交错，手写优化就删掉了。这个演进本身挺有意思：编译器追上了手写 kernel 的一个关键技巧。

图 3 · FP8 的 scale 按块给；Tensor Core 只做局部分和，CUDA Core 在每个 BLOCK_K 结束时做一次 promotion。

scale 的形状也是性能问题

scale 张量本身也要被 TMA 搬运，所以它的布局必须满足 TMA 的对齐要求，而且要转置成 MN-major。DeepGEMM 专门有一篇文档 docs/scaling-factor-format.md 讲这个：

• SM90：FP32 scale，要求 TMA-aligned + transposed；
• SM100：UE8M0 格式，4 个 8 位指数打包进一个 int32。

UE8M0 的含义是：scale 只存 2 的幂次，没有尾数。这样乘 scale 等价于调整指数，硬件实现更便宜。代价是 scale 必须向上取整到 2 的幂（文档里 round_sf=True），精度略损。

最终成绩：2025 年 4 月，DeepGEMM 在 H800 上做到 1550 TFLOPS（README.md:25），是 FP8 峰值 1979 的 78%。

---

四、MoE 把 GEMM 变成了「分组 GEMM」

MoE 层里，每个 token 被路由到 top-k 个专家，每个专家有自己的权重矩阵。理想情况下每个专家都做一次 GEMM，但专家的 token 数是动态的、不均衡的。

DeepGEMM 支持三种形状：

接口	用途	特点
m_grouped_fp8_gemm_nt_contiguous	训练前向 / 推理 prefill	按 M 轴分组，token 拼成连续张量，N/K 固定
m_grouped_fp8_gemm_nt_masked	推理 decode + CUDA Graph	用 mask 标出有效行，只算有效部分
k_grouped_fp8_gemm_tn_contiguous	MoE 权重反向	按 K 轴分组，M/N 固定

masked 这个设计值得多说一句。decode 阶段开 CUDA Graph 时，CPU 不能根据实际 token 数动态改 kernel 参数——图是静态的。所以 kernel 必须按最大形状启动，再靠一个 mask 张量跳过空行。这不是性能优化，是正确性约束逼出来的接口。

---

五、注意力：两个 GEMM 加一个 softmax

公式

单头注意力：

$$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d}}\right)V$$

拆开看：

1. $S = QK^\top$：$[s_q, d] \times [d, s_k] \to [s_q, s_k]$，一个 GEMM；
2. $P = \text{softmax}(S)$：逐行归一化，非矩阵运算；
3. $O = PV$：$[s_q, s_k] \times [s_k, d] \to [s_q, d]$，第二个 GEMM。

多头就是并行做 $h$ 份，再把结果拼起来过一个输出投影。

朴素实现（注意 S 矩阵是完整落地的）：

function attention(Q, K, V, d) {
  const s = Q.map(() => new Float32Array(K.length));

  // 1. S = Q K^T
  for (let i = 0; i < Q.length; i++)
    for (let j = 0; j < K.length; j++) {
      let acc = 0;
      for (let t = 0; t < d; t++) acc += Q[i][t] * K[j][t];
      s[i][j] = acc / Math.sqrt(d);
    }

  // 2. softmax，逐行
  for (let i = 0; i < Q.length; i++) {
    const row = s[i];
    const m = Math.max(...row);
    let sum = 0;
    for (let j = 0; j < row.length; j++) {
      row[j] = Math.exp(row[j] - m);
      sum += row[j];
    }
    for (let j = 0; j < row.length; j++) row[j] /= sum;
  }

  // 3. O = P V
  const out = Q.map(() => new Float32Array(d));
  for (let i = 0; i < Q.length; i++)
    for (let j = 0; j < K.length; j++)
      for (let t = 0; t < d; t++) out[i][t] += s[i][j] * V[j][t];
  return out;
}

问题在第 1 步：S 的大小是 $s_q \times s_k$。decode 时 $s_k$ 可能是 128K，每个头每个请求都要存一份，显存直接爆掉。

FlashAttention 的观察是：softmax 和 $PV$ 可以边算边合并。维护三个数——running max $m$、running sum $l$、部分输出 $\text{acc}$——每读一个 KV token 就更新一次：

function flashDecode(q, K, V) {
  const d = q.length;
  let m = -Infinity,
    l = 0;
  const acc = new Float32Array(d);

  for (let j = 0; j < K.length; j++) {
    let s = 0;
    for (let t = 0; t < d; t++) s += q[t] * K[j][t];
    s /= Math.sqrt(d);

    const mNew = Math.max(m, s);
    const alpha = Math.exp(m - mNew); // 旧输出的修正系数
    const p = Math.exp(s - mNew); // 新 token 的权重
    l = l * alpha + p;
    for (let t = 0; t < d; t++) acc[t] = acc[t] * alpha + p * V[j][t];
    m = mNew;
  }
  for (let t = 0; t < d; t++) acc[t] /= l;
  return acc;
}

显存从 $O(s_q s_k)$ 降到 $O(d)$。数学结果和朴素实现完全一致，只是改变了计算顺序。这就是 FlashAttention 的核心：它不是近似算法，是 IO 算法。

图 4 · 注意力的两个 GEMM；FlashAttention 把 S 的物化改成 (m, l, acc) 三个状态的流式合并。

KV cache 的账

自回归 decode 时，每生成一个 token，都要和前面所有 token 做注意力。K 和 V 不重算，缓存起来。每个 token 每层的 cache 大小：

$$2 \times n_h \times d_h \times \text{sizeof(dtype)}$$

DeepSeek-V3 如果按标准 MHA 算（128 头、每头 128 维、BF16）：

$$128 \times 128 \times 2 \times 2 = 65536\ \text{B} = 64\ \text{KiB}$$

每 token 每层 64 KiB。61 层就是约 3.9 MiB 每 token，128K 上下文要 500 GB。不可能。

GQA / MQA：少存几份

GQA 让多个 query 头共享一组 KV 头。8 组就是 8 KiB，省 8 倍。MQA 更极端，所有 query 头共享一组 KV，只剩 128 维的 K 和 V，省 128 倍。

代价是表达能力下降。DeepSeek 想要 MHA 的质量、又要 MQA 的 cache，于是有了 MLA。

---

六、MLA：把 KV 压成潜向量

MLA（Multi-head Latent Attention）的核心不是 kernel 技巧，是一组线性代数恒等式。公式来自 DeepSeek-V2 论文（式 9-19），V3 沿用。

压缩

对输入 $h_t$，先压到一个低维潜向量：

$$c_t^{KV} = W^{DKV} h_t, \qquad c_t^{KV} \in \mathbb{R}^{d_c}$$

再上投影出 K 和 V：

$$k_t^C = W^{UK} c_t^{KV}, \qquad v_t^C = W^{UV} c_t^{KV}$$

$W^{UK}, W^{UV}$ 的形状是 $d_h n_h \times d_c$。DeepSeek-V2 取 $d_c = 4 d_h = 512$，而 $d_h n_h = 128 \times 128 = 16384$。也就是说，KV cache 里只需要存 512 维的 $c_t^{KV}$，不需要存 16384 维的 K 和 V。

位置编码必须解耦

RoPE 是旋转位置编码，作用在每个头的 K 上。如果对压缩后的 $c_t^{KV}$ 应用 RoPE，再上投影，那么 $W^{UK}$ 就会变成位置相关的矩阵——不同 token 的权重不一样，后面的吸收技巧就失效了。

所以 MLA 单独留了一路解耦的 RoPE 维度：

$$k_t^R = \text{RoPE}(W^{KR} h_t), \qquad k_t^R \in \mathbb{R}^{d_h^R}$$

每个头的最终 K 是拼接：

$$k_{t,i} = [k_{t,i}^C;\ k_t^R]$$

$d_h^R = d_h / 2 = 64$，所有头共享同一个 $k_t^R$。cache 里除 $c_t^{KV}$ 外，还要存这 64 维。

吸收：把上投影折进别的矩阵

注意力分数是 $q_{t,i}^\top k_{j,i}$。展开 K 的部分：

$$q_{t,i}^{C\top} k_{j,i}^C = q_{t,i}^{C\top} (W^{UK}_i c_j^{KV}) = (W^{UK\top}_i q_{t,i}^C)^\top c_j^{KV}$$

矩阵乘法满足结合律。$W^{UK}_i$ 可以提前吸收进 query 侧：为每个头预计算一个「吸收后的 query」$\tilde q_{t,i} = W^{UK\top}_i q_{t,i}^C$，之后直接用 $\tilde q_{t,i}^\top c_j^{KV}$ 打分。

同理，$W^{UV}$ 可以吸收进输出投影 $W^O$：

$$u_t = W^O [o_{t,1}; \dots; o_{t,n_h}], \qquad o_{t,i} = \sum_j p_{ij} W^{UV}_i c_j^{KV}$$

把 $W^{UV}_i$ 和 $W^O$ 的对应块乘在一起，注意力里就只需要对 $c_j^{KV}$ 加权求和。

结果：decode 变成 MQA

吸收之后，每个头的注意力变成：

$$\text{score}_{ij} = \tilde q_{t,i}^\top c_j^{KV} + q_{t,i}^{R\top} k_j^R$$

所有头共享同一份 $c_j^{KV}$ 和 $k_j^R$，只是 query 不同。这正是 MQA 的形状：128 个 query 头，1 个 KV 头。

这就是为什么 FlashMLA 的 head_dim_k 是 576（512 潜向量 + 64 RoPE），head_dim_v 是 512（潜向量本身），而不是常见的 128。也是为什么 k_cache 的 num_heads_k 必须为 1（flash_mla_interface.py:81-82）。

cache 账本

方案	每 token 每层	相对 MHA
MHA（128 头 × 128 维，BF16）	64 KiB	1×
GQA（8 组）	4 KiB	1/16
MQA（1 组）	512 B	1/128
MLA（512 + 64，BF16）	1.125 KiB	1/57
MLA（V4.1，FP8，528 B/token）	528 B	1/124

省下来的显存直接换成 batch size，而 decode 的吞吐几乎线性依赖 batch。这是 MLA 真正的收益来源。

图 5 · 压缩后 cache 只存 512 维潜向量 + 64 维 RoPE；吸收让 decode 形状退化成 128 Q 头 / 1 KV 头的 MQA。

---

七、FlashMLA：带宽和算力分别榨干

先算清楚瓶颈

FlashMLA 的 deep dive（docs/20250422-new-kernel-deep-dive.md:9-13）给了一个很干净的分析。decode 时每请求 $s_q$ 个 query token（不开投机解码就是 1），$h_q$ 个 query 头，KV 长度 $s_k \gg h_q s_q$。

FLOPs：

$$2(h_q s_q \cdot d_k \cdot s_k + h_q s_q \cdot s_k \cdot d_v) = 2 h_q s_q s_k (d_k + d_v)$$

访存（BF16；MLA 里 K 和 V 是同一份潜向量，只读一遍，忽略 query）：

$$2 \times s_k d_k \quad \text{字节}$$

算术强度：

$$\frac{2 h_q s_q s_k (d_k + d_v)}{2 s_k d_k} = h_q s_q \cdot \frac{d_k + d_v}{d_k} \approx 2 h_q s_q$$

H800 的机器平衡点（考虑降频后实际算力约 865 TFLOPS）：

$$\frac{865\ \text{TFLOPS}}{3.35\ \text{TB/s}} \approx 258\ \text{FLOP/Byte}$$

所以 $h_q s_q \ge 128$ 时是算力受限，否则是带宽受限。DeepSeek 线上推理的 decode 实例不做张量并行，$h_q = 128$，正好落在算力受限区。

这解释了 FlashMLA 那组看起来矛盾的数字：

• 3000 GB/s：小 batch、长上下文时，bandwidth-bound，能把 HBM 跑到 90%；
• 580 TFLOPS：大 batch 时，compute-bound，能到降频后峰值的 67%；
• 2025 年 4 月的更新把 compute-bound 场景推到 660 TFLOPS（约 76%）。

数字与推导见 FlashMLA deep dive 第 3-13 行。

同一个 kernel，两种工况，优化的重点完全不同。

图 6 · 横轴是算术强度，纵轴是可达 TFLOPS。小 batch 贴着带宽斜线，大 batch 贴着算力水平线。

split-KV：把长序列切开

decode 时一个请求只有一个 query token，却要和 $s_k$ 个 KV token 做注意力。如果让一个 CTA 串行扫完整个 KV，SM 利用率会很低——$s_k$ 很长时只有一个 CTA 在干活。

FlashMLA 的做法是 split-KV：把 KV 切成若干段，每段一个 CTA 并行算部分注意力（各自维护 running max、running sum、部分输出），最后用一个 combine kernel 合并。

合并的公式还是 online softmax 的拼接：

$$m = \max(m_1, m_2), \quad l = l_1 e^{m_1 - m} + l_2 e^{m_2 - m}, \quad o = \frac{o_1 l_1 e^{m_1 - m} + o_2 l_2 e^{m_2 - m}}{l}$$

get_mla_metadata() 返回的 tile_scheduler_metadata 就是给这件事用的：它在第一次调用时生成每个 SM 处理哪些请求、哪些 KV 块、切几份，后续所有层、所有 step 复用。接口注释里特意警告：形状或 topk_length 变了还复用旧 metadata，会静默用错切分（flash_mla_interface.py:88-92）。

seesaw 调度：一个输出矩阵，两个 warpgroup

这是 2025 年 4 月那版 kernel 最漂亮的部分。

Hopper 上 attention 的常规做法是 FlashAttention-3 的 ping-pong：两个 warpgroup 各持有一个输出矩阵，一个算 GEMM 时另一个做 softmax，交替推进。但 MLA 的输出矩阵太大：

$$64 \times 512 \times 4\ \text{B} = 131072\ \text{B} = 32768\ \text{个 32-bit 寄存器}$$

一个 SM 总共 65536 个寄存器。一个输出就占一半，两个根本放不下。ping-pong 直接不可行。

FlashMLA 的解法是把输出矩阵竖着劈成两半：$O_L$ 和 $O_R$，各 $64 \times 256$。$O_L$ 放 warpgroup 0，$O_R$ 放 warpgroup 1。然后把两个 KV block（$K_0, K_1$）的处理交错开，让同一时刻一个 warpgroup 在用 Tensor Core、另一个在用 CUDA Core 做 softmax。

调度骨架（原文编号 0-11，下面按 WG0 / WG1 重新表述，docs/20250422-new-kernel-deep-dive.md:24-42）：

1. WG0 算 $p_0 = qK_0^\top$；
2. WG1 算 $p_1 = qK_1^\top$；
3. WG0 求 $p_0$ 的行最大值，更新全局 $m$，算 $\text{scale}_0 = e^{m_{\text{new}} - m}$；
4. WG0 对 $p_0$ 做 softmax：$p_0 \leftarrow e^{p_0 - m_{\text{new}}}$；
5. WG0 更新 $O_L \leftarrow O_L \cdot \text{scale}_0 + p_0 V_{0L}$；
6. WG1 求 $p_1$ 的最大值，更新全局 $m$，算 $\text{scale}_1$；
7. WG1 对 $p_1$ 做 softmax；
8. WG1 更新 $O_R \leftarrow O_R \cdot (\text{scale}_0 \cdot \text{scale}_1) + p_1 V_{1R}$；
9. WG0 把 $p_0$ 乘上 $\text{scale}_1$；
10. WG1 补上 $p_0 V_{0R}$ 到 $O_R$；
11. WG0 用 $p_1 V_{1L}$ 更新 $O_L$，并乘上 $\text{scale}_1$。

数学上等价于标准 online softmax。妙处在于：两个 warpgroup 交替使用 Tensor Core 和 CUDA Core，谁都不闲着；而且一旦某块 K/V 用完，TMA 立刻可以发起下一块，访存和计算也重叠上了。

这被作者叫做 seesaw scheduling——跷跷板。一个输出矩阵，两头轮流压。

图 7 · O 劈成 O_L / O_R 后，两个 warpgroup 交替占用 Tensor Core 与 CUDA Core；缩放因子必须交叉作用。

剩下三个优化

• 细粒度 TMA 流水：一个 $64 \times 576$ 的 K block 拆成 9 次 TMA（每次 $64 \times 64$），第一块到就开始算第一段 GEMM，不等整块。这是把「块级流水」细化到「子块级」。
• cache hint：TMA 拷贝加 EVICT_FIRST，让一次性使用的 K/V 不要污染 L2，提高其他数据的命中率。
• PDL：splitkv_mla 和 combine 两个 kernel 用 Programmatic Dependent Launch 重叠，combine 不必等 split 完全退出。

配合 tile scheduler 做 SM 间负载均衡，最终在 H800 上拿到 80% 的 Tensor Core 利用率。

---

八、FP8 KV cache：反量化成了新瓶颈

2025 年 9 月 DeepSeek-V3.2 把上下文从 64K 拉到 128K，KV cache 压力剧增，于是引入 FP8 KV cache。

格式（docs/20250929-hopper-fp8-sparse-deep-dive.md:9）：

• 前 512 维（潜向量）按 $1 \times 128$ 分块量化：512 个 e4m3 + 4 个 FP32 scale；
• 后 64 维（RoPE）对精度敏感，不量化，保留 BF16。

每 token 656 字节。格式与 cycle 计算见 FP8 解码 deep dive 第 9-25 行。

反量化在 CUDA Core 上做，代价是：

$$\left(\frac{1}{64} + \frac{1}{64} + \frac{1}{16} + \frac{1}{256}\right) \times 512 \approx 50\ \text{cycle}$$

而对应的 MMA 只需要约 34 cycle。反量化比矩阵乘还慢，kernel 从 compute-bound 变成了 dequantization-bound。

解法很妙，叫 crossover：既然是 MQA，同一个 query token 的所有头看的是同一份 K/V。把两个 CTA 组成一个 cluster，每个 CTA 只反量化一半 K/V，然后用 Hopper 的 Distributed Shared Memory（st.async）把结果写进对方的 shared memory。每个 CTA 只付一半反量化的代价，拿到完整数据。

410 TFLOPS，比不用 crossover 的 250 TFLOPS 提升 64%。

图 8 · 656 B 的 token 布局；反量化 50 cycle 超过 MMA 34 cycle，crossover 用 DSM 把每侧代价减半。

到 V4.1，格式又变了：528 字节/token，512 个 FP8 + 16 个 UE8M0 scale，连 RoPE 部分也量化了。还多了 FP4 选项：288 字节/token，512 个 e2m1（两位一个字节）+ 32 个 e4m3 scale（flash_mla_interface.py:114-117）。

当前 master 已经是稀疏注意力

现在的 FlashMLA 主分支实现的是 DeepSeek Sparse Attention（DSA）：用 indices 给出每个 query 要看的 top-k token，不再扫全量 KV。indices 的编码是 页块索引 * block_size + 块内偏移，无效值必须是 -1——注释明确说 kernel 不做上界检查，越界值会直接导致 TMA 非法访存（flash_mla_interface.py:99-104）。

block_table 和 cache_seqlens 参数还在，但注释写着 “currently ignored”，纯粹为了兼容旧接口。读代码时别被它们误导。

版本漂移提醒

时间	版本特征	代表性能
2025-02	初版，BF16/FP16，paged KV，block 64，Hopper	3000 GB/s，580 TFLOPS（H800）
2025-04	seesaw 调度 + TMA 细粒度流水	660 TFLOPS（H800）
2025-09	FP8 KV cache + 稀疏注意力（V3.2）	640 TFLOPS prefill / 410 decode（H800）
2026-09	V4.1，全量化 KV，SM100 / 昇腾 950，融合 norm-RoPE-attn-RoPE-cast	1350 TFLOPS prefill / 1024 decode（B200）

你 clone 下来的 master 是最后一行。想复现文章里讨论的经典 kernel，得 checkout 到对应 tag 或 commit。

---

九、前端能拿走什么

同一套思维模型

GPU 概念	前端类比	差异
global memory	ArrayBuffer / storage buffer	带宽差两个数量级
shared memory	手动管理的 scratch buffer	16 KB（WebGPU）vs 228 KB（Hopper）
register	局部变量	寄存器压力决定 occupancy
warp / warpgroup	subgroup	32 / 128 线程
TMA	copyBufferToBuffer + 异步	粒度、维度描述能力差很多
tensor core	无（核心 WebGPU 没有矩阵指令）	WebNN 只在 API 层暴露矩阵/卷积算子
kernel 融合	减少 pass	同样的 IO 逻辑

WebGPU 的 var<workgroup> 就是 shared memory，workgroup_size 就是 block 维度。你写一个 tiled matmul，用 workgroupBarrier() 同步，代码结构和 CUDA 几乎一一对应——只是没有 Tensor Core，吞吐差三四个数量级。

三个可迁移的原则

第一，先算算术强度，再谈优化。 一个操作是算力受限还是带宽受限，决定了你该压 FLOP 还是压 IO。FlashMLA 的 3000 GB/s 和 660 TFLOPS 是两套完全不同的优化路径，硬套会南辕北辙。

第二，布局比算法重要。 DeepGEMM 的 scale 要 TMA-aligned + transposed；FlashMLA 的 KV cache 按 token 连续排列，528 字节里数据在前、scale 在后；MLA 的 $W^{UK}$ 能不能吸收，取决于 RoPE 有没有解耦。这些都是数据布局层面的决策，一旦定错，kernel 再优化也救不回来。

第三，异步和流水线是性能的上限。 朴素的实现是「读-算-写」串行，硬件大部分时间在等。DeepGEMM 用多 stage + 双 barrier，FlashMLA 用 9 段 TMA + seesaw 交错，本质都是让搬运、计算、softmax 三个不同的执行单元同时忙起来。

代码阅读路线

如果只读一个文件：

1. DeepGEMM/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh
   从第 39 行看函数签名，第 171 行看 TMA warp，第 250 行看 math warpgroup，第 316 行看 promotion。这是完整的 200 行流水线，没有花活。
2. FlashMLA/docs/20250422-new-kernel-deep-dive.md
   先读理论分析（第 9-13 行），再读 seesaw 的 11 步（第 24-42 行）。这篇文档比代码好懂，代码可以对照着看。
3. FlashMLA/docs/20250929-hopper-fp8-sparse-deep-dive.md
   看反量化瓶颈的 cycle 计算和 crossover 方案，理解「数据格式变化如何改变 kernel 的瓶颈类型」。
4. DeepGEMM/docs/scaling-factor-format.md
   看似枯燥，但解释了为什么一个 scale 张量要满足一堆 stride 约束——TMA 的对齐要求会一路传导到 Python API。

---

十、数字速查

指标	数值	出处
H800 SXM5 BF16 峰值	989 TFLOPS	NVIDIA 规格
H800 SXM5 FP8 峰值	1979 TFLOPS	NVIDIA 规格
H800 HBM 带宽	3.35 TB/s	NVIDIA 规格
降频后实际 BF16 算力	~865 TFLOPS	FlashMLA deep dive
DeepGEMM FP8 GEMM（H800）	1550 TFLOPS（78%）	DeepGEMM README
FlashMLA decode，带宽受限	3000 GB/s（90%）	初版 README / deep dive
FlashMLA decode，算力受限	580 → 660 TFLOPS	FlashMLA deep dive
MLA 每 token 每层 cache（BF16）	576 × 2 = 1152 B	512 潜向量 + 64 RoPE
MLA 每 token 每层 cache（V3.2 FP8）	656 B	512 fp8 + 4 fp32 + 64 bf16
MLA 每 token 每层 cache（V4.1 FP8）	528 B	512 fp8 + 16 ue8m0
单个 $64 \times 512$ 输出占寄存器	32768	32768 / 65536 = 半个 SM
decode 算术强度	$\approx 2 h_q s_q$	128 头时进入算力受限

---

参考资料

本地固定版本（本文所有行号以此为准）：

• ~/msyGIT/DeepGEMM：v2.8.2，commit 5081f8513af57e61696de4f4f04b0e707ffb503e
• ~/msyGIT/FlashMLA：commit 0f11e178da1f67becb087a324b400854931ec3e3

逐条核对表：文中数字 → 原始出处

文中位置	结论	原文链接	本地文件
一、算术强度	朴素 ~2、分块 ~32 FLOP/B	FlashMLA deep dive L9-L13	FlashMLA/docs/20250422-new-kernel-deep-dive.md:9
一、机器平衡点	989 TFLOPS ÷ 3.35 TB/s ≈ 295	NVIDIA H800 规格	—
二、kernel 结构	TMA 128 线程、BLOCK_K=128	sm90_fp8_gemm_1d1d.cuh L39-L51	DeepGEMM/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh:39
二、寄存器分配	24 / 240、PDL	同文件 L153-L157	同文件 :153-157
二、持久化调度	get_next_block()	scheduler/gemm.cuh L185	DeepGEMM/deep_gemm/include/deep_gemm/scheduler/gemm.cuh:185
三、WGMMA + promotion	两级累加代码	sm90_fp8_gemm_1d1d.cuh L302-L325	同文件 :302-325
三、scale 布局	TMA-aligned + transposed、UE8M0	scaling-factor-format.md	DeepGEMM/docs/scaling-factor-format.md
三、1550 TFLOPS	H800 实测	DeepGEMM README L25	DeepGEMM/README.md:25
五、FlashAttention	在线 softmax、IO 算法	FlashAttention 论文	—
六、MLA 公式	式 (9)-(19)：压缩、解耦 RoPE、吸收	DeepSeek-V2 论文 · ar5iv 全文	—
七、3000 GB/s / 580 TFLOPS	初版 H800 数字	初版 README（tag b31bfe7）	git 历史
七、580 → 660	2025-04 kernel 优化	deep dive L3	FlashMLA/docs/20250422-new-kernel-deep-dive.md:3
七、seesaw 调度	11 步、寄存器占用	deep dive L24-L42	同文件 :24-42
七、9 段 TMA / PDL	细粒度流水、EVICT_FIRST	deep dive L54-L61	同文件 :54-61
八、656 B 布局	V3.2 FP8 KV 格式	sparse deep dive L9	FlashMLA/docs/20250929-hopper-fp8-sparse-deep-dive.md:9
八、50 vs 34 cycle	反量化成为瓶颈	sparse deep dive L25	同文件 :25
八、crossover / 410 TFLOPS	DSM 交换、提升 64%	sparse deep dive L31-L48	同文件 :31-48
八、V4.1 528 B / 288 B	全量化、FP4 格式	flash_mla_interface.py L113-L117	FlashMLA/flash_mla/flash_mla_interface.py:113
八、稀疏 indices 编码	页块 × block_size + offset，-1 无效	flash_mla_interface.py L99-L104	同文件 :99-104
八、block_table 被忽略	currently ignored	flash_mla_interface.py L85	同文件 :85
图 8、token 布局	512 + 4 + 64 = 656 B	sparse deep dive L9	FlashMLA/docs/20250929-hopper-fp8-sparse-deep-dive.md:9

论文与官方文档

• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
• DeepSeek-V3 Technical Report
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
• NVIDIA Hopper Architecture In-Depth
• NVIDIA H800 产品页

自己复核

cd ~/msyGIT/DeepGEMM && git checkout 5081f8513af57e61696de4f4f04b0e707ffb503e
cd ~/msyGIT/FlashMLA && git checkout 0f11e178da1f67becb087a324b400854931ec3e3

# 例：确认 promotion 确实发生在每个 K 块之后
sed -n '296,326p' ~/msyGIT/DeepGEMM/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh

# 例：确认 656 B 的构成
sed -n '9,25p' ~/msyGIT/FlashMLA/docs/20250929-hopper-fp8-sparse-deep-dive.md