A high-throughput and memory-efficient inference and serving engine for LLMs
[Cohere] Fix Cohere2MoE weight loading when using Transformers ≥5.10 (#44747)
Signed-off-by: Terrencezzj <terrence@cohere.ai>
Terrence Zhao committed
Jun 9, 2026 at 09:27 UTC
3e8afdf78598afc8be999a6a049be3a5fe182e48
1 file changed
+27
-9
vllm/model_executor/models/cohere2_moe.py
+27
-9
@@ -51,6 +51,13 @@ from .utils import (
51
)
52
53
54
+def is_prefix_dense_layer(config: CohereConfig, layer_idx: int) -> bool:
55
+ """True when layer_idx lies in the contiguous dense MLP prefix."""
56
+ if layer_idx >= len(config.mlp_layer_types):
57
+ return False
58
+ return all(t == "dense" for t in config.mlp_layer_types[: layer_idx + 1])
59
+
60
+
61
@torch.compile(backend=current_platform.simple_compile_backend)
62
def token_choice_with_bias(
63
hidden_states: torch.Tensor,
@@ -207,17 +214,15 @@ class Cohere2MoeAttention(nn.Module):
214
):
215
self.sliding_window = config.sliding_window
216
210
- # Prefix-dense layers (layer_idx < first_k_dense_replace) have full
211
- # attention (no sliding window). When prefix_dense_sliding_window_pattern
212
- # == 1, they keep RoPE even though they are not sliding-window layers.
213
- first_k_dense_replace = getattr(config, "first_k_dense_replace", 0)
217
+ # Prefix-dense layers have full attention (no sliding window). When
218
+ # prefix_dense_sliding_window_pattern == 1, they keep RoPE even though
219
+ # they are not sliding-window layers.
220
prefix_dense_sliding_window_pattern = getattr(
221
config, "prefix_dense_sliding_window_pattern", 1
222
)
223
self.force_rope = bool(
218
- first_k_dense_replace
224
+ is_prefix_dense_layer(config, self.layer_idx)
225
and prefix_dense_sliding_window_pattern == 1
220
- and self.layer_idx < first_k_dense_replace
226
)
227
228
self.attn = Attention(
@@ -346,9 +351,7 @@ class Cohere2MoeDecoderLayer(nn.Module):
351
prefix=f"{prefix}.self_attn",
352
)
353
349
- # Layers before first_k_dense_replace use a dense MLP instead of MoE.
350
- first_k_dense_replace = getattr(config, "first_k_dense_replace", 0)
351
- if self.layer_idx < first_k_dense_replace:
354
+ if config.mlp_layer_types[self.layer_idx] == "dense":
355
self.mlp = Cohere2MoeMLP(
356
config=config,
357
intermediate_size=getattr(
@@ -402,6 +405,21 @@ class Cohere2MoeModel(nn.Module):
405
self.embed_tokens = VocabParallelEmbedding(
406
config.vocab_size, config.hidden_size
407
)
408
+
409
+ # Decoder layers read per-layer MLP layout from config.mlp_layer_types
410
+ # (dense MLP vs MoE) and use it for weight loading. Transformers >=5.10
411
+ # populates this field; older versions only expose first_k_dense_replace.
412
+ # Normalize here so layer construction below sees a consistent layout.
413
+ if getattr(config, "mlp_layer_types", None) is None:
414
+ first_k_dense_replace = getattr(config, "first_k_dense_replace", None)
415
+ n = config.num_hidden_layers
416
+ if first_k_dense_replace is not None:
417
+ config.mlp_layer_types = ["dense"] * first_k_dense_replace + [
418
+ "sparse"
419
+ ] * (n - first_k_dense_replace)
420
+ else:
421
+ config.mlp_layer_types = ["sparse"] * n
422
+
423
self.start_layer, self.end_layer, self.layers = make_layers(
424
config.num_hidden_layers,
425
lambda prefix: Cohere2MoeDecoderLayer(