[Model] Support speculative decoding method for PLaMo3 (#54239)
Signed-off-by: Shinichi Hemmi <shemmi@preferred.jp>
Shinichi Hemmi committed
Aug 29, 2026 at 07:31 UTC
df14152ac6b09f76345dae05b99541c3e87f8d35
2 files changed
+101
-10
tests/model_executor/test_plamo3.py
new
+72
@@ -0,0 +1,72 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+from torch import nn
9
+
10
+from vllm.model_executor.models import plamo3 as plamo3_mod
11
+from vllm.model_executor.models.interfaces import supports_eagle3
12
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
13
+ set_eagle3_aux_hidden_state_layers,
14
+)
15
+
16
+
17
+class DummyDecoderLayer(nn.Module):
18
+ def forward(self, positions, hidden_states, residual):
19
+ del positions
20
+ return hidden_states + 1, hidden_states + 2
21
+
22
+
23
+class DummyNorm(nn.Module):
24
+ def forward(self, hidden_states, residual):
25
+ return hidden_states + residual, None
26
+
27
+
28
+@pytest.mark.cpu_test
29
+def test_plamo3_returns_dflash_selected_auxiliary_hidden_states(monkeypatch):
30
+ decoder = plamo3_mod.Plamo3Decoder.__new__(plamo3_mod.Plamo3Decoder)
31
+ nn.Module.__init__(decoder)
32
+ decoder.start_layer = 0
33
+ decoder.end_layer = 2
34
+ decoder.layers = nn.ModuleList([DummyDecoderLayer() for _ in range(8)])
35
+
36
+ model = plamo3_mod.Plamo3Model.__new__(plamo3_mod.Plamo3Model)
37
+ nn.Module.__init__(model)
38
+ model.layers = decoder
39
+ model.norm = DummyNorm()
40
+ model.do_not_compile = True
41
+
42
+ target = plamo3_mod.Plamo3ForCausalLM.__new__(plamo3_mod.Plamo3ForCausalLM)
43
+ nn.Module.__init__(target)
44
+ target.model = model
45
+
46
+ monkeypatch.setattr(
47
+ plamo3_mod,
48
+ "get_pp_group",
49
+ lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
50
+ )
51
+
52
+ assert supports_eagle3(target)
53
+ assert target.get_eagle3_default_aux_hidden_state_layers() == (2, 4, 5)
54
+
55
+ spec_config = SimpleNamespace(
56
+ draft_model_config=SimpleNamespace(
57
+ hf_config=SimpleNamespace(dflash_config={"target_layer_ids": [0, 1]})
58
+ )
59
+ )
60
+ set_eagle3_aux_hidden_state_layers(target, spec_config)
61
+ assert decoder.aux_hidden_state_layers == (1, 2)
62
+
63
+ output, aux_hidden_states = target(
64
+ input_ids=None,
65
+ positions=torch.tensor([0]),
66
+ inputs_embeds=torch.tensor([[1.0, 2.0]]),
67
+ )
68
+
69
+ torch.testing.assert_close(output, torch.tensor([[7.0, 9.0]]))
70
+ assert len(aux_hidden_states) == 2
71
+ torch.testing.assert_close(aux_hidden_states[0], torch.tensor([[5.0, 7.0]]))
72
+ torch.testing.assert_close(aux_hidden_states[1], torch.tensor([[7.0, 9.0]]))
vllm/model_executor/models/plamo3.py
+29
-10
@@ -35,7 +35,12 @@ from vllm.model_executor.model_loader.weight_utils import (
35
composed_weight_loader,
36
default_weight_loader,
37
)
38
-from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
38
+from vllm.model_executor.models.interfaces import (
39
+ EagleModelMixin,
40
+ SupportsEagle3,
41
+ SupportsLoRA,
42
+ SupportsPP,
43
+)
44
from vllm.model_executor.models.utils import (
45
AutoWeightsLoader,
46
extract_layer_index,
@@ -285,7 +290,7 @@ class Plamo3DecoderLayer(nn.Module):
290
return hidden_states, residual
291
292
288
-class Plamo3Decoder(torch.nn.Module):
293
+class Plamo3Decoder(torch.nn.Module, EagleModelMixin):
294
def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
295
super().__init__()
296
num_hidden_layers = vllm_config.model_config.hf_config.num_hidden_layers
@@ -296,23 +301,32 @@ class Plamo3Decoder(torch.nn.Module):
301
prefix=f"{prefix}.layers",
302
)
303
304
+ def __len__(self) -> int:
305
+ return len(self.layers)
306
+
307
def forward(
308
self,
309
positions: torch.Tensor,
310
hidden_states: torch.Tensor,
311
residual: torch.Tensor | None,
304
- ) -> tuple[torch.Tensor, torch.Tensor | None]:
305
- for layer in islice(self.layers, self.start_layer, self.end_layer):
312
+ ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
313
+ aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
314
+ for idx, layer in enumerate(
315
+ islice(self.layers, self.start_layer, self.end_layer)
316
+ ):
317
hidden_states, residual = layer(
318
positions=positions,
319
hidden_states=hidden_states,
320
residual=residual,
321
)
311
- return hidden_states, residual
322
+ self._maybe_add_hidden_state(
323
+ aux_hidden_states, idx + 1, hidden_states, residual
324
+ )
325
+ return hidden_states, residual, aux_hidden_states
326
327
328
@support_torch_compile
315
-class Plamo3Model(nn.Module):
329
+class Plamo3Model(nn.Module, EagleModelMixin):
330
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
331
super().__init__()
332
config = vllm_config.model_config.hf_config
@@ -340,13 +354,16 @@ class Plamo3Model(nn.Module):
354
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
355
return self.embed_tokens(input_ids)
356
357
+ def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
358
+ self.layers._set_aux_hidden_state_layers(layers)
359
+
360
def forward(
361
self,
362
input_ids: torch.Tensor | None,
363
positions: torch.Tensor,
364
intermediate_tensors: IntermediateTensors | None = None,
365
inputs_embeds: torch.Tensor | None = None,
349
- ) -> torch.Tensor:
366
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
367
if get_pp_group().is_first_rank:
368
if inputs_embeds is not None:
369
hidden_states = inputs_embeds
@@ -358,7 +375,7 @@ class Plamo3Model(nn.Module):
375
hidden_states = intermediate_tensors["hidden_states"]
376
residual = intermediate_tensors["residual"]
377
361
- hidden_states, residual = self.layers(
378
+ hidden_states, residual, aux_hidden_states = self.layers(
379
positions=positions, hidden_states=hidden_states, residual=residual
380
)
381
if not get_pp_group().is_last_rank:
@@ -366,10 +383,12 @@ class Plamo3Model(nn.Module):
383
{"hidden_states": hidden_states, "residual": residual}
384
)
385
hidden_states, _ = self.norm(hidden_states, residual)
386
+ if len(aux_hidden_states) > 0:
387
+ return hidden_states, aux_hidden_states
388
return hidden_states
389
390
372
-class Plamo3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
391
+class Plamo3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
392
packed_modules_mapping = {
393
"qkv_proj": ["qkv_proj"],
394
"gate_up_proj": ["gate_up_proj"],
@@ -416,7 +435,7 @@ class Plamo3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
435
positions: torch.Tensor,
436
intermediate_tensors: IntermediateTensors | None = None,
437
inputs_embeds: torch.Tensor | None = None,
419
- ) -> torch.Tensor:
438
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
439
hidden_states = self.model(
440
input_ids, positions, intermediate_tensors, inputs_embeds
441
)