@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0bce411a07

[Bugfix] Annotate MTP draft KV cache groups positionally on the hybrid grouping path (#55390)

Signed-off-by: Navjot Singh <navjot.singh@shopify.com> Co-authored-by: Codex <noreply@openai.com>

Navjot Singh committed Sep 21, 2026 at 20:31 UTC 0bce411a073e0da2c9bded9bb6ced1cbee51c564

2 files changed +133 -21

tests/v1/core/test_kv_cache_utils.py

+90 -6

@@ -4275,12 +4275,16 @@ def _deepseek_v4_specs(model_version="deepseek_v4"):

4275 }

4276

4277

4278 -def test_deepseek_v4_draft_group_annotated_on_packed_path():

4278 +@pytest.mark.parametrize(

4279 + ("method", "model_type"),

4280 + [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")],

4281 +)

4282 +def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type):

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

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

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

4286 groups = get_kv_cache_groups(

4283 - _spec_decode_grouping_config(method="mtp", model_type="deepseek_v4"),

4287 + _spec_decode_grouping_config(method=method, model_type=model_type),

4288 _deepseek_v4_specs(model_version=None),

4289 )

4290

@@ -4289,13 +4293,93 @@ def test_deepseek_v4_draft_group_annotated_on_packed_path():

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

4294

4295

4292 -def test_deepseek_v4_annotation_requires_model_type():

4293 - # The positional rule is only sound for DeepseekV4, where the draft layer

4294 - # is known to be registered last. Without that model gate nothing may be

4295 - # flagged, however the grouping happens to fall out.

4296 +def test_trailing_layer_fallback_applies_to_any_mtp_model():

4297 + # The positional rule is sound for every MTP drafter, not just DeepseekV4:

4298 + # MTP blocks reuse the target's decoder layer (no spec marker) and always

4299 + # register after every target layer. The model_type must not gate it.

4300 groups = get_kv_cache_groups(

4301 _spec_decode_grouping_config(method="mtp", model_type="other"),

4302 _deepseek_v4_specs(),

4303 )

4304

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

4306 + assert len(flagged) == 1

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

4308 +

4309 +

4310 +def _qwen3_5_hybrid_specs(with_mtp_layer: bool):

4311 + """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with

4312 + the MTP drafter's full-attn layer (spec-identical to the target's)

4313 + registered last."""

4314 + specs = {}

4315 + idx = 0

4316 + for _ in range(2):

4317 + for _ in range(3):

4318 + specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec(

4319 + mamba_cache_mode="align"

4320 + )

4321 + idx += 1

4322 + specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec()

4323 + idx += 1

4324 + if with_mtp_layer:

4325 + specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec()

4326 + return specs

4327 +

4328 +

4329 +def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm):

4330 + # A hybrid mamba + full-attention model with an MTP drafter that is

4331 + # spec-indistinguishable from the target reaches the general multi-group

4332 + # path. The trailing-layer rule must locate the draft group there so the

4333 + # Mamba groups are not swept up by the flag-all consumer fallback.

4334 + groups = get_kv_cache_groups(

4335 + _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"),

4336 + _qwen3_5_hybrid_specs(with_mtp_layer=True),

4337 + )

4338 +

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

4340 + assert len(flagged) == 1

4341 + assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names

4342 + for group in groups:

4343 + if any(

4344 + isinstance(spec, MambaSpec)

4345 + for spec in iter_layer_specs(group.kv_cache_spec)

4346 + ):

4347 + assert not group.is_eagle_group

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

4349 +

4350 +

4351 +@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"])

4352 +def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method):

4353 + # Other drafters are not covered by the trailing-layer rule, so an

4354 + # unidentifiable hybrid draft must still warn.

4355 + groups = get_kv_cache_groups(

4356 + _spec_decode_grouping_config(method=method, model_type="qwen3_5"),

4357 + _qwen3_5_hybrid_specs(with_mtp_layer=True),

4358 + )

4359 +

4360 assert not any(g.is_eagle_group for g in groups)

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

4362 +

4363 +

4364 +def test_trailing_layer_fallback_requires_exact_partition():

4365 + # If the groups do not partition the layers exactly (e.g. a caller that

4366 + # dropped or duplicated layers), the positional rule is meaningless and

4367 + # must not fire.

4368 + from vllm.v1.core.kv_cache_utils import _annotate_eagle_groups

4369 +

4370 + specs = _qwen3_5_hybrid_specs(with_mtp_layer=True)

4371 + config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5")

4372 + groups = get_kv_cache_groups(config, specs)

4373 + for g in groups:

4374 + g.is_eagle_group = False

4375 + # Remove one layer from its group: no longer an exact partition.

4376 + trimmed = [

4377 + KVCacheGroupSpec(

4378 + [n for n in g.layer_names if n != "model.layers.0.linear_attn"],

4379 + g.kv_cache_spec,

4380 + )

4381 + for g in groups

4382 + ]

4383 + _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True)

4384 +

4385 + assert not any(g.is_eagle_group for g in trimmed)

vllm/v1/core/kv_cache_utils.py

+43 -15

@@ -2117,16 +2117,24 @@ def _get_packed_kv_cache_groups(

2117 vllm_config,

2118 kv_cache_spec,

2119 groups,

2120 - use_deepseek_v4_fallback=_is_deepseek_v4_eagle(vllm_config),

2120 + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config),

2121 )

2122 _warn_if_unannotated_eagle_mamba(vllm_config, groups)

2123 return groups

2124

2125

2126 -def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool:

2126 +def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool:

2127 + """Whether the drafter's KV layers can be located positionally.

2128 +

2129 + MTP drafters register their KV layers after the target layers, without a

2130 + distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too.

2131 + The annotator separately checks that the groups partition the layers exactly.

2132 + """

2133 spec_config = vllm_config.speculative_config

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

2135 return False

2136 + if spec_config.method == "mtp":

2137 + return True

2138 model_config = vllm_config.model_config

2139 return model_config is not None and model_config.hf_config.model_type in (

2140 "deepseek_v4",

@@ -2134,11 +2142,24 @@ def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool:

2142 )

2143

2144

2145 +def _groups_partition_layers_exactly(

2146 + kv_cache_spec: dict[str, KVCacheSpec],

2147 + kv_cache_groups: list[KVCacheGroupSpec],

2148 +) -> bool:

2149 + """Whether the groups cover every layer of ``kv_cache_spec`` exactly once.

2150 +

2151 + The trailing-layer draft fallback is only meaningful when the last

2152 + registered layer is guaranteed to appear in exactly one group.

2153 + """

2154 + names = [name for group in kv_cache_groups for name in group.layer_names]

2155 + return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec)

2156 +

2157 +

2158 def _annotate_eagle_groups(

2159 vllm_config: VllmConfig,

2160 kv_cache_spec: dict[str, KVCacheSpec],

2161 kv_cache_groups: list[KVCacheGroupSpec],

2141 - use_deepseek_v4_fallback: bool = False,

2162 + use_trailing_layer_fallback: bool = False,

2163 ) -> None:

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

2165

@@ -2151,14 +2172,14 @@ def _annotate_eagle_groups(

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

2173 not necessary: a drafter whose spec is indistinguishable from the

2174 target's cannot be found this way.

2154 - 2. Model-scoped positional fallback for DeepseekV4/V4.1, whose MTP block

2155 - reuses the target's own decoder layer and so carries no spec marker. Its

2156 - draft attention layer is always the last registered layer, so flag whichever

2157 - group holds it. This rule is only valid where the groups partition

2158 - exactly the layers of ``kv_cache_spec``, which is true on the packed

2159 - grouping path and not in general; other callers must leave

2160 - ``use_deepseek_v4_fallback`` False. The caller gates this fallback on

2161 - the configured model type.

2175 + 2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose

2176 + MTP block reuses the target's own decoder layer and so carries no spec

2177 + marker. Their draft attention layers always register after every

2178 + target layer, so flag whichever group holds the last registered layer.

2179 + This rule is only valid where the groups partition exactly the layers

2180 + of ``kv_cache_spec``, which is re-checked here before applying it.

2181 + When a drafter's trailing caches span several groups, this rule flags

2182 + only the group holding the very last layer and must be generalized.

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

2184

2185 Args:

@@ -2166,8 +2187,8 @@ def _annotate_eagle_groups(

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

2188 registration order. Only read by rule 2.

2189 kv_cache_groups: Groups to annotate in place.

2169 - use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4/V4.1 packed

2170 - group.

2190 + use_trailing_layer_fallback: Enable rule 2. Callers gate this on

2191 + ``_uses_trailing_mtp_layers``.

2192

2193 """

2194 spec_config = vllm_config.speculative_config

@@ -2181,7 +2202,9 @@ def _annotate_eagle_groups(

2202 ):

2203 group.is_eagle_group = True

2204

2184 - if not use_deepseek_v4_fallback:

2205 + if not use_trailing_layer_fallback:

2206 + return

2207 + if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups):

2208 return

2209 last_layer = next(reversed(kv_cache_spec))

2210 for group in kv_cache_groups:

@@ -2352,7 +2375,12 @@ def get_kv_cache_groups(

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

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

2377

2355 - _annotate_eagle_groups(vllm_config, kv_cache_spec, groups)

2378 + _annotate_eagle_groups(

2379 + vllm_config,

2380 + kv_cache_spec,

2381 + groups,

2382 + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config),

2383 + )

2384 _warn_if_unannotated_eagle_mamba(vllm_config, groups)

2385 return groups

2386