| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """AITER FlashAttention backend for MLA prefill (ROCm). |
| 4 | |
| 5 | This backend calls ``aiter.flash_attn_varlen_func`` directly, which natively |
| 6 | supports different q/k and v head dims (qk headdim 192, v headdim 128) without |
| 7 | padding V, and dispatches to the fast ``aiter::fmha_fwd_`` kernel on |
| 8 | gfx942/gfx950 (fp16/bf16). |
| 9 | """ |
| 10 | |
| 11 | from typing import TYPE_CHECKING |
| 12 | |
| 13 | import torch |
| 14 | |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend |
| 17 | |
| 18 | if TYPE_CHECKING: |
| 19 | from vllm.config import VllmConfig |
| 20 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 21 | MLACommonPrefillMetadata, |
| 22 | ) |
| 23 | from vllm.platforms.interface import DeviceCapability |
| 24 | |
| 25 | |
| 26 | class AiterFlashAttnPrefillBackend(MLAPrefillBackend): |
| 27 | """AITER FlashAttention backend for MLA prefill.""" |
| 28 | |
| 29 | @staticmethod |
| 30 | def get_name() -> str: |
| 31 | return "ROCM_AITER_FA" |
| 32 | |
| 33 | @classmethod |
| 34 | def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: |
| 35 | if not current_platform.is_rocm(): |
| 36 | return False |
| 37 | from vllm.platforms.rocm import on_mi3xx |
| 38 | |
| 39 | return on_mi3xx() |
| 40 | |
| 41 | @classmethod |
| 42 | def is_available(cls) -> bool: |
| 43 | from vllm._aiter_ops import rocm_aiter_ops |
| 44 | |
| 45 | return rocm_aiter_ops.is_enabled() |
| 46 | |
| 47 | def __init__( |
| 48 | self, |
| 49 | num_heads: int, |
| 50 | scale: float, |
| 51 | kv_lora_rank: int, |
| 52 | qk_nope_head_dim: int, |
| 53 | qk_rope_head_dim: int, |
| 54 | v_head_dim: int, |
| 55 | vllm_config: "VllmConfig", |
| 56 | ) -> None: |
| 57 | super().__init__( |
| 58 | num_heads=num_heads, |
| 59 | scale=scale, |
| 60 | kv_lora_rank=kv_lora_rank, |
| 61 | qk_nope_head_dim=qk_nope_head_dim, |
| 62 | qk_rope_head_dim=qk_rope_head_dim, |
| 63 | v_head_dim=v_head_dim, |
| 64 | vllm_config=vllm_config, |
| 65 | ) |
| 66 | |
| 67 | from aiter import flash_attn_varlen_func |
| 68 | |
| 69 | self.flash_attn_varlen_func = flash_attn_varlen_func |
| 70 | |
| 71 | def run_prefill_new_tokens( |
| 72 | self, |
| 73 | q: torch.Tensor, |
| 74 | k: torch.Tensor, |
| 75 | v: torch.Tensor, |
| 76 | return_softmax_lse: bool, |
| 77 | out: torch.Tensor | None = None, |
| 78 | output_scale: torch.Tensor | None = None, |
| 79 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 80 | assert output_scale is None, ( |
| 81 | "AiterFlashAttnPrefillBackend does not support fused quantized output." |
| 82 | ) |
| 83 | result = self.flash_attn_varlen_func( |
| 84 | q=q, |
| 85 | k=k, |
| 86 | v=v, |
| 87 | cu_seqlens_q=self._prefill_metadata.query_start_loc, |
| 88 | cu_seqlens_k=self._prefill_metadata.query_start_loc, |
| 89 | max_seqlen_q=self._prefill_metadata.max_query_len, |
| 90 | max_seqlen_k=self._prefill_metadata.max_query_len, |
| 91 | softmax_scale=self.scale, |
| 92 | causal=True, |
| 93 | return_lse=return_softmax_lse, |
| 94 | out=out, |
| 95 | ) |
| 96 | |
| 97 | # aiter returns the bare output tensor when return_lse is False, and |
| 98 | # (out, softmax_lse) when it is True. |
| 99 | if return_softmax_lse: |
| 100 | return result[0], result[1] |
| 101 | return result |
| 102 | |
| 103 | def run_prefill_context_chunk( |
| 104 | self, |
| 105 | chunk: "MLACommonPrefillMetadata.ContextChunk", |
| 106 | q: torch.Tensor, |
| 107 | k: torch.Tensor, |
| 108 | v: torch.Tensor, |
| 109 | out: torch.Tensor | None = None, |
| 110 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 111 | assert out is None, ( |
| 112 | "AiterFlashAttnPrefillBackend does not report supports_out(), so it " |
| 113 | "is never given a context-chunk `out` to write into." |
| 114 | ) |
| 115 | out, lse = self.flash_attn_varlen_func( |
| 116 | q=q, |
| 117 | k=k, |
| 118 | v=v, |
| 119 | cu_seqlens_q=chunk.query_start_loc, |
| 120 | cu_seqlens_k=chunk.cu_seq_lens, |
| 121 | max_seqlen_q=chunk.max_query_len, |
| 122 | max_seqlen_k=chunk.max_seq_len, |
| 123 | softmax_scale=self.scale, |
| 124 | causal=False, |
| 125 | return_lse=True, |
| 126 | ) |
| 127 | return out, lse |