[Feature] Support EAGLE3 for Sarvam (#53052)
Signed-off-by: mohit-sarvam <mohit@sarvam.ai> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Codex <noreply@openai.com>
Mohit Singla committed
Sep 8, 2026 at 05:38 UTC
54da70c1eb6976fe655efd535f14a01378cd08bc
1 file changed
+40
-6
vllm/model_executor/models/sarvam.py
+40
-6
@@ -54,7 +54,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
54
from vllm.sequence import IntermediateTensors
55
56
from .bailing_moe import BailingMoeForCausalLM
57
-from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
57
+from .interfaces import (
58
+ EagleModelMixin,
59
+ MixtureOfExperts,
60
+ SupportsEagle3,
61
+ SupportsLoRA,
62
+ SupportsPP,
63
+)
64
from .utils import (
65
AutoWeightsLoader,
66
PPMissingLayer,
@@ -442,7 +448,9 @@ class SarvamMLABlock(nn.Module):
448
return hidden_states, residual
449
450
445
-class SarvamMLAModel(nn.Module):
451
+class SarvamMLAModel(nn.Module, EagleModelMixin):
452
+ """Sarvam MLA backbone with stage-local EAGLE3 auxiliary capture."""
453
+
454
hf_to_vllm_mapper = WeightsMapper(
455
orig_to_new_stacked={
456
# .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
@@ -508,7 +516,17 @@ class SarvamMLAModel(nn.Module):
516
positions: torch.Tensor,
517
intermediate_tensors: IntermediateTensors | None,
518
inputs_embeds: torch.Tensor | None = None,
511
- ) -> torch.Tensor | IntermediateTensors:
519
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
520
+ """Run this stage and optionally return its auxiliary hidden states.
521
+
522
+ Auxiliary captures are local to this stage, matching Qwen3 MoE.
523
+ Pipeline transport carries only hidden states and residual; EAGLE3
524
+ capture across pipeline stages is not supported by this model.
525
+
526
+ Returns:
527
+ Intermediate tensors on non-final stages. On the final stage,
528
+ normalized hidden states, paired with auxiliary states if captured.
529
+ """
530
if get_pp_group().is_first_rank:
531
if inputs_embeds is not None:
532
hidden_states = inputs_embeds
@@ -521,12 +539,22 @@ class SarvamMLAModel(nn.Module):
539
hidden_states = intermediate_tensors["hidden_states"]
540
residual = intermediate_tensors["residual"]
541
524
- for layer in islice(self.layers, self.start_layer, self.end_layer):
542
+ aux_hidden_states = self._maybe_add_hidden_state(
543
+ [], self.start_layer, hidden_states, residual
544
+ )
545
+ for layer_idx, layer in enumerate(
546
+ islice(self.layers, self.start_layer, self.end_layer),
547
+ start=self.start_layer,
548
+ ):
549
hidden_states, residual = layer(
550
hidden_states,
551
positions,
552
residual,
553
)
554
+ self._maybe_add_hidden_state(
555
+ aux_hidden_states, layer_idx + 1, hidden_states, residual
556
+ )
557
+
558
if not get_pp_group().is_last_rank:
559
return IntermediateTensors(
560
{"hidden_states": hidden_states, "residual": residual}
@@ -535,6 +563,9 @@ class SarvamMLAModel(nn.Module):
563
hidden_states = self.norm(hidden_states)
564
else:
565
hidden_states, _ = self.norm(hidden_states, residual)
566
+
567
+ if len(aux_hidden_states) > 0:
568
+ return hidden_states, aux_hidden_states
569
return hidden_states
570
571
def load_weights(
@@ -591,7 +622,9 @@ class SarvamMixtureOfExperts(MixtureOfExperts):
622
moe.set_eplb_state(eplb_state)
623
624
594
-class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfExperts):
625
+class SarvamMLAForCausalLM(
626
+ nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3, SarvamMixtureOfExperts
627
+):
628
packed_modules_mapping = {
629
"q_proj": ["q_proj"],
630
"q_a_proj": ["q_a_proj"],
@@ -658,7 +691,8 @@ class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfE
691
positions: torch.Tensor,
692
intermediate_tensors: IntermediateTensors | None = None,
693
inputs_embeds: torch.Tensor | None = None,
661
- ) -> torch.Tensor | IntermediateTensors:
694
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
695
+ """Return backbone outputs, including auxiliary states when captured."""
696
return self.model(
697
input_ids=input_ids,
698
positions=positions,