MAX 推测解码配置指南解读 max.pipelines.speculative 模块SpeculativeConfig、拒绝采样策略与 Ragged Token 合并【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo本文基于 Modular Platform 开源仓库MAX Mojo中 max.python/docs/pipelines.speculative.rst 文档及其对应源码模块max/python/max/pipelines/speculative/系统讲解 MAX 推理流水线中推测解码speculative decoding的配置体系包括SpeculativeConfig的全部字段与 CLI 参数、RejectionSamplingStrategy四种拒绝采样策略、VerifyWidthRange批量调度以及RaggedTokenMerger/ragged_token_merger的 token 合并图。读完本文你将能够直接通过命令行或 Python API 为 MAX 流水线启用并调优 EAGLE、MTP、DFlash 等推测解码方法理解其底层验证与接受流程并能解释各配置项在源码中的真实作用。模块总览speculative 子包结构max.python/docs/pipelines.speculative.rst是 Sphinx autodoc 文档页它通过automodule与autosummary指令聚合了max.pipelines.speculative模块的公开 API。文档将其组织为两大部分Configuration配置RejectionSamplingStrategy、SpeculativeConfig、SpeculativeMethodToken mergingToken 合并RaggedTokenMerger类与ragged_token_merger函数。对应的源码实现位于 max/python/max/pipelines/speculative/max/python/max/pipelines/speculative/ ├── __init__.py # 导出全部公开符号 ├── config.py # SpeculativeConfig / SpeculativeMethod / RejectionSamplingStrategy / VerifyWidthRange ├── depth_schedule.py # 按 batch size 调度验证深度的工具 ├── driver.py # SequentialDriver顺序推测解码EAGLE / MTP的统一驱动 ├── ragged_token_merger.py # RaggedTokenMerger / ragged_token_merger ├── spec_input_types.py # 统一 spec-decode 图输入签名 ├── spec_target.py # SpecDecodeTarget 适配器协议 └── unified_graph_ops.py # merge / verify / mask / accept / shift / propose / pack 七阶段图算子其中init.py 明确导出了MAGIC_DRAFT_TOKEN_ID、RaggedTokenMerger、RejectionSamplingStrategy、SpeculativeConfig、SpeculativeMethod、VerifyWidthRange与ragged_token_merger。SpeculativeMethod四种推测解码方法SpeculativeMethod在 config.py 中定义为字面量类型SpeculativeMethod Literal[eagle, mtp, dflash, dflash2]即 MAX 支持四种推测解码方法方法含义每次提案的 token 数eagleEAGLE草稿模型与目标共享 embedding 与lm_head直接读取目标的隐藏状态进行外推每步 1 个 token默认验证宽度 2mtpMulti-Token Prediction多 token 预测每步 1 个 token默认验证宽度 2dflashDFlash块式草稿block draft一次生成一整块 token由草稿 checkpoint 的训练宽度决定dflash2DFlash2DFlash 的演进版本同上源码中_ONE_TOKEN_PER_STEP (eagle, mtp)标记了两类每步仅草拟 1 个 token的方法它们在num_speculative_tokens未显式设置时默认解析为 2见下文字段详解。SpeculativeConfig提供了三个判别方法is_eagle()、is_mtp()、is_dflash()与is_dflash2()。值得注意的是is_dflash()对dflash与dflash2均返回True——因为 v2 保留了 v1 的融合图形态与块式草稿契约凡是只需要草稿按块到达的消费方流水线类选择、KV cache 容量计算可统一对待两者只有架构选择architecture selection例外因为每个融合图只为一种草稿器构建。MAGIC_DRAFT_TOKEN_ID 42是一个哨兵草稿 token id用于 prefill 阶段与伪草稿图捕获步骤当draft_tokens的某一行每个位置都等于该值时表示没有真实草稿预测需要验证。统一 DFlash 图检测到它会在 prefill 时清零接受数overlap 流水线在真实草稿产生前用它填充草稿槽位。将其定义在配置模块中是为了让图侧architectures与运行时侧lib对同一值保持一致。RejectionSamplingStrategy四种拒绝采样策略RejectionSamplingStrategy同样定义于 config.pyRejectionSamplingStrategy Literal[ greedy, residual, typical-acceptance, logit-comparison ]四种策略的语义如下来自源码 docstringgreedy仅当草稿 token 与目标模型在该位置的 argmax 完全一致时才接受。最严格、分布无偏差但接受率最低。residual从目标分布减去草稿分布后的残差分布中采样这是标准拒绝采样规则用于精确匹配目标分布即保证采样分布无偏。typical-acceptance接受落在目标模型典型集typical set内的草稿 token以轻微的分布不匹配换取更高的接受率。logit-comparison直接比较目标与草稿的 logits 来决定是否接受。重要事实源码明确声明当前没有任何推测解码路径实际读取rejection_sampling_strategy字段——该字段是惰性inert的。每个统一推测架构都会自行构建 AcceptanceSampler其实际分派依据是synthetic_acceptance_rate与use_greedy_acceptance。SpeculativeConfig上的uses_greedy_rejection()、uses_typical_acceptance()、uses_logit_comparison()只是对该字段的便捷判读方法。因此真正决定接受规则的唯一入口是AcceptanceSampler.acceptance_rule属性它在 rejection_sampler.py 中按优先级返回synthetic、stochastic或greedy设置了synthetic_acceptance_rate优先进入合成模式否则use_stochasticTrue走随机拒绝采样否则走 greedy。SpeculativeConfig完整字段与 CLI 映射SpeculativeConfig继承自ConfigFileModelmax.config 的配置文件模型基类其核心作用是选择推测解码方法、每步草拟的 token 数以及目标模型验证草稿的旋钮。文档与源码共同列出的 CLI 映射为配置字段CLI 参数默认值说明speculative_method--speculative-methodNone禁用推测解码取eagle/mtp/dflash/dflash2之一num_speculative_tokens--num-speculative-tokensNone自动解析每步草拟的 token 数num_speculative_tokens_per_batch_size--num-speculative-tokens-per-batch-sizeNone验证全部草稿按 decode batch size 调度的验证宽度num_speculative_tokens_mixed_batch--num-speculative-tokens-mixed-batchNone混合 prefilldecode 批次上的验证数rejection_sampling_strategy--rejection-sampling-strategyNone拒绝采样策略当前惰性synthetic_acceptance_rate--synthetic-acceptance-rateNone合成接受率仅用于基准测试0.01.0除此之外SpeculativeConfig还包含若干高级字段无独立 CLI 映射通过配置文件或 Python 直接构造设置use_relaxed_acceptance_for_thinking、relaxed_topk、relaxed_delta、use_greedy_acceptance、draft_proposal。核心字段详解speculative_method为None时推测解码被禁用否则指定方法。is_dflash()对dflash与dflash2均返回True块式草稿契约相同is_dflash2()仅对dflash2返回True。num_speculative_tokens每次验证通过时草稿提出的 token 数。None表示未设置eagle与mtp在构造时通过字段校验器_resolve_autoregressive_draft_width解析为2而dflash风格的块式草稿由架构从草稿 checkpoint 的训练宽度解析。更大的值可能提高平均草稿接受长度与峰值加速但也可能降低靠后位置的接受率并因额外 token 增加 kernel 延迟。该字段也被draft_width属性使用注意该属性断言字段非空config 未经PipelineConfig.from_args()构建时会断言失败。num_speculative_tokens_per_batch_sizelist[VerifyWidthRange] | None。VerifyWidthRange是一个包含batch_start、batch_end、num_tokens三个字段的 Pydantic 模型表示一个两端闭区间inclusive的 decode batch size 范围及其上要验证的草稿数。其核心思想是每步总是草拟num_speculative_tokens个提案但目标模型验证其中多少个可以由 batch size 决定。例如[ {batch_start: 1, batch_end: 16, num_tokens: 3}, {batch_start: 17, batch_end: 64, num_tokens: 1} ]表示 batch size 116 时验证 3 个草稿1764 时只验证 1 个。约束规则由 depth_schedule.py 的normalize_depth_schedule强制执行并在配置读取时而非流水线构建时校验第一个区间必须从 batch size 1 开始保证每个运行时 batch size 都能解析出验证数区间不得重叠、必须正序start endnum_tokens非负最终每一项都会被max_depth即num_speculative_tokens封顶因为一步不可能验证多于其携带的草稿数区间之间的间隙以及最后一个区间之后的尾部沿用前一个区间的计数。verify_width_schedule属性将配置转换为已排序的(start, end, count)三元组列表None表示未配置。该方法适用于所有推测方法块式草稿器dflash每步仍草拟完整 checkpoint 固定块只是目标验证其中多少块被收窄。build_depth_lookup函数进一步把调度展开为稠密的batch_size - depth查找表索引 0 留空以便直接用运行时 batch size 索引保证对1..max_batch_size全程有定义。num_speculative_tokens_mixed_batch目标模型在混合 prefilldecode 批次上验证的草稿数取值 0。它的收窄理由与 batch-size 调度不同num_speculative_tokens_per_batch_size是在大 batch 下用接受率换解码吞吐而此字段是用它换共享该步的 prefill 行的延迟——prefill 行处于 prompt 首 token 时延TTFT的关键路径上。该字段在读取处被num_speculative_tokens封顶而非在校验处因为dflash的宽度由架构从草稿 checkpoint 解析配置校验时尚不知道上限。它仅在--enable-spec-decode-mixed-batches开启时才可达否则混合批次什么都不验证。synthetic_acceptance_rate仅用于基准测试的覆盖项0.01.0。设置后拒绝采样器绕过真实的草稿/目标比较以校准过的概率接受每个草稿位置使跨num_speculative_tokens个位置的平均联合接受率匹配该值。校验器_validate_synthetic_acceptance_rate保证其必须落在[0.0, 1.0]。用途是不换草稿模型即可建模假想的加速比真实服务时保持未设置。在AcceptanceSampler中它通过compute_synthetic_acceptance_base_rate计算出每位置基准接受率并优先于随机模式分派。rejection_sampling_strategy如前所述当前惰性。启动时配置转储config dump会报告实际生效的AcceptanceSampler.acceptance_rule以及决定它的字段。高级接受规则字段use_relaxed_acceptance_for_thinking默认False启用宽松接受——当草稿位置处于think.../think块内时将目标的 top-N 候选按概率阈值top1_prob - relaxed_delta过滤与草稿 token 比较命中任一候选即接受思考区间之外仍使用既有严格接受规则。要求draft_proposalargmax。relaxed_topk默认10宽松接受激活时从目标分布考虑的 top-N 候选数。relaxed_delta默认0.6低于 top-1 候选的概率间隙间隙内候选保持可接受资格——草稿 token 命中任一概率不低于top1_prob - relaxed_delta的 top-N 候选即被接受。use_greedy_acceptance默认False使用贪婪argmax草稿接受替代随机采样器。贪婪路径无图中间分配mid-graph allocation因此融合推测图可被 CUDA-graph 捕获。仅对贪婪服务temperature 0、top_k 1有效与宽松接受及合成接受不兼容。draft_proposal默认argmax草稿模型如何提案 token。argmax确定性提案sampled使草稿自行采样提案并保留其采样分布从而验证执行真正的推测采样true speculative sampling而非典型接受。sampled与use_relaxed_acceptance_for_thinking不兼容_validate_draft_proposal模型校验器会抛出ValueError且除非服务架构支持否则无效。在AcceptanceSampler中draft_proposalsampled还要求提供vocab_size。relaxed_topk校验器要求 1relaxed_delta校验器要求落在[0.0, 1.0]。不可变性与程序化构造SpeculativeConfig设置了model_config ConfigDict(frozenTrue)实例不可变构造后赋值任何字段都会抛出异常。文档给出的程序化构造示例from max.pipelines.speculative import SpeculativeConfig spec SpeculativeConfig( speculative_methodeagle, num_speculative_tokens3, )_config_file_section_name speculative_config表明配置文件中的对应段落名。在 pipeline_args.py 中(speculative, SpeculativeConfig)注册使--speculative-method、--num-speculative-tokens等参数被嵌套到{speculative: {...}}子配置段落PipelineConfig中speculative: SpeculativeConfig | None字段与draft_model: MAXModelConfig | None字段并存后者承载草稿模型的配置见 pipeline_args.py。这意味着 CLI 上还有一个--draft-model相关的配置入口用于指定草稿模型。顺序推测解码的运行驱动SequentialDriver虽然文档只列出配置类与 token 合并类但 driver.py 是理解这些配置如何落地的关键它实现了 EAGLE / MTP 这类顺序推测解码的模型无关驱动。树中每个统一 spec-decode 模块都运行相同的七个阶段merge - verify - mask - accept - shift - propose - packSequentialDriver拥有这些阶段一个模型贡献一个SpecDecodeTarget目标适配器与一个SequentialProposer草稿提议器。SequentialProposer是 Protocol声明了prefill草稿第 0 步跑完整条目标校正序列与step草稿第index步每个 batch 元素一个 token。这两个适配器刻意不是Module子类驱动在相同属性名下注册目标与草稿模块self.target、self.merger、self.draft使权重加载、state_dict键与权重注册表保持不变。一次 spec-decode 迭代的输入包括一维 ragged prompt token[total_seq_len]段边界来自input_row_offsets、上一轮提案的draft_tokens[batch, K]、每个设备的信号缓冲、目标缓存、采样参数seed、temperature、top_k、top_p等、in_thinking_phase逐行标志启用宽松接受、可选的专家并行输入与结构化输出 bitmask输出为(num_accepted, next_tokens, next_draft_tokens)三元组。驱动内部还处理草稿 KV cache 的q1重定向DecodeKVSwap、跨设备 hidden state 的携带与 DP 分片切片、以及按步推进缓存长度等簿记。AcceptanceSampler由驱动以use_stochasticTrue构建并传入synthetic_acceptance_rate与宽松接受参数仅当 proposer 声明uses_thinking_phase且配置开启use_relaxed_acceptance_for_thinking时。num_draft_steps直接取自speculative_config.num_speculative_tokens配置为None时回退为 1。Token 合并RaggedTokenMerger 与 ragged_token_merger文档的第二部分聚焦 token 合并。推测解码需要把prompt 序列与草稿 token拼成一次前向传播可处理的 ragged 批次。ragged_token_merger.py 提供了两种使用方式。函数式图构建ragged_token_merger(device)ragged_token_merger(device: DeviceRef) - Graph构建一个把 prompt 与草稿 token 合并为单一 ragged 序列的图图输入为三个张量prompt_tokens[batch_prompt_seq_len]int64、prompt_row_offsets[offsets_len]uint32、draft_tokens[batch_size, draft_seq_len]int64图名为merge_prompt_draft_tokens输出为merged_tokens与merged_row_offsets两个张量。模块封装RaggedTokenMergerRaggedTokenMerger是max.nn.layer.Module子类__call__接收三个张量并返回合并结果参数形状含义prompt_tokens[S]prompt tokenprompt_offsets[B1]prompt 的 row offsets独占前缀和draft_tokens[B, K]草稿 tokenK 为每请求草稿数返回merged_tokens形状[SB*K]merged_offsets形状[B1]。实现细节源码可验证K draft_tokens.shape[1]通过_shape_to_scalar在 GPU 上从 shape 转出标量——这是为兼容 CUDA Graphs 而绕开shape_to_tensor的 h2d 拷贝源码注释称之为 especially cursed way并留有 TODO 希望在 shape_to_tensor 支持设备放置后删除。草稿被reshape展平为(-1,)然后借助max.nn.kernels.merge_ragged_tensors与手工构造的draft_offsets [0, K, 2K, ..., N*K]完成合并。构造draft_offsets时特意使用range(step1) * num_steps而非range(stepnum_steps)因为 prefill 阶段draft_seq_len0时 step0 未定义。配套的compute_host_merged_offsets在 CPU 上计算合并偏移merged_offsets[i] host_input_row_offsets[i] i * K镜像 GPU 侧逻辑但留在 CPU 上从而避免 D2H 拷贝阻塞 CUDA graph 捕获。在SequentialDriver中merge_tokens_and_host_offsets一次计算 GPU 侧与主机侧两组偏移随后ops.distributed_broadcast将合并偏移广播到各设备复用目标验证、草稿第 0 步、接受位置 gather 共享同一份广播。配置文件与 CLI 综合示例综合以上通过 CLI 启用 EAGLE 推测解码的典型命令形态字段与 pipeline_args.py 的注册保持一致max-serve \ --model-path /path/to/target-model \ --draft-model /path/to/draft-model \ --speculative-method eagle \ --num-speculative-tokens 3 \ --enable-spec-decode-mixed-batches \ --num-speculative-tokens-mixed-batch 1等价地通过配置文件speculative_config段落speculative: speculative_method: eagle num_speculative_tokens: 3 num_speculative_tokens_per_batch_size: - {batch_start: 1, batch_end: 16, num_tokens: 3} - {batch_start: 17, batch_end: 64, num_tokens: 1} num_speculative_tokens_mixed_batch: 1 synthetic_acceptance_rate: null注意 DFlash 系列架构的宽度约束例如unified_dflash_llama3的model_config.py会校验--num-speculative-tokens与 checkpoint 固定的块宽度一致不匹配时报错提示显式指定而unified_dflash2_qwen3_5允许省略参数、由架构自动填充见 unified_dflash2_qwen3_5/model_config.py 与 unified_dflash_llama3/model_config.py 的注释。实践建议与注意事项rejection_sampling_strategy当前不生效源码明确注释没有任何推测路径读取该选项实际接受规则由AcceptanceSampler依synthetic_acceptance_rate→use_stochastic→ greedy 的优先级决定rejection_sampler.py。迁移/调试时请以启动配置转储中报告的acceptance_rule为准。synthetic_acceptance_rate只用于基准测试设置后草稿验证不再依赖真实 logits适合在固定接受率假设下估算加速上限真实服务务必留空。use_greedy_acceptance与 CUDA Graphs只有贪婪路径没有图中间分配融合推测图才能被 CUDA-graph 捕获它仅适用于贪婪服务temperature 0、top_k 1且与宽松接受、合成接受互斥。验证宽度调度尽早失败num_speculative_tokens_per_batch_size在配置读取时即由normalize_depth_schedule校验首区间必须从 1 开始、不重叠、正序、非负而不是等到模型加载数分钟后才报错config.py 注释明确说明这一设计意图。SpeculativeConfig不可变frozenTrue程序化构造时一次性传入所有字段读取draft_width前请确认配置经由PipelineConfig.from_args()构建。延伸阅读配置基类与嵌套段落机制max/python/max/pipelines/lib/pipeline_args.py接受采样器实现合成/随机/贪婪三分派与 RNG 种子域设计max/python/max/nn/sampling/rejection_sampler.py验证宽度调度校验与稠密查找表max/python/max/pipelines/speculative/depth_schedule.py顺序推测解码七阶段驱动max/python/max/pipelines/speculative/driver.py应用侧参考DFlash2 架构的宽度约束与使用示例见 unified_dflash2_qwen3_5/batch_processor.py其中提到--speculative-method dflash2 --draft-model ...的运行方式以及 max/python/max/pipelines/architectures/speculators_common/draft_config.py 中草稿模型配置的通用处理。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考