[II] Use pre-norm AttnRes states for Kimi-K3 DFlash

tests/models/kimi_k3/test_aux_attn_res_stream.py

@@ -0,0 +1,203 @@
+# 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.
+    """
+    model = SimpleNamespace(
+        _aux_attn_res_stream=enabled,
+        use_attn_res=use_attn_res,
+        end_layer=END_LAYER,
+    )
+    if not use_attn_res:
+        return model
+
+    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,
+            )
+        )
+    model.layers = consumers
+    model.output_attn_res_norm = _weights(99.0)
+    model.output_attn_res_proj = SimpleNamespace(weight=torch.full((1, 2), 99.0))
+    model.num_attn_res_blocks = 99
+    return model
+
+
+@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))
+    torch.testing.assert_close(call.proj_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))
+    torch.testing.assert_close(recorder[0].proj_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

@@ -22,6 +22,7 @@ def _make_kimi_linear_model() -> KimiLinearModel:
     object.__setattr__(model, "aux_hidden_state_layers", (2,))
     object.__setattr__(model, "use_sequence_parallel", False)
     object.__setattr__(model, "reuse_attn_res_output", True)
+    object.__setattr__(model, "use_attn_res", False)
     return model
 
 
@@ -190,6 +191,63 @@ def finish_auxiliary_stream(self):
     ]
     assert len(aux_hidden_states) == 1
     assert aux_hidden_states[0] is projected
+
+
+def test_attn_res_stream_capture_receives_layer_outputs_in_order(monkeypatch):
+    """Verify the positional contract between ``forward`` and the capture tap."""
+    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)
+
+
 def test_kimi_attn_res_workspace_is_reused_and_sliced():
     model = _make_kimi_linear_model()
     object.__setattr__(model, "num_attn_res_blocks", 3)

vllm/envs.py

@@ -247,6 +247,7 @@
     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_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True
     VLLM_USE_FLASHINFER_MOE_INT4: bool = False
     VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None
@@ -1844,6 +1845,13 @@ def _resolve_rust_cli_path() -> str | None:
     "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"))
+    ),
     # Allow use of FlashInfer FP8 block-scale GEMM for linear layers.
     # This uses TensorRT-LLM kernels and requires SM90+ (Hopper).
     "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(