diff --git a/tests/v1/spec_decode/test_dflash_causality.py b/tests/v1/spec_decode/test_dflash_causality.py index 0f1de86cf9..dc13e55700 100644 --- a/tests/v1/spec_decode/test_dflash_causality.py +++ b/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] diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index d657b78a5d..98e649dec6 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/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(