ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

Eagle3 + DeepseekV3 投机解码架构:MAX 中 `eagle3_deepseekV3` 模块的完整解析

Eagle3 + DeepseekV3 投机解码架构:MAX 中 `eagle3_deepseekV3` 模块的完整解析 Eagle3 DeepseekV3 投机解码架构MAX 中eagle3_deepseekV3模块的完整解析【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo导读本文围绕 Modular MAX 仓库中max/python/max/pipelines/architectures/eagle3_deepseekV3模块其 API 文档入口为max/python/docs/pipelines.architectures.eagle3_deepseekV3.rst展开讲解如何在一个 DeepseekV3 系列MLA 架构目标模型上运行 EAGLE3 投机解码目标模型与草稿模型如何被融合进同一张编译图、草稿权重如何加载与共享、KV 缓存如何拆分、以及如何通过 CLI 与SpeculativeConfig启用该能力。读完本文你将掌握该模块的类层次、图构建流程、权重适配规则与两种草稿变体MLA 草稿与 MHA 草稿的差异并能在 MAX 中按需配置--speculative-method eagle与--draft-model。1. 模块定位一个面向 DeepseekV3 的 EAGLE3 投机解码流水线max/python/docs/pipelines.architectures.eagle3_deepseekV3.rst是 Sphinx 的automoduleAPI 文档条目声明了模块max.pipelines.architectures.eagle3_deepseekV3并通过:members:、:imported-members:、:show-inheritance:自动展开其全部公开类与函数。该文档条目的实体内容由源码目录承载max/python/max/pipelines/architectures/eagle3_deepseekV3/下共 8 个文件__init__.py—— 模块公共导出arch.py—— 两个Speculator架构注册eagle3_draft.py——Eagle3DeepseekV3草稿模型别名model.py—— MLA 草稿版统一PipelineModelmha_pipeline.py—— MHA 草稿版统一PipelineModelbatch_processor.py—— 两个批处理类unified_eagle.py—— 融合图驱动SequentialDriver子类weight_adapters.py—— 草稿 checkpoint 权重键映射__init__.py的模块 docstring 直接点明其定位“DeepseekV3 Eagle3 speculator pipeline”并导出了 8 个公开符号__init__.py__all__ [ Eagle3DeepseekV3, Eagle3DeepseekV3Inputs, Eagle3DeepseekV3Model, Eagle3DeepseekV3Unified, Eagle3MHADeepseekV3Inputs, Eagle3MHADeepseekV3Model, convert_eagle3_draft_state_dict, eagle3_deepseekV3_speculator, eagle3_mha_deepseekV3_speculator, ]提示eagle3_deepseekV3目录名与 Python 模块名中的deepseekV3采用了与目标架构目录deepseekV3一致的驼峰拼写引用时需保持大小写一致。2. 架构注册Speculator与懒加载注册表arch.py是模块的架构入口。它从max.pipelines.lib引入Speculator从deepseekV3.arch引入目标架构deepseekV3_arch然后声明了两个互补的 speculatorarch.py# Eagle3 with an MLA draft shipped as its own repo. eagle3_deepseekV3_speculator Speculator( nameEagle3DeepseekV3ForCausalLM, basedeepseekV3_arch, draft_archEagle3DeepseekV2ForCausalLM, methodeagle, pipeline_modelEagle3DeepseekV3Model, batchingEagle3DeepseekV3BatchProcessor, example_repo_ids[austinpowers/Kimi-K2.5-NVFP4-DeepseekV3], opt_out_cascadeTrue, ) # Eagle3 with a Llama-style MHA draft, which runs Kimis fused graph. eagle3_mha_deepseekV3_speculator Speculator( nameEagle3MHADeepseekV3ForCausalLM, basedeepseekV3_arch, draft_archLlamaForCausalLMEagle3, methodeagle, pipeline_modelEagle3MHADeepseekV3Model, batchingEagle3MHADeepseekV3BatchProcessor, example_repo_ids[austinpowers/Kimi-K2.5-NVFP4-DeepseekV3], opt_out_cascadeTrue, )对照Speculator的定义arch_lookup.py可以理解每个字段的作用字段含义name融合架构的名称出现在日志与生成文档中base被加速的目标架构此处为deepseekV3_archdraft_arch草稿 checkpoint 的huggingface_config.architectures[0]用于与草稿仓库匹配method投机解码机制必须与配置的speculative_methodeagle一致pipeline_model融合流水线模型目标与草稿运行在同一张图batching批处理类example_repo_ids替换目标架构自带的示例仓库注释说明该融合图由不同的 checkpoint 驱动opt_out_cascade清除基架构的cascade_pipeline_factory因为融合投机图没有 cascade 路径Speculator本质上不是独立架构而是目标架构的“有界增量”derive()只覆盖上述字段tokenizer、config 类、编码、内存规划器、工具/推理解析器等全部继承自base从而避免与目标架构漂移。在全局架构注册表中这两个 speculator 通过_LazyArch懒加载登记并显式声明speculates_onDeepseekV3ForCausalLMarchitectures/__init__.py_LazyArch( Eagle3DeepseekV3ForCausalLM, .eagle3_deepseekV3, eagle3_deepseekV3_speculator, speculates_onDeepseekV3ForCausalLM, ), _LazyArch( Eagle3MHADeepseekV3ForCausalLM, .eagle3_deepseekV3, eagle3_mha_deepseekV3_speculator, speculates_onDeepseekV3ForCausalLM, ),选择逻辑位于select_speculatorarch_lookup.py方法method与草稿架构draft_arch必须同时匹配——仅匹配草稿会运行“草稿恰好适配的机制”而非用户要求的方法仅匹配方法则无法区分使用不同草稿的两个 speculator。若目标声明了 speculator 但两者均不匹配会抛出ValueError并列出已声明的组合。此外_apply_speculative_draft_architectureconfig.py会在配置阶段重写草稿模型的 HF 架构名例如把LlamaForCausalLM改写为LlamaForCausalLMEagle以适配 EAGLE 系列投机解码。3. 统一图模型目标与草稿编译进同一张图3.1Eagle3DeepseekV3ModelMLA 草稿版Eagle3DeepseekV3Model继承自_UnifiedSpecDecodeModelMixin与DeepseekV3Modelmodel.pydocstring 明确其形态“target draft in one compiled graph”——目标权重来自 DeepseekV3 形状的主 checkpoint草稿权重来自独立的 Eagle3 checkpointpipeline_config.draft_model。其关键流程分三步_load_state_dict目标权重通过parse_state_dict_from_weights解析草稿权重通过draft_model_config.resolved_weight_paths()load_weights加载再经convert_eagle3_draft_state_dict转换为模块键model.py。_create_model_config目标 HF 配置通常不带eagle_config因此从草稿 HF 配置中提取eagle_aux_hidden_state_layer_ids并写入目标配置同时把 KV 缓存参数包装成MultiKVCacheParams.from_params({target: ..., draft: ...})的统一树——草稿 KV 缓存固定为 1 层replace(target_kv_params, num_layers1)model.py。_build_graph_for_compile构造Eagle3DeepseekV3Unified融合模型在加载权重前把draft.embed_tokens与target.embed_tokens共享为同一个Weight对象随后target.load_state_dict(strictTrue)、draft.load_state_dict(strictFalse)model.py。图中以eagle3_deepseekV3_graph为名输入包括 token、行偏移、信号 buffer、统一 KV 树、draft_tokens、采样参数seed/temperature/top_k/max_k/top_p/min_top_p以及结构化输出开启时的 bitmask 三元组。3.2 多设备约束_create_model_config中有一个硬性校验当设备数 1且runtime.ep_size ! n_devices时抛出ValueError(Only the EP strategy is supported.)。也就是说Eagle3 DeepseekV3 的多卡运行目前只支持专家并行EP策略这一限制在 MHA 草稿版mha_pipeline.py中同样存在。4. Eagle3 草稿模型Eagle3MLADraft的实现eagle3_draft.py中的Eagle3DeepseekV3是一个“薄命名别名”实现全部位于共享的Eagle3MLADrafteagle_common/eagle_mla_draft.py。之所以保留子类是为了让既有导入from .eagle3_draft import Eagle3DeepseekV3与Eagle3DeepseekV3Unified的图接线保持不变。Kimi K2.5 的Eagle3KimiK25同样是该基类的别名kimik2_5/eagle3_kimi_k25.py因此两个 MLA 目标共享同一草稿实现。Eagle3MLADraft的 docstring 精确描述了其前向结构eagle_mla_draft.pyEagle3 fuses 3 target hidden states (early, middle, last layer) via a linear projection (fc:3*H - H), concatenates the result with the token embedding (H H 2*H), and feeds that through a single MLA decoder layer configured tohidden_size 2*H, followed by a dense MLP and the output projection.对应到构造函数中的模块装配eagle_mla_draft.py组件说明embed_tokensVocabParallelEmbedding与目标共享加载前被别名指向target.embed_tokensfc线性投影3*H - H无 bias复制分片策略ropeDeepseekYarnRotaryEmbedding必需rope_scaling配置否则构造时报错decoder_layer单个DeepseekV3DecoderLayerhidden_size 2*Hwide_configlayer_idx0稠密 MLP 区不走 EP-MoEep_managerNonehidden_norm/ 各 LayerNorm全部复制分片dtype 由config.norm_dtype决定o_proj替换为Linear(num_heads * v_head_dim - H)BF16dense_mlp稠密MLP非量化、非 MoEnormlm_head输出侧 RMSNorm 与ColumnParallelLinear独立加载不共享前向__call__的关键路径eagle_mla_draft.py通过fuse_captured_hidden_states融合目标捕获的隐藏状态若输入宽度为3*H则再经fc_shards投影到H。h_embed embed_tokens(tokens)随后对两个流分别做 RMSNorminput_layernorm与hidden_norm按最后一维拼接得到2*H输入。送入单层 MLA 注意力prefill 阶段构造MLAPrefillMetadata以buffer_lengths batch_context_lengths填充。残差 → post-attention norm → 稠密 MLP → 残差。按input_row_offsets收集最后 tokenDP 模式先 gather 再 allgather经normlm_head输出 logits。依据return_logits/return_hidden_states决定是否输出可变长度 logitsVARIABLE与隐藏状态LAST/ALL。前向签名中的split_prefix参数默认eagle3_draft用于为符号维度命名且要求每次图调用唯一——prefillstep 0与 decodestep 1会传入不同前缀以避免维度冲突。5. 权重适配从 Eagle3 checkpoint 到模块键Eagle3 草稿 checkpoint例如nvidia/Kimi-K2.5-Thinking-Eagle3的键形如fc.*、layers.0.*、norm.*、lm_head.*。convert_eagle3_draft_state_dictweight_adapters.py通过前缀映射将其改写为Eagle3DeepseekV3模块层级_EAGLE3_KEY_MAP: dict[str, str] { layers.0.hidden_norm.: hidden_norm., layers.0.input_layernorm.: decoder_layer.input_layernorm., layers.0.self_attn.: decoder_layer.self_attn., layers.0.post_attention_layernorm.: decoder_layer.post_attention_layernorm., layers.0.mlp.: decoder_layer.mlp., }映射未命中的键fc.*、norm.*、lm_head.*直接透传。权重加载语义所有权重都加载norm与lm_head相对目标保持独立只有embed_tokens与目标共享。在_build_graph_for_compile中草稿的非共享权重会被重命名为draft.*前缀例如draft.norm.weight与目标的norm.weight区分从而保证图级名称唯一同时会校验draft_expected - draft_provided中是否存在非共享缺失权重缺失即报错多余仅告警。另外_create_draft_configmodel.py会从草稿 HF 配置覆盖rope_scaling例如草稿可能使用beta_fast1.0而目标为32.0并检测量化状态若目标配置了 FP4 但草稿权重中没有weight_scale键则判定草稿为 BF16 稠密权重自动关闭草稿的量化配置并把 dtype/EP 分发类型切回bfloat16。6. 隐藏状态捕获eagle_aux_hidden_state_layer_idsEAGLE3 需要目标模型在指定中间层输出隐藏状态。extract_eagle_aux_layer_idsmodel.py从草稿 HF 配置的eagle_config子对象中读取eagle_aux_hidden_state_layer_ids该列表要求全为正整数i 0直接报错因为“捕获第 0 层输入即原始 token 嵌入”尚未接线返回值统一- 1转为 0 基索引。目标侧的执行路径位于deepseekV3.py当return_hidden_states ReturnHiddenStates.SELECTED_LAYERS时capture_for_eagle3会在层输出回调中收集指定层的隐藏状态deepseekV3.pydef capture_for_eagle3(idx: int, h_out: list[TensorValue]) - None: if idx in eagle3_capture_ids: eagle3_captured.append(list(h_out))如果目标配置缺失这些 ID 且草稿 HF 配置里也没有Eagle3DeepseekV3Model._create_model_config会抛出ValueError“eagle_aux_hidden_state_layer_ids must be present in the draft models eagle_config for EAGLE3 hidden-state capture”。7. KV 缓存统一{target, draft}树两个统一模型都把 KV 缓存组织为MultiKVCacheParamsmodel.pyself._draft_kv_params replace(target_kv_params, num_layers1) self.kv_params MultiKVCacheParams.from_params( {target: target_kv_params, draft: self._draft_kv_params} )目标侧沿用 DeepseekV3 的 MLA KV 缓存KVCacheParams草稿侧是单层缓存num_layers1。在_build_graph_for_compile中通过kv_params.unflatten_basic_kv_tree(...)解包为kv_collections目标与draft_kv_collections草稿两组PagedCacheValues。MHA 草稿版的缓存几何则完全独立mha_pipeline.py草稿 KV 由pipeline_config.model.kv_cache.to_params(...)单独构造is_mlaFalse头数/头维度取自草稿 HF 配置num_key_value_heads、head_dim、num_attention_heads且dtype 强制为bfloat16——因为 MHA flash attention 要求 q/k/v 与 KV 缓存 dtype 一致而草稿模块输出 BF16。若目标 KV 不是 BF16会打印告警。8. 批处理与采样输入batch_processor.py定义_Eagle3DeepseekV3BatchProcessorBase继承自DeepseekV3BatchProcessor并在其上追加 EAGLE3 特有的两个输入batch_processor.pyseed_next_seed()维护一个单调递增的uint64[1]计数器装载到 device 0作为采样种子draft_tokensNoneprefill 首步还没有任何草稿预测置空交给图中处理。Eagle3DeepseekV3Inputs的buffers属性把UnifiedSpecDecodeInputs._spec_decode_tail_buffers(include_in_thinking_phaseFalse)追加到DeepseekV3Inputs之后MHA 版则传入include_in_thinking_phaseTrue因为其图额外绑定了逐行的in_thinking_phase标志用于思考阶段松弛接受。9. 融合驱动SequentialDriver与七阶段投机迭代Eagle3DeepseekV3Unifiedunified_eagle.py继承SequentialDriverdocstring 将其描述为“Fused nn.Module: merge target forward rejection shift”——融合、目标前向、拒绝采样、移位全部内建class Eagle3DeepseekV3Unified(SequentialDriver): target: DeepseekV3 draft: Eagle3DeepseekV3 def __init__(self, config, draft_configNone, speculative_configNone, enable_structured_outputFalse): target DeepseekV3(config) draft Eagle3DeepseekV3(draft_config) super().__init__( DeepseekV3Target(target), Eagle3DeepseekV3Proposer(draft), target_modeltarget, draft_modeldraft, ... )按照spec_adapters.py的说明max.pipelines.speculative.driver中的驱动负责一次顺序投机解码迭代的七个阶段每类模型只需提供“如何调用目标并读取输出元组”与“如何调用草稿”这一对适配器。DeepseekV3 的两个 speculator——MTPNextN 头内置于目标 checkpoint与 Eagle3独立草稿 checkpoint——共享同一个DeepseekV3MLAProposer仅头部类属性不同。Eagle3DeepseekV3Proposerunified_eagle.py的关键属性class Eagle3DeepseekV3Proposer(DeepseekV3MLAProposer): split_prefix eagle3 carry_dim_prefix draft_step step_hidden_mode ReturnHiddenStates.ALLstep 1..K 运行在 decode 模式每 batch 元素一个 token此时 ALL-hs LAST-hs选择ReturnHiddenStates.ALL可避免 LAST 路径的 allgather且无需按副本切片DeepseekV3MLAProposer的prefill阶段强制草稿使用ReturnHiddenStates.ALLReturnLogits.VARIABLE为已接受位置做逐元素 gather、为草稿 argmax 提供可变 logitsstep阶段则切回step_hidden_modeLAST_TOKENDeepseekV3Target.verify读取目标的输出元组[last_logits, logits, offsets, hidden per device...]DeepseekV3 没有emit_last_token_logits属性。10. 配置与 CLI如何启用 EAGLE3 投机解码投机解码的配置由SpeculativeConfig承载speculative/config.py其 docstring 给出了 CLI 与编程两种方式CLI 参数与max.pipelines的命令行入口对应--speculative-method eagle选择 EAGLE 方法可选eagle、mtp、dflash、dflash2--num-speculative-tokens N每步草稿 token 数。EAGLE/MTP 未设置时默认解析为2_ONE_TOKEN_PER_STEP--num-speculative-tokens-per-batch-size按解码 batch 大小区间含端点决定验证多少个草稿 token例如[{batch_start: 1, batch_end: 16, num_tokens: 3}, ...]--num-speculative-tokens-mixed-batch混合 prefilldecode batch 的验证宽度用于权衡 prefill 行的 TTFT--rejection-sampling-strategygreedy/residual/typical-acceptance/logit-comparison。注意该字段目前是惰性的——统一投机架构各自构建自己的AcceptanceSamplermax/python/max/nn/sampling/rejection_sampler.py实际接受规则由其acceptance_rule决定--synthetic-acceptance-rate0.0~1.0 的基准测试专用覆盖按校准概率接受草稿位置用于在不更换草稿模型的前提下模拟加速比。编程方式from max.pipelines.speculative import SpeculativeConfig spec SpeculativeConfig( speculative_methodeagle, num_speculative_tokens3, )与投机解码配套的还有草稿模型配置pipeline_config.draft_modelMAXModelConfig指向独立草稿 checkpoint统一模型通过draft_model_config.resolved_weight_paths()解析权重路径、通过draft_model_config.huggingface_config读取草稿的 HF 配置。仓库其他模块的用法示例如unified_dflash2_qwen3_5/batch_processor.py中出现的mach/tools/gen-mef --speculative-method dflash2 --draft-model ...说明--draft-model是配套的 CLI 参数。另外两个与接受规则相关的配置项use_relaxed_acceptance_for_thinking在think.../think块内启用松弛接受目标 top-N 候选中任一匹配即接受要求draft_proposalargmax相关阈值relaxed_topk默认 10与relaxed_delta默认 0.6use_greedy_acceptance贪心 argmax 接受无图中途分配使融合投机图可被 CUDA graph 捕获仅适用于 temperature 0 / top_k 1 的贪心服务。11. MHA 草稿变体Eagle3MHADeepseekV3Modelmha_pipeline.py是 MLA 草稿版的姊妹实现针对“Llama 风格 MHA 的 Eagle3 头LlamaForCausalLMEagle3叠加在裸DeepseekV3ForCausalLM目标上”的场景直接复用 Kimi K2.5 的融合图Eagle3MHAKimiK25Unifiedmha_pipeline.py。与 MLA 版的关键差异融合宽度推断_infer_fc_input_multiplier根据草稿状态字典与hidden_size推断fc是 2 路融合[hidden*2, hidden]的fc.weight与EagleLlama3布局一致还是 3 路融合aux 层 ID 兜底若草稿 HF 配置缺失eagle_aux_hidden_state_layer_ids_default_aux_layer_idsmha_pipeline.py会按fc_input_multiplier生成均匀分布的默认值——2 路为[1, num_layers-4]3 路为[1, num_layers//2 - 1, num_layers-4]注释说明相比 vLLM 的SupportsEagle3默认整体下移一层对modularai/kimi-k2.5-eagle3checkpoint 经验上更匹配随后校验len(aux_ids) fc_input_multiplier不相等直接报错图绑定MHA 图额外绑定in_thinking_phase与结构化输出 bitmask 三元组图名为eagle3_mha_deepseekV3_graph。12. 构建依赖BUILD.bazeleagle3_deepseekV3/BUILD.bazel展示了该模块的依赖面numpy、transformers、typing-extensions以及max.driver、max.dtype、max.engine、max.graph、max.nn、max.pipelines/architectures/deepseekV3、max.pipelines/architectures/eagle_common、max.pipelines/architectures/kimik2_5、max.pipelines/context、max.pipelines/lib、max.pipelines/speculative。该包已被architectures/all_arches.bzl收录//max/python/max/pipelines/architectures/eagle3_deepseekV3随架构注册表一起参与构建。13. 总结max.pipelines.architectures.eagle3_deepseekV3模块为 MAX 中的 DeepseekV3 系列目标提供了一条完整的 EAGLE3 投机解码路径其核心设计可归纳为单一融合图目标与草稿在同一个编译图eagle3_deepseekV3_graph/eagle3_mha_deepseekV3_graph中运行SequentialDriver负责融合、验证、拒绝与移位草稿复用Eagle3MLADraft是 DeepseekV3 与 Kimi K2.5 两个 MLA 目标共享的草稿实现Eagle3DeepseekV3仅为保持导入兼容的别名权重共享规则仅embed_tokens与目标共享norm/lm_head独立加载非共享草稿权重以draft.*前缀入图两种草稿变体MLA 草稿Eagle3DeepseekV2ForCausalLM架构与 Llama 风格 MHA 草稿LlamaForCausalLMEagle3架构后者复用 Kimi K2.5 的融合实现并支持 2/3 路fc融合宽度推断配置驱动通过--speculative-method eagle--draft-model启用num_speculative_tokens默认 2并支持按 batch 大小的验证宽度调度与思考块内松弛接受等高级选项。实际部署时请务必确保目标 checkpoint 为 DeepseekV3 架构、草稿 checkpoint 的 HF 配置包含或可推导出eagle_aux_hidden_state_layer_ids并注意多卡场景目前仅支持 EP 策略这一前提条件。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表