| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Triton image-mask and FlashAttention causal attention composite.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.config.cache import CacheDType |
| 8 | from vllm.platforms.interface import DeviceCapability |
| 9 | from vllm.v1.attention.backends.composite import ( |
| 10 | MMPrefixAttentionRouting, |
| 11 | create_composite_attention_backend, |
| 12 | ) |
| 13 | from vllm.v1.attention.backends.fa_utils import get_flash_attn_version |
| 14 | from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend |
| 15 | from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend |
| 16 | |
| 17 | |
| 18 | class _FABackend(FlashAttentionBackend): |
| 19 | @classmethod |
| 20 | def supports_combination( |
| 21 | cls, |
| 22 | head_size: int, |
| 23 | dtype: torch.dtype, |
| 24 | kv_cache_dtype: CacheDType | None, |
| 25 | block_size: int | None, |
| 26 | use_mla: bool, |
| 27 | has_sink: bool, |
| 28 | use_sparse: bool, |
| 29 | use_mm_prefix: bool, |
| 30 | device_capability: DeviceCapability, |
| 31 | ) -> str | None: |
| 32 | reason = super().supports_combination( |
| 33 | head_size, |
| 34 | dtype, |
| 35 | kv_cache_dtype, |
| 36 | block_size, |
| 37 | use_mla, |
| 38 | has_sink, |
| 39 | use_sparse, |
| 40 | use_mm_prefix, |
| 41 | device_capability, |
| 42 | ) |
| 43 | if reason is not None: |
| 44 | return reason |
| 45 | fa_version = get_flash_attn_version( |
| 46 | head_size=head_size, |
| 47 | has_sinks=has_sink, |
| 48 | kv_cache_block_size=block_size, |
| 49 | supports_fa4_hd256=True, |
| 50 | ) |
| 51 | if fa_version not in (3, 4): |
| 52 | return "causal route requires FlashAttention v3 or v4" |
| 53 | return None |
| 54 | |
| 55 | |
| 56 | TritonFlashAttentionBackend = create_composite_attention_backend( |
| 57 | TritonAttentionBackend, |
| 58 | _FABackend, |
| 59 | name="TritonFlashAttentionBackend", |
| 60 | backend_name="TRITON_FLASH_ATTN", |
| 61 | module=__name__, |
| 62 | routing_policy=MMPrefixAttentionRouting, |
| 63 | head_sizes=(256, 512), |
| 64 | device_major=9, |
| 65 | ) |