diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index ed9f0d3016..329a4aacfb 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -177,7 +177,7 @@ Priority is **1 = highest** (tried first). | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any | -| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A | | `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 | | `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 | diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 2c9b939fa2..5a8b690433 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2394,6 +2394,7 @@ class rocm_aiter_ops: alibi_slopes: torch.Tensor | None = None, return_lse: bool = False, out: torch.Tensor | None = None, + sink_ptr: torch.Tensor | None = None, ): """ Flash attention with variable length sequences. @@ -2422,6 +2423,7 @@ class rocm_aiter_ops: alibi_slopes=alibi_slopes, return_lse=return_lse, out=out, + sink_ptr=sink_ptr, ) @staticmethod diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index d0cc011fe4..a9fa45debc 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -709,6 +709,11 @@ class AiterFlashAttentionMetadataBuilder( class AiterFlashAttentionBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] + + @classmethod + def supports_sink(cls) -> bool: + return True + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", "float16", @@ -788,6 +793,7 @@ class AiterFlashAttentionImpl(AttentionImpl): logits_soft_cap: float | None = None, attn_type: AttentionType = AttentionType.DECODER, kv_sharing_target_layer_name: int | None = None, + sinks: torch.Tensor | None = None, ) -> None: self.num_heads = num_heads self.head_size = head_size @@ -806,6 +812,7 @@ class AiterFlashAttentionImpl(AttentionImpl): logits_soft_cap = 0.0 self.logits_soft_cap = logits_soft_cap self.kv_sharing_target_layer_name = kv_sharing_target_layer_name + self.sinks = sinks assert self.num_heads % self.num_kv_heads == 0 self.num_queries_per_kv = self.num_heads // self.num_kv_heads @@ -878,6 +885,7 @@ class AiterFlashAttentionImpl(AttentionImpl): alibi_slopes=self.alibi_slopes, return_lse=False, out=output, + sink_ptr=self.sinks, ) def extend_forward( @@ -927,6 +935,7 @@ class AiterFlashAttentionImpl(AttentionImpl): window_size=self.sliding_window, alibi_slopes=self.alibi_slopes, return_lse=True, + sink_ptr=self.sinks, ) assert attn_metadata.extend_metadata is not None chunk_context_metadata = attn_metadata.extend_metadata.chunk_context_metadata @@ -974,6 +983,7 @@ class AiterFlashAttentionImpl(AttentionImpl): window_size=self.sliding_window, alibi_slopes=self.alibi_slopes, return_lse=True, + sink_ptr=self.sinks, ) if chunked_output is None: chunked_output = suf_out @@ -1092,6 +1102,7 @@ class AiterFlashAttentionImpl(AttentionImpl): window_size=self.sliding_window, alibi_slopes=self.alibi_slopes, out=output_actual_tokens[num_decode_tokens + num_extend_tokens :], + sink_ptr=self.sinks, ) # calculate for extends @@ -1136,11 +1147,17 @@ class AiterFlashAttentionImpl(AttentionImpl): assert attn_metadata.decode_metadata is not None decode_max_query_len = attn_metadata.decode_metadata.max_query_len - # Multi-token speculative decode path. - if decode_max_query_len > 1: + # Use unified_attention for speculative decoding (multi-token), + # sliding window, or sinks + # (pa_fwd_asm and paged_attention_v1 don't support sinks) + if ( + self.sliding_window[0] != -1 + or decode_max_query_len > 1 + or self.sinks is not None + ): assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( - "Shuffle KV cache layout is not supported with " - "speculative decoding (multi-token decode)." + "Shuffle KV cache layout is not supported with sliding " + "window, sinks, or speculative decoding (multi-token decode)." ) if not attn_metadata.causal: from aiter.ops.triton.attention.mha_v3 import ( @@ -1207,6 +1224,7 @@ class AiterFlashAttentionImpl(AttentionImpl): q_descale=None, k_descale=layer._k_scale.expand(descale_shape), v_descale=layer._v_scale.expand(descale_shape), + sinks=self.sinks, ) return