karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[Bugfix][V1] Split V2 model-runner attention groups on num_heads_q (#45564)

Browse Source 
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>
This commit is contained in:
Roger Wang
2026-06-14 21:49:46 -07:00
committed by GitHub
co-authored by Claude Opus 4.8 Nick Hill
parent e8d3e22c88
commit b8336c3c7c
1 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/v1/worker/gpu/attn_utils.py
+7 -3
View File
@@ -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