A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 03:45 UTC
Laguna xs dflash support (#41880)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
MeganEFlynn committed
May 7, 2026 at 13:31 UTC
969fbfb4a985ad763a6109851bad9dfc94fea289
2 files changed
+21
-4
vllm/config/speculative.py
+1
@@ -1003,6 +1003,7 @@ class SpeculativeConfig:
1003
"kimi_k25",
1004
"minimax_m2",
1005
"gemma4",
1006
+ "laguna",
1007
]
1008
if (
1009
self.method in ("eagle3", "extract_hidden_states", "dflash")
vllm/model_executor/models/laguna.py
+20
-4
@@ -39,7 +39,12 @@ from vllm.model_executor.model_loader.weight_utils import (
39
default_weight_loader,
40
maybe_remap_kv_scale_name,
41
)
42
-from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
42
+from vllm.model_executor.models.interfaces import (
43
+ EagleModelMixin,
44
+ SupportsEagle3,
45
+ SupportsLoRA,
46
+ SupportsPP,
47
+)
48
from vllm.model_executor.models.utils import (
49
AutoWeightsLoader,
50
PPMissingLayer,
@@ -554,7 +559,7 @@ class LagunaDecoderLayer(nn.Module):
559
560
561
@support_torch_compile
557
-class LagunaModel(nn.Module):
562
+class LagunaModel(nn.Module, EagleModelMixin):
563
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
564
super().__init__()
565
@@ -633,8 +638,17 @@ class LagunaModel(nn.Module):
638
hidden_states = intermediate_tensors["hidden_states"]
639
residual = intermediate_tensors["residual"]
640
636
- for layer in islice(self.layers, self.start_layer, self.end_layer):
641
+ aux_hidden_states = self._maybe_add_hidden_state(
642
+ [], self.start_layer, hidden_states, residual
643
+ )
644
+ for layer_idx, layer in enumerate(
645
+ islice(self.layers, self.start_layer, self.end_layer),
646
+ start=self.start_layer,
647
+ ):
648
hidden_states, residual = layer(positions, hidden_states, residual)
649
+ self._maybe_add_hidden_state(
650
+ aux_hidden_states, layer_idx + 1, hidden_states, residual
651
+ )
652
653
if not get_pp_group().is_last_rank:
654
return IntermediateTensors(
@@ -642,6 +656,8 @@ class LagunaModel(nn.Module):
656
)
657
658
hidden_states, _ = self.norm(hidden_states, residual)
659
+ if len(aux_hidden_states) > 0:
660
+ return hidden_states, aux_hidden_states
661
return hidden_states
662
663
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
@@ -821,7 +837,7 @@ class LagunaModel(nn.Module):
837
return loaded_params
838
839
824
-class LagunaForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
840
+class LagunaForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3):
841
fall_back_to_pt_during_load = False
842
843
packed_modules_mapping = {