执行摘要

• 一句话：新增 Gemma4 MTP 推测解码，含 centroids masking 优化
• 推荐动作：值得精读。本 PR 展示了一条完整的 speculative decoding 集成路径：从模型定义、提议器封装到配置注册和 runner 集成。其中 centroids masking 和跨模型 KV 共享的设计尤其值得借鉴。建议重点关注 gemma4_mtp.py 的 _select_and_score 和 gemma4.py 的 build_per_group_and_layer_attn_metadata。

功能与动机

Gemma4 家族包含多种规模的 checkpoints，官方提供了对应的 assistant 模型用于推测解码。本 PR 旨在将这些 assistant 模型接入 vLLM 的 speculative decoding 框架，利用其轻量级 Q-only 注意力和 centroids 词汇表掩码技术，在大幅减少每步计算的同时保持高接受率，从而显著提升 gemma4 模型的推理吞吐量。

实现拆解

1. 定义 Gemma4MTP 模型（vllm/model_executor/models/gemma4_mtp.py）：
   • 实现 Gemma4MTPMaskedEmbedder：将 hidden states 投影到 centroids 分数，选出 top-k centroids，仅对属于这些 centroids 的 token 计算 logits，实现稀疏 lm_head。
   • 实现 Gemma4MTPAttention：Q-only 注意力层，无 K/V 投影和 norm，运行时从目标模型 KV 缓存读取键值，通过 kv_sharing_target_layer_name 完成跨模型 KV 共享。
   • 实现 Gemma4MultiTokenPredictor：组合 embed_tokens → pre_projection → decoder layers → norm → post_projection → logits_processor。支持可选择的全量或 masked lm_head，并提供 get_top_tokens 供 CUDA graph 加速。
2. 实现 Gemma4 提议器（vllm/v1/spec_decode/gemma4.py）：
   • Gemma4Proposer 继承 SpecDecodeBaseProposer，重写关键方法：
     • model_returns_tuple 返回 True，使 forward 输出 (draft_hidden, backbone_hidden)。
     • build_per_group_and_layer_attn_metadata 为每个 attention group 替换正确的 block table，适配 Gemma4 的滑动/全注意力多组 KV cache。
     • _greedy_sample 支持 centroids CUDA graph 回放，快速选择 token。
     • _setup_centroids_cuda_graphs 为常用 batch size 捕获 CUDA graphs。
     • _create_draft_vllm_config 强制保留 TRITON_ATTN 后端，保证 KV 共享兼容。
3. 增强基础提议器（vllm/v1/spec_decode/llm_base_proposer.py）：
   • 添加 constant_draft_positions 标志，使所有 draft 步骤重用同一位置（用于 Q-only 注意力）。
   • 重构 propose 方法：提取 _update_positions_dependent_metadata，支持 constant positions 时仅写入位置一次。
4. 配置与注册（vllm/config/speculative.py、vllm/model_executor/models/registry.py、vllm/transformers_utils/model_arch_config_convertor.py）：
   • 在 hf_config_override 中将 gemma4_assistant 映射为 gemma4_mtp，设置 n_predict=1、num_kv_shared_layers=0。
   • 添加 use_gemma4_mtp() 方法判断是否使用 Gemma4 MTP。
   • 注册 Gemma4MTPModelArchConfigConvertor，确保 get_hidden_size 返回 backbone 的 hidden size。
5. 集成到 GPU 模型运行器（vllm/v1/worker/gpu_model_runner.py）：
   • 导入 Gemma4Proposer，在初始化时根据配置创建实例。
   • 在 _build_attn_group_metadata 中，为 Gemma4Proposer 透传 per_group_block_table。
   • 在多处类型检查中添加 Gemma4Proposer 分支。
6. 测试与验证（tests/v1/e2e/spec_decode/test_spec_decode.py、tests/models/registry.py）：
   • 在 spec decode 端到端测试中添加 MTP 正确性参数化用例，阈值设为 0.5。
   • 在模型注册表中添加 Gemma4MTPModel 示例信息，包含 min_transformers_version。

关键文件：

• vllm/model_executor/models/gemma4_mtp.py（模块 模型定义；类别 source；类型 data-contract；符号 Gemma4MTPMaskedEmbedder, init, _select_and_score, forward）: 核心模型实现：包含 centroids masking 的稀疏 logits 计算、Q-only attention 层、多 token 预测器以及权重加载逻辑。是 Gemma4 MTP 的基础。
• vllm/v1/spec_decode/gemma4.py（模块 推测解码；类别 source；类型 core-logic；符号 Gemma4Proposer, init, set_per_group_block_table, model_returns_tuple）: 提议器实现：处理 KV 共享、多组 attention metadata、centroids CUDA graphs 捕获以及 draft 配置覆盖。是连接模型和推理引擎的关键。
• vllm/v1/spec_decode/llm_base_proposer.py（模块 推测解码；类别 source；类型 core-logic；符号 _update_positions_dependent_metadata）: 基础提议器修改：添加 constant_draft_positions 支持；重构 propose 方法提取 _update_positions_dependent_metadata；确保所有 draft 模型都能受益。
• vllm/config/speculative.py（模块 配置；类别 source；类型 core-logic；符号 use_gemma4_mtp）: 配置注册：识别 gemma4_assistant 模型类型并映射为 gemma4_mtp；添加 use_gemma4_mtp 方法。
• vllm/v1/worker/gpu_model_runner.py（模块 模型运行器；类别 source；类型 core-logic）: 集成点：导入并实例化 Gemma4Proposer；传递 per-group block tables；在多个类型联合中添加 Gemma4Proposer 以便正确调度。
• vllm/transformers_utils/model_arch_config_convertor.py（模块 转换器；类别 source；类型 data-contract；符号 Gemma4MTPModelArchConfigConvertor, get_hidden_size, get_num_hidden_layers）: 架构配置转换：Gemma4MTPModelArchConfigConvertor 确保 hidden_size 返回 backbone 值。
• tests/v1/e2e/spec_decode/test_spec_decode.py（模块 测试；类别 test；类型 test-coverage）: 端到端测试：验证 MTP 正确性和接受率。
• vllm/model_executor/models/registry.py（模块 模型注册；类别 source；类型 data-contract）: 模型注册：将 Gemma4MTPModel 加入可用模型列表。
• tests/models/registry.py（模块 测试；类别 test；类型 test-coverage）: 测试模型注册信息，添加 Gemma4MTPModel 示例。

关键符号：Gemma4MTPMaskedEmbedder._select_and_score, Gemma4MTPMaskedEmbedder.forward, Gemma4MTPMaskedEmbedder.get_top_tokens, Gemma4MultiTokenPredictor.init, Gemma4MultiTokenPredictor.forward, Gemma4Proposer.init, Gemma4Proposer.set_per_group_block_table, Gemma4Proposer.model_returns_tuple, Gemma4Proposer.build_per_group_and_layer_attn_metadata, Gemma4Proposer._greedy_sample, Gemma4Proposer._setup_centroids_cuda_graphs, Gemma4Proposer._create_draft_vllm_config, SpecDecodeBaseProposer._update_positions_dependent_metadata, SpeculativeConfig.use_gemma4_mtp, Gemma4MTPModelArchConfigConvertor.get_hidden_size, Gemma4MTPModelArchConfigConvertor.get_num_hidden_layers

关键源码片段

vllm/model_executor/models/gemma4_mtp.py

核心模型实现：包含 centroids masking 的稀疏 logits 计算、Q-only attention 层、多 token 预测器以及权重加载逻辑。是 Gemma4 MTP 的基础。

# vllm/model_executor/models/gemma4_mtp.py
# _select_and_score: 实现 centroid 掩码的核心逻辑
​
class Gemma4MTPMaskedEmbedder(nn.Module):
    """稀疏 logit 计算：通过 centroid 词汇表掩码减少计算量。
​
    与传统全词汇表 logits 不同，该方法将 hidden states 投影到
    centroid 分数，选择 top-k centroids，然后只对属于这些
    centroids 的 token 计算 logits，大幅降低 lm_head 开销。
    """
​
    def _select_and_score(
        self,
        hidden_states: torch.Tensor, # shape: (num_tokens, hidden_size)
        lm_head_weight: torch.Tensor, # shape: (vocab_size, hidden_size)
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Centroid 选择 + 稀疏点积。
​
        Returns:
            logits: (num_tokens, num_selected) 稀疏 logits。
            indices: (num_tokens, num_selected) 对应的词汇表索引。
        """
        num_tokens = hidden_states.shape[0]
​
        # Step 1: 将 hidden states 投影到 centroid 分数，取 top-k
        _, top_k_indices = torch.topk(
            self.centroids(hidden_states), # => (num_tokens, num_centroids)
            k=self.centroid_intermediate_top_k, # 例如 32
            dim=-1,
        )
​
        # Step 2: 根据 token_ordering 查找每个 centroid 对应的 token 集合
        clusters = self.token_ordering.view(
            self.num_centroids,
            self.vocab_size_per_centroid, # vocab_size // num_centroids
        )
        selected = clusters[top_k_indices] # (num_tokens, top_k, vocab_per_centroid)
​
        # Step 3: 仅对选中的 token 计算 logits（稀疏点积）
        embeddings = lm_head_weight[selected.reshape(-1)].view(
            num_tokens, self.num_selected, self.hidden_size,
        )
        logits = torch.einsum("td,tsd->ts", hidden_states, embeddings)
        return logits, selected.view(num_tokens, -1)

vllm/v1/spec_decode/gemma4.py

提议器实现：处理 KV 共享、多组 attention metadata、centroids CUDA graphs 捕获以及 draft 配置覆盖。是连接模型和推理引擎的关键。

# vllm/v1/spec_decode/gemma4.py
# _greedy_sample: 使用 centroids CUDA graphs 加速贪婪采样
​
class Gemma4Proposer(SpecDecodeBaseProposer):
    def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """使用预捕获的 CUDA graphs 或 fallback 进行贪婪采样。"""
        if self._centroids_sizes:
            T = hidden_states.shape[0]
            # 查找第一个 >= T 的缓存 size（已排序）
            for size in self._centroids_sizes:
                if size >= T:
                    # 复制输入到静态缓冲区
                    self._centroids_inputs[size][:T].copy_(hidden_states)
                    # 回放 CUDA graph
                    self._centroids_graphs[size].replay()
                    # 取前 T 行并 clone 返回（避免后续修改影响 graph 输出）
                    return self._centroids_outputs[size][:T].clone()
            # 如果没有匹配的 size，直接调用模型的 get_top_tokens
            return self.model.get_top_tokens(hidden_states)
        # 如果未启用 centroids，回退到基类方法
        return super()._greedy_sample(hidden_states)
​
    def _setup_centroids_cuda_graphs(self) -> None:
        """为常见 batch size 捕获 centroids CUDA graphs。"""
        masked_emb = self.model.masked_embedding
        lm_head_weight = self.model._get_full_lm_head_weight()
​
        for size in [1, 2, 4, 8, 16, 32, 64]:
            static_input = torch.zeros(
                size, masked_emb.hidden_size, device=self.device
            )
            static_output = torch.zeros(
                size, dtype=torch.long, device=self.device
            )
            # 捕获 graph：给定 hidden_states，输出 argmax token id
            self._centroids_graphs[size] = torch.cuda.CUDAGraph()
            with torch.cuda.graph(self._centroids_graphs[size]):
                logits, indices = masked_emb._select_and_score(
                    static_input, lm_head_weight
                )
                static_output.copy_(logits.argmax(dim=-1))
            self._centroids_inputs[size] = static_input
            self._centroids_outputs[size] = static_output
        self._centroids_sizes = sorted(self._centroids_inputs.keys())

评论区精华

• benchislett 要求添加测试用例和基准接受率，初始占位符 0.0 经修改为实际阈值 0.5。
• DarkLight1337 建议在测试模型注册表中添加 min_transformers_version，并使用 model_info.check_transformers_version 自动跳过旧版本 CI。
• hospedales 报告了 intermediate_size 不匹配和 quant_config 从目标模型传播到 draft 模型导致加载崩溃的问题，作者随后修复。
• dssugar 指出 Gemma4MTP.__init__ 中的 prefix 可能与 Gemma4ForCausalLM 目标冲突，作者采纳建议调整。

• benchislett 最终批准（"LGTM, Thanks!"），其余审核者无进一步异议。

• 测试添加与接受率阈值 (testing): 添加了 GSM8k 正确性测试，接受率阈值设为 0.5。

• Transformers 版本兼容性 (question): 在 _HfExamplesInfo 中添加了 min_transformers_version 属性。
• intermediate_size 不匹配与 quant_config 传播 (bugfix): 修复 intermediate_size 获取路径，并跳过 quant_config 传播。
• 模型文件前缀冲突 (design): 修改了前缀使用以避免冲突。
• NVFP4 量化目标兼容性 (question): 虽然在量化目标上接受率略有下降，但端到端仍获得加速。

风险与影响

• 风险：
  • 核心路径变更：修改了 llm_base_proposer.py 中的 propose 方法逻辑，可能影响现有的 Eagle/DFlash 提议器，需回归测试。
  • 仅 V1 支持：集成仅发生在 V1 engine 的 gpu_model_runner.py，V2 或其他执行器未覆盖。
  • Centroids CUDA Graphs：假设可捕获 CUDA graphs，若环境不支持（如 CPU、某些 ROCm），会退回到非图路径，但 _setup_centroids_cuda_graphs 无错误处理。
  • Transformer 版本依赖：Gemma4 模型需要较新 transformers 版本，旧版本 CI 可能失败；已添加版本检查但仍需注意。
  • 组合爆炸：新 speculative method 的加入增加了配置空间的复杂度，可能与其他功能（如 prefix caching、multimodal）互斥。
• 影响：
  • 用户：所有 Gemma4 模型用户可直接通过 --speculative-config 启用推测解码，获得 1.3x ~ 3.2x 加速。TP>1 也支持。
  • 系统：增加约 1K 行代码，新文件 gemma4_mtp.py 和 gemma4.py。基类 llm_base_proposer.py 新增 'constant_draft_positions' 机制，对其它提议器透明。
  • 团队：需维护 Gemma4 特有逻辑（centroids masking、多组 block table）；新增配置项 use_gemma4_mtp 使条件判断更清晰。
  • 风险标记：核心路径变更, 仅支持 V1, CUDA Graph 依赖, Transformers 版本依赖, 新配置项增加复杂性

关联脉络

• PR #41880 Laguna xs dflash support: 同样是在 speculative decoding 框架中集成新模型的 PR，涉及 DFlash 和模型注册，与本 PR 共享配置和 runner 集成模式。
• PR #40891 [Core] Avoid using extra thread in UniProcExecutor: 对 V1 executor 的改进，本 PR 新增的 Gemma4Proposer 运行在 V1 engine 上，依赖该 executor 的基础设施。
• PR #41965 [Compressed Tensors] Allow configs with non-explicit ignores: 量化配置的改进，可能影响 Gemma4 MTP 与压缩张量组合使用时的行为。