caijizhuo

发表于 更新于 分类于
本文字数： 3.1k 阅读时长 ≈ 11 分钟

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

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

一、发展历史

时间 工作 一句话贡献
2023.01 Leviathan 等，Fast Inference from Transformers via Speculative Decoding 提出「草稿-验证」框架，并证明它无损
2023.02 Chen 等，Accelerating LLM Decoding with Speculative Sampling 独立提出同一思路，给出拒绝采样算法
2024.01 EAGLE 草稿从「token 级」改成「特征级 + 超前一步 token」，模型小 2~3 个数量级
2024.01 Medusa 不训独立草稿模型，直接在骨干上挂多个解码头
2024.02 Lookahead Jacobi 迭代 + n-gram 池，完全免训练
2024.06 EAGLE-2 草稿树从静态改成动态，按置信度分配节点预算
2025.03 EAGLE-3 training-time test + 多层特征融合，去掉特征预测约束

主线是 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）：

草稿方式 加速比
token 级自回归 1.5×
特征级自回归 1.9×
特征级 + 超前一步 token 2.8×

草稿模型有多小

EAGLE 草稿模型结构

  • Embedding 和 LM Head 直接复用目标模型的，冻结不动，只训练中间一个 FC + 一层 decoder。
  • 特征取自 LM Head 的输入（最后一层 Transformer 的输出，代码里是 hidden_states[-1]）。论文叫它 second-to-top-layer feature，是把 LM Head 当作「顶层」来数的，别误读成「倒数第二层」。
  • 可训练参数量约 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 行拷到一起、丢弃其余分支），属于代码实现细节。

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

一轮 EAGLE 解码的完整流程

五个阶段循环执行：取特征 → 草稿建树 → 一次前向验证 → 沿树接受 → 更新 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

五、三代之间的变化

EAGLE-1 静态树与 EAGLE-2 动态树对比

EAGLE-3 的两个改动

EAGLE（2024.01） EAGLE-2（2024.06） EAGLE-3（2025.03）
草稿输入 特征 + 超前一步 token 同左 同左，但改成多层特征融合
预测目标 特征 + token（两个损失） 同左 只预测 token
草稿树 静态：每层都铺 top-k 动态：按置信度全局分配节点 沿用动态树，深度 6 → 8
训练方式 单步 teacher forcing 同左 training-time test：训练时就跑多步
解决的瓶颈 草稿模型不够「准」 静态树把预算浪费在低概率分支 特征预测约束封死了「堆数据」的收益
加速比 2.78× ~ 3.07× 3.05× ~ 4.26× 最高 6.5×

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 size 权衡

投机解码主要优化延迟，不是吞吐。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

发表于 更新于
本文字数： 6.9k 阅读时长 ≈ 25 分钟

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 简要总结

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

问题 典型表现 SGLang 中相关技术
模型权重大 单卡放不下权重，单层 GEMM 过大 TP、PP、EP、量化
上下文长 prefill TTFT 高，kv cache 存储压力大，长请求阻塞短请求 Chunked Prefill、CP、DP、prefix cache、调度预算、PD 分离
MoE 负载不均衡 expert 参数巨大，token 路由不均衡，decode 阶段 all-to-all 延迟明显 EP、DeepEP、TBO/SBO(计算和通信重叠)、EPLB

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

技术 定位 主要收益 主要代价
Chunked Prefill 调度和显存水位控制技术 不开chunk，长输入的峰值显存高导致oom服务崩溃。开chunk后，分块多轮forward降低峰值显存，解决oom问题，同时可以设置mem-fraction-static更大，增加kv cache上下文存储空间。 ttft升高：单个长请求需要多轮 forward；增加调度、kernel launch 开销
TP 最通用的权重切分手段 切分模型权重，降低单 rank 存储和计算压力 小 GEMM 利用率下降，引入额外通信代价(all-reduce)
PP 跨设备/跨节点承载模型的 stage 并行 降低单卡权重显存压力，并通过流水线方式提高多卡利用率 引入额外通信代价(p2p)；存在 pipeline bubble；
DP(Attention) attention 并行手段 复制注意力权重，切分请求到不同 rank，主要降低单 rank 存储压力 引入额外通信代价(all-gather)
EP MoE 专家并行 切分专家权重，降低单 rank 存储和计算压力 引入额外通信代价(all-to-all)
CP 长上下文 prefill 的 sequence/context 并行 长上下文激活和 KVCache 可以分摊到多卡，主要降低单 rank 存储压力 引入额外通信代价(all-gather)

2. SGLang 并行推理概览

2.1 SGLang 请求的推理流程

1

  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

2
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)

全局 rank Attention 坐标 MoE 坐标
0 DPA0, CP0, AttnTP0 EP0
1 DPA0, CP0, AttnTP1 EP1
2 DPA0, CP1, AttnTP0 EP2
3 DPA0, CP1, AttnTP1 EP3
4 DPA0, CP2, AttnTP0 EP4
5 DPA0, CP2, AttnTP1 EP5
6 DPA0, CP3, AttnTP0 EP6
7 DPA0, CP3, AttnTP1 EP7
8 DPA1, CP0, AttnTP0 EP8
9 DPA1, CP0, AttnTP1 EP9
10 DPA1, CP1, AttnTP0 EP10
11 DPA1, CP1, AttnTP1 EP11
12 DPA1, CP2, AttnTP0 EP12
13 DPA1, CP2, AttnTP1 EP13
14 DPA1, CP3, AttnTP0 EP14
15 DPA1, CP3, AttnTP1 EP15

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 层才按列切分

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

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

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

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

RowParallelLinear:
数学原理：
为了匹配矩阵乘法规则，输入 XX 也必须按列切分为 [X1​,X2​,...,Xp​][X_1​,X_2​,...,X_p​]（这正好是 ColumnParallelLinear 的输出格式）。每个 GPU 计算 Yi=XiAiY_i=X_iA_i​。注意，YiY_i 的形状与最终输出 YY 相同，但它只是部分和。最终输出 Y=∑YiY = \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。

4
其目的是将一个大的张量并行（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 归约的策略。

5

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 状态）

7

重叠点：
当微批次 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 计算重叠

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

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

Flag H 卡值 含义
enable_combine_down_gemm_two_stream_overlap True（如果 backend 是 deep_gemm） combine 和 down_gemm 双 stream 重叠
enable_dispatch_shared_one_stream_overlap True dispatch 和 shared experts 同 stream 重叠
enable_combine_shared_two_stream_overlap False —
fuse_shared_experts_inside_sbo True shared experts 在 SBO 窗口内执行

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：

9

  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，流水线效率大打折扣。（上图中白色部分随着轮次越来越少）

10
11
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。

12
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 拼接成一维后的张量进行取模
13

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

约束差异

round-robin-split in-seq-split
batch 约束 支持 B ≥ 1 强制 B=1
切分约束 S % N == 0 S ≥ 2N
负载均衡 token 均衡 使用 zigzag 保证 Attention 计算均衡

4. 总结

4.1 策略对指标的影响

TTFT / Prefill

并行 通常影响 原因
TP 降低 TTFT，但收益与通信开销相关 **优势：**GEMM 权重读取和计算大致分到多个 GPU；
**代价：**o_proj/MLP 结果需要 all-reduce/gather。
长 prefill 中 GEMM 足够大，TP 通常有收益；小 batch/过大 TP 时通信和小 GEMM 效率会反噬。
PP 单请求 TTFT 不一定下降，常可能上升 **优势：**GEMM 权重读取和计算大致分到多个 GPU；
**代价：**stage 边界通信和排队可能抵消收益。
DP Attention 高并发 prefill 更有价值，单请求收益有限 **优势：**多请求 continuous batching 下，每个 attention DP rank 处理更少 KV/token ；
**代价：**dp_gather_partial / dp_scatter / reduce_scatter 会引入同步。
EP MoE prefill 常降低显存和专家权重读压力 **优势：**experts 分布到不同 GPU，每张卡只放/读部分 expert 权重。
代价：all-to-all 通信；若 batch 很小，dispatch/combine 延迟会更明显。
CP 长上下文 prefill 开启降低每个 rank 的 kv cache 压力 **优势：**Q token 被切到不同 CP rank，单 rank attention query 数下降；SGLang 用 zigzag/round-robin 缓解 causal mask 下前后块负载不均。
代价： K/V/index K 的 CP all-gather。

TPOT / Decode

并行 通常影响 原因
TP 降低 TPOT，但收益与通信开销相关 **优势：**GEMM 权重读取和计算大致分到多个 GPU；
**代价：**o_proj/MLP 结果需要 all-reduce/gather。
PP 不适合降低单请求 TPOT 每生成一个 token 都要按 stage 顺序过完整模型，并跨 stage 传 hidden。对单请求 decode latency，PP 往往增加 pipeline stage 等待和通信；它更偏容量扩展/多请求吞吐。
DP Attention 对高并发 TPS 提升有帮助，对单请求有限 **优势：**多请求 continuous batching 下，每个 attention DP rank 处理更少 KV/token ；
**代价：**dp_gather_partial / dp_scatter / reduce_scatter 会引入同步。
EP 高并发/大 batch decode 有收益，使用 tbo/sbo 实现通信和计算重叠 **优势：**experts 分布到不同 GPU，每张卡只放/读部分 expert 权重。
代价：all-to-all 通信；若 batch 很小，dispatch/combine 延迟会更明显。

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。

发表于 更新于 分类于
本文字数： 1.4k 阅读时长 ≈ 5 分钟

mHC Manifold-Constrained Hyper-Connections

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

计算流程

以下代码来自 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 做预处理。

hc_pre

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

参数 类型 维度 说明
x torch.Tensor […, hidden_size] 输入激活值，通常是上一层输出
hc_fn torch.Tensor [hc_mult, hidden_size] 混合函数权重，定义了如何将输入映射到混合空间
hc_scale torch.Tensor [hc_mult] Sinkhorn 迭代的缩放因子
hc_base torch.Tensor [hc_mult] Sinkhorn 迭代的基础偏置项

返回值:

参数 类型 维度 说明
y torch.Tensor […, hidden_size] 预处理输出：对 x 按混合权重加权后的结果
post torch.Tensor […, hc_mult] 后处理权重：用于后续 hc_post 阶段
comb torch.Tensor […, hc_mult, hc_mult] 组合矩阵：表示各混合通道间的连接关系

具体计算分为两个步骤：

  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

mixes=(Whc⋅x)⋅RMSNorm(x)\text{mixes} = (W_{hc} \cdot x) \cdot \text{RMSNorm}(x)

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

  1. 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)

输入参数：

参数名 类型 维度 输入/输出 说明
mixes float32 (n, mix_hc) 输入 混合参数，其中 mix_hc = (2 + hc) * hc
hc_scale float32 (3,) 输入 三部分（pre/post/comb）的缩放系数
hc_base float32 (mix_hc,) 输入 基础偏置参数
hc_mult int 标量 输入 隐藏维度参数，控制 pre/post/comb 的输出维度
hc_sinkhorn_iters int 标量 输入 Sinkhorn 归一化的迭代次数
hc_eps float 标量 输入 数值稳定性参数，防止除零，用于 sigmoid 和归一化

输出参数：

参数名 类型 维度 输入/输出 说明
pre float32 (n, hc) 输出 前置参数，sigmoid 激活结果
post float32 (n, hc) 输出 后置参数，sigmoid 激活结果乘以 2
comb float32 (n, hc, hc) 输出 组合矩阵，经过 Sinkhorn 归一化的双随机矩阵

其中输出的 pre 和 post 是由 mixes、hc_scale 以及 hc_base 计算得到的：

sigmoid(x)=11+e−x\text sigmoid(x) = \dfrac{1}{1 + e^{-x}}

pre=sigmoid(mixes[0,j]⋅scale[0]+base[0，mult])\text{pre} = \text sigmoid\left(mixes[0,j] \cdot scale[0] + base[0，mult]\right)

post=2∗sigmoid(mixes[mult,2∗mult]⋅scale[1]+base[mult,2∗mult])\text{post} = \text 2 * sigmoid\left(mixes[mult,2*mult] \cdot scale[1] + base[mult,2*mult]\right)

comb=mixes[2∗mult,2∗mult+mult∗mult]⋅scale[2]+base[2∗mult,2∗mult+mult∗mult]\text{comb} = mixes[2*mult,2*mult+mult*mult] \cdot scale[2] + base[2*mult,2*mult+mult*mult]

在分出这三个 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)

输入参数：

参数 形状 含义
x (n, hidden_dim) 当前层输出
residual (n, hc_mult, hidden_dim) 来自 hc_pre 的多路残差
post (n, hc_mult) 后处理权重
comb (n, hc_mult, hc_mult) 组合权重矩阵

输出参数：

参数 形状 含义
output (n, hc_mult, hidden_dim) 当前层输出
公式：

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

总结

mhc 诞生主要为了解决：

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

发表于 更新于 分类于
本文字数： 5.9k 阅读时长 ≈ 21 分钟

大模型Attention发展历程

1

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 的计算，也就是我们说的公式:

Attention(Q,K,V)=softmax(QKTdk)V\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 更多的关注？除以 dk\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

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
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)
  1. 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)
  1. 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 一下。

  1. 位置编码与拼接
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 的时候，有一个小技巧叫做“矩阵吸收”。

xWdownQWup,iQ(Wup,iK)T(cKV)Tx W_{down}^Q W_{up,i}^Q (W_{up,i}^K{})^T (c^{KV}{})^T

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

xWdownQWup,iQRmRnT(Wup,iK)T(cKV)Tx 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 分开处理，这样位置编码和内容编码就可以分开处理，互不干扰。

  1. 注意力计算
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

公式如下：

It,s=∑j=1HIwt,jI⋅ReLU⁡(qt,jI⋅ksI)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)

总体框架如下：
3
其中，绿色部分是 indexer 实现。而框住的部分是函数 fp8_index 部分。可以看到，虽然说 attention 的计算复杂度降低了，但是实际上选择不同 token 重要性这个 indexer 模块，也算是个 attention。还是需要计算查询相对于之前每个 token 的 attention 的重要性分数，复杂度还是 O(n2)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

DeepSeek V4 Attention: Sliding Window Attention + Compressed Sparse Attention + Heavily Compressed Attention

以下代码来自官方模型内部的 model.py

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
def forward(self, x: torch.Tensor, start_pos: int):
    bsz, seqlen, _ = x.size()
    freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
    win = self.window_size
    ratio = self.compress_ratio
    rd = self.rope_head_dim
    if self.compress_ratio and self.compressor.kv_cache is None:
        self.compressor.kv_cache = self.kv_cache[:, win:]
        self.compressor.freqs_cis = self.freqs_cis
        if self.indexer is not None:
            self.indexer.freqs_cis = self.freqs_cis
    # q
    qr = q = self.q_norm(self.wq_a(x))
    q = self.wq_b(q).unflatten(-1, (self.n_local_heads, self.head_dim))
    q *= torch.rsqrt(q.square().mean(-1, keepdim=True) + self.eps)
    apply_rotary_emb(q[..., -rd:], freqs_cis)

    # win kv & topk_idxs
    kv = self.wkv(x)
    kv = self.kv_norm(kv)
    apply_rotary_emb(kv[..., -rd:], freqs_cis)
    # FP8-simulate non-rope dims to match QAT; rope dims stay bf16 for positional precision
    act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True)
    topk_idxs = get_window_topk_idxs(win, bsz, seqlen, start_pos)
    if self.compress_ratio:
        offset = kv.size(1) if start_pos == 0 else win
        if self.indexer is not None:
            compress_topk_idxs = self.indexer(x, qr, start_pos, offset)
        else:
            compress_topk_idxs = get_compress_topk_idxs(ratio, bsz, seqlen, start_pos, offset)
        topk_idxs = torch.cat([topk_idxs, compress_topk_idxs], dim=-1)
    topk_idxs = topk_idxs.int()

    # compress kv & attn
    if start_pos == 0:
        if seqlen <= win:
            self.kv_cache[:bsz, :seqlen] = kv
        else:
            cutoff = seqlen % win
            self.kv_cache[:bsz, cutoff: win], self.kv_cache[:bsz, :cutoff] = kv[:, -win:].split([win - cutoff, cutoff], dim=1)
        if self.compress_ratio:
            if (kv_compress := self.compressor(x, start_pos)) is not None:
                kv = torch.cat([kv, kv_compress], dim=1)
        # We performed QAT here, kv could also use fp8 format, though current implementation uses bf16
        o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
    else:
        self.kv_cache[:bsz, start_pos % win] = kv.squeeze(1)
        if self.compress_ratio:
            self.compressor(x, start_pos)
        o = sparse_attn(q, self.kv_cache[:bsz], self.attn_sink, topk_idxs, self.softmax_scale)
    apply_rotary_emb(o[..., -rd:], freqs_cis, True)

    # o
    o = o.view(bsz, seqlen, self.n_local_groups, -1)
    wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
    # NOTE: wo_a is FP8 in checkpoint; could do FP8 einsum here for better perf,
    # but using BF16 for simplicity.
    o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
    x = self.wo_b(o.flatten(2))
    return x

对于 SWA，没有 indexer 也没有 compressor, compress_ratio = 0。
对于 CSA，既有 indexer 也有 compressor, compress_ratio = 4。
对于 HCA，没有 indexer 只有 compressor, compress_ratio = 128。
关于 attention sink：
如何让大模型处理超出其预训练长度的对话？直接的方案是使用“滑动窗口”，即只保留最近一部分token的键值缓存（KV Cache），丢掉最旧的。这种方法虽然高效，但实验结果却出乎意料地失败了。一旦最开始的几个 token（哪怕只是一个起始符或 the 之类的普通词）被从缓存中移除，模型的困惑度（Perplexity）就会急剧飙升，输出的文本瞬间变得语无伦次。这个现象令人费解，因为这些初始token在语义上似乎并不重要，但移除它们却导致了灾难性的后果。
后来原因找到了：在模型的多数层中，无论当前模型在预测哪个词，总有大量的注意力会“回头看”，集中在序列最开始的几个 token 上。因为 softmax 函数有个硬性约束，概率和为1.这意味着注意力机制不能“弃权”，必须把注意力分配到某个地方。当一个注意力头找不到特别相关的上下文时，它就会把这份“预算”投向那些因在训练中频繁出现而具有轻微偏见的初始位置，这些位置逐渐演变成专门存放闲置注意力的“泄压阀”。
这个发现让之前滑动窗口失败的原因豁然开朗。当移除初始 token 时，我们不只是丢掉了旧的上下文，而是破坏了 softmax 函数分母中的一个重要部分，从而摧毁了模型维持注意力稳定的基础机制。解决方案因此变得异常简单：既然模型如此依赖这些注意力池，那我们就永远不要丢掉它们。

对于 compressor 部分，代码如下：

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
def forward(self, x: torch.Tensor, start_pos: int):
    assert self.kv_cache is not None
    bsz, seqlen, _ = x.size()
    ratio, overlap, d, rd = self.compress_ratio, self.overlap, self.head_dim, self.rope_head_dim
    dtype = x.dtype
    # compression need fp32
    x = x.float()
    kv = self.wkv(x) # 每个 token 要被压缩的值表示
    score = self.wgate(x) # 每个 token 在压缩池化中的 gate 分数
    if start_pos == 0:
        should_compress = seqlen >= ratio
        remainder = seqlen % ratio
        cutoff = seqlen - remainder
        offset = ratio if overlap else 0
        if overlap and cutoff >= ratio:
            self.kv_state[:bsz, :ratio] = kv[:, cutoff-ratio : cutoff]
            self.score_state[:bsz, :ratio] = score[:, cutoff-ratio : cutoff] + self.ape
        if remainder > 0:
            kv, self.kv_state[:bsz, offset : offset+remainder] = kv.split([cutoff, remainder], dim=1)
            self.score_state[:bsz, offset : offset+remainder] = score[:, cutoff:] + self.ape[:remainder]
            score = score[:, :cutoff]
        kv = kv.unflatten(1, (-1, ratio)) # 按照比例压缩
        score = score.unflatten(1, (-1, ratio)) + self.ape
        if overlap:
            kv = self.overlap_transform(kv, 0)
            score = self.overlap_transform(score, float("-inf"))
        kv = (kv * score.softmax(dim=2)).sum(dim=2) # 压缩关键步骤！
    else:
        should_compress = (start_pos + 1) % self.compress_ratio == 0
        score += self.ape[start_pos % ratio]
        if overlap:
            self.kv_state[:bsz, ratio + start_pos % ratio] = kv.squeeze(1)
            self.score_state[:bsz, ratio + start_pos % ratio] = score.squeeze(1)
            if should_compress:
                kv_state = torch.cat([self.kv_state[:bsz, :ratio, :d], self.kv_state[:bsz, ratio:, d:]], dim=1)
                score_state = torch.cat([self.score_state[:bsz, :ratio, :d], self.score_state[:bsz, ratio:, d:]], dim=1)
                kv = (kv_state * score_state.softmax(dim=1)).sum(dim=1, keepdim=True)
                self.kv_state[:bsz, :ratio] = self.kv_state[:bsz, ratio:]
                self.score_state[:bsz, :ratio] = self.score_state[:bsz, ratio:]
        else:
            self.kv_state[:bsz, start_pos % ratio] = kv.squeeze(1)
            self.score_state[:bsz, start_pos % ratio] = score.squeeze(1)
            if should_compress:
                kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True)
    if not should_compress:
        return
    kv = self.norm(kv.to(dtype))
    if start_pos == 0:
        freqs_cis = self.freqs_cis[:cutoff:ratio]
    else:
        freqs_cis = self.freqs_cis[start_pos + 1 - self.compress_ratio].unsqueeze(0)
    apply_rotary_emb(kv[..., -rd:], freqs_cis)
    if self.rotate:
        kv = rotate_activation(kv)
        fp4_act_quant(kv, fp4_block_size, True)
    else:
        act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True)
    if start_pos == 0:
        self.kv_cache[:bsz, :seqlen // ratio] = kv
    else:
        self.kv_cache[:bsz, start_pos // ratio] = kv.squeeze(1)
    return kv

总体框架如下：
5

对于 compressor，主要将多个 kv cache 按照压缩率压缩：

1
2
3
4
5
token  0 1 2 3 | 4 5 6 7 | 8 9 10 11  
        ↓          ↓           ↓  
        score softmax+加权求和  
        ↓          ↓           ↓  
kv_cache[0]   kv_cache[1]   kv_cache[2]

对于 indexer部分 代码如下：

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
def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, offset: int):
    bsz, seqlen, _ = x.size()
    freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
    ratio = self.compress_ratio
    rd = self.rope_head_dim
    end_pos = start_pos + seqlen
    if self.compressor.kv_cache is None:
        self.compressor.kv_cache = self.kv_cache
        self.compressor.freqs_cis = self.freqs_cis
    q = self.wq_b(qr)
    q = q.unflatten(-1, (self.n_local_heads, self.head_dim))
    apply_rotary_emb(q[..., -rd:], freqs_cis)
    q = rotate_activation(q)
    # use fp4 simulation for q and kv in indexer
    fp4_act_quant(q, fp4_block_size, True)
    self.compressor(x, start_pos)
    weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)
    # We performed QAT here, kv could also use fp8 format, though current implementation uses bf16
    index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, :end_pos // ratio])
    index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)
    if world_size > 1:
        dist.all_reduce(index_score)
    if start_pos == 0:
        mask = torch.arange(seqlen // ratio).repeat(seqlen, 1) >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio
        index_score += torch.where(mask, float("-inf"), 0)
    topk_idxs = index_score.topk(min(self.index_topk, end_pos // ratio), dim=-1)[1]
    if start_pos == 0:
        mask = topk_idxs >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio
        topk_idxs = torch.where(mask, -1, topk_idxs + offset)
    else:
        topk_idxs += offset
    return topk_idxs

总体框架如下：
6

indexer 作用和 DSA 完全相同，不同的事，筛选前的 kv cache 是经过 compressor 压缩的 kv cache。即：

1
2
3
4
5
256 个原始 token
  ↓ 压缩 4 倍 (ratio=4)
64 个压缩 KV
  ↓ Indexer 打分 + topk
从这 64 个里选出 index_topk 个 id

发表于 更新于 分类于
本文字数： 1.8k 阅读时长 ≈ 6 分钟

AWQ Activation-aware Weight Quantization（感知激活的权重量化）

该量化方式为 MIT 韩松团队 2023 年提出的大模型训练后 （PTQ Post Training Quantization）量化方案，核心方法是通过激活值分布识别重要的权重通道，用数学缩放来减小量化误差。

背景

AWQ 的设计基于两个重要的实验发现：

  1. 权重重要性并非对等的，而是不均衡的：大模型中只有 0.1% - 1% 的显著权重通道对于模型的精度有决定性作用。这部分权重如果被扰动，会造成极大的精度损失。而其他的权重对结果影响很小。他们在权重量化时，保留了部分重要权重使得该部分保持原来的精度，而其他的权重进行量化，模型的推理结果精度误差很小。
    1
  2. 权重的重要性由激活值决定：基于权重绝对值大小而挑选出的重要权重效果极差（说明并不是这么个重要法），而基于输入激活值的分布（激活能量）来识别重要的重要权重，能实现近乎无损的量化。原因是激活值越高的权重通道，量化的误差会在矩阵乘法中被持续放大，对模型的影响更高。
    2

对于第 jj 个输入通道的 NN 个元素，激活能量 sX,js_{X,j} ​的官方标准公式：

sX,j=E[X2]=1N∑i=1Nxi,j2s_{X,j} = E[X^2] = \frac{1}{N} \sum_{i=1}^{N} x_{i,j}^2

为什么用均方值 E[X²]E[X²]？
AWQ 论文对比了多种统计量，最终确定 均方值（二阶矩） 效果最优，原因如下：

统计量 缺陷 均方值的优势
均值（E[x]） 正负激活值会相互抵消，无法反映真实幅值 平方后消除符号影响，精准体现激活幅值
绝对值均值（L1） 对高幅值异常值不敏感，无法匹配误差放大规律 平方操作会放大高幅值激活的权重，完美贴合量化误差的平方级传播特性
最大值 单样本极值会导致统计结果不稳定，泛化性差 基于全量校准集取平均，统计结果稳定，适配任务分布

但是对于保留部分权重原来精度的方式，系统设计上实现很复杂。需要修改整个结构，还要设计对应的 kernel。于是他们提出，可以通过放大重要权重，与此同时缩小激活值，可以达到相同的效果。

量化步骤

公式

下面的公式时对称量化的核心公式：

Q(w)=Δ⋅round(wΔ),Δ=max⁡(∣w∣)2N−1.Q(\mathbf{w}) = \Delta \cdot \text{round}\left(\frac{\mathbf{w}}{\Delta}\right), \Delta = \frac{\max(|\mathbf{w}|)}{2^{N-1}}.

Q(w)⋅x=Δ⋅Round(wΔ)⋅xQ(\mathbf{w}) \cdot \mathbf{x} = { \Delta } \cdot \text{Round}\left(\frac{\mathbf{w}}{\Delta}\right) \cdot \mathbf{x}

其中，WW 是量化前的 fp16 权重，Δ{\Delta} 是基础量化步长（缩放因子）。NN 是量化比特数。

而 AWQ 量化算法，则是在原版基础上，加了个 sclae 缩放参数，对权重进行了进一步缩放。AWQ 的公式：

Q(w⋅s)⋅xs=Δ′⋅Round(w⋅sΔ′)⋅x⋅1sQ(\mathbf{w} \cdot s) \cdot \frac{\mathbf{x}}{s} = \Delta' \cdot \text{Round}\left(\frac{\mathbf{w} \cdot s}{\Delta'}\right) \cdot \mathbf{x} \cdot \frac{1}{s}

为什么量化误差会减小？

对于这两种算法，量化误差为：

Q(w)⋅x=Δ⋅Round(wΔ)⋅x→Err(Q(w)⋅x)=Δ⋅Err(Round(wΔ))⋅xQ(w⋅s)⋅xs=Δ′⋅Round(w⋅sΔ′)⋅x⋅1s→Err(Q(w⋅s)⋅xs)=Δ′⋅Err(Round(w⋅sΔ′))⋅x⋅1s\begin{aligned} Q(w) \cdot x = \Delta \cdot \text{Round}\left(\frac{w}{\Delta}\right) \cdot x \quad &\xrightarrow{} \quad \text{Err}\left(Q(w) \cdot x\right) = \Delta \cdot \text{Err}\left(\text{Round}\left(\frac{w}{\Delta}\right)\right) \cdot x \\ & \\ Q(w \cdot s) \cdot \frac{x}{s} = \Delta' \cdot \text{Round}\left(\frac{w \cdot s}{\Delta'}\right) \cdot x \cdot \frac{1}{s} \quad &\xrightarrow{} \quad \text{Err}\left(Q(w \cdot s) \cdot \frac{x}{s}\right) = \Delta' \cdot \text{Err}\left(\text{Round}\left(\frac{w \cdot s}{\Delta'}\right)\right) \cdot x \cdot \frac{1}{s} \end{aligned}

其中，WW 是量化前的 fp16 权重，Δ{\Delta} 是基础量化步长，Δ′{\Delta'} 是缩放权重后的量化步长，NN 是量化比特数。
可以看到在这两个公式中，对于 RoundRound 而言，他是舍入误差，误差是常数。对于量化前和量化后的步长而言，如果 ss 不是特别大的话，那么Δ≈Δ′{\Delta} \approx {\Delta'}。所以 AWQ 的量化误差是原版的 1/s1/s。

那么为什么不是 ss 越大越好？
因为我们之前假设 Δ≈Δ′{\Delta} \approx {\Delta'}。如果 ss 过大，那么 Δ′≈Δ{\Delta'} \approx {\Delta} 的假设就不成立了，则由 Δ′{\Delta'} 引起的量化误差就会增大。

如何获取合适的 ss 值？

L(s)=∥Q(W⋅s)(s−1⋅X)−WX∥\mathcal{L}(\mathbf{s}) = \left\| Q(\mathbf{W} \cdot \mathbf{s}) \left( \mathbf{s}^{-1} \cdot \mathbf{X} \right) - \mathbf{W}\mathbf{X} \right\|

s=sXα,α∗=arg⁡min⁡αL(sXα)\mathbf{s} = \mathbf{s}_\mathbf{X}^\alpha, \quad \alpha^* = \arg\min_\alpha \mathcal{L}(\mathbf{s}_\mathbf{X}^\alpha)

找到能让损失函数最小的 α\alpha 值，通过网格搜索来解决。其中，sXs_{X} 是上面提到的激活能量。
AWQ 的 α\alpha 网格搜索是逐层独立执行的（Transformer 的每个 Linear 层单独搜索最优 α\alpha，不同层的激活分布差异大，独立搜索精度更高）。
寻找流程如下：

  1. 搜索前的固定准备（不可动态调整）
    在启动网格搜索前，必须先固定所有影响量化结果的配置。量化比特数（如 4bit）、分组大小（group_size，主流 128）、量化类型（对称 / 非对称）、截断策略，这些参数会直接影响量化误差，搜索全程保持不变。

  2. 准备校准数据集：选取 128-512 条与目标任务分布一致的短文本样本（如通用对话、代码），数据量远小于 QAT/GPTQ，且无需标注，仅用于前向传播、统计激活分布。

  3. 逐通道计算固定激活统计量 sXs_{X}​：将校准数据输入模型，对目标 Linear 层执行前向传播，逐输入通道统计激活能量 sXs_{X}，得到固定的sX​向量，全程不再更新。

  4. 设定 α\alpha 的搜索区间：α∈[0,1]\alpha\in[0,1]。
    α=0，s​=1\alpha=0，s​=1，无任何通道缩放，退化为普通就近舍入量化，作为精度基线；
    α=1\alpha=1：缩放强度最大，ss ​与激活能量完全成正比，高激活通道被最大程度放大保护；
    区间外的 α\alpha（如 > 1）会导致非显著通道的量化噪声被过度放大，整体误差不降反升，因此无需搜索。

  5. 网格搜索:
    标准配置为 20 个网格点，对应步长 0.05，即 α\alpha 候选值为：0, 0.05, 0.10, 0.15, …, 0.95, 1.00。
    工程上可根据需求调整：点数越多，精度上限越高，但搜索耗时线性增加；20 个点是精度与效率的最优平衡。

  6. 最优 α\alpha 的选择与落地
    遍历完所有 α\alpha 候选值后，选择损失值最小的 α，作为当前层的最优 α；
    用最优 α\alpha 重新计算最终的通道缩放因子 s=sXα\mathbf{s} = \mathbf{s}_\mathbf{X}^\alpha，对权重执行正式的预缩放、量化、存储，完成该层的 AWQ 量化；
    对 Transformer 的所有 Linear 层，重复上述完整流程，完成整个模型的量化。

发表于 更新于 分类于
本文字数： 1.6k 阅读时长 ≈ 6 分钟

cuda性能分析工具nsys使用

安装 nsys 命令行工具

一般这个工具是随着 cuda toolkit 一起安装的。如果没有安装，可以查看网站：https://docs.nvidia.com/nsight-systems/InstallationGuide/index.html 。
ubuntu可以用如下安装：

1
2
3
4
5
6
apt update
apt install -y --no-install-recommends gnupg
echo "deb http://developer.download.nvidia.com/devtools/repos/ubuntu$(source /etc/lsb-release; echo "$DISTRIB_RELEASE" | tr -d .)/$(dpkg --print-architecture) /" | tee /etc/apt/sources.list.d/nvidia-devtools.list
apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
apt update
apt install nsight-systems-cli

安装后可以尝试查看：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
@└────> # nsys status -e
Timestamp counter supported: Yes

CPU Profiling Environment Check
Root privilege: disabled
Linux Kernel Paranoid Level = 4
Linux Distribution = Ubuntu
Linux Kernel Version = 5.15.0-105-generic: OK
Linux perf_event_open syscall available: Fail
Sampling trigger event available: Fail
Intel(c) Last Branch Record support: Not Available
CPU Profiling Environment (process-tree): Fail
CPU Profiling Environment (system-wide): Fail

See the product documentation at https://docs.nvidia.com/nsight-systems for more information,
including information on how to set the Linux Kernel Paranoid Level.

在可以使用 nsys 工具后，可以使用 nsys 来查看一些 kernel 的性能。
举个例子，我们有以下 cuda 代码：

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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
#include <bits/stdc++.h>
#include <cuda.h>
#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include <time.h>
#include <sys/time.h>

#define THREAD_PER_BLOCK 256

// baseline
__global__ void reduce0(float* d_in, float* d_out) {
    __shared__ float sdata[THREAD_PER_BLOCK];

    // each thread loads one element from global to shared mem
    unsigned int tid = threadIdx.x;
    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
    sdata[tid] = d_in[i];
    __syncthreads();

    // do reduction in shared mem
    for (unsigned int s = 1; s < blockDim.x; s *= 2) {
        if (tid % (2 * s) == 0) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    // write result for this block to global mem
    if (tid == 0) {
        d_out[blockIdx.x] = sdata[0];
    }
}

// bank conflict
__global__ void reduce1(float* d_in, float* d_out) {
    __shared__ float sdata[THREAD_PER_BLOCK];

    // each thread loads one element from global to shared mem
    unsigned int tid = threadIdx.x;
    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
    sdata[tid] = d_in[i];
    __syncthreads();

    // do reduction in shared mem
    for (unsigned int s = 1; s < blockDim.x; s *= 2) {
        int index = 2 * s * tid;
        if (index < blockDim.x) {
            sdata[index] += sdata[index + s];
        }
        __syncthreads();
    }

    // write result for this block to global mem
    if (tid == 0) {
        d_out[blockIdx.x] = sdata[0];
    }
}

bool check(float* out, float* res, int n) {
    for (int i = 0; i < n; i++) {
        if (out[i] != res[i]) {
            return false;
        }
    }
    return true;
}

int main() {
    const int N = 32 * 1024 * 1024;
    float* a = (float*)malloc(N * sizeof(float));
    float* d_a;
    cudaMalloc((void**)&d_a, N * sizeof(float));

    int block_num = N / THREAD_PER_BLOCK;
    float* out = (float*)malloc((N / THREAD_PER_BLOCK) * sizeof(float));
    float* d_out;
    cudaMalloc((void**)&d_out, (N / THREAD_PER_BLOCK) * sizeof(float));
    float* res = (float*)malloc((N / THREAD_PER_BLOCK) * sizeof(float));

    for (int i = 0; i < N; i++) {
        a[i] = 1;
    }

    for (int i = 0; i < block_num; i++) {
        float cur = 0;
        for (int j = 0; j < THREAD_PER_BLOCK; j++) {
            cur += a[i * THREAD_PER_BLOCK + j];
        }
        res[i] = cur;
    }

    cudaMemcpy(d_a, a, N * sizeof(float), cudaMemcpyHostToDevice);

    dim3 Grid(N / THREAD_PER_BLOCK, 1);
    dim3 Block(THREAD_PER_BLOCK, 1);

    reduce0<<<Grid, Block>>>(d_a, d_out);
    cudaMemcpy(out, d_out, block_num * sizeof(float), cudaMemcpyDeviceToHost);
    if (check(out, res, block_num)) { 
        printf("the ans is right\n");
    } else {
        printf("the ans is wrong\n");
        for (int i = 0; i < block_num; i++) {
            printf("%lf ", out[i]);
        }
        printf("\n");
    }

    reduce1<<<Grid, Block>>>(d_a, d_out);
    cudaMemcpy(out, d_out, block_num * sizeof(float), cudaMemcpyDeviceToHost);
    if (check(out, res, block_num)) { 
        printf("the ans is right\n");
    } else {
        printf("the ans is wrong\n");
        for (int i = 0; i < block_num; i++) {
            printf("%lf ", out[i]);
        }
        printf("\n");
    }

    cudaFree(d_a);
    cudaFree(d_out);
}

以上是一个 reduce_sum 的例子，有两个性能不同的核函数 reduce0 和 reduce1。可以使用命令行工具来看这两个工具的性能。

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
# 一个过渡方式，沿用之前 nvprof 的用法。
@└────> # nsys nvprof ./a.out

# 后续支持的方式
@└────> # nsys profile --stats=true ./a.out
...
[4/8] Executing 'osrtsum' stats report

 Time (%)  Total Time (ns)  Num Calls    Avg (ns)       Med (ns)     Min (ns)     Max (ns)      StdDev (ns)            Name         
 --------  ---------------  ---------  -------------  -------------  ---------  -------------  -------------  ----------------------
     49.8    6,560,704,940         41  160,017,193.7  100,127,248.0      2,946  3,656,490,088  561,147,066.6  poll                  
     40.5    5,331,726,596      1,591    3,351,179.5       34,054.0      1,200    405,442,288   30,016,785.5  ioctl                 
      4.7      614,593,088         60   10,243,218.1   10,295,123.0  1,532,148     14,425,405    2,880,533.4  waitpid               
      2.6      341,072,143         60    5,684,535.7    5,470,557.0    576,832     10,285,481    1,599,344.7  fork                  
      2.4      312,589,312        113    2,766,277.1       14,265.0      4,724    267,144,344   25,109,089.4  open64                
      0.0        5,684,320        144       39,474.4       11,704.5      1,108      3,826,795      317,952.4  fopen                 
      0.0        2,034,643         38       53,543.2       10,234.5      3,945      1,236,239      198,417.6  mmap64                
      0.0          607,662         10       60,766.2       56,248.0     42,036        112,182       19,745.1  sem_timedwait         
      0.0          405,925        123        3,300.2        2,289.0      1,006         77,642        7,014.0  fclose                
      0.0          391,343          4       97,835.8       80,675.0     58,047        171,946       53,486.5  pthread_create        
      0.0          154,726         19        8,143.5        4,985.0      1,004         47,713       10,904.9  mmap                  
      0.0           82,936          1       82,936.0       82,936.0     82,936         82,936            0.0  pthread_cond_wait     
      0.0           78,328          8        9,791.0        4,877.5      2,074         40,677       12,827.6  munmap                
      0.0           71,358          7       10,194.0        9,879.0      3,968         14,551        3,520.2  open                  
      0.0           51,719          3       17,239.7       13,793.0      3,372         34,554       15,874.2  fread                 
      0.0           42,644         29        1,470.5        1,295.0      1,000          5,372          801.5  fcntl                 
      0.0           41,205          1       41,205.0       41,205.0     41,205         41,205            0.0  fgets                 
      0.0           36,755         15        2,450.3        2,072.0      1,083          7,000        1,497.0  read                  
      0.0           33,701         12        2,808.4        2,360.0      1,384          6,014        1,224.1  write                 
      0.0           30,013          3       10,004.3       11,976.0      5,640         12,397        3,785.5  pipe2                 
      0.0           26,410          2       13,205.0       13,205.0     10,057         16,353        4,451.9  socket                
      0.0           12,546          2        6,273.0        6,273.0      5,777          6,769          701.4  fwrite                
      0.0           10,654          2        5,327.0        5,327.0      4,160          6,494        1,650.4  pthread_cond_broadcast
      0.0           10,412          1       10,412.0       10,412.0     10,412         10,412            0.0  pthread_mutex_trylock 
      0.0            9,158          1        9,158.0        9,158.0      9,158          9,158            0.0  connect               
      0.0            5,650          1        5,650.0        5,650.0      5,650          5,650            0.0  bind                  
      0.0            3,085          1        3,085.0        3,085.0      3,085          3,085            0.0  listen                

[5/8] Executing 'cudaapisum' stats report

 Time (%)  Total Time (ns)  Num Calls    Avg (ns)       Med (ns)      Min (ns)      Max (ns)      StdDev (ns)            Name         
 --------  ---------------  ---------  -------------  -------------  -----------  -------------  -------------  ----------------------
     77.1    1,961,601,754          2  980,800,877.0  980,800,877.0  674,626,715  1,286,975,039  432,995,652.3  cudaMalloc            
     22.2      564,147,432          2  282,073,716.0  282,073,716.0  275,230,367    288,917,065    9,677,957.0  cudaFree              
      0.7       17,105,648          3    5,701,882.7      657,298.0      434,348     16,014,002    8,931,253.0  cudaMemcpy            
      0.0          415,480          2      207,740.0      207,740.0       42,497        372,983      233,688.9  cudaLaunchKernel      
      0.0            1,368          1        1,368.0        1,368.0        1,368          1,368            0.0  cuModuleGetLoadingMode

[6/8] Executing 'gpukernsum' stats report

 Time (%)  Total Time (ns)  Instances  Avg (ns)   Med (ns)   Min (ns)  Max (ns)  StdDev (ns)      GridXYZ          BlockXYZ               Name           
 --------  ---------------  ---------  ---------  ---------  --------  --------  -----------  ----------------  --------------  -------------------------
     62.5          556,575          1  556,575.0  556,575.0   556,575   556,575          0.0  131072    1    1   256    1    1  reduce0(float *, float *)
     37.5          334,111          1  334,111.0  334,111.0   334,111   334,111          0.0  131072    1    1   256    1    1  reduce1(float *, float *)

[7/8] Executing 'gpumemtimesum' stats report

 Time (%)  Total Time (ns)  Count    Avg (ns)      Med (ns)     Min (ns)    Max (ns)   StdDev (ns)      Operation     
 --------  ---------------  -----  ------------  ------------  ----------  ----------  -----------  ------------------
     99.7       15,820,690      1  15,820,690.0  15,820,690.0  15,820,690  15,820,690          0.0  [CUDA memcpy HtoD]
      0.3           46,176      2      23,088.0      23,088.0      23,040      23,136         67.9  [CUDA memcpy DtoH]

[8/8] Executing 'gpumemsizesum' stats report

 Total (MB)  Count  Avg (MB)  Med (MB)  Min (MB)  Max (MB)  StdDev (MB)      Operation     
 ----------  -----  --------  --------  --------  --------  -----------  ------------------
    134.218      1   134.218   134.218   134.218   134.218        0.000  [CUDA memcpy HtoD]
      1.049      2     0.524     0.524     0.524     0.524        0.000  [CUDA memcpy DtoH]

可以看见，步骤 4 中调用的是 osrtsum（OS Runtime Summary），关注操作系统层面的性能数据。而如果在命令行中加入

1
2
3
4
5
6
7
8
9
10
11
@└────> # nsys profile --stats=true --trace=cuda,nvtx,cudnn,cublas ./a.out
...
[4/7] Executing 'cudaapisum' stats report

 Time (%)  Total Time (ns)  Num Calls    Avg (ns)       Med (ns)     Min (ns)   Max (ns)     StdDev (ns)            Name         
 --------  ---------------  ---------  -------------  -------------  --------  -----------  -------------  ----------------------
     94.6      282,480,873          2  141,240,436.5  141,240,436.5   269,662  282,211,211  199,362,781.2  cudaMalloc            
      5.2       15,504,039          3    5,168,013.0      826,866.0   541,742   14,135,431    7,767,320.2  cudaMemcpy            
      0.1          392,989          2      196,494.5      196,494.5    96,115      296,874      141,958.1  cudaFree              
      0.1          318,607          2      159,303.5      159,303.5    28,009      290,598      185,678.5  cudaLaunchKernel      
      0.0            1,184          1        1,184.0        1,184.0     1,184        1,184            0.0  cuModuleGetLoadingMode

则使用的是 cudaapisum（CUDA API Summary），关注的是 CUDA API 层面的性能数据。
除此之外，可以看到 gpukernsum 中核函数执行时间。reduce1 函数性能优于 reduce0，他们的执行次数，执行时间最大值最小值和平均值。

若算子是用 pybind 绑定，用 python 调用的，可以使用 torch 的函数来只监控该算子。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import torch
# warmup
for i in range(10):
    _ = torch.matmul(x, y) # x 和 y 是矩阵，这里就不展开了

# start profiling
torch.cuda.cudart().cudaProfilerStart()

### benchmarking
for i in range(100): # 测试 100 次
    torch.cuda.nvtx.range_push("your_ops_name")
    _ = your_ops(x, y)
    torch.cuda.nvtx.range_pop()
torch.cuda.synchronize()

# stop profiling
torch.cuda.cudart().cudaProfilerStop()

安装 Nsight Systems 可视化工具

在进行完命令行分析后，会生成一个报告文件，结尾是 .nsys-rep。这个文件可以下载下来，丢进 nsight-system 可视化软件，在软件中可以看到更加详细的数据以及程序执行的时间线。
1从上图可以看出 cuda hardware 的函数执行时间情况和 cpu 侧的执行时间情况。下面是核函数的发射时间，上面 device 侧是核函数实际执行时间，和 gpukernsum 统计的时间一致。
2上图是在软件中可以看见的较为详细的统计数据，和命令行结果一致。

发表于 更新于 分类于
本文字数： 179 阅读时长 ≈ 1 分钟

模型打印

已Llama-7B hugging face版本为例：

1
2
3
4
5
6
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
device = torch.device("cuda:{}".format(gpu))
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code = True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code = True).half().to(device)
print(model)
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
LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaFlashAttention2(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)

从结构可以看出来，模型参数量为 32,000∗4,096+32∗(4,096∗4,096∗4+4,096∗11,008∗3)+4,096∗32,000=6,738,149,37632,000 * 4,096 + 32 * (4,096 * 4,096 * 4 + 4,096 * 11,008 * 3) + 4,096 * 32,000 = 6,738,149,376。所以约为7B7B。

模型图解

以输入为 10 个 token 为例：
1

发表于 更新于 分类于
本文字数： 470 阅读时长 ≈ 2 分钟

attribute ((visibility(“”)))

是 gcc 的编译器指令，用于设置在 shared object 中所修饰的符号对外的可见性。该修饰对 .a 文件不生效，只对 .so 库生效。

attribute ((visibility(“default”)))

该修饰用于修饰符号的可见性为默认对外可见。意思是通过该符号修饰的函数可以在 so 文件外访问到。

func.cpp:

1
2
3
__attribute__ ((visibility("default"))) void func1(int a) {
        cout << a << endl;
}

main.cpp:

1
2
3
4
5
6
#include <iostream>
using namespace std;
extern void func1(int a);
int main() {
    func1(10);
}
1
2
@└────> g++ func.cpp --shared -fPIC -o libfunc.so
@└────> g++ main.cpp -L./ -lfunc

之后发现是可以编译成功的。因为该符号是可见的。

1
2
@└────> nm libfunc.so | grep func
0000000000001179 T _Z5func1i

大写的 T 表示定义在 text 段，并且可被外部引用。如果你是通过编译 .o 文件再链接为 .so 文件的，还可以使用 readelf -s 查看 .o 文件的可见性.

attribute ((visibility(“hidden”)))

该修饰用于修饰符号的可见性为默认对外不可见。意思是通过该符号修饰的函数不可以在 so 文件外访问到,只能在 so 文件内部访问到。

func.cpp:

1
2
3
__attribute__ ((visibility("hidden"))) void func2(int a) {
        cout << a << endl;
}

main.cpp:

1
2
3
4
5
6
#include <iostream>
using namespace std;
extern void func2(int a);
int main() {
    func2(10);
}
1
2
3
4
5
@└────> g++ func.cpp --shared -fPIC -o libfunc.so
@└────> g++ main.cpp -L./ -lfunc
/usr/bin/ld: /tmp/cc7GABC5.o: in function `main':
fstream.cpp:(.text+0xe): undefined reference to `func2(int)'
collect2: error: ld returned 1 exit status

之后发现是可以编译失败,因为符号不可见.

1
2
@└────> nm libfunc.so | grep func
00000000000011ef T _Z5func2i

其他

1
2
__attribute__ ((visibility("internal")))
__attribute__ ((visibility("protected")))

上述两种一样是用于修饰符号, internal 对外不可见，而 protected 对外可见。
此外，在编译 so 文件时可以通过指定 -fvisibility=xxx 来指定默认的没有给出修饰的符号属性。
如：

1
@└────> gcc -fPIC -shared -o libtest.so -fvisibility=hidden test.c

这样在 test.c 中没用经过修饰的符号对外都不可见，而修饰为 default 的依旧对外可见。

发表于 更新于 分类于
本文字数： 1.7k 阅读时长 ≈ 6 分钟

概念

在 ELF 文件中，查看可以获得它的节的名字。其中有几个带有 plt 和 got 的节。

在此处，给出各节的定义如下：

  • .got：Global Offset Table，全局偏移表。这是链接器为外部符号填充的实际偏移表。
  • .plt：Procedure Linkage Table，程序链接表。他有两个作用，要么在 .got.plt 中拿到链接地址跳转，要么触发链接器去寻找地址。
  • .got.plt：是 .got 的一部分（但是是两个不同的节），是 got 专门为 plt 准备的节，包含了 plt 表需要的地址。（新版 gcc 可能将他叫为 .plt.got）
  • .rela.plt：程序链接表的重定位表，记录所有全局函数的动态链接信息，用于在程序加载时修正 plt 表中的跳转指针，使它们指向正确的地址。

实验

接下来将使用 gdb 一步一步跟着汇编走完动态链接的过程。

准备工作

实验代码如下：

1
2
3
4
5
6
#include <stdio.h>
int main() {
    puts("hello");
    printf("hello");
    return 0;
}

查看节的地址与大小：

1
2
3
4
5
6
@└────> # objdump -h plt | grep -E "plt|got"
plt:     file format elf64-x86-64
  9 .rela.plt     00000030  0000000000400468  0000000000400468  00000468  2**3
 11 .plt          00000030  00000000004004c0  00000000004004c0  000004c0  2**4
 20 .got          00000020  0000000000600fe0  0000000000600fe0  00000fe0  2**3
 21 .got.plt      00000028  0000000000601000  0000000000601000  00001000  2**3

查看需要动态链接的符号：

1
2
3
4
5
6
7
8
9
10
11
12
13
@└────> # readelf -r plt

Relocation section '.rela.dyn' at offset 0x408 contains 4 entries:
  Offset          Info           Type           Sym. Value    Sym. Name + Addend
000000600fe0  000100000006 R_X86_64_GLOB_DAT 0000000000000000 _ITM_deregisterTMClone + 0
000000600fe8  000400000006 R_X86_64_GLOB_DAT 0000000000000000 __libc_start_main@GLIBC_2.2.5 + 0
000000600ff0  000500000006 R_X86_64_GLOB_DAT 0000000000000000 __gmon_start__ + 0
000000600ff8  000600000006 R_X86_64_GLOB_DAT 0000000000000000 _ITM_registerTMCloneTa + 0

Relocation section '.rela.plt' at offset 0x468 contains 2 entries:
  Offset          Info           Type           Sym. Value    Sym. Name + Addend
000000601018  000200000007 R_X86_64_JUMP_SLO 0000000000000000 puts@GLIBC_2.2.5 + 0
000000601020  000300000007 R_X86_64_JUMP_SLO 0000000000000000 printf@GLIBC_2.2.5 + 0

反汇编查看 plt 相关函数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
@└────> # objdump -d plt
Disassembly of section .plt:

00000000004004c0 <.plt>:
  4004c0:       ff 35 42 0b 20 00       pushq  0x200b42(%rip)        # 601008 <_GLOBAL_OFFSET_TABLE_+0x8>
  4004c6:       ff 25 44 0b 20 00       jmpq   *0x200b44(%rip)        # 601010 <_GLOBAL_OFFSET_TABLE_+0x10>
  4004cc:       0f 1f 40 00             nopl   0x0(%rax)

00000000004004d0 <puts@plt>:
  4004d0:       ff 25 42 0b 20 00       jmpq   *0x200b42(%rip)        # 601018 <puts@GLIBC_2.2.5>
  4004d6:       68 00 00 00 00          pushq  $0x0
  4004db:       e9 e0 ff ff ff          jmpq   4004c0 <.plt>

00000000004004e0 <printf@plt>:
  4004e0:       ff 25 3a 0b 20 00       jmpq   *0x200b3a(%rip)        # 601020 <printf@GLIBC_2.2.5>
  4004e6:       68 01 00 00 00          pushq  $0x1
  4004eb:       e9 d0 ff ff ff          jmpq   4004c0 <.plt>

开始

  1. 首先断点到 puts 函数，查看调用处：
1
2
3
4
5
6
7
8
9
10
11
12
13
@(gdb) disassemble main
Dump of assembler code for function main:
   0x00000000004005d6 <+0>:     push   %rbp
   0x00000000004005d7 <+1>:     mov    %rsp,%rbp
=> 0x00000000004005da <+4>:     mov    $0x400698,%edi
   0x00000000004005df <+9>:     callq  0x4004d0 <puts@plt>
   0x00000000004005e4 <+14>:    mov    $0x400698,%edi
   0x00000000004005e9 <+19>:    mov    $0x0,%eax
   0x00000000004005ee <+24>:    callq  0x4004e0 <printf@plt>
   0x00000000004005f3 <+29>:    mov    $0x0,%eax
   0x00000000004005f8 <+34>:    pop    %rbp
   0x00000000004005f9 <+35>:    retq
End of assembler dump.

可以看到，调用处实际上是使用 call 指令走到 puts 的代码段。下面的 printf 也是如出一辙。

  1. 查看 puts@plt 的汇编指令
1
2
3
4
5
6
@(gdb) disassemble
Dump of assembler code for function puts@plt:
=> 0x00000000004004d0 <+0>:     jmpq   *0x200b42(%rip)        # 0x601018 <puts@got.plt>
   0x00000000004004d6 <+6>:     pushq  $0x0
   0x00000000004004db <+11>:    jmpq   0x4004c0
End of assembler dump.

可以看到，在汇编中，他首先要跳转到 0x601018 地址的位置。这个地址内容是个全局变量，实际上根据节的地址位置和大小可以判断，是处于 .got.plt 的位置内（ 0x601000 ~ 0x601028）。所以可以认为，在 .got.plt 中，存在了 puts 函数的地址。

  1. 查看 .got.plt
1
2
3
4
5
@(gdb) x/16x 0x601018
0x601018 <puts@got.plt>:        0x004004d6      0x00000000      0x004004e6      0x00000000
0x601028:       0x00000000      0x00000000      0x00000000      0x00000000
0x601038:       0x00000000      0x00000000      0x00000000      0x00000000
0x601048:       0x00000000      0x00000000      0x00000000      0x00000000

查看表中内容，发现跳转的地址是 0x4004d6，这不就是我们跳转之前的下一个地址吗！（puts@plt 的第二条指令） 同理，printf 函数也是如此（0x4004e6）。这是因为，之前没有调用过 puts 函数，第一次查找的时候，.got.plt 表中找不到函数的地址，那就先返回继续执行去调用链接器获取地址。

  1. 准备调用链接器
1
2
3
4
00000000004004d0 <puts@plt>:
  4004d0:       ff 25 42 0b 20 00       jmpq   *0x200b42(%rip)        # 601018 <puts@GLIBC_2.2.5>
  4004d6:       68 00 00 00 00          pushq  $0x0
  4004db:       e9 e0 ff ff ff          jmpq   4004c0 <.plt>

首先 pushq $0x0，这个是在 got.plt 中的编号，如 puts 是 0，printf 是 1。这个参数是给后续链接器使用的。然后跳到了 .plt 的位置执行（0x4004c0）。可以看到，printf@plt 函数最后也是跳到这个位置执行。

  1. 调用链接器
1
2
3
4
00000000004004c0 <.plt>:
  4004c0:       ff 35 42 0b 20 00       pushq  0x200b42(%rip)        # 601008 <_GLOBAL_OFFSET_TABLE_+0x8>
  4004c6:       ff 25 44 0b 20 00       jmpq   *0x200b44(%rip)        # 601010 <_GLOBAL_OFFSET_TABLE_+0x10>
  4004cc:       0f 1f 40 00             nopl   0x0(%rax)

首先 push 了 0x601008 到栈中，这是 .got.plt 表中的一个地址。之后跳转到 0x601010 所存储的地址去执行相应的代码。不难看出，0x601010 也是存储在 .got.plt 表中的。查看一下存储的内容：

1
2
3
4
@(gdb) x/10x 0x601010
0x601010:       0xf7de64a0      0x00007fff      0x004004d6      0x00000000
0x601020 <printf@got.plt>:      0x004004e6      0x00000000      0x00000000      0x00000000
0x601030:       0x00000000      0x00000000

可以看到，是让我们跳转到 0x00007ffff7de64a0 去执行相应的代码。那么这块代码是什么呢？

1
2
3
4
5
@(gdb) info sharedlibrary
From                To                  Syms Read   Shared Object Library
0x00007ffff7dd0fa0  0x00007ffff7df2cd4  Yes (*)     /lib64/ld-linux-x86-64.so.2
0x00007ffff7a2cb90  0x00007ffff7b798ad  Yes (*)     /lib64/libc.so.6
(*): Shared library is missing debugging information.

可以看到，该地址是 ld-linux-x86-64.so 加载的位置。说明执行的是链接器的代码。

1
2
3
4
5
6
7
8
9
10
1: x/5i $pc
=> 0x7ffff7de64a0 <_dl_runtime_resolve_xsavec>: endbr64
   0x7ffff7de64a4 <_dl_runtime_resolve_xsavec+4>:       push   %rbx
   0x7ffff7de64a5 <_dl_runtime_resolve_xsavec+5>:       mov    %rsp,%rbx
   0x7ffff7de64a8 <_dl_runtime_resolve_xsavec+8>:       and    $0xffffffffffffffc0,%rsp
   0x7ffff7de64ac <_dl_runtime_resolve_xsavec+12>:
    sub    0x21616d(%rip),%rsp        # 0x7ffff7ffc620 <_rtld_local_ro+384>
@(gdb) bt
#0  0x00007ffff7de64a0 in _dl_runtime_resolve_xsavec () from /lib64/ld-linux-x86-64.so.2
#1  0x00000000004005e4 in main () at plt.c:3

可以看到这里代码执行的是 ld 中的 _dl_runtime_resolve_xsavec 函数是第一次函数调用时用于查找函数符号的，并且在结尾处会直接去调用找到的函数符号（本文中为 puts 函数）。

  1. 写回 .got.plt 表
    在 puts 上打个断点，这样继续的话就是执行完 _dl_runtime_resolve_xsavec 还未执行 puts 的状态了。
1
2
3
4
5
6
7
@(gdb) bt
#0  0x00007ffff7a7d8c0 in puts () from /lib64/libc.so.6
#1  0x00000000004005e4 in main () at plt.c:3
@(gdb) x/10x 0x601018
0x601018 <puts@got.plt>:        0xf7a7d8c0      0x00007fff      0x004004e6      0x00000000
0x601028:       0x00000000      0x00000000      0x00000000      0x00000000
0x601038:       0x00000000      0x00000000

可以看到，此时，got.plt 表中的地址已经被写为 puts 函数实际的地址了（0x00007ffff7a7d8c0 在 0x00007ffff7a2cb90 ~ 0x00007ffff7b798ad 范围内，属于 /lib64/libc.so.6），这样下次调用 puts 就不用再次调用链接器了。

题外话

其实看一下 .got.plt 表的内容，会发现明明 puts 是第一个需要被链接的函数，为什么第一个却不是它呢？

1
2
3
4
5
6
@(gdb) x/10x 0x601000
0x601000:       0x0000000000600e10      0x00007ffff7ffe1d0
0x601010:       0x00007ffff7de64a0      0x00007ffff7a7d8c0
0x601020 <printf@got.plt>:      0x00000000004004e6      0x0000000000000000
0x601030:       0x0000000000000000      0x0000000000000000
0x601040:       0x0000000000000000      0x0000000000000000

puts 地址实际上是 got[3]:0x00007ffff7a7d8c0，前面还有 3 项。其中：

  • got[0]:0x0000000000600e10 自身模块 dynamic 段地址
1
2
@(gdb) info symbol 0x0000000000600e10
_DYNAMIC in section .dynamic of /root/xxx/plt

  • got[1]:0x00007ffff7ffe1d0 本模块的 link_map 的地址。编译期间会初始化为 0。link_map 是一个双向链表的入口，链接进程所有加载的动态库。当链接器查找符号时，通过遍历该链表找到对应的符号。

  • got[2]:0x00007ffff7de64a0 _dl_runtime_resolve_xsavec 的地址。

1
2
@(gdb) info symbol 0x00007ffff7de64a0
_dl_runtime_resolve_xsavec in section .text of /lib64/ld-linux-x86-64.so.2

_dl_runtime_resolve 格式：

1
2
3
4
//调用形式为：
_dl_runtime_resolve((link_map*)(got[1]), 0);
// 第二个参数 0，为 <puts@plt>:中的 pushq $0x0；
// 同理如果是 printf，就是<printf@plt>:中 pushq $0x1；

总结

虚拟地址空间内流程图：
1

第二次调用：
2

发表于 更新于 分类于
本文字数： 469 阅读时长 ≈ 2 分钟

为什么要使用内嵌汇编？

内嵌汇编通常用于在程序中实现一些高效、精确的操作。例如，在嵌入式平台上运行的程序，如果需要代码占用内存更小、程序运行的效率更高或需要准确地操作寄存器时，嵌入汇编会是不错的选择。

基本语法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
asm("assembly code"        /* 汇编代码 */
    :output_operand        /* 输出参数列表 */          
    :input_operand         /* 输入参数列表 */          
    :clobbered_operand     /* 被改变的操作对象列表 */      
);

// 举例
static int value_assignment(int input) {
    int ret = 0;
    asm volatile(
        "movl %1, %0\n"   // 超过一条指令就要用 \n 来分割，排版整齐还要加 \t
        :"=r"(ret)
        :"r"(input)
    );
    return ret;
}

被改变的操作对象列表

在被改变的参数列表 clobbered_operand 中有一个比较有用的标识符：memory。指定 memory，相当于对编译器形成了一个内存读写的屏障，保证在内联汇编执行前，编译器将某些寄存器里的值刷新进内存，同时在内联汇编执行后，编译器重新加载相关变量的值。
所以我们可以见到这样的代码：

1
asm volatile ("" ::: "memory");

作为内存屏障，保证编译器的优化不会跨过这道屏障。加上 volatile 告诉编译器不要优化汇编。

修饰符

修饰符一般跟在参数列表前面。

修饰符 含义
= 只写，常用于修饰所有输出操作数
无 只读
+ 可读可写
r 可以是任意通用寄存器存储其值
m 一个有效的内存地址
i 是立即数
% 被修饰的操作数可以和下一个互换
& 只能做输出，一般和 “=” 一起使用，如 “=&r(val)”
x 只能做输入

占位符

%0 表示输入和输出列表合并的第 1 个操作数，%1 表示第 2 个，以此类推。