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 7, 22:45 UTC
[ROCm][AITER] Fix aiter paged_attention_v1 decode for sliding window and head_size < 64 (#34570)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Feb 20, 2026 at 22:25 UTC
cf93c1a12849693d6bcee3e2c917c02e1fc9a47f
1 file changed
+44
-1
vllm/v1/attention/backends/rocm_aiter_fa.py
+44
-1
@@ -1114,7 +1114,50 @@ class AiterFlashAttentionImpl(AttentionImpl):
1114
)
1115
return
1116
1117
- if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1117
+ # The ll4mi kernel in paged_attention_v1 requires
1118
+ # HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4).
1119
+ # For smaller head sizes or sliding window attention,
1120
+ # fall back to the unified_attention triton kernel which
1121
+ # handles both correctly.
1122
+ _MIN_HEAD_SIZE_FOR_LL4MI = 64
1123
+ use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI
1124
+
1125
+ if use_unified_attention:
1126
+ assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
1127
+ "unified_attention fallback with shuffle layout "
1128
+ "is not supported yet."
1129
+ )
1130
+ from aiter.ops.triton.unified_attention import (
1131
+ unified_attention,
1132
+ )
1133
+
1134
+ decode_cu_seqlens_q = attn_metadata.query_start_loc[
1135
+ : num_decodes + 1
1136
+ ]
1137
+ descale_shape = (
1138
+ num_decodes,
1139
+ key_cache.shape[2],
1140
+ )
1141
+ unified_attention(
1142
+ q=query[:num_decode_tokens],
1143
+ k=key_cache,
1144
+ v=value_cache,
1145
+ out=output[:num_decode_tokens],
1146
+ cu_seqlens_q=decode_cu_seqlens_q,
1147
+ max_seqlen_q=1,
1148
+ seqused_k=attn_metadata.seq_lens[:num_decodes],
1149
+ max_seqlen_k=attn_metadata.max_seq_len,
1150
+ softmax_scale=self.scale,
1151
+ causal=True,
1152
+ alibi_slopes=self.alibi_slopes,
1153
+ window_size=self.sliding_window,
1154
+ block_table=attn_metadata.block_table[:num_decodes],
1155
+ softcap=self.logits_soft_cap,
1156
+ q_descale=None,
1157
+ k_descale=layer._k_scale.expand(descale_shape),
1158
+ v_descale=layer._v_scale.expand(descale_shape),
1159
+ )
1160
+ elif rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1161
num_blocks, block_size, num_kv_heads, head_size = key_cache.shape
1162
x = 16 // key_cache.element_size()
1163
k_cache_template = torch.empty(