[Model][Spec Decode] Enable EAGLE3/DSpark pipeline parallelism for Sarvam MLA (#55902)
Signed-off-by: mohit-sarvam <mohit@sarvam.ai> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
Mohit Singla committed
Oct 2, 2026 at 09:58 UTC
0cbac6cd1305f710e12193596b27488397bcb205
1 file changed
+23
-13
vllm/model_executor/models/sarvam.py
+23
-13
@@ -73,6 +73,7 @@ from .utils import (
73
make_empty_intermediate_tensors_factory,
74
make_layers,
75
maybe_prefix,
76
+ spec_decode_needs_target_embed,
77
)
78
79
@@ -460,7 +461,9 @@ class SarvamMLABlock(nn.Module):
461
}
462
)
463
class SarvamMLAModel(nn.Module, EagleModelMixin):
463
- """Sarvam MLA backbone with stage-local EAGLE3 auxiliary capture."""
464
+ """Sarvam MLA backbone with EAGLE3 auxiliary capture across pipeline stages."""
465
+
466
+ supports_aux_hidden_states_over_pp = True
467
468
hf_to_vllm_mapper = WeightsMapper(
469
orig_to_new_stacked={
@@ -487,8 +490,10 @@ class SarvamMLAModel(nn.Module, EagleModelMixin):
490
self.vocab_size = config.vocab_size
491
self.embed_dim = config.hidden_size
492
self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)
490
- if get_pp_group().is_first_rank or (
491
- self.tie_word_embeddings and get_pp_group().is_last_rank
493
+ if (
494
+ get_pp_group().is_first_rank
495
+ or (self.tie_word_embeddings and get_pp_group().is_last_rank)
496
+ or spec_decode_needs_target_embed(vllm_config)
497
):
498
self.embed_tokens = VocabParallelEmbedding(
499
self.vocab_size,
@@ -530,13 +535,10 @@ class SarvamMLAModel(nn.Module, EagleModelMixin):
535
) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
536
"""Run this stage and optionally return its auxiliary hidden states.
537
533
- Auxiliary captures are local to this stage, matching Qwen3 MoE.
534
- Pipeline transport carries only hidden states and residual; EAGLE3
535
- capture across pipeline stages is not supported by this model.
536
-
538
Returns:
538
- Intermediate tensors on non-final stages. On the final stage,
539
- normalized hidden states, paired with auxiliary states if captured.
539
+ Intermediate tensors with local auxiliary states on non-final stages.
540
+ On the final stage, normalized hidden states paired with the ordered
541
+ auxiliary states from all stages if captured.
542
543
"""
544
if get_pp_group().is_first_rank:
@@ -551,9 +553,12 @@ class SarvamMLAModel(nn.Module, EagleModelMixin):
553
hidden_states = intermediate_tensors["hidden_states"]
554
residual = intermediate_tensors["residual"]
555
554
- aux_hidden_states = self._maybe_add_hidden_state(
555
- [], self.start_layer, hidden_states, residual
556
- )
556
+ remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
557
+ aux_hidden_states: list[torch.Tensor] = []
558
+ if get_pp_group().is_first_rank:
559
+ self._maybe_add_hidden_state(
560
+ aux_hidden_states, self.start_layer, hidden_states, residual
561
+ )
562
for layer_idx, layer in enumerate(
563
islice(self.layers, self.start_layer, self.end_layer),
564
start=self.start_layer,
@@ -569,13 +574,18 @@ class SarvamMLAModel(nn.Module, EagleModelMixin):
574
575
if not get_pp_group().is_last_rank:
576
return IntermediateTensors(
572
- {"hidden_states": hidden_states, "residual": residual}
577
+ {
578
+ "hidden_states": hidden_states,
579
+ "residual": residual,
580
+ **self.pack_local_aux_hidden_states(aux_hidden_states),
581
+ }
582
)
583
if residual is None:
584
hidden_states = self.norm(hidden_states)
585
else:
586
hidden_states, _ = self.norm(hidden_states, residual)
587
588
+ aux_hidden_states = remote_aux + aux_hidden_states
589
if len(aux_hidden_states) > 0:
590
return hidden_states, aux_hidden_states
591
return hidden_states