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 3, 13:00 UTC
[FlashAttn] Fix supports_kv_cache_dtype() accepting unhandled fp8 kv-cache dtype variants (#42685)
Signed-off-by: Lanze Liu <lanzetech@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Lanze Liu committed
May 15, 2026 at 12:34 UTC
b2c58ee9427f15563210e184c57a6e530f37e464
6 files changed
+54
-52
tests/kernels/attention/test_attention_selector.py
+32
@@ -514,3 +514,35 @@ def test_non_causal_autoselect_backend():
514
kv_cache_dtype=None,
515
)
516
assert backend.supports_non_causal()
517
+
518
+
519
+@pytest.mark.parametrize(
520
+ "kv_cache_dtype",
521
+ [
522
+ "fp8_e5m2",
523
+ "fp8_ds_mla",
524
+ "fp8_inc",
525
+ "nvfp4",
526
+ "fp8_per_token_head",
527
+ "int8_per_token_head",
528
+ ],
529
+)
530
+def test_flash_attn_rejects_unhandled_kv_cache_dtypes(kv_cache_dtype: str):
531
+ """FlashAttentionBackend must not claim support for kv_cache dtypes
532
+ that it cannot handle."""
533
+ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
534
+
535
+ assert not FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
536
+
537
+
538
+@pytest.mark.parametrize("kv_cache_dtype", ["fp8", "fp8_e4m3"])
539
+def test_flash_attn_accepts_handled_fp8_variants(
540
+ kv_cache_dtype: str, monkeypatch: pytest.MonkeyPatch
541
+):
542
+ """FlashAttentionBackend must accept the two fp8 dtypes it can actually
543
+ handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'."""
544
+ import vllm.v1.attention.backends.flash_attn as fa_mod
545
+ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
546
+
547
+ monkeypatch.setattr(fa_mod.current_platform, "is_xpu", lambda: True)
548
+ assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
tests/models/quantization/test_fp8.py
+8
-2
@@ -9,8 +9,8 @@ Note: these tests will only pass on L4 GPU.
9
import pytest
10
11
from tests.quantization.utils import is_quant_method_supported
12
-from vllm.v1.attention.backends.fa_utils import flash_attn_supports_fp8
12
from vllm.platforms import current_platform
13
+from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
14
from ..utils import check_logprobs_close
15
16
@@ -68,7 +68,13 @@ def test_models(
68
if kv_cache_dtype == "fp8_e5m2" and current_platform.is_rocm():
69
pytest.skip(f"{kv_cache_dtype} is currently not supported on ROCm/HIP.")
70
71
- if not flash_attn_supports_fp8():
71
+ if not (
72
+ current_platform.is_xpu()
73
+ or (
74
+ get_flash_attn_version() == 3
75
+ and current_platform.is_device_capability_family(90)
76
+ )
77
+ ):
78
pytest.skip(
79
f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention."
80
)
tools/pre_commit/generate_attention_backend_docs.py
+1
-13
@@ -878,7 +878,7 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
878
return {}
879
880
# Analyze the functions to determine FA3/FA4-specific features
881
- fa3_supports_fp8 = False
881
+ fa3_supports_fp8 = True
882
fa3_supports_sinks = False
883
fa4_supports_sinks = False
884
fa3_compute_cap: str | None = None
@@ -888,18 +888,6 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
888
if not isinstance(node, ast.FunctionDef):
889
continue
890
891
- # Check flash_attn_supports_fp8 - looks for `get_flash_attn_version() == 3`
892
- if node.name == "flash_attn_supports_fp8":
893
- for n in ast.walk(node):
894
- if (
895
- isinstance(n, ast.Compare)
896
- and isinstance(n.left, ast.Call)
897
- and isinstance(n.left.func, ast.Name)
898
- and n.left.func.id == "get_flash_attn_version"
899
- ):
900
- fa3_supports_fp8 = True
901
- break
902
-
891
# Check flash_attn_supports_sinks - looks for `fa_version == 3/4`
892
# or `get_flash_attn_version() == 3/4` (also accepts `in (3, 4)`)
893
if node.name == "flash_attn_supports_sinks":
vllm/v1/attention/backends/fa_utils.py
-9
@@ -191,15 +191,6 @@ def is_fa_version_supported(fa_version: int) -> bool:
191
return False
192
193
194
-def flash_attn_supports_fp8() -> bool:
195
- if current_platform.is_xpu():
196
- return True
197
- return (
198
- get_flash_attn_version() == 3
199
- and current_platform.is_device_capability_family(90)
200
- )
201
-
202
-
194
def flash_attn_supports_quant_query_input() -> bool:
195
return not current_platform.is_xpu()
196
vllm/v1/attention/backends/flash_attn.py
+10
-23
@@ -22,7 +22,6 @@ from vllm.v1.attention.backend import (
22
MultipleOf,
23
)
24
from vllm.v1.attention.backends.fa_utils import (
25
- flash_attn_supports_fp8,
25
flash_attn_supports_quant_query_input,
26
get_flash_attn_version,
27
is_fa_version_supported,
@@ -170,13 +169,6 @@ class FlashAttentionBackend(AttentionBackend):
169
raise ValueError(f"Unknown cache layout format {cache_layout}.")
170
return stride_order
171
173
- @staticmethod
174
- def get_fp8_dtype_for_flashattn(kv_cache_dtype: str) -> torch.dtype:
175
- if kv_cache_dtype in ("fp8", "fp8_e4m3"):
176
- return torch.float8_e4m3fn
177
- else:
178
- raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
179
-
172
@classmethod
173
def supports_head_size(cls, head_size: int) -> bool:
174
if head_size % 8 != 0:
@@ -191,8 +183,13 @@ class FlashAttentionBackend(AttentionBackend):
183
def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
184
if kv_cache_dtype is None:
185
return True
194
- if is_quantized_kv_cache(kv_cache_dtype):
195
- return flash_attn_supports_fp8()
186
+ if kv_cache_dtype in ("fp8", "fp8_e4m3"):
187
+ if current_platform.is_xpu():
188
+ return True
189
+ return (
190
+ get_flash_attn_version() == 3
191
+ and current_platform.is_device_capability_family(90)
192
+ )
193
return kv_cache_dtype in ["auto", "float16", "bfloat16"]
194
195
@classmethod
@@ -450,9 +447,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
447
):
448
cache_dtype = self.cache_config.cache_dtype
449
if is_quantized_kv_cache(cache_dtype):
453
- qkv_dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(
454
- cache_dtype
455
- )
450
+ qkv_dtype = current_platform.fp8_dtype()
451
else:
452
qkv_dtype = self.kv_cache_dtype
453
if aot_schedule:
@@ -645,11 +640,6 @@ class FlashAttentionImpl(AttentionImpl):
640
# Cache the batch invariant result for use in forward passes
641
self.batch_invariant_enabled = envs.VLLM_BATCH_INVARIANT
642
648
- if is_quantized_kv_cache(self.kv_cache_dtype) and not flash_attn_supports_fp8():
649
- raise NotImplementedError(
650
- "FlashAttention does not support fp8 kv-cache on this device."
651
- )
652
-
643
self.sinks = sinks
644
if self.sinks is not None:
645
assert flash_attn_supports_sinks(), (
@@ -762,11 +752,8 @@ class FlashAttentionImpl(AttentionImpl):
752
753
if is_quantized_kv_cache(self.kv_cache_dtype):
754
# queries are quantized in the attention layer
765
- dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(
766
- self.kv_cache_dtype
767
- )
768
- key_cache = key_cache.view(dtype)
769
- value_cache = value_cache.view(dtype)
755
+ key_cache = key_cache.view(current_platform.fp8_dtype())
756
+ value_cache = value_cache.view(current_platform.fp8_dtype())
757
758
if not attn_metadata.use_cascade:
759
cu_seqlens_q = attn_metadata.query_start_loc
vllm/v1/attention/backends/flash_attn_diffkv.py
+3
-5
@@ -5,6 +5,7 @@
5
import torch
6
7
from vllm.logger import init_logger
8
+from vllm.platforms import current_platform
9
from vllm.utils.torch_utils import (
10
canonicalize_singleton_dim_strides,
11
is_quantized_kv_cache,
@@ -230,11 +231,8 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):
231
232
if is_quantized_kv_cache(self.kv_cache_dtype):
233
# queries are quantized in the attention layer
233
- dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(
234
- self.kv_cache_dtype
235
- )
236
- key_cache = key_cache.view(dtype)
237
- value_cache = value_cache.view(dtype)
234
+ key_cache = key_cache.view(current_platform.fp8_dtype())
235
+ value_cache = value_cache.view(current_platform.fp8_dtype())
236
237
if not attn_metadata.use_cascade:
238
cu_seqlens_q = attn_metadata.query_start_loc