[Model] Add Cohere2MoE Eagle3 auxiliary hidden states (#49819)
Signed-off-by: Doug Brown <doug@github.douggo.com>
Doug Brown committed
Sep 17, 2026 at 16:09 UTC
2b02c6c29b72147de18905ba4e7dc8f284655b40
2 files changed
+29
-6
tests/v1/spec_decode/test_dflash_causality.py
+12
@@ -103,6 +103,18 @@ def test_dflash_fc_uses_aux_layer_count():
103
assert _get_dflash_fc_input_size(vllm_config) == 3 * 4096
104
105
106
+def test_dflash_target_layers_map_to_post_layer_boundaries():
107
+ vllm_config = _vllm_config(
108
+ dflash_config={"target_layer_ids": [0, 16, 31]},
109
+ )
110
+
111
+ assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == (
112
+ 1,
113
+ 17,
114
+ 32,
115
+ )
116
+
117
+
118
@pytest.mark.parametrize("config_name", ["dflash_config", "eagle_config"])
119
def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name):
120
layer_ids = [1, 17, 32]
vllm/model_executor/models/cohere2_moe.py
+17
-6
@@ -35,7 +35,7 @@ from vllm.platforms import current_platform
35
from vllm.sequence import IntermediateTensors
36
37
from .commandr import LayerNorm
38
-from .interfaces import SupportsPP, SupportsQuant
38
+from .interfaces import EagleModelMixin, SupportsEagle3, SupportsPP, SupportsQuant
39
from .utils import (
40
AutoWeightsLoader,
41
WeightsMapper,
@@ -385,7 +385,7 @@ class Cohere2MoeDecoderLayer(nn.Module):
385
386
387
@support_torch_compile
388
-class Cohere2MoeModel(nn.Module):
388
+class Cohere2MoeModel(nn.Module, EagleModelMixin):
389
"""Transformer decoder for Cohere2Moe."""
390
391
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -439,7 +439,7 @@ class Cohere2MoeModel(nn.Module):
439
positions: torch.Tensor,
440
intermediate_tensors: IntermediateTensors | None = None,
441
inputs_embeds: torch.Tensor | None = None,
442
- ) -> torch.Tensor | IntermediateTensors:
442
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
443
if get_pp_group().is_first_rank:
444
if inputs_embeds is not None:
445
hidden_states = inputs_embeds
@@ -450,17 +450,28 @@ class Cohere2MoeModel(nn.Module):
450
assert intermediate_tensors is not None
451
hidden_states = intermediate_tensors["hidden_states"]
452
residual = intermediate_tensors["residual"]
453
- for layer in islice(self.layers, self.start_layer, self.end_layer):
453
+
454
+ aux_hidden_states = []
455
+ if self.start_layer in self.aux_hidden_state_layers:
456
+ aux_hidden_states.append(hidden_states)
457
+ for layer_idx, layer in enumerate(
458
+ islice(self.layers, self.start_layer, self.end_layer),
459
+ start=self.start_layer,
460
+ ):
461
hidden_states, residual = layer(positions, hidden_states, residual)
462
+ if layer_idx + 1 in self.aux_hidden_state_layers:
463
+ aux_hidden_states.append(hidden_states)
464
if not get_pp_group().is_last_rank:
465
return IntermediateTensors(
466
{"hidden_states": hidden_states, "residual": residual}
467
)
468
hidden_states, _ = self.norm(hidden_states, residual)
469
+ if aux_hidden_states:
470
+ return hidden_states, aux_hidden_states
471
return hidden_states
472
473
463
-class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
474
+class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant, SupportsEagle3):
475
is_text_generation_model = True
476
477
hf_to_vllm_mapper = WeightsMapper(
@@ -520,7 +531,7 @@ class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
531
positions: torch.Tensor,
532
intermediate_tensors: IntermediateTensors | None = None,
533
inputs_embeds: torch.Tensor | None = None,
523
- ) -> torch.Tensor | IntermediateTensors:
534
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
535
return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
536
537
def compute_logits(