From b8336c3c7c298e0878f22a7bf70f4e295b2f4e01 Mon Sep 17 00:00:00 2001 From: Roger Wang Date: Sun, 14 Jun 2026 21:49:46 -0700 Subject: [PATCH] [Bugfix][V1] Split V2 model-runner attention groups on num_heads_q (#45564) Signed-off-by: Roger Wang Signed-off-by: Nick Hill Co-authored-by: Claude Opus 4.8 Co-authored-by: Nick Hill --- vllm/v1/worker/gpu/attn_utils.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 35c40a1c229..74158f92bf8 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -85,8 +85,8 @@ def init_attn_backend( layer_type = cast(type[Any], AttentionLayerBase) attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names) - group_map: dict[tuple[tuple[str, str], KVCacheSpec], AttentionGroup] = {} - group_order: list[tuple[tuple[str, str], KVCacheSpec]] = [] + group_map: dict[tuple[tuple[str, str], KVCacheSpec, int], AttentionGroup] = {} + group_order: list[tuple[tuple[str, str], KVCacheSpec, int]] = [] for layer_name in layer_names: attn_backend = attn_layers[layer_name].get_attn_backend() @@ -95,7 +95,11 @@ def init_attn_backend( if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs): layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name] - key = (attn_backend.full_cls_name(), layer_kv_cache_spec) + # Split on per-rank num_heads_q so layers with different Q-head + # counts (e.g. a spec-decode draft head and its target) get separate + # metadata builders. + num_heads_q = getattr(attn_layers[layer_name], "num_heads", 0) + key = (attn_backend.full_cls_name(), layer_kv_cache_spec, num_heads_q) if key not in group_map: group_map[key] = AttentionGroup( attn_backend, [layer_name], layer_kv_cache_spec, kv_cache_group_id