| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Which value the DFlash drafter is fed under AttnRes. |
| 4 | |
| 5 | `_capture_aux_hidden_stream` picks the weights it mixes against from one of |
| 6 | three places depending on where the tapped layer sits, and returns the plain |
| 7 | running prefix when the feature is off. The mixture itself is the kernel's |
| 8 | job and is covered by ``test_attn_res.py``; what is asserted here is the |
| 9 | selection, which is the part that can silently feed the drafter the wrong |
| 10 | tensor. |
| 11 | """ |
| 12 | |
| 13 | from types import SimpleNamespace |
| 14 | |
| 15 | import pytest |
| 16 | import torch |
| 17 | |
| 18 | from vllm.models.kimi_k3.nvidia import model as k3_model |
| 19 | |
| 20 | END_LAYER = 4 |
| 21 | |
| 22 | |
| 23 | def _weights(tag: float) -> SimpleNamespace: |
| 24 | """A norm/projection pair that is identifiable by value.""" |
| 25 | return SimpleNamespace( |
| 26 | weight=torch.full((2,), tag), |
| 27 | variance_epsilon=tag, |
| 28 | ) |
| 29 | |
| 30 | |
| 31 | def _stub_model(*, enabled: bool, use_attn_res: bool = True) -> SimpleNamespace: |
| 32 | """A stand-in carrying only what the tap reads. |
| 33 | |
| 34 | Constructing the real model needs a distributed init and weights, and none |
| 35 | of it participates in the selection under test. |
| 36 | """ |
| 37 | consumers = [] |
| 38 | for i in range(END_LAYER): |
| 39 | consumers.append( |
| 40 | SimpleNamespace( |
| 41 | self_attention_res_norm=_weights(float(i)), |
| 42 | self_attention_res_proj=SimpleNamespace( |
| 43 | weight=torch.full((1, 2), float(i)) |
| 44 | ), |
| 45 | prev_valid_blocks=i, |
| 46 | ) |
| 47 | ) |
| 48 | return SimpleNamespace( |
| 49 | _aux_attn_res_stream=enabled, |
| 50 | use_attn_res=use_attn_res, |
| 51 | end_layer=END_LAYER, |
| 52 | layers=consumers, |
| 53 | output_attn_res_norm=_weights(99.0), |
| 54 | output_attn_res_proj=SimpleNamespace(weight=torch.full((1, 2), 99.0)), |
| 55 | num_attn_res_blocks=99, |
| 56 | ) |
| 57 | |
| 58 | |
| 59 | @pytest.fixture |
| 60 | def recorder(monkeypatch): |
| 61 | """Replace the kernel so the call it would have made is inspectable.""" |
| 62 | calls = [] |
| 63 | |
| 64 | def _fake_attn_res( |
| 65 | prefix, |
| 66 | delta, |
| 67 | block_residual, |
| 68 | norm_weight, |
| 69 | proj_weight, |
| 70 | output_norm_weight, |
| 71 | **kwargs, |
| 72 | ): |
| 73 | calls.append( |
| 74 | SimpleNamespace( |
| 75 | prefix=prefix, |
| 76 | delta=delta, |
| 77 | block_residual=block_residual, |
| 78 | norm_weight=norm_weight, |
| 79 | proj_weight=proj_weight, |
| 80 | kwargs=kwargs, |
| 81 | ) |
| 82 | ) |
| 83 | return torch.full_like(prefix, -1.0) |
| 84 | |
| 85 | monkeypatch.setattr(k3_model, "attn_res", _fake_attn_res) |
| 86 | return calls |
| 87 | |
| 88 | |
| 89 | def _set_last_rank(monkeypatch, is_last: bool): |
| 90 | monkeypatch.setattr( |
| 91 | k3_model, |
| 92 | "get_pp_group", |
| 93 | lambda: SimpleNamespace(is_last_rank=is_last), |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | def _call(stub, layer_idx, prefix_sum, pending_mlp_out, block_residual): |
| 98 | return k3_model.KimiLinearModel._capture_aux_hidden_stream( |
| 99 | stub, layer_idx, prefix_sum, pending_mlp_out, block_residual |
| 100 | ) |
| 101 | |
| 102 | |
| 103 | @pytest.mark.parametrize( |
| 104 | "enabled,use_attn_res", [(False, True), (True, False), (False, False)] |
| 105 | ) |
| 106 | def test_disabled_reproduces_the_plain_residual_sum( |
| 107 | recorder, monkeypatch, enabled, use_attn_res |
| 108 | ): |
| 109 | """Off, the tap must be exactly the sum it replaced. |
| 110 | |
| 111 | Both conditions matter. `use_attn_res` is what constructs the norm and |
| 112 | projection weights, so without it the lookups below would raise rather |
| 113 | than fall back. |
| 114 | """ |
| 115 | _set_last_rank(monkeypatch, True) |
| 116 | prefix_sum = torch.tensor([1.0, 2.0]) |
| 117 | pending = torch.tensor([0.5, 0.25]) |
| 118 | |
| 119 | got = _call( |
| 120 | _stub_model(enabled=enabled, use_attn_res=use_attn_res), |
| 121 | 0, |
| 122 | prefix_sum, |
| 123 | pending, |
| 124 | torch.zeros(2), |
| 125 | ) |
| 126 | |
| 127 | torch.testing.assert_close(got, prefix_sum + pending) |
| 128 | assert not recorder, "the kernel must not run when the tap is off" |
| 129 | |
| 130 | |
| 131 | def test_taps_the_consumer_layer_when_one_follows(recorder, monkeypatch): |
| 132 | """The value the next layer reads is the mixture against *its* weights, |
| 133 | so the tap has to reach forward rather than use the current layer's.""" |
| 134 | _set_last_rank(monkeypatch, True) |
| 135 | |
| 136 | _call(_stub_model(enabled=True), 1, torch.zeros(2), None, torch.zeros(2)) |
| 137 | |
| 138 | assert len(recorder) == 1 |
| 139 | call = recorder[0] |
| 140 | # Layer 2's weights, not layer 1's. |
| 141 | torch.testing.assert_close(call.norm_weight, torch.full((2,), 2.0)) |
| 142 | assert call.kwargs["num_blocks"] == 2 |
| 143 | |
| 144 | |
| 145 | def test_last_layer_on_the_final_rank_uses_the_output_aggregation( |
| 146 | recorder, monkeypatch |
| 147 | ): |
| 148 | """Nothing downstream but the model's own output-side mixture.""" |
| 149 | _set_last_rank(monkeypatch, True) |
| 150 | |
| 151 | _call( |
| 152 | _stub_model(enabled=True), END_LAYER - 1, torch.zeros(2), None, torch.zeros(2) |
| 153 | ) |
| 154 | |
| 155 | assert len(recorder) == 1 |
| 156 | torch.testing.assert_close(recorder[0].norm_weight, torch.full((2,), 99.0)) |
| 157 | assert recorder[0].kwargs["num_blocks"] == 99 |
| 158 | |
| 159 | |
| 160 | def test_aux_layer_at_non_final_pp_boundary_is_rejected(monkeypatch): |
| 161 | model = k3_model.KimiLinearModel.__new__(k3_model.KimiLinearModel) |
| 162 | torch.nn.Module.__init__(model) |
| 163 | model.use_attn_res = True |
| 164 | model.end_layer = 72 |
| 165 | monkeypatch.setattr( |
| 166 | "vllm.distributed.parallel_state.model_parallel_is_initialized", lambda: False |
| 167 | ) |
| 168 | _set_last_rank(monkeypatch, False) |
| 169 | monkeypatch.setattr( |
| 170 | k3_model.KimiLinearModel, |
| 171 | "_aux_attn_res_stream", |
| 172 | property(lambda self: True), |
| 173 | ) |
| 174 | |
| 175 | with pytest.raises(ValueError, match="Auxiliary layer 72"): |
| 176 | model._set_aux_hidden_state_layers((3, 24, 48, 72, 90)) |
| 177 | |
| 178 | |
| 179 | def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta( |
| 180 | recorder, monkeypatch |
| 181 | ): |
| 182 | """The kernel writes an applied delta back into the prefix in place, which |
| 183 | would double-add it into the live residual stream, so the pending output |
| 184 | has to arrive already summed into the prefix with `delta` left None.""" |
| 185 | _set_last_rank(monkeypatch, True) |
| 186 | prefix_sum = torch.tensor([1.0, 2.0]) |
| 187 | pending = torch.tensor([0.5, 0.25]) |
| 188 | |
| 189 | _call(_stub_model(enabled=True), 0, prefix_sum, pending, torch.zeros(2)) |
| 190 | |
| 191 | assert len(recorder) == 1 |
| 192 | assert recorder[0].delta is None |
| 193 | torch.testing.assert_close(recorder[0].prefix, prefix_sum + pending) |
| 194 | # And the caller's tensor is not mutated on the way. |
| 195 | torch.testing.assert_close(prefix_sum, torch.tensor([1.0, 2.0])) |