iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Core][Spec Decode] Replace EAGLE group heuristics with explicit draft-layer tags (#61455)

Browse Source 
Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Dakai An
2026-10-17 08:22:33 +00:00
committed by GitHub
co-authored by Yifan Qiao Claude Opus 5.5
parent 08a65e287e
commit c3de0b37e2
10 changed files with 160 additions and 354 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/core/test_kv_cache_planning.py
+82 -228
View File
@@ -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(
tests/v1/core/test_kv_cache_utils.py
+16
View File
@@ -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
+3 -8
View File
@@ -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])
tests/v1/worker/test_attn_utils.py
+24
View File
@@ -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
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+3 -4
View File
@@ -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
vllm/v1/core/kv_cache_planning.py
+12 -111
View File
@@ -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.
vllm/v1/kv_cache_interface.py
+11 -1
View File
@@ -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
vllm/v1/spec_decode/llm_base_proposer.py
+4 -1
View File
@@ -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
}
vllm/v1/worker/gpu/attn_utils.py
+3 -1
View File
@@ -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)
vllm/v1/worker/gpu_model_runner.py
+2
View File
@@ -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