[ROCm] Fix shuffled KV-cache writes for hybrid attention layouts

vllm/v1/attention/backends/rocm_aiter_fa.py

@@ -225,11 +225,18 @@ def reshape_and_cache_shuffle_kernel(
         x,
         k_stride0,
         v_stride0,
-        k_cache_block_stride,
-        v_cache_block_stride,
+        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,
     ):
@@ -243,20 +250,19 @@ def reshape_and_cache_shuffle_kernel(
             return
         block_id = slot_id // block_size
         block_offset = slot_id % block_size
-        k_dst_offset = (
-            block_id * k_cache_block_stride + head_id * head_size * block_size
-        )
         dst_k_shuffle_offset = (
-            k_dst_offset + offset // x * block_size * x + block_offset * x + offset % x
-        )
-        v_dst_offset = (
-            block_id * v_cache_block_stride + head_id * head_size * block_size
+            block_id * kc_stride0
+            + head_id * kc_stride1
+            + (offset // x) * kc_stride2
+            + block_offset * kc_stride3
+            + (offset % x) * kc_stride4
         )
         dst_v_shuffle_offset = (
-            v_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)
@@ -284,6 +290,22 @@ def reshape_and_cache_shuffle_triton(
         _, num_kv_heads, head_size = key.shape
         num_blocks, block_size, _, _ = key_cache.shape
         x = 16 // key_cache.element_size()
+        k_cache_template = torch.empty(
+            [num_blocks, num_kv_heads, head_size // x, block_size, x],
+            dtype=key_cache.dtype,
+            device="meta",
+        )
+        v_cache_template = torch.empty(
+            [num_blocks, num_kv_heads, block_size // x, head_size, x],
+            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()
+        v_cache_strides = new_value_cache.stride()
         QUANT = False
         if is_quantized_kv_cache(kv_cache_dtype):
             QUANT = True
@@ -294,19 +316,26 @@ def reshape_and_cache_shuffle_triton(
         reshape_and_cache_shuffle_kernel[grid](
             key,
             value,
-            key_cache,
-            value_cache,
+            new_key_cache,
+            new_value_cache,
             slot_mapping,
             k_scales,
             v_scales,
             x,
             key.stride(0),
             value.stride(0),
-            key_cache.stride(0),
-            value_cache.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,
         )
@@ -1408,8 +1437,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,