| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend, AiterMLAImpl |
| 4 | |
| 5 | |
| 6 | class AiterTritonMLABackend(AiterMLABackend): |
| 7 | @staticmethod |
| 8 | def get_name() -> str: |
| 9 | return "AITER_TRITON_MLA" |
| 10 | |
| 11 | @staticmethod |
| 12 | def get_impl_cls() -> type["AiterTritonMLAImpl"]: |
| 13 | return AiterTritonMLAImpl |
| 14 | |
| 15 | |
| 16 | class AiterTritonMLAImpl(AiterMLAImpl): |
| 17 | def __init__( |
| 18 | self, |
| 19 | num_heads: int, |
| 20 | head_size: int, |
| 21 | scale: float, |
| 22 | num_kv_heads: int, |
| 23 | alibi_slopes: list[float] | None, |
| 24 | sliding_window: int | None, |
| 25 | kv_cache_dtype: str, |
| 26 | logits_soft_cap: float | None, |
| 27 | attn_type: str, |
| 28 | kv_sharing_target_layer_name: str | None, |
| 29 | # MLA Specific Arguments |
| 30 | **mla_args, |
| 31 | ) -> None: |
| 32 | super().__init__( |
| 33 | num_heads, |
| 34 | head_size, |
| 35 | scale, |
| 36 | num_kv_heads, |
| 37 | alibi_slopes, |
| 38 | sliding_window, |
| 39 | kv_cache_dtype, |
| 40 | logits_soft_cap, |
| 41 | attn_type, |
| 42 | kv_sharing_target_layer_name, |
| 43 | **mla_args, |
| 44 | ) |
| 45 | from aiter.ops.triton.mha import flash_attn_varlen_func |
| 46 | |
| 47 | self.flash_attn_varlen_func = flash_attn_varlen_func |
| 48 | |
| 49 | def _flash_attn_varlen_diff_headdims( |
| 50 | self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs |
| 51 | ): |
| 52 | result = self.flash_attn_varlen_func( # type: ignore[call-arg] |
| 53 | q, |
| 54 | k, |
| 55 | v, |
| 56 | softmax_scale=softmax_scale, |
| 57 | return_lse=return_softmax_lse, |
| 58 | **kwargs, |
| 59 | ) |
| 60 | # Transpose the LSE if Triton MHA is used: |
| 61 | # (q.shape[0], num_q_heads) to (num_q_heads, q.shape[0]) |
| 62 | if type(result) is tuple and return_softmax_lse: |
| 63 | output, lse = result |
| 64 | lse = lse.T.contiguous() |
| 65 | return (output, lse) |
| 66 | return result |