TensorRT-LLM 源码完全拆解（二十二）：EAGLE-3 一体式草稿——三层隐状态、草稿循环与动态树验证

默认捕获层的选择

_get_eagle3_default_capture_layers 只有一行：

1
def _get_eagle3_default_capture_layers(num_layers: int):
2
    return (1, num_layers // 2 - 1, num_layers - 4)

以 32 层模型为例，捕获层是 1、15、28；以 80 层模型为例，是 1、39、76。三个位置分别落在浅层、中层和接近输出处。这个选择不是随意的：EAGLE-3 论文的消融显示，低层特征保留更多词法信息，高层特征携带语义决策，中间层是两者的折中；只取单层会让草稿模型丢失一类信息。TensorRT-LLM 把这个约定硬编码成默认值，同时允许配置覆盖：如果 EagleDecodingConfig 里显式给了 layers_to_capture，__post_init__ 会排序后直接使用。

这里有个容易忽略的保护：

1
if self.num_layers <= 5:
2
    raise ValueError("Not enough hidden layers for default EAGLE3 capture")

三层捕获要求模型至少 6 层，否则 num_layers - 4 会落到 1 附近，与浅层捕获重叠。与其静默产生重复层，不如在配置阶段直接报错。

捕获缓冲区的形状账

捕获缓冲区在 Eagle3ResourceManager 里一次性分配：

1
max_num_tokens = min(max_num_tokens, max_num_requests * max_seq_len)
2
max_num_tokens += (max_total_draft_tokens + 1) * max_num_requests
3
self.hidden_states = torch.empty(
4
    (max_num_tokens, hidden_size * config.num_capture_layers),
5
    dtype=self.dtype, device='cuda')

第一项是常规调度预算，第二项是投机解码的额外行数：每个请求在本步最多有 max_total_draft_tokens + 1 个待验证位置（树模式下是树节点数，线性模式下是草稿长度加奖励 token），这些位置的目标隐状态也要落进同一个缓冲，供下一步草稿使用。缓冲区只有一份，由 Eagle3OneModelSpecMetadata 直接引用，而不是每个 metadata 各分配一次——这样 CUDA Graph 捕获时地址稳定，也不会随 batch 数增长重复占显存。

列数 hidden_size * num_capture_layers 是三层特征的横向拼接。注意它用的是目标模型的 hidden size 乘以捕获层数，而不是草稿模型的 hidden size。压缩发生在草稿模型自己的 fc 里，缓冲区只负责原样存放。

在解码层里精确落位

捕获动作发生在目标模型的解码层内部。TensorRT-LLM 没有在每个模型文件里重复写捕获逻辑，而是把钩子放在 spec_metadata 上，由模型代码在合适的时机调用：

1
def maybe_capture_hidden_states(self, layer_id, hidden_states, residual=None):
2
    for i, captured_layer_id in enumerate(self.layers_to_capture):
3
        if captured_layer_id == layer_id:
4
            num_tokens = self.num_capture_tokens
5
            to_save = hidden_states[:num_tokens]
6
            if residual is not None:
7
                to_save = to_save + residual[:num_tokens]
8
            inplace_slice_copy(self.hidden_states, to_save,
9
                               i * self.hidden_size,
10
                               (i + 1) * self.hidden_size)
11
            break

三个细节决定了这段代码为什么长这样：

第一，hidden_states + residual 而不是只用 hidden_states。 TensorRT-LLM 的解码层采用”隐状态与残差分离”的写法以支持融合 RMSNorm：hidden_states 保存的是归一化后的主分支，残差流单独存放，直到下一层才相加。如果捕获时只拿主分支，草稿模型看到的就不是目标模型真正的层输出。加法在这里显式补上，代价是一次逐元素 kernel；相比让整个模型退回非融合路径，这个代价可以接受。

第二，inplace_slice_copy 而不是普通赋值。 目标是往大缓冲区的第 $i$ 个列块写。如果写成 self.hidden_states[:, lo:hi] = to_save，在 CUDA Graph 捕获下会产生一个 view 上的切片赋值，Triton/Inductor 有时会把它编译成额外的临时张量；inplace_slice_copy 是自定义算子，语义明确为”按列区间就地拷贝”，图捕获时形状固定、地址固定。

第三，num_capture_tokens 与 num_tokens 的分离。 prepare() 里先快照真实调度 token 数，再把 self.num_tokens 改写成”去掉草稿位之后”的子序列形状：

1
self.num_capture_tokens = self.num_tokens
2
self.num_tokens = self.dp_num_tokens()

后者是给 attention DP 用的：在 DP 模式下，各 rank 只处理自己那部分序列，形状需要按子序列算。但隐状态捕获要的是整个 batch 的真实 token，包括待验证的草稿位。如果捕获时误用被改写后的 num_tokens，最后几个草稿位置的隐状态就会缺行，草稿模型下一步拿到的上下文直接错位。源码在注释里明确写了这一点，并加了两个越界检查：缓冲区容量不足或本次前向的隐状态行数不足都会抛 RuntimeError，而不是静默截断。

num_capture_tokens 还有一层含义：CUDA Graph 会把 batch 填充到固定桶大小，填充行的隐状态是垃圾值。捕获只拷贝前 num_capture_tokens 行，正是为了不把填充行写进缓冲，避免草稿模型下一步读到无意义特征。

MTP Eagle 为什么把捕获层集合置空

__post_init__ 里有一个分支：

1
if self.spec_dec_mode.is_mtp_eagle_one_model():
2
    self.layers_to_capture = ()

MTP Eagle 复用同一个 worker 类，但它的草稿模块是 DeepSeek 式 MTP 层，直接吃目标模型最后一层隐状态，不读 spec_metadata.hidden_states。置空捕获集合后，is_layer_capture() 对所有层返回 False，模型侧就不会插入捕获钩子。捕获层集合为空时，num_capture_layers 变成 0，metadata 干脆不分配隐状态缓冲，hidden_states 保持 None——这一路径上没有任何代码会读它。注释还提到一个副作用：某些模型（DeepSeekV3、GLM）在检测到”有人在捕获隐状态”时会关闭 MLP/MoE 融合，以保证捕获到的是融合前/后的正确值；MTP Eagle 不需要捕获，自然也不该触发这个降级。

第一层与后续层的输入差异

Eagle3DecoderLayer 用一个布尔量区分首层与后续层：

1
self._next_layer_regular = (
2
    (eagle_config.get("next_layer_regular", True) and not is_first_layer)
3
    or eagle_config.get("eh_proj_before_attn", False)
4
)

首层的输入是 [embeds, hidden_states] 拼接后的 $2d$ 向量，后续层是普通 $d$ 向量。因此在 Eagle3Attention 里，首层的 qkv_proj 输入特征数被改成 2 * hidden_size，而 K/V 的输出维度不变：

1
self.qkv_proj = Linear(
2
    2 * self.hidden_size,
3
    tp_size * self.q_size + 2 * tp_size * self.kv_size,
4
    ...
5
    fused_weight_shard_indices_mapping={"q": (0, self.q_size),
6
                                        "k": (self.q_size, self.kv_size),
7
                                        "v": (self.q_size + self.kv_size, self.kv_size)})

这里的关键是 fused_weight_shard_indices_mapping：QKV 是融合成一个权重的，输入维度翻倍只改变权重矩阵的行数（输入特征维），不影响 Q/K/V 三个输出块沿列方向的相对布局。给出显式的 shard 映射后，张量并行切分仍然按 Q、K、V 三段切输出列，且每段携带自己的输入切片规则，而不是把翻倍后的输入维度平均切——后者会让 Q、K、V 各自的权重错位。

MLA 版本同理。Eagle3MLAttention 只把 kv_a_proj_with_mqa 的输入维度改成 2 * hidden_size，因为 MLA 的 Q 走 q_lora_rank 的低秩路径，只有 KV 压缩投影直接吃输入特征。这也解释了为什么 DeepSeek 系模型的 EAGLE-3 草稿层能复用 MLA 的大部分实现：改动面被限制在一个投影层上。

首层的拼接发生在 attention 之前

Eagle3DecoderLayer.forward 的顺序值得逐行看：

1
residual = hidden_states
2
hidden_states = self.hidden_norm(hidden_states)
3
if not self._next_layer_regular:
4
    embeds = self.input_layernorm(embeds)
5
    hidden_states = torch.cat([embeds, hidden_states], dim=-1)
6
hidden_states = self.self_attn(position_ids=..., hidden_states=hidden_states, ...)
7
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
8
hidden_states = self.mlp(hidden_states)
9
spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, residual)

注意三点：

• token embedding 走独立的 input_layernorm，隐状态走 hidden_norm，两者归一化后再拼接。直接拼接未归一化的 embedding 会让两路特征的尺度不一致，attention 的 QKV 投影会把 embedding 那一路放大。
• residual 只保存隐状态这一路，不包含 embedding。这是 EAGLE 草稿模型的约定：embedding 是”当前 token 的身份”，不应该进入残差流被后续层反复累加。
• 层末再次调用 maybe_capture_hidden_states，但这里捕获的是草稿模型自己的层输出，服务于多步草稿：第 $i$ 步草稿的隐状态要喂给第 $i+1$ 步。同一个 metadata 对象、同一个缓冲区，目标模型和草稿模型共用捕获接口，靠 layer_id 是否在捕获集合里区分——草稿模型的层号从 start_layer_idx 开始编号，与目标模型的捕获层号不重叠。

这个”共用捕获接口”的设计有利有弊。利是代码路径统一，CUDA Graph 下缓冲区地址一致；弊是草稿模型的层号必须与目标模型错开，否则会误写目标模型的捕获槽。源码通过 start_layer_idx 传入一个偏移来保证这一点。

首步输入的重排

目标模型这一步的输入里混着两类请求：context（prefill）请求和 generation（decode）请求。对 context 请求，草稿从最后一个 prompt token 之后开始；对 generation 请求，草稿从”本步已接受的最后一个 token”之后开始。prepare_1st_drafter_inputs 负责把这两种情况拼成一次前向的输入。

对 EAGLE-3 路径，先做特征压缩：

1
hidden_size_up = spec_metadata.hidden_size * len(spec_metadata.layers_to_capture)
2
hidden_states = spec_metadata.hidden_states[:num_tokens, :hidden_size_up]
3
hidden_states = draft_model.apply_eagle3_fc(hidden_states)

apply_eagle3_fc 就是前面那个 $\mathrm{FC}$。它的输入是捕获缓冲的前 num_tokens 行，输出是草稿模型的 $d$ 维特征。这里的 num_tokens 取的是 input_ids.shape[0]，即本次前向的真实输入行数，而不是 num_capture_tokens——捕获缓冲里多出来的草稿位行服务于下一轮草稿，本轮草稿的输入只需要与输入 token 对齐的前若干行。

然后拼 generation 部分的输入：

1
input_ids_gen = accepted_tokens[num_contexts:, :runtime_draft_len + 1].flatten()
2
input_ids = torch.concat([input_ids_ctx, input_ids_gen], dim=0)

accepted_tokens 的形状是 [batch, max_draft_len + 1]，第 0 列是本步的根 token，后面是上一轮草稿中被接受的部分。flatten() 之后，每个 generation 请求贡献连续的 runtime_draft_len + 1 行。这一步同时把 attn_metadata 的 seq_lens 设为该长度，使首步 attention 以”投机解码模式”运行：每个请求内部是因果掩码，请求之间互不可见。

首步为什么要跑 $D+1$ 个位置

线性草稿模式（非动态树）下，首步一次前向处理 $D+1$ 个 token，得到 $D+1$ 组隐状态；但只有每组请求的最后一行用于产生第一个草稿 token：

1
if i == 0:
2
    start_ids_gen = batch_indices_cuda[:num_gens] * (runtime_draft_len + 1)
3
    gather_ids_gen = start_ids_gen + num_accepted_tokens[num_contexts:] - 1 + num_ctx_tokens
4
    gather_ids = torch.concat([last_tokens_idx[:num_contexts], gather_ids_gen], dim=0)
5
else:
6
    gather_ids = slice(0, batch_size)

为什么首步要浪费 $D$ 行？因为草稿模型的 KV Cache 需要把这 $D+1$ 个 token 的 K/V 都写进去——后续每一步只喂一个新 token，但它要 attend 到”已接受路径”上的所有位置。首步一次性写入，等于把草稿侧的前缀缓存建好；如果只算最后一行，下一步就得回头补写前 $D$ 行的 KV，反而多一次前向。这与目标模型 prefill 时”一次算完整段、只取最后一行 logits”是同一个思路。

gather_ids_gen 的计算里 num_accepted_tokens - 1 是关键：首步的 $D+1$ 行里，只有前 num_accepted_tokens 行是真实已接受 token，其余是填充。要取的是真实路径的最后一个，所以要加上接受数再减一。num_ctx_tokens 把 context 部分的偏移补上，使 gather 在同一块扁平张量里寻址。

后续步的 KV 长度修正

首步之后，每个请求只喂一个 token。此时 attn_metadata 的 seq_lens 被改成 1，num_contexts 被置 0，所有行都变成 generation 行：

1
attn_metadata._seq_lens[:batch_size].fill_(1)
2
attn_metadata._seq_lens_cuda[:batch_size].fill_(1)
3
attn_metadata.num_contexts = 0
4
attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= (runtime_draft_len - num_accepted_tokens[num_contexts:])
5
attn_metadata.kv_lens_cuda[:num_contexts] += 1
6
attn_metadata.use_spec_decoding = False

kv_lens_cuda 的修正有一个容易读错的符号。目标模型这一步为每个 generation 请求写入了 runtime_draft_len + 1 个位置的 KV，但其中只有 num_accepted_tokens 个是真实被接受的。草稿模型要 attend 的是接受路径，所以要从当前 KV 长度里减掉”多写的草稿位”：runtime_draft_len - num_accepted_tokens。减完之后，草稿侧看到的上下文长度恰好等于目标模型的实际序列长度。

context 请求的 +1 则是另一回事：prefill 请求在草稿开始时还没有”下一个 token”，首步用 prompt 的最后一个 token 作为草稿输入，草稿 KV 里要多算这一个位置。

use_spec_decoding = False 同样重要。首步需要投机解码掩码（一个请求内的 $D+1$ 行互相因果可见），后续步每个请求只有一行，普通因果掩码即可。如果不在首步后关掉这个标志，attention 后端会继续按 $D+1$ 的 query 长度解释只有 1 行的输入，掩码和位置偏移全部错位。

每次前向只取一行 logits

后续步的前向输出是 [batch_size, d]，因为每个请求只有一行。logits_processor 只对这一行做 LM Head 投影，得到 [batch_size, vocab]，再采样出下一个草稿 token。采样结果被压进 next_draft_tokens 列表，循环结束后 torch.stack(..., dim=1) 变成 [batch_size, runtime_draft_len]。

这里有一个 CUDA Graph 友好的设计：循环次数 runtime_draft_len 是 Python 整数，在捕获时固定；每一步的输入形状都是 [batch_size, ...]，不随步数变化。因此整个草稿循环可以被捕获进同一张图，运行时只改指针指向的数据，不改形状。动态树模式打破了这个性质——树深固定但每层节点数按 $K$ 增长，所以它单独实现了一套预分配缓冲，下一节展开。

配置层面的三个约束

打开 use_dynamic_tree=True 之前，EagleDecodingConfig 的校验器已经替工程侧挡掉了三种无效组合：

• dynamic_tree_max_topK 必须显式给出且大于 0。它不参与”是否启用树”的判断，只控制每层分叉数；如果只设 topK 而不设 use_dynamic_tree，校验器会打一条 warning 并把 topK 清空，然后退回线性草稿。
• max_total_draft_tokens 缺省时取 dynamic_tree_max_topK * max_draft_len。这是树上所有候选节点的预算上限。
• 若显式给出 max_total_draft_tokens，它必须落在 [max_draft_len, dynamic_tree_max_topK * max_draft_len] 区间内。下界保证树至少能覆盖线性草稿的深度，上界保证预分配的树缓冲区不会超过”每层最多 $K$ 个节点”这一布局假设。

此外，_max_batch_size 是私有属性，由 py_executor_creator 从全局 max_batch_size 注入。动态树 worker 的所有 CUDA 缓冲都按这个值预分配，运行时用 batch 下标直接寻址、不做越界检查。如果它比全局 batch 小，warmup 或正式生成时一旦 batch 下标超界就是非法访存，而不是一个可恢复的 Python 异常。这也是官方文档强调”不要把 max_batch_size 直接传给 Eagle3DecodingConfig”的原因。

为什么需要树

线性草稿每步只保留一个 token，一旦某步猜错，后面的草稿全部作废。EAGLE-2 的做法是每步保留 top-K 个候选，形成一棵树，让目标模型一次性验证所有分支（EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees）。树的价值在于：验证成本几乎不变（一次前向多算几行），但命中概率显著提高。

图中的两种草稿形态对应 TensorRT-LLM 的两条代码路径。上半部分的链式草稿每次只押注一个 token，验证时只会有”全中”或”从某处截断”两种结果，代码对应 _forward_linear_draft_loop；下半部分的树式草稿在每一层分叉，验证时需要在树上找一条最长的一致路径，代码对应 Eagle3OneModelDynamicTreeWorker。树的分叉越多，同一个 $D$ 能覆盖的候选序列就越多，但掩码与验证内核也要相应改写——下一节就从这个改写开始。

代价是注意力掩码不再是因果掩码。树上的每个节点只能看到自己的祖先链，而不是”位置编号小于自己的所有节点”。位置编号也不再等于深度，因为同一深度的不同分支共享前缀。

树的三套账：父指针、掩码、位置

TensorRT-LLM 用三个并行的数据结构描述一棵树：

其中 $N =$ max_total_draft_tokens + 1，包含根节点。父指针是构建树的输入，掩码和位置是构建树的输出。

草稿循环每一层产生 $K \times K$ 个候选：上一层的 $K$ 个节点各自再取 top-K。_select_topk_draft_tokens 把”父节点分数 × 子节点分数”作为累计分数，在 $K^2$ 个候选里再取 top-K，得到本层保留的 $K$ 个节点及其父下标：

1
all_scores_flat = (new_draft_scores * previous_draft_scores.unsqueeze(2)).reshape(-1, K * K)
2
topk_values, topk_indices = torch.topk(all_scores_flat, k=K, dim=-1)
3
real_tokens = torch.gather(all_tokens_flat, dim=1, index=topk_indices)
4
parents = topk_indices // K

用乘积而不是相加，是 EAGLE-2 动态树的评分约定：路径概率是沿途条件概率的乘积。实现上用对数域相加会更稳定，但这里保留了原始概率的乘积形式，因为分数只用于排序，不参与最终的接受判定；在 bf16 下乘积可能下溢，不过 top-K 的相对顺序在候选分数同量级时仍然保持。源码没有做 log 转换，这是一个已知的数值取舍。

parents = topk_indices // K 把扁平下标还原成”来自上一层哪个父节点”。因为候选排列是 [父0的子0..K-1, 父1的子0..K-1, ...]，整除 $K$ 正好得到父下标。

树构建内核做了什么

DynamicTreeOpsConverter.build_dynamic_tree 调用 torch.ops.trtllm.build_dynamic_tree_op，把父指针、选中下标转换成三样东西：祖先掩码、位置偏移、以及检索用的链表（retrieve_index、retrieve_next_token、retrieve_next_sibling）。

调用前先把输出缓冲清零并填 -1：

1
tree_mask[:bs].zero_()
2
positions[:bs].zero_()
3
retrieve_index[:bs].zero_()
4
retrieve_next_token[:bs].fill_(-1)
5
retrieve_next_sibling[:bs].fill_(-1)

这一步不是可有可无的初始化。缓冲区是跨请求、跨步复用的，CUDA Graph 捕获后每步只覆写”活跃节点”对应的位置；如果上一棵树有更深的节点，残留的 1 位会让本棵树的浅层节点错误地 attend 到不存在的祖先。清零保证了掩码语义只由本步的树决定。

tree_mask_mode 有两个取值：2 表示位压缩（QLEN_ONLY_BITPACKING），1 表示不压缩。位压缩把 $N \times N$ 的布尔掩码压成 $N \times \lceil N/32 \rceil$ 的 int32：每个 32 位整数打包 32 个位置，字节数降到原来的 $1/8$，代价是内核里要多做位移与按位与。TensorRT-LLM 在动态树模式下默认用压缩格式，因为树的 $N$ 等于 tokens_per_gen_step，通常有几十到上百行，不压缩的掩码在每步验证时都要占去可观的显存带宽。

retrieve_next_token / retrieve_next_sibling 是树的孩子-兄弟链表表示。验证内核不需要完整的 $N \times N$ 掩码，它只沿祖先链走：从根开始，比较目标 token 与候选 token，一致就走向第一个孩子，不一致就走向下一个兄弟。链表表示让验证内核的访存从 $O(N^2)$ 降到 $O(\text{路径长度})$。retrieve_index 则记录每个节点在扁平 token 数组里的原始下标，用于把验证结果映射回草稿 token。

每请求一份树，按 slot 存放

树不是全局一份，而是每个请求一棵。DynamicTreeSlotStorage 以 py_seq_slot 为索引，把所有请求的树元数据存成 [S, ...] 的大张量，其中 $S$ 是序列槽数加一（加的那个是 CUDA Graph dummy）：

1
self.position_offsets = no_tree_position_offsets.unsqueeze(0).repeat(S, 1).contiguous()
2
self.packed_mask = no_tree_packed_mask.unsqueeze(0).repeat(S, 1, 1).contiguous()
3
self.has_tree = torch.zeros(S, dtype=torch.bool, device='cuda')

初始值不是零，而是一棵”无树”的线性链。原因写在注释里：bootstrap 请求和复用槽可能在第一次草稿之前就被验证内核读到，如果 retrieve_next_token 全是零或 -1，验证内核会走进非法分支。给一条合法的线性链作为默认值，内核即使误读也能安全退化成”只接受根 token”。

has_tree 是有效性位。mark_invalid 不只是把位清零，还会把该槽的所有元数据恢复成默认线性链：

1
self.has_tree[slot_id] = False
2
self.packed_mask[slot_id] = self._no_tree_packed_mask
3
self.position_offsets[slot_id] = self._no_tree_position_offsets
4
self.retrieve_index[slot_id] = 0
5
self.retrieve_next_token[slot_id] = self._no_tree_next_token
6
self.retrieve_next_sibling[slot_id] = -1

请求释放时调用这个函数，避免下一个请求复用槽时读到上一棵树的残骸。验证内核用 tree_valid 参数接收这个位，无效时直接走单 token 路径。

dummy 槽为什么要造一棵 K 叉树

CUDA Graph 的 dummy 请求没有真实数据，但它的形状必须与真实请求一致，否则图捕获会因形状不匹配失败。_make_kary_tree_metadata 为 dummy 槽生成一棵规则 $K$ 叉树：节点 $i$ 的父节点是 $(i-1)//K$。这棵树的掩码和位置偏移是确定性的，与任何请求状态无关，因此可以在图捕获前就填好。

同一个函数也用于生成”无树”默认链，只是传 top_k=1。两种用途共用一份代码，保证默认链与 dummy 树的掩码格式完全一致——否则验证内核要处理两种格式，分支会变多。

贪心路径的并行遍历

贪心模式下，验证的目标是：从根开始，每一步在候选集合里找与目标 argmax 一致的 token；一致就下探，不一致就停。树的形状让这个”下探”变成沿链表走：

1
accept_index, accept_token_num, accept_token = (
2
    self.tree_ops_converter.verify_dynamic_tree_greedy_out_packed(
3
        candidates, retrieve_packed, target_predict, num_gens, max_path_len,
4
        tree_valid=tree_valid))

candidates 的第 0 列是目标模型在根位置的采样结果，后面是本请求的草稿 token。target_predict 是目标模型对树上每个节点的 argmax。内核沿 retrieve_next_token / retrieve_next_sibling 走：在当前节点的孩子里找 token 等于 target_predict 的那个；找到就接受并下探，找不到就停止。

这条路径的长度最多 max_draft_len + 1，因为每步只能下探一层。即使树里有很多分支，验证只走一条路径，所以复杂度是 $O(D \cdot K)$ 而不是 $O(N)$。这也是树式投机解码相对”一次验证多个独立序列”的优势：验证代价随深度线性增长，而候选数按 $K^D$ 指数增长。

非贪心：只靠目标分布的拒绝采样

非贪心模式下，标准投机采样需要草稿分布与目标分布的概率比。但动态树模式下 TensorRT-LLM 走了一条不同的路：target-only rejection。

1
def _can_use_rejection_sampling(self, spec_metadata):
2
    return spec_metadata.use_rejection_sampling and not spec_metadata.is_all_greedy_sample

注释解释得很直接：target-only 拒绝采样只需要目标 logits，而验证阶段目标 logits 总是有的；草稿分布、top-K 分数下标都不需要保存。这省掉了草稿侧概率张量的显存和一次 gather，代价是接受判定不再利用”草稿分布与目标分布有多接近”这一信息。标准投机采样在两者接近时能接受更多 token，target-only 只用目标分布本身做接受与修正，因此接受率通常略低，换取的是显存与一次前向里少一次全词表概率运算。

实现上，目标 logits 先按每个请求的采样参数（temperature、top-k、top-p、min-p）过滤成概率：

1
target_probs = self._compute_probs_from_logits(
2
    spec_metadata.advanced_sampling_mode, target_logits_tree,
3
    request_temperatures[gen_slice].repeat_interleave(N),
4
    request_top_ks[gen_slice].repeat_interleave(N), ...)

repeat_interleave(N) 把”每请求一组参数”扩展成”每节点一组参数”，因为树上每个节点都要独立过滤。这一步在 CUDA Graph 下是固定形状的，不随树结构变化。

拒绝采样内核接收 retrieve_next_token / retrieve_next_sibling 两条链表和 tree_valid 位，输出接受下标、接受数、接受 token。与贪心路径的区别是：它在每个节点上做一次概率比较，而不是直接比 argmax；并且它会从目标分布里采样一个修正 token，作为拒绝点的替代。

接受之后：KV 从树形搬回线性

树上节点的 KV 是按树形布局写在目标 KV Cache 里的：一个请求的 $N$ 个节点占据 $N$ 个连续槽位。但序列本身是线性的，下一步的 decode 需要这些 KV 按线性顺序排列。所以验证之后要做一次KV 重定位：

1
torch.ops.tensorrt_llm.update_kv_cache_draft_token_location_2d(
2
    accepted_draft_indices, num_accepted_tokens, kv_lens_cuda, True,
3
    num_local_layers, num_kv_heads_per_layer[0], kv_head_dim_bytes,
4
    max_total_draft_tokens, max_attention_window,
5
    kv_cache_pool_pointers[pool_idx], kv_cache_block_offsets[pool_idx],
6
    max_blocks_per_seq, tokens_per_block, None)

accepted_draft_indices 记录被接受路径上每个节点在树里的原始下标。内核按这个下标序列，把对应位置的 K/V 搬到线性位置。参数里有几个硬性约束：

• num_kv_heads_per_layer 必须所有层一致，max_attention_window_vec 也必须一致。因为内核用一个统一的 stride 计算地址，如果不同层的 KV 头数或滑窗大小不同，就得按层分多次调用。源码用 assert len(set(...)) == 1 在 Python 侧提前拦截，避免内核里出现静默的地址错位。
• max_attention_window 取局部层的滑窗大小；如果为 None（全注意力），退化成 max_seq_len。滑窗模型只需要搬窗口内的 KV，超出窗口的旧位置可以直接丢弃。
• 所有 local 层必须落在同一个 KV pool 里。混合注意力模型（全注意力 + 滑窗 + Mamba）有多个 pool，local_pool_ids 集合大小必须为 1。这是动态树目前不支持异构 KV 池的原因。

为什么用 update_kv_cache_draft_token_location_2d 而不是在 Python 里切片拷贝？因为 KV Cache 是分页的，节点的物理位置由块表和块内偏移决定，Python 侧无法高效表达”按树下标序列搬动分页 KV”。自定义内核把块表遍历和块内拷贝合成一次，并且支持 CUDA Graph 捕获（所有指针和形状在捕获时固定）。

CUDA Graph 里的掩码重排

验证内核在 Hopper 上要求掩码是紧凑的一维位压缩布局，而动态树缓冲是按”最大宽度”分配的二维布局。当实际 $N$ 小于缓冲宽度时，行 stride 不一致，内核会读错行。_repack_mask_padded_to_packed 解决这个问题：

1
mask_width = math.ceil(n_tok / 32)
2
total_elems = n_req * n_tok * mask_width
3
scratch = self._mask_repack_buf[:total_elems].view(n_req, n_tok, mask_width)
4
scratch.copy_(mask_buf[:n_req, :n_tok, :mask_width])
5
flat[:total_elems] = scratch.view(-1)

先在 scratch 里拷出有效子块，再把 scratch 展平写回原缓冲的前缀。这样做而不是直接 .contiguous()，是因为后者会分配新张量，地址在 CUDA Graph 捕获前后不一致。self._needs_mask_repack = not is_sm_100f(sm) 说明 Blackwell（SM100 系列）的 prepareCustomMask 直接接受 padded 3D 布局，不需要这一步，所以只在 Hopper 及更早的架构上执行。

草稿 KV 的两种归属

Eagle3OneModelWorker.__init__ 接收 use_separate_draft_kv_cache，它的默认值是 False，但最终是否真的单独开池并不由这个参数单独决定。should_use_separate_draft_kv_cache() 先看是否启用了 _use_shared_kv_cache：这个私有属性默认 False，一旦置为真，草稿只计算 Q、直接复用目标模型的 K/V，对应 _forward_external_shared_target_kv_draft_loop。否则再看 _allow_separate_draft_kv_cache，它默认 True，因此 EAGLE 一体式路径默认让草稿层拥有自己的分页 KV 池。少数特殊后端会在这个判断里被排除：VANILLA/FA4 的 DFlash 草稿、内嵌在目标模型里的 DSpark 草稿，以及 suffix automaton，都不读受管的分页草稿池，也就没有必要为它们再开一套。

两种归属的取舍不同。共享目标 KV 省掉一套池的显存和一次”草稿池 → 目标池”的拷贝，但要求草稿的特征空间与目标模型兼容，且回滚时必须精确扣掉被拒绝的草稿位，不能污染块表。单独开池则允许草稿侧有自己的页大小与量化格式，代价是接受之后要把草稿 KV 搬进目标池，并且调度器要为第二套池单独做准入计算。源码注释里提到，只有 TRTLLM attention 会借用受管的分页上下文池；VANILLA 与 FA4 把 context KV 放在 worker 自己的缓冲里，再开一套池只会预留用不到的页，反而多出一道无关的准入限制。

动态树的三条硬限制

从源码的断言可以直接读出动态树当前不支持什么：

• 不支持 relaxed acceptance 和 suffix automaton 增强。__init__ 里直接 raise ValueError。原因是这两种增强都会改写”接受哪些 token”的判定，而动态树的验证内核把判定固化在 CUDA 里，Python 侧的改写没有插入点。
• 不支持非均匀 KV 头数或滑窗的混合模型。KV 重定位内核要求所有 local 层同池、同头数、同窗口。TensorRT-LLM 官方文档对此有一句更概括的表述：动态树模式目前不支持使用滑窗注意力或 MLA（Multi-Latent Attention）的模型，例如 DeepSeek 与 gpt-oss 系列。这与源码里的三个断言一一对应——滑窗带来 max_attention_window_vec 不为 None，MLA 的 latent 池布局与普通 KV 池不同，二者都会让”一次重定位覆盖所有 local 层”的前提失效。
• 需要 max_batch_size 在构造时已知。所有缓冲按 max_batch_size 预分配，assert batch_size <= self._max_batch_size 在每次前向时检查。动态树不能像线性草稿那样按需增长，因为 CUDA Graph 要求形状固定。

这三条不是实现缺陷，而是”把树结构固化进 GPU 内核”的必然代价：一旦验证逻辑在内核里，任何改变判定语义的 Python 侧扩展都无处安放。

首步变长、后续定长：一个被低估的复杂度来源

线性草稿和动态树都遵循”首步变长、后续定长”的模式，但两者的实现完全不同。线性模式靠 gather_ids 在扁平张量里寻址，后续步退化成 slice(0, batch_size)；动态树模式靠 _gather_repack_step0_kernel 这个 Triton 内核，把树形拓扑重排成”每请求 max_path_len 行”的定长布局，后续层再按 $K$ 增长。

为什么动态树不能复用线性模式的 gather？因为首步之后，树模式每层的输入行数随层数线性增长：第 $i$ 层要重新处理前 $i$ 层累积的 $i \cdot K$ 个节点，从中扩展出 $K^2$ 个候选、再裁回 $K$ 个。线性模式每步都是 1 行，天然定长；树模式每层的输入形状都不同，必须为每层单独准备输入，并且每层的 attention 掩码都要重建。这就是 _build_mask_and_position 和 _gather_parent_hs 存在的原因：它们把”本层节点的祖先掩码”和”本层节点的父隐状态”融合成两个 torch.compile 函数，避免在 Python 循环里逐节点处理。

_build_mask_and_position 的拼接逻辑值得单独看：

1
parent_mask_selected = torch.gather(parent_mask[:, -K:, :], 1,
2
    selected_parents.unsqueeze(-1).expand(batch_size, K, num_tokens_previous_layer))
3
padding = parent_mask.new_zeros(batch_size, num_tokens_previous_layer, num_tokens_current_layer)
4
bottom = torch.cat([parent_mask_selected, tree_mask_init], dim=2)
5
current_mask = torch.cat([padding, bottom], dim=1)

新掩码是分块矩阵：左上角是全零（旧节点看不到新节点），右下角是 tree_mask_init（本层 $K$ 个节点两两之间的可见性），右上角是从父掩码里 gather 出来的祖先关系（新节点能看到父节点的祖先），左下角是旧节点之间的原掩码。这种分块构造让掩码随层数增长时只做一次 cat，而不是每层重算整张 $N \times N$ 矩阵。