DeepInfer 源码拆解（三十一）：模型装配与前向主循环——从 HF 参数壳到一次前向，再到采样器

第一刀：num_total_layers = len(module) - 2

1
num_total_layers = len(model.module.module) - 2  # 除去 emb 和 lm_head
2
self.shared_topology = build_shared_layer_topology(cfg, num_total_layers)

为什么减 2？因为扁平序列的第一项是词嵌入、最后一项是 LM head 相关参数，中间的才是 decoder 层。注意这里统计的是全部层，包括 MTP（NextN）层——共享拓扑需要用完整的层数来做区间校验，第二十五篇讲 CP、第十一篇讲跨层共享时都用到这套拓扑。这里体现了一个设计原则：拓扑是模型级对象，构造一次，显式传给每一层，而不是让每层自己去查配置。

第二刀：主干与预测头分离

1
head_module = model.module.module[cfg.num_hidden_layers + 1]
2
assert isinstance(head_module, MainPredictHeadHC), f"{type(head_module)=}"
3

4
self.model = DeepseekModel(
5
    model.module.module[: cfg.num_hidden_layers + 1],
6
    head_module,
7
    cfg,
8
    self.load_nextn,
9
    self.shared_topology,
10
)

HF 的模块序列里，cfg.num_hidden_layers 之后的下一项就是预测头（MainPredictHeadHC），它同时包含了输出归一化和真正用于 LM head 的权重。这里把它单独摘出来传给 DeepseekModel：DeepseekModel 只取其中的 output_norm 和（可选的）HCHead，而 shared_head 里那份词表矩阵留给下一刀。

第三刀：LM head 与采样器

1
head_module = head_module.shared_head
2
assert isinstance(head_module, GPTHead)
3
self.lm_head = ParallelLMHead(
4
    cfg.vocab_size,
5
    self.hidden_size,
6
    weight=head_module.head.weight,
7
)
8
self.sampler = Sampler(cfg.vocab_size)

ParallelLMHead 继承自 VocabParallelEmbedding，整个文件只有 30 行：

1
class VocabParallelEmbedding(torch.nn.Module):
2
    def __init__(self, num_embeddings, embedding_dim, weight):
3
        super().__init__()
4
        assert weight.shape == (num_embeddings, embedding_dim)
5
        self.weight = reconstruct_param(weight)
6

7
    def forward(self, input_: torch.Tensor) -> torch.Tensor:
8
        output = gather(self.weight, input_)
9
        return output
10

11

12
class ParallelLMHead(VocabParallelEmbedding):
13
    def forward(self, input_: torch.Tensor) -> None:
14
        del input_
15
        raise RuntimeError("LMHead's weights should be used in the sampler.")

ParallelLMHead.forward 故意直接抛异常。这是一个很强的架构约束：词表投影计算不是模型的最后一层，而是采样器的第一步。原因在第八篇已经分析过：投机解码的验证需要在多行 logits 上做，而词表矩阵是一个 12 万级的大矩阵，如果用普通的 nn.Linear 挂在模型末尾，每次前向都会多一次模块调用、多一层 Python 开销，CUDA Graph 捕获时也更容易被切段。DeepSeek 的处理是把“算 logits”这件事连同“挑 token”放进同一个采样内核管线（get_logits → fused_sample），模型层只输出 hidden state。

VocabParallelEmbedding 在 DeepInfer 里其实并不真正“parallel”——它没有 vLLM 那种词表分片 + all_reduce 的实现，而是用 reconstruct_param 把参数重新包一层，然后 forward 直接走 TileLang 的 gather 内核。reconstruct_param 的理由很简单：

1
# 重新构造 nn.Parameter，隔离模型间的参数生命周期，避免释放一个模型时清空另一个模型的参数。
2
def reconstruct_param(param: torch.nn.Parameter) -> torch.nn.Parameter:
3
    p = torch.nn.Parameter(param.data, requires_grad=False)
4
    if hasattr(param, "sf"):
5
        p.sf = param.sf
6
    return p

DeepInfer 在装配阶段会先物化扁平参数壳，再构造运行时模型，然后把扁平壳 del gpt。如果运行时模型直接持有扁平壳里的 Parameter 对象，那么扁平壳释放时，Parameter 的 __del__ 或引用关系可能把新模型的数据一起带走。重新包一层 Parameter，共享同一块 data，但生命周期独立——这是一个只有在这种“两次构造”的装载流程里才需要的小心处理。

第四刀：MTP 模块

1
if self.load_nextn:
2
    self.MTP_module = DFlash(model, cfg, self.model.embed_tokens, self.lm_head)
3
    self.num_mtp_moe_layers = len(self.MTP_module.moe_layer_list())
4
else:
5
    self.MTP_module = None
6
    self.num_mtp_moe_layers = 0

注释写得很清楚：MTP 层的层数和预测个数解耦。load_nextn（来自 envs.DEEPINFER_USE_NEXTN_DECODE）先决定“要不要构造 DFlash、加载 NextN 权重”；至于一次解码实际验证几个草稿 token，由投机解码路径的投机深度控制决定（取值上限默认 5，第五篇讲了它的调度逻辑）。所以 num_mtp_moe_layers 是一个静态装配结果：它只反映“MTP 有几层 MoE”，和每一步跑几次草稿无关。这个数字的唯一消费者是计时器——decode 阶段 MLP 计时按它缩放（见后文），目的是让“单层 MLP 时间 × 层数”对齐一次完整 decode step 的 MLP 总量。

DFlash 拿到了 embed_tokens 与 lm_head 的引用，这两份大权重在 MTP 路径里不再复制——第十四篇讲权重装载时的 tied parameters 覆盖在这里兑现了。MTP 内部细节见第五篇与第二十六篇。

DeepSeek-V3 技术报告里的 MTP 结构图可以说明“为什么 MTP 层不放进主循环”：每个 MTP 模块都在主模型之上再挂一个 Transformer Block 与输出头，各模块的输出头共享同一个词表嵌入，训练时对多个未来 token 分别算交叉熵：

推理时并不需要把这几层“训练用的深监督”全部串在主模型后面——主模型只需要把指定目标层的 hidden 收集出来（dflash_input），草稿头在采样器内部单独跑。图中“输出头共享嵌入”这一点，对应到装配阶段就是 DFlash 直接拿 lm_head.weight 的引用，而不是复制一份。

一个值得警惕的歧义：load_nextn 是整数，不是布尔量

self.load_nextn = envs.DEEPINFER_USE_NEXTN_DECODE 的取值不是 True/False，而是一个整数上限；DeepseekModel 里把它原样存下，并在构造 DFlash 目标层映射时这样用：

1
dflash_target_layer_ids = tuple(dflash_cfg.dspark_target_layer_ids) if load_nextn and dflash_cfg is not None else ()

if load_nextn 是 Python 的真值判断——只要环境变量非零就走构造分支。DFlash 内部并不直接接收这个参数，而是重新读一次环境变量并存成 max_MTP_value，源码注释写明它的含义是“每个 seq 生成的 propose 个数”（即每一步最多生成几个草稿 token）。这种“一个变量同时承担开关与上限”的写法在同类引擎里很常见（vLLM 的 num_speculative_tokens 也兼有这两种语义），但它对阅读者的要求是：看到 if self.load_nextn 与 self.load_nextn > 0 时要意识到它们不是同一个意思。前者等价于“启用 NextN”，后者用于数值比较。装配阶段的所有分支判断用的都是前者。

还有一个隐含约束：dflash_target_layer_ids 只在 load_nextn 为真且 dflash_cfg 存在时才非空；它与 DeepseekModel.dflash_input_size 的关系是 size = hidden_size × len(ids)。主循环开头那条 assert (dflash is not None) == bool(self.dflash_input_size) 就是保证这两个派生量同进同退。

三块构造期元数据的派生

装配过程中有三组数据不是从 HF 权重里“读出来”的，而是从 cfg 和权重形状推导出来的，且只在构造期算一次：

shared_topology 的层数是 len(model.module.module) - 2，包含 NextN 层：共享拓扑的区间校验要覆盖所有层，否则 MTP 层里引用的 source 可能在主模型区间之外。dflash_input_size 的构造则依赖一个断言：

1
assert (dflash is not None) == bool(self.dflash_input_size), "dflash_cfg.dspark_target_layer_ids 与 MTP head 类型不匹配"

这个断言出现在 forward 里而不是 __init__ 里，原因是 DeepseekModel.forward 的 dflash 参数是调用方传进来的——正常执行路径总把 MTP_module 传进去，但测试与离线工具路径可能不传。把断言放在前向入口，是为了保证任何一条路径上，dflash 与 dflash_input_size 同时存在或同时缺失；否则主循环会尝试写入一块不存在的缓冲，或者申请了缓冲却没人写。

Engram 的派生更值得单独说：NgramHashMapping 产出的是哈希乘子、词表大小、偏移量三组常量，EngramRuntime 则按 engram_layer_ids 收集每层的 Engram 模块。两者在构造期分开准备，主循环里只做一次哈希、一次 fetch，然后把结果按层号分发给各层（第十八篇拆过完整链路）。

装配后的第一个动作：分配 userbuffer

1
torch.cuda.empty_cache()
2
self.allocate_userbuffer()

allocate_userbuffer 遍历所有 decoder 层（含 MTP），做三件事：

1. 遍历主模型 decoder 层与（若启用 NextN）MTP 的 MoE 层，对每个 DeepseekMoE 调 mlp.allocate_userbuffer()，给专家并行分配 dispatch/combine 缓冲；
2. 对每个 V4AttentionBase 取 get_max_ag_bytes_per_token()，在初始预算上取最大值，得到 CP all-gather 的每 token 字节；
3. 调 userbuffer.allocate_cp(PG.context_parallel_group(), max_num_tokens=self.cfg.max_num_prefilling_tokens, max_ag_bytes_per_token=...)。
4. 若模型有 Engram，最后再调 engram_runtime.initialize() 分配跨 rank 查表缓冲（第十八篇）。

每 token 的初始预算是 hidden_size * 2（bf16 的 hidden state）；如果 DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact" 且模型存在 YoCo 边界层，则改成 max(hidden_size * 2, hc_mult * hidden_size * 2)——第二十八篇讲过原因：compact 模式下首层要额外传 HC 残差走 CP gather。

allocate_cp 的算术也值得复述一句，因为它把“缓存大小”变成了一个可核对公式：

其中 $t_{\max}$ 是 max_num_prefilling_tokens，$b_{\text{tok}}$ 是每 token 的 all-gather 字节，$S$ 是 CP 组大小，$n_{\text{ag}}=8$ 是一次会话里最多并发的 all-gather 次数（max_num_ags_per_session 的默认值），$A_{\text{rdma}}$ 是 DeepEP 的 RDMA 对齐粒度，$A_{\text{alloc}}$ 是整体分配对齐粒度。公式的含义是：先按“最大预填充 token 数 × 每 token 字节 × 组大小”算出一个会话的通信量，再为最多 8 次并发会话各留一条 RDMA 对齐的余量，最后向上取整到分配对齐。CP 组大小为 1 时直接返回 None，不占任何通信缓冲——这也解释了为什么单机部署的 DeepInfer 显存账单里几乎没有这部分。

嵌入与跨层容器

1
if hidden_states is None:
2
    hidden_states = self.embed_tokens(input_metadata.input_ids)
3

4
residual = None
5
cross_layer_runtime = CrossLayerRuntime()

hidden_states 参数是给“调用方已经准备好嵌入”的路径预留的：传了非 None 就跳过词表 gather，直接把它当第一层输入。主执行路径（ModelRunner.execute_main_model）传的是 None，所以正常 decode 的这一步就是一次 gather；预填充与 decode 共用这段代码，差别只在 input_ids 的行数与元数据。

CrossLayerRuntime 是一个“空壳”——所有字段默认 None——但它是本循环里唯一在层与层之间共享可变状态的对象。它的字段在 cross_layer.py 里有完整注释，共六类：

• shared_hidden：第一个 YoCo consumer 层的 attn-norm bf16 hidden；
• pre_rope_kv：最近的 kv source 层在 rope 之前的压缩 KV；
• topk_indices：最近的 index source 层的 topk；
• block_candidate_idxs：候选块索引；
• loaded_prefill_kv / loaded_prefill_index_kv：预填充阶段已装载的主 KV / 索引 KV。

为什么要一个可变、逐层传引用、写回不通过返回值的容器？注释给出理由：层的共享拓扑是分段校验过的、执行顺序严格递增，所以一个 consumer 需要的值一定是它前面最近的 source 层写入的；不需要记录来源层再去匹配。如果把这些值塞进每层的返回值，那么层与层之间就要传一个越来越大的 tuple，而 CUDA Graph 捕获要求调用栈形状稳定——每层返回 tuple 会让 Python 层的形状依赖层号，破坏一致性。用一个原地写入的容器，主循环的形状对所有层完全相同。

CP 镜像 gather

1
if input_metadata.is_prompt and PG.context_parallel_group().size() > 1:
2
    input_metadata.prepare_for_mirror_cp()
3
    hidden_states = gather(hidden_states, input_metadata.mirror_cp_index)

只在预填充阶段、且 CP 组大于 1 时生效。第二十五篇详细讲过 mirror 布局：每个 rank 只持有序列的一部分，但每一层的注意力需要看到全序列——做法是把所有 rank 的 hidden 先 gather 成完整序列，再按 zigzag 镜像索引重排，让每个 rank 拿到的行既包含本地部分，也包含对端需要的位置。注意索引本身不是模型层算的：prepare_for_mirror_cp() 在模型前向开始时一次性算出 mirror_cp_index 与 mirror_cp_reverse_index（源码注释写明“Call once at the start of the model”），并把 positions、ngram_token_ids、images_seq_mask 等逐行相关的字段也按同一索引重排；模型层随后只是拿这个索引执行 gather。源码注释还说明了为什么能这么做：预填充不走 CUDA Graph，所以这些字段换成新地址是安全的；decode 路径不进入这段代码。

DFlash 特征收集

1
dflash_input = hidden_states.new_empty((hidden_states.size(0), self.dflash_input_size)) if dflash is not None else None
2
dflash_metadata = input_metadata

dflash_input_size = hidden_size × len(dspark_target_layer_ids)。它是一个“拼接”缓冲：DFlash 需要若干目标层的 hidden 作为草稿输入，这些层散落在主模型的不同深度，所以先分配一块连续缓冲，让每个目标层在 forward 时把自己的 hidden 写到对应切片里。写入发生在解码层的 _store_dflash_input：

1
def _store_dflash_input(self, hidden_states, dflash_input):
2
    n_embed = hidden_states.shape[-1]
3
    index = self.dflash_input_index
4
    dflash_input_store(hidden_states, dflash_input[:, index * n_embed : (index + 1) * n_embed])

self.dflash_input_index 在构造层时就由 dflash_target_layer_ids 的位置查好了，非目标层为 None，写入被完全跳过。DFlash 的具体用法见第五篇：投机解码。

Engram 哈希

1
if self.engram_runtime is not None:
2
    hashes = tile_kernels.engram.engram_hash(
3
        input_metadata.ngram_token_ids,
4
        self.ngram_multipliers, self.ngram_vocab_sizes,
5
        self.ngram_offsets,
6
        getattr(input_metadata, "images_seq_mask", None),
7
    )
8
    ngram_embed_hooks = self.engram_runtime.fetch(hashes)

Engram 是“在若干层插入 N-gram 查表记忆”的机制，第十八篇专门拆过。这里只强调一个结构上的选择：哈希计算和查表在循环之前一次性完成，产出一个 {layer_idx: hook} 字典；循环里的每层只是 ngram_embed_hooks.get(i) 拿出自己的 hook。这会带来一个显存代价：所有层的查表结果在同一时刻都要活着。但换来的是循环体里没有条件判断——hook 要么存在要么是 None，传给层的都是同一个参数。

images_seq_mask 是一个 [total_q] 的布尔向量，True 标记图像 token。它在 Engram 里有两条用途：一是哈希时把图像 token 排除在 N-gram 统计之外，避免视觉内容的 token 组合污染语言侧的 n-gram 分布；二是同一个掩码会顺带传给层内的 MoE 门控——DeepInfer 为多模态输入准备了一份单独的 e_score_correction_bias_vl，当 images_seq_mask 存在时用它替代普通偏置（mlp.py 里的 image_bias=... 分支），让视觉 token 与文本 token 走不同的路由偏置。也就是说，同一份掩码在模型层被两个下游模块共用，主循环只需要把它放在元数据里一次。

YoCo 边界：一次 metadata 视图切换

1
for i in range(self.num_layers):
2
    if input_metadata.is_prompt and i == self.yoco_boundary_layer_idx:
3
        if envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact":
4
            consumer_metadata = input_metadata.get_consumer_view()
5
            cross_layer_runtime.consumer_full_metadata = input_metadata
6
            input_metadata = consumer_metadata
7
            if self.dflash_targets_in_consumer:
8
                rows_c = consumer_metadata.cp_mirror_len if PG.context_parallel_group().size() > 1 else consumer_metadata.total_q
9
                dflash_input = hidden_states.new_empty((rows_c, self.dflash_input_size))
10
                dflash_metadata = consumer_metadata
11
        else:
12
            input_metadata = input_metadata.get_mask_view()

这是主循环里唯一在层与层之间改变输入元数据的地方，而且只发生在预填充。YoCo 有界重放是“只重算最后一段 token”：n_prefill_layer 之前的 encoder 半区需要完整 attention 语义，边界之后属于 decoder（consumer）半区，它们的注意力只作用于被保留的尾部窗口，元数据里的序列长度、mask、KV 索引都不同。第三篇拆过截断与等价性验证，这里只关心模型层怎么切换元数据视图。

两种截断模式的区别在注释里写得很直接：

• compact：物理上把行数切小，input_metadata 换成 consumer 视图；hidden_states 的行仍是原始的，在第一个 consumer 层内部才切。full 视图与完整 hidden 存进 cross_layer_runtime，供“需要算到底的”路径使用。
• mask：不做行数切换，只换 mask 视图，计算量还是全量，只是不再算不需要的注意力位置。

DFlash 的额外处理来自一个硬约束：compaction 下 DFlash 目标层要么全在 consumer 半区，要么全不在。构造时就断言了这一点：

1
assert all(in_consumer) or not any(in_consumer), \
2
    f"yoco consumer compaction requires DFlash target layers on the same side of n_prefill_layer"

如果全在 consumer 半区，dflash_input 的缓冲也必须同步切成 compact 行数，否则 dflash_input_store 的行数与 consumer 层 hidden 的行数对不上。这是一个典型的“两个子系统共享同一个行布局假设”的耦合点。

逐层调用

1
hidden_states = Movable(hidden_states)
2
hidden_states, residual = self.layers[i](
3
    hidden_states,
4
    kv_caches[i],
5
    input_metadata,
6
    residual,
7
    embed_hook=ngram_embed_hooks.get(i),
8
    cross_layer_runtime=cross_layer_runtime,
9
    dflash_input=dflash_input,
10
)

每层的参数集是固定的六项，与层号无关。kv_caches[i] 是从 GPUCacheInterface 里取出的第 i 层缓存视图；因为层间可能共享 KV（share-kv），多个 i 可能指向同一块 alias 内存，第三十篇讲过alias 是怎么在 cache engine 里建立的。

Movable 的包装用一句话解释：模拟 C++ std::move 的一次性取出语义。它的约束与第十二篇的 CUDA Graph 捕获直接相关。

1
class Movable[T]:
2
    """A move-only wrapper that simulates C++ std::move semantics."""
3
    def move(self) -> T:
4
        if self._value is None:
5
            raise RuntimeError("Attempt to take from an already moved Movable instance")
6
        value = self._value
7
        self._value = None
8
        return value

Python 没有所有权语义，一个张量被多个局部变量引用时，只有最后一个引用消失才会真正释放。DeepInfer 在每个可能“交出所有权”的边界前包一层 Movable：调用方交出后不能再读，接收方取走后可以尽早 del，从而让显存释放时机确定下来。启动期的显存画像（第二十篇的 profile_run）测出的峰值，正是依赖这种精确的释放时机：如果某个张量因为被隐式共享引用而晚释放，峰值就会被算高，进而折算出的 KV 块数偏少。

尾部处理：DFlash、mHC 折叠、输出归一化

循环结束后：

1
if dflash is not None:
2
    dflash.update_main_cache(dflash_input, dflash_metadata, kv_caches)
3

4
if self.hc_enabled:
5
    if hidden_states.dim() == 2:
6
        hidden_states = HCLayerMixin.hc_post(hidden_states, residual)
7
    assert hidden_states.ndim == 3
8
    if self.hc_shifted:
9
        hidden_states = mhc_pre_apply_mix(hidden_states, residual.pre_layer_mix)
10
    else:
11
        hidden_states = hidden_states.flatten(1, 2)
12
else:
13
    hidden_states.add_(residual)
14

15
if self.hc_enabled and not self.hc_shifted:
16
    hidden_states = hidden_states.unflatten(-1, (self.hc_head.hc_mult, -1))
17
    hidden_states = self.hc_head(hidden_states)
18
hidden_states = self.output_norm(hidden_states)

这一段要分清楚两种残差主干：

• 非 HC 模型：最后一层返回的 residual 还没加进 hidden（融合残差归一化的设计，见第十五篇），所以这里 add_ 一次，补上最后一次残差；
• HC 模型：残差是一条 [n, hc_mult, hidden] 的多路 trunk，最后一个 post 还没做，所以先 hc_post；如果层数是奇数导致 residual 是“懒数据”，hc_post 会先把它实体化。之后按 hc_shifted 走两条不同的折叠路径。

flatten(1, 2) / unflatten(-1, ...) 这一对 reshape 看起来像“压平再还原”的无用往返，实际是给 HCHead 做布局归一化。HCHead.forward 的契约是输入 [n, hc_mult, hidden_size]、输出 [n, hidden_size]，内部调用融合算子 mhc_head：

1
def forward(self, residual: torch.Tensor) -> torch.Tensor:  # [n, hc_mult, hidden_size] -> [n, hidden_size]
2
    out = mhc_head(
3
        residual.unsqueeze(0),  # the kernel requires a batch dim
4
        self.hc_head_fn, self.hc_head_scale, self.hc_head,
5
        norm_weight=None,  # HF already folds the norm weight into hc_head_fn
6
        norm_eps=self.hc_head_norm_eps,
7
        mhc_mult=self.hc_mult, pre_eps=self.hc_pre_eps,
8
    )
9
    return out.squeeze(0)

flatten 先把 trunk 变成 [n, hc_mult*hidden_size] 的二维布局（必要时触发一次连续化），unflatten 再按 (hc_mult, -1) 解释回三维。输入本来就连续时两步都是零拷贝视图；输入带有奇怪 stride 时，flatten 会先把它变成连续布局，保证 mhc_head 拿到的内存排布符合预期。norm_weight=None 那一行说明 HF 已经把折叠用的 norm 权重折进了 hc_head_fn，运行时不再重复传参。

CP 反向 gather与输出

1
if input_metadata.is_prompt and PG.context_parallel_group().size() > 1:
2
    with userbuffer.cp_session():
3
        gathered_hidden_states = userbuffer.cp.all_gather(hidden_states)()
4
        hidden_states = gather(gathered_hidden_states, input_metadata.mirror_cp_reverse_index)
5

6
return CausalLMForwardOutput(hidden_states=hidden_states)

进循环时做过一次 mirror gather，出循环时就要做对应的反向操作：all-gather 拿到所有 rank 的输出，再用 mirror_cp_reverse_index 把行重排回原来的顺序。mirror_cp_reverse_index 与进循环时的索引一样，在前向开始时由 prepare_for_mirror_cp() 一并算出，模型层只按行拷贝。cp_session 是第二十八篇讲的会话式缓冲管理——一次会话里的 all-gather 共用同一组 userbuffer 通信缓冲，退出会话才允许下一次通信复用。CP 为 1 时 cp_session 直接 yield，不产生任何通信。

Engram 作为“主干干预钩子”挂在 post 之后

1
if self.engram is not None:
2
    hc_post_intervention = lambda trunk: self.engram(
3
        trunk, embed_hook, getattr(input_metadata, "images_seq_mask", None)
4
    )
5
else:
6
    hc_post_intervention = None

注意这是一个函数而不是数值：HCLayerMixin.attn_norm_forward 会把当前主干的值交给它，返回值就是干预后的主干。为什么挂在 attn pre/post 融合里而不是直接在 hidden 上加？因为 mHC 的 post 与下一层的 pre 是融合计算的（_hc_fused_post_pre，第十篇），在这两个混合之间插入干预，才能保证下一层看到的是已经加过 Engram 记忆的多路 trunk。

两个计时器：分组按同构、预填充与解码各取不同的采样层

1
second_layer, second_last_layer, same_layers = get_enable_timer_layer(self.layer_idx)
2
attn_timer_scale = len(same_layers)
3
with (
4
    cuda_timer.record(f"attn-layer-{self.layer_idx}", enabled=(self.layer_idx == second_last_layer and input_metadata.is_prompt), scale=attn_timer_scale),
5
    cuda_graph_timer.record(f"attn-part-{self.layer_idx}", enabled=(self.layer_idx == second_layer and not input_metadata.is_prompt), scale=attn_timer_scale),
6
):

计时器的采样策略有三个关键点：预填充取下数第二层、解码取第二层；同构层按组缩放；CUDA Graph 里的计时用另一套探针。第二十九篇详细讲过图内计时器，这里的补充是“按组缩放”的来源：

1
def layer_group_key(self_attn: nn.Module, mlp: nn.Module | None) -> str:
2
    """key of the timer sampling group: module structure repr + attention config"""
3
    sched_meta_key = getattr(self_attn, "sched_meta_key", "")
4
    is_yoco_consumer = getattr(self_attn, "is_yoco_hidden_consumer_layer", False)
5
    return f"{self_attn}|{mlp}|{sched_meta_key}|yoco{int(is_yoco_consumer)}"
6

7

8
def get_enable_timer_layer(layer_idx: int) -> tuple[int, int, list[int]]:
9
    same_layers = DeepseekDecoderLayerType[DeepseekDecoderLayerIdx[layer_idx]]
10
    second_layer = same_layers[1] if len(same_layers) > 1 else same_layers[0]
11
    second_last_layer = same_layers[-2] if len(same_layers) > 1 else same_layers[0]
12
    return second_layer, second_last_layer, same_layers

layer_group_key 的四个分量分别锁定“注意力模块结构、MLP 结构、调度元数据形状、是否 YoCo consumer”。sched_meta_key 在 V4Attention 构造时由 kv_compress_ratio / has_packed_kv / use_swa / index_topk 拼成；结构相同的层被归入同一组，计时器只在组内的“第二层”或“倒数第二层”记一次，再用 scale=len(same_layers) 把这一次的时间放大到整组。分组本身在每层构造时用 create_deepseek_decode_layer 注册，因此组的划分是静态的，不随 batch 改变。

为什么每组选第二层和倒数第二层？V4.1 的层是异构的：不同层可能挂稀疏索引器、可能用 SWA、可能是共享 KV 的 consumer、也可能是 YoCo 的 consumer。第一层通常是“边界层”——它没有前一层可复用 KV、可能承担 expand_to_mhc 的主干初始化；最后一层则可能赶上 YoCo 边界切换或 DFlash 目标层写入。取组内第二层与倒数第二层，是避开首尾边界效应、又保留组内代表性的折中。组只有一层时两个采样点退化为同一层（same_layers[0]），这是为单层组留的后路。

YoCo 边界处 consumer 的 full hidden

1
if cross_layer_runtime.consumer_full_metadata is not None and cross_layer_runtime.consumer_full_hidden is None:
2
    compact_inputs, full_inputs = compact_yoco_boundary(hidden_states, residual, input_metadata, cross_layer_runtime, hc_enabled=self.hc_enabled)
3
    cross_layer_runtime.consumer_full_hidden = self.attn_norm_forward(*full_inputs)[1]
4
    hidden_states, residual = compact_inputs

这是 compact 模式的实现核心：同一个子层要对两种行布局各算一次 attn-norm。第一次在 full 布局上，产物 [1] 即 bf16 normed hidden，存进运行时容器（供 KV 写入使用）；第二次在 compact 布局上，产物作为本层实际要算的输入。为什么不能只用 compact 的 hidden 吗？因为 YoCo 截断只是计算上的——必须写入 KV 的那部分内容仍然是完整语义的，full 布局的 hidden 就是 KV 更新的来源。注释也说明了这一点：在截断语义下“原始路径不更新 kv”，而是“在专门路径上更新”。

注意力与 FFN

1
hidden_states, hidden_states_bf16 = Movable(hidden_states), Movable(hidden_states_bf16)
2
hidden_states = self.self_attn(
3
    hidden_states=hidden_states,
4
    kv_cache=kv_cache,
5
    input_metadata=input_metadata,
6
    hidden_states_bf16=hidden_states_bf16,
7
    cross_layer_runtime=cross_layer_runtime,
8
)
9
hidden_states, hidden_states_bf16, residual = self.ffn_norm_forward(hidden_states, residual)

注意 attn_norm_forward 返回三个值：hidden（可能是量化过的）、hidden_bf16（一份保留 bf16 副本）、residual。为什么同时要两份？第十五篇讲线性层与归一化时提到过，FP8 注意力的输入可以是 FP8，但 MoE 的 gate、以及某些跨层消费（shared_hidden 捕获）需要 bf16。所以归一化一次性算出两份，注意力按需要取用，FFN 拿 bf16——这也正好解释了前面 shared_hidden 捕获时的 hidden_states_bf16.clone()：

1
if self.yoco_captures_shared_hidden and cross_layer_runtime.shared_hidden is None:
2
    # 第一个 consumer 层：attn-norm bf16 hidden 即 shared hidden
3
    cross_layer_runtime.shared_hidden = hidden_states_bf16.clone()

clone 是因为 Movable 交给注意力后，注意力内部可能会原地改写这块 buffer（注释 TODO(YOCO): 去掉这个 clone 承认了这是一笔待偿还的拷贝开销）。这是显存与正确性之间的一次权衡：要省这块拷贝，注意力就不能原地改写 bf16 输入，需要改一整族 kernel 的实现约束。

然后是 FFN，计时器按“实际执行的 MLP 层数”缩放：

1
mlp_timer_scale = self.num_prefill_mlp_layers if input_metadata.is_prompt else self.num_mlp_layers

num_mlp_layers 包含 NextN 层、num_prefill_mlp_layers 不含：预填充阶段不跑 MTP 的 MLP，而 decode 阶段（启用 DFlash 时）每个 step 会顺带执行草稿层的 MLP。计时器记录的是 layer_idx == 1 那一层的 MLP 耗时，再乘上这个层数；两组层数分别对应预填充与 decode 的实际结构，避免把 decode 的草稿层开销摊进预填充的数字里。

DFlashDecoderLayer 是这套结构的变体（第 178 行起）：它多了 update_main_cache 与 forward_drafts，在解码时由采样器在单独路径上调用，负责写 DFlash 的 main cache。第五篇从投机解码的视角讲过它，这里只指出它在模型层的边界：它复用了 HCLayerMixin 的残差语义，但不在 DeepseekModel.forward 的主循环里。

装配时预留一块最大输出

DeepseekModel 构造时并不急着分配输出；输出缓冲由 allocate_output(num_tokens) 在需要时分配：

1
def allocate_output(self, num_tokens: int) -> CausalLMForwardOutput:
2
    hidden_states = self.embed_tokens.weight.new_zeros((num_tokens, self.hidden_size))
3
    return CausalLMForwardOutput(hidden_states=hidden_states)

embed_tokens.weight.new_zeros 而不是 torch.empty 有两个用意：一是复用词表参数的设备与 dtype，避免硬编码；二是让 padding 行有一份确定的初值。图捕获时只捕获前 batch_size 行的内核，重放时超过真实 batch 的行不会被有意义地计算，零值让这些行在调试和数值检查里不会带出未初始化内存的随机垃圾。

CUDA Graph 捕获怎么用这块缓冲

第十二篇拆过捕获流程，这里只把与模型层直接相关的几行摘出来：

1
output_full = model.allocate_output(self.max_batch_size)
2
...
3
for batch_size in reversed(self.batch_sizes):
4
    def _cut(x, _bs=batch_size):
5
        return x[:_bs]
6
    input_metadata_source = input_metadata_source_full.recursive_fn(_cut)
7
    input_metadata_device = input_metadata_device_full.recursive_fn(_cut)
8
    output = output_full.narrow(batch_size)
9
    ...
10
    with torch.cuda.graph(graph, pool=self.graph_memory_pool, stream=self.graph_stream, ...):
11
        model_output = _run_model(model, kv_caches, input_metadata_source, input_metadata_device, ...)
12
        output.copy_(model_output)
13
        del model_output

有五个细节值得展开：

1. 输入元数据的切分也是按行视图。 input_metadata_source_full 与 input_metadata_device_full 同样用 recursive_fn(_cut) 裁到前 batch_size 行——recursive_fn 的作用是遍历对象里所有张量字段并逐一施加 _cut，因此每个档位拿到的元数据对象与其档位的行数一致。这一步说明：图捕获“冻结”的不只是模型代码，还包括元数据的形状；重放时只更新元数据的内容（recursive_copy_ 对每个张量字段做 non_blocking 拷贝，从本步的新元数据拷入设备侧图输入），不改形状；字典里嵌套的一层张量（例如按压缩比分桶的调度元数据）也会逐个 copy_。
2. 一块最大缓冲，多个视图。 output_full 按 max_batch_size 行分配一次，每个档位 batch_size 通过 narrow 得到一个前 batch_size 行的视图。图捕获时只捕获前 N 行对应的内核，重放时这些内核写入的也是同一块内存的前 N 行。
3. 图内的输出张量不能直接返回，因此要 copy_。 _run_model 调用的是普通的 DeepseekForCausalLM.forward，它返回的 hidden_states 是最后一层 RMSNorm 新建的张量；这张量是在捕获期间分配的，来自图共享内存池。图重放确实会往同一块地址写，但所有档位共用内存池，别的档位捕获时可能复用同一块内存。如果每个档位都直接把图内张量返回，两个档位的输出就可能重叠。捕获时多做一步 output.copy_(model_output)，把结果搬进 output_full 这个不属于图内存池的持久缓冲；这条拷贝也被捕获进图里，重放时自动执行。
4. 图内存池复用。 所有档位共享 self.graph_memory_pool，后捕获的档位复用先捕获档位释放的内存。reversed(self.batch_sizes) 让大档位先捕获，内存池的峰值由最大档位决定，小档位在其中切分。持久输出缓冲放在池外，正是为了不受这种复用影响。

copy_ 这一步是模型层与图运行时之间的契约：模型承诺“输出是一个只有 hidden_states 字段的不可变对象”，图运行时承诺“重放时把结果拷到我指定的视图上”。两者互不知道对方的实现细节，只通过这两个方法交互。

一个只为“最小”存在的输出类

1
@dataclass(frozen=True, slots=True)
2
class CausalLMForwardOutput:
3
    hidden_states: torch.Tensor
4
    """Input of the lm_head. Shape: [num_tokens, hidden_size]."""
5

6
    def narrow(self, num_tokens: int) -> Self:
7
        assert 0 < num_tokens <= self.hidden_states.size(0), f"{num_tokens=}, {self.hidden_states.size(0)=}"
8
        return type(self)(hidden_states=self.hidden_states[:num_tokens])
9

10
    def copy_(self, src: Self) -> None:
11
        self.hidden_states.copy_(src.hidden_states)

整个类只有一个字段。为什么不直接返回张量？narrow 和 copy_ 给出了答案：前者给图运行时切视图，后者给图捕获做结果搬运。frozen=True 保证字段引用不被人改掉——一旦有人把 hidden_states 替换成另一块张量，图重放写入的内存就不再是采样器读取的那块，这种错误在 eager 模式下看不出来，只在图重放下表现为“读到旧值”。slots=True 减少每个对象的开销。类名里的 “CausalLM” 暗示它可能源自 HF 因果语言模型的输出约定，但在 DeepInfer 的执行路径上，模型层只把它当作“一块带形状说明的 hidden 缓冲”的载体——真正的采样输出是 SamplerOutput，两者之间隔着 sample 这一道。

narrow 的断言 0 < num_tokens <= size(0) 刻画了一个边界：CUDA Graph 只捕获了离散的档位 batch_sizes，档位由 build_graph_batch_sizes(max_size, interval) 生成：先列 1、2、4……直到 interval 之前的 2 的幂，再以 interval 为步长铺满到 max_size。例如 $\text{max}=512$、$\text{interval}=32$ 时，档位就是 1、2、4、8、16、32、64……512。重放时真实 batch 不一定恰好落在某个档位上。select 用 bisect_left 找到最小的足够大的档位，replay 返回的就是该档位 narrow 过的视图——这个 narrow 发生在捕获时，每个档位的 CudaGraphData 里存的就是自己那一份视图；重放时不重复 narrow。真实 batch 与档位之间的差距，由采样器用 selected_token_indices gather 掉 padding 行来收尾。两次裁剪、两个位置：视图裁剪在捕获时固定，gather 裁剪在采样时执行。

不图重放时的输出

eager 模式下没有 allocate_output 这一环：DeepseekModel.forward 返回的 hidden_states 就是最后一层 RMSNorm 的输出，生命周期由 Python 引用计数决定。allocate_output 只在图捕获（以及 DFlash 草稿图）这条路径上使用。换句话说，这个 API 的存在本身就是“这条路径会被冻成图”的信号：凡是需要固定地址的地方，才会提前把缓冲申请好。

输入：显式张量 vs 批对象 vs 元数据对象

vLLM（以 V1 的 DeepSeek V2 实现为例）的 DeepseekV2Model.forward 签名是显式张量：

1
def forward(self, input_ids: torch.Tensor, positions: torch.Tensor,
2
            intermediate_tensors: Optional[IntermediateTensors],
3
            inputs_embeds: Optional[torch.Tensor] = None) -> ...:
4
    ...
5
    for layer in self.layers[self.start_layer:self.end_layer]:
6
        hidden_states, residual = layer(positions, hidden_states, residual)
7
    ...
8
    hidden_states, _ = self.norm(hidden_states, residual)
9
    return hidden_states

注意这里没有把 KV cache 和注意力元数据作为参数传进每一层。vLLM 的做法是把注意力元数据放进一个线程局部的 ForwardContext，注意力层内部通过 get_forward_context().attn_metadata 取用。这样做的直接后果是：模型的 forward 签名只保留“逐层数据流”（hidden、residual、positions），与 KV 相关的形状信息全部走侧通道。PP（流水线并行）跨阶段时，IntermediateTensors 显式携带 hidden_states 与 residual——这是 vLLM 跨层状态唯一被“显式化”的地方。

SGLang 的 DeepseekV2ForCausalLM.forward 则引入了一个打包对象：

1
def forward(self, input_ids, positions, forward_batch: ForwardBatch,
2
            input_embeds=None, pp_proxy_tensors=None):
3
    with get_attn_tp_context().maybe_input_scattered(forward_batch):
4
        hidden_states = self.model(input_ids, positions, forward_batch, ...)
5
    ...
6
    return self.logits_processor(input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states)

ForwardBatch 是 SGLang 的核心数据结构，字段覆盖一次前向所需的全部批级信息：forward_mode、input_ids、req_pool_indices、seq_lens、out_cache_loc、以及这一轮 KV 写入的 kv_loc_plan 等。它和 vLLM 的 ForwardContext 的差别在于：ForwardBatch 是显式参数，ForwardContext 是隐式上下文。显式参数让调用图更容易被静态分析；隐式上下文对层代码的侵入更小，层内部可以从任意位置取到元数据。

DeepInfer 的 DeepseekModel.forward 收的是 (kv_caches, input_metadata, hidden_states=None, dflash=None)：kv_caches 与 input_metadata 都是显式参数，没有全局上下文；注意力元数据全部在 InputMetadata 对象里。和 SGLang 相比，DeepInfer 没有 ForwardBatch 这种“把采样参数也打包进来”的设计——采样参数走独立的 SamplerInputs，在 sample 入口才出现。

三者的取舍其实是一条线：vLLM 偏隐式上下文、SGLang 偏批对象、DeepInfer 偏元数据对象。隐式上下文最省参数传递，但对 CUDA Graph 捕获与多线程复用最不利；批对象最方便，但容易变成“什么都往里塞”的大背包；元数据对象最纯粹，代价是采样相关的东西要另起一条路径。DeepInfer 选择 CUDA Graph 与 ZOS 都是强约束的系统，隐式上下文显然不合适；把采样参数拆出去，也正好和 execute_main_model / execute_sample 两段式入口对应。

跨层共享：层内查询 vs 可变容器 vs 批字段传递

这是三个框架差别最大的一点，因为它直接决定了“share-kv / share-index”这类跨层复用在代码里怎么落地。

在 vLLM 中，MLA 这类跨层共享通常由具体的注意力层实现自行处理：共享关系写在层对象上，层内代码按需读取并复用另一层已算好的结果。代码层面的共享关系是由层内代码去查的。

DeepInfer 把这件事抽象成了 SharedLayerTopology：模型构造时一次性算出“哪些层是 KV source、哪些层是 index source、哪些层是 candidate 源”，并做分段校验；运行期再用 CrossLayerRuntime 在层间传值。为什么两者都可行？因为共享关系在任何模型里都是静态的、可以在构造期确定的；区别只在于“构造期计算出关系”与“运行期在层内查询关系”——前者的查询开销在启动时一次付清，后者在每一步都要走一次分支。

CrossLayerRuntime 承载六类跨层值，源码注释把它们逐条列了出来：

这张表解释了为什么容器是“原地写入、不通过返回值”的：每一类值在拓扑里都有明确的“最近 source 层”。SharedLayerTopology 构造时已经做了分段校验，层按递增顺序执行，所以 consumer 需要的值一定是它前面最近的 source 层写进去的那一份，不需要记录来源层再匹配。如果这些值走返回值，层间就要传一个随层号增长的 tuple ——这正是 CUDA Graph 捕获最不希望看到的形状依赖。

SGLang 的跨层共享更多依赖 ForwardBatch 字段在层间的传递与修改。这条路径在阅读上最直白，但代价是 ForwardBatch 的字段会持续增长，越往后越难分清“哪个字段在哪些层之间有效”；DeepInfer 用一个显式的 CrossLayerRuntime 把“层间临时状态”和“批级输入”彻底分开，代价是多了一个需要理解的对象。

输出：logits_processor 与模型内采样

vLLM 的 compute_logits(hidden_states, sampling_metadata) 和 SGLang 的 logits_processor(input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states) 都说明了同一件事：采样是模型层之外的一步。两个框架的模型 forward 都返回 hidden，再交给一个独立的 logits 处理器。这与 DeepInfer 的 DeepseekForCausalLM.sample 属于同一种分层，差别只在于：

• vLLM / SGLang 的 logits 处理器通常先用 lm_head 计算全量 logits，采样器再在 logits 上做采样；
• DeepInfer 把“hidden → logits”和“logits → token”串进了同一个采样器管线（get_logits → fused_sample），并把 MTP/DFlash 草稿路径也挂在采样器内部（见第五篇）。

Movable 是 DeepInfer 独有的

vLLM 和 SGLang 的模型层都没有所有权包装，依赖 Python 引用计数自然回收张量；DeepInfer 显式引入 move-only 包装，核心目的是让 CUDA Graph 捕获的静态调用栈里所有权交接的位置也变成静态的：每一处 move() 在代码里都是固定位置，运行时不会因为引用关系的变化而改变释放时机，从而保证显存峰值可预测（第二十篇的 profile_run 依赖这一点）。

设计取向归纳

DeepInfer 的三层结构（参数壳 HFModelParameters → 运行时模型 DeepseekForCausalLM + DeepseekModel → 执行入口 execute_main_model / execute_sample）带来三个具体收益：

1. 构造阶段就做架构约束校验：MainPredictHeadHC 类型、hc_shifted 与 HCHead 的互斥、DFlash 目标层与 YoCo 边界的兼容性（assert all(in_consumer) or not any(in_consumer)）、compact 截断下 Engram 不得落在 consumer 层——这些断言在 __init__ 里全部 fail-fast，而不是等到第 N 个 prefill step 才暴露。
2. 参数壳与运行时模型生命周期解耦：扁平壳的 del 不会带走运行时模型的数据（reconstruct_param），让启动期的显存画像（第二十篇）能分阶段测量。
3. 前向与采样拆成两个入口：GPU 时间线上天然形成前向、采样两段，ZOS 可以在这两段之间插入调度动作。

多模态与词表对齐

build_embed_tokens 的断言里有一条很容易被忽略：

1
assert isinstance(embedding, MultimodalEmbedding)
2
assert self.cfg.attention_dropout == 0.0

第一条对应的设计是：多模态模型的 HF 嵌入层是一个 MultimodalEmbedding 容器，模型层只取其中的 emb.tok_emb 这一份纯文本词表，用作 VocabParallelEmbedding 的权重。视觉侧的编码由独立的 VE 链路完成（第九篇），图像 token 在词表里的位置是预留的特殊 ID；多模态数据的接入细节不在本篇主循环的范围内，这里只强调模型层接收的 input_ids 已经被分词器与 VE 前端处理好，主循环本身不区分文本与图像 token。

第二条 attention_dropout == 0.0 是推理引擎的通性要求。任何形式的随机性会破坏 batch-invariant 与确定性采样（第六篇），也会让 CUDA Graph 重放的结果不可复现。多数推理框架都在启动时检查这一项，DeepInfer 把它写成构造期的强制断言，避免“配置错了但直到出现数值差异才被发现”。