[Model][Spec Decode] Tap the pre-norm AttnRes mixture as the Kimi K3 DFlash aux state (#50487)
Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Co-authored-by: Janelle Cai <janelle.cai@modal.com>
Rahul Chalamala committed
Aug 14, 2026 at 12:03 UTC
03a8d0b1ede68efc20b8bbfc14ee4681a62f8d75
4 files changed
+350
-1
tests/models/kimi_k3/test_aux_attn_res_stream.py
new
+197
@@ -0,0 +1,197 @@
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_last_layer_of_a_non_final_stage_falls_back(recorder, monkeypatch):
161
+ """The consumer lives on the next rank and the output aggregation only
162
+ exists on the last one, so there is nothing here to mix against.
163
+
164
+ This is the case that would otherwise reach for weights this rank never
165
+ constructs. The forward guard is `layer_idx + 1 < end_layer`, where
166
+ `end_layer` is the rank's own exclusive bound from `get_pp_indices`, so a
167
+ `PPMissingLayer` is unreachable by construction -- the fallback below is
168
+ what makes that true rather than merely likely.
169
+ """
170
+ _set_last_rank(monkeypatch, False)
171
+ prefix_sum = torch.tensor([3.0, 4.0])
172
+
173
+ got = _call(
174
+ _stub_model(enabled=True), END_LAYER - 1, prefix_sum, None, torch.zeros(2)
175
+ )
176
+
177
+ torch.testing.assert_close(got, prefix_sum)
178
+ assert not recorder, "no weights exist on this rank to mix against"
179
+
180
+
181
+def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta(
182
+ recorder, monkeypatch
183
+):
184
+ """The kernel writes an applied delta back into the prefix in place, which
185
+ would double-add it into the live residual stream, so the pending output
186
+ has to arrive already summed into the prefix with `delta` left None."""
187
+ _set_last_rank(monkeypatch, True)
188
+ prefix_sum = torch.tensor([1.0, 2.0])
189
+ pending = torch.tensor([0.5, 0.25])
190
+
191
+ _call(_stub_model(enabled=True), 0, prefix_sum, pending, torch.zeros(2))
192
+
193
+ assert len(recorder) == 1
194
+ assert recorder[0].delta is None
195
+ torch.testing.assert_close(recorder[0].prefix, prefix_sum + pending)
196
+ # And the caller's tensor is not mutated on the way.
197
+ torch.testing.assert_close(prefix_sum, torch.tensor([1.0, 2.0]))
tests/models/kimi_k3/test_eagle3.py
+62
@@ -19,6 +19,7 @@ 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
@@ -137,3 +138,64 @@ def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch):
138
torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states)
139
torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states)
140
assert final_attn_res.call_args.args[2] is block_residual
141
+
142
+
143
+def test_attn_res_stream_capture_receives_the_layer_outputs_in_order(monkeypatch):
144
+ """Pin the argument mapping at the call site.
145
+
146
+ The capture helper's own tests invoke it directly by keyword, so they
147
+ cannot catch a swap where `forward` hands it the residual as the pending
148
+ MLP output. Both are tensors of the same shape, so a swap is silent: it
149
+ feeds the drafter a wrong but well-formed tensor.
150
+ """
151
+ model = _make_kimi_linear_model()
152
+ initial_hidden_states = torch.tensor([[1.0, 2.0]])
153
+ layer_hidden_states = torch.tensor([[3.0, 4.0]])
154
+ prefix_sum = torch.tensor([[5.0, 6.0]])
155
+ block_residual = torch.tensor([[[7.0, 8.0]]])
156
+ captured = torch.tensor([[11.0, 12.0]])
157
+
158
+ object.__setattr__(model, "start_layer", 0)
159
+ object.__setattr__(model, "end_layer", 1)
160
+ object.__setattr__(
161
+ model,
162
+ "layers",
163
+ [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))],
164
+ )
165
+ object.__setattr__(model, "aux_hidden_state_layers", (1,))
166
+ object.__setattr__(model, "use_attn_res", True)
167
+ object.__setattr__(model, "num_attn_res_blocks", 1)
168
+ object.__setattr__(
169
+ model,
170
+ "output_attn_res_norm",
171
+ SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5),
172
+ )
173
+ object.__setattr__(
174
+ model,
175
+ "output_attn_res_proj",
176
+ SimpleNamespace(weight=torch.ones(1, 2)),
177
+ )
178
+ monkeypatch.setattr(
179
+ kimi_model,
180
+ "get_pp_group",
181
+ lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
182
+ )
183
+ monkeypatch.setattr(kimi_model, "attn_res", Mock(return_value=torch.zeros(1, 2)))
184
+ monkeypatch.setenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "1")
185
+
186
+ capture = Mock(return_value=captured)
187
+ monkeypatch.setattr(KimiLinearModel, "_capture_aux_hidden_stream", capture)
188
+
189
+ _, aux_hidden_states = model.forward(
190
+ input_ids=None,
191
+ positions=torch.tensor([0]),
192
+ intermediate_tensors=None,
193
+ inputs_embeds=initial_hidden_states,
194
+ )
195
+
196
+ layer_idx, got_prefix, got_pending, got_residual = capture.call_args.args
197
+ assert layer_idx == 0
198
+ assert got_prefix is prefix_sum
199
+ assert got_pending is layer_hidden_states
200
+ assert got_residual is block_residual
201
+ torch.testing.assert_close(aux_hidden_states[0], captured)
vllm/envs.py
+8
@@ -207,6 +207,7 @@ if TYPE_CHECKING:
207
VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True
208
VLLM_MOE_SKIP_PADDING: bool = True
209
VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False
210
+ VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False
211
VLLM_KIMI_K3_GEMM_RS: bool = False
212
VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True
213
VLLM_USE_FLASHINFER_MOE_INT4: bool = False
@@ -1605,6 +1606,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
1606
"VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool(
1607
int(os.getenv("VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", "0"))
1608
),
1609
+ # Kimi K3 only, and unrelated to the MoE flags above. Tap the pre-norm
1610
+ # AttnRes mixture, rather than the post-mixture sum, as the auxiliary
1611
+ # hidden state handed to a DFlash drafter. This changes the numerics the
1612
+ # speculator sees, so it is off by default while the effect is measured.
1613
+ "VLLM_KIMI_K3_AUX_ATTN_RES_STREAM": lambda: bool(
1614
+ int(os.getenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "0"))
1615
+ ),
1616
# Use the SM100 BF16 GEMM-RS kernel for eligible Kimi-K3 sequence-parallel
1617
# row-parallel projections. All TP ranks must belong to one NVLink domain.
1618
"VLLM_KIMI_K3_GEMM_RS": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_RS", "0"))),
vllm/models/kimi_k3/nvidia/model.py
+83
-1
@@ -1198,6 +1198,85 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1198
}
1199
)
1200
1201
+ def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
1202
+ super()._set_aux_hidden_state_layers(layers)
1203
+ if self.use_attn_res:
1204
+ # Emitted once, at configuration time. Which layers are tapped and
1205
+ # which convention is in force are the two things you need to
1206
+ # confirm from a running process, and neither is recoverable from
1207
+ # the served output.
1208
+ logger.info_once(
1209
+ "Kimi-K3 aux hidden capture: layers=%s mode=%s "
1210
+ "(VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=%d)",
1211
+ layers,
1212
+ "attn_res_stream" if self._aux_attn_res_stream else "prefix_only",
1213
+ int(self._aux_attn_res_stream),
1214
+ )
1215
+
1216
+ @property
1217
+ def _aux_attn_res_stream(self) -> bool:
1218
+ return envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM
1219
+
1220
+ def _capture_aux_hidden_stream(
1221
+ self,
1222
+ layer_idx: int,
1223
+ prefix_sum: torch.Tensor,
1224
+ pending_mlp_out: torch.Tensor | None,
1225
+ block_residual: torch.Tensor,
1226
+ ) -> torch.Tensor:
1227
+ """Auxiliary feature tapped after ``layer_idx`` under AttnRes.
1228
+
1229
+ The wire between layers only carries the current block's running prefix;
1230
+ the committed blocks live in the bank. The value the next consumer
1231
+ actually reads is the pre-norm AttnRes mixture over
1232
+ ``bank[:num_blocks] + prefix``, which is what the DFlash drafters were
1233
+ trained against. ``attn_res`` with no delta, no block write and no
1234
+ output norm computes exactly that and leaves both the prefix and the
1235
+ bank untouched.
1236
+
1237
+ Folding the pending MLP output into the prefix rather than passing it as
1238
+ ``delta`` is deliberate: the kernel writes an applied delta back into
1239
+ the prefix in place, which would double-add it into the live residual
1240
+ stream.
1241
+ """
1242
+ prefix = prefix_sum if pending_mlp_out is None else prefix_sum + pending_mlp_out
1243
+ # `use_attn_res` is what constructs the norm and projection weights this
1244
+ # reads; without it there is no mixture to compute and the attribute
1245
+ # lookups below would raise.
1246
+ if not (self._aux_attn_res_stream and self.use_attn_res):
1247
+ return prefix
1248
+
1249
+ if layer_idx + 1 < self.end_layer:
1250
+ consumer = self.layers[layer_idx + 1]
1251
+ score_norm = consumer.self_attention_res_norm
1252
+ score_proj = consumer.self_attention_res_proj
1253
+ num_blocks = consumer.prev_valid_blocks
1254
+ elif get_pp_group().is_last_rank:
1255
+ # Nothing downstream but the model's own output-side aggregation.
1256
+ score_norm = self.output_attn_res_norm
1257
+ score_proj = self.output_attn_res_proj
1258
+ num_blocks = self.num_attn_res_blocks
1259
+ else:
1260
+ # Last layer of a non-final pipeline stage: the consumer lives on
1261
+ # the next rank and the output-side aggregation only exists on the
1262
+ # last one, so there is nothing here to mix against. Falling back
1263
+ # to the running prefix keeps the tap defined rather than reaching
1264
+ # for weights this rank does not construct.
1265
+ return prefix
1266
+
1267
+ return attn_res(
1268
+ prefix,
1269
+ None,
1270
+ block_residual,
1271
+ score_norm.weight,
1272
+ score_proj.weight.squeeze(0),
1273
+ None,
1274
+ num_blocks=num_blocks,
1275
+ block_write_idx=-1,
1276
+ eps=score_norm.variance_epsilon,
1277
+ output_norm_eps=0.0,
1278
+ )
1279
+
1280
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1281
return self.embed_tokens(input_ids)
1282
@@ -1265,7 +1344,10 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1344
if (layer_idx + 1) in self.aux_hidden_state_layers:
1345
if self.use_attn_res:
1346
assert prefix_sum is not None
1268
- aux_hidden_state = prefix_sum + hidden_states
1347
+ assert residual is not None
1348
+ aux_hidden_state = self._capture_aux_hidden_stream(
1349
+ layer_idx, prefix_sum, hidden_states, residual
1350
+ )
1351
else:
1352
assert residual is not None
1353
aux_hidden_state = hidden_states + residual