A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 19:30 UTC
[Bugfix] Split attention groups by num_heads_q for spec-decode drafts (#43543)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Luciano Martins committed
May 26, 2026 at 21:11 UTC
dede691c9536e8a79275d2bd0defbccb93f6ef7f
2 files changed
+26
-5
.buildkite/test_areas/spec_decode.yaml
+1
@@ -32,6 +32,7 @@ steps:
32
source_file_dependencies:
33
- vllm/v1/spec_decode/
34
- vllm/v1/worker/gpu/spec_decode/
35
+ - vllm/v1/attention/backends/
36
- vllm/transformers_utils/configs/speculators/
37
- tests/v1/e2e/spec_decode/
38
commands:
vllm/v1/worker/gpu_model_runner.py
+25
-5
@@ -6556,8 +6556,21 @@ class GPUModelRunner(
6556
assert len(self.attn_groups) == 0, "Attention backends are already initialized"
6557
6558
class AttentionGroupKey(NamedTuple):
6559
+ """Deduplication key for attention groups within a KV cache group.
6560
+
6561
+ Splits on per-rank ``num_heads_q`` in addition to backend + spec
6562
+ so layers with different Q-head counts (e.g. a spec-decode draft
6563
+ with fewer attention heads than its target) get separate metadata
6564
+ builders. The builders' scratch (e.g. ``softmax_segm_*`` in
6565
+ ``triton_attn``, ``num_qo_heads`` in FlashInfer) is sized by
6566
+ ``num_heads_q`` and assumes uniformity within the group; see
6567
+ ``get_num_attention_heads_from_layers`` in
6568
+ ``vllm/v1/attention/backends/utils.py``.
6569
+ """
6570
+
6571
attn_backend: type[AttentionBackend]
6572
kv_cache_spec: KVCacheSpec
6573
+ num_heads_q: int
6574
6575
def get_attn_backends_for_group(
6576
kv_cache_group_spec: KVCacheGroupSpec,
@@ -6586,9 +6599,16 @@ class GPUModelRunner(
6599
layer_kv_cache_spec = kv_cache_group_spec.kv_cache_spec
6600
if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
6601
layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name]
6589
- key = (full_cls_name, layer_kv_cache_spec)
6602
+ # Non-Attention layer types (e.g. Mamba1, ShortConv) do not
6603
+ # expose ``num_heads``; fall back to 0 so they cluster as
6604
+ # before. Such layers never coexist with Attention in a
6605
+ # single KV cache group (different KVCacheSpec), so the
6606
+ # fallback can never spuriously merge them with attention
6607
+ # layers.
6608
+ num_heads_q = getattr(layers[layer_name], "num_heads", 0)
6609
+ key = (full_cls_name, layer_kv_cache_spec, num_heads_q)
6610
attn_backends[key] = AttentionGroupKey(
6591
- attn_backend, layer_kv_cache_spec
6611
+ attn_backend, layer_kv_cache_spec, num_heads_q
6612
)
6613
attn_backend_layers[key].append(layer_name)
6614
return (
@@ -6601,11 +6621,11 @@ class GPUModelRunner(
6621
kv_cache_group_id: int,
6622
) -> list[AttentionGroup]:
6623
attn_groups: list[AttentionGroup] = []
6604
- for (attn_backend, kv_cache_spec), layer_names in attn_backends_map.items():
6624
+ for key, layer_names in attn_backends_map.items():
6625
attn_group = AttentionGroup(
6606
- attn_backend,
6626
+ key.attn_backend,
6627
layer_names,
6608
- kv_cache_spec,
6628
+ key.kv_cache_spec,
6629
kv_cache_group_id,
6630
)
6631