From c3de0b37e25fb4b59eaa3e292df928c61b3dc172 Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Sat, 17 Oct 2026 16:22:33 +0800 Subject: [PATCH] [Core][Spec Decode] Replace EAGLE group heuristics with explicit draft-layer tags (#61455) Signed-off-by: Dakai An Signed-off-by: Yifan Qiao Co-authored-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 --- tests/v1/core/test_kv_cache_planning.py | 310 +++++------------- tests/v1/core/test_kv_cache_utils.py | 16 + tests/v1/spec_decode/test_mtp.py | 11 +- tests/v1/worker/test_attn_utils.py | 24 ++ .../kv_connector/v1/offloading/scheduler.py | 7 +- vllm/v1/core/kv_cache_planning.py | 123 +------ vllm/v1/kv_cache_interface.py | 12 +- vllm/v1/spec_decode/llm_base_proposer.py | 5 +- vllm/v1/worker/gpu/attn_utils.py | 4 +- vllm/v1/worker/gpu_model_runner.py | 2 + 10 files changed, 160 insertions(+), 354 deletions(-) diff --git a/tests/v1/core/test_kv_cache_planning.py b/tests/v1/core/test_kv_cache_planning.py index 6292868081..d60a188780 100644 --- a/tests/v1/core/test_kv_cache_planning.py +++ b/tests/v1/core/test_kv_cache_planning.py @@ -58,7 +58,6 @@ from vllm.v1.kv_cache_interface import ( 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 - - -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(), - ) - - 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), - ) - - 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 +@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]) ] - _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) - assert not any(g.is_eagle_group for g in trimmed) + 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 + ] + + +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) + + +@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( diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 6a0b6393a2..09f365d382 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/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", [ diff --git a/tests/v1/spec_decode/test_mtp.py b/tests/v1/spec_decode/test_mtp.py index 112ff77167..7cc0274298 100644 --- a/tests/v1/spec_decode/test_mtp.py +++ b/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 @@ def test_mtp_load_model_unified(mock_get_model, mock_get_layers, mock_get_pp_gro 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]) diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 0f1537403e..55910201ab 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -92,6 +92,7 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size( layers[name] = SimpleNamespace( get_kv_cache_spec=lambda _, spec=specs[name]: spec, get_attn_backend=lambda backend=backend: backend, + is_draft_layer=False, ) monkeypatch.setattr(attn_utils, "get_layers_from_vllm_config", lambda *_: layers) monkeypatch.setattr( @@ -112,6 +113,29 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size( assert all(spec.block_size == block_size for spec in specs.values()) +def test_get_kv_cache_spec_tags_draft_layers(monkeypatch): + spec = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=128, dtype=torch.bfloat16 + ) + backend = SimpleNamespace( + customize_spec=AttentionBackend.customize_spec, + get_supported_kernel_block_sizes=lambda _=None: [16], + ) + layers = { + name: SimpleNamespace( + get_kv_cache_spec=lambda _: spec, + get_attn_backend=lambda: backend, + is_draft_layer=name == "draft", + ) + for name in ("target", "draft") + } + monkeypatch.setattr(attn_utils, "get_layers_from_vllm_config", lambda *_: layers) + config = SimpleNamespace(attention_config=SimpleNamespace(hisparse_config=None)) + + specs = attn_utils.get_kv_cache_spec(config) + assert specs["draft"].is_draft and not specs["target"].is_draft + + class _FakeMetadataBuilder: def __init__(self, support: AttentionCGSupport, varlen_bound: int | None = None): self.support = support diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index a28090c8f3..cd3469c339 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -236,10 +236,9 @@ class SchedulerOffloadConfig(NamedTuple): ) if use_eagle_block_drop and not eagle_groups: # No group is annotated as holding drafter layers. This happens - # for shared-group MTP models (e.g. Qwen3.5-style), whose drafter - # is a regular decoder layer merged into the target's - # full-attention group, and for separate-draft models that miss - # annotation. Treating every group as a drafter group here makes + # when the drafter owns no KV cache (e.g. KV-sharing drafters) or + # the model runner does not tag draft layers' KV cache specs. + # Treating every group as a drafter group here makes # the volatile-tail pop apply to ALL groups, which can zero the # whole request's offload hit whenever any group's servable # window is a single chunk (issue #52735). Drafter KV served one diff --git a/vllm/v1/core/kv_cache_planning.py b/vllm/v1/core/kv_cache_planning.py index 10016409af..4d73f4303e 100644 --- a/vllm/v1/core/kv_cache_planning.py +++ b/vllm/v1/core/kv_cache_planning.py @@ -9,7 +9,7 @@ from dataclasses import replace from functools import partial, reduce from typing import Any -from vllm.config import VllmConfig +from vllm.config import SpeculativeConfig, VllmConfig from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import format_gib @@ -1038,120 +1038,20 @@ def _get_packed_kv_cache_groups( assert group_spec is not None groups.append(KVCacheGroupSpec(group_layer_names, group_spec)) - _annotate_eagle_groups( - vllm_config, - kv_cache_spec, - groups, - use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), - ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) return groups -def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: - """Whether the drafter's KV layers can be located positionally. - - MTP drafters register their KV layers after the target layers, without a - distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too. - The annotator separately checks that the groups partition the layers exactly. - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return False - if spec_config.method == "mtp": - return True - model_config = vllm_config.model_config - return model_config is not None and model_config.hf_config.model_type in ( - "deepseek_v4", - "deepseek_v41", - ) - - -def _groups_partition_layers_exactly( - kv_cache_spec: dict[str, KVCacheSpec], - kv_cache_groups: list[KVCacheGroupSpec], -) -> bool: - """Whether the groups cover every layer of ``kv_cache_spec`` exactly once. - - The trailing-layer draft fallback is only meaningful when the last - registered layer is guaranteed to appear in exactly one group. - """ - names = [name for group in kv_cache_groups for name in group.layer_names] - return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) - - -def _annotate_eagle_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - kv_cache_groups: list[KVCacheGroupSpec], - use_trailing_layer_fallback: bool = False, -) -> None: - """Flag the KV cache groups that hold drafter attention layers. - - Two detection rules, in order of preference: - - 1. Spec-driven. ``non_causal_multi_token_decode`` is declared on - MLAAttentionSpec and set by drafter attention layers that run a - non-causal multi-token decode (today only Kimi-K3 DSpark). It survives - MLAAttentionSpec.merge, so it still identifies a group after per-group - spec merging, wherever grouping happens to land. It is sufficient but - not necessary: a drafter whose spec is indistinguishable from the - target's cannot be found this way. - 2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose - MTP block reuses the target's own decoder layer and so carries no spec - marker. Their draft attention layers always register after every - target layer, so flag whichever group holds the last registered layer. - This rule is only valid where the groups partition exactly the layers - of ``kv_cache_spec``, which is re-checked here before applying it. - When a drafter's trailing caches span several groups, this rule flags - only the group holding the very last layer and must be generalized. - FIXME(yifan): avoid/generalize this hacky check. - - Args: - vllm_config: Config supplying the speculative method, if any. - kv_cache_spec: The kv cache spec of each attention layer, in layer - registration order. Only read by rule 2. - kv_cache_groups: Groups to annotate in place. - use_trailing_layer_fallback: Enable rule 2. Callers gate this on - ``_uses_trailing_mtp_layers``. - - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle_block_drop(): - return - - for group in kv_cache_groups: - if any( - getattr(spec, "non_causal_multi_token_decode", False) - for spec in iter_layer_specs(group.kv_cache_spec) - ): - group.is_eagle_group = True - - if not use_trailing_layer_fallback: - return - if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): - return - last_layer = next(reversed(kv_cache_spec)) - for group in kv_cache_groups: - if last_layer in group.layer_names: - group.is_eagle_group = True - break - - def _warn_if_unannotated_eagle_mamba( - vllm_config: VllmConfig, + spec_config: SpeculativeConfig, kv_cache_groups: list[KVCacheGroupSpec], ) -> None: """Warn when no KV cache group could be identified as the draft model's. Args: - vllm_config: Config supplying the speculative method, if any. + spec_config: Config supplying the speculative method. kv_cache_groups: Groups as they will be handed to consumers. """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return if any(group.is_eagle_group for group in kv_cache_groups): return mamba_groups = [ @@ -1297,13 +1197,6 @@ def get_kv_cache_groups( aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) groups.append(KVCacheGroupSpec([name], aligned)) - _annotate_eagle_groups( - vllm_config, - kv_cache_spec, - groups, - use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), - ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) return groups @@ -1556,7 +1449,7 @@ def _project_kv_cache_groups_to_worker( KVCacheGroupSpec( worker_layer_names, group_spec, - is_eagle_group=group.is_eagle_group and bool(worker_layer_names), + is_eagle_group=group.is_eagle_group, host_resident=group.host_resident, enable_kv_transfer=group.enable_kv_transfer, role=group.role, @@ -1636,6 +1529,14 @@ def get_kv_cache_configs( global_kv_cache_groups = get_kv_cache_groups_from_builder( builder, vllm_config, merged_kv_cache_specs ) + spec_config = vllm_config.speculative_config + if spec_config is not None and spec_config.use_eagle_block_drop(): + draft_layer_names = { + name for name, spec in merged_kv_cache_specs.items() if spec.is_draft + } + for group in global_kv_cache_groups: + group.is_eagle_group = not draft_layer_names.isdisjoint(group.layer_names) + _warn_if_unannotated_eagle_mamba(spec_config, global_kv_cache_groups) # If original_max_model_len was -1, automatically # determine the maximum model length that fits in available GPU memory. diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 79e88ca91d..b0bc91c5ed 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -175,6 +175,10 @@ class KVCacheSpec: ``MultipleOf(n)``: whole kernel blocks, when blocks are split into kernel blocks of up to ``n`` tokens.""" + is_draft: bool = field(default=False, kw_only=True, compare=False) + """Whether a draft model owns this layer; excluded from equality so it never + affects grouping.""" + def __post_init__(self): alignment = self.block_stride_alignment if isinstance(alignment, int) and alignment <= 0: @@ -660,6 +664,8 @@ class FullAttentionSpec(AttentionSpec): ) for spec in specs: for f in fields(AttentionSpec): + if not f.compare: + continue assert getattr(spec, f.name) == getattr(merged_spec, f.name), ( "All attention layers in the same KV cache group must have " "the same attention spec." @@ -747,6 +753,8 @@ class MLAAttentionSpec(FullAttentionSpec): ) for spec in specs: for f in fields(AttentionSpec): + if not f.compare: + continue assert getattr(spec, f.name) == getattr(merged_spec, f.name), ( "All attention layers in the same KV cache group must have " "the same attention spec." @@ -1208,6 +1216,8 @@ class SinkFullAttentionSpec(FullAttentionSpec): ) for spec in specs: for f in fields(AttentionSpec): + if not f.compare: + continue assert getattr(spec, f.name) == getattr(merged_spec, f.name), ( "All attention layers in the same KV cache group must have " "the same attention spec." @@ -1445,7 +1455,7 @@ class KVCacheGroupSpec: layer_names: list[str] # The KV cache spec of this manager layer kv_cache_spec: KVCacheSpec - # Whether this group contains EAGLE/MTP draft attention layers. + # Whether this group holds draft layers and takes the EAGLE last-block drop. is_eagle_group: bool = False # Host groups use the dedicated HiSparse pool; others share the device pool. host_resident: bool = False diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 1ee04d1861..d481818930 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1330,10 +1330,13 @@ class SpecDecodeBaseProposer: self.vllm_config, AttentionLayerBase, # type: ignore[type-abstract] ) + draft_attn_layer_names = set(all_attn_layers) - target_attn_layer_names + for name in draft_attn_layer_names: + all_attn_layers[name].is_draft_layer = True # Filter to only layers that have KV cache specs. self._draft_attn_layer_names = { name - for name in (set(all_attn_layers.keys()) - target_attn_layer_names) + for name in draft_attn_layer_names if all_attn_layers[name].get_kv_cache_spec(self.vllm_config) is not None } diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 8964c660a1..c10f1916d9 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable, Mapping, Sequence from contextlib import AbstractContextManager, nullcontext -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -148,6 +148,8 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: if spec := attn_module.get_kv_cache_spec(vllm_config): if isinstance(spec, AttentionSpec): spec = customize_attention_spec(attn_module.get_attn_backend(), spec) + if attn_module.is_draft_layer: + spec = replace(spec, is_draft=True) kv_cache_spec[layer_name] = spec if vllm_config.attention_config.hisparse_config is not None: kv_cache_spec = resolve_hisparse_specs(vllm_config, kv_cache_spec, attn_layers) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 3f1c3fe4c0..7a23b12b5d 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -7140,6 +7140,8 @@ class GPUModelRunner( spec = customize_attention_spec( attn_module.get_attn_backend(), spec ) + if attn_module.is_draft_layer: + spec = replace(spec, is_draft=True) kv_cache_spec[layer_name] = spec return kv_cache_spec