[Model] Add Cohere2MoE Eagle3 auxiliary hidden states (#49819)

Signed-off-by: Doug Brown <doug@github.douggo.com>

tests/v1/spec_decode/test_dflash_causality.py

@@ -103,6 +103,18 @@ def test_dflash_fc_uses_aux_layer_count():
     assert _get_dflash_fc_input_size(vllm_config) == 3 * 4096
 
 
+def test_dflash_target_layers_map_to_post_layer_boundaries():
+    vllm_config = _vllm_config(
+        dflash_config={"target_layer_ids": [0, 16, 31]},
+    )
+
+    assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == (
+        1,
+        17,
+        32,
+    )
+
+
 @pytest.mark.parametrize("config_name", ["dflash_config", "eagle_config"])
 def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name):
     layer_ids = [1, 17, 32]

vllm/model_executor/models/cohere2_moe.py

@@ -35,7 +35,7 @@ from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 
 from .commandr import LayerNorm
-from .interfaces import SupportsPP, SupportsQuant
+from .interfaces import EagleModelMixin, SupportsEagle3, SupportsPP, SupportsQuant
 from .utils import (
     AutoWeightsLoader,
     WeightsMapper,
@@ -385,7 +385,7 @@ class Cohere2MoeDecoderLayer(nn.Module):
 
 
 @support_torch_compile
-class Cohere2MoeModel(nn.Module):
+class Cohere2MoeModel(nn.Module, EagleModelMixin):
     """Transformer decoder for Cohere2Moe."""
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -439,7 +439,7 @@ class Cohere2MoeModel(nn.Module):
         positions: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
-    ) -> torch.Tensor | IntermediateTensors:
+    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
         if get_pp_group().is_first_rank:
             if inputs_embeds is not None:
                 hidden_states = inputs_embeds
@@ -450,17 +450,28 @@ class Cohere2MoeModel(nn.Module):
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
-        for layer in islice(self.layers, self.start_layer, self.end_layer):
+
+        aux_hidden_states = []
+        if self.start_layer in self.aux_hidden_state_layers:
+            aux_hidden_states.append(hidden_states)
+        for layer_idx, layer in enumerate(
+            islice(self.layers, self.start_layer, self.end_layer),
+            start=self.start_layer,
+        ):
             hidden_states, residual = layer(positions, hidden_states, residual)
+            if layer_idx + 1 in self.aux_hidden_state_layers:
+                aux_hidden_states.append(hidden_states)
         if not get_pp_group().is_last_rank:
             return IntermediateTensors(
                 {"hidden_states": hidden_states, "residual": residual}
             )
         hidden_states, _ = self.norm(hidden_states, residual)
+        if aux_hidden_states:
+            return hidden_states, aux_hidden_states
         return hidden_states
 
 
-class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
+class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant, SupportsEagle3):
     is_text_generation_model = True
 
     hf_to_vllm_mapper = WeightsMapper(
@@ -520,7 +531,7 @@ class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
         positions: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
-    ) -> torch.Tensor | IntermediateTensors:
+    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
         return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
 
     def compute_logits(