DeepSeek V4 源码剖析

第5章 sparse_attn 与 FlashMLA：稀疏注意力落到 GPU 上

作者 杨艺韬 · 8,552 字 · 发布于 · 更新于

o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) —— inference/model.py:528

V4 的全部 attention 革命（MLA + Compressor + Indexer），最后必须落到一个 GPU kernel 上才能产生工程价值。这一章拆两份实现：官方仓库里那份可读的 TileLang 参考实现，和 FlashMLA 里那份为 H800 / B200 手写的 CUDA。

5.1 引子：从一行 PyTorch 到 GPU 上的真实计算

V4 的 Attention.forward 里 attention 的"实际计算"被压缩成一行：

o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)

参数：

  • q：query，形状 [B, S, n_heads=128, head_dim=512]，BF16（kernel 签名里写死 T.Tensor[(b, m, h, d), BF16]）
  • kv：decode 时是 KV cache（滑窗 + 压缩段拼起来的完整 KV），形状 [B, kv_cache_size, head_dim=512]；prefill 时是本段的 kv 与 Compressor 输出拼起来的张量（inference/model.py:524-528）
  • attn_sink：每 head 一个 float32 标量，形状 [n_heads]
  • topk_idxs：稀疏选取的 KV 位置索引，在 ratio=4 的层上形状最多为 [B, S, window_size + index_topk] = [B, S, 1152]（ratio=128 的层是滑窗 128 条加全部 end_pos // 128 条压缩 KV，ratio=0 的层只有滑窗 128 条），类型为 int32（inference/model.py:515 有一行 topk_idxs = topk_idxs.int()），-1 表示无效槽位
  • softmax_scale：缩放系数 head_dim ** -0.5

sparse_attn 是一个从 kernel 模块导入的函数：

from kernel import act_quant, fp4_act_quant, fp8_gemm, fp4_gemm, sparse_attn, hc_split_sinkhorn

这里要先纠正一个很容易想当然的误解：这个 kernel 不是 FlashMLA。它就是 V4 仓库自己的 inference/kernel.py（536 行），和 model.py 平级——model.py 用的是 from kernel import ... 这种同目录导入。打开它会看到，里面没有一行 CUDA C++，全部是 TileLang：

import tilelang
import tilelang.language as T

@tilelang.jit(pass_configs=pass_configs)
def sparse_attn_kernel(h: int, d: int, scale=None):
    ...

也就是说，官方给的是一份用 DSL 写的、可读的参考实现，运行时由 TileLang JIT 编译成 GPU 代码。而 FlashMLA 是另一件事：它是 DeepSeek 单独维护的手写 CUDA 内核库，README 自述"powering the DeepSeek-V3 and DeepSeek-V3.2-Exp models"，V4 发布时它并没有被 inference/ 目录引用。但——FlashMLA 里已经躺着 V4 的代码路径，这是本章后半段要拆的重点。

flowchart TB
  subgraph Ref["官方参考实现（本章 §5.2-§5.5）"]
    P1["Attention.forward:<br/>sparse_attn(q, kv, sink, idxs, scale)"]
    P2["inference/kernel.py:355<br/>sparse_attn(...) 包装"]
    P3["inference/kernel.py:277<br/>@tilelang.jit sparse_attn_kernel"]
    P4["TileLang JIT → PTX/SASS"]
    P1 --> P2 --> P3 --> P4
  end
  subgraph Prod["生产路径（本章 §5.6-§5.9）"]
    F1["FlashMLA csrc/api/api.cpp<br/>sparse_decode_fwd / sparse_prefill_fwd"]
    F2["csrc/sm90/decode/sparse_fp8/<br/>run_flash_splitkv_mla_fp8_sparse_kernel"]
    F3["csrc/sm100/decode/head64/<br/>KernelTemplate&lt;ModelType&gt;<br/>（h_q=128、d_qk=512 时改走<br/>prefill/sparse/fwd_for_small_topk/head128）"]
    F1 --> F2
    F1 --> F3
  end

本章拆这两条链——先读懂 TileLang 那份把语义讲清楚，再去 FlashMLA 里核对生产实现真正长什么样。

5.2 sparse_attn 的接口语义

sparse_attn(q, kv, sink, idxs, scale) 的语义是：

对每个 query token 的每个 head，仅计算 q · kv[idxs] 的内积，做 softmax + 加权求和，然后跨 head 输出。

伪代码形式：

def sparse_attn(q, kv, attn_sink, topk_idxs, softmax_scale):
    # q: [B, S, H, D]
    # kv: [B, T, D]   (T = window_size + max_seq_len // ratio)
    # attn_sink: [H]
    # topk_idxs: [B, S, K]   (K = window_size + index_topk = 1152)
    # softmax_scale: scalar
    # 返回 o: [B, S, H, D]

    B, S, H, D = q.shape
    K = topk_idxs.shape[-1]
    o = torch.zeros_like(q)

    for b in range(B):
        for s in range(S):
            for h in range(H):
                # gather K 个 KV 位置
                gathered_kv = kv[b, topk_idxs[b, s]]   # [K, D]
                # mask 掉 -1 位置
                valid = topk_idxs[b, s] >= 0
                # logits = q · k / sqrt(D)
                logits = (q[b, s, h] @ gathered_kv.T) * softmax_scale
                logits = torch.where(valid, logits, float("-inf"))
                # 拼接 attn_sink
                logits_with_sink = torch.cat([logits, attn_sink[h:h+1]])
                weights = F.softmax(logits_with_sink, dim=0)
                # 加权求和（不包括 sink）
                o[b, s, h] = (weights[:K, None] * gathered_kv).sum(dim=0)

    return o

这只是语义模型——真实 CUDA 实现要快几个数量级。但理解这个伪代码后，CUDA 实现就是"如何把这段循环并行化、向量化、利用 TensorCore"的工程问题。

5.3 FlashMLA 仓库的代码组织

把 FlashMLA（github.com/deepseek-ai/FlashMLA，本章核对的是 71c7379，2026-03-31，V4 发布前 main 分支上的最后一版）的目录摊开：

FlashMLA/
├── csrc/
│   ├── api/
│   │   ├── api.cpp              # pybind：sparse_decode_fwd / dense_decode_fwd /
│   │   │                        #         sparse_prefill_fwd / dense_prefill_fwd / dense_prefill_bwd
│   │   ├── sparse_decode.h
│   │   └── sparse_fwd.h
│   ├── params.h                 # enum class ModelType { V32, MODEL1 };
│   ├── sm90/
│   │   ├── decode/sparse_fp8/   # splitkv_mla.cuh + instantiations/{model1,v32}_persistent_h{64,128}.cu
│   │   └── prefill/sparse/      # phase1.cuh + instantiations/phase1_k{512,576}[_topklen].cu
│   ├── sm100/
│   │   ├── decode/head64/       # KernelTemplate<ModelType>
│   │   └── prefill/sparse/      # fwd/head{64,128}/ 与 fwd_for_small_topk/head128/
│   └── smxx/decode/combine/     # split-KV 的 combine kernel（attn_sink 在这里参与归一化）
├── flash_mla/                   # Python 包装：get_mla_metadata / flash_mla_with_kvcache /
│                                #              flash_mla_sparse_fwd
├── tests/
└── benchmark/bench_flash_mla.py

三条可以直接 grep 出来的观察：

1. ModelType 只有两个取值：V32 和 MODEL1。

csrc/params.h:5：

enum class ModelType {
    V32,
    MODEL1
};

MODEL1 是什么？看 csrc/sm100/decode/head64/config.h:33 起按 ModelType 分叉的那组常量（节选）：

static constexpr int D_Q    = MODEL_TYPE == ModelType::V32 ? 576 : 512;
static constexpr int D_V    = 512;
static constexpr int D_NOPE = MODEL_TYPE == ModelType::V32 ? 512 : 448;
static constexpr int D_ROPE = 64;
static constexpr int QUANT_TILE_SIZE   = MODEL_TYPE == ModelType::V32 ? 128 : 64;
static constexpr bool V_HAVE_ROPE      = MODEL_TYPE == ModelType::V32 ? false : true;
static constexpr int NUM_SCALES_EACH_TOKEN = MODEL_TYPE == ModelType::V32 ? 4 : 8;

MODEL1 那一列——head_dim 512、nope 448、rope 64、量化块 64、V 也带 RoPE——与 V4 的 config.json（head_dim=512、qk_rope_head_dim=64）和 model.py 里那两行（act_quant(kv[..., :-rd], 64, ...) 的块大小 64、apply_rotary_emb(o[..., -rd:], freqs_cis, True) 的输出反旋转）逐条对上。448/64 = 7 个 scale，padding 到 8，就是 NUM_SCALES_EACH_TOKEN = 8。换句话说：V4 的内核路径在模型发布前近一个月就已经躺在 FlashMLA 里了，只是用了一个不说破的代号。

2. 每个架构一套独立实现，且实例化是按"模型 × head 数"展开的。

csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu 全文只有一行有效代码：

template void run_flash_splitkv_mla_fp8_sparse_kernel<ModelType::MODEL1, 128>(const SparseAttnDecodeParams &params);

同目录下还有 model1_persistent_h64.cu、v32_persistent_h128.cu、v32_persistent_h64.cu——四份组合各编一次。prefill 侧则按 head_dim_k 展开：sm90::fwd::run_fwd_phase1_kernel<576, false> 与 <512, false>，_topklen 后缀那两个对应"每条序列 topk 长度不同"的变体。

3. PyTorch binding 极薄。csrc/api/api.cpp 一共只有五行 m.def(...)，C++ 入口做的是形状 / dtype 检查 + 按架构 dispatch（arch.is_sm90a() / arch.is_sm100f()），业务逻辑全在 kernel 里。

5.4 稀疏 attention 的 GPU 优化挑战

稀疏 attention 比 dense attention 在 GPU 上更难优化。dense attention 的内存访问模式是连续的——每个 query 顺序读 KV cache 从 0 到 T。稀疏 attention 的内存访问模式是索引跳跃的——每个 query 按 topk_idxs 跳着读 KV。

这带来三个挑战：

挑战一：合并访问失效

GPU 的全局内存读取按 32-byte / 128-byte cache line 合并。dense attention 的 KV 读取连续，多个线程一次读取一条 cache line。稀疏 attention 的 KV 读取分散，多个线程可能读到不同的 cache line——内存带宽利用率下降。

官方 TileLang 参考实现的做法是最朴素的一种——按 block=64 一批，把索引对应的 KV 逐元素搬进 shared memory，无效槽位（-1）填 0：

for i in T.Parallel(block):
    idxs[i] = T.if_then_else(t * block + i < topk, topk_idxs[by, bx, t * block + i], -1)
for i, j in T.Parallel(block, d):
    kv_shared[i, j] = T.if_then_else(idxs[i] != -1, kv[by, idxs[i], j], 0)

注意这里的关键：gather 的粒度是"一整条 KV"（d=512 个连续元素），不是单个标量。所以合并访问并没有真的失效——每条 KV 内部仍然是连续的 512×2 字节，只是条与条之间跳着走。这也是 V4 这种 MQA-style 单组 512 维 KV 在稀疏 gather 下的一个好处：head_dim 越大，一次 gather 搬的连续块越大，索引跳跃的相对代价越低。

挑战二：相邻 query 的索引重叠没有被利用

两个相邻 query token 的 topk_idxs 会有重叠（至少滑窗那 128 个位置几乎完全重合），但参考实现里每个 (batch, seq_pos) 是一个独立的 CTA（with T.Kernel(m, b, threads=threads) as (bx, by)），各自把自己那 1152 条 KV 从头搬一遍——重叠部分被重复读取，只能指望 L2 命中。这是"可读的参考实现"与"生产内核"之间最典型的一类差距。

挑战三：tile 大小与 K 的关系

FlashAttention 的标准做法是把 KV 分成 tiles，逐 tile 处理。参考实现里这几个数字是写死的：

num_stages = 2
threads = 256
block = 64
num_blocks = tilelang.cdiv(topk, block)

block=64、K=1152 ⟹ 每个 query 跨 18 个 tile，kernel 用 online softmax 在 tile 间累积分子分母，最后统一归一化。

FlashMLA 那边的分块常量是另一套：csrc/sm90/decode/sparse_fp8/components/config.h:19 的 PAGE_BLOCK_SIZE = 64、同文件 :16 的 QUANT_TILE_SIZE = 128（V3.2 路径）——而 V4（MODEL1）路径在 SM90（csrc/sm90/decode/sparse_fp8/config.h:28）和 SM100 上用的都是 QUANT_TILE_SIZE = 64，正好对应 model.py 里 act_quant(kv[..., :-rd], 64, ...) 那个 64。

5.5 online softmax：稀疏 attention 的核心算法挑战

V4 的 sparse_attn kernel 必须解决一个问题：当一个 query 跨多个 tile 计算时，怎么在 tile 之间正确累积 softmax 的分子分母？

这个问题在 dense FlashAttention 里就已经存在——每个 query 看完整 KV，KV 被分 tile 处理。FlashAttention 的解法是 online softmax 算法（Milakov & Gimelshein 2018）：

对于每个 tile, 维护两个 running quantity:
  m: 当前已见 logits 的最大值
  l: 当前的 sum of exp(logits - m)

收到新 tile 的 logits 时:
  m_new = max(m, max(new_logits))
  l_new = exp(m - m_new) * l + sum(exp(new_logits - m_new))
  o_new = exp(m - m_new) * o + sum(exp(new_logits - m_new) * new_v)

最后归一化:
  o = o / l

这个算法的妙处在于 m 和 l 的更新可以增量进行，且数学上等价于 dense softmax。

V4 的 sparse_attn 的 online softmax 多了一个复杂性：attn_sink。sink 不是某个 tile 里的位置，而是一个 head-level 的常数 logit——必须把它"虚拟"地参与到 softmax 归一化里。

参考实现的处理方式比"把 sink 当成第 0 个 tile"更省事——它把 sink 放到最后一步补进分母：

for t in T.Pipelined(num_blocks, num_stages=num_stages):
    ...                                    # 常规 online softmax，sink 完全不参与
    sum_exp[i] = sum_exp[i] * scores_scale[i] + scores_sum[i]
    ...
for i in T.Parallel(h):
    sum_exp[i] += T.exp(attn_sink[i] - scores_max[i])   # ← 循环结束后才补
for i, j in T.Parallel(h, d):
    acc_o[i, j] /= sum_exp[i]

为什么这样是对的？因为 scores_max 在循环结束时已经是全局最大值，sum_exp 是以它为基准的 ∑esi−m\sum e^{s_i - m}，此时补上 esink−me^{\text{sink} - m} 恰好就是"把 sink 当成一个虚拟位置参与归一化"。而 acc_o 全程不加 sink 项——sink 只进分母，不进分子，这正是"什么都不选"的语义。

FlashMLA 的 prefill kernel 和不拆分的 decode 与参考实现一样，在最后归一化时把 sink 补进分母（如 csrc/sm90/prefill/sparse/phase1.cuh:206 的 1.0f / (rL[i] + exp2f(attn_sink - rM[i]))）；一旦走 split-KV，就留到 combine 阶段处理，用的是 log-sum-exp 形式（csrc/smxx/decode/combine/combine.cu:107）：

float attn_sink = __ldg(params.attn_sink + q_head_idx);
global_lse += log2f(1 + exp2f(attn_sink*CUDART_L2E_F - global_lse));

数学上等价（lse′=log⁡(else+esink)\text{lse}' = \log(e^{\text{lse}} + e^{\text{sink}})），但因为 split-KV 会把一条序列的 KV 拆到多个 CTA 上算、最后再合并，sink 只能在合并那一步加一次，不能每个 split 都加。同一个文件里还处理了两个边界：attn_sink = +inf（所有真实位置权重归零）和 attn_sink = -inf（等价于没有 sink）。

这种"sink-aware online softmax"是 V4 / V3.2 这条稀疏路径的特化部分，标准 FlashAttention 的 online softmax 没有这一项。

flowchart LR
  subgraph TileFlow["online softmax 的 tile 流"]
    T0["tile 0: 64 个 KV"] -->|更新 m,l,o| T1["tile 1: 64 个 KV"]
    T1 -->|更新 m,l,o| T2["tile 2: ..."]
    T2 -->|...| TN["tile N"]
    TN -->|最后除以 l| Out["o 输出"]
  end
  Sink["attn_sink<br/>循环结束后补进分母"] -.参与归一化.-> Out

参考实现在每个 (batch, seq_pos) 上跑这套 sink-aware online softmax，跨 18 个 KV tile（K=1152 / block=64）累积，最后统一归一化输出。

5.6 索引 gather 的硬件路径

稀疏 attention 在 SM90 / SM100 上的"索引 gather"是 V4 工程化的一个具体硬件挑战。让我们看一段简化的 CUDA 伪代码理解它在硬件层做了什么：

// 简化的稀疏 KV gather kernel（C++/CUDA）
__global__ void gather_kv_for_sparse_attn(
    const __nv_fp8_e4m3* kv_global,        // [B, T, D] 全部 KV
    const int* topk_idxs_global,           // [B, S, K] topk 索引
    __nv_fp8_e4m3* kv_gathered_smem,       // shared memory 输出
    int B, int S, int T, int D, int K)
{
    int b = blockIdx.x;
    int s = blockIdx.y;
    int tid = threadIdx.x;

    // 每个 thread 负责 K/blockDim.x 个 KV 位置
    for (int k_local = tid; k_local < K; k_local += blockDim.x) {
        int kv_idx = topk_idxs_global[b * S * K + s * K + k_local];
        if (kv_idx < 0) continue;          // mask -1 invalid
        // gather D 维 (D=512)
        // SM90: prefill 用 cp.async 异步拷贝，decode 用 128-bit 全局 load 后在寄存器里反量化
        // SM100: 用 TMA 的 gather4 模式
        int dst_offset = k_local * D;
        int src_offset = b * T * D + kv_idx * D;
        // ... TMA / cp.async copy ...
    }
}

实际 FlashMLA 的 kernel 比这复杂得多。SM90 的 decode 路径叫 run_flash_splitkv_mla_fp8_sparse_kernel（csrc/sm90/decode/sparse_fp8/splitkv_mla.cuh），名字里的三个词就是它的三件事——splitkv（把一条序列的 KV 拆到多个 CTA，最后由 csrc/smxx/decode/combine/combine.cu 合并）、fp8（KV cache 以 FP8 存，反量化在 components/dequant.h）、sparse（按 indices gather）。SM100 的 csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/ 这份 kernel 用 CUtensorMap 走 TMA，其 config.h 里的 TmaParamsForDecode 结构体给 kv_nope / kv_rope / extra_kv_nope / extra_kv_rope 分别建了 tensor map——nope 段与 rope 段是分开搬的，正对应 V4 里"nope 走量化、rope 保 BF16"的切分。

但核心思想不变：把 topk 选取的 KV 搬到 SMEM，然后正常做 attention。FlashMLA 的工程价值就在于"把这条 gather + GEMM + softmax 流水线优化到接近硬件极限"。

5.7 SM90 vs SM100：两套 kernel 的差异

V4 同时支持 H100/H800（SM90）和 B200（SM100）。FlashMLA 给两套架构写了完全独立的 kernel：

先看 FlashMLA 自己声明的支持矩阵（README 里的表，原样搬过来）：

Kernel GPU Architecture MLA Mode KVCache Format
Dense Decoding SM90 MQA BF16
Sparse Decoding SM90 & SM100 MQA FP8
Dense Prefill SM100 MHA
Sparse Prefill SM90 & SM100 MQA

注意 Dense Decoding 只有 SM90、Dense Prefill 只有 SM100——这不是疏漏，是"哪条路径值得投人力"的结果：稀疏路径两个架构都全须全尾。

代码组织上的差异也很直白：SM90 的 decode 走 sparse_fp8/splitkv_mla.cuh（split-KV + persistent kernel，实例化按 h64/h128 展开），SM100 的 decode 走 head64/KernelTemplate<ModelType>；SM100 的 prefill 还额外分出一条 fwd_for_small_topk/ 路径（h_q=128 且 topk <= 1280 时选它，csrc/api/sparse_fwd.h:225）。这条路径还兼做 decode：h_q=128、d_qk=512（即 V4 Pro 的形状）的稀疏 decode 在 SM100 上调用的正是它的 DecodeWithSplitKV 模式（csrc/api/sparse_decode.h:179）。这些都是可以直接 ls 出来的事实。

性能数字（引自 FlashMLA README 的 Performance 一节，H800 SXM5 + CUDA 12.8；注意两行稀疏数字是 V3.2 的 DSA 内核在发布时公布的，dense 一行是 V3 的 dense MLA 内核，V4 路径没有单独公布过 benchmark）：

内核 H800 SXM5 B200
Sparse MLA decoding（FP8 KV cache，矩阵乘走 BF16） 410 TFLOPS 350 TFlops（README 注明"尚未真正优化"）
Sparse MLA prefill 640 TFlops 1450 TFlops（CUDA 12.9）
Dense MLA decoding 3000 GB/s（访存受限）/ 660 TFLOPS（算力受限） —

这张表里有个容易看反的地方：B200 的稀疏 decode（350）比 H800（410）还低，README 直接写了原因——那条路径还没优化。所以"新卡一定更快"在这里不成立；真正拉开差距的是 prefill（640 → 1450）。本专栏不为 V4 编造任何未公开的吞吐数字。

5.8 V4 sparse_attn 与 FlashAttention v3 的对比

FlashAttention v3 是 dense attention 在 H100 / B200 上的事实标准。把 V4 的 sparse_attn 与 FA3 横向对比：

维度 FlashAttention v3 V4 sparse_attn
KV 访问模式 连续 索引 gather
支持的 head_dim 最大 256 FlashMLA 只接受 d_qk ∈ {512, 576}、d_v = 512（源码 TORCH_CHECK）；参考实现是符号维度，不限死
支持的稀疏度 不支持 top-K
是否支持 sink 不支持 支持
是否支持 grouped O 不支持（O 投影外置） 不支持（O 投影也外置）
主要应用场景 dense attention 全部模型 V3.2（DSA）+ V4
长上下文成本 O(n²) FLOPs O(n) FLOPs (K 固定)

V4 的 sparse_attn 不能替代 FA3——它是为 V4 这种"超大 head_dim + 稀疏选取 + sink"的特定形态量身定做的。反过来 FA3 也不能跑 V4——FA3 的 dense KV 假设与 V4 的 topk_idxs 接口完全不兼容。

这种"特化 kernel"的代价是 V4 必须自带一套内核——要么直接用官方那份 TileLang 参考实现，要么接 FlashMLA。从 FlashMLA README 的数字看，同一张 H800 上稀疏 decode 是 410 TFLOPS、dense decode 在算力受限配置下是 660 TFLOPS：稀疏路径的硬件利用率更低（两者不是同一个 kernel、KV 格式也不同，稀疏路径多了 gather 和 FP8 反量化，README 没有拆分原因），但它只需要算 1152 个位置而 dense 要算全部——总账仍然是稀疏赢。

5.9 vLLM / SGLang 集成 sparse_attn 的工程接缝

把 sparse_attn 集成进 vLLM / SGLang 这类推理引擎，至少要做四件事：

事项一：确定内核来源

有两条路：直接把官方那份 TileLang 参考实现拉进依赖（好处是跨架构、坏处是性能与手写 CUDA 有差距），或者接 FlashMLA（性能好，但只有 SM90 / SM100，且需要 CUDA 12.8+、SM100 还要 12.9+）。选后者的话还要处理"用户硬件不在支持矩阵内"的回退——FlashMLA 自己不提供回退，TORCH_CHECK 会直接抛错。

事项二：传 topk_idxs 给 kernel

引擎需要在每次 forward 时为每个 attention layer 计算 topk_idxs——这意味着引擎要 invoke Indexer，传 query / 中间表示给 Indexer，再把 Indexer 的输出送给 sparse_attn。这个 dataflow 在 vLLM / SGLang 之前的代码里完全不存在——必须新增。

事项三：KV cache 形状改造

vLLM 默认的 PagedAttention KV cache 形状是 [block, block_size, head, dim]。V4 的 KV cache 是 [B, window + n/ratio, dim]——不分 head（MQA-style），不按 block 切。集成时要么扩展 PagedAttention 的形状抽象，要么新增一个"V4-style KV cache" 类型。第 19 章 §19.3 会指出这块该在哪里动刀——那一节写的是推断，不是已存在的 PR。

事项四：与调度器的协调

V4 的 prefill / decode 走不同 codepath，且 prefill 有"压缩 KV 一次性算"的批量步骤。vLLM 的调度器需要识别"这是 V4 模型"，分别调度 prefill 和 decode 阶段，避免把它们错误合并到同一个 CUDA stream。

第 19 章 §19.3 会把 vLLM 侧需要动刀的位置逐条列出（写作时 vLLM 尚无 V4 支持，那一节是推断而非已有 PR）。本章只到这里——给读者建立"sparse_attn 需要外部生态怎么配合"的全局认知。

5.10 动手实验：跑通 FlashMLA 的 V4 测试

# 1. 拉取 FlashMLA（注意有 submodule）
git clone https://github.com/deepseek-ai/FlashMLA.git flash-mla
cd flash-mla
git submodule update --init --recursive

# 2. 编译（README 要求：SM90 / SM100，CUDA 12.8+，SM100 内核需要 CUDA 12.9+）
pip install -v .

# 3. 跑稀疏路径的测试（仓库里真实存在的两个文件）
python tests/test_flash_mla_sparse_decoding.py
python tests/test_flash_mla_sparse_prefill.py

# 4. 性能基准
python benchmark/bench_flash_mla.py

没有 SM90 / SM100 就跑不了——csrc/api/sparse_fwd.h 里第一件事就是 TORCH_CHECK(is_sm90a || is_sm100f, "Sparse Attention Forward Kernel is only supported on SM90a and SM100f architectures.")。 A100（SM80）不在支持矩阵里，仓库也没有 Triton 回退。

想读懂 V4 那条路径，最省事的入口不是跑 benchmark，而是 grep：

grep -rn "MODEL1" csrc/ | head                       # V4 路径的所有分叉点
sed -n '30,45p' csrc/sm100/decode/head64/config.h    # 512 / 448 / 64 / 8 那组常量

把它和 inference/config.json 的 head_dim / rope_head_dim 并排看，比任何 TFlops 数字都更能说明"这套内核到底是给谁写的"。

5.11 sparse_attn 的成本模型：什么随 n 变、什么不变

本专栏不给未经实测的加速倍数。但成本模型是可以从源码直接读出来的，而它比任何倍数都更有用——它告诉你上下文变长时，哪一项会涨、哪一项不涨。

把一个 ratio=4 层里与上下文长度 nn 有关的项列全：

项 复杂度 1M 下的绝对值（每 query token 的乘加）
Indexer 打分（index_score einsum） O(n/ratio)O(n/\text{ratio}) ~2147M —— 唯一随 nn 增长的大项
Compressor 的 wkv / wgate O(1)O(1)（逐 token 恒定） ~18M
sparse_attn 的 QK + PV O(window+topk)O(\text{window} + \text{topk}) ~151M —— 与 nn 无关
Q / O 投影 O(1)O(1) ~296M

三条能直接用来做容量规划的结论：

  1. ratio=4 层上 sparse_attn 本身的成本与上下文无关。 无论 8K 还是 1M，每个 query token 都只看 128 + 1024 = 1152 个位置（ratio=128 的层要看全部 n/128 条压缩 KV，仍随 nn 增长，只是斜率小得多）。这是 V4 敢把 max_position_embeddings 写到 1048576 的底气。
  2. 随 nn 增长的是 Indexer，不是 attention。 想继续往上扩上下文，要动的是打分那一环（比如分层筛选），而不是 kernel。
  3. 短上下文下 V4 的架构红利拿不到。 8K 时 Indexer 只有 ~16M 乘加、sparse_attn 仍然是 151M（上下文不到 4K 时还会因为 min(index_topk, end_pos // ratio) 而更少），但 Compressor / Indexer / gather 这些固定开销一分不少。第 3 章 §3.15 按结构推演过：V4 的甜区在 64K 以上（本专栏的估计，不是实测）。

至于稀疏比例：1M context、ratio=4 的层里候选是 262144 条，选 1024 条，稀疏比 0.39%。index_topk 如果调大，多出来的成本只落在上面那张表的第三行（QK + PV 与 topk 成正比），相对 2147M 的打分开销几乎看不见——这也说明首发当天（04-24）把 index_topk 从 512 翻成 1024 在推理算力上代价不大。

5.12 判断 sparse_attn 是不是瓶颈

生产里最常见的问题不是"sparse_attn 有多快"，而是"它是不是当前的瓶颈"。这个判断不需要任何厂商数字，只需要把两件事对上：

第一件：这一层到底是算力受限还是访存受限。

decode 阶段每条序列每层要读 1152 条 × 512 维 × 1 字节（FP8 KV，按每维 1 字节粗算；FlashMLA 的 MODEL1 格式每条实为 584 字节：448 字节 FP8 + 64 维 BF16 RoPE + 8 字节 scale，见 tests/quant.py:56，合约 673 KB，量级不变）≈ 590 KB，算 128×1152×512×2≈151128 \times 1152 \times 512 \times 2 \approx 151M 次乘加（约 302 MFLOP）。算术强度 ≈ 512 FLOP/Byte。

这个数字之所以这么高，正是 MQA 的功劳：128 个 query head 共用同一份 KV，KV 只读一遍却被用了 128 次。它落在 H800 的机器平衡点（FlashMLA 稀疏 decode 的矩阵乘走 BF16，所以看 BF16 峰值算力 ÷ HBM 带宽，量级约几百 FLOP/Byte）附近——也就是说 sparse_attn 的 decode 既可能算力受限也可能访存受限，取决于 batch、topk 和实际达到的带宽。FlashMLA README 给 dense decode 同时报 "3000 GB/s（访存受限）" 和 "660 TFLOPS（算力受限）" 两个数，就是这个原因：同一个 kernel 在不同配置下落在不同的受限区，报单一数字没有意义。

实际含义：调优前先测出你这档配置落在哪一侧。落在访存侧就去看 KV 的存储精度和 gather 的局部性，落在算力侧才去看 tile 配置和 TensorCore 利用率——两边的优化手段完全不同。

prefill 阶段则相反：一次算 SS 个 query，KV 被复用，算力占比大幅上升——所以 FlashMLA 基本把 prefill 和 decode 拆成了两套独立的 kernel（csrc/*/prefill/sparse/ 与 csrc/sm90/decode/sparse_fp8/、csrc/sm100/decode/head64/）；例外是 SM100 的 fwd_for_small_topk/head128，同一份模板按 SparseAttnFwdMode 实例化出 prefill 与 split-KV decode 两个版本。

第二件：它前面还有一个串行依赖。

Indexer 必须先输出 topk_idxs，sparse_attn 才能开始（inference/model.py:511 → 514 → 528）。§4.6 算过，1M 下 Indexer 的算力是这一层 attention 的四倍多。所以如果你观测到"attention 段很慢"，先排除是不是 Indexer 慢——两者在 profile 里挨着，很容易归错账。

排查顺序：GPU 利用率 → 是 Indexer 还是 sparse_attn → 访存还是算力 → SMEM / tile 配置。第一步利用率就明显偏低的话，问题多半根本不在 attention 上，而在数据加载、通信或 Python 侧调度。

5.13 sparse_attn 的"工程债务"清单

任何工业级 kernel 都有它的工程债务——为快速发布而留下的"以后再优化" 项。FlashMLA 的稀疏 MLA kernel 也有。把可观察到的工程债务列出来：

债务 1：一个仓库里塞了 2 个模型 × 2 个阶段 × 2 个架构的笛卡尔积

FlashMLA 同时维护 dense（V3）与 sparse（V3.2 的 DSA、V4 的 MODEL1）两条路径，每条又分 prefill / decode，每个又分 SM90 / SM100。instantiations/ 目录里那一堆 model1_persistent_h128.cu / v32_persistent_h64.cu / phase1_k512.cu / phase1_k576.cu 就是这个笛卡尔积的显式展开——每个组合单独编译一次。好处是每种配置都能被独立特化到极致，代价是编译时间和二进制体积随组合数爆炸，以及支持矩阵里那几个空格（Dense Decoding 只有 SM90、Dense Prefill 只有 SM100）。

债务 2：SM80（A100）路径直接缺席

FlashMLA 的 README 支持矩阵里只有 SM90 / SM100，csrc/api/sparse_fwd.h 里的 TORCH_CHECK(is_sm90a || is_sm100f, ...) 把这件事写死成了运行时断言——不是"性能差"，是根本跑不了。A100 上想跑 V4 的 sparse_attn，只能退回官方那份 TileLang 参考实现，或者自己写一份。

债务 3：稀疏度是"半动态"的

index_topk 写在 config.json 里，改它不需要重编内核——FlashMLA 的 sparse_attn_prefill_interface 是从 indices.size(2) 读 topk 的，参考实现里 topk 更是 T.symbolic("topk")（inference/kernel.py:284）。但 Q/K 的 head dim 不行（内核侧叫 d_qk，取的就是 q.size(2)，csrc/api/sparse_fwd.h:127）：同文件 :131 的 TORCH_CHECK(d_qk == 576 || d_qk == 512, "Invalid d_qk: ", d_qk) 只认这两个值（576 = V3.2、512 = V4），下一行 d_v 更是硬性 512。能调的是选多少条，不能调的是每条多宽。

这里要把 pin 版与 V4 的 index_topk=1024（首发当天 04-24 就已改成 1024）的关系说清楚。按 71c7379 的源码判断（本机无 GPU，未实跑）：推理前向只从 indices 读 topk，不设取值白名单，只要求它是分块的整数倍（SM90 prefill 要求 128 的倍数，csrc/sm90/prefill/sparse/phase1.cuh:577；其余路径是 64 或 128 的倍数），V4 ratio=4 层的 128 + 1024 = 1152 条都满足。pin 版没有的是另一个参数 indexer_topk——让 prefill kernel 额外输出前 indexer_topk 条的 lse_indexer，提交标题称之为 "DSA backward indexer LSE"，是为训练侧准备的。它在 pin 之后才出现，而且都在 nv_dev 分支上：04-27 的 13d173a 加入时只接受 0 / 512 / 2048，INDEXER_TOPK=1024 要到 04-30 的 b7643bd 才支持。换句话说，本章核对的 pin 版能接 V4 形状的推理调用，但没有为 V4 的 1024 准备这套训练用的 LSE 实例化。

债务 4：调度元数据要手工管理

flash_mla_with_kvcache（flash_mla/flash_mla_interface.py:53）需要先调 get_mla_metadata 拿一个 FlashMLASchedMeta，而这个 meta 的复用条件写在 docstring 里（同文件 :80）且相当苛刻——"只有当张量形状以及 cache_seqlens / topk_length / extra_topk_length 的值都不变时才能跨次复用"。在 continuous batching 这种每步都在变的场景里，这个约束意味着引擎侧必须自己判断什么时候要重建 meta，判断错了就是静默的错误结果。

债务 5：错误诊断信息有限

如果 sparse_attn 跑出错（如 decode 侧 topk_idxs 的 page block 编码错了），错误信息通常是 CUDA error，不直接指向问题位置。FlashMLA 缺一套 debug build——开启后可以打印每个 tile 的状态。这是开源项目的常见短板。

理解这些债务让你在使用 FlashMLA 时心里有底——遇到问题时知道去哪里查、知道哪些限制是"暂时的"，避免被"看似不一致的现象" 困惑。

5.14 sparse_attn 工程师速记

部署或调试 V4 sparse_attn 时最常用的几条速记规则。打印一张贴在工位上：

速记 1：版本要求（FlashMLA README 的 Requirements 一节）

  • SM90 / SM100；CUDA 12.8+（SM100 内核要 12.9+）；PyTorch 2.0+
  • A100（SM80）不在支持矩阵内，TORCH_CHECK 会直接拦下

速记 2：官方公布过的性能数字（FlashMLA README，H800 SXM5 + CUDA 12.8）

  • 稀疏 MLA decoding：410 TFLOPS
  • 稀疏 MLA prefill：640 TFlops（B200 + CUDA 12.9 上 1450）
  • dense MLA decoding：访存受限 3000 GB/s，算力受限 660 TFLOPS
  • 稀疏两条是 V3.2 的 DSA 内核数字，dense 一条是 V3 的 dense MLA 内核；V4 路径（ModelType::MODEL1）没有单独公布过

速记 3：典型形状参数

  • topk_idxs：ratio=4 层上最多 [B, S, 1152]（= window_size 128 + index_topk 1024），int32，-1 为无效
  • KV cache：每层 [B, window_size + max_seq_len // ratio, 512]
  • 单序列 KV cache 总和：BF16 下 ~8.3 GB、混合格式（RoPE 64 维 BF16 + 其余 FP8，不计 scale）下 ~4.7 GB（1M context，不含 Indexer 缓存，见 §3.10）

速记 4：常见错误信号

  • Sparse Attention Forward Kernel is only supported on SM90a and SM100f architectures.（prefill，csrc/api/sparse_fwd.h:115）/ Unsupported architecture for sparse decode fwd（decode，csrc/api/sparse_decode.h:380）：GPU 架构不在支持矩阵内
  • topk (...) mod B_TOPK (...) must be 0（如 csrc/sm100/decode/head64/kernel.cuh:859）：topk 不是分块大小的整数倍
  • 输出 NaN：先查 attn_sink 等权重是否加载（attn_sink 由 torch.empty 创建，inference/model.py:456，没加载就是未初始化内存）

速记 5：优化优先级

  • 检查 GPU 利用率（nvidia-smi）—— 明显偏低多半是 IO/通信瓶颈，sparse_attn 不是元凶
  • 检查 tile size 与 SMEM 余量
  • 检查 stream 配置（DeepGEMM 与 sparse_attn 是否在同一 stream 串行）

速记 6：与 vLLM 集成的关键文件

  • vllm/attention/backends/flashmla.py（MLA 的 attention backend；注意文件名没有下划线）
  • vllm/model_executor/models/deepseek_v2.py（DeepSeek 系列的 model class。本专栏锁定的 vLLM 0.8.5 里没有 deepseek_v4.py——这个文件同时服务 V2 与 V3：DeepseekV3ForCausalLM 继承自 DeepseekV2ForCausalLM，registry.py 里两个名字都映射到它。V4 要落到 vLLM， 按同样的方式扩展这里）
  • vllm/distributed/device_communicators/（设备间通信器所在目录；0.8.5 里还没有 DeepEP 集成，第 16 章讨论接入方式）

5.15 两份实现的分工：TileLang 参考实现 vs FlashMLA

V4 这一条稀疏注意力路径上，同时存在两份实现。它们不是"新旧版本"的关系，而是两种不同的工程目的：

维度 inference/kernel.py 的 TileLang 版 FlashMLA 的 CUDA 版
位置 V4 模型仓库自带，536 行整个文件 独立仓库，csrc/ 下上万行
写法 @tilelang.jit + T.Kernel / T.gemm / T.Parallel CUDA C++ + CUTLASS + PTX 内联
分块 block=64、threads=256、num_stages=2，写死 按架构 / 模型 / head 数展开成一堆 instantiations/*.cu
KV 精度 BF16 进 kernel（QAT 模拟量化在 Python 侧做完） decode 路径原生吃 FP8 KV cache，dequant 在 kernel 内
prefill / decode 同一个 kernel，靠 topk 的符号维度自适应 基本拆成两套 kernel、两套目录（SM100 的 fwd_for_small_topk/head128 兼做两者）
split-KV 没有 有，splitkv_mla.cuh + 独立的 combine.cu
支持架构 凡是 TileLang 能编的 CUDA 卡 只有 SM90 / SM100
目的 把语义讲清楚、让人能读、让结果可复现 把这台机器榨干

对读者的实际建议：先读 TileLang 那份。它把"gather → GEMM → online softmax → sink → 归一化"这条链压在 80 行里，每一步都能对上第 2-4 章讲过的数学；把它读透之后再去看 FlashMLA，你会发现后者多出来的那上万行几乎全部在解决同一件事——怎么让数据搬运不成为瓶颈（TMA、split-KV、persistent kernel、FP8 KV cache、nope/rope 分离搬运），算法本身一行没变。

这也是本专栏反复强调的一点：算法的复杂度在论文里，工程的复杂度在访存里。

5.16 sparse_attn 与"动态批处理（continuous batching）"的协同

vLLM 的核心调度优化是"continuous batching"——多个请求的 prompt / decode 共享 GPU 跑在一个 batch 里。V4 的 sparse_attn 与这套调度有几个微妙的协同点。

协同点 1：每请求独立的 topk_idxs

continuous batching 把多个请求的 query 拼到一起跑 attention。V4 的 sparse_attn 接受 [B, S, K] 的 topk_idxs——每条 sequence 有独立的稀疏选择。

协同点 2：长短 prompt 共存

某些请求是 64K context、某些是 1K——一起 batch 时长 prompt 的 KV cache 占大头，但 sparse_attn 的计算仍可控（ratio=4 层 K 最多 1152）。这是 V4 在 mixed workload 下的红利。

协同点 3：prefill / decode 的混合 batching

vLLM 的 chunked prefill 让 prefill 与 decode 在同一 batch 内跑。V4 的 sparse_attn 需要分别处理两种 phase——prefill 走批量 codepath，decode 走增量 codepath。kernel 需要支持"同一 batch 内同时跑两种 codepath"。

协同点 4：prefix caching 与稀疏 KV

vLLM 的 prefix caching 复用相同前缀的 KV——多个请求共享 KV cache 的某部分。V4 的稀疏 KV 让这种复用更复杂——压缩 KV 段可以共享，滑窗段每请求独立。详见第 19 章关于 SGLang RadixAttention 的讨论。

协同点 5：抢占（preemption）

vLLM 的调度器会抢占 / 恢复请求——把"被抢占请求的 KV cache" 暂存到 CPU 后再恢复。V4 的稀疏 KV 要支持这种 swap，swap 单位会变成"滑窗段 + 压缩段"两块，不再是单一 KV block（本专栏的推断，锁定的 vLLM 0.8.5 尚无 V4 支持）。

5.17 延伸阅读

  • FlashAttention v3（arXiv:2407.08608）：dense attention 的 SM90 实现
  • Native Sparse Attention（arXiv:2502.11089）：稀疏 attention 训练的理论
  • DeepSeek FlashMLA 仓库（github.com/deepseek-ai/FlashMLA）：本章 §5.3、§5.6-§5.8、§5.13-§5.15 的一手依据，核对的是 71c7379（2026-03-31）；README 的 Performance 一节是本章唯一引用的性能数字来源
  • V4 仓库自带的 inference/kernel.py（536 行 TileLang）：本章 §5.1-§5.5 的一手依据
  • TileLang 项目：理解 T.Kernel / T.gemm / T.Pipelined 这套 DSL 的语义
  • 《vLLM 推理内核深度解析》专栏第 4-5 章：PagedAttention 与 V4 的 KV cache 对接
  • 本专栏第 19 章 §19.3：vLLM 适配 V4 要动哪些地方（推断，非已有 PR）

5.18 本章小结

  • V4 的 sparse_attn 在 PyTorch 是一行调用，官方参考实现是 inference/kernel.py 里 80 行 TileLang——不是 CUDA C++，也不依赖 FlashMLA
  • 但 FlashMLA（71c7379，V4 发布前近一个月）里已经有 V4 的路径：enum class ModelType { V32, MODEL1 } 的 MODEL1 分支，其 D_Q=512 / D_NOPE=448 / D_ROPE=64 / QUANT_TILE_SIZE=64 / V_HAVE_ROPE=true 与 V4 的 config 逐条吻合
  • attn_sink 只进 softmax 的分母不进分子：参考实现在循环结束后补 exp(sink - max)，FlashMLA 不拆分时同样在最终归一化里补，走 split-KV 时在 combine.cu 里用 lse += log2(1 + exp2(sink - lse))，三者等价
  • 官方公布过的性能数字全部来自 FlashMLA README 且属于 V3.2 的 DSA 内核（H800 上稀疏 decode 410 TFLOPS、稀疏 prefill 640 TFlops）；V4 路径没有单独公布过 benchmark
  • ratio=4 层上 sparse_attn 的成本与上下文长度无关（只看 1152 个位置），随 nn 增长的是它上游的 Indexer
  • 集成进 vLLM / SGLang 需要四件事：内核依赖、传 topk_idxs、改 KV cache 形状、调度器协调

第 6 章我们离开 attention 内部，来到 V4 长上下文工程的另一支柱：YaRN RoPE——V4 预训练按 4K→16K→64K→1M 逐步扩展上下文，RoPE 怎样配合这一过程。