From bdca715ca80c6a20fe9868c96b4f9f043244dd0b Mon Sep 17 00:00:00 2001 From: Terrencezzj Date: Fri, 8 May 2026 14:32:24 +0000 Subject: [PATCH 1/5] Add Cohere Eagle speculative decoding model and update reasoning parser for MHL v2 - Add EagleCohereForCausalLM model (cohere_eagle.py) for Eagle speculative decoding - Register EagleCohereForCausalLM in model and test registries - Add Cohere2VisionForConditionalGeneration to multimodal spec decode list - Update CohereCommandReasoningParser for MHL v2: support multiple JSON tags per architecture (e.g. MOE uses both START_RESPONSE and START_TEXT delimiters), add Cohere2MoeForCausalLM tag style, fix structural_tag triggers list Co-authored-by: Cursor Signed-off-by: Terrencezzj --- tests/models/registry.py | 7 + vllm/model_executor/models/cohere_eagle.py | 255 ++++++++++++++++++ vllm/model_executor/models/registry.py | 1 + .../cohere_command_reasoning_parser.py | 45 +++- vllm/v1/spec_decode/llm_base_proposer.py | 1 + 5 files changed, 296 insertions(+), 13 deletions(-) create mode 100644 vllm/model_executor/models/cohere_eagle.py diff --git a/tests/models/registry.py b/tests/models/registry.py index e50b0a8de4d9..1786028d8e03 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -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", diff --git a/vllm/model_executor/models/cohere_eagle.py b/vllm/model_executor/models/cohere_eagle.py new file mode 100644 index 000000000000..83abc386ef16 --- /dev/null +++ b/vllm/model_executor/models/cohere_eagle.py @@ -0,0 +1,255 @@ +# 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.distributed.parallel_state import get_pp_group +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 + +logger = init_logger(__name__) + + +class CohereEagleDecoderLayer(CohereDecoderLayer): + """Eagle draft variant of CohereDecoderLayer. + + Optionally replaces ``input_layernorm`` with ``nn.Identity`` on the very + first draft layer to mirror the original EAGLE implementation. + + Reference: + https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427 + """ + + def __init__( + self, + config: CohereConfig, + cache_config=None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + disable_input_layernorm: bool = False, + ) -> None: + super().__init__( + config, + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + ) + if disable_input_layernorm: + del self.input_layernorm + self.input_layernorm = nn.Identity() + + +@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: + # When PP is disabled the draft model shares its token + # embeddings with the target model, so skip loading any + # embed_tokens weights coming from the EAGLE checkpoint. + if get_pp_group().world_size == 1 and "embed_tokens." in name: + continue + + # lm_head is tied with embed_tokens; skip to avoid duplicate + # loads. + if "lm_head.weight" in name: + continue + + 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 + 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]]): + loader = AutoWeightsLoader( + self, + skip_prefixes=( + ["lm_head.", "model.embed_tokens."] + if self.config.tie_word_embeddings + else None + ), + ) + + loaded_weight_names = loader.load_weights(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 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index d38cd63b90b1..70ab5f3e4120 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -582,6 +582,7 @@ "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"), "MiMoV2MTPModel": ("mimo_v2_mtp", "MiMoV2MTP"), "MiMoV2OmniMTPModel": ("mimo_v2_mtp", "MiMoV2OmniMTP"), + "EagleCohereForCausalLM": ("cohere_eagle", "EagleCohereForCausalLM"), "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"), "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"), "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"), diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index c96b21d4e8fb..eaa616db1fea 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -39,16 +39,20 @@ class CohereTagRegistry(NamedTuple): - """A single ``structural_tag`` begin("trigger")/end pair.""" + """A single ``structural_tag`` trigger / end pair (``begin`` uses ``trigger``).""" trigger: str end: str class CohereTagStyle(NamedTuple): - """The structural tags style for a given model architecture.""" + """The structural tags style for a given model architecture. - json: CohereTagRegistry + ``json_tags`` lists every JSON-schema wrapper the model may emit (MOE uses + both response and text delimiters). ``tools`` is the tool-call wrapper. + """ + + json_tags: tuple[CohereTagRegistry, ...] tools: CohereTagRegistry @@ -64,18 +68,30 @@ class CohereNormalizedTool(TypedDict): COMMAND_A_TOOLS_TAG = CohereTagRegistry( - trigger="<|START_ACTION|>", end="<|END_ACTION|>" + trigger="<|START_ACTION|>", + end="<|END_ACTION|>", ) COMMAND_A_JSON_TAG = CohereTagRegistry( - trigger="<|START_RESPONSE|>", end="<|END_RESPONSE|>" + trigger="<|START_RESPONSE|>", + end="<|END_RESPONSE|>", +) +COMMAND_A_PLUS_JSON_TAG = CohereTagRegistry( + trigger="<|START_TEXT|>", + end="<|END_TEXT|>", ) MODEL_TO_TAG_STYLE: dict[str, CohereTagStyle] = { "Cohere2ForCausalLM": CohereTagStyle( - json=COMMAND_A_JSON_TAG, tools=COMMAND_A_TOOLS_TAG + json_tags=(COMMAND_A_JSON_TAG,), + tools=COMMAND_A_TOOLS_TAG, ), "Cohere2VisionForConditionalGeneration": CohereTagStyle( - json=COMMAND_A_JSON_TAG, tools=COMMAND_A_TOOLS_TAG + json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG), + tools=COMMAND_A_TOOLS_TAG, + ), + "CohereMoeForCausalLM": CohereTagStyle( + json_tags=(COMMAND_A_JSON_TAG,), + tools=COMMAND_A_TOOLS_TAG, ), } @@ -211,15 +227,18 @@ def convert_schema_to_structural_tags( style = MODEL_TO_TAG_STYLE[model_architecture] tags: list[dict] = [] + triggers: list[str] = [] def _add_tag(tag: CohereTagRegistry, content: dict) -> None: tags.append({"begin": tag.trigger, "content": content, "end": tag.end}) + triggers.append(tag.trigger) if schema is not None: - # Add the JSON-schema tag both for schema-only requests and for the - # "tools plus JSON mode" case (North use case: follow the schema when - # the model decides not to call any tool). - _add_tag(style.json, {"type": "json_schema", "json_schema": schema}) + # One structural tag per JSON wrapper (e.g. MOE: response + text). + # Same for schema-only and "tools plus JSON mode" (North: schema when + # the model does not call tools). + for jt in style.json_tags: + _add_tag(jt, {"type": "json_schema", "json_schema": schema}) if _has_effective_tools(tools): # ``tools`` may be a JSON string (poseidon / RESPONSE_FORMAT_TOOL_DEFINITIONS) @@ -240,7 +259,7 @@ def _add_tag(tag: CohereTagRegistry, content: dict) -> None: "type": "structural_tag", "format": { "type": "triggered_tags", - "triggers": [t["begin"] for t in tags], + "triggers": triggers, "tags": tags, }, } @@ -505,7 +524,7 @@ def adjust_request( model_architecture=model_architecture, ) if result is None: - # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``; conversion + # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``. raise ValueError( "Failed to build structural_tag guided decoding constraints from " "this request's JSON schema and/or tools. The configured model " diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 8ee349a1cc0d..30a44398ff74 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1382,6 +1382,7 @@ def load_model(self, target_model: nn.Module) -> None: # handle multimodality assert hasattr(target_model, "config") if self.get_model_name(target_model) in [ + "Cohere2VisionForConditionalGeneration", "Exaone4_5_ForConditionalGeneration", "GlmOcrForConditionalGeneration", "HunYuanVLForConditionalGeneration", From 80d10291c58a3f3df0881af5981045f6e3a206e8 Mon Sep 17 00:00:00 2001 From: Terrencezzj Date: Fri, 8 May 2026 15:59:22 +0000 Subject: [PATCH 2/5] renaming Signed-off-by: Terrencezzj --- docs/models/supported_models.md | 2 +- tests/models/registry.py | 2 +- .../fused_moe/router/custom_routing_router.py | 2 +- .../models/{cohere_moe.py => cohere2_moe.py} | 28 +++++++++---------- vllm/model_executor/models/registry.py | 2 +- .../cohere_command_reasoning_parser.py | 2 +- 6 files changed, 19 insertions(+), 19 deletions(-) rename vllm/model_executor/models/{cohere_moe.py => cohere2_moe.py} (96%) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index e79fec8169f2..2b078ca772e2 100644 --- a/docs/models/supported_models.md +++ b/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. | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index 1786028d8e03..3465ce9b9734 100644 --- a/tests/models/registry.py +++ b/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, diff --git a/vllm/model_executor/layers/fused_moe/router/custom_routing_router.py b/vllm/model_executor/layers/fused_moe/router/custom_routing_router.py index 8e35169d9005..a3e0075f2b7d 100644 --- a/vllm/model_executor/layers/fused_moe/router/custom_routing_router.py +++ b/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. diff --git a/vllm/model_executor/models/cohere_moe.py b/vllm/model_executor/models/cohere2_moe.py similarity index 96% rename from vllm/model_executor/models/cohere_moe.py rename to vllm/model_executor/models/cohere2_moe.py index a059d68c9d02..63f8a9ba76c6 100644 --- a/vllm/model_executor/models/cohere_moe.py +++ b/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 = ( diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 70ab5f3e4120..5b77da7aab11 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -89,7 +89,7 @@ "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"), "CohereForCausalLM": ("commandr", "CohereForCausalLM"), "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"), - "CohereMoeForCausalLM": ("cohere_moe", "CohereMoeForCausalLM"), + "Cohere2MoeForCausalLM": ("cohere2_moe", "Cohere2MoeForCausalLM"), "CwmForCausalLM": ("llama", "LlamaForCausalLM"), "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"), "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"), diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index eaa616db1fea..b28a59089e73 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -89,7 +89,7 @@ class CohereNormalizedTool(TypedDict): json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG), tools=COMMAND_A_TOOLS_TAG, ), - "CohereMoeForCausalLM": CohereTagStyle( + "Cohere2MoeForCausalLM": CohereTagStyle( json_tags=(COMMAND_A_JSON_TAG,), tools=COMMAND_A_TOOLS_TAG, ), From 344552d96754a5ee9f0ee1582ef41e20bf0566b5 Mon Sep 17 00:00:00 2001 From: Terrencezzj Date: Fri, 8 May 2026 16:51:11 +0000 Subject: [PATCH 3/5] address comment Signed-off-by: Terrencezzj --- vllm/model_executor/models/cohere_eagle.py | 41 ++++++++-------------- 1 file changed, 15 insertions(+), 26 deletions(-) diff --git a/vllm/model_executor/models/cohere_eagle.py b/vllm/model_executor/models/cohere_eagle.py index 83abc386ef16..f2e1ecbb2940 100644 --- a/vllm/model_executor/models/cohere_eagle.py +++ b/vllm/model_executor/models/cohere_eagle.py @@ -9,7 +9,6 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig -from vllm.distributed.parallel_state import get_pp_group from vllm.logger import init_logger from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -22,20 +21,13 @@ LayerNorm, ) -from .utils import AutoWeightsLoader, get_draft_quant_config, maybe_prefix +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. - - Optionally replaces ``input_layernorm`` with ``nn.Identity`` on the very - first draft layer to mirror the original EAGLE implementation. - - Reference: - https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427 - """ + """Eagle draft variant of CohereDecoderLayer.""" def __init__( self, @@ -43,7 +35,6 @@ def __init__( cache_config=None, quant_config: QuantizationConfig | None = None, prefix: str = "", - disable_input_layernorm: bool = False, ) -> None: super().__init__( config, @@ -51,9 +42,6 @@ def __init__( quant_config=quant_config, prefix=prefix, ) - if disable_input_layernorm: - del self.input_layernorm - self.input_layernorm = nn.Identity() @support_torch_compile @@ -180,17 +168,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader(param, loaded_weight, shard_id) break else: - # When PP is disabled the draft model shares its token - # embeddings with the target model, so skip loading any - # embed_tokens weights coming from the EAGLE checkpoint. - if get_pp_group().world_size == 1 and "embed_tokens." in name: - continue - - # lm_head is tied with embed_tokens; skip to avoid duplicate - # loads. - if "lm_head.weight" in name: - continue - if name.endswith(".bias") and name not in params_dict: continue @@ -205,6 +182,11 @@ 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 ) @@ -236,6 +218,11 @@ def forward( 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=( @@ -245,7 +232,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): ), ) - loaded_weight_names = loader.load_weights(weights) + 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 From 4f43e372320d44be614fdf2fdab2c309ec7a7d87 Mon Sep 17 00:00:00 2001 From: Terrencezzj Date: Fri, 8 May 2026 17:41:11 +0000 Subject: [PATCH 4/5] pre-commit Signed-off-by: Terrencezzj --- vllm/model_executor/models/cohere_eagle.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/cohere_eagle.py b/vllm/model_executor/models/cohere_eagle.py index f2e1ecbb2940..5c22d6e34dd5 100644 --- a/vllm/model_executor/models/cohere_eagle.py +++ b/vllm/model_executor/models/cohere_eagle.py @@ -21,7 +21,12 @@ LayerNorm, ) -from .utils import AutoWeightsLoader, get_draft_quant_config, maybe_prefix, process_eagle_weight +from .utils import ( + AutoWeightsLoader, + get_draft_quant_config, + maybe_prefix, + process_eagle_weight, +) logger = init_logger(__name__) @@ -232,9 +237,7 @@ def _track_and_forward(inputs): ), ) - loaded_weight_names = loader.load_weights( - map(_track_and_forward, weights) - ) + 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 From 7792da45d8aa7e991224628cbc38a434bd79b3c0 Mon Sep 17 00:00:00 2001 From: Terrencezzj Date: Fri, 8 May 2026 20:59:26 +0000 Subject: [PATCH 5/5] model Signed-off-by: Terrencezzj --- vllm/model_executor/models/cohere2_moe.py | 83 ++++++++++++++++++++--- 1 file changed, 72 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index 63f8a9ba76c6..aa8adff188f7 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -30,7 +30,9 @@ from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, + row_parallel_weight_loader, ) +from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -65,6 +67,37 @@ def token_choice_with_bias( return topk_weights.to(torch.float32), topk_ids.to(torch.int32) +@torch.compile(backend=current_platform.simple_compile_backend) +def rms_norm_func(hidden_states, weight, variance_epsilon): + 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): + 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]: + """Returns (norm_class, eps). Uses RMSNorm when config.rms_norm_eps is set, + otherwise falls back to LayerNorm with config.layer_norm_eps.""" + rms_eps = getattr(config, "rms_norm_eps", None) + if rms_eps is not None: + return RMSNorm, rms_eps + return LayerNorm, config.layer_norm_eps + + class Cohere2MoeMLP(nn.Module): """Cohere MLP used as shared experts in the MoE block.""" @@ -73,6 +106,7 @@ def __init__( config: CohereConfig, intermediate_size: int | None = None, quant_config: QuantizationConfig | None = None, + reduce_results: bool = False, prefix: str = "", ): super().__init__() @@ -95,7 +129,7 @@ def __init__( self.hidden_size, bias=False, quant_config=quant_config, - reduce_results=False, + reduce_results=reduce_results, prefix=f"{prefix}.down_proj", ) self.act_fn = SiluAndMul() @@ -170,6 +204,19 @@ def __init__( ): self.sliding_window = config.sliding_window + # Prefix-dense layers (layer_idx < first_k_dense_replace) have full + # attention (no sliding window). When prefix_dense_sliding_window_pattern + # == 1, they keep RoPE even though they are not sliding-window layers. + first_k_dense_replace = getattr(config, "first_k_dense_replace", 0) + prefix_dense_sliding_window_pattern = getattr( + config, "prefix_dense_sliding_window_pattern", 1 + ) + self.force_rope = bool( + first_k_dense_replace + and prefix_dense_sliding_window_pattern == 1 + and self.layer_idx < first_k_dense_replace + ) + self.attn = Attention( self.num_heads, self.head_dim, @@ -188,7 +235,7 @@ def forward( ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - if self.sliding_window: + if self.sliding_window or self.force_rope: q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) @@ -287,6 +334,7 @@ def __init__( super().__init__() self.config = config self.hidden_size = config.hidden_size + self.layer_idx = extract_layer_index(prefix) self.self_attn = Cohere2MoeAttention( config, @@ -294,12 +342,26 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.self_attn", ) - self.mlp = Cohere2Moe( - config=config, quant_config=quant_config, prefix=f"{prefix}.mlp" - ) - self.input_layernorm = LayerNorm( - param_shape=(config.hidden_size,), eps=config.layer_norm_eps - ) + + # Layers before first_k_dense_replace use a dense MLP instead of MoE. + first_k_dense_replace = getattr(config, "first_k_dense_replace", 0) + if self.layer_idx < first_k_dense_replace: + self.mlp = Cohere2MoeMLP( + config=config, + intermediate_size=getattr( + config, "prefix_dense_intermediate_size", config.intermediate_size + ), + quant_config=quant_config, + reduce_results=True, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = Cohere2Moe( + config=config, quant_config=quant_config, prefix=f"{prefix}.mlp" + ) + + norm_cls, norm_eps = select_norm_impl(config) + self.input_layernorm = norm_cls(param_shape=(config.hidden_size,), eps=norm_eps) def forward( self, @@ -344,9 +406,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) - self.norm = LayerNorm( - param_shape=(config.hidden_size,), eps=config.layer_norm_eps - ) + norm_cls, norm_eps = select_norm_impl(config) + self.norm = norm_cls(param_shape=(config.hidden_size,), eps=norm_eps) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size )