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 8, 08:45 UTC
Use paged_attention_v1 for sliding window decode in rocm_aiter_fa (#34378)
Signed-off-by: Martin Yuan <myuan@meta.com> Co-authored-by: Martin Yuan <myuan@meta.com>
Mengtao (Martin) Yuan committed
Feb 12, 2026 at 16:14 UTC
9ea1f598ce48da3054d073e74b9e51e8d0de945a
1 file changed
+2
-29
vllm/v1/attention/backends/rocm_aiter_fa.py
+2
-29
@@ -1075,35 +1075,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
1075
assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
1076
"Sliding window with shuffle layout is not supported yet."
1077
)
1078
- from aiter.ops.triton.unified_attention import (
1079
- unified_attention,
1080
- )
1081
-
1082
- descale_shape = (
1083
- attn_metadata.query_start_loc[:num_decodes].shape[0] - 1,
1084
- key_cache.shape[2],
1085
- )
1086
- unified_attention(
1087
- q=query[:num_decode_tokens],
1088
- k=key_cache,
1089
- v=value_cache,
1090
- out=output[:num_decode_tokens],
1091
- cu_seqlens_q=attn_metadata.query_start_loc[:num_decodes],
1092
- max_seqlen_q=1, # optimize this
1093
- seqused_k=attn_metadata.seq_lens[:num_decodes],
1094
- max_seqlen_k=attn_metadata.max_seq_len,
1095
- softmax_scale=self.scale,
1096
- causal=True,
1097
- alibi_slopes=self.alibi_slopes,
1098
- window_size=self.sliding_window,
1099
- block_table=attn_metadata.block_table[:num_decodes],
1100
- softcap=self.logits_soft_cap,
1101
- q_descale=None,
1102
- k_descale=layer._k_scale.expand(descale_shape),
1103
- v_descale=layer._v_scale.expand(descale_shape),
1104
- )
1105
- return
1106
- assert attn_metadata.decode_metadata is not None
1078
1079
if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1080
num_blocks, block_size, num_kv_heads, head_size = key_cache.shape
@@ -1172,6 +1143,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1143
layer._v_scale,
1144
None,
1145
_PARTITION_SIZE_ROCM,
1146
+ 1,
1147
+ self.sliding_window[0] + 1,
1148
)
1149
else:
1150
raise NotImplementedError(