From e2eaefdec9f1fc39da30a8390679055aec171f1a Mon Sep 17 00:00:00 2001 From: MOVZX Date: Wed, 23 Sep 2026 07:33:48 -0400 Subject: [PATCH] [KV] Restore PP fix #54793 -- no KV tensors for another stage's layers A KV cache group with no layer on a pipeline stage was still given tensors there, and allocate_kv_cache raised StopIteration on the second stage. The fix was lost in the v0.30.0 rebase. Test drives the real get_kv_cache_configs with the production layout. Co-Authored-By: Hinata Signed-off-by: MOVZX --- vllm/v1/core/kv_cache_utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 46934d530b..7da377146d 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1806,7 +1806,14 @@ def get_kv_cache_config_from_groups( layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) if isinstance(group_spec, UniformTypeKVCacheSpecs): for layer_name, spec in group_spec.kv_cache_specs.items(): - layers_by_spec[spec].append(layer_name) + # Fork (E2 PP fix #54793, dropped in the v0.30.0 rebase, restored + # the vLLM 0.30 rebase): _project_kv_cache_groups_to_worker keeps the GLOBAL + # UniformTypeKVCacheSpecs for a group with no layer on this PP + # stage, so without this filter the worker gets a tensor for + # another stage's layer and allocate_kv_cache's group lookup + # raises StopIteration (Flash-Next PP1: the layer-2 PLE state). + if layer_name in group.layer_names: + layers_by_spec[spec].append(layer_name) elif group.layer_names: layers_by_spec[group_spec].extend(group.layer_names)