A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 02:15 UTC
[Misc] Log the reason for falling back to FlexAttention (#20699)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Jul 14, 2025 at 19:16 UTC
e8cc53af5e17205470c04f442e67f276e08623a1
10 files changed
+104
-32
vllm/attention/selector.py
+40
-9
@@ -3,6 +3,7 @@
3
4
import os
5
from contextlib import contextmanager
6
+from dataclasses import dataclass
7
from functools import cache
8
from typing import Generator, Optional, Union
9
@@ -79,31 +80,61 @@ def get_global_forced_attn_backend() -> Optional[_Backend]:
80
return forced_attn_backend
81
82
82
-def supports_head_size(
83
+@dataclass(frozen=True)
84
+class _IsSupported:
85
+ can_import: bool
86
+ head_size: bool
87
+ dtype: bool
88
+
89
+ def __bool__(self) -> bool:
90
+ return self.can_import and self.head_size and self.dtype
91
+
92
+
93
+def is_attn_backend_supported(
94
attn_backend: Union[str, type[AttentionBackend]],
95
head_size: int,
85
-) -> bool:
96
+ dtype: torch.dtype,
97
+ *,
98
+ allow_import_error: bool = True,
99
+) -> _IsSupported:
100
if isinstance(attn_backend, str):
101
try:
102
attn_backend = resolve_obj_by_qualname(attn_backend)
103
except ImportError:
90
- return False
104
+ if not allow_import_error:
105
+ raise
106
+
107
+ return _IsSupported(can_import=False, head_size=False, dtype=False)
108
109
assert isinstance(attn_backend, type)
110
111
# TODO: Update the interface once V0 is removed
112
if get_supported_head_sizes := getattr(attn_backend,
113
"get_supported_head_sizes", None):
97
- return head_size in get_supported_head_sizes()
98
- if validate_head_size := getattr(attn_backend, "validate_head_size", None):
114
+ is_head_size_supported = head_size in get_supported_head_sizes()
115
+ elif validate_head_size := getattr(attn_backend, "validate_head_size",
116
+ None):
117
try:
118
validate_head_size(head_size)
101
- return True
119
+ is_head_size_supported = True
120
except Exception:
103
- return False
121
+ is_head_size_supported = False
122
+ else:
123
+ raise NotImplementedError(f"{attn_backend.__name__} does not support "
124
+ "head size validation")
125
+
126
+ if get_supported_dtypes := getattr(attn_backend, "get_supported_dtypes",
127
+ None):
128
+ is_dtype_supported = dtype in get_supported_dtypes()
129
+ else:
130
+ raise NotImplementedError(f"{attn_backend.__name__} does not support "
131
+ "dtype validation")
132
105
- raise NotImplementedError(f"{attn_backend.__name__} does not support "
106
- "head size validation")
133
+ return _IsSupported(
134
+ can_import=True,
135
+ head_size=is_head_size_supported,
136
+ dtype=is_dtype_supported,
137
+ )
138
139
140
def get_attn_backend(
vllm/platforms/cuda.py
+35
-22
@@ -259,43 +259,56 @@ class CudaPlatformBase(Platform):
259
logger.info_once("Using Flash Attention backend on V1 engine.")
260
return FLASH_ATTN_V1
261
262
- from vllm.attention.selector import supports_head_size
262
+ from vllm.attention.selector import is_attn_backend_supported
263
264
# Default backends for V1 engine
265
- # FP32 is only supported by FlexAttention
266
- if dtype not in (torch.float16, torch.bfloat16):
267
- logger.info_once(
268
- "Using FlexAttention backend for %s on V1 engine.",
269
- dtype,
270
- )
271
- return FLEX_ATTENTION_V1
272
-
265
# Prefer FlashInfer for Blackwell GPUs if installed
274
- if cls.is_device_capability(100) and \
275
- supports_head_size(FLASHINFER_V1, head_size):
276
- try:
277
- import flashinfer # noqa: F401
278
-
266
+ if cls.is_device_capability(100):
267
+ if is_default_backend_supported := is_attn_backend_supported(
268
+ FLASHINFER_V1, head_size, dtype):
269
from vllm.v1.attention.backends.utils import (
270
set_kv_cache_layout)
271
+
272
logger.info_once(
273
"Using FlashInfer backend with HND KV cache layout on "
274
"V1 engine by default for Blackwell (SM 10.0) GPUs.")
275
set_kv_cache_layout("HND")
276
+
277
return FLASHINFER_V1
286
- except ImportError:
287
- logger.info_once(
278
+
279
+ if not is_default_backend_supported.can_import:
280
+ logger.warning_once(
281
"FlashInfer failed to import for V1 engine on "
282
"Blackwell (SM 10.0) GPUs; it is recommended to "
283
"install FlashInfer for better performance.")
291
- pass
284
+
285
# FlashAttention is the default for SM 8.0+ GPUs
293
- if cls.has_device_capability(80) and \
294
- supports_head_size(FLASH_ATTN_V1, head_size):
295
- logger.info_once("Using Flash Attention backend on V1 engine.")
296
- return FLASH_ATTN_V1
286
+ if cls.has_device_capability(80):
287
+ if is_default_backend_supported := is_attn_backend_supported(
288
+ FLASH_ATTN_V1, head_size, dtype,
289
+ allow_import_error=False):
290
+ logger.info_once("Using Flash Attention backend on "
291
+ "V1 engine.")
292
+ return FLASH_ATTN_V1
293
+
294
+ # FlexAttention is the default for older GPUs
295
+ else:
296
+ logger.info_once("Using FlexAttention backend on V1 engine.")
297
+ return FLEX_ATTENTION_V1
298
+
299
+ assert not is_default_backend_supported
300
+
301
+ use_flex_attention_reason = {}
302
+ if not is_default_backend_supported.head_size:
303
+ use_flex_attention_reason["head_size"] = head_size
304
+ if not is_default_backend_supported.dtype:
305
+ use_flex_attention_reason["dtype"] = dtype
306
298
- logger.info_once("Using FlexAttention backend on V1 engine.")
307
+ logger.info_once(
308
+ "Using FlexAttention backend for %s on V1 engine.",
309
+ ", ".join(f"{k}={v}"
310
+ for k, v in use_flex_attention_reason.items()),
311
+ )
312
return FLEX_ATTENTION_V1
313
314
# Backends for V0 engine
vllm/reasoning/hunyuan_a13b_reasoning_parser.py
+1
-1
@@ -1,10 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import re
4
from collections.abc import Sequence
5
from typing import Optional, Union
6
7
+import regex as re
8
from transformers import PreTrainedTokenizerBase
9
10
from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
vllm/v1/attention/backends/cpu_attn.py
+4
@@ -37,6 +37,10 @@ logger = init_logger(__name__)
37
class TorchSDPABackend(AttentionBackend):
38
accept_output_buffer: bool = False
39
40
+ @classmethod
41
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
42
+ return [torch.float16, torch.bfloat16, torch.float32]
43
+
44
@classmethod
45
def validate_head_size(cls, head_size: int) -> None:
46
attn_impl = _get_paged_attn_impl()
vllm/v1/attention/backends/flash_attn.py
+4
@@ -44,6 +44,10 @@ class FlashAttentionBackend(AttentionBackend):
44
45
accept_output_buffer: bool = True
46
47
+ @classmethod
48
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
49
+ return [torch.float16, torch.bfloat16]
50
+
51
@classmethod
52
def get_supported_head_sizes(cls) -> list[int]:
53
return [32, 64, 96, 128, 160, 192, 224, 256]
vllm/v1/attention/backends/flashinfer.py
+4
@@ -42,6 +42,10 @@ class FlashInferBackend(AttentionBackend):
42
accept_output_buffer: bool = True
43
cached_sm100a_supported: Optional[bool] = None
44
45
+ @classmethod
46
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
47
+ return [torch.float16, torch.bfloat16]
48
+
49
@classmethod
50
def get_supported_head_sizes(cls) -> list[int]:
51
# https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157
vllm/v1/attention/backends/flex_attention.py
+4
@@ -42,6 +42,10 @@ def _offsets_to_doc_ids_tensor(offsets: torch.Tensor) -> torch.Tensor:
42
class FlexAttentionBackend(AttentionBackend):
43
accept_output_buffer: bool = True
44
45
+ @classmethod
46
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
47
+ return [torch.float16, torch.bfloat16, torch.float32]
48
+
49
@classmethod
50
def validate_head_size(cls, head_size: int) -> None:
51
return # FlexAttention supports any head size
vllm/v1/attention/backends/mla/common.py
+4
@@ -262,6 +262,10 @@ class MLACommonBackend(AttentionBackend):
262
) -> tuple[int, ...]:
263
return (num_blocks, block_size, head_size)
264
265
+ @classmethod
266
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
267
+ return [torch.float16, torch.bfloat16]
268
+
269
@classmethod
270
def get_supported_head_sizes(cls) -> list[int]:
271
return [576]
vllm/v1/attention/backends/rocm_aiter_fa.py
+4
@@ -314,6 +314,10 @@ class AiterFlashAttentionBackend(AttentionBackend):
314
315
accept_output_buffer: bool = True
316
317
+ @classmethod
318
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
319
+ return [torch.float16, torch.bfloat16]
320
+
321
@classmethod
322
def get_supported_head_sizes(cls) -> list[int]:
323
return [32, 64, 96, 128, 160, 192, 224, 256]
vllm/v1/attention/backends/triton_attn.py
+4
@@ -190,6 +190,10 @@ class TritonAttentionBackend(AttentionBackend):
190
191
accept_output_buffer: bool = True
192
193
+ @classmethod
194
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
195
+ return [torch.float16, torch.bfloat16]
196
+
197
@classmethod
198
def get_supported_head_sizes(cls) -> list[int]:
199
return [32, 64, 96, 128, 160, 192, 224, 256]