From e7b6addc22317e1deeeef4741628dac2901dc904 Mon Sep 17 00:00:00 2001 From: Tuukka Sarvi Date: Tue, 7 Apr 2026 08:07:10 -0500 Subject: [PATCH 1/2] [ROCm] Fix shuffled KV-cache writes for hybrid attention layouts Hybrid attention+mamba models can expose non-contiguous KV cache views via as_strided_. The shuffled cache writer assumed a contiguous layout and could corrupt shuffle=1 decode inputs, so compute shuffled write offsets from the actual cache view strides. This fixes for example Qwen3.5 for shuffle=1 for low batch sizes. Signed-off-by: tuukkjs --- vllm/v1/attention/backends/rocm_aiter_fa.py | 43 +++++++++++++++------ 1 file changed, 32 insertions(+), 11 deletions(-) diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index a897b33fcab3..59f35e327701 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -223,9 +223,18 @@ def reshape_and_cache_shuffle_kernel( x, k_stride0, v_stride0, + kc_stride0, + kc_stride1, + kc_stride2, + kc_stride3, + kc_stride4, + vc_stride0, + vc_stride1, + vc_stride2, + vc_stride3, + vc_stride4, block_size, head_size, - num_kv_heads, BLOCK_SIZE: tl.constexpr, QUANT: tl.constexpr, ): @@ -239,18 +248,19 @@ def reshape_and_cache_shuffle_kernel( return block_id = slot_id // block_size block_offset = slot_id % block_size - dst_offset = ( - block_id * num_kv_heads * head_size * block_size - + head_id * head_size * block_size - ) dst_k_shuffle_offset = ( - dst_offset + offset // x * block_size * x + block_offset * x + offset % x + block_id * kc_stride0 + + head_id * kc_stride1 + + (offset // x) * kc_stride2 + + block_offset * kc_stride3 + + (offset % x) * kc_stride4 ) dst_v_shuffle_offset = ( - dst_offset - + block_offset // x * head_size * x - + offset * x - + block_offset % x + block_id * vc_stride0 + + head_id * vc_stride1 + + (block_offset // x) * vc_stride2 + + offset * vc_stride3 + + (block_offset % x) * vc_stride4 ) k_val = tl.load(key_ptr + src_offset_k + offset) v_val = tl.load(value_ptr + src_offset_v + offset) @@ -290,6 +300,8 @@ def reshape_and_cache_shuffle_triton( ) new_key_cache = key_cache.view_as(k_cache_template) new_value_cache = value_cache.view_as(v_cache_template) + k_cache_strides = new_key_cache.stride() + v_cache_strides = new_value_cache.stride() QUANT = False if is_quantized_kv_cache(kv_cache_dtype): QUANT = True @@ -308,9 +320,18 @@ def reshape_and_cache_shuffle_triton( x, key.stride(0), value.stride(0), + k_cache_strides[0], + k_cache_strides[1], + k_cache_strides[2], + k_cache_strides[3], + k_cache_strides[4], + v_cache_strides[0], + v_cache_strides[1], + v_cache_strides[2], + v_cache_strides[3], + v_cache_strides[4], block_size, head_size, - num_kv_heads, BLOCK_SIZE=head_size, QUANT=QUANT, ) From 98e6f7dc20eaa538e470f13ffbc6aaaabfff9c7d Mon Sep 17 00:00:00 2001 From: Tuukka Sarvi Date: Wed, 8 Apr 2026 07:45:53 -0500 Subject: [PATCH 2/2] Update comments Signed-off-by: tuukkjs Signed-off-by: Tuukka Sarvi --- vllm/v1/attention/backends/rocm_aiter_fa.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 59f35e327701..eb4f826fe667 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -298,6 +298,8 @@ def reshape_and_cache_shuffle_triton( dtype=value_cache.dtype, device="meta", ) + # Hybrid attention+mamba can expose non-contiguous KV cache views, + # so the shuffle writer must use the actual view strides. new_key_cache = key_cache.view_as(k_cache_template) new_value_cache = value_cache.view_as(v_cache_template) k_cache_strides = new_key_cache.stride() @@ -1409,8 +1411,6 @@ def do_kv_cache_update( assert k_scale is not None and v_scale is not None, ( "k_scale and v_scale are required for shuffled update" ) - # TODO: Add correct KV cache handling for hybrid model. KV cache - # may not be contiguous if mamba state exists. reshape_and_cache_shuffle_triton( key, value,