[Core][Spec Decode] Replace EAGLE group heuristics with explicit draft-layer tags

tests/v1/core/test_kv_cache_planning.py

@@ -58,7 +58,6 @@
     SlidingWindowSpec,
     SparseCacheRole,
     UniformTypeKVCacheSpecs,
-    iter_layer_specs,
 )
 from vllm.v1.kv_cache_layout import KVCacheLayout
 
@@ -1945,247 +1944,102 @@ def test_check_enough_kv_cache_memory_reserves_null_block():
     )
 
 
-def _spec_decode_grouping_config(method="dspark", model_type=None):
-    """Grouping config with an EAGLE-family speculative method enabled."""
-    return SimpleNamespace(
-        scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
-        cache_config=SimpleNamespace(
-            get_resolved_kv_cache_layout=lambda: SimpleNamespace(
-                is_block_outermost=True
-            )
-        ),
-        model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)),
-        speculative_config=SimpleNamespace(
-            method=method,
-            use_eagle=lambda: True,
-            use_eagle_block_drop=lambda: True,
-        ),
+def _spec_decode_kv_cache_groups(worker_specs, use_eagle_block_drop=True):
+    config = VllmConfig(model_config=ModelConfig(max_model_len=64))
+    config.cache_config.kv_cache_layout = "BLHNC"
+    config.speculative_config = SimpleNamespace(
+        method="dflash",
+        use_eagle=lambda: True,
+        use_eagle_block_drop=lambda: use_eagle_block_drop,
+        use_multi_module_mtp=lambda: False,
+        parallel_drafting=False,
     )
+    configs = get_kv_cache_configs(config, worker_specs, [2**30] * len(worker_specs))
+    return generate_scheduler_kv_cache_config(configs).kv_cache_groups
 
 
-def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False):
-    """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a
-    DSpark-style draft MLA layer marked non_causal_multi_token_decode.
+def _as_draft(spec):
+    return replace(spec, is_draft=True)
 
-    The target's fp8 KV dtype is one difference from the draft on Kimi-K3
-    (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass
-    draft_shares_target_spec to match that dtype while retaining the draft
-    marker as a group boundary.
-    """
-    target_dtype = None if draft_shares_target_spec else "fp8_e4m3"
-    specs = {
-        "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype),
-        "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype),
-        "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"),
-        "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"),
-    }
+
+def _hybrid_specs(draft=True):
+    """Qwen3.5-shaped [GDN x3, full-attn] x2, plus an MTP full-attn layer."""
+    specs = {}
+    for i in range(8):
+        specs[f"layers.{i}"] = (
+            new_kv_cache_spec()
+            if i % 4 == 3
+            else new_mamba_spec(mamba_cache_mode="align")
+        )
     if draft:
-        draft_spec = new_mla_spec(block_size=64)
-        specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True)
+        specs["mtp.0"] = _as_draft(new_kv_cache_spec())
     return specs
 
 
-def test_draft_group_annotated_on_hybrid_general_path():
-    # A drafter's MLA layer carries non_causal_multi_token_decode, so its group
-    # is identifiable without keying off a model version.
-    groups = get_kv_cache_groups(
-        _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True)
-    )
-
-    flagged = [g for g in groups if g.is_eagle_group]
-    assert len(flagged) == 1
-    assert "draft.attn.0" in flagged[0].layer_names
-
-
-def test_mamba_groups_never_flagged_even_when_draft_shares_a_group():
-    # Packed uniform-type groups can contain distinct target and draft layer
-    # specs; the combined group still holds volatile draft KV and must be
-    # flagged. What must never happen is a Mamba group being flagged: that
-    # widens its lookup window to two consecutive chunks, which align-mode
-    # checkpointing never produces, zeroing every lookup.
-    groups = get_kv_cache_groups(
-        _spec_decode_grouping_config(),
-        _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True),
-    )
-
-    for group in groups:
-        if "draft.attn.0" in group.layer_names:
-            assert group.is_eagle_group
-        if isinstance(group.kv_cache_spec, MambaSpec):
-            assert not group.is_eagle_group
-
-
-def test_draft_group_not_annotated_without_spec_decode():
-    # The marker alone must not flag anything; the eagle semantics only apply
-    # when a speculative method is actually enabled.
-    config = _spec_decode_grouping_config()
-    config.speculative_config = None
-    groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True))
-
-    assert not any(g.is_eagle_group for g in groups)
-
-
-def test_unidentifiable_draft_with_mamba_warns(caplog_vllm):
-    # No group carries the draft marker, so consumers fall back to
-    # conservative behavior that silently breaks reuse for Mamba groups.
-    # That must at least be visible.
-    groups = get_kv_cache_groups(
-        _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False)
-    )
-
-    assert not any(g.is_eagle_group for g in groups)
-    assert "no KV cache group could be identified as the draft model's" in (
-        caplog_vllm.text
-    )
-
-
-def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm):
-    # Pure-attention models degrade gracefully under the consumers'
-    # conservative fallback (a one-block hit drop at most), so the warning
-    # stays silent to avoid noise on every unannotated EAGLE deployment.
-    specs = {
-        "target.attn.0": new_mla_spec(block_size=64),
-        "target.attn.1": new_mla_spec(block_size=64),
-    }
-    groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs)
-
-    assert not any(g.is_eagle_group for g in groups)
-    assert "could be identified as the draft model's" not in caplog_vllm.text
-
-
-def test_no_warning_when_draft_group_is_identified(caplog_vllm):
-    get_kv_cache_groups(
-        _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True)
-    )
-
-    assert "could be identified as the draft model's" not in caplog_vllm.text
-
-
-def _deepseek_v4_specs(model_version="deepseek_v4"):
-    """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers
-    at differing page sizes, with the MTP draft layer registered last."""
-    return {
-        "model.layers.0.self_attn.attn": new_mla_spec(),
-        "model.layers.1.self_attn.attn": new_mla_spec(),
-        "model.layers.2.self_attn.attn": new_swa_mla_spec(
-            head_size=1024, model_version=model_version
-        ),
-        # The MTP block registers last, and its sliding-window size differs, so
-        # it lands in a group of its own.
-        "model.layers.3.self_attn.attn": new_swa_mla_spec(
-            head_size=1024, sliding_window=256, model_version=model_version
-        ),
-    }
+_DRAFT_SPEC_SHAPES = {
+    "uniform": lambda: {
+        "layers.0": new_kv_cache_spec(),
+        "draft.0": _as_draft(new_kv_cache_spec()),
+    },
+    "hybrid": _hybrid_specs,
+    "draft_spans_groups": lambda: {
+        "layers.0": new_mla_spec(),
+        "layers.1": new_swa_mla_spec(head_size=1024),
+        "mtp.0": _as_draft(new_mla_spec()),
+        "mtp.1": _as_draft(new_swa_mla_spec(head_size=1024)),
+    },
+}
 
 
 @pytest.mark.parametrize(
-    ("method", "model_type"),
-    [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")],
+    "make_specs", _DRAFT_SPEC_SHAPES.values(), ids=list(_DRAFT_SPEC_SHAPES)
 )
-def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type):
-    # DeepseekV4's MTP block reuses the target's decoder layer, so its spec
-    # carries no draft marker and only the positional rule can find it. This
-    # pins the pre-existing behaviour that the unified annotator must preserve.
-    groups = get_kv_cache_groups(
-        _spec_decode_grouping_config(method=method, model_type=model_type),
-        _deepseek_v4_specs(model_version=None),
-    )
-
-    flagged = [g for g in groups if g.is_eagle_group]
-    assert len(flagged) == 1
-    assert "model.layers.3.self_attn.attn" in flagged[0].layer_names
-
+@pytest.mark.parametrize("pp", [1, 2])
+def test_eagle_groups_are_the_groups_holding_draft_layers(make_specs, pp):
+    specs = make_specs()
+    draft_layers = {name for name, spec in specs.items() if spec.is_draft}
+    untagged = {name: replace(spec, is_draft=False) for name, spec in specs.items()}
+    groups = _spec_decode_kv_cache_groups([specs])
+    assert [(g.layer_names, g.kv_cache_spec) for g in groups] == [
+        (g.layer_names, g.kv_cache_spec)
+        for g in _spec_decode_kv_cache_groups([untagged])
+    ]
 
-def test_trailing_layer_fallback_applies_to_any_mtp_model():
-    # The positional rule is sound for every MTP drafter, not just DeepseekV4:
-    # MTP blocks reuse the target's decoder layer (no spec marker) and always
-    # register after every target layer. The model_type must not gate it.
-    groups = get_kv_cache_groups(
-        _spec_decode_grouping_config(method="mtp", model_type="other"),
-        _deepseek_v4_specs(),
+    worker_specs = [
+        {n: s for n, s in specs.items() if n not in draft_layers},
+        {n: s for n, s in specs.items() if n in draft_layers},
+    ]
+    scheduler_groups = _spec_decode_kv_cache_groups(
+        [specs] if pp == 1 else worker_specs
     )
+    assert [g.is_eagle_group for g in scheduler_groups] == [
+        bool(draft_layers & set(g.layer_names)) for g in groups
+    ]
 
-    flagged = [g for g in groups if g.is_eagle_group]
-    assert len(flagged) == 1
-    assert "model.layers.3.self_attn.attn" in flagged[0].layer_names
-
-
-def _qwen3_5_hybrid_specs(with_mtp_layer: bool):
-    """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with
-    the MTP drafter's full-attn layer (spec-identical to the target's)
-    registered last."""
-    specs = {}
-    idx = 0
-    for _ in range(2):
-        for _ in range(3):
-            specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec(
-                mamba_cache_mode="align"
-            )
-            idx += 1
-        specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec()
-        idx += 1
-    if with_mtp_layer:
-        specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec()
-    return specs
-
-
-def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm):
-    # A hybrid mamba + full-attention model with an MTP drafter that is
-    # spec-indistinguishable from the target reaches the general multi-group
-    # path. The trailing-layer rule must locate the draft group there so the
-    # Mamba groups are not swept up by the flag-all consumer fallback.
-    groups = get_kv_cache_groups(
-        _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"),
-        _qwen3_5_hybrid_specs(with_mtp_layer=True),
-    )
-
-    flagged = [g for g in groups if g.is_eagle_group]
-    assert len(flagged) == 1
-    assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names
-    for group in groups:
-        if any(
-            isinstance(spec, MambaSpec)
-            for spec in iter_layer_specs(group.kv_cache_spec)
-        ):
-            assert not group.is_eagle_group
-    assert "could be identified as the draft model's" not in caplog_vllm.text
-
-
-@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"])
-def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method):
-    # Other drafters are not covered by the trailing-layer rule, so an
-    # unidentifiable hybrid draft must still warn.
-    groups = get_kv_cache_groups(
-        _spec_decode_grouping_config(method=method, model_type="qwen3_5"),
-        _qwen3_5_hybrid_specs(with_mtp_layer=True),
-    )
 
+def test_no_eagle_groups_without_eagle_block_drop():
+    groups = _spec_decode_kv_cache_groups([_hybrid_specs()], use_eagle_block_drop=False)
     assert not any(g.is_eagle_group for g in groups)
-    assert "could be identified as the draft model's" in caplog_vllm.text
-
 
-def test_trailing_layer_fallback_requires_exact_partition():
-    # If the groups do not partition the layers exactly (e.g. a caller that
-    # dropped or duplicated layers), the positional rule is meaningless and
-    # must not fire.
-    from vllm.v1.core.kv_cache_planning import _annotate_eagle_groups
 
-    specs = _qwen3_5_hybrid_specs(with_mtp_layer=True)
-    config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5")
-    groups = get_kv_cache_groups(config, specs)
-    for g in groups:
-        g.is_eagle_group = False
-    # Remove one layer from its group: no longer an exact partition.
-    trimmed = [
-        KVCacheGroupSpec(
-            [n for n in g.layer_names if n != "model.layers.0.linear_attn"],
-            g.kv_cache_spec,
-        )
-        for g in groups
-    ]
-    _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True)
-
-    assert not any(g.is_eagle_group for g in trimmed)
+@pytest.mark.parametrize(
+    "make_specs,warns",
+    [
+        (lambda: _hybrid_specs(draft=False), True),
+        (
+            lambda: {
+                "layers.0": new_kv_cache_spec(),
+                "draft.0": new_kv_cache_spec(),
+            },
+            False,
+        ),
+    ],
+    ids=["untagged_hybrid", "untagged_uniform_attention"],
+)
+def test_warns_only_when_untagged_draft_meets_mamba(make_specs, warns, caplog_vllm):
+    _spec_decode_kv_cache_groups([make_specs()])
+    warned = "could be identified as the draft model's" in caplog_vllm.text
+    assert warned == warns
 
 
 _GQA_SPEC = FullAttentionSpec(

tests/v1/core/test_kv_cache_utils.py

@@ -1303,6 +1303,22 @@ def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla):
         type(draft).merge([draft, replace(draft, dcp_sharded=True)])
 
 
+@pytest.mark.parametrize(
+    "target",
+    [
+        new_kv_cache_spec(),
+        new_mla_spec(),
+        SinkFullAttentionSpec(
+            block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32, sink_len=4
+        ),
+    ],
+)
+def test_draft_layer_tag_is_not_part_of_the_cache_format(target):
+    draft = replace(target, is_draft=True)
+    assert draft == target and hash(draft) == hash(target)
+    assert type(target).merge([target, draft]) == target
+
+
 @pytest.mark.parametrize(
     "layer_type,dcp_size,expected_width",
     [

tests/v1/spec_decode/test_mtp.py

@@ -103,15 +103,8 @@ def test_mtp_load_model_unified(mock_get_model, mock_get_layers, mock_get_pp_gro
 
     target_attn_layers = {"target_attn_1": mock.MagicMock()}
     all_attn_layers = {**target_attn_layers, "draft_attn_1": mock.MagicMock()}
-    target_indexer_layers: dict = {}
-    all_indexer_layers: dict = {}
 
-    mock_get_layers.side_effect = [
-        target_attn_layers,
-        target_indexer_layers,
-        all_attn_layers,
-        all_indexer_layers,
-    ]
+    mock_get_layers.side_effect = [target_attn_layers, all_attn_layers]
 
     mock_pp_group = mock.MagicMock()
     mock_pp_group.world_size = 1
@@ -138,6 +131,8 @@ class _TargetModelStub(LlamaForCausalLM):
     assert proposer.model.lm_head == target_model.lm_head
     # MTP shares embed_tokens with target model
     assert proposer.model.model.embed_tokens == target_model.model.embed_tokens
+    assert proposer._draft_attn_layer_names == {"draft_attn_1"}
+    assert all_attn_layers["draft_attn_1"].is_draft_layer is True
 
 
 @pytest.mark.parametrize("num_speculative_tokens", [1])