| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Triton image-mask and FlashInfer causal attention composite.""" |
| 4 | |
| 5 | from vllm.v1.attention.backends.composite import ( |
| 6 | MMPrefixAttentionRouting, |
| 7 | create_composite_attention_backend, |
| 8 | ) |
| 9 | from vllm.v1.attention.backends.flashinfer import FlashInferBackend |
| 10 | from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend |
| 11 | |
| 12 | TritonFlashInferBackend = create_composite_attention_backend( |
| 13 | TritonAttentionBackend, |
| 14 | FlashInferBackend, |
| 15 | name="TritonFlashInferBackend", |
| 16 | backend_name="TRITON_FLASHINFER", |
| 17 | module=__name__, |
| 18 | routing_policy=MMPrefixAttentionRouting, |
| 19 | head_sizes=(256, 512), |
| 20 | kernel_block_sizes=(64,), |
| 21 | device_major=10, |
| 22 | ) |