A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][V1] Split V2 model-runner attention groups on num_heads_q (#45564)
Signed-off-by: Roger Wang <hey@rogerw.io> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Roger Wang committed
Jun 14, 2026 at 21:49 UTC
b8336c3c7c298e0878f22a7bf70f4e295b2f4e01
1 file changed
+7
-3
vllm/v1/worker/gpu/attn_utils.py
+7
-3
@@ -85,8 +85,8 @@ def init_attn_backend(
85
layer_type = cast(type[Any], AttentionLayerBase)
86
attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
87
88
- group_map: dict[tuple[tuple[str, str], KVCacheSpec], AttentionGroup] = {}
89
- group_order: list[tuple[tuple[str, str], KVCacheSpec]] = []
88
+ group_map: dict[tuple[tuple[str, str], KVCacheSpec, int], AttentionGroup] = {}
89
+ group_order: list[tuple[tuple[str, str], KVCacheSpec, int]] = []
90
91
for layer_name in layer_names:
92
attn_backend = attn_layers[layer_name].get_attn_backend()
@@ -95,7 +95,11 @@ def init_attn_backend(
95
if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
96
layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name]
97
98
- key = (attn_backend.full_cls_name(), layer_kv_cache_spec)
98
+ # Split on per-rank num_heads_q so layers with different Q-head
99
+ # counts (e.g. a spec-decode draft head and its target) get separate
100
+ # metadata builders.
101
+ num_heads_q = getattr(attn_layers[layer_name], "num_heads", 0)
102
+ key = (attn_backend.full_cls_name(), layer_kv_cache_spec, num_heads_q)
103
if key not in group_map:
104
group_map[key] = AttentionGroup(
105
attn_backend, [layer_name], layer_kv_cache_spec, kv_cache_group_id