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 4, 21:15 UTC
[Feat] dflash support for ROCm (#39703)
Signed-off-by: Hang Yang <hangy@amd.com>
hangy-amd committed
Apr 21, 2026 at 14:58 UTC
f95c11a848c72c8e10e8c5bb14d0a744f9ae2a82
1 file changed
+75
-29
vllm/v1/attention/backends/rocm_aiter_fa.py
+75
-29
@@ -389,6 +389,7 @@ class AiterFlashAttentionMetadata:
389
seq_lens: torch.Tensor
390
slot_mapping: torch.Tensor
391
block_table: torch.Tensor
392
+ causal: bool
393
394
# prefill and decode split
395
num_decodes: int
@@ -676,6 +677,7 @@ class AiterFlashAttentionMetadataBuilder(
677
max_seq_len=common_attn_metadata.max_seq_len,
678
seq_lens=common_attn_metadata.seq_lens,
679
block_table=common_attn_metadata.block_table_tensor,
680
+ causal=common_attn_metadata.causal,
681
slot_mapping=common_attn_metadata.slot_mapping,
682
num_decodes=num_decodes,
683
num_decode_tokens=num_decode_tokens,
@@ -724,6 +726,7 @@ class AiterFlashAttentionMetadataBuilder(
726
max_seq_len=common_attn_metadata.max_seq_len,
727
seq_lens=common_attn_metadata.seq_lens,
728
block_table=common_attn_metadata.block_table_tensor,
729
+ causal=common_attn_metadata.causal,
730
slot_mapping=common_attn_metadata.slot_mapping,
731
num_decodes=num_reqs,
732
num_decode_tokens=num_tokens,
@@ -808,6 +811,10 @@ class AiterFlashAttentionBackend(AttentionBackend):
811
# more reliable.
812
return on_mi3xx()
813
814
+ @classmethod
815
+ def supports_non_causal(cls) -> bool:
816
+ return True
817
+
818
819
class AiterFlashAttentionImpl(AttentionImpl):
820
def __init__(
@@ -1122,7 +1129,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
1129
min_seqlen_q=1,
1130
dropout_p=0.0,
1131
softmax_scale=self.scale,
1125
- causal=True,
1132
+ causal=attn_metadata.causal,
1133
window_size=self.sliding_window,
1134
alibi_slopes=self.alibi_slopes,
1135
out=output_actual_tokens[num_decode_tokens + num_extend_tokens :],
@@ -1170,39 +1177,78 @@ class AiterFlashAttentionImpl(AttentionImpl):
1177
assert attn_metadata.decode_metadata is not None
1178
decode_max_query_len = attn_metadata.decode_metadata.max_query_len
1179
1173
- # Use unified_attention for speculative decoding (multi-token)
1180
+ # Multi-token speculative decode path.
1181
if decode_max_query_len > 1:
1182
assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
1183
"Shuffle KV cache layout is not supported with "
1184
"speculative decoding (multi-token decode)."
1185
)
1179
- from aiter.ops.triton.unified_attention import (
1180
- unified_attention,
1181
- )
1182
-
1183
- descale_shape = (
1184
- num_decodes,
1185
- key_cache.shape[2],
1186
- )
1187
- unified_attention(
1188
- q=query[:num_decode_tokens],
1189
- k=key_cache,
1190
- v=value_cache,
1191
- out=output[:num_decode_tokens],
1192
- cu_seqlens_q=attn_metadata.query_start_loc[: num_decodes + 1],
1193
- max_seqlen_q=decode_max_query_len,
1194
- seqused_k=attn_metadata.seq_lens[:num_decodes],
1195
- max_seqlen_k=attn_metadata.max_seq_len,
1196
- softmax_scale=self.scale,
1197
- causal=True,
1198
- alibi_slopes=self.alibi_slopes,
1199
- window_size=self.sliding_window,
1200
- block_table=attn_metadata.block_table[:num_decodes],
1201
- softcap=self.logits_soft_cap,
1202
- q_descale=None,
1203
- k_descale=layer._k_scale.expand(descale_shape),
1204
- v_descale=layer._v_scale.expand(descale_shape),
1205
- )
1186
+ if not attn_metadata.causal:
1187
+ from aiter.ops.triton.attention.mha_v3 import (
1188
+ flash_attn_with_kvcache,
1189
+ )
1190
+
1191
+ descale_shape = (num_decodes, key_cache.shape[2])
1192
+ decode_query = query[:num_decode_tokens].reshape(
1193
+ num_decodes,
1194
+ decode_max_query_len,
1195
+ query.shape[1],
1196
+ query.shape[2],
1197
+ )
1198
+ decode_out = flash_attn_with_kvcache(
1199
+ q=decode_query,
1200
+ k_cache=key_cache,
1201
+ v_cache=value_cache,
1202
+ cache_seqlens=attn_metadata.seq_lens[:num_decodes],
1203
+ softmax_scale=self.scale,
1204
+ causal=attn_metadata.causal,
1205
+ window_size=self.sliding_window,
1206
+ softcap=self.logits_soft_cap,
1207
+ q_descale=None,
1208
+ k_descale=layer._k_scale.expand(descale_shape),
1209
+ v_descale=layer._v_scale.expand(descale_shape),
1210
+ page_table=attn_metadata.block_table[:num_decodes],
1211
+ )
1212
+ output[:num_decode_tokens].copy_(
1213
+ decode_out.reshape(
1214
+ num_decode_tokens,
1215
+ query.shape[1],
1216
+ query.shape[2],
1217
+ )
1218
+ )
1219
+ else:
1220
+ # Non-uniform query lengths can appear in real serving
1221
+ # traffic (e.g. mixed datasets). Fall back to varlen
1222
+ # unified_attention instead of asserting.
1223
+ from aiter.ops.triton.unified_attention import (
1224
+ unified_attention,
1225
+ )
1226
+
1227
+ descale_shape = (
1228
+ num_decodes,
1229
+ key_cache.shape[2],
1230
+ )
1231
+ unified_attention(
1232
+ q=query[:num_decode_tokens],
1233
+ k=key_cache,
1234
+ v=value_cache,
1235
+ out=output[:num_decode_tokens],
1236
+ cu_seqlens_q=attn_metadata.query_start_loc[
1237
+ : num_decodes + 1
1238
+ ],
1239
+ max_seqlen_q=decode_max_query_len,
1240
+ seqused_k=attn_metadata.seq_lens[:num_decodes],
1241
+ max_seqlen_k=attn_metadata.max_seq_len,
1242
+ softmax_scale=self.scale,
1243
+ causal=True,
1244
+ alibi_slopes=self.alibi_slopes,
1245
+ window_size=self.sliding_window,
1246
+ block_table=attn_metadata.block_table[:num_decodes],
1247
+ softcap=self.logits_soft_cap,
1248
+ q_descale=None,
1249
+ k_descale=layer._k_scale.expand(descale_shape),
1250
+ v_descale=layer._v_scale.expand(descale_shape),
1251
+ )
1252
return
1253
1254
# The ll4mi kernel in paged_attention_v1 requires