[Cohere] Misc changes to cohere model definitions (#50156)
Co-authored-by: Cursor <cursoragent@cursor.com>
Komal Kumar Teru committed
Aug 18, 2026 at 00:59 UTC
aa9903490c616dc6871e5acc62cec7bb1e5e9434
2 files changed
+87
-26
vllm/model_executor/models/cohere2_moe.py
+1
-1
@@ -208,7 +208,7 @@ class Cohere2MoeAttention(nn.Module):
208
layer_types is not None
209
and layer_types[self.layer_idx] == "sliding_attention"
210
):
211
- self.sliding_window = config.sliding_window
211
+ self.sliding_window = config.sliding_window + 1
212
213
# Prefix-dense layers have full attention (no sliding window). When
214
# prefix_dense_sliding_window_pattern == 1, they keep RoPE even though
vllm/model_executor/models/commandr.py
+86
-25
@@ -32,7 +32,11 @@ from transformers import Cohere2Config, CohereConfig
32
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
35
+from vllm.distributed import (
36
+ get_pp_group,
37
+ get_tensor_model_parallel_world_size,
38
+ tensor_model_parallel_all_reduce,
39
+)
40
from vllm.model_executor.layers.activation import SiluAndMul
41
from vllm.model_executor.layers.attention import Attention
42
from vllm.model_executor.layers.linear import (
@@ -44,9 +48,7 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
48
from vllm.model_executor.layers.quantization import QuantizationConfig
49
from vllm.model_executor.layers.rotary_embedding import get_rope
50
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
47
-from vllm.model_executor.model_loader.weight_utils import (
48
- row_parallel_weight_loader,
49
-)
51
+from vllm.model_executor.model_loader.weight_utils import row_parallel_weight_loader
52
from vllm.model_executor.utils import set_weight_attrs
53
from vllm.platforms import current_platform
54
from vllm.sequence import IntermediateTensors
@@ -87,18 +89,61 @@ class LayerNorm(nn.Module):
89
return hidden_states, residuals
90
91
92
+@torch.compile(backend=current_platform.simple_compile_backend)
93
+def rms_norm_func(hidden_states, weight, variance_epsilon):
94
+ input_dtype = hidden_states.dtype
95
+ hidden_states = hidden_states.to(torch.float32)
96
+
97
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
98
+ hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
99
+
100
+ hidden_states = weight.to(torch.float32) * hidden_states
101
+ return hidden_states.to(input_dtype)
102
+
103
+
104
+class RMSNorm(nn.Module):
105
+ def __init__(self, param_shape=None, eps=1e-6):
106
+ super().__init__()
107
+ self.weight = nn.Parameter(torch.ones(param_shape))
108
+ self.variance_epsilon = eps
109
+ set_weight_attrs(self.weight, {"weight_loader": row_parallel_weight_loader})
110
+
111
+ def forward(self, hidden_states, residuals=None):
112
+ hidden_states = rms_norm_func(hidden_states, self.weight, self.variance_epsilon)
113
+ return hidden_states, residuals
114
+
115
+
116
+def select_norm_impl(config: CohereConfig) -> tuple[type[nn.Module], float]:
117
+ """
118
+ Returns the normalization layer class and epsilon value to use.
119
+ If `config.rms_norm_eps` is present, use RMSNorm.
120
+ Otherwise default to LayerNorm with `config.layer_norm_eps`.
121
+ """
122
+ rms_eps = getattr(config, "rms_norm_eps", None)
123
+ if rms_eps is not None:
124
+ return RMSNorm, rms_eps
125
+
126
+ return LayerNorm, config.layer_norm_eps
127
+
128
+
129
# Copied from transformers.models.llama.modeling_llama.LlamaMLP Llama->Cohere
130
class CohereMLP(nn.Module):
131
def __init__(
132
self,
94
- config: CohereConfig | Cohere2Config,
133
+ config: CohereConfig,
134
+ intermediate_size: int | None = None,
135
quant_config: QuantizationConfig | None = None,
136
+ # so we can reduce the attention and MLP outputs together
137
+ reduce_results: bool = False,
138
prefix: str = "",
139
):
140
super().__init__()
141
self.config = config
142
self.hidden_size = config.hidden_size
101
- self.intermediate_size = config.intermediate_size
143
+ if intermediate_size is None:
144
+ self.intermediate_size = config.intermediate_size
145
+ else:
146
+ self.intermediate_size = intermediate_size
147
self.gate_up_proj = MergedColumnParallelLinear(
148
self.hidden_size,
149
[self.intermediate_size] * 2,
@@ -111,6 +156,7 @@ class CohereMLP(nn.Module):
156
self.hidden_size,
157
bias=False,
158
quant_config=quant_config,
159
+ reduce_results=reduce_results,
160
prefix=f"{prefix}.down_proj",
161
)
162
self.act_fn = SiluAndMul()
@@ -133,11 +179,15 @@ class CohereAttention(nn.Module):
179
super().__init__()
180
tp_size = get_tensor_model_parallel_world_size()
181
self.config = config
182
+ self.layer_idx = extract_layer_index(prefix)
183
self.attention_dropout = config.attention_dropout
184
self.hidden_size = config.hidden_size
185
self.total_num_heads = config.num_attention_heads
186
self.num_heads = self.total_num_heads // tp_size
140
- self.head_dim = self.hidden_size // self.total_num_heads
187
+ if hasattr(config, "head_dim"):
188
+ self.head_dim = config.head_dim
189
+ else:
190
+ self.head_dim = self.hidden_size // self.total_num_heads
191
self.total_num_kv_heads = config.num_key_value_heads
192
if self.total_num_kv_heads >= tp_size:
193
# Number of KV heads is greater than TP size, so we partition
@@ -171,6 +221,9 @@ class CohereAttention(nn.Module):
221
bias=False,
222
quant_config=quant_config,
223
prefix=f"{prefix}.o_proj",
224
+ # NOTE: reduction will happen in the decoder layer forward
225
+ # so we can combine the attention and MLP outputs together
226
+ reduce_results=False,
227
)
228
self.rotary_emb = get_rope(
229
self.head_dim,
@@ -182,11 +235,14 @@ class CohereAttention(nn.Module):
235
# Model v2 has interleaved sliding windows, v1 does not
236
self.v1 = isinstance(config, CohereConfig)
237
238
+ # cohere swa layer sees [pos - sliding_window, pos], i.e. sliding_window + 1
239
+ # tokens. vLLM's FlashAttention backend does (value - 1, 0), so we pass
240
+ # sliding_window + 1 here to match the training convention. The same +1
241
+ # propagates into the KV-cache eviction formula
242
+ # (single_type_kv_cache_manager.py), keeping both paths consistent.
243
self.sliding_window = None
186
- if not self.v1:
187
- layer_idx = extract_layer_index(prefix)
188
- if config.layer_types[layer_idx] == "sliding_attention":
189
- self.sliding_window = config.sliding_window
244
+ if not self.v1 and config.layer_types[self.layer_idx] == "sliding_attention":
245
+ self.sliding_window = config.sliding_window + 1
246
247
self.attn = Attention(
248
self.num_heads,
@@ -198,13 +254,13 @@ class CohereAttention(nn.Module):
254
per_layer_sliding_window=self.sliding_window,
255
prefix=f"{prefix}.attn",
256
)
257
+ norm_cls, norm_eps = select_norm_impl(config)
258
if self.use_qk_norm:
202
- self.q_norm = LayerNorm(
203
- param_shape=(self.num_heads, self.head_dim), eps=config.layer_norm_eps
259
+ self.q_norm = norm_cls(
260
+ param_shape=(self.num_heads, self.head_dim), eps=norm_eps
261
)
205
- self.k_norm = LayerNorm(
206
- param_shape=(self.num_kv_heads, self.head_dim),
207
- eps=config.layer_norm_eps,
262
+ self.k_norm = norm_cls(
263
+ param_shape=(self.num_kv_heads, self.head_dim), eps=norm_eps
264
)
265
266
def _apply_qk_norm(self, q, k):
@@ -242,18 +298,16 @@ class CohereDecoderLayer(nn.Module):
298
):
299
super().__init__()
300
self.hidden_size = config.hidden_size
245
-
301
+ self.tp_size = get_tensor_model_parallel_world_size()
302
self.self_attn = CohereAttention(
303
config,
304
cache_config,
305
quant_config=quant_config,
306
prefix=f"{prefix}.self_attn",
307
)
252
-
308
self.mlp = CohereMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp")
254
- self.input_layernorm = LayerNorm(
255
- param_shape=(config.hidden_size), eps=config.layer_norm_eps
256
- )
309
+ norm_cls, norm_eps = select_norm_impl(config)
310
+ self.input_layernorm = norm_cls(param_shape=(config.hidden_size), eps=norm_eps)
311
312
def forward(
313
self,
@@ -269,8 +323,14 @@ class CohereDecoderLayer(nn.Module):
323
hidden_states=hidden_states,
324
)
325
hidden_states_mlp = self.mlp(hidden_states)
326
+ parallel_block_output = hidden_states_attention + hidden_states_mlp
327
# Add everything together
273
- hidden_states = residual + hidden_states_attention + hidden_states_mlp
328
+ if self.tp_size > 1:
329
+ # do the reduction in 1 shot instead of 2 separate all reduce
330
+ parallel_block_output = tensor_model_parallel_all_reduce(
331
+ parallel_block_output
332
+ )
333
+ hidden_states = residual + parallel_block_output
334
335
return hidden_states, residual
336
@@ -299,9 +359,10 @@ class CohereModel(nn.Module):
359
),
360
prefix=f"{prefix}.layers",
361
)
302
- self.norm = LayerNorm(
303
- param_shape=(config.hidden_size), eps=config.layer_norm_eps
304
- )
362
+
363
+ norm_cls, norm_eps = select_norm_impl(config)
364
+ self.norm = norm_cls(param_shape=(config.hidden_size), eps=norm_eps)
365
+
366
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
367
["hidden_states", "residual"], config.hidden_size
368
)