A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 23:30 UTC
[ROCM][Bugfix] Use correct stride in cp_mha_gather_cache_kernel for hybrid model (#37228) (#37228)
Signed-off-by: jennyyyyzhen <yzhen@hmc.edu> Co-authored-by: yZhen <yZhen@fb.com>
jennyyyyzhen committed
Mar 26, 2026 at 10:33 UTC
a4cf9b22ba8f08df8d7977c2b0753fe20231cffc
1 file changed
+25
-6
vllm/v1/attention/backends/rocm_aiter_fa.py
+25
-6
@@ -59,6 +59,12 @@ if current_platform.is_rocm():
59
head_size,
60
x,
61
max_block_num,
62
+ k_cache_stride0,
63
+ k_cache_stride1,
64
+ k_cache_stride2,
65
+ v_cache_stride0,
66
+ v_cache_stride1,
67
+ v_cache_stride2,
68
DEQUANT: tl.constexpr,
69
PAGE_SIZE: tl.constexpr,
70
CACHE_FORMAT: tl.constexpr,
@@ -90,15 +96,15 @@ if current_platform.is_rocm():
96
# V: [num_blocks, page_size, num_head, head_dim]
97
key_cache_ptr_offset = (
98
key_cache_ptr
93
- + block_id * num_heads * head_size * PAGE_SIZE
94
- + slot_id * num_heads * head_size
95
- + head_id * head_size
99
+ + block_id * k_cache_stride0
100
+ + slot_id * k_cache_stride1
101
+ + head_id * k_cache_stride2
102
)
103
value_cache_ptr_offset = (
104
value_cache_ptr
99
- + block_id * num_heads * head_size * PAGE_SIZE
100
- + slot_id * num_heads * head_size
101
- + head_id * head_size
105
+ + block_id * v_cache_stride0
106
+ + slot_id * v_cache_stride1
107
+ + head_id * v_cache_stride2
108
)
109
k_reg = tl.load(key_cache_ptr_offset + col_offsets)
110
v_reg = tl.load(value_cache_ptr_offset + col_offsets)
@@ -171,6 +177,11 @@ if current_platform.is_rocm():
177
page_size = key_cache.shape[1]
178
num_heads = key_cache.shape[2]
179
180
+ # Pass actual tensor strides so the kernel works correctly
181
+ # even when the cache is non-contiguous (e.g, for hybrid model)
182
+ k_strides = key_cache.stride()
183
+ v_strides = value_cache.stride()
184
+
185
grid = lambda meta: (total_tokens, num_heads)
186
cp_mha_gather_cache_kernel[grid](
187
key_cache,
@@ -187,6 +198,12 @@ if current_platform.is_rocm():
198
head_dim,
199
x,
200
block_tables.size(1),
201
+ k_strides[0],
202
+ k_strides[1],
203
+ k_strides[2],
204
+ v_strides[0],
205
+ v_strides[1],
206
+ v_strides[2],
207
DEQUANT=dequant,
208
PAGE_SIZE=page_size,
209
CACHE_FORMAT=kv_cache_layout,
@@ -1337,6 +1354,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1354
assert k_scale is not None and v_scale is not None, (
1355
"k_scale and v_scale are required for shuffled update"
1356
)
1357
+ # TODO: Add correct KV cache handling for hybrid model. KV cache
1358
+ # may not be contiguous if mamba state exists.
1359
reshape_and_cache_shuffle_triton(
1360
key,
1361
value,