@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 93ab92be0c

[Bugfix][AMD] Annotate draft KV cache groups on the hybrid grouping path (#52047)

Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com> Co-authored-by: okorzh-amd <okorzh-amd@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>

Oxana Korzh committed Aug 28, 2026 at 20:53 UTC 93ab92be0cde1f33a0532e0096d5a42d3007f244

2 files changed +227 -6

tests/v1/core/test_kv_cache_utils.py

+143 -1

@@ -4,6 +4,7 @@ import copy

4 import hashlib

5 import importlib

6 from collections.abc import Callable

7 +from dataclasses import replace

8 from types import SimpleNamespace

9 from typing import Any

@@ -2214,13 +2215,14 @@ def new_mla_spec(cache_dtype_str=None, block_size=16):

2215 )

2216

2217

2217 -def new_swa_mla_spec(head_size=576, sliding_window=128):

2218 +def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None):

2219 return SlidingWindowMLASpec(

2220 block_size=16,

2221 num_kv_heads=1,

2222 head_size=head_size,

2223 dtype=torch.float32,

2224 sliding_window=sliding_window,

2225 + model_version=model_version,

2226 )

2227

2228

@@ -3277,3 +3279,143 @@ def test_iter_layer_specs_returns_group_members():

3279 block_size=4, kv_cache_specs={"a": full, "b": mla}

3280 )

3281 assert list(iter_layer_specs(wrapped)) == [full, mla]

3282 +

3283 +

3284 +def _spec_decode_grouping_config(method="dspark"):

3285 + """Grouping config with an EAGLE-family speculative method enabled."""

3286 + return SimpleNamespace(

3287 + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),

3288 + speculative_config=SimpleNamespace(

3289 + method=method,

3290 + use_eagle=lambda: True,

3291 + ),

3292 + )

3293 +

3294 +

3295 +def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False):

3296 + """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a

3297 + DSpark-style draft MLA layer marked non_causal_multi_token_decode.

3298 +

3299 + The target's fp8 KV dtype is what keeps the draft in its own bucket, as it

3300 + does on Kimi-K3 (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass

3301 + draft_shares_target_spec to collapse them into one group instead.

3302 + """

3303 + target_dtype = None if draft_shares_target_spec else "fp8_e4m3"

3304 + specs = {

3305 + "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype),

3306 + "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype),

3307 + "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"),

3308 + "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"),

3309 + }

3310 + if draft:

3311 + draft_spec = new_mla_spec(block_size=64)

3312 + specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True)

3313 + return specs

3314 +

3315 +

3316 +def test_draft_group_annotated_on_hybrid_general_path():

3317 + # A drafter's MLA layer carries non_causal_multi_token_decode, so its group

3318 + # is identifiable without keying off a model version.

3319 + groups = get_kv_cache_groups(

3320 + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True)

3321 + )

3322 +

3323 + flagged = [g for g in groups if g.is_eagle_group]

3324 + assert len(flagged) == 1

3325 + assert flagged[0].layer_names == ["draft.attn.0"]

3326 +

3327 +

3328 +def test_mamba_groups_never_flagged_even_when_draft_shares_a_group():

3329 + # MLAAttentionSpec.merge ORs non_causal_multi_token_decode rather than

3330 + # requiring equality, so a draft layer can share a group with target

3331 + # layers; that combined group still holds volatile draft KV and must be

3332 + # flagged. What must never happen is a Mamba group being flagged: that

3333 + # widens its lookup window to two consecutive chunks, which align-mode

3334 + # checkpointing never produces, zeroing every lookup.

3335 + groups = get_kv_cache_groups(

3336 + _spec_decode_grouping_config(),

3337 + _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True),

3338 + )

3339 +

3340 + for group in groups:

3341 + if "draft.attn.0" in group.layer_names:

3342 + assert group.is_eagle_group

3343 + if isinstance(group.kv_cache_spec, MambaSpec):

3344 + assert not group.is_eagle_group

3345 +

3346 +

3347 +def test_draft_group_not_annotated_without_spec_decode():

3348 + # The marker alone must not flag anything; the eagle semantics only apply

3349 + # when a speculative method is actually enabled.

3350 + config = SimpleNamespace(

3351 + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),

3352 + speculative_config=None,

3353 + )

3354 + groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True))

3355 +

3356 + assert not any(g.is_eagle_group for g in groups)

3357 +

3358 +

3359 +def test_unidentifiable_draft_with_mamba_warns(caplog_vllm):

3360 + # No group carries the draft marker, so every consumer falls back to

3361 + # flagging all groups -- including Mamba ones, which then can never report

3362 + # a hit. That is silent today; it must at least be visible.

3363 + groups = get_kv_cache_groups(

3364 + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False)

3365 + )

3366 +

3367 + assert not any(g.is_eagle_group for g in groups)

3368 + assert "no KV cache group could be identified as the draft model's" in (

3369 + caplog_vllm.text

3370 + )

3371 + assert "Mamba groups" in caplog_vllm.text

3372 +

3373 +

3374 +def test_no_warning_when_draft_group_is_identified(caplog_vllm):

3375 + get_kv_cache_groups(

3376 + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True)

3377 + )

3378 +

3379 + assert "could be identified as the draft model's" not in caplog_vllm.text

3380 +

3381 +

3382 +def _deepseek_v4_specs(model_version="deepseek_v4"):

3383 + """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers

3384 + at differing page sizes, with the MTP draft layer registered last."""

3385 + return {

3386 + "model.layers.0.self_attn.attn": new_mla_spec(),

3387 + "model.layers.1.self_attn.attn": new_mla_spec(),

3388 + "model.layers.2.self_attn.attn": new_swa_mla_spec(

3389 + head_size=1024, model_version=model_version

3390 + ),

3391 + # The MTP block registers last, and its sliding-window size differs, so

3392 + # it lands in a group of its own.

3393 + "model.layers.3.self_attn.attn": new_swa_mla_spec(

3394 + head_size=1024, sliding_window=256, model_version=model_version

3395 + ),

3396 + }

3397 +

3398 +

3399 +def test_deepseek_v4_draft_group_annotated_on_group_and_unify_path():

3400 + # DeepseekV4's MTP block reuses the target's decoder layer, so its spec

3401 + # carries no draft marker and only the positional rule can find it. This

3402 + # pins the pre-existing behaviour that the unified annotator must preserve.

3403 + groups = get_kv_cache_groups(

3404 + _spec_decode_grouping_config(method="mtp"), _deepseek_v4_specs()

3405 + )

3406 +

3407 + flagged = [g for g in groups if g.is_eagle_group]

3408 + assert len(flagged) == 1

3409 + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names

3410 +

3411 +

3412 +def test_deepseek_v4_annotation_requires_model_version():

3413 + # The positional rule is only sound for DeepseekV4, where the draft layer

3414 + # is known to be registered last. Without that model gate nothing may be

3415 + # flagged, however the grouping happens to fall out.

3416 + groups = get_kv_cache_groups(

3417 + _spec_decode_grouping_config(method="mtp"),

3418 + _deepseek_v4_specs(model_version=None),

3419 + )

3420 +

3421 + assert not any(g.is_eagle_group for g in groups)

vllm/v1/core/kv_cache_utils.py

+84 -5

@@ -1742,23 +1742,56 @@ def _get_kv_cache_groups_uniform_groups(

1742 return [full_mla_group, *swa_mla_groups]

1743

1744

1745 -def _annotate_eagle_groups_deepseek_v4(

1745 +def _annotate_eagle_groups(

1746 vllm_config: VllmConfig,

1747 kv_cache_spec: dict[str, KVCacheSpec],

1748 kv_cache_groups: list[KVCacheGroupSpec],

1749 + use_deepseek_v4_fallback: bool = False,

1750 ) -> None:

1751 + """Flag the KV cache groups that hold drafter attention layers.

1752 +

1753 + Two detection rules, in order of preference:

1754 +

1755 + 1. Spec-driven. ``non_causal_multi_token_decode`` is declared on

1756 + MLAAttentionSpec and set by drafter attention layers that run a

1757 + non-causal multi-token decode (today only Kimi-K3 DSpark). It survives

1758 + MLAAttentionSpec.merge, so it still identifies a group after per-group

1759 + spec merging, wherever grouping happens to land. It is sufficient but

1760 + not necessary: a drafter whose spec is indistinguishable from the

1761 + target's cannot be found this way.

1762 + 2. Model-scoped positional fallback for DeepseekV4, whose MTP block reuses

1763 + the target's own decoder layer and so carries no spec marker. Its draft

1764 + attention layer is always the last registered layer, so flag whichever

1765 + group holds it. This rule is only valid where the groups partition

1766 + exactly the layers of ``kv_cache_spec``, which is true on the

1767 + group_and_unify path and not in general; other callers must leave

1768 + ``use_deepseek_v4_fallback`` False.

1769 + FIXME(yifan): avoid/generalize this hacky check.

1770 +

1771 + Args:

1772 + vllm_config: Config supplying the speculative method, if any.

1773 + kv_cache_spec: The kv cache spec of each attention layer, in layer

1774 + registration order. Only read by rule 2.

1775 + kv_cache_groups: Groups to annotate in place.

1776 + use_deepseek_v4_fallback: Enable rule 2. Only the group_and_unify path

1777 + may set this.

1778 + """

1779 spec_config = vllm_config.speculative_config

1780 if spec_config is None or not spec_config.use_eagle():

1781 return

1782 +

1783 + for group in kv_cache_groups:

1784 + if getattr(group.kv_cache_spec, "non_causal_multi_token_decode", False):

1785 + group.is_eagle_group = True

1786 +

1787 + if not use_deepseek_v4_fallback:

1788 + return

1789 # Detection uses the merged MLA spec's model_version.

1790 if not any(

1791 getattr(spec, "model_version", None) == "deepseek_v4"

1792 for spec in kv_cache_spec.values()

1793 ):

1794 return

1759 - # DeepseekV4's MTP attention layer is always the last layer, and we flag whichever

1760 - # group contains it.

1761 - # FIXME(yifan): avoid/generalize this hacky check.

1795 last_layer = next(reversed(kv_cache_spec))

1796 for group in kv_cache_groups:

1797 if last_layer in group.layer_names:

@@ -1766,6 +1799,45 @@ def _annotate_eagle_groups_deepseek_v4(

1799 break

1800

1801

1802 +def _warn_if_unannotated_eagle_mamba(

1803 + vllm_config: VllmConfig,

1804 + kv_cache_groups: list[KVCacheGroupSpec],

1805 +) -> None:

1806 + """Warn when the flag-all eagle fallback will silently disable reuse.

1807 +

1808 + With no group annotated, consumers flag every group as a draft group. That

1809 + widens a Mamba group's required lookup window to two consecutive chunks,

1810 + which align-mode checkpointing never produces, so reuse drops to zero with

1811 + no error and no metric to show it.

1812 +

1813 + Args:

1814 + vllm_config: Config supplying the speculative method, if any.

1815 + kv_cache_groups: Groups as they will be handed to consumers.

1816 + """

1817 + spec_config = vllm_config.speculative_config

1818 + if spec_config is None or not spec_config.use_eagle():

1819 + return

1820 + if any(group.is_eagle_group for group in kv_cache_groups):

1821 + return

1822 + mamba_groups = [

1823 + idx

1824 + for idx, group in enumerate(kv_cache_groups)

1825 + if isinstance(group.kv_cache_spec, MambaSpec)

1826 + ]

1827 + if not mamba_groups:

1828 + return

1829 + logger.warning(

1830 + "Speculative decoding (method=%s) is enabled but no KV cache group "

1831 + "could be identified as the draft model's, so every group -- "

1832 + "including Mamba groups %s -- will be treated as a draft group. A "

1833 + "Mamba group cannot satisfy the widened lookup window that implies, "

1834 + "so prefix-cache reuse across requests will be disabled and any "

1835 + "external KV offload tier will store without ever serving a hit.",

1836 + spec_config.method,

1837 + mamba_groups,

1838 + )

1839 +

1840 +

1841 def _largest_divisor_at_most(value: int, limit: int) -> int:

1842 for candidate in range(min(value, limit), 0, -1):

1843 if value % candidate == 0:

@@ -1810,7 +1882,12 @@ def get_kv_cache_groups(

1882 # attention in different sizes. Need to group layers into multiple

1883 # UniformTypeKVCacheSpecs.

1884 kv_cache_groups = _get_kv_cache_groups_uniform_groups(grouped_specs)

1813 - _annotate_eagle_groups_deepseek_v4(vllm_config, kv_cache_spec, kv_cache_groups)

1885 + _annotate_eagle_groups(

1886 + vllm_config,

1887 + kv_cache_spec,

1888 + kv_cache_groups,

1889 + use_deepseek_v4_fallback=True,

1890 + )

1891 return kv_cache_groups

1892

1893 # Pull HiddenStateCacheSpec layers out before the general multi-group

@@ -1855,6 +1932,8 @@ def get_kv_cache_groups(

1932 aligned = replace(spec, block_size=new_bs, page_size_padded=common_page)

1933 groups.append(KVCacheGroupSpec([name], aligned))

1934

1935 + _annotate_eagle_groups(vllm_config, kv_cache_spec, groups)

1936 + _warn_if_unannotated_eagle_mamba(vllm_config, groups)

1937 return groups

1938

1939