DeppInfer系列之 - 矩阵、注意力与 GPU：拆解 DeepGEMM 与 FlashMLA

LLM 推理的瓶颈：矩阵乘

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

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

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

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

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

了解更多我的Deepseek解析文章点这里

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

先统一记号。DeepGEMM 的约定是
$$D=C+ABD\; =\; C\; +\; A\; B$$ D=C+AB

（Python 里写作 A @ B。）其中 $AA$ A 是 $M\times KM\; \backslash times\; K$ M×K， $BB$ B 是 $K\times NK\; \backslash times\; N$ K×N， $DD$ D 是 $M\times NM\; \backslash times\; N$ M×N。默认 NT 布局：A 行主序（normal），B 以转置形式存放（transposed），也就是计算 $A\cdot B\top A\; \backslash cdot\; B^\backslash top$ A⋅B⊤。这是推理里最常见的形状------权重按 [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;
    }
  }
}

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

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

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

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

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

把输出切成 $T\times TT\; \backslash times\; T$ T×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 TT\; \backslash times\; T$ T×T 输出块时，A 的 $TT$ T 行、B 的 $TT$ 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 1012/3.35\times 1012\approx 295989\; \backslash times\; 10^\{12\}\; /\; 3.35\; \backslash times\; 10^\{12\}\; \backslash approx\; 295$ 989×1012/3.35×1012≈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 N64\; \backslash times\; N$ 64×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+2i+2$ i+2 块时，Tensor Core 在算第 $ii$ i 块，CUDA Core 在处理第 $i-1i-1$ 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-42^\{-4\}$ 2−4。整层共用一个 scale 会把小值全部压成 0，所以必须分块量化。

DeepSeek-V3 的训练配方：

• 激活：每 token、每 128 个通道一个 scale（ $1\times 1281\; \backslash times\; 128$ 1×128）；
• 权重：每 $128\times 128128\; \backslash times\; 128$ 128×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 128128\; \backslash times\; 128$ 128×128 配方。

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

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

正确做法是：每个 K 块单独算出一段部分和，乘上 $scaleA\times scaleB\; \backslash text\{scale\}\_A\; \backslash times\; \backslash text\{scale\}\_B$ scaleA×scaleB，再加进最终累加器。这就是 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

公式

单头注意力：
$$Attention(Q,K,V)=softmax\; (\; QK\top d\; )\; V\backslash text\{Attention\}(Q,K,V)\; =\; \backslash text\{softmax\}\backslash left(\backslash frac\{QK^\backslash top\}\{\backslash sqrt\{d\}\}\backslash right)V$$ Attention(Q,K,V)=softmax(d QK⊤)V

拆开看：

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

多头就是并行做 $hh$ 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 的大小是 $sq\times sk\; s\_q\; \backslash times\; s\_k$ sq×sk。decode 时 $sk\; s\_k$ sk 可能是 128K，每个头每个请求都要存一份，显存直接爆掉。

FlashAttention 的观察是：softmax 和 $PVPV$ PV 可以边算边合并。维护三个数------running max $mm$ m、running sum $ll$ l、部分输出 $acc\backslash text\{acc\}$ 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(sqsk)\; O(s\_q\; s\_k)$ O(sqsk) 降到 $O(d)O(d)$ O(d)。数学结果和朴素实现完全一致，只是改变了计算顺序。这就是 FlashAttention 的核心：它不是近似算法，是 IO 算法。

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

KV cache 的账

自回归 decode 时，每生成一个 token，都要和前面所有 token 做注意力。K 和 V 不重算，缓存起来。每个 token 每层的 cache 大小：
$$2\times nh\times dh\times sizeof(dtype)\; 2\; \backslash times\; n\_h\; \backslash times\; d\_h\; \backslash times\; \backslash text\{sizeof(dtype)\}$$ 2×nh×dh×sizeof(dtype)

DeepSeek-V3 如果按标准 MHA 算（128 头、每头 128 维、BF16）：
$$128\times 128\times 2\times 2=65536\; B=64\; KiB128\; \backslash times\; 128\; \backslash times\; 2\; \backslash times\; 2\; =\; 65536\backslash \; \backslash text\{B\}\; =\; 64\backslash \; \backslash text\{KiB\}$$ 128×128×2×2=65536 B=64 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 沿用。

压缩

对输入 $ht\; h\_t$ ht，先压到一个低维潜向量：
$$ctKV\; =WDKVht,\; ctKV\; \in Rdc\; c\_t^\{KV\}\; =\; W^\{DKV\}\; h\_t,\; \backslash qquad\; c\_t^\{KV\}\; \backslash in\; \backslash mathbb\{R\}^\{d\_c\}$$ ctKV=WDKVht,ctKV∈Rdc

再上投影出 K 和 V：
$$ktC=WUK\; ctKV\; ,vtC=WUV\; ctKV\; k\_t^C\; =\; W^\{UK\}\; c\_t^\{KV\},\; \backslash qquad\; v\_t^C\; =\; W^\{UV\}\; c\_t^\{KV\}$$ ktC=WUKctKV,vtC=WUVctKV

$WUK,WUVW^\{UK\},\; W^\{UV\}$ WUK,WUV 的形状是 $dhnh\times dc\; d\_h\; n\_h\; \backslash times\; d\_c$ dhnh×dc。DeepSeek-V2 取 $dc=4dh=512\; d\_c\; =\; 4\; d\_h\; =\; 512$ dc=4dh=512，而 $dhnh=128\times 128=16384\; d\_h\; n\_h\; =\; 128\; \backslash times\; 128\; =\; 16384$ dhnh=128×128=16384。也就是说，KV cache 里只需要存 512 维的 $ctKV\; c\_t^\{KV\}$ ctKV，不需要存 16384 维的 K 和 V。

位置编码必须解耦

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

所以 MLA 单独留了一路解耦的 RoPE 维度：
$$ktR=RoPE(WKRht),ktR\in RdhR\; k\_t^R\; =\; \backslash text\{RoPE\}(W^\{KR\}\; h\_t),\; \backslash qquad\; k\_t^R\; \backslash in\; \backslash mathbb\{R\}^\{d\_h^R\}$$ ktR=RoPE(WKRht),ktR∈RdhR

每个头的最终 K 是拼接：
$$kt,i\; =$$ kt,iC ; ktR k_{t,i} = k_{t,i}\^C;\\ k_t\^R kt,i=kt,iC; ktR

$dhR=dh/2=64\; d\_h^R\; =\; d\_h\; /\; 2\; =\; 64$ dhR=dh/2=64，所有头共享同一个 $ktR\; k\_t^R$ ktR。cache 里除 $ctKV\; c\_t^\{KV\}$ ctKV 外，还要存这 64 维。

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

注意力分数是 $qt,i\top \; kj,i\; q\_\{t,i\}^\backslash top\; k\_\{j,i\}$ qt,i⊤kj,i。展开 K 的部分：
$$qt,iC\top \; kj,iC\; =\; qt,iC\top \; (\; WiUK\; cjKV\; )=(\; WiUK\top \; qt,iC\; )\top \; cjKV\; q\_\{t,i\}^\{C\backslash top\}\; k\_\{j,i\}^C\; =\; q\_\{t,i\}^\{C\backslash top\}\; (W^\{UK\}\_i\; c\_j^\{KV\})\; =\; (W^\{UK\backslash top\}$$i q{t,i}^C)^\top c_j^{KV} qt,iC⊤kj,iC=qt,iC⊤(WiUKcjKV)=(WiUK⊤qt,iC)⊤cjKV

矩阵乘法满足结合律。 $WiUK\; W^\{UK\}$i WiUK 可以提前吸收进 query 侧：为每个头预计算一个「吸收后的 query」 $q~t,i\; =\; WiUK\top \; qt,iC\; \backslash tilde\; q${t,i} = W^{UK\top}i q{t,i}^C q~t,i=WiUK⊤qt,iC，之后直接用 $q~t,i\top \; cjKV\; \backslash tilde\; q\_\{t,i\}^\backslash top\; c\_j^\{KV\}$ q~t,i⊤cjKV 打分。

同理， $WUVW^\{UV\}$ WUV 可以吸收进输出投影 $WOW^O$ WO：
$$ut=WO$$ ot,1 ;... ; o t,nh , ot,i =∑j pij WiUV cjKV 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} ut=WOot,1;...;ot,nh,ot,i=j∑pijWiUVcjKV

把 $WiUV\; W^\{UV\}\_i$ WiUV 和 $WOW^O$ WO 的对应块乘在一起，注意力里就只需要对 $cjKV\; c\_j^\{KV\}$ cjKV 加权求和。

结果：decode 变成 MQA

吸收之后，每个头的注意力变成：
$$scoreij\; =\; q~t,i\top \; cjKV\; +\; qt,iR\top \; kjR\; \backslash text\{score\}$${ij} = \tilde q{t,i}^\top c_j^{KV} + q_{t,i}^{R\top} k_j^R scoreij=q~t,i⊤cjKV+qt,iR⊤kjR

所有头共享同一份 $cjKV\; c\_j^\{KV\}$ cjKV 和 $kjR\; k\_j^R$ kjR，只是 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 时每请求 $sq\; s\_q$ sq 个 query token（不开投机解码就是 1）， $hq\; h\_q$ hq 个 query 头，KV 长度 $sk\gg hqsq\; s\_k\; \backslash gg\; h\_q\; s\_q$ sk≫hqsq。

FLOPs：
$$2(hqsq\cdot dk\cdot sk+hqsq\cdot sk\cdot dv)=2hqsqsk(dk+dv)\; 2(h\_q\; s\_q\; \backslash cdot\; d\_k\; \backslash cdot\; s\_k\; +\; h\_q\; s\_q\; \backslash cdot\; s\_k\; \backslash cdot\; d\_v)\; =\; 2\; h\_q\; s\_q\; s\_k\; (d\_k\; +\; d\_v)$$ 2(hqsq⋅dk⋅sk+hqsq⋅sk⋅dv)=2hqsqsk(dk+dv)

访存（BF16；MLA 里 K 和 V 是同一份潜向量，只读一遍，忽略 query）：
$$2\times skdk字节\; 2\; \backslash times\; s\_k\; d\_k\; \backslash quad\; \backslash text\{字节\}$$ 2×skdk字节

算术强度：
$$2hqsqsk(dk+dv)\; 2skdk\; =hqsq\cdot \; dk+dv\; dk\; \approx 2hqsq\; \backslash frac\{2\; h\_q\; s\_q\; s\_k\; (d\_k\; +\; d\_v)\}\{2\; s\_k\; d\_k\}\; =\; h\_q\; s\_q\; \backslash cdot\; \backslash frac\{d\_k\; +\; d\_v\}\{d\_k\}\; \backslash approx\; 2\; h\_q\; s\_q$$ 2skdk2hqsqsk(dk+dv)=hqsq⋅dkdk+dv≈2hqsq

H800 的机器平衡点（考虑降频后实际算力约 865 TFLOPS）：
$$865\; TFLOPS3.35\; TB/s\; \approx 258\; FLOP/Byte\; \backslash frac\{865\backslash \; \backslash text\{TFLOPS\}\}\{3.35\backslash \; \backslash text\{TB/s\}\}\; \backslash approx\; 258\backslash \; \backslash text\{FLOP/Byte\}$$ 3.35 TB/s865 TFLOPS≈258 FLOP/Byte

所以 $hqsq\ge 128\; h\_q\; s\_q\; \backslash ge\; 128$ hqsq≥128 时是算力受限，否则是带宽受限。DeepSeek 线上推理的 decode 实例不做张量并行， $hq=128\; h\_q\; =\; 128$ hq=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，却要和 $sk\; s\_k$ sk 个 KV token 做注意力。如果让一个 CTA 串行扫完整个 KV，SM 利用率会很低------ $sk\; s\_k$ sk 很长时只有一个 CTA 在干活。

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

合并的公式还是 online softmax 的拼接：
$$m=max(m1,m2),l=l1e\; m1-m\; +l2e\; m2-m\; ,o=\; o1l1e\; m1-m\; +o2l2e\; m2-m\; l\; m\; =\; \backslash max(m\_1,\; m\_2),\; \backslash quad\; l\; =\; l\_1\; e^\{m\_1\; -\; m\}\; +\; l\_2\; e^\{m\_2\; -\; m\},\; \backslash quad\; o\; =\; \backslash frac\{o\_1\; l\_1\; e^\{m\_1\; -\; m\}\; +\; o\_2\; l\_2\; e^\{m\_2\; -\; m\}\}\{l\}$$ m=max(m1,m2),l=l1em1−m+l2em2−m,o=lo1l1em1−m+o2l2em2−m

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\; B=131072\; B=32768\; 个\; 32-bit\; 寄存器64\; \backslash times\; 512\; \backslash times\; 4\backslash \; \backslash text\{B\}\; =\; 131072\backslash \; \backslash text\{B\}\; =\; 32768\backslash \; \backslash text\{个\; 32-bit\; 寄存器\}$$ 64×512×4 B=131072 B=32768 个 32-bit 寄存器

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

FlashMLA 的解法是把输出矩阵竖着劈成两半 ： $OL\; O\_L$ OL 和 $OR\; O\_R$ OR，各 $64\times 25664\; \backslash times\; 256$ 64×256。 $OL\; O\_L$ OL 放 warpgroup 0， $OR\; O\_R$ OR 放 warpgroup 1。然后把两个 KV block（ $K0,K1\; K\_0,\; K\_1$ K0,K1）的处理交错开，让同一时刻一个 warpgroup 在用 Tensor Core、另一个在用 CUDA Core 做 softmax。

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

1. WG0 算 $p0=qK0\top \; p\_0\; =\; qK\_0^\backslash top$ p0=qK0⊤；
2. WG1 算 $p1=qK1\top \; p\_1\; =\; qK\_1^\backslash top$ p1=qK1⊤；
3. WG0 求 $p0\; p\_0$ p0 的行最大值，更新全局 $mm$ m，算 $scale0=e\; mnew-m\; \backslash text\{scale\}$0 = e^{m{\text{new}} - m} scale0=emnew−m；
4. WG0 对 $p0\; p\_0$ p0 做 softmax： $p0\leftarrow e\; p0-mnew\; p\_0\; \backslash leftarrow\; e^\{p\_0\; -\; m\_\{\backslash text\{new\}\}\}$ p0←ep0−mnew；
5. WG0 更新 $OL\leftarrow OL\cdot scale0+p0\; V0L\; O\_L\; \backslash leftarrow\; O\_L\; \backslash cdot\; \backslash text\{scale\}$0 + p_0 V{0L} OL←OL⋅scale0+p0V0L；
6. WG1 求 $p1\; p\_1$ p1 的最大值，更新全局 $mm$ m，算 $scale1\; \backslash text\{scale\}\_1$ scale1；
7. WG1 对 $p1\; p\_1$ p1 做 softmax；
8. WG1 更新 $OR\leftarrow OR\cdot (scale0\cdot scale1)+p1\; V1R\; O\_R\; \backslash leftarrow\; O\_R\; \backslash cdot\; (\backslash text\{scale\}\_0\; \backslash cdot\; \backslash text\{scale\}$1) + p_1 V{1R} OR←OR⋅(scale0⋅scale1)+p1V1R；
9. WG0 把 $p0\; p\_0$ p0 乘上 $scale1\; \backslash text\{scale\}\_1$ scale1；
10. WG1 补上 $p0\; V0R\; p\_0\; V\_\{0R\}$ p0V0R 到 $OR\; O\_R$ OR；
11. WG0 用 $p1\; V1L\; p\_1\; V\_\{1L\}$ p1V1L 更新 $OL\; O\_L$ OL，并乘上 $scale1\; \backslash text\{scale\}\_1$ scale1。

数学上等价于标准 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 57664\; \backslash times\; 576$ 64×576 的 K block 拆成 9 次 TMA（每次 $64\times 6464\; \backslash times\; 64$ 64×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 1281\; \backslash times\; 128$ 1×128 分块量化：512 个 e4m3 + 4 个 FP32 scale；
• 后 64 维（RoPE）对精度敏感，不量化，保留 BF16。

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

反量化在 CUDA Core 上做，代价是：
$$(164+164+116+1256)\; \times 512\approx 50\; cycle\; \backslash left(\backslash frac\{1\}\{64\}\; +\; \backslash frac\{1\}\{64\}\; +\; \backslash frac\{1\}\{16\}\; +\; \backslash frac\{1\}\{256\}\backslash right)\; \backslash times\; 512\; \backslash approx\; 50\backslash \; \backslash text\{cycle\}$$ (641+641+161+2561)×512≈50 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 的 $WUKW^\{UK\}$ WUK 能不能吸收，取决于 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 51264\; \backslash times\; 512$ 64×512 输出占寄存器	32768	32768 / 65536 = 半个 SM
decode 算术强度	$\approx 2hqsq\; \backslash approx\; 2\; h\_q\; s\_q$ ≈2hqsq	128 头时进入算力受限

---

参考资料

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

• /DeepGEMM：v2.8.2，commit 5081f8513af57e61696de4f4f04b0e707ffb503e
• /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 /DeepGEMM && git checkout 5081f8513af57e61696de4f4f04b0e707ffb503e
cd /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

了解更多我的Deepseek解析文章点这里