@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a2812becd6
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 01:15 UTC

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

Signed-off-by: Terrencezzj <terrence@cohere.ai> Co-authored-by: Cursor <cursoragent@cursor.com>

Terrence Zhao committed May 9, 2026 at 00:46 UTC a2812becd6952ab66b1ee54856bb855a0551353b
8 files changed +377 -41
docs/models/supported_models.md
+1 -1
@@ -378,7 +378,7 @@ th {
378 | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
379 | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
380 | `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. | ✅︎ | ✅︎ |
381 -| `CohereMoeForCausalLM` | Command (MoE) | (model checkpoints loaded with `trust_remote_code=True`) | ✅︎ | ✅︎ |
381 +| `Cohere2MoeForCausalLM` | Command (MoE) | (model checkpoints loaded with `trust_remote_code=True`) | ✅︎ | ✅︎ |
382 | `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
383 | `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ |
384 | `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+8 -1
@@ -238,7 +238,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
238 "CohereLabs/c4ai-command-r7b-12-2024",
239 trust_remote_code=True,
240 ),
241 - "CohereMoeForCausalLM": _HfExamplesInfo(
241 + "Cohere2MoeForCausalLM": _HfExamplesInfo(
242 "/host/engines/cohere-moe",
243 trust_remote_code=True,
244 is_available_online=False,
@@ -1406,6 +1406,13 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1406 max_num_seqs=32,
1407 ),
1408 # [Eagle]
1409 + "EagleCohereForCausalLM": _HfExamplesInfo(
1410 + "/host/engines/cohere-moe",
1411 + speculative_model="/host/engines/cohere-moe/eagle",
1412 + tokenizer="/host/engines/cohere-moe",
1413 + trust_remote_code=True,
1414 + is_available_online=False,
1415 + ),
1416 "EagleDeepSeekMTPModel": _HfExamplesInfo(
1417 "eagle618/deepseek-v3-random",
1418 speculative_model="eagle618/eagle-deepseek-v3-random",
vllm/model_executor/layers/fused_moe/router/custom_routing_router.py
+1 -1
@@ -34,7 +34,7 @@ class CustomRoutingRouter(BaseRouter):
34
35 @property
36 def routing_method_type(self) -> RoutingMethodType:
37 - from vllm.model_executor.models.cohere_moe import token_choice_with_bias
37 + from vllm.model_executor.models.cohere2_moe import token_choice_with_bias
38 from vllm.model_executor.models.llama4 import Llama4MoE
39
40 # NOTE: FLASHINFER_TRTLLM support the Llama4 router.
vllm/model_executor/models/cohere2_moe.py renamed
+85 -24
@@ -30,7 +30,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmb
30 from vllm.model_executor.model_loader.weight_utils import (
31 default_weight_loader,
32 maybe_remap_kv_scale_name,
33 + row_parallel_weight_loader,
34 )
35 +from vllm.model_executor.utils import set_weight_attrs
36 from vllm.platforms import current_platform
37 from vllm.sequence import IntermediateTensors
38
@@ -53,7 +55,7 @@ def token_choice_with_bias(
55 topk: int,
56 renormalize: bool,
57 ):
56 - """Sigmoid -> top-k (-> renormalize) custom routing for CohereMoe."""
58 + """Sigmoid -> top-k (-> renormalize) custom routing for Cohere2Moe."""
59 assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
60
61 scores = gating_output.float().sigmoid()
@@ -65,7 +67,38 @@ def token_choice_with_bias(
67 return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
68
69
68 -class CohereMoeMLP(nn.Module):
70 +@torch.compile(backend=current_platform.simple_compile_backend)
71 +def rms_norm_func(hidden_states, weight, variance_epsilon):
72 + input_dtype = hidden_states.dtype
73 + hidden_states = hidden_states.to(torch.float32)
74 + variance = hidden_states.pow(2).mean(-1, keepdim=True)
75 + hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
76 + hidden_states = weight.to(torch.float32) * hidden_states
77 + return hidden_states.to(input_dtype)
78 +
79 +
80 +class RMSNorm(nn.Module):
81 + def __init__(self, param_shape=None, eps=1e-6):
82 + super().__init__()
83 + self.weight = nn.Parameter(torch.ones(param_shape))
84 + self.variance_epsilon = eps
85 + set_weight_attrs(self.weight, {"weight_loader": row_parallel_weight_loader})
86 +
87 + def forward(self, hidden_states, residuals=None):
88 + hidden_states = rms_norm_func(hidden_states, self.weight, self.variance_epsilon)
89 + return hidden_states, residuals
90 +
91 +
92 +def select_norm_impl(config: CohereConfig) -> tuple[type[nn.Module], float]:
93 + """Returns (norm_class, eps). Uses RMSNorm when config.rms_norm_eps is set,
94 + otherwise falls back to LayerNorm with config.layer_norm_eps."""
95 + rms_eps = getattr(config, "rms_norm_eps", None)
96 + if rms_eps is not None:
97 + return RMSNorm, rms_eps
98 + return LayerNorm, config.layer_norm_eps
99 +
100 +
101 +class Cohere2MoeMLP(nn.Module):
102 """Cohere MLP used as shared experts in the MoE block."""
103
104 def __init__(
@@ -73,6 +106,7 @@ class CohereMoeMLP(nn.Module):
106 config: CohereConfig,
107 intermediate_size: int | None = None,
108 quant_config: QuantizationConfig | None = None,
109 + reduce_results: bool = False,
110 prefix: str = "",
111 ):
112 super().__init__()
@@ -95,7 +129,7 @@ class CohereMoeMLP(nn.Module):
129 self.hidden_size,
130 bias=False,
131 quant_config=quant_config,
98 - reduce_results=False,
132 + reduce_results=reduce_results,
133 prefix=f"{prefix}.down_proj",
134 )
135 self.act_fn = SiluAndMul()
@@ -107,7 +141,7 @@ class CohereMoeMLP(nn.Module):
141 return x
142
143
110 -class CohereMoeAttention(nn.Module):
144 +class Cohere2MoeAttention(nn.Module):
145 """Cohere MoE attention with sliding-window interleave."""
146
147 def __init__(
@@ -170,6 +204,19 @@ class CohereMoeAttention(nn.Module):
204 ):
205 self.sliding_window = config.sliding_window
206
207 + # Prefix-dense layers (layer_idx < first_k_dense_replace) have full
208 + # attention (no sliding window). When prefix_dense_sliding_window_pattern
209 + # == 1, they keep RoPE even though they are not sliding-window layers.
210 + first_k_dense_replace = getattr(config, "first_k_dense_replace", 0)
211 + prefix_dense_sliding_window_pattern = getattr(
212 + config, "prefix_dense_sliding_window_pattern", 1
213 + )
214 + self.force_rope = bool(
215 + first_k_dense_replace
216 + and prefix_dense_sliding_window_pattern == 1
217 + and self.layer_idx < first_k_dense_replace
218 + )
219 +
220 self.attn = Attention(
221 self.num_heads,
222 self.head_dim,
@@ -188,15 +235,15 @@ class CohereMoeAttention(nn.Module):
235 ) -> torch.Tensor:
236 qkv, _ = self.qkv_proj(hidden_states)
237 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
191 - if self.sliding_window:
238 + if self.sliding_window or self.force_rope:
239 q, k = self.rotary_emb(positions, q, k)
240 attn_output = self.attn(q, k, v)
241 output, _ = self.o_proj(attn_output)
242 return output
243
244
198 -class CohereMoe(nn.Module):
199 - """Tensor-parallel MoE block for CohereMoe with shared experts."""
245 +class Cohere2Moe(nn.Module):
246 + """Tensor-parallel MoE block for Cohere2Moe with shared experts."""
247
248 def __init__(
249 self,
@@ -234,7 +281,7 @@ class CohereMoe(nn.Module):
281 )
282
283 if hasattr(config, "num_shared_experts") and config.num_shared_experts > 0:
237 - self.shared_experts = CohereMoeMLP(
284 + self.shared_experts = Cohere2MoeMLP(
285 config=config,
286 intermediate_size=config.intermediate_size * config.num_shared_experts,
287 quant_config=quant_config,
@@ -276,7 +323,7 @@ class CohereMoe(nn.Module):
323 return final_hidden_states.view(orig_shape)
324
325
279 -class CohereMoeDecoderLayer(nn.Module):
326 +class Cohere2MoeDecoderLayer(nn.Module):
327 def __init__(
328 self,
329 config: CohereConfig,
@@ -287,19 +334,34 @@ class CohereMoeDecoderLayer(nn.Module):
334 super().__init__()
335 self.config = config
336 self.hidden_size = config.hidden_size
337 + self.layer_idx = extract_layer_index(prefix)
338
291 - self.self_attn = CohereMoeAttention(
339 + self.self_attn = Cohere2MoeAttention(
340 config,
341 cache_config,
342 quant_config=quant_config,
343 prefix=f"{prefix}.self_attn",
344 )
297 - self.mlp = CohereMoe(
298 - config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
299 - )
300 - self.input_layernorm = LayerNorm(
301 - param_shape=(config.hidden_size,), eps=config.layer_norm_eps
302 - )
345 +
346 + # Layers before first_k_dense_replace use a dense MLP instead of MoE.
347 + first_k_dense_replace = getattr(config, "first_k_dense_replace", 0)
348 + if self.layer_idx < first_k_dense_replace:
349 + self.mlp = Cohere2MoeMLP(
350 + config=config,
351 + intermediate_size=getattr(
352 + config, "prefix_dense_intermediate_size", config.intermediate_size
353 + ),
354 + quant_config=quant_config,
355 + reduce_results=True,
356 + prefix=f"{prefix}.mlp",
357 + )
358 + else:
359 + self.mlp = Cohere2Moe(
360 + config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
361 + )
362 +
363 + norm_cls, norm_eps = select_norm_impl(config)
364 + self.input_layernorm = norm_cls(param_shape=(config.hidden_size,), eps=norm_eps)
365
366 def forward(
367 self,
@@ -320,8 +382,8 @@ class CohereMoeDecoderLayer(nn.Module):
382
383
384 @support_torch_compile
323 -class CohereMoeModel(nn.Module):
324 - """Transformer decoder for CohereMoe."""
385 +class Cohere2MoeModel(nn.Module):
386 + """Transformer decoder for Cohere2Moe."""
387
388 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
389 super().__init__()
@@ -339,14 +401,13 @@ class CohereMoeModel(nn.Module):
401 )
402 self.start_layer, self.end_layer, self.layers = make_layers(
403 config.num_hidden_layers,
342 - lambda prefix: CohereMoeDecoderLayer(
404 + lambda prefix: Cohere2MoeDecoderLayer(
405 config, cache_config, quant_config, prefix=prefix
406 ),
407 prefix=f"{prefix}.layers",
408 )
347 - self.norm = LayerNorm(
348 - param_shape=(config.hidden_size,), eps=config.layer_norm_eps
349 - )
409 + norm_cls, norm_eps = select_norm_impl(config)
410 + self.norm = norm_cls(param_shape=(config.hidden_size,), eps=norm_eps)
411 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
412 ["hidden_states", "residual"], config.hidden_size
413 )
@@ -471,7 +532,7 @@ class CohereMoeModel(nn.Module):
532 return loaded_params
533
534
474 -class CohereMoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
535 +class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
536 is_text_generation_model = True
537
538 packed_modules_mapping = {
@@ -498,7 +559,7 @@ class CohereMoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
559 self.logits_processor = LogitsProcessor(
560 self.unpadded_vocab_size, config.vocab_size, scale=self.logits_scale
561 )
501 - self.model = CohereMoeModel(
562 + self.model = Cohere2MoeModel(
563 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
564 )
565 self.make_empty_intermediate_tensors = (
vllm/model_executor/models/cohere_eagle.py new
+247
@@ -0,0 +1,247 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Iterable
5 +
6 +import torch
7 +import torch.nn as nn
8 +from transformers import CohereConfig
9 +
10 +from vllm.compilation.decorators import support_torch_compile
11 +from vllm.config import VllmConfig
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.layers.linear import ReplicatedLinear
14 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
15 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
16 +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
17 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
18 +from vllm.model_executor.models.commandr import (
19 + CohereDecoderLayer,
20 + CohereForCausalLM,
21 + LayerNorm,
22 +)
23 +
24 +from .utils import (
25 + AutoWeightsLoader,
26 + get_draft_quant_config,
27 + maybe_prefix,
28 + process_eagle_weight,
29 +)
30 +
31 +logger = init_logger(__name__)
32 +
33 +
34 +class CohereEagleDecoderLayer(CohereDecoderLayer):
35 + """Eagle draft variant of CohereDecoderLayer."""
36 +
37 + def __init__(
38 + self,
39 + config: CohereConfig,
40 + cache_config=None,
41 + quant_config: QuantizationConfig | None = None,
42 + prefix: str = "",
43 + ) -> None:
44 + super().__init__(
45 + config,
46 + cache_config=cache_config,
47 + quant_config=quant_config,
48 + prefix=prefix,
49 + )
50 +
51 +
52 +@support_torch_compile
53 +class CohereEagleModel(nn.Module):
54 + def __init__(
55 + self,
56 + *,
57 + vllm_config: VllmConfig,
58 + prefix: str = "",
59 + start_layer_id: int = 0,
60 + ) -> None:
61 + super().__init__()
62 + self.config = vllm_config.speculative_config.draft_model_config.hf_config
63 + self.quant_config = get_draft_quant_config(vllm_config)
64 +
65 + # Cohere2-targeted EAGLE drafts inherit the target's sliding-window
66 + # attention pattern. ``CohereAttention`` resolves per-layer behavior
67 + # via ``config.layer_types[layer_idx]`` and the eagle layers use
68 + # absolute indices (target_layer_num + i), so prepend the target's
69 + # ``layer_types`` to the draft's so the lookup succeeds.
70 + target_text_config = vllm_config.model_config.hf_text_config
71 + if hasattr(target_text_config, "layer_types") and hasattr(
72 + self.config, "layer_types"
73 + ):
74 + self.config.layer_types = list(target_text_config.layer_types) + list(
75 + self.config.layer_types
76 + )
77 +
78 + self.vocab_size = self.config.vocab_size
79 + self.embed_tokens = VocabParallelEmbedding(
80 + self.config.vocab_size,
81 + self.config.hidden_size,
82 + prefix=maybe_prefix(prefix, "embed_tokens"),
83 + )
84 +
85 + self.layers = nn.ModuleList(
86 + [
87 + CohereEagleDecoderLayer(
88 + self.config,
89 + cache_config=vllm_config.cache_config,
90 + quant_config=self.quant_config,
91 + prefix=maybe_prefix(prefix, f"layers.{i + start_layer_id}"),
92 + )
93 + for i in range(self.config.num_hidden_layers)
94 + ]
95 + )
96 +
97 + # Cohere EAGLE checkpoints include a bias term on the input fusion
98 + # projection (unlike LLaMA EAGLE which uses bias=False).
99 + self.fc = ReplicatedLinear(
100 + input_size=self.config.hidden_size * 2,
101 + output_size=self.config.hidden_size,
102 + bias=True,
103 + params_dtype=vllm_config.model_config.dtype,
104 + quant_config=self.quant_config,
105 + prefix=maybe_prefix(prefix, "fc"),
106 + return_bias=False,
107 + )
108 +
109 + # Cohere EAGLE applies an explicit final LayerNorm to the draft
110 + # hidden states before they are consumed by the logits processor.
111 + self.norm = LayerNorm(
112 + param_shape=(self.config.hidden_size),
113 + eps=self.config.layer_norm_eps,
114 + )
115 +
116 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
117 + return self.embed_tokens(input_ids)
118 +
119 + def forward(
120 + self,
121 + input_ids: torch.Tensor,
122 + positions: torch.Tensor,
123 + hidden_states: torch.Tensor,
124 + ) -> tuple[torch.Tensor, torch.Tensor]:
125 + input_embeds = self.embed_tokens(input_ids)
126 + hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1))
127 + residual = None
128 + for layer in self.layers:
129 + hidden_states, residual = layer(
130 + positions,
131 + hidden_states,
132 + residual,
133 + )
134 + hidden_states, _ = self.norm(hidden_states, residual)
135 + return hidden_states, hidden_states
136 +
137 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
138 + stacked_params_mapping = [
139 + # (param_name, shard_name, shard_id)
140 + (".qkv_proj", ".q_proj", "q"),
141 + (".qkv_proj", ".k_proj", "k"),
142 + (".qkv_proj", ".v_proj", "v"),
143 + (".gate_up_proj", ".gate_proj", 0),
144 + (".gate_up_proj", ".up_proj", 1),
145 + ]
146 + params_dict = dict(self.named_parameters())
147 + loaded_params: set[str] = set()
148 +
149 + for name, loaded_weight in weights:
150 + if "rotary_emb.inv_freq" in name:
151 + continue
152 +
153 + if self.quant_config is not None and (
154 + scale_name := self.quant_config.get_cache_scale(name)
155 + ):
156 + param = params_dict[scale_name]
157 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
158 + loaded_weight = (
159 + loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
160 + )
161 + weight_loader(param, loaded_weight)
162 + loaded_params.add(scale_name)
163 + continue
164 +
165 + for param_name, weight_name, shard_id in stacked_params_mapping:
166 + if weight_name not in name:
167 + continue
168 + name = name.replace(weight_name, param_name)
169 + if name.endswith(".bias") and name not in params_dict:
170 + continue
171 + param = params_dict[name]
172 + weight_loader = param.weight_loader
173 + weight_loader(param, loaded_weight, shard_id)
174 + break
175 + else:
176 + if name.endswith(".bias") and name not in params_dict:
177 + continue
178 +
179 + param = params_dict[name]
180 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
181 + weight_loader(param, loaded_weight)
182 + loaded_params.add(name)
183 + return loaded_params
184 +
185 +
186 +class EagleCohereForCausalLM(CohereForCausalLM):
187 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
188 + nn.Module.__init__(self)
189 + self.config = vllm_config.speculative_config.draft_model_config.hf_config
190 + # Flags checked by the speculative proposer to decide whether to share
191 + # embed_tokens / lm_head with the target model. Cohere EAGLE checkpoints
192 + # use tied embeddings so these weights are absent from the draft file.
193 + self.has_own_embed_tokens = False
194 + self.has_own_lm_head = False
195 + target_layer_num = vllm_config.model_config.get_num_layers(
196 + vllm_config.parallel_config
197 + )
198 + self.model = CohereEagleModel(
199 + vllm_config=vllm_config,
200 + prefix=maybe_prefix(prefix, "model"),
201 + start_layer_id=target_layer_num,
202 + )
203 +
204 + logit_scale = getattr(self.config, "logit_scale", 1.0)
205 + self.logits_processor = LogitsProcessor(
206 + self.config.vocab_size, scale=logit_scale
207 + )
208 +
209 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
210 + return self.model.embed_input_ids(input_ids)
211 +
212 + def forward(
213 + self,
214 + input_ids: torch.Tensor,
215 + positions: torch.Tensor,
216 + hidden_states: torch.Tensor,
217 + inputs_embeds: torch.Tensor | None = None,
218 + ) -> tuple[torch.Tensor, torch.Tensor]:
219 + if inputs_embeds is not None:
220 + raise NotImplementedError(
221 + f"{type(self).__name__} does not support multimodal inputs yet."
222 + )
223 + return self.model(input_ids, positions, hidden_states)
224 +
225 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
226 + def _track_and_forward(inputs):
227 + name, weight = inputs
228 + process_eagle_weight(self, name)
229 + return name, weight
230 +
231 + loader = AutoWeightsLoader(
232 + self,
233 + skip_prefixes=(
234 + ["lm_head.", "model.embed_tokens."]
235 + if self.config.tie_word_embeddings
236 + else None
237 + ),
238 + )
239 +
240 + loaded_weight_names = loader.load_weights(map(_track_and_forward, weights))
241 +
242 + # Embed tokens are tied with the target model and therefore not
243 + # present in the EAGLE checkpoint; mark them as loaded explicitly to
244 + # avoid a spurious "weight not found" warning from the default
245 + # weight loader.
246 + loaded_weight_names.add("model.embed_tokens.weight")
247 + return loaded_weight_names
vllm/model_executor/models/registry.py
+2 -1
@@ -89,7 +89,7 @@ _TEXT_GENERATION_MODELS = {
89 "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
90 "CohereForCausalLM": ("commandr", "CohereForCausalLM"),
91 "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
92 - "CohereMoeForCausalLM": ("cohere_moe", "CohereMoeForCausalLM"),
92 + "Cohere2MoeForCausalLM": ("cohere2_moe", "Cohere2MoeForCausalLM"),
93 "CwmForCausalLM": ("llama", "LlamaForCausalLM"),
94 "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
95 "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
@@ -582,6 +582,7 @@ _SPECULATIVE_DECODING_MODELS = {
582 "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
583 "MiMoV2MTPModel": ("mimo_v2_mtp", "MiMoV2MTP"),
584 "MiMoV2OmniMTPModel": ("mimo_v2_mtp", "MiMoV2OmniMTP"),
585 + "EagleCohereForCausalLM": ("cohere_eagle", "EagleCohereForCausalLM"),
586 "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
587 "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
588 "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
vllm/reasoning/cohere_command_reasoning_parser.py
+32 -13
@@ -39,16 +39,20 @@ REPLACEMENT_CHAR = "\ufffd"
39
40
41 class CohereTagRegistry(NamedTuple):
42 - """A single ``structural_tag`` begin("trigger")/end pair."""
42 + """A single ``structural_tag`` trigger / end pair (``begin`` uses ``trigger``)."""
43
44 trigger: str
45 end: str
46
47
48 class CohereTagStyle(NamedTuple):
49 - """The structural tags style for a given model architecture."""
49 + """The structural tags style for a given model architecture.
50
51 - json: CohereTagRegistry
51 + ``json_tags`` lists every JSON-schema wrapper the model may emit (MOE uses
52 + both response and text delimiters). ``tools`` is the tool-call wrapper.
53 + """
54 +
55 + json_tags: tuple[CohereTagRegistry, ...]
56 tools: CohereTagRegistry
57
58
@@ -64,18 +68,30 @@ class CohereNormalizedTool(TypedDict):
68
69
70 COMMAND_A_TOOLS_TAG = CohereTagRegistry(
67 - trigger="<|START_ACTION|>", end="<|END_ACTION|>"
71 + trigger="<|START_ACTION|>",
72 + end="<|END_ACTION|>",
73 )
74 COMMAND_A_JSON_TAG = CohereTagRegistry(
70 - trigger="<|START_RESPONSE|>", end="<|END_RESPONSE|>"
75 + trigger="<|START_RESPONSE|>",
76 + end="<|END_RESPONSE|>",
77 +)
78 +COMMAND_A_PLUS_JSON_TAG = CohereTagRegistry(
79 + trigger="<|START_TEXT|>",
80 + end="<|END_TEXT|>",
81 )
82
83 MODEL_TO_TAG_STYLE: dict[str, CohereTagStyle] = {
84 "Cohere2ForCausalLM": CohereTagStyle(
75 - json=COMMAND_A_JSON_TAG, tools=COMMAND_A_TOOLS_TAG
85 + json_tags=(COMMAND_A_JSON_TAG,),
86 + tools=COMMAND_A_TOOLS_TAG,
87 ),
88 "Cohere2VisionForConditionalGeneration": CohereTagStyle(
78 - json=COMMAND_A_JSON_TAG, tools=COMMAND_A_TOOLS_TAG
89 + json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG),
90 + tools=COMMAND_A_TOOLS_TAG,
91 + ),
92 + "Cohere2MoeForCausalLM": CohereTagStyle(
93 + json_tags=(COMMAND_A_JSON_TAG,),
94 + tools=COMMAND_A_TOOLS_TAG,
95 ),
96 }
97
@@ -211,15 +227,18 @@ def convert_schema_to_structural_tags(
227 style = MODEL_TO_TAG_STYLE[model_architecture]
228
229 tags: list[dict] = []
230 + triggers: list[str] = []
231
232 def _add_tag(tag: CohereTagRegistry, content: dict) -> None:
233 tags.append({"begin": tag.trigger, "content": content, "end": tag.end})
234 + triggers.append(tag.trigger)
235
236 if schema is not None:
219 - # Add the JSON-schema tag both for schema-only requests and for the
220 - # "tools plus JSON mode" case (North use case: follow the schema when
221 - # the model decides not to call any tool).
222 - _add_tag(style.json, {"type": "json_schema", "json_schema": schema})
237 + # One structural tag per JSON wrapper (e.g. MOE: response + text).
238 + # Same for schema-only and "tools plus JSON mode" (North: schema when
239 + # the model does not call tools).
240 + for jt in style.json_tags:
241 + _add_tag(jt, {"type": "json_schema", "json_schema": schema})
242
243 if _has_effective_tools(tools):
244 # ``tools`` may be a JSON string (poseidon / RESPONSE_FORMAT_TOOL_DEFINITIONS)
@@ -240,7 +259,7 @@ def convert_schema_to_structural_tags(
259 "type": "structural_tag",
260 "format": {
261 "type": "triggered_tags",
243 - "triggers": [t["begin"] for t in tags],
262 + "triggers": triggers,
263 "tags": tags,
264 },
265 }
@@ -505,7 +524,7 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
524 model_architecture=model_architecture,
525 )
526 if result is None:
508 - # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``; conversion
527 + # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``.
528 raise ValueError(
529 "Failed to build structural_tag guided decoding constraints from "
530 "this request's JSON schema and/or tools. The configured model "
vllm/v1/spec_decode/llm_base_proposer.py
+1
@@ -1168,6 +1168,7 @@ class SpecDecodeBaseProposer:
1168 # handle multimodality
1169 assert hasattr(target_model, "config")
1170 if self.get_model_name(target_model) in [
1171 + "Cohere2VisionForConditionalGeneration",
1172 "Exaone4_5_ForConditionalGeneration",
1173 "GlmOcrForConditionalGeneration",
1174 "HunYuanVLForConditionalGeneration",