[Model][Spec Decode] Refactor #41834 with useful improvements

tests/kernels/mhc/test_mhc_kernels.py

@@ -10,6 +10,8 @@
 import vllm.model_executor.kernels.mhc  # noqa: F401
 import vllm.model_executor.layers.mhc as mhc_layers
 from vllm.model_executor.kernels.mhc.tilelang import (
+    mhc_fused_post_pre_tilelang,
+    mhc_fused_post_pre_tilelang_reuse_residual,
     mhc_post_tilelang,
     mhc_pre_delayed_tilelang,
 )
@@ -1158,6 +1160,69 @@ def run_ref():
     torch.testing.assert_close(x, layer_input_ref, atol=1e-2, rtol=1e-2)
 
 
+def _mhc_fused_post_pre_inputs(
+    num_tokens: int, hidden_size: int, hc_mult: int = 4
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        float,
+        float,
+        float,
+        float,
+        int,
+    ],
+    torch.Tensor,
+]:
+    hc_mult3 = hc_mult * 2 + hc_mult * hc_mult
+    x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
+    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
+    mixes = (
+        torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32),
+        torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32),
+        torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4,
+        torch.randn((3,), dtype=torch.float32) * 0.1,
+        torch.randn((hc_mult3,), dtype=torch.float32) * 0.1,
+        1e-6,
+        1e-6,
+        1e-6,
+        1.0,
+        20,
+    )
+    norm_weight = torch.randn(hidden_size, dtype=torch.bfloat16)
+    return x, residual, mixes, norm_weight
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize("num_tokens", [8, 33, 1024])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
+def test_mhc_fused_post_pre_reuse_residual_is_exact(num_tokens, hidden_size):
+    """Writing the next residual into the dead input residual changes no
+    output, and the unfused path then allocates no new residual."""
+    torch.set_default_device(DEVICE)
+    set_random_seed(0)
+    x, residual, mixes, norm_weight = _mhc_fused_post_pre_inputs(
+        num_tokens, hidden_size
+    )
+
+    untouched = residual.clone()
+    expected = mhc_fused_post_pre_tilelang(x, residual, *mixes, norm_weight=norm_weight)
+    assert torch.equal(residual, untouched)
+
+    actual = mhc_fused_post_pre_tilelang_reuse_residual(
+        x, residual, *mixes, norm_weight=norm_weight
+    )
+    for got, want in zip(actual, expected):
+        assert torch.equal(got, want)
+    unfused = mhc_fused_post_pre_split_config(num_tokens, hidden_size, 4) is None
+    assert (actual[0].data_ptr() == residual.data_ptr()) == unfused
+
+
 def _rocm_mhc_inputs(num_tokens=2, hidden_size=256, hc_mult=4):
     residual = torch.randn(
         (num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16, device=DEVICE
@@ -1676,6 +1741,58 @@ def _patch_first_rank_pp_group(monkeypatch):
     )
 
 
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize("carried", [False, True])
+def test_deepseek_v4_layer_reuse_residual_keeps_caller_input(carried):
+    """reuse_residual changes no output and never writes the caller's x, even
+    when a first layer starts from ``residual = x``."""
+    torch.set_default_device(DEVICE)
+    set_random_seed(0)
+    hc_mult, hidden_size, num_tokens = 4, 4096, 64
+    layer = _make_mhc_decoder_layer(hc_mult=hc_mult, hidden_size=hidden_size)
+    _, _, mixes, norm_weight = _mhc_fused_post_pre_inputs(num_tokens, hidden_size)
+    _, _, fn, hc_scale, hc_base, *_ = mixes
+    layer.hc_attn_fn.data.copy_(fn)
+    layer.hc_attn_scale, layer.hc_attn_base = hc_scale, hc_base
+    layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base = fn, hc_scale, hc_base
+    norm = SimpleNamespace(weight=norm_weight, variance_epsilon=1e-6)
+    layer.attn_norm = layer.ffn_norm = norm
+    layer.rms_norm_eps, layer.hc_eps, layer.hc_post_alpha = 1e-6, 1e-6, 1.0
+    layer.hc_sinkhorn_iters = 20
+    layer.use_sequence_parallel = False
+    layer.attn = lambda positions, x, _: x * 0.5
+    layer.ffn = lambda x, input_ids, metadata: x * 2.0
+    positions = torch.arange(num_tokens)
+
+    x = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
+    if carried:
+        residual = torch.randn_like(x)
+        x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
+        post_mix = torch.rand((num_tokens, hc_mult, 1), dtype=torch.float32)
+        res_mix = torch.rand((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
+    else:
+        residual = post_mix = res_mix = None
+
+    def run(reuse_residual):
+        return DeepseekV4DecoderLayer.forward(
+            layer,
+            x.clone() if carried else x,
+            positions,
+            None,
+            post_mix,
+            res_mix,
+            None if residual is None else residual.clone(),
+            reuse_residual=reuse_residual,
+        )
+
+    x_before = x.clone()
+    expected = run(False)
+    actual = run(True)
+    assert torch.equal(x, x_before)
+    for got, want in zip(actual, expected):
+        assert torch.equal(got, want)
+
+
 def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch):
     """First finalize (at the end of load_weights) allocates
     hc_attn_fn_broadcast as hc_attn_fn summed over hc streams."""

tests/v1/spec_decode/test_eagle.py

@@ -34,6 +34,7 @@
 from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
+from vllm.v1.spec_decode.utils import extend_all_queries_by_N
 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 
 model_dir = "meta-llama/Llama-3.1-8B-Instruct"
@@ -1004,6 +1005,124 @@ def create_deterministic_logits(token_ids):
     assert torch.equal(result, expected_tokens)
 
 
+@pytest.mark.skipif(current_platform.is_rocm(), reason="uses FLASH_ATTN")
+def test_propose_refreshes_layout_caches_each_draft_step():
+    """Builders called inside the draft loop must see token_to_req_indices()
+    and compute_num_computed_tokens() for the current layout, not values
+    cached during the ragged first pass or an earlier draft step."""
+    device = torch.device(DEVICE_TYPE)
+    num_speculative_tokens = 3
+    seq_lens = [5, 3]
+    batch_size, total_tokens, vocab_size = len(seq_lens), sum(seq_lens), 100
+
+    proposer = _create_proposer(
+        "eagle", num_speculative_tokens, attention_backend="FLASH_ATTN"
+    )
+    hidden_size = proposer.hidden_size
+    model_mock = mock.MagicMock()
+    model_mock.side_effect = [
+        (torch.zeros(n, hidden_size, device=device),) * 2
+        for n in [total_tokens] + [batch_size] * (num_speculative_tokens - 1)
+    ]
+    model_mock.compute_logits.side_effect = [
+        torch.zeros(batch_size, vocab_size, device=device)
+    ] * num_speculative_tokens
+    proposer.model = model_mock
+    proposer._draft_attn_layer_names = {"layer.0"}
+
+    builder_cls, _ = try_get_attention_backend(AttentionBackendEnum.FLASH_ATTN)
+    builder = builder_cls(
+        kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
+        layer_names=["layer.0"],
+        vllm_config=proposer.vllm_config,
+        device=device,
+    )
+    req_index_buffer = torch.zeros(64, dtype=torch.int32, device=device)
+    observed = []
+
+    def build_for_drafting(common_attn_metadata, draft_index):
+        query_lens = common_attn_metadata.query_start_loc.diff()
+        observed.append(
+            (
+                common_attn_metadata.token_to_req_indices(req_index_buffer).tolist(),
+                torch.repeat_interleave(
+                    torch.arange(len(query_lens), device=device), query_lens
+                ).tolist(),
+                common_attn_metadata.compute_num_computed_tokens().tolist(),
+                (common_attn_metadata.seq_lens - query_lens).tolist(),
+            )
+        )
+        return builder.build_for_drafting(
+            common_attn_metadata=common_attn_metadata, draft_index=draft_index
+        )
+
+    attn_group = mock.MagicMock()
+    attn_group.get_metadata_builder.return_value = builder
+    attn_group.build_metadata_for_drafting = build_for_drafting
+    attn_group.layer_names = ["layer.0"]
+    attn_group.kv_cache_spec = builder.kv_cache_spec
+    proposer.draft_attn_groups = [attn_group]
+
+    proposer.propose(
+        num_speculative_tokens=num_speculative_tokens,
+        target_token_ids=torch.randint(0, vocab_size, (total_tokens,), device=device),
+        target_positions=torch.cat([torch.arange(n, device=device) for n in seq_lens]),
+        target_hidden_states=torch.randn(total_tokens, hidden_size, device=device),
+        next_token_ids=torch.randint(
+            0, vocab_size, (batch_size,), dtype=torch.int32, device=device
+        ),
+        token_indices_to_sample=None,
+        common_attn_metadata=create_common_attn_metadata(
+            BatchSpec(seq_lens=seq_lens, query_lens=seq_lens),
+            block_size=BLOCK_SIZE,
+            device=device,
+        ),
+        sampling_metadata=mock.MagicMock(),
+    )
+
+    assert len(observed) == num_speculative_tokens
+    for step, (req, want_req, computed, want_computed) in enumerate(observed):
+        assert req == want_req, f"draft step {step}"
+        assert computed == want_computed, f"draft step {step}"
+
+
+def test_extend_all_queries_by_n_refreshes_token_to_req_indices():
+    """The extended layout must not reuse the token->request mapping cached
+    for the layout it was extended from."""
+    device = torch.device(DEVICE_TYPE)
+    common_attn_metadata = create_common_attn_metadata(
+        BatchSpec(seq_lens=[8, 4, 4], query_lens=[3, 1, 1]),
+        block_size=BLOCK_SIZE,
+        device=device,
+    )
+    req_index_buffer = torch.zeros(16, dtype=torch.int32, device=device)
+    assert common_attn_metadata.token_to_req_indices(req_index_buffer).tolist() == [
+        0,
+        0,
+        0,
+        1,
+        2,
+    ]
+
+    extended = extend_all_queries_by_N(
+        common_attn_metadata,
+        N=1,
+        arange=torch.arange(8, dtype=torch.int32, device=device),
+        new_slot_mapping=torch.zeros(8, dtype=torch.int64, device=device),
+    )
+
+    assert extended.token_to_req_indices(req_index_buffer).tolist() == [
+        0,
+        0,
+        0,
+        0,
+        1,
+        1,
+        2,
+        2,
+    ]
+
+
 @pytest.mark.parametrize(
     "attn_backend",
     ["ROCM_ATTN", "TRITON_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"],

tests/v1/spec_decode/test_mtp.py

@@ -24,6 +24,7 @@
 )
 from vllm.config.load import LoadConfig
 from vllm.model_executor.models.llama import LlamaForCausalLM
+from vllm.models.deepseek_v4.common.spec_decode import needs_mtp_hidden_buffer
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.spec_decode.eagle import EagleProposer
@@ -237,3 +238,24 @@ def create_deterministic_logits(batch_size, vocab_size, token_offset):
     assert model_mock.called
     # Verify output shape
     assert result.shape == (batch_size, num_speculative_tokens)
+
+
+@pytest.mark.parametrize(
+    ("method", "expected"),
+    [
+        (None, False),
+        ("ngram", False),
+        ("mtp", True),
+        ("eagle3", True),
+        ("draft_model", True),
+        ("dspark", False),
+    ],
+)
+def test_deepseek_v4_mtp_hidden_buffer_skips_dspark(method, expected):
+    """Only drafters that read DeepSeek V4's pre-hc_head residual get the
+    buffer; DSpark drafts from the target's aux hidden states."""
+    spec_config = None
+    if method is not None:
+        spec_config = SpeculativeConfig.__new__(SpeculativeConfig)
+        object.__setattr__(spec_config, "method", method)
+    assert needs_mtp_hidden_buffer(spec_config) is expected