DSA 稀疏注意力完全拆解：从 DeepSeek-V3.2 的闪电索引器到 vLLM 的算子实现

问题的实质：注意力贵在哪里，稀疏能省掉哪一项

要判断稀疏注意力这套思路值不值，先把注意力在推理时的开销拆开。

在自回归生成中，每一步的注意力可以写成这样的形式：给定当前 token（或当前一批 token）的 query $\mathbf{q}_t$，以及历史所有位置的 key $\mathbf{k}_s$ 和 value $\mathbf{v}_s$（$s \le t$），输出是

这个式子有两个层次的开销。第一个层次是算分：$\mathbf{q}_t \cdot \mathbf{k}_s$ 对每个历史位置 $s$ 都要算一次，一共 $t$ 次点积，每次点积涉及 $d$ 维乘加。第二个层次是取值：softmax 之后要对 $t$ 个 value 向量做加权求和。

在 prefill 阶段，一次要处理 $L$ 个 query，所以算分这一步的总计算量是 $L^2 d$，也就是那个老生常谈的 $\mathcal{O}(L^2)$。在 decode 阶段，每次只有一个新 query，看起来算分只要 $t d$，但真正的瓶颈换成了读数据：为了算出这 $t$ 个分数，必须把 $t$ 个 key 从显存里读进来；为了做加权求和，又必须把 $t$ 个 value 读进来。在 batch size 不大、算术强度很低的解码场景里，显存带宽是决定性因素，而 KV 缓存的大小和读写量都随 $t$ 线性增长——序列越长，每一步解码需要搬运的数据越多，吞吐越低。

所以注意力有两条独立的成本曲线：一条是 prefill 的 $\mathcal{O}(L^2)$ 计算，一条是 decode 的 $\mathcal{O}(L)$ 字节搬运（乘以层数、头数、精度字节数）。任何”长上下文加速”的方案都必须说清楚它打算砍哪一条。

上面这张图来自 Attention Is All You Need（Vaswani 等人，arXiv:1706.03762）中的多头注意力结构图，它标出了后面所有讨论的坐标系：左边是一组线性投影把输入拆成 Q、K、V 三条通路，中间是缩放点积注意力在多个头上并行发生，右边把各头的输出拼接再投影一次。本文后面反复出现的”主注意力的头数 $H$“指的就是中间那一叠并行通路；而 DSA 的做法，是在这套结构之外再长出一条小一号的注意力通路专门用来打分。

DSA 砍的是这样一个关键事实：对某个 query 而言，绝大多数历史位置对最终的输出几乎没有贡献。 注意力分数经过 softmax 之后高度集中在少数几个位置上，其余位置分到的权重接近零；把它们读进来、算出来、再乘上去，只是在浪费算力和带宽。如果把每个 query 真正需要的历史位置数固定为 $k$（$k \ll L$），那么 prefill 的算分从 $L^2 d$ 降到 $L \cdot k \cdot d$，decode 的读取从 $t$ 个 key/value 降到 $k$ 个。这就是 $\mathcal{O}(L^2)\to\mathcal{O}(Lk)$ 的来历。

代价也很清楚：你必须先知道”哪 $k$ 个位置重要”。这个”先知道”本身需要一次全序列的打分——如果不做任何优化，它就退化成一个完整的 $\mathcal{O}(L^2)$。DSA 的闪电索引器就是这次打分本身，它的全部设计目标可以概括为一句话：用足够便宜的代价算出足够准的近似分数，让”找 Top-K”这一步的总成本远小于它省下来的一次完整注意力。

上图是 DeepSeek-V3.2 论文中的 Figure 2，也是理解整篇文章的总图。图中左侧是标准的 MLA 通路：query token 的隐状态 $\mathbf{h}_t$ 经过下投影得到压缩的 query 表示 $\mathbf{q}_t$（在 MLA 中 query 本身也可能被低秩压缩，见下文），历史 token 的隐状态经过 KV 联合投影得到压缩的 KV 条目 $\mathbf{c}_s$。图中绿色的那一路则是 DSA 新增的旁路：索引器读取 $\mathbf{h}_t$ 和历史 KV，算出一组索引分数 $I_{t,s}$，据此挑出贡献最大的若干条目。真正被送进注意力的是被挑选的那部分 $\mathbf{c}_s$，而不是全部。

一个容易忽略的细节是：图中被挑选的 $\mathbf{c}_s$ 是被多个 query 头共享的。这与 MLA 的 MQA 模式直接相关，下文会专门讲。

索引器：为什么是这个分数函数

DSA 论文把索引分数定义成（式 (1)）：

逐个符号拆开看：

• $H^I$ 是索引器的头数。注意它和主注意力的头数 $H$ 是两个独立的配置量，而且索引器头数远小于主注意力。每个索引器头有自己的 query 向量 $\mathbf{q}^I_{t,j}$ 和 key 向量 $\mathbf{k}^I_s$，维度是 $d^I$，同样远小于主注意力的 head dim。
• $\mathbf{q}^{I}_{t,j} \in \mathbb{R}^{d^I}$、$w^{I}_{t,j} \in \mathbb{R}$ 都由当前 query token 的隐状态 $\mathbf{h}_t$ 通过线性投影得到。也就是说，每个索引器头既产出一个 query 向量，又产出一个标量权重。
• $\mathbf{k}^{I}_{s} \in \mathbb{R}^{d^I}$ 由历史 token 的隐状态 $\mathbf{h}_s$ 投影得到，每个历史 token 只有一份索引器 key，被所有索引器头共享（注意下标里没有 $j$）。

这个函数形状的每一处都是被”便宜”这两个字逼出来的。

为什么用点积而不是完整的注意力。 主注意力要用 $\mathbf{q}\cdot\mathbf{k}$ 之后接 softmax。索引器去掉了 softmax，也去掉了除以 $\sqrt{d}$ 的缩放。softmax 是一个跨所有位置的归一化操作，意味着算分必须看到完整的行才能出结果——这在分块和流式处理时非常别扭。去掉 softmax 之后，索引分数变成一个逐位置独立的标量：$I_{t,s}$ 只依赖 $\mathbf{q}_t$ 和 $\mathbf{k}_s$，可以任意分块计算、任意并行，不需要归约。这对一个要在 GPU 上以极小代价跑完全序列的打分器来说是最重要的性质。

为什么用 ReLU 而不是 Softmax 或 GeLU。 论文给出了直截了当的理由：吞吐考虑。ReLU 是逐元素的最大值与零比较，没有除法、没有指数、没有跨元素依赖，在 GPU 上是纯粹的 ALU 操作，可以和乘加指令流水线挤在一起；而 exp 属于 SFU（特殊功能单元）指令，吞吐远低于 FMA。对索引器这种每 token 要跑 $L$ 次的算子，把激活函数从 exp 换成 ReLU 带来的吞吐差异是实打实的。

另一个隐含好处是 ReLU 的稀疏性：负数被直接截断为零，意味着”这个（头, 位置）组合贡献为零”。虽然实现上不一定利用这种稀疏，但至少它保证单项贡献不会出现负值相互抵消的怪异情况，让这个”加权求和”更接近一个”投票”语义。

为什么需要那组权重 $w^I_{t,j}$。 如果没有 $w_{t,j}$，式 (1) 就退化成 $\sum_j \text{ReLU}(\mathbf{q}_j\cdot\mathbf{k})$，即所有索引器头等权求和。加上这组由 query 动态生成的权重，索引器就获得了逐头加权的能力：某个头如果在当前上下文里更”敏锐”，它的投票权重就可以被抬高。这是一个典型的轻量级门控设计——门控信号本身来自 query 隐状态，代价只是一个 $d^I \times H^I$ 的投影矩阵。

为什么”每个历史 token 只有一份索引器 key”。 因为索引器的 key 侧如果也按头存多份，那么它在 KV 缓存里占的空间就要乘以 $H^I$，而索引器缓存本来就必须对全部历史位置保持可见（它的打分范围是全序列，不能像滑窗那样只保留最近一段）。把 key 侧做成单头，是让索引器缓存尽量小的关键决策。

把这个函数和标准注意力放在一起对比，可以看出索引器是一个被刻意”降级”的注意力：去掉了归一化、去掉了缩放、换掉了激活函数、缩小了头数和维度、并且把 key 压成单头。它放弃了绝大部分表达能力，换来的是可以在极低成本下跑完整序列。

三个必须搞清楚的数字：$d^I$、$H^I$ 与 $k$

看 DSA 的实现之前，有三个配置量必须先弄明白各自控制什么，因为后面所有内核代码都是围着它们转的。

1
def get_max_prefill_buffer_size(vllm_config: VllmConfig):
2
    max_model_len = vllm_config.model_config.max_model_len
3
    factor = vllm_config.attention_config.sparse_mla_prefill_workspace_factor
4
    # Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes.
5
    # The flashmla_sparse backend uses a workspace size of factor * max_model_len.
6
    # The memory usage of the workspace there is 576 * 2 bytes; so we size this as
7
    # (576 * 2 // 132) * factor (40 at factor 5) to maximize this workspace
8
    # size while still fitting within the flashmla_sparse workspace.
9
    # For DeepSeek-V3.2, the max_model_len is 163840.
10
    #   40 * 163840 * 132 = 865075200 bytes = 825 MB
11
    return max_model_len * (576 * 2 // 132) * factor

1
        self.topk_tokens = config.index_topk
2
        self.n_head = config.index_n_heads  # 64
3
        self.head_dim = config.index_head_dim  # 128

实例化：DSA 为什么必须挂在 MLA 的 MQA 模式上

论文在 “Instantiate DSA Under MLA” 一节里给出了一个乍看之下很技术化、但背后是硬件约束的决定：DSA 基于 MLA（Multi-head Latent Attention）实例化，并且选用 MLA 的 MQA 模式。

先看这个选择的动机。论文的措辞是：在 kernel 层面，每个 key-value 条目必须被多个 query 共享，才能获得计算效率。 这句话点出了稀疏注意力在 GPU 上的核心矛盾。

设想如果 DSA 挂在标准 MHA 上：每个 query 头都有自己的一套 key/value。那么”选出 Top-K 位置”这个动作选出的是位置，但每个位置在每个头下都有不同的 key/value 数据。于是稀疏注意力内核在读数据时，需要按照”位置 × 头”的笛卡尔积去 gather——这会让索引数组从”位置列表”变成”位置与头的乘积”，gather 的粒度和地址计算都变得极其零碎，访存几乎无法做宽向量化，稀疏带来的带宽收益被 gather 开销吃光。

MLA 的 MQA 模式恰好解决了这个问题。MLA 的根本设计是把 KV 压到一个低维的潜向量（latent vector）$\mathbf{c}_s$ 上：历史 token 的隐状态经过一个下投影得到一个维度远小于 $H \times d$ 的压缩表示，所有头部共享这同一个潜向量。这个潜向量就是论文里说的”MLA 的 key-value 条目”。在 MQA 模式下，一个潜向量 $\mathbf{c}_s$ 被该 token 的所有 query 头共享。

于是”选 Top-K 位置”就等价于”选 Top-K 潜向量”。索引数组描述的是”要读哪些潜向量”，而这些潜向量对所有头都是同一份数据。内核对一个被选中的位置只需要做一次 gather，读进来的 $d_c$ 维潜向量可以同时在所有头上复用。稀疏性落到了”数据搬运”这个维度上，才是真正有价值的稀疏性。

站内的 MLA 完全拆解：DeepSeek 低秩 KV 压缩注意力的原理与全部实现细节 把 MLA 的低秩压缩、RoPE 解耦与矩阵吸收讲得很透，GQA 完全拆解：分组查询注意力，用 1/8 的 KV 缓存保住 MHA 的质量 则解释了 GQA/MQA 这类”共享 KV 头”的设计为什么能同时省缓存和带宽。本文不再重复这两条线，只强调它们在这里的作用：MQA 不是为了让 DSA 省缓存才选的，而是为了让稀疏索引落到一个粒度足够粗、能被所有头共享的寻址单位上。

顺带一个推论：这个选择也解释了为什么 DSA 和 MLA 是一体的，而很难单独迁移到别的注意力结构上。稀疏注意力需要一个”多个消费者共享一份数据”的最小寻址单元，MLA 的潜向量天然提供了这个单元。

实测口径：每 token 成本随序列位置怎么变

机制讲完，接下来需要一个校准数量的锚点：这套设计在真实硬件上到底把成本改成了什么形状。

论文的 Figure 3 给出的是一组”每 token 成本随序列中位置变化”的曲线，对比对象是 DeepSeek-V3.1-Terminus（纯 MLA 稠密注意力）和 DeepSeek-V3.2（DSA）。数据的口径论文写得很清楚：在 H800 上对实际部署的服务做基准测试，算力单价按 2 美元每 GPU 小时折算，因此纵轴是货币化的成本，横轴是 token 在序列中的位置。

两张图对照着看，可以读出三件事。

第一，稠密注意力的成本随位置增长，稀疏的成本近似平坦。 这是 $\mathcal{O}(L^2)$ 与 $\mathcal{O}(Lk)$ 在成本曲线上最直接的区别：在稠密路径上，越靠后的 token 要看的上下文越长，单位成本越高；在稀疏路径上，无论 token 排在第几千位还是第几十万位，它要读的 KV 条目数都是 $k=2048$——不变。曲线的形状本身就验证了复杂度量级的改变。

第二，索引器的开销体现在”平坦但非零”上。 稀疏路径的曲线不是零，它是一条有高度的近似水平线。这条线的高度就是索引器打分、Top-K 选择、地址翻译加上 $k$ 个条目注意力这一整套的成本之和。“平坦”意味着它不随上下文增长，“有高度”意味着它在短上下文下可能反而是负收益——这正是下文要讨论的上下文长度门槛的来源。

第三，论文专门提到短序列预填充的处理方式。 论文里的一句话值得单独拎出来：对短序列的预填充，我们专门实现了一个 masked MHA 模式来模拟 DSA，这样在短上下文条件下能获得更高的效率。 这句话说明了 DSA 的实现者在工程上做了一件很务实的事：为短上下文单独准备一条稠密路径。 当上下文很短时，索引器那一整套开销（打分、Top-K、翻译）相对收益显得太重，不如直接用 masked MHA 把注意力算完——反正短上下文下稠密注意力的成本已经不大了。这与前文提到的 K->ne[1] >= max(4096, 2*n_gather) 判据是同一件事的两个层次：论文在模型侧准备了稠密回退，框架在 kernel 侧准备了稀疏门槛。

有了这三条校准，后面讨论所有”稀疏的代价”才有量纲：不是”稀疏有什么理论缺陷”，而是”这条平坦的线要高到什么程度，才会在某个上下文长度上超过稠密路径”。

训练侧：为什么它只需要继续预训练

虽然本文的主体是推理侧，但”索引器是怎么被训练出来的”这个问题必须简短交代，因为它决定了索引器在推理时能做多少事。

论文把 DSA 的引入分成两个阶段，都建立在 DeepSeek-V3.1-Terminus 的基座 checkpoint（上下文已扩展到 128K）之上，训练数据分布与 128K 长上下文扩展阶段完全对齐。

第一阶段是 dense warm-up。 在这个阶段索引器已经参与，但注意力的行为与稠密注意力一致——每个 query 仍然看完整的上下文。这一步的目的是让新初始化的索引器参数先”热身”，在主模型行为尽量不受扰动的前提下，把索引器的分数分布先训到一个合理的范围。

第二阶段是 sparse training。 从这个阶段开始，索引器真正控制注意力看到哪些位置（每 query 选 2048 个 KV），并且索引器的优化目标与主模型分离：索引器只由索引器自己的损失 $\mathcal{L}^I$ 驱动，主模型只由语言建模损失驱动。论文给出的这一阶段学习率是 $7.3\times10^{-6}$，训练 15000 步，每步 480 条 128K 长度的序列，总计 943.7B token。

这里最值得停下来想的一点是索引器有自己的损失。这不是一个”用主模型的梯度顺便带出来”的附属模块，而是一个有独立监督信号的模块。它的监督信号从哪来？直观上可以理解成：既然最终注意力只在稀疏选出的集合上计算，那么”如果让 query 看完整上下文，注意力权重会集中在哪些位置”就是索引器应该学会预测的答案。索引器的损失就是在逼近这个回答。这种”用稠密注意力当老师、用稀疏索引器当学生”的做法，是稀疏注意力能被接进已有模型而不掉点的关键。

论文对这一阶段的表述很克制，没有展开损失的具体形式，只说明了它与主模型损失分离。对于本文的读者，需要记住的结论只有两条：索引器是一个被专门训练过的模块，不是随手设计的启发式打分器；它的训练目标与主模型解耦，因此可以在不扰动主干的情况下被替换或继续优化。

落地形态之一：vLLM 里的索引器算子

接下来进入本文的重点：这套机制在真实推理框架里长什么样。以 vLLM 的实现为主线，因为它把索引器拆成了若干个可以分别理解的内核，层次最清楚。

索引器在 vLLM 里的位置是 vllm/model_executor/layers/sparse_attn_indexer.py（算子层）与 vllm/v1/attention/backends/mla/indexer.py（后端与元数据层）。前者是 SparseAttnIndexer 这个 CustomOp 的实现，后者负责为它准备元数据、管理索引器自己的 KV 缓存。

先看整体数据流。一次索引器的调用需要吃进来这样几样东西：

• 当前这批 token 的 hidden_states（用来算索引器的 query 和权重）；
• q_quant（已经量化成 FP8 的索引器 query）；
• k（当前批次新产生的索引器 key，需要写进索引器缓存）；
• weights（每个 query 在每个索引器头上的标量权重）；
• 以及最关键的一份外部状态：索引器自己的 KV 缓存 kv_cache。

输出则是一个 topk_indices_buffer：每行对应一个 query token，行内是长度为 $k$ 的整数数组，元素是该 query 选中的历史位置编号。

把这条链路拆成”打分 → 选 Top-K → 翻译成物理地址”三步，正好对应三个不同的内核。

1
    pos = tl.load(positions + token)
2
    cos = tl.load(cos_sin_cache + pos * cos_sin_s0 + offs32).to(tl.float32)
3
    sin = tl.load(cos_sin_cache + pos * cos_sin_s0 + 32 + offs32).to(tl.float32)
4
    q_base = q + token * q_s0 + head * q_s1
5
    out_base = q_fp8 + token * q_fp8_s0 + head * q_fp8_s1
6

7
    if is_neox:
8
        # NeoX layout, x0 = q[0:32], x1 = q[32:64]
9
        x0 = tl.load(q_base + offs32).to(tl.float32)
10
        x1 = tl.load(q_base + 32 + offs32).to(tl.float32)
11
    else:
12
        # interleaved layout
13
        # x0 = q[0, 2, 4, ...], x1 = q[1, 3, 5, ...]
14
        x0 = tl.load(q_base + offs32 * 2).to(tl.float32)
15
        x1 = tl.load(q_base + offs32 * 2 + 1).to(tl.float32)
16
    r0 = (x0 * cos - x1 * sin).to(tl.bfloat16).to(tl.float32)
17
    r1 = (x1 * cos + x0 * sin).to(tl.bfloat16).to(tl.float32)
18
    amax = tl.maximum(tl.max(tl.abs(r0)), tl.max(tl.abs(r1)))
19

20
    q_nope = tl.load(q_base + 64 + offs64).to(tl.float32)
21
    amax = tl.maximum(amax, tl.max(tl.abs(q_nope)))
22
    scale_raw = tl.maximum(amax, 1e-10) * (1.0 / fp8_max)
23
    # e8m0 format
24
    q_scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))

1
    r0 = (x0 * cos - x1 * sin).to(tl.bfloat16).to(tl.float32)
2
    r1 = (x1 * cos + x0 * sin).to(tl.bfloat16).to(tl.float32)

1
    scale_raw = tl.maximum(amax, 1e-10) * (1.0 / fp8_max)
2
    q_scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))

1
    weight = tl.load(weights + token * weights_s0 + head * weights_s1).to(tl.float32)
2
    tl.store(
3
        weights_out + token * weights_out_s0 + head * weights_out_s1,
4
        weight * q_scale * softmax_scale * head_scale,
5
    )

1
                row_ends = cu_seqlen_ke - cu_seqlen_ks
2
                logits = as_fp32_logits(
3
                    fp8_fp4_mqa_logits(
4
                        (q_slice_cast, q_scale_slice),
5
                        (k_quant_cast, k_scale_cast),
6
                        chunk_weights,
7
                        cu_seqlen_ks,
8
                        cu_seqlen_ke,
9
                        chunk.max_seqlen_k,
10
                    ),
11
                    row_ends,
12
                )

1
            # DeepGEMM returns compressed logits: row i starts at K position
2
            # cu_seqlen_ks[i]. None means every row starts at column zero.

1
                ops.top_k_per_row_prefill(
2
                    logits,
3
                    torch.zeros_like(row_ends) if row_starts is None else row_starts,
4
                    row_ends,
5
                    topk_indices,
6
                    num_rows,
7
                    logits.stride(0),
8
                    logits.stride(1),
9
                    topk_tokens,
10
                )

1
            topk_indices, topk_length = triton_convert_req_index_to_global_index(
2
                req_id_per_token,
3
                block_table,
4
                topk_indices,
5
                BLOCK_SIZE=attn_metadata.block_size,
6
                BLOCK_STRIDE_ROWS=block_stride_rows,
7
                NUM_TOPK_TOKENS=topk_indices.shape[1],
8
                return_valid_counts=True,
9
            )

稀疏注意力内核：索引数组怎么被消费

翻译出物理索引之后，真正的稀疏注意力计算才发生。这里以 flashmla_sparse 后端的路径为例，把”索引数组如何被内核消费”讲清楚。

整体的调用形态是这样的（省略了预填充与解码分流的细节）：

1
        attn_out, lse = self._bf16_flash_mla_kernel(
2
            q,
3
            kv_rows,
4
            topk_indices,
5
            topk_length,
6
            actual_num_heads,
7
        )

传进去的 kv_rows 是一个二维行视图——把分页的 KV 缓存摊平成”每行一个潜向量”的排布，这样物理索引就可以直接当作行号使用。转换这一步由 flat_kv_row_view 完成。这个视图化是 GPU 上处理分页缓存的标准手法：把间接寻址问题转化成”在一个规则二维张量上做行 gather”，从而避免在 kernel 里做多级指针追踪。站内 FlashInfer 完全拆解：把分页 KV、动态批处理与定制注意力编译成服务级 Kernel 里也用了同样的思路处理稀疏 KV 表示，两篇可以对照着看。

1
                int64_t i_KV;
2
                if constexpr (use_sparse) {
3
                    // padded slots gather row 0, the -inf mask removes their contribution
4
                    const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : 0;
5
                    i_KV = index >= 0 ? index : 0;
6
                } else {
7
                    i_KV = k_VKQ_0 + i;
8
                }

1
        // the sparse mask values are gathered per element, always load them synchronously
2
        constexpr bool use_cp_async = nstages == 1 && !use_sparse;

1
    // sparse gather is not implemented for multi-stage loading
2
    return use_sparse && nstages_target > 1 ? 1 : nstages_target;

1
template <int ncols1, bool oob>
2
__launch_bounds__(256, 1)
3
static __global__ void flash_attn_mask_to_sparse_indices(
4
        const half * mask_ptr, int32_t * indices_ptr, int32_t * counts_ptr, const int ne30, const int n_queries,
5
        const int n_kv_max, const int64_t s31, const int64_t s33) {
6
    ...
7
    const half * mask = mask_ptr + sequence*s33 + q0*s31;
8
    int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + group)*n_kv_max;
9

10
    __shared__ int warp_offsets[256/WARP_SIZE];
11
    __shared__ int row_count;
12
    __shared__ int chunk_count;

1
    const int count = min(row_count, n_kv_max);
2
    for (int i = count + tid; i < n_kv_max; i += blockDim.x) {
3
        indices[i] = -1;
4
    }
5
    if (tid == 0) {
6
        counts_ptr[int64_t(sequence)*gridDim.x + group] = count;
7
    }

掩码重写：vLLM 选择的那条路

与 llama.cpp 的”掩码压缩成索引”相反，vLLM 在 DSA 路径上走的是另一条路：用 Top-K 索引去重写掩码，然后仍然调用标准的 FlashAttention 类内核，只不过这个内核知道”只需要遍历前 $k$ 个槽位”。

在 DSA 的 build_attn 里，这个过程写得非常明确：

1
    const auto & kq_mask = inp->get_kq_mask_mla();
2

3
    // prepare new kq mask - starts filled with -INFINITY
4
    ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
5

6
    // reshape KQ mask into tensor with rows of size 1:
7
    // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
8
    kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], ...);
9

10
    // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
11
    ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, ...);
12

13
    // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
14
    ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);
15
    zeros = ggml_fill(ctx0, zeros, 0.0f);
16

17
    // modify KQ mask by unmasking elements that are in top_k indices
18
    ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
19

20
    // reshape to restore the original shape of KQ mask
21
    kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, ...);
22

23
    // combine with the original kq mask
24
    kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
25

26
    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, top_k->ne[0], kq_scale, il);

（这段代码来自 llama.cpp 的 llama-graph.cpp，它在 DSA 路径上采用了与 vLLM 相同的”掩码重写”思路，同样因为注释写得清楚而适合用来讲解。llama.cpp 侧的计算图与注意力算子的组织方式，站内 llama.cpp 源码完全拆解（二）：ggml 计算图与内存分配器——惰性张量、生命周期复用与跨后端切图 有完整讲解。）

这套操作的语义可以这样读：

1. 先造一个全 $-\infty$ 的掩码 kq_mask_all——把它理解成”所有位置都被屏蔽”的初始状态。
2. 造一个全 0 的 zeros 张量，形状是 [1, n_top_k, n_batch, n_stream]——每个被选中的槽位对应一个 0。
3. 用 ggml_set_rows 把 zeros 的值按 Top-K 索引散射写回到 kq_mask_all 里。写进去之后，被 Top-K 选中的位置掩码值从 $-\infty$ 变成了 0，也就是”解除屏蔽”。
4. 再把原始的 kq_mask 加上去——原始掩码里本身就带因果性和序列边界的约束，两者相加之后，一个位置要真正被看到，必须既被 Top-K 选中，又满足原始的因果/边界条件。

相加这一步是不能省的。Top-K 选出的 2048 个位置是”分数上贡献最大的”，但索引器打分的范围是全序列，它可能选出一些本来就不该被看到的槽位（例如其他序列的槽位、或未来位置）。原始的 kq_mask 承载了这些约束，必须叠加回去。

为什么 vLLM 选这条”重写掩码”的路而不是”压缩成索引数组”的路？ 因为稀疏注意力内核通常已经具备”按照掩码跳过整块”的能力（站内 BLASST：一个标量阈值实现的动态块稀疏注意力（MLSys 2026） 与 vLLM 源码完全拆解（七）：PagedAttention Kernel 如何遍历分页 KV 并完成稳定 Softmax 都讲过块级跳过）。如果内核本身就能根据掩码跳过贡献为零的块，那么把索引写成掩码的形式，不需要为稀疏单独写一个内核，复用现有的 FlashAttention 实现即可。

代价是掩码张量本身仍然是 $\mathcal{O}(n_{kv} \times n_{tokens})$ 大小——也就是平方级。所以”掩码重写”这条路并没有把显存占用降到稀疏的级别，它省掉的是计算和KV 读取，省不掉掩码本身。而 llama.cpp 那条”索引数组”的路，索引数组是 $\mathcal{O}(n_{tokens} \times k)$，连掩码的平方级占用也一并省掉了。

这解释了 n_kv_max 这个参数的存在。在 llama.cpp 的稀疏路径里，它把 top_k->ne[0]（即 $k$）传给注意力算子：

1
    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, top_k->ne[0], kq_scale, il);

算子拿到这个数之后就知道”最多只有前 $k$ 个槽位需要真正读取”，于是内核的 gather 循环只需要跑 $k$ 次而不是 $n_{kv}$ 次。而 CUDA 侧的稀疏分流逻辑正是拿它当判据：

1
    const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4);
2

3
    // the dense kernel handles up to 64/ncols2 queries per K/V pass, the single-query gather has to beat that
4
    const int64_t n_gather = (ncols1 == 1 ? std::min<int64_t>(Q->ne[1], 64/ncols2) : ncols1) * (int64_t) n_kv_max;
5

6
    return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) &&
7
        mask != nullptr && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f &&
8
        mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 &&
9
        K->ne[1] >= std::max<int64_t>(4096, 2*n_gather);

这段判据很值得读第二遍。它说的是：只有当 KV 总长度至少是”需要 gather 的元素数”的两倍（并且不低于 4096）时，才值得切换到稀疏路径。

这是一个非常务实的工程判据。它承认了 sparse 路径有固定的额外开销（查表、失去异步拷贝、失去多级流水），因此只有在”扫过整个 KV 的成本”显著大于”gather 一小部分”的成本时，切换才划算。2*n_gather 这个两倍的系数就是这条经验分界线——稀疏不是无条件的优化，它对上下文长度有门槛。

索引器缓存：为什么不能和注意力共用一块 KV

前面提到索引器有自己独立的 KV 缓存。这个”独立”是必须的，原因有三层。

第一层：存的东西不一样。 索引器缓存存的是 $\mathbf{k}^I_s$——由 $\mathbf{h}_s$ 经索引器 key 投影得到的 $d^I$ 维向量，单头。注意力缓存存的是 MLA 的潜向量 $\mathbf{c}_s$ 加上解耦的 RoPE 分量 $\mathbf{k}^{rope}_s$。两者的维度、投影矩阵、甚至精度都不一样（索引器走 FP8，主注意力走 bf16 或 FP8 的另一套 scale 口径）。合在一起没有意义。

第二层：生命周期不一样。 索引器缓存必须覆盖全部历史位置，因为索引器的打分范围是全序列——这是它的职责所在。而注意力侧的 KV 缓存可能带滑窗优化、可能被压缩、可能有不同的分块策略。在 vLLM 的实现里，索引器后端甚至有自己的 get_supported_kernel_block_sizes，与主注意力后端分开配置。如果强行合并，任何一方的存储策略变化都会牵动另一方。

第三层：读写时机不一样。 索引器 key 需要在打分之前就位，注意力的 KV 需要在打分之后按选出的位置被读取。两者在一次前向里的时序位置不同，混在一个池子里会让分配与回收的记账变得复杂。

vLLM 用了一套专门的代码路径来管理这件事。索引器后端 DeepseekV32IndexerBackend 是一个独立的 AttentionBackend，注册在注意力后端的体系里，有自己的一套元数据构造器 DeepseekV32IndexerMetadataBuilder。它甚至有自己的 CUDA Graph 支持策略判断：

1
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
2
    # The indexer opts out of the shared reorder-threshold vote (see __init__),
3
    # so this is None; its own split uses self.decode_threshold.
4
    reorder_batch_threshold: int | None = None
5
    requires_block_table_width = True
6

7
    @classmethod
8
    def get_cudagraph_support(
9
        cls,
10
        vllm_config: VllmConfig,
11
        kv_cache_spec: KVCacheSpec,
12
    ) -> AttentionCGSupport:
13
        if _supports_varlen_paged_mqa_logits() or _use_flattening(vllm_config):
14
            return AttentionCGSupport.ALWAYS

（这段代码在 vLLM 的 indexer.py 里。vLLM 的 CUDA Graph 机制本身，站内 vLLM 源码完全拆解（十二）：piecewise compilation 与 CUDA Graph——消灭 kernel 调度开销的两层机制 有系统讲解。）

这里有两处细节值得注意。一是”选择退出共享的 reorder-threshold 投票”——reorder_batch_threshold 被显式设为 None，注释说明索引器自己用 decode_threshold 来做预填充与解码的切分。为什么？因为索引器后端对”一行多长才算解码”的判断口径和主注意力不同：索引器的打分粒度、它对变长行的处理方式都不一样，强行共用同一个阈值会导致切分不合理。二是 CUDA Graph 支持策略是条件性的——只有当平台支持变长分页 MQA logits、或者启用了 flattening 时才返回 ALWAYS。这反映了一个现实：索引器内核的形状对 CUDA Graph 的静态形状要求比较敏感，能不能整图捕获取决于底层算子是否支持变长输入。

在 Blackwell 平台上，vLLM 还走了一条额外的路径。当满足条件时，它会用 CuteDSL 实现的 pack 内核与稳定 Top-K 选择器：

1
def _assert_cutedsl_dcp_merge_supported(logits, topk_indices, topk_tokens):
2
    # The DCP merge only supports the CuteDSL path (Triton pack kernel + CuteDSL
3
    # stable-topk selector); there is no PyTorch fallback. The first cut targets
4
    # Blackwell/Hopper with index_topk in (512, 1024, 2048) (the selector's radix
5
    # sizing); the Triton pack itself has no shape/topk constraints.
6
    ...
7
    if topk_tokens not in (512, 1024, 2048):
8
        raise ...

（这段来自 sparse_attn_indexer.py。）注释把约束讲得很清楚：CuteDSL 的 Top-K 选择器用基数排序（radix）实现，它的排序桶划分是按 $k$ 的大小预先定死的，所以只支持 512 / 1024 / 2048 三档 $k$；而 Triton 那个 pack 内核本身没有形状约束。这是一个很典型的”为了性能而收窄可配置空间”的工程取舍——用限制 $k$ 的取值，换取一个更快的 Top-K 实现。DeepSeek-V3.2 用的正是 2048，恰好落在支持档位里。

这段代码里还有一个概念值得单独点一下：DCP。注释提到”这个合并只在 CuteDSL 路径下支持，没有 PyTorch 回退”，说明索引器在解码上下文并行（Decode Context Parallelism）场景下，每个 rank 只持有 1/N 的 KV 缓存，因此每个 rank 算出的 Top-K 是”我自己这 1/N KV 里的 Top-K”。要得到全局正确的 Top-K，必须把各 rank 的候选合并起来——合并的方式是每个 rank 打包自己的 $(score, global\_id)$ 候选，all-gather 之后再做一次全局稳定 Top-K。_merge_dcp_topk_global 做的就是这件事。

这暴露了一个容易被忽略的事实：稀疏注意力在分布式部署下需要额外的通信。稠密注意力下，每个 rank 各自算自己的那一份注意力分数，不存在”必须先知道全局哪几个位置最重要”这个跨 rank 的依赖。而稀疏注意力必须先做一次全局 Top-K，这就引入了 all-gather。长上下文场景下，稀疏本身省下的通信/带宽和这次额外通信要放在一起算总账。

稀疏到底省掉了什么

到这里，可以回过头把”省了什么、没省什么”列清楚了。

省掉的：

1. 主注意力的算分计算量。 从 $L^2 d H$ 降到 $L k d H$。在 $k=2048$、$L$ 远大于 $k$ 时，这是数量级的差异。
2. 主注意力的 KV 读取量。 每次 decode 不再需要把全部历史 key/value 从显存读进来，只需要读 $k$ 个潜向量。
3. 主注意力的 value 加权求和量。 同理，从 $t$ 个降到 $k$ 个。

没省掉的（或者说省得有限的）：

1. 索引器自身的 $\mathcal{O}(L^2)$ 打分。 论文明确承认这一点。索引器仍然是每对 (query, 位置) 都要算 $H^I$ 次 $d^I$ 维点积。它省的是常数系数——通过缩小 $H^I$ 和 $d^I$、去掉 softmax、把激活换成 ReLU、跑在 FP8 上——但复杂度阶数没有变。
2. 索引器 key 缓存的显存。 它必须对全部历史位置可见，$d^I$ 字节（加 scale）每 token。128K 上下文下这是几百 MB 量级的固定开销。
3. “逻辑位置 → 物理槽位”的翻译开销。 这是一次随 $n_{tokens} \times k$ 规模增长的索引运算，虽然比读 KV 便宜，但不是零。
4. 掩码的开销（如果走掩码重写这条路）。 掩码张量本身仍是平方级大小。
5. gather 访存带来的效率损失。 失去异步拷贝、失去多级流水、失去部分向量化。
6. 分布式下的全局 Top-K 通信。

第 1 条和第 2 条是理解 DSA 收益边界的关键。DSA 并没有消灭长上下文的平方项，它把平方项从”昂贵的完整注意力”转移到了”廉价的索引器”。 所以 DSA 的收益能不能兑现，本质上取决于索引器的常数系数相对于主注意力的常数系数小多少倍。论文给出的数字是索引器”只需要远少于 MLA 的计算量”，而 vLLM 的实现里能看到这个”少”是怎么被做出来的：更小的头数、更小的 head dim、FP8 量化、以及一个把 RoPE 和量化融成一步的专用内核。

代价的工程账本

把上面这些代价按工程维度重新组织一遍，就得到一份”DSA 落地时你需要关心的清单”。

工作区显存。 索引器的预填充需要一个大小正比于 max_model_len 的工作区。前面那段 get_max_prefill_buffer_size 的代码给出的量级是：128K 上下文下约 825 MB。这个数字不是凭空来的，而是为了让索引器工作区恰好塞进 flashmla_sparse 自身工作区（$576\times2$ 字节每 token）的预算里。部署时要预留的第一笔预算就是它。

索引器缓存的容量。 每 token 132 字节（$d^I=128$ 的 FP8 数据加 4 字节 scale），且必须覆盖全序列。这是第二笔预算，且它不可被滑窗或驱逐策略缩减——减了会直接破坏索引器的正确性。

$k$ 的可选范围。 在 Blackwell 上走 CuteDSL 路径时，$k$ 被限制在 512/1024/2048。想用别的 $k$ 就只能走 Triton 路径。DeepSeek-V3.2 选了 2048，正好在支持范围内，这个巧合背后显然是协同设计的结果。

上下文长度门槛。 稀疏路径的判据是 K->ne[1] >= max(4096, 2*n_gather)。短上下文下切换到稀疏路径是负优化，所以实现里必须有稠密路径兜底（论文也提到，短序列预填充时他们”专门实现了一个 masked MHA 模式来模拟 DSA”，以获得更高的短上下文效率）。

平台差异。 变长分页 MQA logits 需要 SM100 级别的能力；flattening 路径需要 SM90；ROCm 上只对特定架构开放。同一个索引器，在不同平台上的可用路径和 CUDA Graph 支持等级都不同。

分布式下的 Top-K 合并。 用 DCP 时，每个 rank 的局部 Top-K 必须 all-gather 后做全局稳定 Top-K，且只支持 CuteDSL 路径（无回退）。

精度口径的对齐。 索引器的 query 在旋转后刻意降到 bf16 再升回 fp32，是为了让索引器的数值精度与主通路的精度口径一致。改动这个精度锚点会改变索引器分数分布，进而可能改变 Top-K 的选择结果。

从 DSA 到 vLLM：一次完整的调用链

把前面的所有片段串起来，DSA 在 vLLM 里一次调用的完整链路是这样的：

准备阶段。 元数据构造器 DeepseekV32IndexerMetadataBuilder 为本批次算好索引器需要的元数据：分块计划（预填充阶段会把长序列切成若干 chunk）、每行的起止位置（cu_seqlen_ks/cu_seqlen_ke）、每行的局部序列长度等。索引器自己的 KV 缓存的块表与槽位映射在这一步准备好。

打分准备。 fused_indexer_q_rope_quant 用 _fused_indexer_q_rope_quant_kernel 把 bf16 的索引器 query 一次性完成旋转、量化、scale 提取，并把 q_scale 折进索引器权重，产出 FP8 的 query 与调整后的权重。

写入索引器缓存。 当前批次新产生的 $\mathbf{k}^I$ 按槽位写进索引器自己的 KV 缓存。

打分。 逐 chunk 调用 DeepGEMM 的 fp8_fp4_mqa_logits，得到压缩形式的逐行 logits（每行从 cu_seqlen_ks[i] 开始）。

Top-K。 top_k_per_row_prefill 把每行的 logits 转成 $k$ 个下标，写进 topk_indices_buffer。这一步可能带两级的候选块筛选（在某些架构上）。

（分布式时）跨 rank 合并。 _merge_dcp_topk_global 做候选打包 + all-gather + 全局稳定 Top-K。

地址翻译。 triton_convert_req_index_to_global_index 把逻辑位置翻译成物理槽位，同时产出每行的有效计数 topk_length。

稀疏注意力。 _bf16_flash_mla_kernel 拿着物理索引和有效计数，在摊平成行视图的 KV 缓存上做 gather 式读取与注意力计算。

这条链路上有六个以上的内核，涉及三个阶段（打分、索引、注意力），横跨三个内存结构（索引器缓存、注意力缓存、索引缓冲区）。DSA 在论文里是三个公式，在框架里是一条完整的数据管线。

顺带说一句，站内已经有一篇 NSA 原生稀疏注意力（二）：Triton Kernel 设计与 27B 模型实验 讨论过稀疏注意力的 Triton 内核设计。NSA 和 DSA 的稀疏内核共享大量工程手法（分块、掩码、跳跃），但两者的稀疏来源不同：NSA 的稀疏来自训练时就固定的三分支结构，DSA 的稀疏来自推理前动态算出的索引。前者是”结构决定的稀疏”，后者是”数据决定的稀疏”，这导致 DSA 多出了整整一条”索引器 + 地址翻译”的管线，而 NSA 不需要。这个差异是理解两篇文档时最值得抓住的一点。

小结

把这篇的要点重新串一遍。

• 问题的形状。 注意力有两条独立的成本曲线：prefill 的 $\mathcal{O}(L^2)$ 算分，和 decode 的 $\mathcal{O}(L)$ 字节搬运。DSA 砍的是”绝大多数历史位置贡献接近于零”这个事实带来的浪费，把主注意力的复杂度从 $\mathcal{O}(L^2)$ 降到 $\mathcal{O}(Lk)$，$k=2048$。
• 索引器的分数函数 $I_{t,s}=\sum_j w^I_{t,j}\cdot\text{ReLU}(\mathbf{q}^I_{t,j}\cdot\mathbf{k}^I_s)$ 的每一处设计都是为”便宜”服务的：去掉 softmax 让它逐位置独立、可任意分块；ReLU 换掉 exp 换吞吐；$w^I$ 提供逐头加权；key 侧单头让缓存最小。
• 三个关键数字：$d^I=128$ 决定索引器缓存每 token 的字节数（加 scale 共 132 B）；$H^I$ 决定打分的计算常数，也是索引器 $\mathcal{O}(L^2)$ 的来源；$k=2048$ 决定主注意力看多少历史。
• DSA 必须挂在 MLA 的 MQA 模式上，因为稀疏需要一个”多个 query 头共享同一份数据”的寻址单元，MLA 的潜向量正是这个单元。稀疏只有落在数据搬运的粒度上才有价值。
• 索引器是被单独训练的模块，有自己的损失 $\mathcal{L}^I$，与主模型损失分离；引入方式是在 DeepSeek-V3.1-Terminus 上做两阶段继续预训练（dense warm-up + sparse training），不重新预训练主干。
• vLLM 的实现链路：融合内核（RoPE + FP8 量化 + scale 折叠）→ DeepGEMM 压缩 logits 打分 → 逐行 Top-K → 逻辑位置到物理槽位的翻译（含有效计数）→ 稀疏注意力内核的 gather 式读取。
• 稀疏在 kernel 里的样子：把”顺序读 tile”换成”查表读”，代价是失去 cp.async、失去多级流水、掩码只能同步加载；负索引统一落到行 0、靠掩码抹掉贡献，以避免分支发散。
• 两条实现路线：llama.cpp 用”掩码压缩成索引数组”，索引数组是 $\mathcal{O}(n_{tokens}\times k)$；vLLM 用”Top-K 重写掩码”，复用现有 FlashAttention 内核，省掉计算与 KV 读取但不省掩码本身的平方级占用。
• 收益的边界：DSA 没有消灭平方项，只是把它从贵的注意力转移到了便宜的索引器；它省不掉索引器缓存、地址翻译、gather 的效率损失、以及分布式下的全局 Top-K 通信；稀疏路径有上下文长度门槛（K->ne[1] >= max(4096, 2*n_gather)）。
• 部署清单：约 825 MB 的索引器预填充工作区、132 B/token 的索引器缓存、$k$ 在 CuteDSL 路径上只能取 512/1024/2048、平台能力决定哪些路径可用。

参考资料

1. DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models
2. DeepSeek-V3.2 论文 HTML 版（含 DSA 架构图与推理成本图）
3. DeepSeek-V3.2-Exp 官方推理参考实现
4. vLLM vllm/v1/attention/backends/mla/indexer.py：索引器后端与元数据构造
5. vLLM vllm/model_executor/layers/sparse_attn_indexer.py：索引器算子与融合量化内核
6. vLLM vllm/v1/attention/backends/mla/sparse_utils.py：逻辑索引到物理槽位的转换
7. vLLM vllm/v1/attention/backends/mla/flashmla_sparse.py：FlashMLA 稀疏后端
8. vLLM vllm/model_executor/models/deepseek_v2.py：DeepSeek 模型的索引器接入
9. llama.cpp src/llama-graph.cpp：DSA 路径的掩码重写与稀疏注意力构建
10. llama.cpp ggml/src/ggml-cuda/fattn.cu：稀疏索引压缩与稀疏路径判据
11. llama.cpp ggml/src/ggml-cuda/fattn-mma-f16.cuh：稀疏路径的 KV tile gather
12. FlashMLA：MLA 解码的高效 kernel 实现
13. Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
14. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model（MLA 原始论文）
15. Fast Transformer Decoding: One Write-Head is All You Need（MQA 原始论文）