| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | from unittest.mock import Mock |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.model_executor.models.interfaces import supports_eagle3 |
| 10 | from vllm.models.kimi_k3.nvidia import model as kimi_model |
| 11 | from vllm.models.kimi_k3.nvidia.model import ( |
| 12 | KimiK3ForConditionalGeneration, |
| 13 | KimiLinearForCausalLM, |
| 14 | KimiLinearModel, |
| 15 | ) |
| 16 | |
| 17 | |
| 18 | def _make_kimi_linear_model() -> KimiLinearModel: |
| 19 | model = object.__new__(KimiLinearModel) |
| 20 | object.__setattr__(model, "aux_hidden_state_layers", (2,)) |
| 21 | object.__setattr__(model, "use_sequence_parallel", False) |
| 22 | object.__setattr__(model, "use_attn_res", False) |
| 23 | return model |
| 24 | |
| 25 | |
| 26 | def test_kimi_k3_advertises_eagle3_support(): |
| 27 | assert supports_eagle3(KimiK3ForConditionalGeneration) |
| 28 | |
| 29 | |
| 30 | def test_kimi_linear_advertises_eagle3_support(): |
| 31 | # The text-only architecture serves the same inner KimiLinearModel, which |
| 32 | # already carries the EagleModelMixin tap machinery - only the interface |
| 33 | # declaration was missing, so EAGLE3-family speculative decoding (dspark) |
| 34 | # was rejected at startup with "Model does not support EAGLE3 interface". |
| 35 | assert supports_eagle3(KimiLinearForCausalLM) |
| 36 | |
| 37 | |
| 38 | def test_kimi_k3_uses_shared_eagle3_layer_configuration(): |
| 39 | target = object.__new__(KimiK3ForConditionalGeneration) |
| 40 | torch.nn.Module.__init__(target) |
| 41 | model = _make_kimi_linear_model() |
| 42 | object.__setattr__(model, "layers", [None] * 93) |
| 43 | language_model = SimpleNamespace( |
| 44 | embed_input_ids=lambda _: None, |
| 45 | forward=lambda input_ids, positions: None, |
| 46 | model=model, |
| 47 | ) |
| 48 | object.__setattr__(target, "language_model", language_model) |
| 49 | object.__setattr__(target, "_language_model_names", ["language_model"]) |
| 50 | |
| 51 | target.set_aux_hidden_state_layers((2, 46, 90)) |
| 52 | |
| 53 | assert model.aux_hidden_state_layers == (2, 46, 90) |
| 54 | assert target.get_eagle3_default_aux_hidden_state_layers() == ( |
| 55 | 2, |
| 56 | 46, |
| 57 | 90, |
| 58 | ) |
| 59 | |
| 60 | |
| 61 | def test_kimi_linear_forward_extracts_standard_aux_hidden_states(monkeypatch): |
| 62 | model = _make_kimi_linear_model() |
| 63 | initial_hidden_states = torch.tensor([[1.0, 2.0]]) |
| 64 | layer_hidden_states = torch.tensor([[3.0, 4.0]]) |
| 65 | layer_residual = torch.tensor([[5.0, 6.0]]) |
| 66 | |
| 67 | object.__setattr__(model, "start_layer", 0) |
| 68 | object.__setattr__(model, "end_layer", 1) |
| 69 | object.__setattr__( |
| 70 | model, |
| 71 | "layers", |
| 72 | [Mock(return_value=(layer_hidden_states, None, layer_residual))], |
| 73 | ) |
| 74 | object.__setattr__(model, "aux_hidden_state_layers", (0, 1)) |
| 75 | object.__setattr__(model, "use_attn_res", False) |
| 76 | monkeypatch.setattr( |
| 77 | kimi_model, |
| 78 | "get_pp_group", |
| 79 | lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), |
| 80 | ) |
| 81 | |
| 82 | output, aux_hidden_states = model.forward( |
| 83 | input_ids=None, |
| 84 | positions=torch.tensor([0]), |
| 85 | intermediate_tensors=None, |
| 86 | inputs_embeds=initial_hidden_states, |
| 87 | ) |
| 88 | |
| 89 | expected_layer_output = layer_hidden_states + layer_residual |
| 90 | torch.testing.assert_close(output, expected_layer_output) |
| 91 | torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) |
| 92 | torch.testing.assert_close(aux_hidden_states[1], expected_layer_output) |
| 93 | |
| 94 | |
| 95 | def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch): |
| 96 | model = _make_kimi_linear_model() |
| 97 | initial_hidden_states = torch.tensor([[1.0, 2.0]]) |
| 98 | layer_hidden_states = torch.tensor([[3.0, 4.0]]) |
| 99 | prefix_sum = torch.tensor([[5.0, 6.0]]) |
| 100 | block_residual = torch.tensor([[[7.0, 8.0]]]) |
| 101 | final_hidden_states = torch.tensor([[9.0, 10.0]]) |
| 102 | |
| 103 | object.__setattr__(model, "start_layer", 0) |
| 104 | object.__setattr__(model, "end_layer", 1) |
| 105 | object.__setattr__( |
| 106 | model, |
| 107 | "layers", |
| 108 | [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))], |
| 109 | ) |
| 110 | object.__setattr__(model, "aux_hidden_state_layers", (0, 1)) |
| 111 | object.__setattr__(model, "use_attn_res", True) |
| 112 | object.__setattr__(model, "num_attn_res_blocks", 1) |
| 113 | object.__setattr__( |
| 114 | model, |
| 115 | "output_attn_res_norm", |
| 116 | SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), |
| 117 | ) |
| 118 | object.__setattr__( |
| 119 | model, |
| 120 | "output_attn_res_proj", |
| 121 | SimpleNamespace(weight=torch.ones(1, 2)), |
| 122 | ) |
| 123 | monkeypatch.setattr( |
| 124 | kimi_model, |
| 125 | "get_pp_group", |
| 126 | lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), |
| 127 | ) |
| 128 | final_attn_res = Mock(return_value=final_hidden_states) |
| 129 | monkeypatch.setattr(kimi_model, "attn_res", final_attn_res) |
| 130 | |
| 131 | output, aux_hidden_states = model.forward( |
| 132 | input_ids=None, |
| 133 | positions=torch.tensor([0]), |
| 134 | intermediate_tensors=None, |
| 135 | inputs_embeds=initial_hidden_states, |
| 136 | ) |
| 137 | |
| 138 | torch.testing.assert_close(output, final_hidden_states) |
| 139 | torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) |
| 140 | torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states) |
| 141 | assert final_attn_res.call_args.args[2] is block_residual |
| 142 | |
| 143 | |
| 144 | def test_attn_res_stream_capture_receives_the_layer_outputs_in_order(monkeypatch): |
| 145 | """Pin the argument mapping at the call site. |
| 146 | |
| 147 | The capture helper's own tests invoke it directly by keyword, so they |
| 148 | cannot catch a swap where `forward` hands it the residual as the pending |
| 149 | MLP output. Both are tensors of the same shape, so a swap is silent: it |
| 150 | feeds the drafter a wrong but well-formed tensor. |
| 151 | """ |
| 152 | model = _make_kimi_linear_model() |
| 153 | initial_hidden_states = torch.tensor([[1.0, 2.0]]) |
| 154 | layer_hidden_states = torch.tensor([[3.0, 4.0]]) |
| 155 | prefix_sum = torch.tensor([[5.0, 6.0]]) |
| 156 | block_residual = torch.tensor([[[7.0, 8.0]]]) |
| 157 | captured = torch.tensor([[11.0, 12.0]]) |
| 158 | |
| 159 | object.__setattr__(model, "start_layer", 0) |
| 160 | object.__setattr__(model, "end_layer", 1) |
| 161 | object.__setattr__( |
| 162 | model, |
| 163 | "layers", |
| 164 | [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))], |
| 165 | ) |
| 166 | object.__setattr__(model, "aux_hidden_state_layers", (1,)) |
| 167 | object.__setattr__(model, "use_attn_res", True) |
| 168 | object.__setattr__(model, "num_attn_res_blocks", 1) |
| 169 | object.__setattr__( |
| 170 | model, |
| 171 | "output_attn_res_norm", |
| 172 | SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), |
| 173 | ) |
| 174 | object.__setattr__( |
| 175 | model, |
| 176 | "output_attn_res_proj", |
| 177 | SimpleNamespace(weight=torch.ones(1, 2)), |
| 178 | ) |
| 179 | monkeypatch.setattr( |
| 180 | kimi_model, |
| 181 | "get_pp_group", |
| 182 | lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), |
| 183 | ) |
| 184 | monkeypatch.setattr(kimi_model, "attn_res", Mock(return_value=torch.zeros(1, 2))) |
| 185 | monkeypatch.setenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "1") |
| 186 | |
| 187 | capture = Mock(return_value=captured) |
| 188 | monkeypatch.setattr(KimiLinearModel, "_capture_aux_hidden_stream", capture) |
| 189 | |
| 190 | _, aux_hidden_states = model.forward( |
| 191 | input_ids=None, |
| 192 | positions=torch.tensor([0]), |
| 193 | intermediate_tensors=None, |
| 194 | inputs_embeds=initial_hidden_states, |
| 195 | ) |
| 196 | |
| 197 | layer_idx, got_prefix, got_pending, got_residual = capture.call_args.args |
| 198 | assert layer_idx == 0 |
| 199 | assert got_prefix is prefix_sum |
| 200 | assert got_pending is layer_hidden_states |
| 201 | assert got_residual is block_residual |
| 202 | torch.testing.assert_close(aux_hidden_states[0], captured) |