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 2, 08:30 UTC
[ROCm] Add attention sink support to AITer flash attention backend (#43817)
Signed-off-by: Xiaoran Chen <xiaoran@fb.com> Co-authored-by: Xiaoran Chen <xiaoran@fb.com>
Xiaoran committed
May 30, 2026 at 03:13 UTC
3becc5db4034a65802c7d7b867fd236655c0ebcc
3 files changed
+25
-5
docs/design/attention_backends.md
+1
-1
@@ -177,7 +177,7 @@ Priority is **1 = highest** (tried first).
177
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
178
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
179
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
180
-| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
180
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
181
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
182
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
183
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
vllm/_aiter_ops.py
+2
@@ -2394,6 +2394,7 @@ class rocm_aiter_ops:
2394
alibi_slopes: torch.Tensor | None = None,
2395
return_lse: bool = False,
2396
out: torch.Tensor | None = None,
2397
+ sink_ptr: torch.Tensor | None = None,
2398
):
2399
"""
2400
Flash attention with variable length sequences.
@@ -2422,6 +2423,7 @@ class rocm_aiter_ops:
2423
alibi_slopes=alibi_slopes,
2424
return_lse=return_lse,
2425
out=out,
2426
+ sink_ptr=sink_ptr,
2427
)
2428
2429
@staticmethod
vllm/v1/attention/backends/rocm_aiter_fa.py
+22
-4
@@ -709,6 +709,11 @@ class AiterFlashAttentionMetadataBuilder(
709
710
class AiterFlashAttentionBackend(AttentionBackend):
711
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
712
+
713
+ @classmethod
714
+ def supports_sink(cls) -> bool:
715
+ return True
716
+
717
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
718
"auto",
719
"float16",
@@ -788,6 +793,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
793
logits_soft_cap: float | None = None,
794
attn_type: AttentionType = AttentionType.DECODER,
795
kv_sharing_target_layer_name: int | None = None,
796
+ sinks: torch.Tensor | None = None,
797
) -> None:
798
self.num_heads = num_heads
799
self.head_size = head_size
@@ -806,6 +812,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
812
logits_soft_cap = 0.0
813
self.logits_soft_cap = logits_soft_cap
814
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
815
+ self.sinks = sinks
816
817
assert self.num_heads % self.num_kv_heads == 0
818
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
@@ -878,6 +885,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
885
alibi_slopes=self.alibi_slopes,
886
return_lse=False,
887
out=output,
888
+ sink_ptr=self.sinks,
889
)
890
891
def extend_forward(
@@ -927,6 +935,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
935
window_size=self.sliding_window,
936
alibi_slopes=self.alibi_slopes,
937
return_lse=True,
938
+ sink_ptr=self.sinks,
939
)
940
assert attn_metadata.extend_metadata is not None
941
chunk_context_metadata = attn_metadata.extend_metadata.chunk_context_metadata
@@ -974,6 +983,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
983
window_size=self.sliding_window,
984
alibi_slopes=self.alibi_slopes,
985
return_lse=True,
986
+ sink_ptr=self.sinks,
987
)
988
if chunked_output is None:
989
chunked_output = suf_out
@@ -1092,6 +1102,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
1102
window_size=self.sliding_window,
1103
alibi_slopes=self.alibi_slopes,
1104
out=output_actual_tokens[num_decode_tokens + num_extend_tokens :],
1105
+ sink_ptr=self.sinks,
1106
)
1107
1108
# calculate for extends
@@ -1136,11 +1147,17 @@ class AiterFlashAttentionImpl(AttentionImpl):
1147
assert attn_metadata.decode_metadata is not None
1148
decode_max_query_len = attn_metadata.decode_metadata.max_query_len
1149
1139
- # Multi-token speculative decode path.
1140
- if decode_max_query_len > 1:
1150
+ # Use unified_attention for speculative decoding (multi-token),
1151
+ # sliding window, or sinks
1152
+ # (pa_fwd_asm and paged_attention_v1 don't support sinks)
1153
+ if (
1154
+ self.sliding_window[0] != -1
1155
+ or decode_max_query_len > 1
1156
+ or self.sinks is not None
1157
+ ):
1158
assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
1142
- "Shuffle KV cache layout is not supported with "
1143
- "speculative decoding (multi-token decode)."
1159
+ "Shuffle KV cache layout is not supported with sliding "
1160
+ "window, sinks, or speculative decoding (multi-token decode)."
1161
)
1162
if not attn_metadata.causal:
1163
from aiter.ops.triton.attention.mha_v3 import (
@@ -1207,6 +1224,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
1224
q_descale=None,
1225
k_descale=layer._k_scale.expand(descale_shape),
1226
v_descale=layer._v_scale.expand(descale_shape),
1227
+ sinks=self.sinks,
1228
)
1229
return
1230