[Models] Cohere Eagle + fix to Cohere MoE

docs/models/supported_models.md

@@ -378,7 +378,7 @@ th {
 | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
 | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
 | `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
-| `CohereMoeForCausalLM` | Command (MoE) | (model checkpoints loaded with `trust_remote_code=True`) | ✅︎ | ✅︎ |
+| `Cohere2MoeForCausalLM` | Command (MoE) | (model checkpoints loaded with `trust_remote_code=True`) | ✅︎ | ✅︎ |
 | `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
 | `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ |
 | `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ |

tests/models/registry.py

@@ -238,7 +238,7 @@ def check_available_online(
         "CohereLabs/c4ai-command-r7b-12-2024",
         trust_remote_code=True,
     ),
-    "CohereMoeForCausalLM": _HfExamplesInfo(
+    "Cohere2MoeForCausalLM": _HfExamplesInfo(
         "/host/engines/cohere-moe",
         trust_remote_code=True,
         is_available_online=False,
@@ -1406,6 +1406,13 @@ def check_available_online(
         max_num_seqs=32,
     ),
     # [Eagle]
+    "EagleCohereForCausalLM": _HfExamplesInfo(
+        "/host/engines/cohere-moe",
+        speculative_model="/host/engines/cohere-moe/eagle",
+        tokenizer="/host/engines/cohere-moe",
+        trust_remote_code=True,
+        is_available_online=False,
+    ),
     "EagleDeepSeekMTPModel": _HfExamplesInfo(
         "eagle618/deepseek-v3-random",
         speculative_model="eagle618/eagle-deepseek-v3-random",

vllm/model_executor/layers/fused_moe/router/custom_routing_router.py

@@ -34,7 +34,7 @@ def __init__(
 
     @property
     def routing_method_type(self) -> RoutingMethodType:
-        from vllm.model_executor.models.cohere_moe import token_choice_with_bias
+        from vllm.model_executor.models.cohere2_moe import token_choice_with_bias
         from vllm.model_executor.models.llama4 import Llama4MoE
 
         # NOTE: FLASHINFER_TRTLLM support the Llama4 router.

vllm/model_executor/models/cohere_moe.py → vllm/model_executor/models/cohere2_moe.py

@@ -53,7 +53,7 @@ def token_choice_with_bias(
     topk: int,
     renormalize: bool,
 ):
-    """Sigmoid -> top-k (-> renormalize) custom routing for CohereMoe."""
+    """Sigmoid -> top-k (-> renormalize) custom routing for Cohere2Moe."""
     assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
 
     scores = gating_output.float().sigmoid()
@@ -65,7 +65,7 @@ def token_choice_with_bias(
     return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
 
 
-class CohereMoeMLP(nn.Module):
+class Cohere2MoeMLP(nn.Module):
     """Cohere MLP used as shared experts in the MoE block."""
 
     def __init__(
@@ -107,7 +107,7 @@ def forward(self, x):
         return x
 
 
-class CohereMoeAttention(nn.Module):
+class Cohere2MoeAttention(nn.Module):
     """Cohere MoE attention with sliding-window interleave."""
 
     def __init__(
@@ -195,8 +195,8 @@ def forward(
         return output
 
 
-class CohereMoe(nn.Module):
-    """Tensor-parallel MoE block for CohereMoe with shared experts."""
+class Cohere2Moe(nn.Module):
+    """Tensor-parallel MoE block for Cohere2Moe with shared experts."""
 
     def __init__(
         self,
@@ -234,7 +234,7 @@ def __init__(
         )
 
         if hasattr(config, "num_shared_experts") and config.num_shared_experts > 0:
-            self.shared_experts = CohereMoeMLP(
+            self.shared_experts = Cohere2MoeMLP(
                 config=config,
                 intermediate_size=config.intermediate_size * config.num_shared_experts,
                 quant_config=quant_config,
@@ -276,7 +276,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return final_hidden_states.view(orig_shape)
 
 
-class CohereMoeDecoderLayer(nn.Module):
+class Cohere2MoeDecoderLayer(nn.Module):
     def __init__(
         self,
         config: CohereConfig,
@@ -288,13 +288,13 @@ def __init__(
         self.config = config
         self.hidden_size = config.hidden_size
 
-        self.self_attn = CohereMoeAttention(
+        self.self_attn = Cohere2MoeAttention(
             config,
             cache_config,
             quant_config=quant_config,
             prefix=f"{prefix}.self_attn",
         )
-        self.mlp = CohereMoe(
+        self.mlp = Cohere2Moe(
             config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
         )
         self.input_layernorm = LayerNorm(
@@ -320,8 +320,8 @@ def forward(
 
 
 @support_torch_compile
-class CohereMoeModel(nn.Module):
-    """Transformer decoder for CohereMoe."""
+class Cohere2MoeModel(nn.Module):
+    """Transformer decoder for Cohere2Moe."""
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -339,7 +339,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         )
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
-            lambda prefix: CohereMoeDecoderLayer(
+            lambda prefix: Cohere2MoeDecoderLayer(
                 config, cache_config, quant_config, prefix=prefix
             ),
             prefix=f"{prefix}.layers",
@@ -471,7 +471,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loaded_params
 
 
-class CohereMoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
+class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
     is_text_generation_model = True
 
     packed_modules_mapping = {
@@ -498,7 +498,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.logits_processor = LogitsProcessor(
             self.unpadded_vocab_size, config.vocab_size, scale=self.logits_scale
         )
-        self.model = CohereMoeModel(
+        self.model = Cohere2MoeModel(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
         self.make_empty_intermediate_tensors = (

vllm/model_executor/models/cohere_eagle.py

@@ -0,0 +1,247 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+from transformers import CohereConfig
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.commandr import (
+    CohereDecoderLayer,
+    CohereForCausalLM,
+    LayerNorm,
+)
+
+from .utils import (
+    AutoWeightsLoader,
+    get_draft_quant_config,
+    maybe_prefix,
+    process_eagle_weight,
+)
+
+logger = init_logger(__name__)
+
+
+class CohereEagleDecoderLayer(CohereDecoderLayer):
+    """Eagle draft variant of CohereDecoderLayer."""
+
+    def __init__(
+        self,
+        config: CohereConfig,
+        cache_config=None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__(
+            config,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
+
+@support_torch_compile
+class CohereEagleModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+        start_layer_id: int = 0,
+    ) -> None:
+        super().__init__()
+        self.config = vllm_config.speculative_config.draft_model_config.hf_config
+        self.quant_config = get_draft_quant_config(vllm_config)
+
+        # Cohere2-targeted EAGLE drafts inherit the target's sliding-window
+        # attention pattern. ``CohereAttention`` resolves per-layer behavior
+        # via ``config.layer_types[layer_idx]`` and the eagle layers use
+        # absolute indices (target_layer_num + i), so prepend the target's
+        # ``layer_types`` to the draft's so the lookup succeeds.
+        target_text_config = vllm_config.model_config.hf_text_config
+        if hasattr(target_text_config, "layer_types") and hasattr(
+            self.config, "layer_types"
+        ):
+            self.config.layer_types = list(target_text_config.layer_types) + list(
+                self.config.layer_types
+            )
+
+        self.vocab_size = self.config.vocab_size
+        self.embed_tokens = VocabParallelEmbedding(
+            self.config.vocab_size,
+            self.config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                CohereEagleDecoderLayer(
+                    self.config,
+                    cache_config=vllm_config.cache_config,
+                    quant_config=self.quant_config,
+                    prefix=maybe_prefix(prefix, f"layers.{i + start_layer_id}"),
+                )
+                for i in range(self.config.num_hidden_layers)
+            ]
+        )
+
+        # Cohere EAGLE checkpoints include a bias term on the input fusion
+        # projection (unlike LLaMA EAGLE which uses bias=False).
+        self.fc = ReplicatedLinear(
+            input_size=self.config.hidden_size * 2,
+            output_size=self.config.hidden_size,
+            bias=True,
+            params_dtype=vllm_config.model_config.dtype,
+            quant_config=self.quant_config,
+            prefix=maybe_prefix(prefix, "fc"),
+            return_bias=False,
+        )
+
+        # Cohere EAGLE applies an explicit final LayerNorm to the draft
+        # hidden states before they are consumed by the logits processor.
+        self.norm = LayerNorm(
+            param_shape=(self.config.hidden_size),
+            eps=self.config.layer_norm_eps,
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_embeds = self.embed_tokens(input_ids)
+        hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1))
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(
+                positions,
+                hidden_states,
+                residual,
+            )
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states, hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+
+            if self.quant_config is not None and (
+                scale_name := self.quant_config.get_cache_scale(name)
+            ):
+                param = params_dict[scale_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                loaded_weight = (
+                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
+                )
+                weight_loader(param, loaded_weight)
+                loaded_params.add(scale_name)
+                continue
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class EagleCohereForCausalLM(CohereForCausalLM):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        nn.Module.__init__(self)
+        self.config = vllm_config.speculative_config.draft_model_config.hf_config
+        # Flags checked by the speculative proposer to decide whether to share
+        # embed_tokens / lm_head with the target model. Cohere EAGLE checkpoints
+        # use tied embeddings so these weights are absent from the draft file.
+        self.has_own_embed_tokens = False
+        self.has_own_lm_head = False
+        target_layer_num = vllm_config.model_config.get_num_layers(
+            vllm_config.parallel_config
+        )
+        self.model = CohereEagleModel(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+            start_layer_id=target_layer_num,
+        )
+
+        logit_scale = getattr(self.config, "logit_scale", 1.0)
+        self.logits_processor = LogitsProcessor(
+            self.config.vocab_size, scale=logit_scale
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if inputs_embeds is not None:
+            raise NotImplementedError(
+                f"{type(self).__name__} does not support multimodal inputs yet."
+            )
+        return self.model(input_ids, positions, hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        def _track_and_forward(inputs):
+            name, weight = inputs
+            process_eagle_weight(self, name)
+            return name, weight
+
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(
+                ["lm_head.", "model.embed_tokens."]
+                if self.config.tie_word_embeddings
+                else None
+            ),
+        )
+
+        loaded_weight_names = loader.load_weights(map(_track_and_forward, weights))
+
+        # Embed tokens are tied with the target model and therefore not
+        # present in the EAGLE checkpoint; mark them as loaded explicitly to
+        # avoid a spurious "weight not found" warning from the default
+        # weight loader.
+        loaded_weight_names.add("model.embed_tokens.weight")
+        return loaded_weight_names