[Models] Cohere Eagle + fix to Cohere MoE (#42078)

Signed-off-by: Terrencezzj <[email protected]>
Co-authored-by: Cursor <[email protected]>

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 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "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 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         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 @@ class CustomRoutingRouter(BaseRouter):
 
     @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/cohere2_moe.py

@@ -30,7 +30,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmb
 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
 
@@ -53,7 +55,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 +67,38 @@ def token_choice_with_bias(
     return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
 
 
-class CohereMoeMLP(nn.Module):
+@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."""
 
     def __init__(
@@ -73,6 +106,7 @@ class CohereMoeMLP(nn.Module):
         config: CohereConfig,
         intermediate_size: int | None = None,
         quant_config: QuantizationConfig | None = None,
+        reduce_results: bool = False,
         prefix: str = "",
     ):
         super().__init__()
@@ -95,7 +129,7 @@ class CohereMoeMLP(nn.Module):
             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()
@@ -107,7 +141,7 @@ class CohereMoeMLP(nn.Module):
         return x
 
 
-class CohereMoeAttention(nn.Module):
+class Cohere2MoeAttention(nn.Module):
     """Cohere MoE attention with sliding-window interleave."""
 
     def __init__(
@@ -170,6 +204,19 @@ class CohereMoeAttention(nn.Module):
         ):
             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,15 +235,15 @@ class CohereMoeAttention(nn.Module):
     ) -> 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)
         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 +281,7 @@ class CohereMoe(nn.Module):
         )
 
         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 +323,7 @@ class CohereMoe(nn.Module):
         return final_hidden_states.view(orig_shape)
 
 
-class CohereMoeDecoderLayer(nn.Module):
+class Cohere2MoeDecoderLayer(nn.Module):
     def __init__(
         self,
         config: CohereConfig,
@@ -287,19 +334,34 @@ class CohereMoeDecoderLayer(nn.Module):
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
+        self.layer_idx = extract_layer_index(prefix)
 
-        self.self_attn = CohereMoeAttention(
+        self.self_attn = Cohere2MoeAttention(
             config,
             cache_config,
             quant_config=quant_config,
             prefix=f"{prefix}.self_attn",
         )
-        self.mlp = CohereMoe(
-            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,
@@ -320,8 +382,8 @@ class CohereMoeDecoderLayer(nn.Module):
 
 
 @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,14 +401,13 @@ class CohereMoeModel(nn.Module):
         )
         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",
         )
-        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
         )
@@ -471,7 +532,7 @@ class CohereMoeModel(nn.Module):
         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 +559,7 @@ class CohereMoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
         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

vllm/model_executor/models/registry.py

@@ -89,7 +89,7 @@ _TEXT_GENERATION_MODELS = {
     "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"),
@@ -582,6 +582,7 @@ _SPECULATIVE_DECODING_MODELS = {
     "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"),

vllm/reasoning/cohere_command_reasoning_parser.py

@@ -39,16 +39,20 @@ REPLACEMENT_CHAR = "\ufffd"
 
 
 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,
+    ),
+    "Cohere2MoeForCausalLM": 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 convert_schema_to_structural_tags(
             "type": "structural_tag",
             "format": {
                 "type": "triggered_tags",
-                "triggers": [t["begin"] for t in tags],
+                "triggers": triggers,
                 "tags": tags,
             },
         }
@@ -505,7 +524,7 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
             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 "

vllm/v1/spec_decode/llm_base_proposer.py

@@ -1168,6 +1168,7 @@ class SpecDecodeBaseProposer:
             # handle multimodality
             assert hasattr(target_model, "config")
             if self.get_model_name(target_model) in [
+                "Cohere2VisionForConditionalGeneration",
                 "Exaone4_5_ForConditionalGeneration",
                 "GlmOcrForConditionalGeneration",
                 "HunYuanVLForConditionalGeneration",