diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index 9b5a784cef..3ce94b3f32 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -54,7 +54,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( from vllm.sequence import IntermediateTensors from .bailing_moe import BailingMoeForCausalLM -from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP +from .interfaces import ( + EagleModelMixin, + MixtureOfExperts, + SupportsEagle3, + SupportsLoRA, + SupportsPP, +) from .utils import ( AutoWeightsLoader, PPMissingLayer, @@ -442,7 +448,9 @@ class SarvamMLABlock(nn.Module): return hidden_states, residual -class SarvamMLAModel(nn.Module): +class SarvamMLAModel(nn.Module, EagleModelMixin): + """Sarvam MLA backbone with stage-local EAGLE3 auxiliary capture.""" + hf_to_vllm_mapper = WeightsMapper( orig_to_new_stacked={ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP @@ -508,7 +516,17 @@ class SarvamMLAModel(nn.Module): positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + """Run this stage and optionally return its auxiliary hidden states. + + Auxiliary captures are local to this stage, matching Qwen3 MoE. + Pipeline transport carries only hidden states and residual; EAGLE3 + capture across pipeline stages is not supported by this model. + + Returns: + Intermediate tensors on non-final stages. On the final stage, + normalized hidden states, paired with auxiliary states if captured. + """ if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -521,12 +539,22 @@ class SarvamMLAModel(nn.Module): 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 = self._maybe_add_hidden_state( + [], self.start_layer, hidden_states, residual + ) + for layer_idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): hidden_states, residual = layer( hidden_states, positions, residual, ) + self._maybe_add_hidden_state( + aux_hidden_states, layer_idx + 1, hidden_states, residual + ) + if not get_pp_group().is_last_rank: return IntermediateTensors( {"hidden_states": hidden_states, "residual": residual} @@ -535,6 +563,9 @@ class SarvamMLAModel(nn.Module): hidden_states = self.norm(hidden_states) else: hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states return hidden_states def load_weights( @@ -591,7 +622,9 @@ class SarvamMixtureOfExperts(MixtureOfExperts): moe.set_eplb_state(eplb_state) -class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfExperts): +class SarvamMLAForCausalLM( + nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3, SarvamMixtureOfExperts +): packed_modules_mapping = { "q_proj": ["q_proj"], "q_a_proj": ["q_a_proj"], @@ -658,7 +691,8 @@ class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfE 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 backbone outputs, including auxiliary states when captured.""" return self.model( input_ids=input_ids, positions=positions,