MAX 中 GLM-5.2 (DeepSeek-V3.2 Sparse) 的统一 MTP 推测解码架构解析-北京尧图网络

MAX 中 GLM-5.2 (DeepSeek-V3.2 Sparse) 的统一 MTP 推测解码架构解析【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo本文讲解 MAX 平台中UnifiedMTPGlm5_2模块的设计与实现。该模块将 DeepSeek-V3.2 稀疏 MoE 目标模型和单层稀疏 NextN Draft 模型、贪心拒绝采样及 prefill shift 融合为单一可编译图结构是 GLM-5.2 系列zai-org/GLM-5.2在 MAX Pipeline 中进行推测解码的核心架构组件。读完本文你将理解其双层 KV 缓存设计、index_share_for_mtp_iteration优化策略、权重量化适配逻辑以及完整的编译与运行时流程。架构概述UnifiedMTPGlm5_2定义在 max/python/max/pipelines/architectures/unified_mtp_glm5_2/unified_mtp_glm5_2.py 中类似于UnifiedMTPDeepseekV3之于 DeepSeek-V3 的关系是 V3.2 sparse 对应的 MTPMulti-Token Prediction版本。与标准 V3.2 的 MTP 有两个关键的结构差异双层稀疏 KV 缓存目标网络target和草稿网络draft均使用稀疏 MLAlightning indexer因此各自携带一对{mla, indexer}KV 缓存而非单一的 MLA 缓存。index_share_for_mtp_iteration草稿网络的 lightning indexer 仅在 step 0 执行一次 top-k 选择之后各 draft step 通过 gather 已被接受的 token 位置来复用该选择结果避免重复计算。该模块继承自Module将 token merging、V3.2 target 前向传播、贪心拒绝采样和稀疏 draft 前向传播融合为一个端到端的图。模块结构与注册架构注册在 arch.py 中该架构注册为unified_mtp_glm5_2_arch SupportedArchitecture( nameUnifiedMTPGlmMoeDsaForCausalLM, taskPipelineTask.TEXT_GENERATION, example_repo_ids[zai-org/GLM-5.2-FP8], default_encodingfloat8_e4m3fn, supported_encodings{float4_e2m1fnx2, float8_e4m3fn, bfloat16}, multi_gpu_supportedTrue, pipeline_modelUnifiedMTPGlm5_2Model, tokenizerGlmTokenizer, context_typeTextContext, default_weights_formatWeightsFormat.safetensors, weight_adapters{WeightsFormat.safetensors: convert_with_mtp_state_dict}, supports_empty_batchesTrue, requires_max_batch_context_lengthTrue, configGlm5_1Config, memory_plannerDeepseekV3_2MemoryPlanner, batchingUnifiedMTPGlm5_2BatchProcessor, tool_parserglm45, reasoning_parserglm45, default_structured_output_backendxgrammar, default_structured_output_any_whitespaceTrue, )关键配置说明配置项值说明nameUnifiedMTPGlmMoeDsaForCausalLM架构标识名用于 Pipeline 自动匹配example_repo_ids[zai-org/GLM-5.2-FP8]在 docs/max/models.mdx 的模型表格中GLM-5.1 条目下同样列出了zai-org/GLM-5.2、zai-org/GLM-5.2-FP8等 Model IDdefault_encodingfloat8_e4m3fn默认权重量化编码supported_encodingsfloat4_e2m1fnx2,float8_e4m3fn,bfloat16支持的量化精度multi_gpu_supportedTrue支持多 GPU 分布式部署default_weights_formatsafetensorsHuggingFace safetensors 格式weight_adaptersconvert_with_mtp_state_dict权重 key 映射适配器模块依赖依据 BUILD.bazel 的依赖列表该模块的核心依赖包括//max/python/max/nn— 神经网络层基类//max/python/max/pipelines/architectures/deepseekV3— DeepSeek-V3 权重映射//max/python/max/pipelines/architectures/deepseekV3_2— V3.2 目标模型//max/python/max/pipelines/architectures/deepseekV3_2_nextn— NextN Draft 模型//max/python/max/pipelines/architectures/glm5_1— GLM-5.1 基类Tokenizer、ReasoningParser、ToolParser//max/python/max/pipelines/speculative— 推测解码配置与统一图操作UnifiedMTPGlm5_2 前向传播详解__init__与配置初始化UnifiedMTPGlm5_2.__init__接受三个核心配置def __init__( self, config: DeepseekV3_2Config, # 目标模型配置 draft_config: DeepseekV3_2NextNConfig | None, # 草稿模型配置 speculative_config: SpeculativeConfig | None, # 推测解码配置 enable_structured_output: bool False, ) - None在初始化中num_draft_steps从speculative_config.num_speculative_tokens读取每次生成的草稿 token 数量默认为 1。AcceptanceSampler初始化接受采样器支持宽松接受relaxed acceptance——在 thinking 阶段若use_relaxed_acceptance_for_thinkingTrue则relaxed_topk和relaxed_delta参数放宽拒绝条件加速 thinking 阶段的 token 吞吐。该采样器定义在 max/nn/sampling/rejection_sampler.py 中。target DeepseekV3_2(config)初始化稀疏 V3.2 目标模型emit_last_token_logits False以抑制最后一个 token 的 logits 输出。merger RaggedTokenMerger初始化 ragged token 合并器用于拼接用户输入 token 和草稿 token。draft DeepseekV3_2NextN(draft_config)初始化单层稀疏 NextN 草稿模型。__call__前向流程前向传播分为 Step 0 和后续步骤整体流程如下阶段 1Token 合并merged_tokens, merged_offsets, host_merged_offsets merge_tokens_and_host_offsets( self.merger, tokens, input_row_offsets, draft_tokens, host_input_row_offsets, )将用户输入的tokens和预先准备的draft_tokens按 batch 合并同时合并对应的 row offsets。这一步由RaggedTokenMerger完成。阶段 2目标模型前向target_outputs self.target( merged_tokens, signal_buffers, target_mla_kv, target_indexer_kv, # 目标模型的双层 KV 缓存 return_n_logits, merged_offsets, host_merged_offsets, data_parallel_splits, batch_context_lengths, ep_inputs, )目标模型使用两级缓存target_mla_kvMLAMulti-head Latent Attention的 KV 缓存target_indexer_kvLightning Indexer 的 KV 缓存用于稀疏注意力返回(logits, offsets, hs_0..hs_{n-1})其中 hidden states 为 ALL_NORMALIZED 模式已在层内完成最终归一化。阶段 3拒绝采样与 Bitmaskeffective_bitmasks apply_overlap_bitmask( pinned_bitmask, wait_payload, device_bitmask_scratch, num_stepsdraft_tokens.shape[1], devicedevice0, ) num_accepted_draft_tokens, recovered, bonus, next_tokens accept_and_pick_next_tokens( self.acceptance_sampler, draft_tokens, logits, seedseed[0], temperaturetemperature, top_ktop_k, max_kmax_k, top_ptop_p, min_top_pmin_top_p, in_thinking_phasein_thinking_phase, token_bitmaskseffective_bitmasks, )如果enable_structured_outputTrue则通过pinned_bitmask、wait_payload和device_bitmask_scratch对特定 token 位置施加掩码约束结构化输出约束。accept_and_pick_next_tokens执行贪心拒绝采样返回num_accepted_draft_tokens被接受的草稿 token 数量recovered被拒绝位置恢复的 tokenbonusbonus token从目标分布中额外采样next_tokens下一轮的输入 token阶段 4Draft Step 0带 index_share 初始化self.draft.return_hidden_states ReturnHiddenStates.ALL self.draft.return_logits ReturnLogits.VARIABLE self.draft.emit_last_token_logits False # 抑制 lm_head 的大词汇表投影 draft_outputs self.draft( shifted_corrected, hidden_states, signal_buffers, draft_mla_kv, draft_indexer_kv, return_n_logits, merged_offsets_per_dev, host_merged_offsets, data_parallel_splits, batch_context_lengths, ep_inputs, prev_topk_indicesNone, reuse_prev_topkFalse, )Step 0 的特殊之处设置return_hidden_states ALL返回所有层的 hidden statesdraft 仅 1 层返回全部设置return_logits VARIABLE返回每个 token 位置的 logits用于计算 draft argmax设置emit_last_token_logits False抑制最后一个 token 的 lm_head 投影该位置不参与 draft 自回归prev_topk_indicesNone/reuse_prev_topkFalse在此步骤中计算lightning indexer 的 top-k并保存供后续复用Step 0 输出布局为(logits, offsets, hs[n], topk[n])阶段 5Draft 后续步骤复用 top-k在进入循环之前切换 draft 的返回模式self.draft.return_hidden_states ReturnHiddenStates.LAST_PER_DEVICE self.draft.return_logits ReturnLogits.LAST_TOKEN self.draft.emit_last_token_logits True同时切换 draft MLA 缓存的分发元数据draft_mla_kv [ replace(kv, max_prompt_lengthone, attention_dispatch_metadatakv.draft_attention_dispatch_metadata, mla_num_partitionskv.draft_mla_num_partitions, ) for kv in draft_mla_kv ]index_share_for_mtp_iteration核心优化在 Draft Step 0 中已经计算了 lightning indexer 的 top-k 选择结果step0_topk后续步骤通过gather_accepted_hidden_states收集已接受位置的 top-k 索引后在迭代中作为prev_topk_indices传入并设置reuse_prev_topkTrue跳过重复的 top-k 计算。reuse_topk gather_accepted_hidden_states( step0_topk, merged_offsetsmerged_offsets, merged_offsets_per_devmerged_offsets_per_dev, num_acceptednum_accepted_draft_tokens, num_draft_tokensdraft_tokens.shape[1], data_parallel_degree..., data_parallel_splits..., signal_buffers..., devicedevice0, split_prefixmtp_topk, ) for step in range(1, self.num_draft_steps): step_outputs self.draft( next_draft_tokens, draft_hs, signal_buffers, step_mla_kv, draft_indexer_kv, draft_return_n_logits, decode_offsets_per_dev, host_decode_offsets, data_parallel_splits, batch_context_lengths, ep_inputs, prev_topk_indicesreuse_topk, # 复用 step 0 的 top-k reuse_prev_topkTrue, # 跳过重复计算 split_prefixfmtp_draft_step{step}, )每个后续步骤的缓存mla_cache_lengths_per_dev会递增1而 indexer 缓存长度保持不变因为 indexer 仅在 step 0 参与。阶段 6输出组装if len(all_draft_tokens) 1: new_token ops.stack(all_draft_tokens, axis-1) else: new_token ops.unsqueeze(all_draft_tokens[0], -1) return (num_accepted_draft_tokens, next_tokens, new_token)最终返回三元组(被接受的草稿数量, 下一轮主 token, 拼接的新 draft tokens)。PipelineModel 编译与运行时UnifiedMTPGlm5_2Model该 PipelineModel 定义在 model.py 中继承自_UnifiedSpecDecodeModelMixin和Glm5_1Model。权重加载_load_state_dict从 checkpoint 解析target.*和draft.*前缀self._draft_state_dict { k[len(draft.):]: v for k, v in raw_state_dict.items() if k.startswith(draft.) } # 某些 checkpoint 共享 shared_head_norm 与 final norm if (shared_head_norm.weight not in self._draft_state_dict and target.norm.weight in raw_state_dict): self._draft_state_dict[shared_head_norm.weight] raw_state_dict[target.norm.weight]KV 缓存树_create_model_config构建嵌套的{target: {mla, indexer}, draft: {mla, indexer}}KV 缓存树。draft 的缓存仅有 1 层num_layers1且mla和indexer分开管理draft_kv MultiKVCacheParams.from_params({ mla: replace(target_mla_params, num_layers1), indexer: replace(target_indexer_params, num_layers1), }) self.kv_params MultiKVCacheParams.from_params( {target: target_kv, draft: draft_kv} )分布式专家并行EP_init_distributed_runtime处理专家并行初始化。对于 NVFP4 量化检查点其 MTP 层的 routed experts 以 bf16 精度存储无.weight_scale因此 draft 的 EP 分发精度必须从 NVFP4 提升到 bf16draft_moe_dispatches_bf16 not _subtree_quantized( self._draft_state_dict, .mlp.experts. ) if draft_moe_dispatches_bf16: ep_alloc_config replace(model_config.ep_config, dispatch_dtypeDType.bfloat16, dispatch_quant_configNone, fused_shared_expertmodel_config.n_shared_experts 1, )图编译_build_graph_for_compile方法实例化UnifiedMTPGlm5_2模型权重共享将 draft 的embed_tokens和lm_head别名为 target 的对应层strictFalse加载时跳过已共享的 key通过nn_model.input_types(kv_params)构建图输入类型签名用Graph上下文管理器构建完整的glm5_2_with_mtp_graph计算图从输入中解包四组 KV 缓存target_mla、target_indexer、draft_mla、draft_indexer提取采样超参数seed、temperature、top_k、max_k、top_p、min_top_p和in_thinking_phase标志调用nn_model(...)完成前向传播绑定图输出输入 batchingUnifiedMTPGlm5_2BatchProcessor定义在 batch_processor.py 中继承自DeepseekV3BatchProcessor。它在标准 batch 输入基础上扩展了draft_tokens字段初始化为None由 overlap pipeline 在实际执行时填充。输入结构定义在UnifiedMTPGlm5_2Inputs中继承UnifiedSpecDecodeInputs和DeepseekV3Inputs其buffers属性除了父类输入外额外包含in_thinking_phase标志位。权重适配器convert_with_mtp_state_dict定义在 weight_adapters.py 中负责将 HuggingFace safetensors 格式的 checkpoint 转换为 MAX 内部格式。权重 key 映射规则如下常规层通过DEEPSEEK_SAFETENSOR_MAP来自deepseekV3.weight_adapters将 HuggingFace key 转换为 MAX key丢弃 KV 缩放因子跳过以.k_scale或.v_scale结尾的 keyMAX 从独立配置路径读取 KV 缓存缩放MTP 层重映射MTP 层在 checkpoint 中位于layers.{num_hidden_layers}.索引处映射到draft.*。特定子模块映射如下Checkpoint Key 前缀MAX Key 路径说明layers.N.shared_head.norm.draft.shared_head_norm.共享头部归一化layers.N.enorm.draft.enorm.专家归一化layers.N.hnorm.draft.hnorm.头部归一化layers.N.eh_proj.draft.eh_proj.专家头部投影layers.N.其他draft.decoder_layer.解码器层self-attention / MLP非layers.N.前缀target.*目标模型权重权重共享embed_tokens和shared_head.headlm_head仅在 target 前缀下保存一份draft 通过模块别名共享state_dict()自动去重。Draft 配置创建_create_draft_config方法model.py从 draft 的 state_dict 推导配置验证 NextN 层存在decoder_layer.self_attn.kv_a_layernorm.weight基于 target 的基础配置创建DeepseekV3_2NextNConfig设置indexer_types []空调度让单层 MTP 的 indexer 保持满计算避免引用 target 的 78 层调度方案检测 NVFP4 量化的子树范围.mlp.experts.和.self_attn.将 MTP 层索引添加到正确的量化层集合中如果 draft 的 routed experts 未量化NVFP4 下 MTP 层为 bf16则修正 EP 分发配置为bfloat16使用方式要使用该架构在 MAX 中加载 GLM-5.2-FP8 模型并启用 MTP 推测解码需要在 Pipeline 配置中指定python -m max.pipelines.run \ --model-path zai-org/GLM-5.2-FP8 \ --speculative-config.num-speculative-tokens 3 \ --speculative-config.use-relaxed-acceptance-for-thinking \ --speculative-config.relaxed-topk 5 \ --speculative-config.relaxed-delta 0.1关键推测解码配置项说明源自 max/python/max/pipelines/speculative/config.py配置项类型默认值说明num_speculative_tokensint \| NoneNone每次生成的草稿 token 数量num_speculative_tokens_per_batch_sizelist[VerifyWidthRange] \| NoneNone按 batch size 分级的草稿数量调度synthetic_acceptance_ratefloat \| NoneNone合成接受率0.0~1.0用于绕过真实分布模拟use_relaxed_acceptance_for_thinkingboolFalsethinking 阶段是否启用宽松接受relaxed_topkint—宽松接受的 top-k 范围需 1relaxed_deltafloat—宽松接受的 delta 阈值0.0~1.0支持的量化编码可通过--dtype float8_e4m3fn默认、--dtype bfloat16或--dtype float4_e2m1fnx2指定。多 GPU 训练可通过--tensor-parallel-size N启用。该架构已默认注册到 MAX Pipeline 中通过架构名UnifiedMTPGlmMoeDsaForCausalLM自动匹配zai-org/GLM-5.2-FP8等模型 ID无需额外的手动注册步骤。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

MAX 中 GLM-5.2 (DeepSeek-V3.2 Sparse) 的统一 MTP 推测解码架构解析

相关推荐

Cloudflare Pipelines 配置全指南：从 Worker 绑定、Streams/Sinks 到不可变 SQL 管道的完整实战手册

go2rtc 上手实录：3 分钟把任意 RTSP 摄像头变成零延迟 WebRTC 直播源

深度探讨ROS在多机器人系统中的通信机制：核心解析与实践指南

更多资讯

Label Studio 标注一致性（Agreement）计算完整指南：Consensus 与 Pairwise 方法论、指标与权重配置

Task Master Auto-Implement：Claude Code 插件中的智能任务自动实现工作流

Actual 怎么套用官方自定义规则示例处理收款人与转账自动化

如何用 MSAL 在 Refine 中接入 Azure AD B2C 登录？

Beads 的 MCP Server 集成指南：在无 Shell 环境中让 Coding Agent 使用 bd 管理任务

量化交易数据源怎么选？四个常用工具一张表对比

今日更新

SpacetimeDB C++ 模块库实战：用 C++20 在数据库内构建 WebAssembly 模块

别踩雷！不是所有 AI 写作工具都靠谱，2026 高校认可工具精选

SSM+MySQL幼儿园管理系统源码解析：从部署到优化

本周精选

超人会飞不算本事：系统稳定依赖清晰规则与边界设计

超人VS蜘蛛侠：拆解超级IP的影响力与传播方法论

基于CNN的调制信号识别：MATLAB实现时频图分类实战

本月热门

基于YOLO的焊缝缺陷检测：从数据标注到工业部署全流程实践

分布式服务的决策记录

创业产品的首版范围

读完文章，想聊聊您的网站？