MLX-VLM Gemma 4 Assistant DrafterMTP深度解析原理、源码与性能调优【免费下载链接】mlx-vlmMLX-VLM is a package for inference and fine-tuning of Vision Language Models (VLMs) on your Mac using MLX.项目地址: https://gitcode.com/GitHub_Trending/ml/mlx-vlm本文基于 mlx-vlm 仓库中mlx_vlm/speculative/drafters/gemma4_assistant/目录的官方文档结合其源码实现、配置与测试完整讲解 Google Gemma 4Multi-Token PredictionMTP草稿模型assistant drafter在 MLX 上的移植原理、如何通过mlx_vlm.generate启用、如何以编程方式调用以及其 KV-cache 共享、质心路由稀疏 LM Head 与批处理轮转循环等底层机制并给出实测性能数据与使用注意事项。什么是 Gemma 4 MTP Assistant DrafterGemma 4 MTP Assistant Drafter 是一个为投机解码speculative decoding而训练的小型助手模型。它是一个仅 4 层的轻量 Transformer其职责不是独立生成答案而是在每一轮中起草draft多个候选 token然后由完整的 Gemma 4 目标模型target在单次前向传播中一次性验证这些候选 token。被目标模型接受的 token 继续前进被拒绝的 token 以及其后的所有候选 token 会被整体丢弃在温度 0贪心temp0下投机解码输出与不带草稿模型的目标模型输出逐字节一致byte-identical即质量无损。该实现是 Google Gemma 4 MTP 官方方案在 MLX 上的移植参考自官方 MTP 文档gemma4_assistant/README.md中的外部参考链接仓库内实现以官方方案为蓝本。草稿模型与目标模型的紧耦合与通用小模型充当草稿模型不同Gemma 4 assistant drafter深度耦合目标模型的内部结构具体体现在三个关键设计上KV-cache 共享KV-cache sharing草稿模型的每一个 Transformer 层都被标记为is_kv_shared_layerTrue直接读取目标模型最后一个全注意力层full-attention与最后一个滑动注意力层sliding-attention的 K/V。草稿模型没有自己的 KV cache见 gemma4_assistant.py 中make_cache()直接返回[]的实现它唯一的循环状态是目标模型的最后一层隐藏状态last hidden通过post_projection投影后传入下一轮。恒定位置的交叉注意力Cross-attention from constant position草稿模型的 query 在bonus token 的绝对位置上进行 RoPE 旋转并且在一个 block 内的所有起草步骤中保持该位置不变见draft_block中position_ids全程固定的实现。隐藏状态 token 拼接Hiddentoken concatenation草稿模型每一步的输入是concat([target_embed(last_token), last_hidden_state], dim-1) 形状为[B, 1, 2 * backbone_hidden_size]随后由pre_projection线性投影到草稿模型的隐藏维度。这与 HuggingFaceSinglePositionMultiTokenCandidateGenerator的输入顺序一致token embedding 在前、hidden state 在后可从 draft_block 的mx.concatenate([tok_embed, h_prev], axis-1)中直接看到。支持的目标模型配对Supported pairings仓库文档给出四组官方已验证的目标模型 草稿模型配对目标模型Target草稿模型DrafterLM Head 类型mlx-community/gemma-4-E2B-it-bf16mlx-community/gemma-4-E2B-it-assistant-bf16centroid稀疏mlx-community/gemma-4-E4B-it-bf16mlx-community/gemma-4-E4B-it-assistant-bf16centroid稀疏mlx-community/gemma-4-26B-A4B-it-bf16mlx-community/gemma-4-26B-A4B-it-assistant-bf16tied dense权重绑定稠密mlx-community/gemma-4-31B-it-bf16mlx-community/gemma-4-31B-it-assistant-bf16tied dense权重绑定稠密其中 E2B / E4B 草稿模型使用use_ordered_embeddingsTrue其 LM Head 是centroid-routed 稀疏 softmaxMaskedEmbedder而 26B-A4B / 31B 草稿模型则直接复用embed_tokens的转置作为 LM Headtie_word_embeddingsTrue。从源码看这两种路径在 bind() 中被自动选择有masked_embedding就走稀疏头否则走 tied embedding。源码文件结构mlx_vlm/speculative/drafters/gemma4_assistant/目录包含以下文件config.py —Gemma4AssistantConfigHF 兼容、扁平化的配置类gemma4_assistant.py —Gemma4AssistantDraftModelforward、bind、set_shared_kv、draft_block、sanitize等核心方法masked_embedder.py — E2B / E4B 使用的 centroid-routed 稀疏 LM Headmasks.py — 草稿模型前向所需的双向 full / SWAsliding-window attention掩码parity_check.py — 使用伪造目标fake target的冒烟测试脚本init.py — 导出Gemma4AssistantDraftModel、Gemma4AssistantConfig与TextConfig。配置类核心字段Gemma4AssistantConfig继承自BaseModelConfig默认值对齐官方gg-hf-am/gemma-4-26B-A4B-it-assistant检查点关键字段如下字段默认值含义model_typegemma4_assistantHF 模型类型标识用于 drafter 自动发现backbone_hidden_size1536目标模型骨干隐藏维度决定拼接输入宽度use_ordered_embeddingsFalse是否使用 centroid-routed 稀疏 LM HeadE2B / E4B 为Truenum_centroids2048稀疏头的质心cluster数量centroid_intermediate_top_k32每步保留的 top-K 质心数量tie_word_embeddingsTrueLM Head 是否与输入 embedding 权重绑定block_size4草稿块大小每轮候选 token 数target_layer_ids[]MTP 未使用草稿模型消费共享 K/V仅为与 DFlash 的 round-loop API 对齐而保留text_configNone嵌套的 Gemma 4 文本配置值得注意的是__post_init__中的一条自动修正逻辑config.py当text_config.num_kv_shared_layers未设置为 0时会自动将其设置为num_hidden_layers即assistant 模型默认在所有层间共享 K/V。稀疏 LM HeadMaskedEmbedder的工作原理对于词表高达262144的 E2B / E4B 目标模型如果让一个隐藏维度仅 256 的小草稿模型计算hidden embed.T的全词表 logits代价过高。MaskedEmbedder的方案是用一个centroids线性层nn.Linear(hidden_size, num_centroids)对隐藏状态打分得到2048 个 token 簇cluster的分数用mx.argpartition选取分数最高的top-K默认 32个簇每个簇对应vocab_size // num_centroids 128个连续排列的规范 token ID由静态缓冲区token_ordering保存这是一个不可学习的检查点缓冲区在_freeze_static_buffers中被冻结只对选中的32 × 128 4096个 token 稠密计算 logitshidden E.T将选中位置的 logits 通过mx.put_along_axis散布回[B, L, 262144]的全词表张量未选中的位置填充min(selected_logits) - 1哨兵值使它们在 argmax 或采样竞争中必然落败。贪心路径还有一个专门优化argmax()方法masked_embedder.py只对选中的 ~4096 个 logits 取 argmax完全不物化全词表 logits直接返回对应簇内的规范 token ID。如何使用命令行与编程接口自动发现机制草稿模型通过 HFmodel_type gemma4_assistant被自动发现。在 drafters/init.py 的DRAFTER_KIND_BY_MODEL_TYPE映射中gemma4_assistant被显式映射到mtpround-loop 类型resolve_drafter_kind会在用户忘记传--draft-kind时自动检测并覆盖为mtp并给出 warning避免在draft_block深处报出不透明错误。如果用户显式传入错误的 kind同样会被强制纠正为mtpdrafters/init.py。此外validate_drafter_compatibility会校验backbone_hidden_size与目标模型的 hidden size 是否一致不一致时直接抛出ValueError提示请使用同一目标家族与尺寸的草稿检查点。命令行用法只需在mlx_vlm.generate中传入--draft-model与--draft-kind mtpuv run python -m mlx_vlm.generate \ --model mlx-community/gemma-4-31B-it-bf16 \ --draft-model mlx-community/gemma-4-31B-it-assistant-bf16 \ --draft-kind mtp \ --draft-block-size 4 \ --prompt Explain speculative decoding in 3 sentences. \ --max-tokens 256 --temp 0关键参数说明--draft-block-size每轮投机起草的 token 数量Google 官方称之为num_assistant_tokens。注意块内第一个 token 是最近一次被接受的 bonus token因此草稿模型每轮实际新生成的候选数为block_size - 1这从 draft_block 的for _ in range(block_size - 1)循环可以直接印证--temp 0贪心模式输出与无草稿基线逐字节一致。编程方式调用from mlx_vlm.utils import load from mlx_vlm.speculative.drafters import load_drafter from mlx_vlm.generate import generate_step model, processor load(mlx-community/gemma-4-31B-it-bf16) drafter load_drafter(mlx-community/gemma-4-31B-it-assistant-bf16, kindmtp) for tok, _ in generate_step( input_ids, model, None, None, max_tokens256, draft_modeldrafter, draft_kindmtp, draft_block_size4, ): ...这里load_drafter返回(model, resolved_kind)二元组drafters/init.pyresolved_kind是自动检测后的最终 kind如果直接调用load_drafter(..., kindmtp)务必把draft_kindmtp同步传给generate_step否则 MTP round-loop 不会运行。服务端Server用法在服务端场景可通过环境变量MLX_VLM_DRAFT_KINDmtp指定draft_block中的报错信息明确提示了这一通道。如果误用了 DFlash round-loop例如忘记指定 kinddraft_block会抛出带指引信息的RuntimeError提示改用--draft-kind mtp或MLX_VLM_DRAFT_KINDmtp。MTP Round-Loop目标如何验证候选MTP 的验证轮转循环位于 mtp.py其中_mtp_rounds_batchmtp.py是批处理版本B ≥ 1每轮先由草稿模型draft_block自回归生成block_size - 1个候选 token目标模型对全部候选执行单次前向验证确定接受长度每行row的positions记录其有效目标 KV 长度每轮按accepted 1推进草稿模型每轮通过set_shared_kv重新绑定shared_kv_states批量路径下还会把共享 K/V归一化回非批处理的 prefix-valid 布局normalize_batched_shared_kv_states见 masks.py。两种前向路径在 gemma4_assistant.py 的draft_block中贪心路径greedy当greedyTrue且使用MaskedEmbedder时直接调用masked_embedding.argmax跳过全词表 logits 物化只做稀疏 argmax采样路径调用self(...)得到全词表或稀疏散布后的logits交给sampler(logits)采样。掩码生成masks.py 中的make_drafter_masks按层类型生成掩码全注意力层full_attentionbidirectional_full_mask在无 padding 的非批处理场景下退化为NoneSDPA 直接处理滑动注意力层sliding_attentionbidirectional_swa_mask对每个 query 位置q只允许关注k ∈ (q - window, q window)的 KV 位置当kv_len sliding_window时同样短路返回None。正如 README Caveats 中说明的由于RotatingKVCache只会产生kv_len sliding_window的短 KV 场景长提示long-prompt下的掩码路径目前实际上是死代码dead code该限制值得实现时注意。冒烟测试与自动校验parity_check.py 提供了一个不需要真实目标模型的冒烟测试用随机权重伪造目标 embedding_FakeTargetEmbed构建全注意力 滑动注意力两套假共享 K/V然后依次执行单步 forward 与多步draft_block验证 logits 与 token 形状是否正确。uv run python -m mlx_vlm.speculative.drafters.gemma4_assistant.parity_check \ --drafter gg-hf-am/gemma-4-26B-A4B-it-assistant此外仓库测试 test_speculative.py 中覆盖了gemma4_assistant的 kind 自动检测test_gemma4_assistant_overrides_dflash_to_mtp、test_kind_none_autodetects_mtp_for_gemma4_assistant等test_gemma4_assistant_masks_static.py 则对masks.py做了静态校验服务端测试 test_server.py 也验证了model_typegemma4_assistant场景。实测性能仓库文档给出的性能数据测量环境为Apple SiliconM3 Max96GB RAM17-token promptmax_tokens64–96贪心模式temp0输出与无草稿基线逐字节一致。各目标, 批大小组合下的最优block_size目标B最优 bs总 tok/s相对无草稿加速比26B-A4B4385.53.94×26B-A4B83165.11.55×31B4317.12.29×31B8221.41.41×E4B4462.11.56×E4B82115.91.07×E4B16——草稿模型更慢≤1.0×结论与选型建议草稿模型在大而慢的目标26B-A4B、31B上收益最明显因为此时目标前向时间占主导而在较小的 E4B 目标上目标前向本身已经很便宜高批大小下草稿模型的单步开销会超过它买来的加速甚至出现回退。实测数据表明 26B-A4B 在批大小 4 时可获得接近 4 倍的吞吐提升。使用注意事项Caveats采样Sampling贪心temp 0已验证逐字节一致随机采样也可用但接受率会下降因为草稿模型与目标模型的采样分布会发生分歧。多模态提示Multimodal prompts图像 / 音频 prefill 仍由目标模型原样处理投机解码只作用于文本解码尾部因此多模态可用但草稿模型只会看到文本 token。滑动窗口掩码Sliding-window masksmasks.py中的双向 SWA 掩码在kv_len sliding_window时短路为None而这是RotatingKVCache唯一会产生的场景长提示的掩码路径当前实际是死代码。批处理生成Batched generation连续批处理支持在_mtp_rounds_batchmtp.py。对于 KV cache 未实现.filter()的目标已结束的行会保留在批次中只是停止输出吞吐不会随已退出行数而收缩——这是实现批处理投机解码时需要理解的行为。小结Gemma 4 MTP Assistant Drafter 通过 KV-cache 共享、恒定位置交叉注意力与 hiddentoken 拼接三个设计将一个小型 4 层模型无缝嵌入 Gemma 4 目标的解码流程在苹果芯片MLX上实现了贪心输出无损的投机解码加速。本文从 官方文档 出发结合 gemma4_assistant.py、masked_embedder.py、masks.py、config.py 与 mtp.py 等源码完整还原了其配置、调用方式与底层机制。对于拥有 M 系列芯片且需要推理大参数 Gemma 4 模型的开发者这是一个可直接落地的加速方案——只需一行--draft-kind mtp再按目标模型尺寸选择合适的--draft-block-size建议从 3–4 起步并实测调优。【免费下载链接】mlx-vlmMLX-VLM is a package for inference and fine-tuning of Vision Language Models (VLMs) on your Mac using MLX.项目地址: https://gitcode.com/GitHub_Trending/ml/mlx-vlm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
MLX-VLM Gemma 4 Assistant Drafter（MTP）深度解析：原理、源码与性能调优
MLX-VLM Gemma 4 Assistant DrafterMTP深度解析原理、源码与性能调优【免费下载链接】mlx-vlmMLX-VLM is a package for inference and fine-tuning of Vision Language Models (VLMs) on your Mac using MLX.项目地址: https://gitcode.com/GitHub_Trending/ml/mlx-vlm本文基于 mlx-vlm 仓库中mlx_vlm/speculative/drafters/gemma4_assistant/目录的官方文档结合其源码实现、配置与测试完整讲解 Google Gemma 4Multi-Token PredictionMTP草稿模型assistant drafter在 MLX 上的移植原理、如何通过mlx_vlm.generate启用、如何以编程方式调用以及其 KV-cache 共享、质心路由稀疏 LM Head 与批处理轮转循环等底层机制并给出实测性能数据与使用注意事项。什么是 Gemma 4 MTP Assistant DrafterGemma 4 MTP Assistant Drafter 是一个为投机解码speculative decoding而训练的小型助手模型。它是一个仅 4 层的轻量 Transformer其职责不是独立生成答案而是在每一轮中起草draft多个候选 token然后由完整的 Gemma 4 目标模型target在单次前向传播中一次性验证这些候选 token。被目标模型接受的 token 继续前进被拒绝的 token 以及其后的所有候选 token 会被整体丢弃在温度 0贪心temp0下投机解码输出与不带草稿模型的目标模型输出逐字节一致byte-identical即质量无损。该实现是 Google Gemma 4 MTP 官方方案在 MLX 上的移植参考自官方 MTP 文档gemma4_assistant/README.md中的外部参考链接仓库内实现以官方方案为蓝本。草稿模型与目标模型的紧耦合与通用小模型充当草稿模型不同Gemma 4 assistant drafter深度耦合目标模型的内部结构具体体现在三个关键设计上KV-cache 共享KV-cache sharing草稿模型的每一个 Transformer 层都被标记为is_kv_shared_layerTrue直接读取目标模型最后一个全注意力层full-attention与最后一个滑动注意力层sliding-attention的 K/V。草稿模型没有自己的 KV cache见 gemma4_assistant.py 中make_cache()直接返回[]的实现它唯一的循环状态是目标模型的最后一层隐藏状态last hidden通过post_projection投影后传入下一轮。恒定位置的交叉注意力Cross-attention from constant position草稿模型的 query 在bonus token 的绝对位置上进行 RoPE 旋转并且在一个 block 内的所有起草步骤中保持该位置不变见draft_block中position_ids全程固定的实现。隐藏状态 token 拼接Hiddentoken concatenation草稿模型每一步的输入是concat([target_embed(last_token), last_hidden_state], dim-1) 形状为[B, 1, 2 * backbone_hidden_size]随后由pre_projection线性投影到草稿模型的隐藏维度。这与 HuggingFaceSinglePositionMultiTokenCandidateGenerator的输入顺序一致token embedding 在前、hidden state 在后可从 draft_block 的mx.concatenate([tok_embed, h_prev], axis-1)中直接看到。支持的目标模型配对Supported pairings仓库文档给出四组官方已验证的目标模型 草稿模型配对目标模型Target草稿模型DrafterLM Head 类型mlx-community/gemma-4-E2B-it-bf16mlx-community/gemma-4-E2B-it-assistant-bf16centroid稀疏mlx-community/gemma-4-E4B-it-bf16mlx-community/gemma-4-E4B-it-assistant-bf16centroid稀疏mlx-community/gemma-4-26B-A4B-it-bf16mlx-community/gemma-4-26B-A4B-it-assistant-bf16tied dense权重绑定稠密mlx-community/gemma-4-31B-it-bf16mlx-community/gemma-4-31B-it-assistant-bf16tied dense权重绑定稠密其中 E2B / E4B 草稿模型使用use_ordered_embeddingsTrue其 LM Head 是centroid-routed 稀疏 softmaxMaskedEmbedder而 26B-A4B / 31B 草稿模型则直接复用embed_tokens的转置作为 LM Headtie_word_embeddingsTrue。从源码看这两种路径在 bind() 中被自动选择有masked_embedding就走稀疏头否则走 tied embedding。源码文件结构mlx_vlm/speculative/drafters/gemma4_assistant/目录包含以下文件config.py —Gemma4AssistantConfigHF 兼容、扁平化的配置类gemma4_assistant.py —Gemma4AssistantDraftModelforward、bind、set_shared_kv、draft_block、sanitize等核心方法masked_embedder.py — E2B / E4B 使用的 centroid-routed 稀疏 LM Headmasks.py — 草稿模型前向所需的双向 full / SWAsliding-window attention掩码parity_check.py — 使用伪造目标fake target的冒烟测试脚本init.py — 导出Gemma4AssistantDraftModel、Gemma4AssistantConfig与TextConfig。配置类核心字段Gemma4AssistantConfig继承自BaseModelConfig默认值对齐官方gg-hf-am/gemma-4-26B-A4B-it-assistant检查点关键字段如下字段默认值含义model_typegemma4_assistantHF 模型类型标识用于 drafter 自动发现backbone_hidden_size1536目标模型骨干隐藏维度决定拼接输入宽度use_ordered_embeddingsFalse是否使用 centroid-routed 稀疏 LM HeadE2B / E4B 为Truenum_centroids2048稀疏头的质心cluster数量centroid_intermediate_top_k32每步保留的 top-K 质心数量tie_word_embeddingsTrueLM Head 是否与输入 embedding 权重绑定block_size4草稿块大小每轮候选 token 数target_layer_ids[]MTP 未使用草稿模型消费共享 K/V仅为与 DFlash 的 round-loop API 对齐而保留text_configNone嵌套的 Gemma 4 文本配置值得注意的是__post_init__中的一条自动修正逻辑config.py当text_config.num_kv_shared_layers未设置为 0时会自动将其设置为num_hidden_layers即assistant 模型默认在所有层间共享 K/V。稀疏 LM HeadMaskedEmbedder的工作原理对于词表高达262144的 E2B / E4B 目标模型如果让一个隐藏维度仅 256 的小草稿模型计算hidden embed.T的全词表 logits代价过高。MaskedEmbedder的方案是用一个centroids线性层nn.Linear(hidden_size, num_centroids)对隐藏状态打分得到2048 个 token 簇cluster的分数用mx.argpartition选取分数最高的top-K默认 32个簇每个簇对应vocab_size // num_centroids 128个连续排列的规范 token ID由静态缓冲区token_ordering保存这是一个不可学习的检查点缓冲区在_freeze_static_buffers中被冻结只对选中的32 × 128 4096个 token 稠密计算 logitshidden E.T将选中位置的 logits 通过mx.put_along_axis散布回[B, L, 262144]的全词表张量未选中的位置填充min(selected_logits) - 1哨兵值使它们在 argmax 或采样竞争中必然落败。贪心路径还有一个专门优化argmax()方法masked_embedder.py只对选中的 ~4096 个 logits 取 argmax完全不物化全词表 logits直接返回对应簇内的规范 token ID。如何使用命令行与编程接口自动发现机制草稿模型通过 HFmodel_type gemma4_assistant被自动发现。在 drafters/init.py 的DRAFTER_KIND_BY_MODEL_TYPE映射中gemma4_assistant被显式映射到mtpround-loop 类型resolve_drafter_kind会在用户忘记传--draft-kind时自动检测并覆盖为mtp并给出 warning避免在draft_block深处报出不透明错误。如果用户显式传入错误的 kind同样会被强制纠正为mtpdrafters/init.py。此外validate_drafter_compatibility会校验backbone_hidden_size与目标模型的 hidden size 是否一致不一致时直接抛出ValueError提示请使用同一目标家族与尺寸的草稿检查点。命令行用法只需在mlx_vlm.generate中传入--draft-model与--draft-kind mtpuv run python -m mlx_vlm.generate \ --model mlx-community/gemma-4-31B-it-bf16 \ --draft-model mlx-community/gemma-4-31B-it-assistant-bf16 \ --draft-kind mtp \ --draft-block-size 4 \ --prompt Explain speculative decoding in 3 sentences. \ --max-tokens 256 --temp 0关键参数说明--draft-block-size每轮投机起草的 token 数量Google 官方称之为num_assistant_tokens。注意块内第一个 token 是最近一次被接受的 bonus token因此草稿模型每轮实际新生成的候选数为block_size - 1这从 draft_block 的for _ in range(block_size - 1)循环可以直接印证--temp 0贪心模式输出与无草稿基线逐字节一致。编程方式调用from mlx_vlm.utils import load from mlx_vlm.speculative.drafters import load_drafter from mlx_vlm.generate import generate_step model, processor load(mlx-community/gemma-4-31B-it-bf16) drafter load_drafter(mlx-community/gemma-4-31B-it-assistant-bf16, kindmtp) for tok, _ in generate_step( input_ids, model, None, None, max_tokens256, draft_modeldrafter, draft_kindmtp, draft_block_size4, ): ...这里load_drafter返回(model, resolved_kind)二元组drafters/init.pyresolved_kind是自动检测后的最终 kind如果直接调用load_drafter(..., kindmtp)务必把draft_kindmtp同步传给generate_step否则 MTP round-loop 不会运行。服务端Server用法在服务端场景可通过环境变量MLX_VLM_DRAFT_KINDmtp指定draft_block中的报错信息明确提示了这一通道。如果误用了 DFlash round-loop例如忘记指定 kinddraft_block会抛出带指引信息的RuntimeError提示改用--draft-kind mtp或MLX_VLM_DRAFT_KINDmtp。MTP Round-Loop目标如何验证候选MTP 的验证轮转循环位于 mtp.py其中_mtp_rounds_batchmtp.py是批处理版本B ≥ 1每轮先由草稿模型draft_block自回归生成block_size - 1个候选 token目标模型对全部候选执行单次前向验证确定接受长度每行row的positions记录其有效目标 KV 长度每轮按accepted 1推进草稿模型每轮通过set_shared_kv重新绑定shared_kv_states批量路径下还会把共享 K/V归一化回非批处理的 prefix-valid 布局normalize_batched_shared_kv_states见 masks.py。两种前向路径在 gemma4_assistant.py 的draft_block中贪心路径greedy当greedyTrue且使用MaskedEmbedder时直接调用masked_embedding.argmax跳过全词表 logits 物化只做稀疏 argmax采样路径调用self(...)得到全词表或稀疏散布后的logits交给sampler(logits)采样。掩码生成masks.py 中的make_drafter_masks按层类型生成掩码全注意力层full_attentionbidirectional_full_mask在无 padding 的非批处理场景下退化为NoneSDPA 直接处理滑动注意力层sliding_attentionbidirectional_swa_mask对每个 query 位置q只允许关注k ∈ (q - window, q window)的 KV 位置当kv_len sliding_window时同样短路返回None。正如 README Caveats 中说明的由于RotatingKVCache只会产生kv_len sliding_window的短 KV 场景长提示long-prompt下的掩码路径目前实际上是死代码dead code该限制值得实现时注意。冒烟测试与自动校验parity_check.py 提供了一个不需要真实目标模型的冒烟测试用随机权重伪造目标 embedding_FakeTargetEmbed构建全注意力 滑动注意力两套假共享 K/V然后依次执行单步 forward 与多步draft_block验证 logits 与 token 形状是否正确。uv run python -m mlx_vlm.speculative.drafters.gemma4_assistant.parity_check \ --drafter gg-hf-am/gemma-4-26B-A4B-it-assistant此外仓库测试 test_speculative.py 中覆盖了gemma4_assistant的 kind 自动检测test_gemma4_assistant_overrides_dflash_to_mtp、test_kind_none_autodetects_mtp_for_gemma4_assistant等test_gemma4_assistant_masks_static.py 则对masks.py做了静态校验服务端测试 test_server.py 也验证了model_typegemma4_assistant场景。实测性能仓库文档给出的性能数据测量环境为Apple SiliconM3 Max96GB RAM17-token promptmax_tokens64–96贪心模式temp0输出与无草稿基线逐字节一致。各目标, 批大小组合下的最优block_size目标B最优 bs总 tok/s相对无草稿加速比26B-A4B4385.53.94×26B-A4B83165.11.55×31B4317.12.29×31B8221.41.41×E4B4462.11.56×E4B82115.91.07×E4B16——草稿模型更慢≤1.0×结论与选型建议草稿模型在大而慢的目标26B-A4B、31B上收益最明显因为此时目标前向时间占主导而在较小的 E4B 目标上目标前向本身已经很便宜高批大小下草稿模型的单步开销会超过它买来的加速甚至出现回退。实测数据表明 26B-A4B 在批大小 4 时可获得接近 4 倍的吞吐提升。使用注意事项Caveats采样Sampling贪心temp 0已验证逐字节一致随机采样也可用但接受率会下降因为草稿模型与目标模型的采样分布会发生分歧。多模态提示Multimodal prompts图像 / 音频 prefill 仍由目标模型原样处理投机解码只作用于文本解码尾部因此多模态可用但草稿模型只会看到文本 token。滑动窗口掩码Sliding-window masksmasks.py中的双向 SWA 掩码在kv_len sliding_window时短路为None而这是RotatingKVCache唯一会产生的场景长提示的掩码路径当前实际是死代码。批处理生成Batched generation连续批处理支持在_mtp_rounds_batchmtp.py。对于 KV cache 未实现.filter()的目标已结束的行会保留在批次中只是停止输出吞吐不会随已退出行数而收缩——这是实现批处理投机解码时需要理解的行为。小结Gemma 4 MTP Assistant Drafter 通过 KV-cache 共享、恒定位置交叉注意力与 hiddentoken 拼接三个设计将一个小型 4 层模型无缝嵌入 Gemma 4 目标的解码流程在苹果芯片MLX上实现了贪心输出无损的投机解码加速。本文从 官方文档 出发结合 gemma4_assistant.py、masked_embedder.py、masks.py、config.py 与 mtp.py 等源码完整还原了其配置、调用方式与底层机制。对于拥有 M 系列芯片且需要推理大参数 Gemma 4 模型的开发者这是一个可直接落地的加速方案——只需一行--draft-kind mtp再按目标模型尺寸选择合适的--draft-block-size建议从 3–4 起步并实测调优。【免费下载链接】mlx-vlmMLX-VLM is a package for inference and fine-tuning of Vision Language Models (VLMs) on your Mac using MLX.项目地址: https://gitcode.com/GitHub_Trending/ml/mlx-vlm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考