From 03a8d0b1ede68efc20b8bbfc14ee4681a62f8d75 Mon Sep 17 00:00:00 2001 From: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:03:42 -0400 Subject: [PATCH] [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 --- .../kimi_k3/test_aux_attn_res_stream.py | 197 ++++++++++++++++++ tests/models/kimi_k3/test_eagle3.py | 62 ++++++ vllm/envs.py | 8 + vllm/models/kimi_k3/nvidia/model.py | 84 +++++++- 4 files changed, 350 insertions(+), 1 deletion(-) create mode 100644 tests/models/kimi_k3/test_aux_attn_res_stream.py diff --git a/tests/models/kimi_k3/test_aux_attn_res_stream.py b/tests/models/kimi_k3/test_aux_attn_res_stream.py new file mode 100644 index 0000000000..7227ed6fd9 --- /dev/null +++ b/tests/models/kimi_k3/test_aux_attn_res_stream.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Which value the DFlash drafter is fed under AttnRes. + +`_capture_aux_hidden_stream` picks the weights it mixes against from one of +three places depending on where the tapped layer sits, and returns the plain +running prefix when the feature is off. The mixture itself is the kernel's +job and is covered by ``test_attn_res.py``; what is asserted here is the +selection, which is the part that can silently feed the drafter the wrong +tensor. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.kimi_k3.nvidia import model as k3_model + +END_LAYER = 4 + + +def _weights(tag: float) -> SimpleNamespace: + """A norm/projection pair that is identifiable by value.""" + return SimpleNamespace( + weight=torch.full((2,), tag), + variance_epsilon=tag, + ) + + +def _stub_model(*, enabled: bool, use_attn_res: bool = True) -> SimpleNamespace: + """A stand-in carrying only what the tap reads. + + Constructing the real model needs a distributed init and weights, and none + of it participates in the selection under test. + """ + consumers = [] + for i in range(END_LAYER): + consumers.append( + SimpleNamespace( + self_attention_res_norm=_weights(float(i)), + self_attention_res_proj=SimpleNamespace( + weight=torch.full((1, 2), float(i)) + ), + prev_valid_blocks=i, + ) + ) + return SimpleNamespace( + _aux_attn_res_stream=enabled, + use_attn_res=use_attn_res, + end_layer=END_LAYER, + layers=consumers, + output_attn_res_norm=_weights(99.0), + output_attn_res_proj=SimpleNamespace(weight=torch.full((1, 2), 99.0)), + num_attn_res_blocks=99, + ) + + +@pytest.fixture +def recorder(monkeypatch): + """Replace the kernel so the call it would have made is inspectable.""" + calls = [] + + def _fake_attn_res( + prefix, + delta, + block_residual, + norm_weight, + proj_weight, + output_norm_weight, + **kwargs, + ): + calls.append( + SimpleNamespace( + prefix=prefix, + delta=delta, + block_residual=block_residual, + norm_weight=norm_weight, + proj_weight=proj_weight, + kwargs=kwargs, + ) + ) + return torch.full_like(prefix, -1.0) + + monkeypatch.setattr(k3_model, "attn_res", _fake_attn_res) + return calls + + +def _set_last_rank(monkeypatch, is_last: bool): + monkeypatch.setattr( + k3_model, + "get_pp_group", + lambda: SimpleNamespace(is_last_rank=is_last), + ) + + +def _call(stub, layer_idx, prefix_sum, pending_mlp_out, block_residual): + return k3_model.KimiLinearModel._capture_aux_hidden_stream( + stub, layer_idx, prefix_sum, pending_mlp_out, block_residual + ) + + +@pytest.mark.parametrize( + "enabled,use_attn_res", [(False, True), (True, False), (False, False)] +) +def test_disabled_reproduces_the_plain_residual_sum( + recorder, monkeypatch, enabled, use_attn_res +): + """Off, the tap must be exactly the sum it replaced. + + Both conditions matter. `use_attn_res` is what constructs the norm and + projection weights, so without it the lookups below would raise rather + than fall back. + """ + _set_last_rank(monkeypatch, True) + prefix_sum = torch.tensor([1.0, 2.0]) + pending = torch.tensor([0.5, 0.25]) + + got = _call( + _stub_model(enabled=enabled, use_attn_res=use_attn_res), + 0, + prefix_sum, + pending, + torch.zeros(2), + ) + + torch.testing.assert_close(got, prefix_sum + pending) + assert not recorder, "the kernel must not run when the tap is off" + + +def test_taps_the_consumer_layer_when_one_follows(recorder, monkeypatch): + """The value the next layer reads is the mixture against *its* weights, + so the tap has to reach forward rather than use the current layer's.""" + _set_last_rank(monkeypatch, True) + + _call(_stub_model(enabled=True), 1, torch.zeros(2), None, torch.zeros(2)) + + assert len(recorder) == 1 + call = recorder[0] + # Layer 2's weights, not layer 1's. + torch.testing.assert_close(call.norm_weight, torch.full((2,), 2.0)) + assert call.kwargs["num_blocks"] == 2 + + +def test_last_layer_on_the_final_rank_uses_the_output_aggregation( + recorder, monkeypatch +): + """Nothing downstream but the model's own output-side mixture.""" + _set_last_rank(monkeypatch, True) + + _call( + _stub_model(enabled=True), END_LAYER - 1, torch.zeros(2), None, torch.zeros(2) + ) + + assert len(recorder) == 1 + torch.testing.assert_close(recorder[0].norm_weight, torch.full((2,), 99.0)) + assert recorder[0].kwargs["num_blocks"] == 99 + + +def test_last_layer_of_a_non_final_stage_falls_back(recorder, monkeypatch): + """The consumer lives on the next rank and the output aggregation only + exists on the last one, so there is nothing here to mix against. + + This is the case that would otherwise reach for weights this rank never + constructs. The forward guard is `layer_idx + 1 < end_layer`, where + `end_layer` is the rank's own exclusive bound from `get_pp_indices`, so a + `PPMissingLayer` is unreachable by construction -- the fallback below is + what makes that true rather than merely likely. + """ + _set_last_rank(monkeypatch, False) + prefix_sum = torch.tensor([3.0, 4.0]) + + got = _call( + _stub_model(enabled=True), END_LAYER - 1, prefix_sum, None, torch.zeros(2) + ) + + torch.testing.assert_close(got, prefix_sum) + assert not recorder, "no weights exist on this rank to mix against" + + +def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta( + recorder, monkeypatch +): + """The kernel writes an applied delta back into the prefix in place, which + would double-add it into the live residual stream, so the pending output + has to arrive already summed into the prefix with `delta` left None.""" + _set_last_rank(monkeypatch, True) + prefix_sum = torch.tensor([1.0, 2.0]) + pending = torch.tensor([0.5, 0.25]) + + _call(_stub_model(enabled=True), 0, prefix_sum, pending, torch.zeros(2)) + + assert len(recorder) == 1 + assert recorder[0].delta is None + torch.testing.assert_close(recorder[0].prefix, prefix_sum + pending) + # And the caller's tensor is not mutated on the way. + torch.testing.assert_close(prefix_sum, torch.tensor([1.0, 2.0])) diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py index 5f1575efef..652f6652b3 100644 --- a/tests/models/kimi_k3/test_eagle3.py +++ b/tests/models/kimi_k3/test_eagle3.py @@ -19,6 +19,7 @@ def _make_kimi_linear_model() -> KimiLinearModel: model = object.__new__(KimiLinearModel) object.__setattr__(model, "aux_hidden_state_layers", (2,)) object.__setattr__(model, "use_sequence_parallel", False) + object.__setattr__(model, "use_attn_res", False) return model @@ -137,3 +138,64 @@ def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch): torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states) assert final_attn_res.call_args.args[2] is block_residual + + +def test_attn_res_stream_capture_receives_the_layer_outputs_in_order(monkeypatch): + """Pin the argument mapping at the call site. + + The capture helper's own tests invoke it directly by keyword, so they + cannot catch a swap where `forward` hands it the residual as the pending + MLP output. Both are tensors of the same shape, so a swap is silent: it + feeds the drafter a wrong but well-formed tensor. + """ + model = _make_kimi_linear_model() + initial_hidden_states = torch.tensor([[1.0, 2.0]]) + layer_hidden_states = torch.tensor([[3.0, 4.0]]) + prefix_sum = torch.tensor([[5.0, 6.0]]) + block_residual = torch.tensor([[[7.0, 8.0]]]) + captured = torch.tensor([[11.0, 12.0]]) + + object.__setattr__(model, "start_layer", 0) + object.__setattr__(model, "end_layer", 1) + object.__setattr__( + model, + "layers", + [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))], + ) + object.__setattr__(model, "aux_hidden_state_layers", (1,)) + object.__setattr__(model, "use_attn_res", True) + object.__setattr__(model, "num_attn_res_blocks", 1) + object.__setattr__( + model, + "output_attn_res_norm", + SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), + ) + object.__setattr__( + model, + "output_attn_res_proj", + SimpleNamespace(weight=torch.ones(1, 2)), + ) + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + monkeypatch.setattr(kimi_model, "attn_res", Mock(return_value=torch.zeros(1, 2))) + monkeypatch.setenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "1") + + capture = Mock(return_value=captured) + monkeypatch.setattr(KimiLinearModel, "_capture_aux_hidden_stream", capture) + + _, aux_hidden_states = model.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=initial_hidden_states, + ) + + layer_idx, got_prefix, got_pending, got_residual = capture.call_args.args + assert layer_idx == 0 + assert got_prefix is prefix_sum + assert got_pending is layer_hidden_states + assert got_residual is block_residual + torch.testing.assert_close(aux_hidden_states[0], captured) diff --git a/vllm/envs.py b/vllm/envs.py index e68e1a335e..07a22dc7c7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -207,6 +207,7 @@ if TYPE_CHECKING: VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True VLLM_MOE_SKIP_PADDING: bool = True VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False + VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False VLLM_KIMI_K3_GEMM_RS: bool = False VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False @@ -1605,6 +1606,13 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool( int(os.getenv("VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", "0")) ), + # Kimi K3 only, and unrelated to the MoE flags above. Tap the pre-norm + # AttnRes mixture, rather than the post-mixture sum, as the auxiliary + # hidden state handed to a DFlash drafter. This changes the numerics the + # speculator sees, so it is off by default while the effect is measured. + "VLLM_KIMI_K3_AUX_ATTN_RES_STREAM": lambda: bool( + int(os.getenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "0")) + ), # Use the SM100 BF16 GEMM-RS kernel for eligible Kimi-K3 sequence-parallel # row-parallel projections. All TP ranks must belong to one NVLink domain. "VLLM_KIMI_K3_GEMM_RS": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_RS", "0"))), diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 80d0106669..d7a08a6da2 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1198,6 +1198,85 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): } ) + def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + super()._set_aux_hidden_state_layers(layers) + if self.use_attn_res: + # Emitted once, at configuration time. Which layers are tapped and + # which convention is in force are the two things you need to + # confirm from a running process, and neither is recoverable from + # the served output. + logger.info_once( + "Kimi-K3 aux hidden capture: layers=%s mode=%s " + "(VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=%d)", + layers, + "attn_res_stream" if self._aux_attn_res_stream else "prefix_only", + int(self._aux_attn_res_stream), + ) + + @property + def _aux_attn_res_stream(self) -> bool: + return envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM + + def _capture_aux_hidden_stream( + self, + layer_idx: int, + prefix_sum: torch.Tensor, + pending_mlp_out: torch.Tensor | None, + block_residual: torch.Tensor, + ) -> torch.Tensor: + """Auxiliary feature tapped after ``layer_idx`` under AttnRes. + + The wire between layers only carries the current block's running prefix; + the committed blocks live in the bank. The value the next consumer + actually reads is the pre-norm AttnRes mixture over + ``bank[:num_blocks] + prefix``, which is what the DFlash drafters were + trained against. ``attn_res`` with no delta, no block write and no + output norm computes exactly that and leaves both the prefix and the + bank untouched. + + Folding the pending MLP output into the prefix rather than passing it as + ``delta`` is deliberate: the kernel writes an applied delta back into + the prefix in place, which would double-add it into the live residual + stream. + """ + prefix = prefix_sum if pending_mlp_out is None else prefix_sum + pending_mlp_out + # `use_attn_res` is what constructs the norm and projection weights this + # reads; without it there is no mixture to compute and the attribute + # lookups below would raise. + if not (self._aux_attn_res_stream and self.use_attn_res): + return prefix + + if layer_idx + 1 < self.end_layer: + consumer = self.layers[layer_idx + 1] + score_norm = consumer.self_attention_res_norm + score_proj = consumer.self_attention_res_proj + num_blocks = consumer.prev_valid_blocks + elif get_pp_group().is_last_rank: + # Nothing downstream but the model's own output-side aggregation. + score_norm = self.output_attn_res_norm + score_proj = self.output_attn_res_proj + num_blocks = self.num_attn_res_blocks + else: + # Last layer of a non-final pipeline stage: the consumer lives on + # the next rank and the output-side aggregation only exists on the + # last one, so there is nothing here to mix against. Falling back + # to the running prefix keeps the tap defined rather than reaching + # for weights this rank does not construct. + return prefix + + return attn_res( + prefix, + None, + block_residual, + score_norm.weight, + score_proj.weight.squeeze(0), + None, + num_blocks=num_blocks, + block_write_idx=-1, + eps=score_norm.variance_epsilon, + output_norm_eps=0.0, + ) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -1265,7 +1344,10 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): if (layer_idx + 1) in self.aux_hidden_state_layers: if self.use_attn_res: assert prefix_sum is not None - aux_hidden_state = prefix_sum + hidden_states + assert residual is not None + aux_hidden_state = self._capture_aux_hidden_stream( + layer_idx, prefix_sum, hidden_states, residual + ) else: assert residual is not None aux_hidden_state = hidden_states + residual