DFlash 一次并行出整块 draft，但块内各位置互不商量，而且整块都要送进 target model。DSpark 只做两件事——draft 侧加一个很轻的串行采样循环让位置之间互相看得见，target 侧加一个裁剪决定这一轮到底验多长。采样规则和接受规则都没变，所以依然无损。

归属：DeepSeek-AI 与北京大学，DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation，arXiv:2607.05147，作者含 Wenfeng Liang（梁文锋）。代码 DeepSpec，权重 DeepSeek-V4-Pro-DSpark，CC BY 4.0。

前置：只讲相对 DFlash 的增量。块扩散 draft、KV 注入、拒绝采样见 DFlash 投机解码。

一、出发点

DFlash 的链条是：draft model 一次前向算出整块 γ 个位置的分数 → 每位置各自取 top-1 → 整块送去验证。这条链有两个口子，正好在两头：

• draft 侧：每个位置独立取 top-1。它是在对「不确定的前驱」取期望，所以可能拼出并不存在的组合，越靠块尾越容易错。
• target 侧：整块都验。尾部位置本来大概率被拒，验它们等于白占 target model 的 batch 容量。

DSpark 的两处改动分别堵这两个口子。

图 1：一轮五步。标绿（串行采样）和标橙（裁剪）是全部新增，其余沿用 DFlash。

二、串行采样：把块内各位置串起来

2.1 它改了什么

draft model 里那层完全不动——仍然是块内双向注意力、一次前向出整块的分数，延迟与 γ 几乎无关。变的只有「怎么从这些分数得到候选」：

• DFlash：每个位置各自取 top-1，位置之间不通信；
• DSpark：从左到右逐位采样，采到第 k 位时，用刚刚采出的第 k−1 位去修正第 k 位的分数，再采样。

修正量就是一个加在分数上的偏置：

$$p_k(v)\;\propto\;\exp\big(U_k(v)+B_k(v)\big)$$

$U_k$ 是 draft model 在每位给出的分数，$B_k$ 是 draft 侧串行头给的偏置。这是全文唯一的算法公式，读法就一句：draft model 的分数 + 前文修正，再 softmax 采样。

偏置怎么算出来的？查表 + 一次小矩阵乘：$B_k$ 只依赖紧邻的前一位 token，所以维护一张「token → 向量」的表，查出来乘个小矩阵，就得到加在词表上的偏置。没有额外的 Transformer 前向——这就是它「轻」的原因。

论文还试了能记住整块前缀的递归版本（RNN 头），结论是收益只多一点、主要出现在很长的块上，但实现复杂不少，所以生产用的是只看前一位那版。

2.2 例子

假设上文已定，接下来该写「很大程度」。draft model 给位置 2 的 top-1 是「可」0.40，「大」只拿到 0.31 —— 因为它不知道位置 1 会采出什么：若位置 1 是「稍」，这里本该是「微」。两个模式被平均掉了，所以纯并行 draft 在这里容易采偏。

串行采样把「不确定的前驱」换成「已经采出的前驱」：

• 第 1 步没有前文，直接用原始分数，采到 很；
• 第 2 步用「很」查表得到偏置，抬高「大」、压低「可」，采到 大（0.62 vs 0.19）；
• 第 3 步用「大」的偏置抬高「程」，采到 程（0.88）；
• 第 4 步同理采到 度 —— 但这一位「猜得准不准」是另一回事，见第三节。

一旦「很」被真正采出来，「稍」那条分支就不存在了，概率质量自然集中回「大」。DFlash 在这里可能采到「可」或「慢」，DSpark 不会。

同时这个循环顺手产出每个位置的条件存活概率 $c_k$——在第 k 位之前全部被接受的前提下、第 k 位还能活下来的概率。注意它的输入只有 draft model 隐状态和前一位 token。四个位置依次给出 $c_1=0.95$、$c_2=0.92$、$c_3=0.85$，而 $c_4$ 只有 $0.30$。

图 2：上例的完整推演（上半部分是串行采样，下半部分是接着做的裁剪）。

三、裁剪：决定这一轮验多长

3.1 为什么要裁

DFlash 把整块送进 target model。但尾部位置本来大概率被拒，验它们纯属浪费；而且在高并发下，这些浪费的验证位会挤占别的请求的 batch 容量。

论文还有一个观察：结构化任务（数学、代码）的接受率明显高于开放式闲聊，可是验证花费是一样的。所以「写死一个验证长度」天然不合理。

3.2 三个量各自的算法

裁剪要最大化系统总吞吐。所有参与决策的量只有三个，且都能从上一步的输入算出来：

其中 $B$ 是这次前向的 batch 大小（$B = 1$ 个 anchor + 已准入的 draft 位数），$\mathrm{SPS}(B)$ 是引擎在该 batch 下的 step 速度（在每次前向B个token时，每秒跑多少次前向）——它随 $B$ 单调不增，引擎初始化时 profile 一次存成表。

于是判据是：逐个纳入，$\Theta$ 不再上升就停。

3.3 接上例子，把每个数字算一遍

沿用上例，四位 draft 的存活概率取 $c_1,c_2,c_3,c_4 = 0.95,\,0.92,\,0.85,\,0.30$，并取一条示意的 step 速度曲线（真实曲线由引擎实测得到、有数千档，这里取整数方便逐格验算）：

$$\mathrm{SPS}(1,2,3,4,5,\dots)=1000,\ 520,\ 400,\ 340,\ 280,\ \dots\ \text{steps/s}$$

它随 $B$ 单调不增。先把连乘算出来：$s_1=0.950$，$s_2=0.950\times0.92=0.874$，$s_3=0.874\times0.85=0.743$，$s_4=0.743\times0.30=0.223$。然后逐个纳入：

判决列只做一件事：新的 $\Theta$ 有没有超过目前的最好值。

四、采样流程

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

def dspark_cycle(prefix, target, draft, head, scheduler, gamma=16):
    """一个解码轮次。gamma 为上界，本轮实际验多长由 scheduler 决定。"""

    # ① target model 前向（本来就做）：产出 anchor，并留下上下文特征
    x0, H_ctx = target.forward(prefix)

    # ② 并行段：一次前向算出整块 draft 分数（与 DFlash 相同，也是最重的一步）
    #    输入 anchor + (gamma-1) 个 mask，得出 gamma 个位置的 draft model 分数
    hidden, U = draft.parallel_backbone([x0] + [MASK] * (gamma - 1), H_ctx)

    # ③ 串行段：左→右逐位采样，每位只用「刚刚采出的前一位」修正自己
    cand, conf = [], []
    prev = None
    for k in range(gamma):
        bias = head.bias(prev) if prev is not None else 0
        token = sample(softmax(U[k] + bias))     # p_k ∝ exp(U_k + B_k)
        cand.append(token)
        conf.append(head.confidence(hidden[k], prev_token=token))
        prev = token          # 只传一位：循环是串行的，但每步极轻

    # ④ 裁剪：按累积存活概率贪心准入，吞吐不再上升就停
    keep = scheduler.admit(conf)

    # ⑤ target model 只并行验证裁剪后的前缀（接受规则与 DFlash 完全一致）
    #    对齐关系：第 i 个 logits 预测的是 cand[i]；accepted = 被接受的前缀长度
    logits = target.forward_parallel(prefix + [x0] + cand[:keep])
    accepted = longest_accepted_prefix(cand[:keep], logits)   # 从左碰到第一个拒绝

    # ⑥ 提交：接受前缀 + 被拒位置的残差重采（无拒绝时即 bonus token，也是下一轮 anchor）
    anchor_next = sample_residual(logits[accepted])
    prefix += [x0] + cand[:accepted] + [anchor_next]
    target.rollback_kv(keep=len(prefix))   # 丢弃未接受 draft 留下的 KV
    return prefix

调度器本体：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

def admit(conf, base_batch=1):
    """
    conf        : 逐位条件存活概率 c_1..c_gamma
    base_batch  : 不验任何 draft 时占的 batch 大小（anchor 自己算 1 个位置）
    返回        : 本轮该验的 draft 位数

    依赖一个全局函数 speed(B)：验证 batch 为 B 时的相对 step 速度，
    引擎初始化时 profile 一次存成表，对 B 单调不增。
    """
    # 位置 k 能被验到，前提是前面每位都被接受 → 累积连乘
    survival, s = [], 1.0
    for c in conf:
        s *= c
        survival.append(s)

    # 起点：一位 draft 都不验，期望接受 = anchor 自身
    batch, expect = base_batch, float(base_batch)
    best = speed(batch) * expect

    # 贪心：多验第 k 位，多赚 survival[k]，代价是 batch +1
    # 注意 batch / expect 随循环累加，就是当前已准入的状态
    keep = 0
    for k in range(len(conf)):
        batch += 1
        expect += survival[k]
        score = speed(batch) * expect
        if score <= best:      # 吞吐不再上升就停（见 §3.4：这一步是正确性要求）
            break
        best, keep = score, k + 1

    return keep

五、效果

参考资料

DSpark 本体

• DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation — arXiv:2607.05147（CC BY 4.0）
  • DFlash 背景：§2.2；半自回归与两种串行头：§3.1；置信头与 STS：§3.2.1；调度器与 Algorithm 1：§3.2.2
  • 生产适配（异步、零开销调度、变长 kernel）：§5.2–5.3；线上结果与局限：§5.4；选择偏差反例：Appendix A
• 代码 deepseek-ai/DeepSpec　权重 deepseek-ai/DeepSeek-V4-Pro-DSpark

前代与对照

• Chen 等，DFlash: Block Diffusion for Flash Speculative Decoding — arXiv:2602.06036
• Li 等，EAGLE-3 — arXiv:2503.01840
• Leviathan 等，Fast Inference from Transformers via Speculative Decoding — arXiv:2211.17192
• DeepSeek-AI，DeepSeek-V3 Technical Report（MTP-1 基线）— arXiv:2412.19437

站内相关

• DFlash 投机解码：块扩散草稿与 KV 注入 —— 本文前置
• 投机解码发展 —— 三步骤与 EAGLE 三代差异

一句话：把草稿模型从「自回归生成器」换成「块扩散填充器」—— 一次前向并行解出一整块 mask 位置；同时把目标模型的融合隐状态作为 KV 注入到草稿的每一层。输出分布与目标模型逐字解码完全一致。

归属：DFlash 出自 Z Lab / UCSD（Chen, Liang, Liu），论文 arXiv:2602.06036，ICML 2026。

符号约定

一、瓶颈与成本模型

Decode 阶段每个 token 都要把全部权重从 HBM 搬进计算单元一次，而这一步算术强度极低 —— 计算单元大部分时间在等数据（memory-bound）。关键观察是：权重搬运这一趟的成本已经付了，却只换来 1 个 token。

投机解码的基本形式：用便宜的草稿模型 $M_d$ 提出若干候选，交给目标模型 $M_t$ 一次并行验证。DFlash 论文（§3.1）沿用的平均每 token 延迟与加速比为

$$L = T_{\text{draft}} + \frac{T_{\text{verify}}}{\tau}, \qquad \eta = \frac{L_{\text{target}}}{L}$$

其中 $T_{\text{verify}}$ 是验证开销，$\tau \in [1, \gamma+1]$ 是每轮期望接受长度（含目标模型顺带产出的 bonus token），$\eta$ 是整体加速比。$L_{\text{target}}$ 是逐字解码的每 token 延迟。这个式子把设计空间一次说清，论文原文也强调：要提高加速比，只能抬高 $\tau$ 或压低 $T_{\text{draft}}$（DFlash §3.1）。

而在自回归草稿里，这两件事互相冲突。

自回归草稿：草稿是逐 token 生成的，第 $k$ 个 token 必须等第 $k-1$ 个算完，于是

$$T_{\text{draft}} = \gamma \cdot t_{\text{step}}$$

成本随 $\gamma$ 线性增长。论文指出：为了让延迟可控，自回归草稿被迫使用极浅的架构（EAGLE-3 只有一层），这严重限制了草稿质量 —— $\tau$ 随 $\gamma$ 增加会迅速饱和。

块扩散草稿：$\gamma$ 个位置在同一次前向内并行解出，

$$T_{\text{draft}} = t_{\text{parallel}}$$

论文的表述是：对中等块大小，$t_{\text{parallel}} \ll \gamma \cdot t_{\text{step}}$，$T_{\text{draft}}$ 基本对 $\gamma$ 不敏感。这改变了设计空间 —— 草稿模型可以做得更深而不显著增加延迟。论文实测：一个 5 层 DFlash 草稿生成 16 个 token，延迟比单层 EAGLE-3 生成 8 个 token 更低，同时接受长度更高。

二、块扩散草稿：为什么能一次猜好几个

这是全文核心。先看清自回归为什么必须串行。

自回归的串行是结构性的

自回归草稿第 $k$ 步的输入是 $(h_{k-1}, \text{embed}(x_{k-1}))$ —— 第 $k$ 步依赖第 $k-1$ 步的输出。这是一条函数依赖链，硬件再并行也打不断。

扩散把它换成「条件」而非「依赖」

块扩散把「预测下一个 token」改写成「填补被 mask 的位置」。给定块

$$[\,a,\ m_1,\ m_2,\ \dots,\ m_{B-1}\,]$$

其中 $a$ 是 anchor，$m_i$ 是 mask。所有 mask 位置共享同一份条件：

• 目标上下文特征 $H_c$（见第三节）；
• 块内所有位置的嵌入 $H_b$（含 $a$ 与其他 mask）。

关键在于：$m_i$ 之间没有互相依赖，它们只是共同以彼此为条件。注意力天然能同时互相可见 —— 只要块内做双向注意力、不加因果掩码。

于是 $B-1$ 个位置变成一次并行求解，而不是 $B-1$ 次串行求解。

注意力模式

设草稿第 $i$ 层，块内嵌入为 $H_b$，上下文特征为 $H_c$。该层的 Q 与拼接后的 KV 为：

$$Q = W_Q\,H_b, \qquad K = [\,K_c \;;\; K_b\,], \qquad V = [\,V_c \;;\; V_b\,]$$

其中

$$K_c = W_c^{K} H_c,\quad K_b = W_b^{K} H_b,\qquad V_c = W_c^{V} H_c,\quad V_b = W_b^{V} H_b$$

三个要点：

• $Q$ 只来自块内（上下文特征不产生 query）；
• $H_c$ 只作为额外的 KV 条目，绕过草稿的 Q 投影、输出投影、self-attention 更新与 FFN（论文 §A.3）；
• 块内双向可见（推理时无因果掩码），因此 $B-1$ 个位置互为条件、同时出结果。

一个常见的误解：这里没有「多步去噪」

经典扩散需要多步迭代去噪，这也是扩散模型做端到端生成时慢的原因。但 DFlash 的草稿是一步出整块的。因此草稿延迟是 $t_{\text{parallel}}$ 而非若干步去噪之和。

代价：位置之间「不商量」

并行填充的固有代价是各位置独立决策，块内可能不连贯。这是「用并行换一致性」的显式权衡，不是实现瑕疵 —— 它直接引出了 DFlash 2（第六节）。

三、KV 注入：让目标信号贯穿每一层

草稿模型很小，猜不准的根本原因是缺乏上下文信息。前代方法（EAGLE 系列）已经会利用目标模型隐状态，但做法是输入融合：把目标特征与草稿 token 嵌入拼接后作为输入喂进去一次。

论文对这种做法的批评很直接：信号只在入口出现，随着草稿层数增加会被逐步稀释，导致加深草稿带来的接受长度收益递减。所以，DFlash 改为 KV 注入：从目标模型均匀抽取 $L$ 层隐状态（论文实现取第 2 层到倒数第 3 层之间均匀 5 层），拼接后一次投影：

$$H_c = \text{RMSNorm}\big(W_c\,[\,H^{(l_1)} ; \dots ; H^{(l_L)}\,]\big)$$

然后把 $H_c$ 注入到每一个草稿层的 K/V（即上节的 $K_c, V_c$）。论文的解释是：这让每一层都能重新看到目标上下文，信号不被深度稀释，因此接受长度可以随草稿层数持续增长。

复杂度与参数量（论文 §A.3）：唯一新增的参数化组件是共享投影 $W_c \in \mathbb{R}^{D \times LD}$ 加若干草稿层；论文给出的量级是约占目标模型 1.6%。$W_c$ 本身很小 —— $D{=}2048$、$L{=}5$、bf16 下约 42 MB，相对 70 GB 级目标模型可忽略。词嵌入表与 LM head 与目标模型共享并冻结。

消融证据

论文 Table 9 用同一套设定换掉注入方式（Qwen3-4B，$N{=}5$，$B{=}8$，单元格为 τ / 加速）：

论文还做了一个更干净的交叉验证：把 DFlash 草稿改回自回归、只保留 KV 注入（DFlash-AR），接受长度依然全面高于 EAGLE-3-5L。这说明「并行」与「注入」两个贡献是独立可测的，不是互相依赖的巧合。

四、推理流程

一轮四个动作：draft → verify → accept → rewind。

伪代码

草稿模块前向（KV 注入）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

def draft_forward(block_ids, H_c):
    """block_ids: [a, MASK, ..., MASK]，长度 B
       H_c: 目标上下文融合特征，所有草稿层共享"""
    H_b = embed(block_ids)                    # [B, D]

    for layer in draft_layers:                # N 层，每层都注入
        Q   = layer.W_Q(H_b)                  # [B, D]，只有块内产生 query
        K_c = layer.W_c_K(H_c)                # 上下文只进 K/V
        V_c = layer.W_c_V(H_c)
        K_b = layer.W_b_K(H_b)
        V_b = layer.W_b_V(H_b)
        # 块内双向：不加因果掩码
        H_b = layer.attn(Q, concat(K_c, K_b), concat(V_c, V_b)) + layer.ffn(H_b)

    return lm_head(H_b)                       # 共享目标模型的 LM head

一个完整推理周期：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

def dflash_cycle(prefix, target, draft, B=8):
    # ① 目标前向（本来就要做）：产出 anchor a，并缓存 L 层隐状态
    a, hidden = target.forward(prefix)
    H_c = fuse(hidden[selected_layers])       # 提前物化，便于前缀共享

    # ② 草稿：一次前向填满整块
    block = [a] + [MASK] * (B - 1)
    cand = argmax(draft_forward(block, H_c), dim=-1)[1:]     # B-1 个候选

    # ③ 验证：目标一次并行前向，同时得到每个位置的真实分布
    logits = target.forward_parallel(prefix + [a] + cand)

    # ④ 取最长一致前缀（贪心情形即逐位比对）
    n = longest_common_prefix(cand, argmax(logits, dim=-1))

    # 被拒位置用目标分布重采样；末位顺带产出 bonus token
    prefix += [a] + cand[:n] + [sample(logits[n])]

    # 回滚被拒草稿的 KV，避免污染下一轮上下文
    draft.rewind_kv(keep=len(prefix))
    target.rewind_kv(keep=len(prefix))
    return prefix

一个工程细节（有明确出处）

DFlash 把 $H_c$ 提前投影成 KV，然后将这份投影后的 KV 存入草稿模型的 KV Cache 中，而不是把原始$H_c$存进 cache。SGLang / Modal / Z Lab 的联合博客给出的两个理由是：不占用 KV cache 空间去存原始隐状态；且前缀相同的请求可以直接复用这份 cache（radix cache）。

六、改进

DFlash 2：补上并行的两个固有短板

2026 年 8 月，Inco AI（DFlash 原班人马）发布 DFlash 2。它不改变并行范式，而是修补并行草稿的两个弱点。以下数据均来自该博客。

短板一：块内不连贯。 每个位置独立取 top-1，可能出现语义打架的块（例如相邻两位都选中同一个词）。但博客给出的证据表明问题在选择而非预测：

也就是说，正确答案几乎总在候选表里；一个总能从 top-16 里选对的「oracle」能把接受长度从 4.27 抬到 6.79 —— 这段差距是纯粹的选择余量。

DFlash 2 的做法：保留每位置 top-$k$ 候选，对相邻候选对打一个低秩双线性分数

$$S_t(a,b) = U_t(b) + \big\langle A(a) \odot H(h_t),\; B(b) \big\rangle$$

第一项 $U_t(b)$ 是草稿自身的 logit，只依赖于当前位置 t 和候选 b ，与前一个候选 a 无关。
它通常表示 b 本身在该位置有多好，第二项衡量 $b$ 跟在 $a$ 后面有多合理。

短板二：块尾衰减（suffix decay）。 博客的诊断是：DFlash 的注意力身兼两职 —— 读块外上下文 + 建块内依赖，而后者被逐渐挤掉（块内注意力占比从 Layer 1 的 30% 降到 Layer 5 的 8%）。于是即使选择器完美，块尾的候选本身也在变差。

直觉解法是加深（15 层确实有效），但博客指出 15 层要多 3 倍参数、增加 15.2% 的周期延迟，太不划算。DFlash 2 改为分工：注意力继续读上下文，另加一个专管块内一致性的卷积模块，约 +3% 参数、+0.7% 延迟，块尾命中率接近 15 层模型。

合起来：每次验证多产出 16~25% 的 token，额外周期延迟约 1%。

七、小结

1. 成本模型决定设计空间：加速要么抬高 $\tau$，要么压低 $T_{\text{draft}}$。自回归草稿让这两者冲突，块扩散让它们解耦。
2. 块扩散为什么能一次出多个：它把「token 之间的函数依赖」换成「mask 位置之间的共同条件」。依赖是串行的，共同条件是可并行的 —— 块内双向注意力把这件事表达出来。
3. KV 注入解决加深的收益问题：目标特征作为额外 KV 条目贯穿每一层，信号不被深度稀释，接受长度才能随层数增长。论文的消融显示它与「并行」的贡献互相独立。
4. 这里没有多步去噪：块内所有 mask 位置在单次前向里一起解出，这是它能当草稿的前提。
5. 并行的代价是块内不一致，需要选择器或专用模块兜底。DFlash 2 证明这块短板可以用极小的参数代价补掉 —— 补短板比堆容量划算。

参考资料

DFlash 本体

• Chen, Liang, Liu，DFlash: Block Diffusion for Flash Speculative Decoding（Z Lab / UCSD，ICML 2026）— arXiv:2602.06036
  • 成本模型公式：§3.1；注意力与并行草稿：§4.1；KV 注入与内存开销：§4.1、§A.3；无目标特征消融：§A.2
  • 结果表：Table 1（主结果）、Table 3（SGLang 并发）、Table 6（层数）、Table 8（块大小）、Table 9（注入方式）
• 代码与模型：github.com/z-lab/dflash；HuggingFace DFlash 集合
• LMSYS / Modal / Z Lab，The next generation of speculative decoding: DFlash and Spec V2 — lmsys.org（SGLang 落地细节：KV 物化、前缀共享、融合 kernel）

DFlash 2

• Inco AI，DFlash 2: Keep Drafting Parallel — inco.ai/blog/dflash2
• 对照方法：Domino — arXiv:2605.29707；DSpark — arXiv:2607.05147

前代与背景

• Leviathan 等，Fast Inference from Transformers via Speculative Decoding — arXiv:2211.17192
• Cai 等，Medusa — arXiv:2401.10774
• Li 等，EAGLE — arXiv:2401.15077；EAGLE-2 — arXiv:2406.16858；EAGLE-3 — arXiv:2503.01840
• An 等，PARD — arXiv:2504.18583；Li 等，DiffuSpec — arXiv:2510.02358；Sandler 等，SpecDiff-2 — arXiv:2511.00606
• Arriola 等，Block Diffusion — arXiv:2503.09573
• Tencent，FastMTP — arXiv:2509.18362

工程落地

• 小米 MiMo × TileRT，MiMo-V2.5-Pro-UltraSpeed — mimo.xiaomi.com
• 开源权重：XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash

投机解码与 EAGLE：步骤、发展历史与三代差异

一句话原理：用便宜的小模型先猜几个 token，再让大模型用一次前向把它们全部验完。 猜对了就白赚，猜错了就从出错的地方重新采样 —— 最终输出的分布与大模型自己逐字生成完全一致。

一、发展历史

主线是 EAGLE 这一条，演进逻辑非常好记：

1. EAGLE 解决「草稿用什么表示」—— token 换成特征；
2. EAGLE-2 解决「节点预算怎么花」—— 静态树换成动态树；
3. EAGLE-3 解决「训练和推理不一致」—— 单步训练换成 training-time test。

二、基础三步骤

1. 草稿（Draft）：小模型 Mq 自回归猜 4 个候选 token。一次只出一个，但模型小，4 次前向也不心疼。
2. 验证（Verify）：把「上文 + 4 个候选」整段喂给大模型 Mp，一次前向并行算出每个位置的真实概率分布。原本要 4 次串行前向，现在 1 次。
3. 接受（Accept）：从左往右逐个判定，遇到第一个不一致就截断；如果全部通过，再从大模型的分布里补采 1 个 token（bonus token）。

以「中国的首都是」为例：草稿给出「北京 / 是 / 中国 / 的」，大模型验证后接受「北京」「是」，拒绝「中国」，并从残差分布补出「一座」—— 一轮产出 3 个 token，只花了 1 次大模型前向。

为什么这是无损的

判定只有两条规则：

• 大模型比小模型更想要它（p ≥ q）→ 直接收下；
• 小模型高估了它（p < q）→ 以 p/q 的概率收下，否则拒绝。

拒绝之后不是重来，而是从「残差分布」采样：把小模型多给出去的那部分概率扣掉，补给大模型真正想要的 token。一进一出刚好抵消，所以每个位置输出各 token 的概率，与大模型自己生成的分布完全相同。

注意：无损指的是分布一致，不是逐位一致，两次运行不会输出完全相同的结果。

什么时候能加速

一轮能多产出多少 token，只看两个量：

• 接受率 α：草稿猜对的比例，越高越好；
• 草稿相对开销 c：草稿一次前向比大模型一次前向便宜多少，越小越好。

草稿 token 数不是越多越好：多猜一个，能被全部接受的概率按 α 连乘迅速衰减，而草稿开销线性增长，所以存在一个最优长度。直觉上的临界条件是 α > c —— 接受率还不如草稿的便宜程度，就没得赚。

这也解释了为什么朴素的「7B 草稿 + 13B 目标」几乎零提升：草稿不够便宜，接受率也不够高。

三、EAGLE 的关键改动：从猜 token 到猜特征

朴素的草稿模型要单独训一个，还得让它和目标模型分布接近，很难同时做到「小」和「准」。EAGLE 的观察是：在特征层面做自回归，比在 token 层面更容易。 token 经过 softmax 之后信息已经被「压硬」了，而连续特征保留了更丰富的语义。

但特征级自回归有个麻烦：

位置 t+1 的特征长什么样，取决于位置 t+1 的那个 token 是什么 —— 而这个 token 还没定。只拿 f(t) 去猜 f(t+1)，是一对多的病态映射，学不准。

EAGLE 的解法很直接：把「超前一步」的 token 也一起喂进去。输入是「上一位特征的向量」加上「已经确定的真 token 的 embedding」，输出下一个特征和目标位置的 token 分布。条件一给，特征基本就被钉住了。

这个改动效果很大。论文消融实验（Vicuna-7B，MT-bench）：

草稿模型有多小

• Embedding 和 LM Head 直接复用目标模型的，冻结不动，只训练中间一个 FC + 一层 decoder。
• 特征取自 LM Head 的输入（最后一层 Transformer 的输出）。论文叫它 second-to-top-layer feature，是把 LM Head 当作「顶层」来数的，别误读成「倒数第二层」。FC([embed(xt),ft])
• 可训练参数量约 0.24B（7B 目标）到 0.99B（70B 目标）；70B 的草稿头在 4×A100 上训 1~2 天。

为什么要树，以及树怎么一次验完

只猜一条链的话，中间错一个后面全废。EAGLE 让每个节点保留 top-k 个分支，长成一棵树。

难点是：这棵树怎么用一次前向验证？ 答案是树形注意力掩码 —— 把树展平成一条序列、给每个节点编号，然后规定每个节点只能看见自己 + 自己的所有祖先（看不见兄弟节点）。这样一次前向就能同时算出所有候选路径的概率。

EAGLE-1 默认 25 个节点、深度 5。因为同层节点互不依赖、可以并成一个 batch，这棵树只花 5 次草稿前向 + 1 次目标前向。

容易搞错的一点：网上流传的「EAGLE 用 KV Cache 重排让祖先相邻」并不在论文里。论文只有树形注意力掩码；真正的 KV 处理发生在接受之后（把被接受路径对应的 KV 行拷到一起、丢弃其余分支），属于代码实现细节。

四、一轮完整流程与伪代码

五个阶段循环执行：取特征 → 草稿建树 → 一次前向验证 → 沿树接受 → 更新 KV。

下面用伪代码把整条链路串起来（去掉工程细节，只保留主干）。

1
2
3
4
5
6

# ============ 对照：基线自回归解码 ============
def baseline_decode(prompt, Mp, n_tokens):
    for _ in range(n_tokens):
        p = Mp.forward(prompt)        # 大模型一次前向
        prompt.append(sample(p))      # 一次只能产出 1 个 token
    return prompt                     # n 个 token = n 次串行前向

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

# ============ 朴素投机解码：单链草稿 ============
def spec_decode(prefix, Mp, Mq, gamma=4):
    # 1) 草稿：小模型自回归猜 gamma 个 token
    draft, q_prob = [], []
    for _ in range(gamma):
        q = Mq.forward(prefix + draft)
        tok = sample(q)
        draft.append(tok)
        q_prob.append(q[tok])

    # 2) 验证：大模型一次前向，算出 gamma+1 个位置的真实分布
    p_prob = Mp.forward_multi(prefix + draft)

    # 3) 接受：从左往右，第一个不通过就截断
    out = []
    for i, tok in enumerate(draft):
        if uniform() < min(1.0, p_prob[i][tok] / q_prob[i]):
            out.append(tok)                                     # 收下
        else:
            out.append(sample(residual(p_prob[i], q_prob[i])))  # 拒绝 -> 残差重采样
            return out
    out.append(sample(p_prob[gamma]))    # 全部接受 -> 白赚一个 bonus token
    return out

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

# ============ EAGLE：特征级草稿 + 树 + 一次验证 ============
def eagle_round(prefix, Mp, Mq, depth=5, topk=10, budget=60):
    # 1) 取特征：大模型最后一次前向留下的隐状态（也就是 LM Head 的输入）
    f = Mp.feature_cache[prefix]

    # 2) 草稿建树：同一层的节点拼成一个 batch，所以 depth 次前向就能长 depth 层
    tree = Tree(root=prefix[-1])
    frontier = [Node(token=prefix[-1], feature=f, score=1.0)]
    for _ in range(depth):
        # 关键：输入 = (token embedding 并上 上一位的特征向量)
        f_hat, logits = Mq.forward(
            [(embed(n.token), n.feature) for n in frontier],
            mask=ancestor_mask(frontier),       # 只看自己 + 自己的祖先
        )
        for node, feat, dist in zip(frontier, f_hat, logits):
            for tok in topk_of(dist, topk):
                # 路径分数 约等于 整条路径被全部接受的概率
                tree.add_child(node, tok, feat, node.score * dist[tok])
        # 动态树（EAGLE-2）：全局挑分数最高的节点继续扩张，超预算就剪掉最低分
        frontier = tree.pick_top_nodes(count=topk, budget=budget)

    # 3) 验证：大模型一次前向，算完整棵树每个节点的真实分布
    p_all = Mp.forward(tree.flatten(), mask=tree.attention_mask())

    # 4) 沿树接受：找出最长可接受前缀
    accepted, node = [], tree.root
    while True:
        child = tree.best_child(node)
        if child is None:
            break
        p, q = p_all[child], child.draft_prob
        if uniform() < min(1.0, p[child.token] / q):
            accepted.append(child.token)              # 接受，继续往下走
            node = child
        else:
            accepted.append(sample(residual(p, q)))   # 拒绝 -> 残差重采样后结束
            break

    # 5) 收尾：补 1 个 bonus token，压缩双方的 KV Cache，回到第 1 步
    accepted.append(sample(p_all[node]))
    Mp.compact_kv(keep=node)
    Mq.compact_kv(keep=node)
    return accepted

一句话，每次目标模型的forward可以验证草稿模型产生的一棵树。

五、三代之间的变化

EAGLE-2：换成动态树。 静态树隐含假设「接受率只取决于位置」，但实测发现同一位置上不同分支的方差很大 —— 接受率是上下文相关的。论文举的例子：10+2= 后面几乎必然是 1，再铺一堆候选纯属浪费；而 10+2 很难猜，这时才该多分配节点。

EAGLE-2 用草稿模型的置信度来估计接受率（实测校准得很好：置信度低于 0.05 的候选，真实接受率约 0.04；高于 0.95 的约 0.98），给每个节点算一个「整条路径被全部接受」的分数，然后在整棵树里按分数挑最高的节点去扩张，超出预算就全局剪枝。同样 60 个节点的预算，树能扎得更深，平均接受长度更长。

注意：EAGLE-2 没有改变单 token 的接受率，涨的是每轮能捞回来的 token 数。

EAGLE-3：换训练方式。 作者发现把训练数据翻倍，收益却很有限，根因是特征预测约束：损失里要求回归大模型的特征，这是额外约束，限制了草稿模型的表达能力；而且顶层特征本身只包含「下一个 token」的信息，用它预测「下下个 token」天然吃亏。于是 EAGLE-3 直接砍掉特征回归，改成纯 token 预测，同时把单一层特征换成低/中/高三层（代码里取第 2、L/2、L−3 层）拼接后压回原维度。

但只砍损失会出新问题：训练时输入永远是真特征，推理从第二步起输入是草稿自己猜的特征，两边分布对不上。training-time test 就是让训练模拟推理 —— 把自己的输出喂回去跑多步（代码里 7 步），每步都算损失。大模型只在最开始算一次特征，反复跑的只有那一层草稿模型，所以并不贵。

六、工程上唯一要记住的事

投机解码主要优化延迟，不是吞吐。batch 一大，GPU 本来就吃饱了，多出来的草稿计算只会挤占资源：

• vLLM 实测（GPT-OSS-20B）：并发 1 时 1.55×，并发 64 时只剩 1.05×；
• SGLang 默认在 batch size ≥ 32 时直接关掉多步投机。

其余几条经验：temperature 越高收益越低；代码补全、摘要这类「输入高度决定输出」的任务收益最大；MoE 模型（如 Mixtral 8x7B）收益明显更小。

框架支持：vLLM 用 method: "eagle" 或 "eagle3"；SGLang 用 --speculative-algorithm EAGLE3；TensorRT-LLM 用 decoding_type: Eagle3（EAGLE v1/v2 的权重不兼容）；HuggingFace Transformers 目前没有 EAGLE。

参考资料

• Leviathan 等，Fast Inference from Transformers via Speculative Decoding
• Chen 等，Accelerating Large Language Model Decoding with Speculative Sampling
• Li 等，EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty（ICML 2024）
• Li 等，EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees（EMNLP 2024）
• Li 等，EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test
• 官方代码：https://github.com/SafeAILab/EAGLE

1. 为什么需要并行和 Chunked Prefill

从存储和计算两个方面看：用资源/通信换并行处理

1.1 存储压力

大模型推理时，模型显存中会存储：

• 模型权重
• KV Cache
• 中间激活值

当我们使用比较大的模型进行长下文推理时，单卡显存会存不下，我们需要并行技术可以进行多卡推理，如：

TP：按 head 切分模型权重
PP：按 layer 切分模型权重
EP：按 MoE expert 切分专家权重
CP：按 token 切分，每个 rank 推理不同的 token，降低 kv cache
DP：KV Cache 按请求分摊，MLA 模型不会复制冗余 kv cache
Chunked Prefill：降低单次 Prefill 的临时显存峰值

1.2 计算压力

我们需要使用并行等技术分摊计算量：

TP：多卡分摊矩阵计算
PP：多卡分摊不同层计算
EP：多卡分摊不同 expert 计算
CP：多卡分摊长上下文 attention / KV 读取
DP：多张 GPU 并行处理不同请求，提高整体吞吐
Chunked Prefill：单机推理时，降低单次 Prefill 计算压力，避免长 Prefill 独占 GPU

1.3 简要总结

长上下文推理在推理引擎中面临的问题可以概括成三类：

针对上面的问题，目前已有的应对方案有以下几种：

2. SGLang 并行推理概览

2.1 SGLang 请求的推理流程

1. 请求排队：当前 Scheduler 中无请求，该 req 放入 waiting_queue 中

2. 控制 batch 内请求数量：创建 PrefillAdder，从 waiting_queue 中取出该请求，调用 PrefillAdder::add_one_req() prefill 会被截断分块，can_run_list 中只有这个 req，单个请求被打包成 ScheduleBatch

3. 分配空闲 kv cache：调用 ScheduleBatch::prepare_for_extend() ，实际上调用 alloc_for_extend()，分配 req_pool_indices 以及为需要 extend 的长度申请真实的 KV 显存 out_cache_loc，并把对应关系记录到的 req_to_token_pool

4. 模型推理 (TpModelWorker & ModelRunner)

   • 进入模型层 forward()。
   • Attention 执行 : 调用不同的 attention backend 根据 metadata 进行 attention 计算
   • MLP/MoE：进行 MLP/MoE 计算
   • lm_head：得到最终的 logits

5. Sample：根据 req 中的 sample params 进行采样

6. 最后得到 batch_result，流式返回给 Detokenizer

2.2 SGLang 中的 Chunked Prefill

SGLang 参数里有 max_prefill_tokens、prefill_max_requests、max_running_requests、chunked_prefill_size 等，它们共同决定一次 prefill 能放多少 token 和多少请求。

SGLang 的 chunked_prefill_size 控制一个请求一次 prefill forward 最多处理多少新 token。调度器会将完整 req tokens 截断到当前剩余 budget。

chunked prefill size 实际上是 MoE 层看到的 token 数，除以 dp_size 得到每个 DP rank 实际使用的 chunk size

最开始提出是为了解决TPOT过高（因为continous batching将prefill和decode拼在一起算，导致prefill很长而decode空等待出现尖刺）问题以及长文本下激活值占用过大导致 OOM 问题

• chunked prefill size增大，TTFT 减小，但需要预留更多显存给中间激活值，支持最大上下文减少
• chunked prefill size减小，TTFT 增大，但可以支持更大上下文

2.3 SGLang 中并行推理示例

1
2
3
4
5
6
7
8
9

python -m sglang.launch_server \
  --model-path <model> \
  --tp 16 \
  --dp 2 \
  --cp 4 \
  --ep 16 \
  --enable-dp-attention \
  --enable-prefill-context-parallel \
  --moe-a2a-backend deepep # 开 deepep 后 shared_experts 不切分（tp_size=1）

核心公式：tp_size = attn_dp_size × attn_cp_size × attn_tp_size

1
2
3
4
5
6
7
8
9
10
11
12
13

world_size  = tp_size * pp_size# 全局卡数
tp_size # 一个 pipeline stage 内参与模型并行的 rank 数
pp_size # 有多少个 pipeline stage

# attention 并行参数
attn_dp_size = attention_data_parallel_size
attn_cp_size = attention_context_model_parallel_size
attn_tp_size = tensor_model_parallel_size // attn_cp_size // attn_dp_size

# moe 并行参数
moe_ep_size = expert_model_parallel_size
moe_dp_size = moe_data_model_parallel_size
moe_tp_size = tensor_model_parallel_size // moe_ep_size // moe_dp_size

2.3.1 Rank 布局

单个 rank layout (dp, cp, tp)

2.3.2 参数设置（开启 DP, TP, CP, EP）

模型：NSA 模型(nsa indexer + MLA)

16 卡示例参数：

1
2
3
4
5
6
7
8
9

python -m sglang.launch_server \
  --model-path <model> \
  --tp 16 \
  --dp 2 \
  --cp 4 \
  --ep 16 \
  --enable-dp-attention \
  --enable-prefill-context-parallel \
  --moe-a2a-backend deepep # 开 deepep 后 shared_experts 不切分（tp_size=1）

3. 并行技术详细介绍

3.1 TP：Tensor Parallelism

解决单卡计算瓶颈和单卡显存瓶颈
当前模型都比较大，单张卡不能放得下，需要将模型权重按 hidden_dim 切分到不同卡上，通过通信来聚合中间结果 hidden_states，同时减少计算量和中间值的存储空间。

主要有两种并行方式：Column Parallel (列并行) 和 Row Parallel (行并行)。通常它们会成对出现（例如在 MLP 中：先 Column 后 Row），以最小化通信开销。

以 TP=2 为例：单卡计算量、中间激活和权重都减半，代价是 row parallel 层后多一次 All-Reduce；输入在 column parallel 层是复制的（不切分），到 row parallel 层才按列切分

shared_experts内部用的是 MergedColumnParallelLinear 和 RowParallelLinear，实际上会走 tp 切分。

如果设置 moe-a2a-backend shared_experts tp_size=1 不切分，复制到 rank

ColumnParallelLinear:
数学原理：
输入 $X$ 是完整的（复制在所有 GPU 上）。每个 GPU 计算 $Y_i=XA_i$​。输出 $Y$被 切分为 $[Y_1,Y_2,...,Y_p]$，即每个 GPU 得到输出向量的一部分特征。
典型应用:

• Attention 的 QKV Projection (QKVParallelLinear)。
• MLP 的 Gate / Up Projection (MergedColumnParallelLinear)。

RowParallelLinear:
数学原理：
为了匹配矩阵乘法规则，输入 $X$ 也必须按列切分为 $[X_1​,X_2​,...,X_p​]$（这正好是 ColumnParallelLinear 的输出格式）。每个 GPU 计算 $Y_i=X_iA_i$​。注意，$Y_i$ 的形状与最终输出 $Y$ 相同，但它只是部分和。最终输出 $Y = \sum Y_i$​ ，需要一次 All-Reduce (Sum) 操作
典型应用：

• Attention 的 Output Projection (o_proj)。
• MLP 的 Down Projection (down_proj)。

3.2 PP：Pipeline Parallelism

解决单卡显存瓶颈和计算瓶颈
pipeline schedule:

• stage 0 处理 req 的输入和 embedding 层
• 中间的 stage i th 从 i-1 th 接受上一轮的 output 并发送：
• 非最后一个 PP rank forward 结束后返回 hidden states；最后一个 PP rank 才做 logits processor / lm_head。

3.3 DP Attention：MLA 系列模型使用

解决 MLA 模型的显存瓶颈：kv cache 大量冗余
MLA 的 KV Cache 与 head 数量无关（因为 KV 只有一个头，无法切分所以只能复制多份），使用常规的 TP 会重复存多份 KV Cache，造成显存浪费。DP attention 就是为了解决这个问题。

• DP Attention 每个 rank 复制 attention 权重，每个 attention 只存储自己分到这部分 token 的 KV Cache(用权重存储换 KV Cache 存储)

SGLang 的 DP Attention：attention 层按 DP 切请求/token；MoE/MLP 仍然可以用 TP/EP。

其目的是将一个大的张量并行（TP）组重组为几个较小的 TP 组，这些较小的 TP 组又形成一个专门用于注意力层的新的数据并行（DP）组。如 TP=8，拆分成DP=2，TP=4。

1. 对于模型中除 MLP 层以外的部分（如 Embedding, Self-Attention），每个数据并行单元（DP_Rank）内部的张量并行规模（TP_Size）设置为 1，即每个 DP_Rank 独立计算这些部分。
2. 对于 MLP 层，所有 DP_Rank 则共同组成一个大的张量并行组（TP_Group），该组的大小等于原始的 tp_size（图中示例 attn_tp_size = 1，所以恰好等于 DP 规模）。

也就是说，原始 tp_size 被重新解释成：tp_size = attn_dp_size × attn_cp_size × attn_tp_size

3.3.1 DP Attention 的 gather/scatter

由于 attention 只处理本 DP rank 的 token，但 MLP/MoE 需要全局 token 布局，所以 SGLang 需要在 attention/MLP 边界做 gather/scatter。

没有 MoE all-to-all 后端时，MoE 不能直接把 token 发给专家，只能先在 DP 组内把各 rank 的 token 收集成全局布局。为了满足 all-gather / all-reduce 的形状要求，需要根据 prefill/decode 和 token 分布，选择把每个 rank 的 batch padding 到 max_len 还是 sum_len，对应使用 all-gather 或 all-reduce；MoE 算完后再 scatter 回各 DP rank。常见情况是 prefill 用 SUM，decode 用 MAX，但 decode 分布很不均时会退化成 SUM。

3.4 EP：Expert Parallelism

现在参数量比较大的 MoE 模型，专家数量可能是几十到几百个。如果每张 GPU 都用TP做切分，在总卡数相同时不会增加每卡显存；主要问题是每个激活专家都要一次 TP all-reduce，不同专家无法并行，导致通信开销大、计算效率低、扩展性差。

EP 一般流程

MoE 层的 Dispatcher

1
2
3
4
5
6
7
8
9

hidden_states
  -> gate/router logits
  -> top-k expert ids + top-k weights
  -> dispatch (all-to-all)：按 expert id 把 token 发到专家所在 GPU
  -> 本地按专家分组 -> grouped GEMM（专家 FFN）
  -> [若专家内部有 TP：TP all-reduce / reduce-scatter + all-gather]
  -> combine (反向 all-to-all)：把专家输出送回原 token 所在 GPU
  -> 按 top-k weights 加权求和
  -> output

3.4.1 Standard 模式（无 moe-a2a-backend）

Standard 模式 不做 All-to-All 通信，而是采用 AllGather 全量 token + 本地只算自己负责的专家 + AllReduce 归约的策略。

3.4.1.1 特点

1. 无 token 物理搬移：dispatch 阶段不移动 token 数据，只重映射 expert ID
2. 每个 rank 处理全部 token：但只在本地专家上计算，其余跳过（id=-1 的 token 贡献为 0）
3. 靠 AllReduce 聚合：各 rank 的部分结果相加得到最终结果

• 路由计算 (TopK Router)

  • 计算 Token 路由，输出全局专家 ID (topk_ids) 及其对应权重。

• 本地分发 (Dispatch)

  • 将全局专家 ID 映射为本地专家 ID（属于当前 Rank 的分配 Local ID，非本 Rank 的标记为 -1）。

• 专家执行 (MoE Runner)

  • 重排 (Pre-permute): 按专家 ID 对 Token 排序，直接过滤/跳过 ID 为 -1 的 Token。
  • 计算 (Grouped GEMM): 仅对分配到本地的专家执行核心矩阵运算 (W13 + W2)。
  • 还原 (Post-permute): 将计算后的 Token 恢复至原始输入顺序。

• 状态合并 (Combine)

  • 输出当前 Rank 计算完成的隐藏层状态 (Hidden States)。

• 全局聚合 (AllReduce)

  • 跨所有 EP/TP Rank 执行通信 (AllReduce)，将各节点分散计算的结果相加，得到最终完整输出。

3.4.2 DeepEP 模式（–moe-a2a-backend deepep）

DeepEP 模式使用真正的 All-to-All 通信：将 token 物理发送到目标专家所在的 rank，每个 rank 只接收和处理路由到自己的 token。
传统固定形状 All-to-All：
Rank0 → Rank1: 5 个 token（其中 3 个是 padding）
Rank0 → Rank2: 5 个 token（其中 5 个是 padding，Rank2 根本不需要）
Rank0 → Rank3: 5 个 token（其中 4 个是 padding）

DeepEP 动态 All-to-All：
Rank0 → Rank1: 2 个真实 token
Rank0 → Rank2: 0 个（根本不发）
Rank0 → Rank3: 1 个真实 token

分为两种模式 Normal 和 LowLatency 模式

• prefill 使用 normal 模式
• decode 使用 low-latency 模式

这是两者最根本的区别，源于对硬件拓扑的不同利用方式。

Normal 模式（两级接力）：它利用了节点内 (Intra-node) NVLink 的高带宽（~160 GB/s） 和节点间 (Inter-node) RDMA 的相对低带宽（~50 GB/s） 的不对称性。其策略是：

第一级 (RDMA)：源 GPU 只将一份数据通过 RDMA 发送到目标节点上编号相同的 GPU（例如，0号卡发给0号卡）。这大大减少了昂贵的跨节点数据传输量。
第二级 (NVLink)：接收节点上的 GPU 再通过高速 NVLink，将数据“扇出”给节点内所有需要此 token 的 GPU。
这种设计以牺牲一点路径长度为代价，换取了对稀缺 RDMA 带宽的高效利用，非常适合需要传输大量 token 的 Prefill 阶段。
Low-Latency 模式（一级直达）：它完全放弃了 NVLink 转发，所有通信都通过纯 RDMA 完成。其策略是：

直接发送：每个 token 的通信请求，都通过 RDMA 直接发送到目标专家所在的 GPU，不经过中间节点转发。
虽然这可能增加 RDMA 的跳数，但路径最短，没有中间环节的等待和开销，因此单次通信延迟极低，完美契合 Decode 阶段每次只处理少量 token 的场景。

3.4.3 通信计算 overlap 机制

目的：为了提高 GPU 利用率

3.4.3.1 TBO（Two-Batch Overlap 双批次重叠）

TBO 工作原理：

• 将一个 batch 拆成 2 个 sub-batch，在 attention 和 dispatch/combine 之间交错执行，隐藏通信延迟
• Sub-batch 0 的 MoE compute 与 Sub-batch 1 的 dispatch 通信重叠
• 需要 2 个独立 dispatcher（各自有独立的 handle 和 buffer 状态）

重叠点：
当微批次 0 在进行 MoE 的 dispatch（发送 token）通信时，微批次 1 正在执行其注意力（Attention）的计算。反之亦然。

效果：
计算流（Compute Stream）和通信流（Comm Stream）被交织在一起，原本串行的“计算-通信-计算”链条被打破，通信时间被有效利用起来做有用功。

TBO 并非在所有场景下都能带来收益，它的效果与批次大小和部署环境密切相关。

最佳场景：
大规模专家并行（EP）：特别是在多节点、跨节点带宽受限的环境下，通信瓶颈越突出，TBO 隐藏延迟的效果越好。
足够的批次大小（Batch Size）：TBO 只有在批次足够大时才能有效覆盖通信开销。在 SGLang 的测试中，当每个 DP rank 的批次大小超过 320 时，TBO 能带来约 5%~10% 的性能提升；但在小批次（如 bs=60）下，TBO 反而可能导致性能显著下降。

默认策略：
在多数实现中，TBO 默认仅用于 Prefill 阶段。因为 Prefill 阶段计算密集，适合用计算来隐藏通信。若强行用于 Decode 阶段，由于 Decode 本身计算量小，拆分批次带来的开销可能超过收益，导致吞吐量下降。

3.4.3.2 SBO（Single-Batch Overlap）

SBO 工作原理：

TBO 在超高并发下也可能遇到瓶颈。有实践发现，在特定硬件（如 H20）上，当并发极高时，TBO 的双批次策略可能导致延迟“爆炸”，反而变慢。为此，社区（如蚂蚁集团与 SGLang、DeepSeek 合作）提出了 SBO（Single-Batch Overlap） 作为改进方案。SBO 的优化在于，它不再等待整个批次完全准备好，而是在单个批次内部进行更细粒度的重叠，例如在计算 shared experts 的同时就启动 dispatch 通信，或者在计算完一个 block 后就立即发送，从而实现更极致的流水线。

在单 batch 的 MoE 层内部，用多 stream 将通信（dispatch/combine）与计算（gemm/shared expert）重叠。

• combine 操作与 down_gemm 重叠
• dispatch 与 shared expert 计算重叠

combine 以 block_m=64 为单位，down_gemm 每产出一块，combine 立刻消费一块，两者在 GPU 上并行执行（不同 SM 集合，不同 CUDA stream）

上面是非 blackwell 的实现，Blackwell 在 _pre_combine_hook 中运行 shared expert：

3.4.4 EPLB（Expert Parallelism Load Balancer）

在 MoE 的专家并行中，不同的“专家”被放置在不同的 GPU 上。但由于输入数据的差异，某些“热门专家”会被大量 token 选中，其所在的 GPU 就会过载；而“冷门专家”所在的 GPU 则相对空闲。整个系统的速度取决于最慢的那块 GPU，这会造成严重的资源浪费。

EPLB 的核心思想是：
动态地为高负载的“热门专家”创建冗余副本，并将这些副本智能地分配到不同 GPU 上，从而将负载从过载的 GPU 分摊出去，最终最小化所有 GPU 中的最大负载。

EPLB 的工作流程分为三个主要步骤：

1. 收集负载数据：在模型运行时，持续统计每个专家被路由到的 token 数量，作为其“热度”指标。这些数据通常通过历史统计的移动平均值来估计。
2. 执行均衡算法：基于收集到的负载数据，EPLB 算法会计算出一个新的专家放置方案。这个方案决定了哪些专家需要被复制，以及所有专家（包括副本）应该被放置在哪些 GPU 上。
3. 动态重排专家：根据算法计算出的方案，在运行时动态地重新排列专家及其副本在 GPU 上的分布，以实现负载均衡。

我们引用 DeepSeek 官方仓库中的一个经典例子来具体说明。

场景设定：
一个两层的 MoE 模型。
每层有 12 个逻辑专家（编号 0-11）。
部署在 2 个节点上，每个节点有 4 个 GPU，共 8 个 GPU。
为每层引入 4 个冗余专家，因此每层总共有 12 + 4 = 16 个“物理专家副本”需要放置。
假设的负载情况：经过数据收集，假设在第1层中，专家 4、5、1、10 是“热门专家”，负载显著高于其他专家。

EPLB 的分配过程（采用分层负载均衡）：
节点间均衡 (Step 1)：EPLB 先将 12 个逻辑专家分成 4 组，例如 (0-2), (3-5), (6-8), (9-11)。然后，将这 4 个组分配到 2 个节点上，确保每个节点分到的专家总“基础负载”大致相等。这是一个“背包问题”，可用贪心算法解决。
复制热门专家 (Step 2)：在节点内部，EPLB 识别出热门专家。假设节点A分到了专家组 (3-5)，其中专家 4 和 5 是热门的；节点B分到了专家组 (9-11)，其中专家 10 是热门的（专家 1 可能被分到了另一个组）。EPLB 会为这些热门专家创建副本。
节点内均衡 (Step 3)：最后，EPLB 将所有物理专家副本（包括原始专家和冗余副本）打包到该节点内的 4 个 GPU 上。它同样将此视为“背包问题”，目标是最小化每个 GPU 上的总负载。例如，它可能会将专家 4 的一个副本放在 GPU0，另一个副本放在 GPU2，从而分摊其负载，避免任何单个 GPU 过载。

通过逻辑专家和实际物理专家的映射，原本因专家 4 和 5 而可能过载的某个 GPU，其负载被多个副本分摊到了不同 GPU 上，实现了节点内和节点间的负载均衡。

动态统计专家热度 → 为热门专家复制多份完整权重副本 → 把所有专家副本重新排布到不同 GPU → 让路由把请求分散到副本上，从而最小化最大 GPU 负载。
缺点：需要事先申请冗余专家显存，降低了显存上限。

3.5 CP：Context Parallelism

将长上下文分发到多个 GPU 上计算，所以用了特殊的attention：

1. 基础版：Ring Attention（环状注意力）
   QKV序列分割：把长序列切成 N 块，分给 N 个 GPU（Rank）。
   设备环状结构：GPU 之间形成一个环。每个 GPU 计算当前块的 Q 与当前块的 KV 的注意力，然后把自己的 KV 块发给下一个 GPU，同时接收上一个 GPU 发来的 KV 块。
   完成条件：经过 N 次传递，每个 GPU 都看到了序列中所有的 KV 部分，计算完成。
   开销隐藏：由于是异步通信，GPU 可以在等待接收下一个 KV 块的同时，计算当前块的注意力，从而把通信时间藏在计算时间背后（Overlap），实现近零开销。
   核心痛点：Ring Attention 假设序列是均匀切分的（比如序列长 12，4 个 GPU，则 GPU0 拿 0-2，GPU1 拿 3-5，以此类推）。但在因果注意力（Causal Attention） 中，前面的 token 只能看自己，后面的 token 要看前面所有 token。这就导致排在后面的 GPU 计算量远大于排在前面的 GPU，产生严重的负载不均衡，环上的其他 GPU 必须等最慢的那个 GPU，流水线效率大打折扣。（上图中白色部分随着轮次越来越少）

2. 改进版一：Striped Attention（条纹注意力）
Striped Attention 是为了解决上述因果掩码导致的负载不均衡而提出的。
它的做法不是按连续块切分，而是按“条纹”交错切分：
假设序列长 16，4 个 GPU。
GPU0 拿 token：0, 4, 8, 12
GPU1 拿 token：1, 5, 9, 13
GPU2 拿 token：2, 6, 10, 14
GPU3 拿 token：3, 7, 11, 15
效果：在因果注意力下，每个 GPU 拿到的 token 分布在序列的各个位置，每个 GPU 需要计算的注意力范围（即它需要看到的前序 KV 数量）变得大致相等。这就极大地缓解了负载不均衡，让环上的计算更加同步，提升了整体吞吐量。（可以看到图上白色部分每一轮大致相等）
任意两张卡上的 Q 与 K 之间的 attention mask，都是一个完整的方形，而且每张卡上的计算量也是相对平均的
每个 DP 组分配 chunked_prefill_size // dp_size，每个 CP rank 在组内又拿到 1/cp_size token。

3. 改进版二：Zigzag Attention（锯齿注意力）
Zigzag Attention 是另一种解决因果掩码负载不均衡的切分策略，常见于后来的一些长上下文优化框架中。
原始块顺序：block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7
Zigzag 重排后：block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4

3.5.1 两种 CP 方式

主要为了降低 TTFT

• CP 只用于 prefill（extend）阶段，decode 阶段不使用 CP（每次只生成 1 个 token，无法按 seq 切分）。
• Decode 时每个 rank 需要读取完整的历史 KV cache，因此 prefill 结束后每个 rank 的 kv_pool 必须包含完整序列的 KV。

权重在 CP 维度上被复制了 cp_size 倍；但 attn 计算时每张 rank 只算 seq/cp_size 份 token（按 seq 切），CP 组内的 ranks 处理不同的 token（seq 按 CP 切分），它们的 o_proj 输出通过最后的 AllGather 合并，而不是 AllReduce
支持两种 CP 方式

round-robin-split：token_idx % cp_size，rank i拿到所有下标模 cp_size 等于 i 的 token
* batch 里所有 req 的 token 拼接成一维后的张量进行取模

in-seq-split：序列切成 2*cp_size 个连续 block，rank i 取 block i 和 block 2*cp_size-1-i 只支持单个 req 进行切分

约束差异

4. 总结

4.1 策略对指标的影响

TTFT / Prefill

TPOT / Decode

decode cp 不使用，设置参数也会被忽略

4.2 选用何种并行组合

4.2.1 Dense 模型

• 选型推荐：TP (节点内) + PP (跨节点)
• 依据：
  • Dense 模型的计算分布均匀，没有专家路由机制。首要解决的是单卡显存（装不下权重和 KV Cache）与单层算力不足的问题。
  • TP 是切分 Dense 模型最直接的方式，有效降低单次矩阵运算的维度。

4.2.2 MoE + NSA 模型

• 选型推荐：EP + TP + DP Attention
• 依据：
  • **EP：**通过 EP 将海量专家参数分散到不同 GPU，避免全量复制带来的显存问题，同时提高 expert 的计算并行度。
  • DP Attention： 在多请求并发时大幅降低单个 Rank 的 KV Cache 压力。
  • TP： 用于处理 MoE 模型中未被路由的稠密层（如 QKV 投影、Dense MLP 层）。

4.2.3 降低 TPOT

• 选型推荐：最大化单机 TP + (MoE场景下引入 DP Attention) + 大并发下 EP
• 依据：
  • TP 提升聚合显存带宽，显著压低单次 GEMV 耗时。
  • DP Attention 允许更大的 Batch Size，摊薄每个 Token 的访存成本。
  • 大并发下 EP 可以分摊权重读取成本

4.2.4 降低 TTFT

• **选型推荐：TP + EP + 按需引入 CP (针对长文本) **
• 核心依据：
  • **EP **将 Token 分发到不同 GPU 的特定专家上并行计算，实现通信和计算的 overlap
  • TP 能够将大矩阵切块，利用多卡算力直接并行计算
  • CP 可以多卡并行计算 req。

mHC Manifold-Constrained Hyper-Connections

流形约束超连接，是 DeepSeek 团队在 2025 年 12 月提出，在 2026 年 4 月发布的 DeepSeek V4 中应用的新的残差技术，是 DeepSeek 的核心创新之一。他的目标是为了解决传统残差连接表达能力不足以及使用原始超链接（Hyper-Connections）深层次训练崩溃的问题。

1. 从 residual 到 mHC

mHC 是 Hyper-Connections（HC）的改进版，要解决的问题只有一句话：HC 让残差变强了，但丢掉了残差连接最关键的"恒等通路"。

图 1：残差从「一个向量」变成「一个矩阵」。标准残差只有加法；HC 用 H_pre / H_post / H_res 三个矩阵管理读写与混合；mHC 只多做一件事——把 H_res 约束成双随机矩阵。

1.1 标准残差

$$\text{out} = x + F(x)$$

$x$ 这一项就是恒等映射：浅层信号不加修改地直达深层，反向梯度里也存在一条系数为 1 的通路，所以深层网络不会梯度消失。

代价是：透传比例写死为 1，$F(x)$ 的强度无法自适应，信息通路单一。

1.2 Hyper-Connections：残差变成矩阵

HC 把残差从 1 路向量 $x$ 拓宽成 n 路矩阵（n 称为 expansion rate），并用三个可学矩阵管理它：

$$\text{out} = Hres \cdot x + Hpost \cdot F(Hpre \cdot x)$$

图 2：HC 单层更新里的三个矩阵、三步运算。子层 $F$ 自己仍然只吃 1 路、跑在正常 hidden 宽度上，被加宽的是它外面那个残差矩阵；三个矩阵的系数 = 输入的动态部分 + 全局静态偏置。mHC 复用的是同一套流程，不同之处只有 $Hres$ 被约束成双随机矩阵，见图 3。

这样残差的宽度就和层的输入宽度解耦了：加宽残差几乎不增加 $F$ 的 FLOPs。消融显示三个矩阵里 $Hres$ 贡献最大，也就是性能提升主要来自残差流内部的信息交换。

1.3 HC 为什么会在深层崩掉

把 HC 递归展开，浅层信号到深层要乘上一长串 $Hres$：

$$\text{out} = \left(\prod Hres\right) x + \sum \left(\prod Hres\right) Hpost \cdot F(Hpre \cdot x)$$

标准残差里 $x$ 前面的系数恒为 1，HC 里它变成了 $\prod Hres$。而 $Hres$ 是自由学习的矩阵，没有任何约束保证这个连乘接近恒等：奇异值略大于 1 就指数放大，略小于 1 就指数衰减。论文实测该连乘的增益峰值到 3000 量级（理想值应为 1），同步出现的就是 loss 尖峰与梯度范数失控。

所以 HC 深层不稳定的根源不是某一层出了问题，而是"保范数的恒等通路"被换成了任意线性映射。

1.4 mHC：把 H_res 投影成双随机矩阵

mHC 不允许 $Hres$ 自由生长，而是用 Sinkhorn-Knopp 算法把它投影到 Birkhoff 多面体——双随机矩阵构成的集合，即元素非负、每行每列之和都为 1。这就是"流形约束（Manifold-Constrained）"的含义。

图 3：约束只作用在 comb 这个 $n \times n$ 矩阵上。先指数化，再交替做行归一化与列归一化（V4 发布配置迭代 20 次），收敛到双随机矩阵。

双随机正好补上 HC 缺的两条性质：

1. 行和为 1 ⇒ 守恒：$Hres \cdot x$ 是各路特征的凸组合，均值被保存、信号不被放大。这就是论文说的"恢复恒等映射性质"——恒等矩阵本身也是双随机矩阵。
2. 对乘法封闭 ⇒ 任意深度都守恒：双随机矩阵相乘仍是双随机，所以 $\prod Hres$ 不会失控。这是"流形约束"和"调初始化 / 加正则"的本质区别。

另外两个矩阵用更松的约束，都是为了消除正负抵消：

对照下文 hc_pre 的实现：mixes 被拆成三段，前两段过 sigmoid 得到 pre / post，第三段 reshape 成 $n \times n$ 后交给 sinkhorn_normalize 做"行归一化 → 列归一化"的交替迭代，收敛后得到的正是双随机的 comb。

1.5 小结

mHC 保留了 HC"多路残差 + 内部混合"的表达力，只把混合矩阵限制在双随机流形上，把 HC 丢掉的守恒性拿了回来。V4 取 $n = 4$（Pro 每路 7168 维，Flash 每路 4096 维）；加宽的残差要多读写显存，所以还要靠算子融合、选择性重计算、DualPipe 通信重叠压开销，端到端训练时间只增加约 6.7%。

2. 计算流程

以下代码来自 sglang 中的 deepseek v4 的 forward函数，省略了部分细节。

1
2
3
4
5
6
7
8
9
10
11
12
13
14

residual = hidden_states
hidden_states, post, comb = self.hc_pre(
    hidden_states, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
)

# 模块计算，以 attention 为例
hidden_states = self.input_layernorm(hidden_states)
hidden_states = self.self_attn(
    x=hidden_states,
    positions=positions,
    forward_batch=forward_batch,
)

hidden_states = self.hc_post(hidden_states, residual, post, comb)

可以看见，在最开始的时候会保存原始的 hidden_states（residual），之后对这部分 hidden_states 做预处理。这里的 residual 就是第 1 节里的 $x$（n 路残差矩阵），hc_pre 负责算出 $Hpre$、$Hpost$、$Hres$ 三个矩阵，hc_post 完成的正是 $\text{out} = Hres \cdot x + Hpost \cdot F(Hpre \cdot x)$ 这一步。

hc_pre

对于 hc_pre 函数，可以看到有四个入参。
输入参数：

返回值:

具体计算分为两个步骤：

1. 计算输入的归一化（rmsnorm），并与混合函数矩阵乘，映射得到混合权重。

1
2
3
4
5
6
7

def hc_pre_torch_impl(x, hc_fn):
    x_flat = x.flatten(1).float()
    rsqrt = torch.rsqrt(
        x_flat.square().mean(-1, keepdim=True) + self.rms_norm_eps
    )
    mixes = (F.linear(x_flat, hc_fn) * rsqrt).unsqueeze(1)
    return x_flat, mixes

$$\text{mixes} = (Whc \cdot x) \cdot \text{RMSNorm}(x)$$

这个函数将返回 mixes 交给下一步处理。

2. Sinkhorn 归一化
   Sinkhorn 归一化是一种快速迭代算法，用来将任意非负矩阵转换成双随机矩阵（每行和每列的和都等于 1）。它是 DeepSeek V4 mHC（流形约束超连接）架构中最核心的数学创新，解决了原始超连接训练崩溃的问题。hc_split_sinkhorn 的作用是将输入的 mixes 张量拆分为三部分（pre、post、comb），并对 comb 部分执行 Sinkhorn 归一化，使其收敛为双随机矩阵（doubly stochastic matrix）。
   核心思想：交替对行和列进行归一化，直到收敛。

1
2
3
4
5
6
7
8
9

pre, post, comb = hc_split_sinkhorn(
    mixes,
    hc_scale,
    hc_base,
    self.hc_mult,
    self.hc_sinkhorn_iters,
    self.hc_eps,
)
y = (pre.squeeze(1).unsqueeze(-1) * x_flat.view(shape)).sum(dim=1)

输入参数：

输出参数：

其中输出的 pre 和 post 是由 mixes、hc_scale 以及 hc_base 计算得到的（把 mixes 按 mult、mult、mult × mult 切成三段，分别记作 mixes₁ / mixes₂ / mixes₃）：

$$\operatorname{sigmoid}(x) = \frac{1}{1 + e^{-x}}$$

$$\text{pre} = \operatorname{sigmoid}\big(\text{mixes1} \cdot scale1 + base1\big)$$

$$\text{post} = 2 \cdot \operatorname{sigmoid}\big(\text{mixes2} \cdot scale2 + base2\big)$$

$$\text{comb} = \text{mixes3} \cdot scale3 + base3$$

三段的形状分别是 $mult$、$mult$、$mult^2$，合计 $(2 + mult)\times mult$，与上表里 mix_hc 的定义一致。comb 这一段 reshape 成 $mult \times mult$ 之后才进 Sinkhorn 迭代，最后得到的 $n \times n$ 矩阵每行、每列之和都是 1——这正是第 1.4 节说的"投影到双随机矩阵流形"。

在分出这三个 tensor 后，对最后一个 combination tensor 进行迭代的归一化：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

def sinkhorn_normalize(comb_frag: torch.Tensor, sinkhorn_iters: int, eps: float = 1e-6):
    """
    comb_frag: (hc, hc)
    """
    # --- 第一次行归一化（带 softmax 数值稳定）---
    row_max = torch.max(comb_frag, dim=1).values          # (hc,)
    comb_frag = torch.exp(comb_frag - row_max.unsqueeze(1))
    row_sum = torch.sum(comb_frag, dim=1)                 # (hc,)
    comb_frag = comb_frag / row_sum.unsqueeze(1) + eps    # 注意：eps 加在结果上

    # --- 第一次列归一化 ---
    col_sum = torch.sum(comb_frag, dim=0)                 # (hc,)
    comb_frag = comb_frag / (col_sum.unsqueeze(0) + eps)  # 注意：eps 加在分母上

    # --- 交替迭代 sinkhorn_iters - 1 次 ---
    for _ in range(sinkhorn_iters - 1):
        row_sum = torch.sum(comb_frag, dim=1)
        comb_frag = comb_frag / (row_sum.unsqueeze(1) + eps)

        col_sum = torch.sum(comb_frag, dim=0)
        comb_frag = comb_frag / (col_sum.unsqueeze(0) + eps)

    return comb_frag

hc_post

将注意力/FFN 的输出与残差进行加权组合，是 hc_pre 的逆操作。该部分是对 mhc 部分的后处理函数：

1
2
3
4
5

def hc_post_torch_impl(x, residual, post, comb):
    return (
        post.unsqueeze(-1) * x.unsqueeze(1)
        + (comb.unsqueeze(-1) * residual.unsqueeze(2)).sum(dim=1)
    ).type_as(x)

输入参数：

输出参数：

$$\text{out} = \text{post} \cdot x + \text{comb} \cdot \text{residual}$$

拆开看就是两件事：post 把"子层算出来的那 1 路结果"按权重摊回每一路残差，comb 则决定"上一层的 n 路残差"如何重新混合成这一层的 n 路。前者对应第 1.2 节的 $Hpost$，后者对应 $Hres$（也就是经过 Sinkhorn 归一化的那个双随机矩阵）。

3. 总结

mhc 诞生主要为了解决：

1. 传统残差连接表达能力不足
   标准残差连接只是简单的 out = x + f(x)，信息通路单一。
   mHC 通过 hc_pre 将输入映射到多路混合空间（hc_mult 路），在 hc_post 阶段用 post 和 comb 对当前层输出和残差做加权组合，相当于同时维护了多条残差通路，增强了表征能力。
2. 原始 Hyper-Connections 深层训练崩溃
   纯超连接在深层网络中容易出现数值不稳定或训练崩溃。
   mHC 的核心创新是引入 Sinkhorn 归一化：将组合矩阵 comb 约束为双随机矩阵（每行、每列和均为 1）。
   这保证了多路残差之间的信息流动是均衡且稳定的，避免了权重爆炸或消失，从而解决了深层训练崩溃问题。

大模型Attention发展历程

MHA Multi Head Attention

MHA 相比于传统的 Self-Attention 提升了模型的表达能力。例如一个专家组对比一个专家。
先看代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_values: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor, torch.Tensor]:
    input_shape = hidden_states.shape[:-1]
    hidden_shape = (*input_shape, -1, self.head_dim)

    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

    cos, sin = position_embeddings
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

    attn_output, attn_weights = attention_interface(
        self,
        query_states,
        key_states,
        value_states,
        attention_mask,
        dropout=0.0 if not self.training else self.attention_dropout,
        scaling=self.scaling,
        **kwargs,
    )

    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
    attn_output = self.o_proj(attn_output)
    return attn_output, attn_weights

这个 attention 是 Multi-head self-attention，注意力有头这个概念。其中，Q K V 这三个向量分别是由 hidden_states 乘上一个 Linear 层来计算的。这里面由于是多头注意力机制，所以需要让 Q K V 的 shape 变为 (batch, seq_len, num_head, head_dim)。之后 transpose 成 (batch, num_head, seq_len, head_dim)，这么做的目的是为了让后续计算时，每个头可以独立进行注意力计算，让 num_head 变为“批次”的一部分，注意力仅在 seq_len 维度计算，而不把不同的 head 混在一起算。
再算完这几个向量后，需要对 Q K 向量添加位置信息，这就是 apply_rotary_pos_emb 的作用。之后就是 attention 的计算，也就是我们说的公式:

$$\begin{aligned} \text{Attention}(Q, K, V) &= \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V \end{aligned}$$

其中，QK 的意思是 Q 和 K 的点积，表示 Q 和 K 之间的相似性，更进一步，这个向量是词向量，是词在高维空间的数值映射。词向量之间相关度高表示什么？是不是在一定程度上（不是完全）表示，在关注词 A 的时候，应当给予词 B 更多的关注？除以 $\sqrt{d_k}$ 的作用是训练时防止梯度消失。
算完 attention 后，再被 o_projection 映射，结果就算出来了。

• MHA (Multi-Head Attention)：假设头个数为32，
  • 配置：每个 Query 头都有自己独立的 Key 头和 Value 头。
  • 数量：32 个 Q 头，对应 32 个 K 头 和 32 个 V 头。
  • 特点：完全独立，信息隔离性好，模型表达能力强。
    总结：MHA 多头自注意力机制是效果最好的也是最慢的，而 MQA 和 GQA 则是通过不同的方式来提升计算效率。
    一些比较出名的模型，如 LLama2-7B、LLama2-13B 使用的都是这种架构。

MQA Multi Query Attention

相比 MHA，MQA 允许一个 Query 头同时关注多个 Key 头和 Value 头，从而在保证信息隔离性的前提下，提升模型的计算效率。

• MQA (Multi-Query Attention)：假设头个数为32，
  • 配置：所有的 Query 头共享同一组 Key 和 Value 头。
  • 数量：32 个 Q 头，对应 1 个 K 头 和 1 个 V 头。
  • 特点：极端压缩，推理速度极快，但模型效果下降明显（因为所有头被迫看同样的 K/V 信息）。
    目前没听说哪个模型全程使用这个架构。但是 DeepSeek V3.2 的 decode 阶段使用了该方式。

GQA Group Query Attention

GQA 相比 MHA，其主要区别在于，GQA 是先进行“组内”注意力计算，再进行“组间”注意力计算。也就是说，GQA 首先将 Q K V 分成多组，然后在组内进行注意力计算，最后在组间进行注意力计算。推理速度更快，普遍用于 70B 及更大的模型中。MHA 则用于稍微小一点的模型中。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
    bsz, q_len, _ = hidden_states.size()

    query_states = self.q_proj(hidden_states)
    key_states = self.k_proj(hidden_states)
    value_states = self.v_proj(hidden_states)

    query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
    key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
    value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)

    cos, sin = position_embeddings
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

    # repeat k/v heads if n_kv_heads < n_heads
    key_states = repeat_kv(key_states, self.num_key_value_groups)
    value_states = repeat_kv(value_states, self.num_key_value_groups)

    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

    if attention_mask is not None:  # no matter the length, we just slice it
        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
        attn_weights = attn_weights + causal_mask

    # upcast attention to fp32
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
    attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
    attn_output = torch.matmul(attn_weights, value_states)

    if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
        raise ValueError(
            f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
            f" {attn_output.size()}"
        )

    attn_output = attn_output.transpose(1, 2).contiguous()
    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

    attn_output = self.o_proj(attn_output)

    if not output_attentions:
        attn_weights = None

    return attn_output, attn_weights

可以看见代码里有一个 repeat_kv 的操作，作用是将较少的 Key/Value 头复制扩展，使它们能和 Query 的头数一一对应进行矩阵乘法。这正是 GQA 的典型操作：“分组共享，通过复制来对齐”。

• GQA (Grouped Query Attention)：
  • 配置：将 Query 头分组，每组 Query 头共享一组 Key 和 Value。
  • 数量：例如分为 8 组（Group=8）。32 个 Q 头，对应 8 个 K 头 和 8 个 V 头（每4个Q头共享1个K/V头）。
  • 特点：介于 MHA 和 MQA 之间。比 MQA 效果好（保留了部分独立性），比 MHA 速度快（KV Cache 大小减少了 4 倍）。
    GQA 相比于 MHA，效果上几乎相近，而速度却能达到 MQA 的水平，被广泛的应用在现在的大模型上。
    LLama2-70B、LLama3、Qwen2 全系列使用的都是这种架构。

MLA Multi Latent Attention

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_values: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
    batch_size, seq_length = hidden_states.shape[:-1]
    query_shape = (batch_size, seq_length, -1, self.qk_head_dim)
    key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim)

    if self.q_lora_rank is None:
        q_states = self.q_proj(hidden_states)
    else:
        q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
    q_states = q_states.view(query_shape).transpose(1, 2)
    q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

    compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
    k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

    k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2)
    k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)

    k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)

    cos, sin = position_embeddings
    if self.config.rope_interleave:  # support using interleaved weights for efficiency
        q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
    else:
        q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
    k_rot = k_rot.expand(*k_pass.shape[:-1], -1)

    query_states = torch.cat((q_pass, q_rot), dim=-1)
    key_states = torch.cat((k_pass, k_rot), dim=-1)

    if past_key_values is not None:
        # sin and cos are specific to RoPE models; cache_position needed for the static cache
        cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
        key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)

    if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim:
        value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])

    attention_interface: Callable = eager_attention_forward
    if self.config._attn_implementation != "eager":
        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

    attn_output, attn_weights = attention_interface(
        self,
        query_states,
        key_states,
        value_states,
        attention_mask,
        dropout=0.0 if not self.training else self.attention_dropout,
        scaling=self.scaling,
        **kwargs,
    )

    if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim:
        attn_output = attn_output[:, :, :, : self.v_head_dim]

    attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
    attn_output = self.o_proj(attn_output)
    return attn_output, attn_weights

上面代码是 MLA 的流程。可以分为以下几步：

1. Query 的生成与压缩
   DeepSeek 对 Q 进行了压缩，减少计算量

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

# 这一部分对应 Query 的生成
if self.q_lora_rank is None:
    # 如果没有设置压缩秩，则使用标准投影
    q_states = self.q_proj(hidden_states)
else:
    # 【核心 MLA 特征】：Query 低秩压缩
    # 1. q_a_proj: 降维 (hidden -> low_rank)
    # 2. q_a_layernorm: 训练时加了个layernorm，防止梯度爆炸，归一化
    # 3. q_b_proj: 升维 (low_rank -> num_heads * qk_head_dim)
    q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))

q_states = q_states.view(query_shape).transpose(1, 2)

# 【核心 MLA 特征】：解耦 RoPE
# 将 Query 分割为两部分：一部分不带位置信息，一部分带位置信息
q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

2. KV 的压缩与潜在向量生成

1
2
3
4
5
6
7
8
9

# 【核心 MLA 特征】：KV 压缩
# 将输入投影到一个非常小的压缩向量 (Latent Vector)
# 这就是推理时缓存的 "KV Cache"，体积非常小！
compressed_kv = self.kv_a_proj_with_mqa(hidden_states)

# 将压缩向量拆分：
# k_pass: 压缩后的 KV 内容特征 (用于后续还原)
# k_rot: 专门用于位置编码的部分
k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

3. KV 的解压与还原

1
2
3
4
5
6

# 1. 对压缩的内容特征进行 LayerNorm
# 2. kv_b_proj: 升维投影，将低维 Latent 还原成高维的 Key 和 Value
k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2)

# 分离还原后的 Key 和 Value
k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)

虽然我们在 Cache 里存的是压缩后的向量，但在计算当前时刻的 Attention 分数时，需要将其还原成完整的多头 Key 和 Value。虽然计算量增加了一些，但换来了显存带宽的大幅节省。做了个标准的工程化优化，k 和
v 的上投影矩阵被拼在一起了，所以出来需要 split 一下。

4. 位置编码与拼接

1
2
3
4
5
6
7
8
9
10

# 对专门的位置部分应用 RoPE
cos, sin = position_embeddings
q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin)

# 扩展 k_rot 维度以匹配多头

# 【核心 MLA 特征】：重组
# 将内容部分 和位置部分 拼接
query_states = torch.cat((q_pass, q_rot), dim=-1)
key_states = torch.cat((k_pass, k_rot), dim=-1)

最终参与计算的 query_states 和 key_states 是由“非旋转部分”和“旋转部分”拼接而成的。这样做的好处是，位置信息不会破坏压缩向量的低秩结构。为什么位置编码要这么做呢？因为在实现 MLA 的时候，有一个小技巧叫做“矩阵吸收”。

$$x W_{down}^Q W_{up,i}^Q (W_{up,i}^K{})^T (c^{KV}{})^T$$

在这个公式中，可以将 kv 的上投影矩阵$W_{up,i}^K$看成是$W_{up,i}^Q$，这样子，位置信息部分就可以通过矩阵吸收的方式，被$W_{down}^Q$“吸收”进去了。但是如果加了常规旋转位置编码的话，位置信息部分就无法被$W_{down}^Q$“吸收”进去了。可以看见被隔断了。无法吸收，会导致推理效率下降。

$$x W_{down}^Q W_{up,i}^Q R_{m} R_{n}^T (W_{up,i}^K{})^T (c^{KV}{})^T$$

所以引入了 Decoupled RoPE (解耦 RoPE)。q_rot 和 k_rot 分开处理，这样位置编码和内容编码就可以分开处理，互不干扰。

5. 注意力计算

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

# 更新 KV Cache (这里存储的通常是还原后的 K 和 V，或者是压缩向量，取决于具体的 Cache 实现)
if past_key_values is not None:
    key_states, value_states = past_key_values.update(...)

# 【Flash Attention 兼容性处理】
# 如果 QK 的维度和 V 的维度不一样，需要 Padding 对齐
if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim:
    value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])

# 调用标准注意力接口计算
attn_output, attn_weights = attention_interface(...)

# 去除 Padding
if ...:
    attn_output = attn_output[:, :, :, : self.v_head_dim]

# 输出投影
attn_output = self.o_proj(attn_output)

• MLA (Multi-Head Latent Attention)
  • 配置：通过生成 compressed_kv，将 KV Cache 的大小压缩了 90% 以上。存储上 compressed_kv 只有一份，接近 MQA。
  • 数量：通过 kv_b_proj 动态还原 KV，保留了多头注意力的表达能力。每个头一个 KV 数量，这一部分接近 MHA。
  • 特点：通过 split 和 cat 操作，把位置信息和内容信息分开处理，使得压缩机制能与 RoPE 完美兼容。

NSA Natural Sparse Attention

该注意力机制首次被提出是在 2025 年的 2 月。由 DeepSeek 针对长序列处理提出的原生稀疏注意力机制，属于 block-wise 的稀疏方案。这个目前只存在于论文中，没有实际用于某个模型。直到 Deepseek-V3.2 中，他们提出了 DSA，可以说是 NSA 的一个变种实现， 属于 token-wise 的稀疏方案。

1. 压缩块，计算块间注意力
   标准自注意力需要计算序列中每个词元（Token）与其他所有词元的关系。对于一个长度为 N 的序列，需要计算 N*N 个注意力权重，随长度二次增长。
   解决方案是，NSA 不再以单个 token 为基本计算单元，而是将连续的多个 token 聚合成一个 block。例如，将 4096 个词元的序列划分为 64 个块，每个块包含 64 个 token。
   计算流程转变：
   之前（标准注意力）： 处理 4096 个 token -> 需要处理 4096 * 4096≈16.7M 个关系对。
   之后（块压缩）： 处理 64 个 block -> 先计算块与块之间的注意力。此时，需要处理的关系对数量骤降至 64 * 64=4096 个。
2. 重要性筛选：选出重要的 K 个块
   在压缩后的 block 中，筛选出需要详细看的部分：
   块压缩是假设所有 block 都同等重要。但实际上，对于当前要处理的 token 来说，某些 block 是关键的而其他的是次要的。NSA 引入动态机制，根据当前的 Q 内容，评估并筛选出最相关的少量关键 block。
   系统会为每个 block 计算一个“重要性分数”。这个分数通常基于当前 Q 与每个 block 的“摘要向量”（通常是 block 内 token 的均值或通过一个小型网络生成）的相似度。
   在块压缩的基础上增加重要性筛选，系统从64个块中选出最相关的 K 个块（例如 K=8），额外计算当前查询向量与这 K 个关键块内全部 token 的注意力。
3. 滑动窗口
   为了保证局部上下文信息的完整性，token 不但要关注遥远的关键块，还要关注紧挨着它的其他 token。NSA 强制性地规定，无论重要性筛选的结果如何，每个 token 都必须关注以其自身为中心的一个局部窗口内的所有其他token。这个窗口是“滑动”的，因为它随着 token 位置的变化而移动。例如，假设滑动窗口大小为 256。那么对于序列中第 500 个 token，无论如何它都会关注第 500-255 至第 500 这个范围内的所有 token。

DSA DeepSeek Sparse Attention

DSA 首次被应用在 Deepseek V3.2 Exp 中，属于 token-wise 的稀疏方案。DeepSeek-V3.2-Exp 仅在 DeepSeek-V3 的基础上新增了 Lightning Indexer 模块，用于选择参与 attention 的 token。
该模块的主要输入是 compressed_q 与 MLA 的输入矩阵 hidden_states，输出则是每个 token 所对应的 2048 个可参与 attention 计算的历史 token 的 index。
被选中的 2048 个 token 会在 attention 的 mask 阶段发挥作用：通过将未选中的位置的 mask 值设为 inf，在经过 softmax 之后，就能有效去除不需要参与 attention 的 token。这样实现了高效的稀疏化，显著降低了计算量。从O(L^2)降到了O(LK)，其中 K<<L 为 2048。
Indexer 代码（来自于官方 DeepSeek V3.2 仓库）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

def fp8_index(
    q: torch.Tensor,
    q_s: torch.Tensor,
    k: torch.Tensor,
    k_s: torch.Tensor,
) -> torch.Tensor:
    """
    Perform index score using FP8 precision.

    Args:
        q (torch.Tensor): The Q tensor, must be contiguous.
        q_s (torch.Tensor): The scaling factor for Q (float), must be contiguous.
        k (torch.Tensor): The K tensor, must be contiguous.
        k_s (torch.Tensor): The scaling factor for K (e8m0 here), must be contiguous.

        fp8 q @ fp8 k -> fp32 logits
        relu(fp32 logits) * q_s (weights) -> fp32 logits
        fp32 logits -> fp32 logits_sum
        fp32 logits_sum * k_s (e8m0) -> fp32 index_score
    """
    return fp8_index_kernel(q.shape[2], q.shape[3])(q, q_s, k, k_s)

class Indexer(torch.nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.dim: int = args.dim
        self.n_heads: int = args.index_n_heads
        self.n_local_heads = args.index_n_heads // world_size
        self.head_dim: int = args.index_head_dim
        self.rope_head_dim: int = args.qk_rope_head_dim
        self.index_topk: int = args.index_topk
        self.q_lora_rank: int = args.q_lora_rank
        self.wq_b = Linear(self.q_lora_rank, self.n_heads * self.head_dim)
        self.wk = Linear(self.dim, self.head_dim)
        self.k_norm = LayerNorm(self.head_dim)
        # weights_proj in the checkpoint is stored in bf16, while the parameters here are stored in fp32 for convenient.
        self.weights_proj = Linear(self.dim, self.n_heads, dtype=torch.float32)
        self.softmax_scale = self.head_dim ** -0.5
        self.scale_fmt = args.scale_fmt

        self.register_buffer("k_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.head_dim, dtype=torch.float8_e4m3fn), persistent=False)
        self.register_buffer("k_scale_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.head_dim // block_size, dtype=torch.float32), persistent=False)

    def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
        bsz, seqlen, _ = x.size()
        end_pos = start_pos + seqlen
        q = self.wq_b(qr)
        q = q.view(bsz, seqlen, self.n_heads, self.head_dim)
        q_pe, q_nope = torch.split(q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1)
        # rope in indexer is not interleaved
        q_pe = apply_rotary_emb(q_pe, freqs_cis, False)
        q = torch.cat([q_pe, q_nope], dim=-1)
        k = self.wk(x)
        k = self.k_norm(k)
        k_pe, k_nope = torch.split(k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1)
        # rope in indexer is not interleaved
        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis, False).squeeze(2)
        k = torch.cat([k_pe, k_nope], dim=-1)
        q = rotate_activation(q)
        k = rotate_activation(k)
        q_fp8, q_scale = act_quant(q, block_size, self.scale_fmt)
        k_fp8, k_scale = act_quant(k, block_size, self.scale_fmt)
        self.k_cache[:bsz, start_pos:end_pos] = k_fp8
        self.k_scale_cache[:bsz, start_pos:end_pos] = k_scale
        weights = self.weights_proj(x.float()) * self.n_heads ** -0.5
        weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
        index_score = fp8_index(q_fp8.contiguous(), weights, self.k_cache[:bsz, :end_pos].contiguous(), self.k_scale_cache[:bsz, :end_pos].contiguous())
        if mask is not None:
            index_score += mask
        topk_indices = index_score.topk(min(self.index_topk, end_pos), dim=-1)[1]
        topk_indices_ = topk_indices.clone()
        dist.broadcast(topk_indices_, src=0)
        assert torch.all(topk_indices == topk_indices_), f"{topk_indices=} {topk_indices_=}"
        return topk_indices

公式如下：

$$I_{t,s} = \sum_{j=1}^{H_{I}} w_{t,j}^{I} \cdot \operatorname{ReLU}\left( \boldsymbol{q}_{t,j}^{I} \cdot \boldsymbol{k}_{s}^{I} \right)$$

总体框架如下：

其中，绿色部分是 indexer 实现。而框住的部分是函数 fp8_index 部分。可以看到，虽然说 attention 的计算复杂度降低了，但是实际上选择不同 token 重要性这个 indexer 模块，也算是个 attention。还是需要计算查询相对于之前每个 token 的 attention 的重要性分数，复杂度还是 $O(n^2)$。只不过这个过程是在 fp8 精度下进行的，所以效率会比较高。

为什么需要哈达玛变换"rotate_activation"？
直接量化 q 和 k 向量可能会造成精度损失。引入旋转变换操作：将大范围数值打散到小范围。可以与随机的正交矩阵相乘来做这件事，使得向量数值分布均匀且模长不变。不过矩阵乘法代价较大。于是使用小代价的哈达玛变换。

最后得出的topk_indices，会作为 mask 输入到 MLA 中，从而实现稀疏化。可以看下面代码。
对于 Deepseek V3.2 的 attention 实现而言：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
    """
    Forward pass for the Multi-Head Latent Attention (MLA) Layer.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim).
        start_pos (int): Starting position in the sequence for caching.
        freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
        mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention.

    Returns:
        torch.Tensor: Output tensor with the same shape as the input.
    """
    bsz, seqlen, _ = x.size()
    end_pos = start_pos + seqlen
    qr = self.q_norm(self.wq_a(x))
    q = self.wq_b(qr)
    q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)
    q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
    q_pe = apply_rotary_emb(q_pe, freqs_cis)
    kv = self.wkv_a(x)
    kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
    kv = self.kv_norm(kv)
    k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)
    # we use fp8 kv cache in actual deployment, so here we simulate the precision by casting kv to fp8 and then back to bf16.
    kv_fp8, kv_scale = act_quant(kv, block_size, self.scale_fmt)
    kv = (kv_fp8.view(-1, block_size).float() * kv_scale.view(-1, 1)).to(kv.dtype).view_as(kv)
    self.kv_cache[:bsz, start_pos:end_pos] = kv
    self.pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2)
    if mask is not None:    # MHA prefill
        q = torch.cat([q_nope, q_pe], dim=-1)
        kv = self.wkv_b(kv)
        kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim)
        k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
        k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1)
        scores = torch.einsum("bshd,bthd->bsht", q, k).mul_(self.softmax_scale)

        # indexer
        topk_indices = self.indexer(x, qr, start_pos, freqs_cis, mask)
        index_mask = torch.full((bsz, seqlen, seqlen), float("-inf"), device=x.device).scatter_(-1, topk_indices, 0)
        index_mask += mask
        scores += index_mask.unsqueeze(2)

        scores = scores.softmax(dim=-1)
        x = torch.einsum("bsht,bthd->bshd", scores, v)
    else:                   # MQA decode
        if self.dequant_wkv_b is None and self.wkv_b.scale is not None:
            self.dequant_wkv_b = weight_dequant(self.wkv_b.weight, self.wkv_b.scale)
        wkv_b = self.wkv_b.weight if self.dequant_wkv_b is None else self.dequant_wkv_b
        wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)
        q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])
        scores = (torch.einsum("bshc,btc->bsht", q_nope, self.kv_cache[:bsz, :end_pos]) +
                    torch.einsum("bshr,btr->bsht", q_pe, self.pe_cache[:bsz, :end_pos])) * self.softmax_scale

        # indexer
        topk_indices = self.indexer(x, qr, start_pos, freqs_cis, mask)
        index_mask = torch.full((bsz, 1, end_pos), float("-inf"), device=x.device).scatter_(-1, topk_indices, 0)
        scores += index_mask.unsqueeze(2)

        scores = scores.softmax(dim=-1)
        x = torch.einsum("bsht,btc->bshc", scores, self.kv_cache[:bsz, :end_pos])
        x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:])
    x = self.wo(x.flatten(2))
    return x