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, 01:15 UTC
[Bugfix][Hardware][AMD] Fix ROCM_AITER_FA speculative decoding support (#32877)
Signed-off-by: c0de128 <kevin.mckay@outlook.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
Kevin McKay committed
Feb 20, 2026 at 00:25 UTC
8de7c636cc02a8306441af868b9c1d0e6d64799f
1 file changed
+35
-2
vllm/v1/attention/backends/rocm_aiter_fa.py
+35
-2
@@ -1076,10 +1076,43 @@ class AiterFlashAttentionImpl(AttentionImpl):
1076
# calculate for decodes
1077
if num_decodes > 0:
1078
assert attn_metadata.decode_metadata is not None
1079
- if self.sliding_window[0] != -1:
1079
+ decode_max_query_len = attn_metadata.decode_metadata.max_query_len
1080
+
1081
+ # Use unified_attention for speculative decoding (multi-token)
1082
+ # or when sliding window is enabled
1083
+ if self.sliding_window[0] != -1 or decode_max_query_len > 1:
1084
assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
1081
- "Sliding window with shuffle layout is not supported yet."
1085
+ "Shuffle KV cache layout is not supported with sliding "
1086
+ "window or speculative decoding (multi-token decode)."
1087
+ )
1088
+ from aiter.ops.triton.unified_attention import (
1089
+ unified_attention,
1090
+ )
1091
+
1092
+ descale_shape = (
1093
+ attn_metadata.query_start_loc[:num_decodes].shape[0] - 1,
1094
+ key_cache.shape[2],
1095
+ )
1096
+ unified_attention(
1097
+ q=query[:num_decode_tokens],
1098
+ k=key_cache,
1099
+ v=value_cache,
1100
+ out=output[:num_decode_tokens],
1101
+ cu_seqlens_q=attn_metadata.query_start_loc[:num_decodes],
1102
+ max_seqlen_q=decode_max_query_len,
1103
+ seqused_k=attn_metadata.seq_lens[:num_decodes],
1104
+ max_seqlen_k=attn_metadata.max_seq_len,
1105
+ softmax_scale=self.scale,
1106
+ causal=True,
1107
+ alibi_slopes=self.alibi_slopes,
1108
+ window_size=self.sliding_window,
1109
+ block_table=attn_metadata.block_table[:num_decodes],
1110
+ softcap=self.logits_soft_cap,
1111
+ q_descale=None,
1112
+ k_descale=layer._k_scale.expand(descale_shape),
1113
+ v_descale=layer._v_scale.expand(descale_shape),
1114
)
1115
+ return
1116
1117
if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1118
num_blocks, block_size, num_kv_heads, head_size = key_cache.shape