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

moshuying
137 阅读4分钟

LLM 推理的瓶颈：矩阵乘

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

前端其实不陌生：CSS matrix3d 是 4×4，canvas 的 transform 是仿射矩阵，图像卷积是滑动窗口的点积。区别只在于规模——GPU 上的 GEMM 是 4096×4096×40964096 \times 4096 \times 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

（Python 里写作 A @ B。）其中 AA 是 M×KM \times K，BB 是 K×NK \times N，DD 是 M×NM \times N。默认 NT 布局：A 行主序（normal），B 以转置形式存放（transposed），也就是计算 A⋅B⊤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;
    }
  }
}

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

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

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

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

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

把输出切成 T×TT \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×TT \times T 输出块时，A 的 TT 行、B 的 TT 列会被反复使用，缓存命中率大幅提升。这叫 cache blocking，GPU 上的「tiling」是同一件事，只是"缓存"换成了片上 shared memory。

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

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

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

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

图 1：GEMM 分块与数据复用

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

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

硬件层次

一张 H800 SXM5：

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

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

三个关键指令

WGMMA（Warpgroup Matrix Multiply-Accumulate）：Hopper 的 Tensor Core 指令，一个 warpgroup 是 128 个线程，一条 WGMMA 可以算出 64×N64 \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+2i+2 块时，Tensor Core 在算第 ii 块，CUDA Core 在处理第 i−1i-1 块的缩放。

另外两个细节值得注意：

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

图 2：Hopper GEMM 的存储层级与流水线

图 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}。整层共用一个 scale 会把小值全部压成 0，所以必须分块量化。

DeepSeek-V3 的训练配方：

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

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

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

正确做法是：每个 K 块单独算出一段部分和，乘上 scaleA×scaleB\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 分块量化与两级累加

图 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_contiguousMoE 权重反向按 K 轴分组，M/N 固定

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

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

公式

单头注意力：

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

拆开看：

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

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

朴素实现（注意 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×sks_q \times s_k。decode 时 sks_k 可能是 128K，每个头每个请求都要存一份，显存直接爆掉。

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

图 4：注意力 = 两个 GEMM + softmax，以及流式合并

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

KV cache 的账

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

2×nh×dh×sizeof(dtype)2 \times n_h \times d_h \times \text{sizeof(dtype)}

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

128×128×2×2=65536 B=64 KiB128 \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 沿用。

压缩

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

ctKV=WDKVht,ctKV∈Rdcc_t^{KV} = W^{DKV} h_t, \qquad c_t^{KV} \in \mathbb{R}^{d_c}

再上投影出 K 和 V：

ktC=WUKctKV,vtC=WUVctKVk_t^C = W^{UK} c_t^{KV}, \qquad v_t^C = W^{UV} c_t^{KV}

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

位置编码必须解耦

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

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

ktR=RoPE(WKRht),ktR∈RdhRk_t^R = \text{RoPE}(W^{KR} h_t), \qquad k_t^R \in \mathbb{R}^{d_h^R}

每个头的最终 K 是拼接：

kt,i=[kt,iC; ktR]k_{t,i} = [k_{t,i}^C;\ k_t^R]

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

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

注意力分数是 qt,i⊤kj,iq_{t,i}^\top k_{j,i}。展开 K 的部分：

qt,iC⊤kj,iC=qt,iC⊤(WiUKcjKV)=(WiUK⊤qt,iC)⊤cjKVq_{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}

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

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

ut=WO[ot,1;… ;ot,nh],ot,i=∑jpijWiUVcjKVu_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}

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

结果：decode 变成 MQA

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

scoreij=q~t,i⊤cjKV+qt,iR⊤kjR\text{score}_{ij} = \tilde q_{t,i}^\top c_j^{KV} + q_{t,i}^{R\top} k_j^R

所有头共享同一份 cjKVc_j^{KV} 和 kjRk_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 KiB1×
GQA（8 组）4 KiB1/16
MQA（1 组）512 B1/128
MLA（512 + 64，BF16）1.125 KiB1/57
MLA（V4.1，FP8，528 B/token）528 B1/124

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

图 5：MLA 的压缩、解耦 RoPE 与吸收

图 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 时每请求 sqs_q 个 query token（不开投机解码就是 1），hqh_q 个 query 头，KV 长度 sk≫hqsqs_k \gg h_q s_q。

FLOPs：

2(hqsq⋅dk⋅sk+hqsq⋅sk⋅dv)=2hqsqsk(dk+dv)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×skdk字节2 \times s_k d_k \quad \text{字节}

算术强度：

2hqsqsk(dk+dv)2skdk=hqsq⋅dk+dvdk≈2hqsq\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）：

865 TFLOPS3.35 TB/s≈258 FLOP/Byte\frac{865\ \text{TFLOPS}}{3.35\ \text{TB/s}} \approx 258\ \text{FLOP/Byte}

所以 hqsq≥128h_q s_q \ge 128 时是算力受限，否则是带宽受限。DeepSeek 线上推理的 decode 实例不做张量并行，hq=128h_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：decode 的 roofline

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

split-KV：把长序列切开

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

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

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

m=max⁡(m1,m2),l=l1em1−m+l2em2−m,o=o1l1em1−m+o2l2em2−mlm = \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×512×4 B=131072 B=32768 个 32-bit 寄存器64 \times 512 \times 4\ \text{B} = 131072\ \text{B} = 32768\ \text{个 32-bit 寄存器}

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

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

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

  1. WG0 算 p0=qK0⊤p_0 = qK_0^\top；
  2. WG1 算 p1=qK1⊤p_1 = qK_1^\top；
  3. WG0 求 p0p_0 的行最大值，更新全局 mm，算 scale0=emnew−m\text{scale}_0 = e^{m_{\text{new}} - m}；
  4. WG0 对 p0p_0 做 softmax：p0←ep0−mnewp_0 \leftarrow e^{p_0 - m_{\text{new}}}；
  5. WG0 更新 OL←OL⋅scale0+p0V0LO_L \leftarrow O_L \cdot \text{scale}_0 + p_0 V_{0L}；
  6. WG1 求 p1p_1 的最大值，更新全局 mm，算 scale1\text{scale}_1；
  7. WG1 对 p1p_1 做 softmax；
  8. WG1 更新 OR←OR⋅(scale0⋅scale1)+p1V1RO_R \leftarrow O_R \cdot (\text{scale}_0 \cdot \text{scale}_1) + p_1 V_{1R}；
  9. WG0 把 p0p_0 乘上 scale1\text{scale}_1；
  10. WG1 补上 p0V0Rp_0 V_{0R} 到 ORO_R；
  11. WG0 用 p1V1Lp_1 V_{1L} 更新 OLO_L，并乘上 scale1\text{scale}_1。

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

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

图 7：seesaw 调度的 warpgroup 时间线

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

剩下三个优化

  • 细粒度 TMA 流水：一个 64×57664 \times 576 的 K block 拆成 9 次 TMA（每次 64×6464 \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×1281 \times 128 分块量化：512 个 e4m3 + 4 个 FP32 scale；
  • 后 64 维（RoPE）对精度敏感，不量化，保留 BF16。

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

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

(164+164+116+1256)×512≈50 cycle\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：FP8 KV cache 布局与 crossover

图 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，Hopper3000 GB/s，580 TFLOPS（H800）
2025-04seesaw 调度 + TMA 细粒度流水660 TFLOPS（H800）
2025-09FP8 KV cache + 稀疏注意力（V3.2）640 TFLOPS prefill / 410 decode（H800）
2026-09V4.1，全量化 KV，SM100 / 昇腾 950，融合 norm-RoPE-attn-RoPE-cast1350 TFLOPS prefill / 1024 decode（B200）

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

九、前端能拿走什么

同一套思维模型

GPU 概念前端类比差异
global memoryArrayBuffer / storage buffer带宽差两个数量级
shared memory手动管理的 scratch buffer16 KB（WebGPU）vs 228 KB（Hopper）
register局部变量寄存器压力决定 occupancy
warp / warpgroupsubgroup32 / 128 线程
TMAcopyBufferToBuffer + 异步粒度、维度描述能力差很多
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} 能不能吸收，取决于 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 TFLOPSNVIDIA 规格
H800 SXM5 FP8 峰值1979 TFLOPSNVIDIA 规格
H800 HBM 带宽3.35 TB/sNVIDIA 规格
降频后实际 BF16 算力~865 TFLOPSFlashMLA deep dive
DeepGEMM FP8 GEMM（H800）1550 TFLOPS（78%）DeepGEMM README
FlashMLA decode，带宽受限3000 GB/s（90%）初版 README / deep dive
FlashMLA decode，算力受限580 → 660 TFLOPSFlashMLA deep dive
MLA 每 token 每层 cache（BF16）576 × 2 = 1152 B512 潜向量 + 64 RoPE
MLA 每 token 每层 cache（V3.2 FP8）656 B512 fp8 + 4 fp32 + 64 bf16
MLA 每 token 每层 cache（V4.1 FP8）528 B512 fp8 + 16 ue8m0
单个 64×51264 \times 512 输出占寄存器3276832768 / 65536 = 半个 SM
decode 算术强度≈2hqsq\approx 2 h_q s_q128 头时进入算力受限

参考资料

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

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

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

文中位置结论原文链接本地文件
一、算术强度朴素 ~2、分块 ~32 FLOP/BFlashMLA deep dive L9-L13FlashMLA/docs/20250422-new-kernel-deep-dive.md:9
一、机器平衡点989 TFLOPS ÷ 3.35 TB/s ≈ 295NVIDIA H800 规格—
二、kernel 结构TMA 128 线程、BLOCK_K=128sm90_fp8_gemm_1d1d.cuh L39-L51DeepGEMM/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 L185DeepGEMM/deep_gemm/include/deep_gemm/scheduler/gemm.cuh:185
三、WGMMA + promotion两级累加代码sm90_fp8_gemm_1d1d.cuh L302-L325同文件 :302-325
三、scale 布局TMA-aligned + transposed、UE8M0scaling-factor-format.mdDeepGEMM/docs/scaling-factor-format.md
三、1550 TFLOPSH800 实测DeepGEMM README L25DeepGEMM/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 → 6602025-04 kernel 优化deep dive L3FlashMLA/docs/20250422-new-kernel-deep-dive.md:3
七、seesaw 调度11 步、寄存器占用deep dive L24-L42同文件 :24-42
七、9 段 TMA / PDL细粒度流水、EVICT_FIRSTdeep dive L54-L61同文件 :54-61
八、656 B 布局V3.2 FP8 KV 格式sparse deep dive L9FlashMLA/docs/20250929-hopper-fp8-sparse-deep-dive.md:9
八、50 vs 34 cycle反量化成为瓶颈sparse deep dive L25同文件 :25
八、crossover / 410 TFLOPSDSM 交换、提升 64%sparse deep dive L31-L48同文件 :31-48
八、V4.1 528 B / 288 B全量化、FP4 格式flash_mla_interface.py L113-L117FlashMLA/flash_mla/flash_mla_interface.py:113
八、稀疏 indices 编码页块 × block_size + offset，-1 无效flash_mla_interface.py L99-L104同文件 :99-104
八、block_table 被忽略currently ignoredflash_mla_interface.py L85同文件 :85
图 8、token 布局512 + 4 + 64 = 656 Bsparse deep dive L9FlashMLA/docs/20250929-hopper-fp8-sparse-deep-dive.md:9

论文与官方文档

自己复核

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解析文章点这里

avatar
文章
阅读
粉丝