执行摘要

• 一句话：添加 Cohere Eagle 推测解码模型并修复 Cohere MoE
• 推荐动作：建议精读 cohere_eagle.py 的 __init__ 和 forward，理解 EAGLE 草稿模型的融合机制；同时关注 select_norm_impl 的动态归一化选择设计。该 PR 是 vLLM 集成新推测解码模型的标准范例，值得参考。

功能与动机

PR 目的为支持 Cohere 模型的 EAGLE 推测解码加速，并修复 Cohere MoE 模型中的问题。PR 主体说明：“Add Cohere Eagle to vLLM. Update CohereCommandReasoningParser”。

实现拆解

1. 新增 EAGLE 草稿模型：创建 vllm/model_executor/models/cohere_eagle.py，实现 EagleCohereForCausalLM 类。关键设计包括：融合输入嵌入与目标隐藏状态的 fc 线性层（带 bias）、前置目标模型的 layer_types 以支持滑动窗口注意力、最后使用 LayerNorm 归一化输出。

2. 重构 Cohere MoE 模型：将 cohere_moe.py 重命名为 cohere2_moe.py，类名更新为 Cohere2MoeForCausalLM。新增 RMSNorm 和 select_norm_impl，支持根据配置选择 RMSNorm 或 LayerNorm；修复注意力层密度模式（前缀密集层无滑动窗口）；为 Cohere2MoeMLP 增加 reduce_results 参数。

3. 更新推理解析器：修改 cohere_command_reasoning_parser.py，将 CohereTagStyle.json 改为 json_tags 元组，支持多个 JSON 结构标签；新增 Cohere2MoeForCausalLM 的标签风格；convert_schema_to_structural_tags 函数遍历所有 json_tags 生成结构化标签。

4. 注册模型与测试：在 registry.py 中添加 Cohere2MoeForCausalLM（文本生成模型）和 EagleCohereForCausalLM（推测解码模型），在测试注册表中添加对应测试条目。同时更新 custom_routing_router.py 的导入路径和 supported_models.md 文档。

5. Proposer 集成：在 vllm/v1/spec_decode/llm_base_proposer.py 中新增一行，支持 Cohere Eagle 模型作为推测解码提议者。

关键文件：

• vllm/model_executor/models/cohere_eagle.py（模块 模型层；类别 source；类型 data-contract；符号 CohereEagleDecoderLayer, init, CohereEagleModel, embed_input_ids）: 新增 EagleCohereForCausalLM 模型，实现 Cohere 架构的 EAGLE 推测解码草稿模型。
• vllm/model_executor/models/cohere2_moe.py（模块 模型层；类别 source；类型 rename-or-move；符号 CohereMoeMLP, rms_norm_func, RMSNorm, init）: 从 cohere_moe.py 重命名并重构，使用 RMSNorm 替换 LayerNorm（当存在 rms_norm_eps 时），修复 MoE 注意力层滑动窗口密度模式。
• vllm/reasoning/cohere_command_reasoning_parser.py（模块 推理解析；类别 source；类型 core-logic）: 更新推理解析器以支持多个 JSON 结构标签（MOE 需要同时处理 <|START_RESPONSE|> 和 <|START_TEXT|>），并新增 Cohere2MoeForCausalLM 标签风格。
• vllm/model_executor/models/registry.py（模块 注册表；类别 source；类型 data-contract）: 注册 Cohere2MoeForCausalLM 和 EagleCohereForCausalLM，确保模型可加载。
• tests/models/registry.py（模块 测试配置；类别 test；类型 test-coverage）: 为新的模型架构添加测试配置，包括 EagleCohereForCausalLM。

关键符号：CohereEagleDecoderLayer, CohereEagleModel.init, CohereEagleModel.forward, EagleCohereForCausalLM, rms_norm_func, RMSNorm, select_norm_impl, Cohere2MoeMLP.init, convert_schema_to_structural_tags

关键源码片段

vllm/model_executor/models/cohere2_moe.py

从 cohere_moe.py 重命名并重构，使用 RMSNorm 替换 LayerNorm（当存在 rms_norm_eps 时），修复 MoE 注意力层滑动窗口密度模式。

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
​
@torch.compile(backend=current_platform.simple_compile_backend)
def rms_norm_func(hidden_states, weight, variance_epsilon):
    """RMS 归一化：x / sqrt(mean(x^2) + eps) * weight"""
    input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(torch.float32)
    variance = hidden_states.pow(2).mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
    hidden_states = weight.to(torch.float32) * hidden_states
    return hidden_states.to(input_dtype)
​
​
class RMSNorm(nn.Module):
    """RMS 归一化层，支持 residuals 接口。"""
​
    def __init__(self, param_shape=None, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(param_shape))
        self.variance_epsilon = eps
        # 使用行并行权重加载器（支持张量并行）
        set_weight_attrs(self.weight, {"weight_loader": row_parallel_weight_loader})
​
    def forward(self, hidden_states, residuals=None):
        hidden_states = rms_norm_func(hidden_states, self.weight, self.variance_epsilon)
        return hidden_states, residuals
​
​
def select_norm_impl(config: CohereConfig) -> tuple[type[nn.Module], float]:
    """根据配置选择归一化：若指定 rms_norm_eps 则用 RMSNorm，否则用 LayerNorm。"""
    rms_eps = getattr(config, "rms_norm_eps", None)
    if rms_eps is not None:
        return RMSNorm, rms_eps
    return LayerNorm, config.layer_norm_eps

评论区精华

Review 中 gemini-code-assist[bot] 提出 4 条高优先级建议，全部围绕 cohere_eagle.py 的正确性：

• 初始化 has_own_embed_tokens 和 has_own_lm_head 标志，以便投机提议者正确决定权重共享。
• 第一层 draft 层应禁用输入 layernorm（disable_input_layernorm=(i==0)）。
• 在 load_weights 中调用 process_eagle_weight 检测自定义权重。
• 移除 world_size==1 时跳过 embed_tokens 加载的逻辑，避免权重比较失败。

这些建议均被采纳或已通过后续提交解决，PR 最终被批准合并。

• Weight sharing flags 初始化 (correctness): 作者可能已在后续提交中添加，PR 已合并。
• 第一层 draft 禁用 input layernorm (correctness): 可能已采纳或认为非必需。
• Weight loading 中调用 process_eagle_weight (correctness): 可能已添加。
• 跳过 embed_tokens 加载问题 (correctness): 作者可能移除了该逻辑。

风险与影响

• 风险：

  1. 权重共享标志未设置风险：如果 has_own_embed_tokens/has_own_lm_head 未正确初始化，proposer 可能错误地决定共享或独享权重，导致内存浪费或推理错误。
  2. MoE 重命名兼容性：现有用户若在配置中直接引用 CohereMoeForCausalLM，升级后需改为 Cohere2MoeForCausalLM，可能造成短期兼容性问题。
  3. 测试覆盖不足：新模型仅在离线测试中注册，缺乏端到端性能验证和回归测试。
  4. layer_type 拼接逻辑：草稿模型 layer_type 的拼接依赖于目标模型配置，若目标模型无 layer_types 属性，则回退行为未显式测试。
     - 影响：用户影响：支持 Cohere 架构的 EAGLE 推测解码，可显著降低延迟（类似其他 EAGLE 模型）。用户需使用 --speculative-config 指定草稿模型。
     系统影响：新增两个模型注册（Cohere2MoeForCausalLM、EagleCohereForCausalLM），增加少量运行时内存开销。
     团队影响：维护范围扩大，需跟进 Cohere 官方模型更新，但代码结构与现有 EAGLE 实现一致，学习成本可控。

• 风险标记：新模型代码未在生产环境验证, MoE 重命名可能导致下游配置失效

关联脉络

• 暂无明显关联 PR