| 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]])) |