[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>

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]))

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)

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"))),

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