| 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 | import torch.nn as nn |
| 9 | |
| 10 | from vllm.config.compilation import CompilationMode |
| 11 | from vllm.model_executor.models import deepseek_v2 as deepseek_mod |
| 12 | from vllm.model_executor.models import mistral_large_3_eagle as eagle_mod |
| 13 | |
| 14 | |
| 15 | class DummyPPGroup: |
| 16 | world_size = 1 |
| 17 | is_first_rank = True |
| 18 | is_last_rank = True |
| 19 | |
| 20 | |
| 21 | class DummyEmbedding(nn.Module): |
| 22 | def __init__(self, vocab_size, hidden_size, *args, **kwargs): |
| 23 | super().__init__() |
| 24 | self.hidden_size = hidden_size |
| 25 | |
| 26 | def forward(self, input_ids): |
| 27 | return torch.zeros( |
| 28 | (*input_ids.shape, self.hidden_size), |
| 29 | dtype=torch.float32, |
| 30 | device=input_ids.device, |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | class DummyLinear(nn.Module): |
| 35 | def __init__(self, in_features, out_features, *args, **kwargs): |
| 36 | super().__init__() |
| 37 | self.out_features = out_features |
| 38 | |
| 39 | def forward(self, x): |
| 40 | return torch.zeros( |
| 41 | (*x.shape[:-1], self.out_features), |
| 42 | dtype=x.dtype, |
| 43 | device=x.device, |
| 44 | ) |
| 45 | |
| 46 | |
| 47 | class DummyNorm(nn.Module): |
| 48 | def __init__(self, *args, **kwargs): |
| 49 | super().__init__() |
| 50 | |
| 51 | def forward(self, hidden_states, residual=None): |
| 52 | return hidden_states, residual |
| 53 | |
| 54 | |
| 55 | class DummyDecoderLayer(nn.Module): |
| 56 | def __init__(self, *args, **kwargs): |
| 57 | super().__init__() |
| 58 | |
| 59 | def forward(self, positions, hidden_states, residual, llama_4_scaling=None): |
| 60 | return hidden_states, residual |
| 61 | |
| 62 | |
| 63 | def make_vllm_config( |
| 64 | *, model_type="mistral3", qk_nope_head_dim=128, qk_rope_head_dim=64 |
| 65 | ): |
| 66 | hf_config = SimpleNamespace( |
| 67 | model_type=model_type, |
| 68 | first_k_dense_replace=0, |
| 69 | vocab_size=32000, |
| 70 | hidden_size=16, |
| 71 | num_hidden_layers=1, |
| 72 | rms_norm_eps=1e-5, |
| 73 | qk_nope_head_dim=qk_nope_head_dim, |
| 74 | qk_rope_head_dim=qk_rope_head_dim, |
| 75 | ) |
| 76 | |
| 77 | return SimpleNamespace( |
| 78 | model_config=SimpleNamespace(hf_config=hf_config), |
| 79 | quant_config=None, |
| 80 | parallel_config=SimpleNamespace( |
| 81 | eplb_config=SimpleNamespace(num_redundant_experts=0), |
| 82 | ), |
| 83 | scheduler_config=SimpleNamespace(max_num_batched_tokens=8), |
| 84 | cache_config=None, |
| 85 | compilation_config=SimpleNamespace(mode=CompilationMode.NONE), |
| 86 | ) |
| 87 | |
| 88 | |
| 89 | @pytest.fixture(autouse=True) |
| 90 | def patch_heavy_modules(monkeypatch): |
| 91 | monkeypatch.setattr(eagle_mod, "get_pp_group", lambda: DummyPPGroup()) |
| 92 | monkeypatch.setattr(deepseek_mod, "get_pp_group", lambda: DummyPPGroup()) |
| 93 | |
| 94 | monkeypatch.setattr(eagle_mod, "VocabParallelEmbedding", DummyEmbedding) |
| 95 | monkeypatch.setattr(eagle_mod, "RowParallelLinear", DummyLinear) |
| 96 | monkeypatch.setattr(eagle_mod, "RMSNorm", DummyNorm) |
| 97 | monkeypatch.setattr(eagle_mod, "DeepseekV2DecoderLayer", DummyDecoderLayer) |
| 98 | |
| 99 | |
| 100 | @pytest.mark.cpu_test |
| 101 | @pytest.mark.parametrize( |
| 102 | ("model_type", "qk_nope_head_dim", "qk_rope_head_dim", "expected_use_mha"), |
| 103 | [ |
| 104 | # MLA-style config: should not use MHA. |
| 105 | ("mistral3", 128, 64, False), |
| 106 | # No MLA dims: should use MHA, matching DeepseekV2Model.__init__ logic. |
| 107 | ("mistral3", 0, 0, True), |
| 108 | # DeepSeek model type always uses MHA by the parent logic. |
| 109 | ("deepseek", 128, 64, True), |
| 110 | ], |
| 111 | ) |
| 112 | def test_eagle_mistral_large3_initializes_deepseek_runtime_attrs( |
| 113 | model_type, |
| 114 | qk_nope_head_dim, |
| 115 | qk_rope_head_dim, |
| 116 | expected_use_mha, |
| 117 | ): |
| 118 | vllm_config = make_vllm_config( |
| 119 | model_type=model_type, |
| 120 | qk_nope_head_dim=qk_nope_head_dim, |
| 121 | qk_rope_head_dim=qk_rope_head_dim, |
| 122 | ) |
| 123 | |
| 124 | model = eagle_mod.EagleMistralLarge3Model(vllm_config=vllm_config) |
| 125 | |
| 126 | assert model.aux_hidden_state_layers == () |
| 127 | assert model.use_mha is expected_use_mha |
| 128 | |
| 129 | # Add this if your fix also copies num_redundant_experts from |
| 130 | # DeepseekV2Model.__init__. |
| 131 | assert model.num_redundant_experts == 0 |
| 132 | |
| 133 | |
| 134 | @pytest.mark.cpu_test |
| 135 | def test_eagle_mistral_large3_forward_reuses_deepseek_parent_forward(): |
| 136 | vllm_config = make_vllm_config() |
| 137 | model = eagle_mod.EagleMistralLarge3Model(vllm_config=vllm_config) |
| 138 | |
| 139 | input_ids = torch.tensor([[1, 2, 3]]) |
| 140 | positions = torch.tensor([[0, 1, 2]]) |
| 141 | hidden_states = torch.zeros((1, 3, 16)) |
| 142 | |
| 143 | output = model(input_ids, positions, hidden_states) |
| 144 | |
| 145 | assert isinstance(output, torch.Tensor) |
| 146 | assert output.shape == hidden_states.shape |