From a2812becd6952ab66b1ee54856bb855a0551353b Mon Sep 17 00:00:00 2001 From: Terrence Zhao <32208165+Terrencezzj@users.noreply.github.com> Date: Sat, 9 May 2026 00:46:26 -0400 Subject: [PATCH] [Models] Cohere Eagle + fix to Cohere MoE (#42078) Signed-off-by: Terrencezzj Co-authored-by: Cursor --- docs/models/supported_models.md | 2 +- tests/models/registry.py | 9 +- .../fused_moe/router/custom_routing_router.py | 2 +- .../models/{cohere_moe.py => cohere2_moe.py} | 109 ++++++-- vllm/model_executor/models/cohere_eagle.py | 247 ++++++++++++++++++ vllm/model_executor/models/registry.py | 3 +- .../cohere_command_reasoning_parser.py | 45 +++- vllm/v1/spec_decode/llm_base_proposer.py | 1 + 8 files changed, 377 insertions(+), 41 deletions(-) rename vllm/model_executor/models/{cohere_moe.py => cohere2_moe.py} (83%) create mode 100644 vllm/model_executor/models/cohere_eagle.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index e79fec8169..2b078ca772 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 e50b0a8de4..3465ce9b97 100644 --- a/tests/models/registry.py +++ b/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", 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 8e35169d90..a3e0075f2b 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 @@ 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. diff --git a/vllm/model_executor/models/cohere_moe.py b/vllm/model_executor/models/cohere2_moe.py similarity index 83% rename from vllm/model_executor/models/cohere_moe.py rename to vllm/model_executor/models/cohere2_moe.py index a059d68c9d..aa8adff188 100644 --- a/vllm/model_executor/models/cohere_moe.py +++ b/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 = ( diff --git a/vllm/model_executor/models/cohere_eagle.py b/vllm/model_executor/models/cohere_eagle.py new file mode 100644 index 0000000000..5c22d6e34d --- /dev/null +++ b/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 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index d38cd63b90..5b77da7aab 100644 --- a/vllm/model_executor/models/registry.py +++ b/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"), diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index c96b21d4e8..b28a59089e 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/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 " diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 08ed9798ed..cc113025c1 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/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",