| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from types import SimpleNamespace |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.distributed import parallel_state |
| 9 | from vllm.model_executor.models.interfaces import EagleModelMixin |
| 10 | from vllm.model_executor.models.mimo import MiMoModel |
| 11 | from vllm.models.minimax_m3.nvidia import model as minimax_m3 |
| 12 | from vllm.sequence import IntermediateTensors |
| 13 | from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( |
| 14 | reserve_aux_intermediate_tensor_slots, |
| 15 | verify_supports_aux_hidden_states_over_pp, |
| 16 | ) |
| 17 | |
| 18 | |
| 19 | def test_aux_layers_are_sorted_and_deduplicated(): |
| 20 | model = EagleModelMixin() |
| 21 | model._set_aux_hidden_state_layers((48, 3, 90, 24, 48)) |
| 22 | assert model.aux_hidden_state_layers == (3, 24, 48, 90) |
| 23 | |
| 24 | |
| 25 | def test_mimo_does_not_inherit_aux_hidden_state_pp_support(): |
| 26 | inner = MiMoModel.__new__(MiMoModel) |
| 27 | target = SimpleNamespace(model=inner) |
| 28 | |
| 29 | assert not inner.supports_aux_hidden_states_over_pp |
| 30 | with pytest.raises(ValueError, match="does not support eagle3"): |
| 31 | verify_supports_aux_hidden_states_over_pp(target, "eagle3") |
| 32 | |
| 33 | |
| 34 | class _PPAuxLayer(torch.nn.Module): |
| 35 | """Aliases and mutates residual in place like MiniMaxM3DecoderLayer.""" |
| 36 | |
| 37 | ffn_all_reduce_deferred = False |
| 38 | |
| 39 | def __init__(self, index): |
| 40 | super().__init__() |
| 41 | self.index = index |
| 42 | |
| 43 | def forward(self, positions, hidden_states, residual): |
| 44 | total = hidden_states if residual is None else hidden_states + residual |
| 45 | residual = hidden_states if residual is None else residual |
| 46 | residual.copy_(total + self.index + 1) |
| 47 | return torch.full_like(hidden_states, 10 * (self.index + 1)), residual |
| 48 | |
| 49 | |
| 50 | class _PPAuxNorm(torch.nn.Module): |
| 51 | """Clobbers its inputs so aux taps that alias them would fail.""" |
| 52 | |
| 53 | def forward(self, hidden_states, residual): |
| 54 | output = hidden_states + residual |
| 55 | hidden_states.fill_(-101) |
| 56 | residual.fill_(-103) |
| 57 | return output, residual |
| 58 | |
| 59 | |
| 60 | @pytest.fixture |
| 61 | def minimax_pp_stage(monkeypatch): |
| 62 | """Use the real MiniMax constructor/forward and aux helpers on CPU tensors.""" |
| 63 | group = SimpleNamespace(world_size=2, is_first_rank=True, is_last_rank=False) |
| 64 | bounds = [0, 3] |
| 65 | monkeypatch.setattr(minimax_m3, "get_pp_group", lambda: group) |
| 66 | monkeypatch.setattr(parallel_state, "get_pp_group", lambda: group) |
| 67 | monkeypatch.setattr(parallel_state, "model_parallel_is_initialized", lambda: True) |
| 68 | monkeypatch.setattr( |
| 69 | minimax_m3, |
| 70 | "VocabParallelEmbedding", |
| 71 | lambda *args, **kwargs: torch.nn.Identity(), |
| 72 | ) |
| 73 | monkeypatch.setattr( |
| 74 | minimax_m3, "MiniMAXGemmaRMSNorm", lambda *args, **kwargs: _PPAuxNorm() |
| 75 | ) |
| 76 | monkeypatch.setattr( |
| 77 | minimax_m3, |
| 78 | "MiniMaxM3DecoderLayer", |
| 79 | lambda **kwargs: _PPAuxLayer(int(kwargs["prefix"].rsplit(".", 1)[1])), |
| 80 | ) |
| 81 | |
| 82 | def make_layers(num_layers, build, prefix): |
| 83 | return ( |
| 84 | *bounds, |
| 85 | torch.nn.ModuleList( |
| 86 | build(f"{prefix}.{index}") |
| 87 | if bounds[0] <= index < bounds[1] |
| 88 | else torch.nn.Identity() |
| 89 | for index in range(num_layers) |
| 90 | ), |
| 91 | ) |
| 92 | |
| 93 | monkeypatch.setattr(minimax_m3, "make_layers", make_layers) |
| 94 | |
| 95 | def make_stage(start, end, taps, world_size=2): |
| 96 | bounds[:] = [start, end] |
| 97 | group.world_size = world_size |
| 98 | group.is_first_rank = start == 0 |
| 99 | group.is_last_rank = end == 6 |
| 100 | config = SimpleNamespace( |
| 101 | vocab_size=16, hidden_size=6144, num_hidden_layers=6, rms_norm_eps=1e-6 |
| 102 | ) |
| 103 | model = minimax_m3.MiniMaxM3Model( |
| 104 | vllm_config=SimpleNamespace( |
| 105 | model_config=SimpleNamespace(hf_text_config=config), |
| 106 | quant_config=None, |
| 107 | speculative_config=None, |
| 108 | use_v2_model_runner=False, |
| 109 | ) |
| 110 | ) |
| 111 | model._set_aux_hidden_state_layers(taps) |
| 112 | target = SimpleNamespace( |
| 113 | model=model, |
| 114 | make_empty_intermediate_tensors=model.make_empty_intermediate_tensors, |
| 115 | ) |
| 116 | reserve_aux_intermediate_tensor_slots(target) |
| 117 | return model, target |
| 118 | |
| 119 | return make_stage |
| 120 | |
| 121 | |
| 122 | # Input 7; _PPAuxLayer i adds 11 * (i + 1) to hidden + residual. |
| 123 | def _expected_aux_value(tap): |
| 124 | return 7 + 11 * tap * (tap + 1) // 2 |
| 125 | |
| 126 | |
| 127 | def _pp_aux_state(value): |
| 128 | return torch.full((2, 6144), float(value), device="cpu") |
| 129 | |
| 130 | |
| 131 | @pytest.mark.parametrize("taps", [(0, 3, 6), (0, 1, 3), (4, 5, 6)]) |
| 132 | def test_minimax_pp2_preserves_ordered_global_aux_taps(minimax_pp_stage, taps): |
| 133 | first, _ = minimax_pp_stage(0, 3, taps) |
| 134 | sent = first(None, None, None, _pp_aux_state(7)) |
| 135 | last, target = minimax_pp_stage(3, 6, taps) |
| 136 | reserved = target.make_empty_intermediate_tensors(2, torch.float32, "cpu") |
| 137 | assert set(sent.tensors) == set(reserved.tensors) |
| 138 | result = last(None, None, sent) |
| 139 | assert isinstance(result, tuple) |
| 140 | output, aux = result |
| 141 | torch.testing.assert_close(output, _pp_aux_state(_expected_aux_value(6))) |
| 142 | for hidden, tap in zip(aux, taps, strict=True): |
| 143 | torch.testing.assert_close(hidden, _pp_aux_state(_expected_aux_value(tap))) |
| 144 | |
| 145 | |
| 146 | def test_minimax_missing_upstream_aux_slot_is_not_silently_dropped(minimax_pp_stage): |
| 147 | model, _ = minimax_pp_stage(3, 6, (0, 3, 6)) |
| 148 | incoming = IntermediateTensors( |
| 149 | { |
| 150 | "hidden_states": _pp_aux_state(30), |
| 151 | "residual": _pp_aux_state(43), |
| 152 | "aux_hidden_states_0": _pp_aux_state(7), |
| 153 | } |
| 154 | ) |
| 155 | with pytest.raises(RuntimeError, match="Missing aux_hidden_states_1"): |
| 156 | model(None, None, incoming) |
| 157 | |
| 158 | |
| 159 | @pytest.mark.parametrize("taps", [(), (0, 3, 6)]) |
| 160 | def test_minimax_pp1_output_and_aux_taps(minimax_pp_stage, taps): |
| 161 | model, _ = minimax_pp_stage(0, 6, taps, world_size=1) |
| 162 | output = model(None, None, None, _pp_aux_state(7)) |
| 163 | if taps: |
| 164 | output, aux = output |
| 165 | for hidden, tap in zip(aux, taps, strict=True): |
| 166 | torch.testing.assert_close(hidden, _pp_aux_state(_expected_aux_value(tap))) |
| 167 | torch.testing.assert_close(output, _pp_aux_state(_expected_aux_value(6))) |
| 168 | |
| 169 | |
| 170 | def test_minimax_pp2_without_aux_keeps_hidden_and_residual_transport(minimax_pp_stage): |
| 171 | first, _ = minimax_pp_stage(0, 3, ()) |
| 172 | sent = first(None, None, None, _pp_aux_state(7)) |
| 173 | assert set(sent.tensors) == {"hidden_states", "residual"} |
| 174 | torch.testing.assert_close(sent["hidden_states"], _pp_aux_state(30)) |
| 175 | torch.testing.assert_close(sent["residual"], _pp_aux_state(43)) |
| 176 | last, target = minimax_pp_stage(3, 6, ()) |
| 177 | assert set( |
| 178 | target.make_empty_intermediate_tensors(12, torch.bfloat16, "cpu").tensors |
| 179 | ) == {"hidden_states", "residual"} |
| 180 | output = last(None, None, sent) |
| 181 | torch.testing.assert_close(output, _pp_aux_state(_expected_aux_value(6))) |