[BugFix] Fix packed HND KV cache reshape for FlashAttention (#47314)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lucas Wilkinson committed
Jul 10, 2026 at 23:22 UTC
9c18e90f6c94b90ecdaa99b2230389ba40e0fc69
1 file changed
+1
-1
vllm/v1/worker/gpu/attn_utils.py
+1
-1
@@ -218,7 +218,7 @@ def _reshape_attention_kv_cache(
218
kv_cache = (
219
kv_raw_tensor.view(-1, block_stride)[:, offset : offset + page_bytes]
220
.view(dtype)
221
- .view(kv_cache_shape)
221
+ .view(permuted_kv_cache_shape)
222
)
223
elif kv_cache_spec.page_size_padded is not None:
224
# Use a strided view to skip the padding between physical pages.