From bd091079cba0800d8b8ee8ed22feab5864d1b101 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 20 Jul 2026 13:53:27 -0400 Subject: [PATCH] [Attention] FlashAttention 4 SM100 FP8 kv cache support (#42569) Signed-off-by: Matthew Bonanni Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- cmake/external_projects/vllm_flash_attn.cmake | 2 +- docs/design/attention_backends.md | 4 +- tests/models/quantization/test_fp8.py | 10 +-- tests/v1/attention/test_attention_backends.py | 71 ++++++++++++++++-- tests/v1/attention/utils.py | 2 + .../generate_attention_backend_docs.py | 73 +++++++++++++++++-- vllm/v1/attention/backends/fa_utils.py | 23 ++++++ vllm/v1/attention/backends/flash_attn.py | 42 +++++++++-- 8 files changed, 194 insertions(+), 33 deletions(-) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 28a6b336e2..89ef9192e8 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8 + GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199 GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 8062ce237d..9497f70e3b 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first). | `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 | | `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x | | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 | -| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x | -| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | +| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x | +| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any | | `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 | diff --git a/tests/models/quantization/test_fp8.py b/tests/models/quantization/test_fp8.py index 6a13794427..e27d52d15b 100644 --- a/tests/models/quantization/test_fp8.py +++ b/tests/models/quantization/test_fp8.py @@ -9,8 +9,8 @@ Note: these tests will only pass on L4 GPU. import pytest from tests.quantization.utils import is_quant_method_supported +from vllm.v1.attention.backends.fa_utils import flash_attn_supports_kv_cache_dtype from vllm.platforms import current_platform -from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from ..utils import check_logprobs_close @@ -70,13 +70,7 @@ def test_models( if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda(): pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.") - if not ( - current_platform.is_xpu() - or ( - get_flash_attn_version() == 3 - and current_platform.is_device_capability_family(90) - ) - ): + if not flash_attn_supports_kv_cache_dtype(kv_cache_dtype): pytest.skip( f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention." ) diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 3110e4b4ee..b190433b25 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -21,6 +21,7 @@ from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import ( STR_DTYPE_TO_TORCH_DTYPE, + is_quantized_kv_cache, is_torch_equal_or_newer, set_random_seed, ) @@ -45,6 +46,11 @@ BACKENDS_TO_TEST = [ DEVICE_TYPE = current_platform.device_type +FP8_KV_CACHE_DTYPES = { + "fp8": torch.float8_e4m3fn, + "fp8_e4m3": torch.float8_e4m3fn, +} + # Remove flashinfer from the list if it's not available try: import flashinfer # noqa: F401 @@ -110,6 +116,7 @@ def create_and_prepopulate_kv_cache( num_blocks: int, common_attn_metadata: CommonAttentionMetadata, randomize_blocks: bool = True, + kv_cache_dtype: str = "auto", ) -> torch.Tensor: """Create and prepopulate a KV cache with context data. @@ -140,8 +147,18 @@ def create_and_prepopulate_kv_cache( block_table = common_attn_metadata.block_table_tensor slot_mapping = common_attn_metadata.slot_mapping + # For an fp8 kv cache, store the cache in the fp8 dtype so that assigning + # the higher-precision context tensors quantizes them, mirroring runtime. + fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype) + storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype + kv_cache = torch.zeros( - num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device + num_blocks, + block_size, + num_kv_heads, + 2 * head_size, + dtype=storage_dtype, + device=device, ) kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size) @@ -195,7 +212,12 @@ def create_and_prepopulate_kv_cache( ] * block_size + token_inter_block_offsets.to(device) # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs) - return kv_cache.transpose(1, 2).contiguous() + kv_cache = kv_cache.transpose(1, 2).contiguous() + + if fp8_kv_cache: + kv_cache = kv_cache.view(torch.uint8) + + return kv_cache class MockAttentionLayer: @@ -224,6 +246,7 @@ def run_attention_backend( kv_cache: torch.Tensor, attn_type: AttentionType = AttentionType.DECODER, sliding_window: int | None = None, + kv_cache_dtype: str = "auto", ) -> torch.Tensor: """Run attention computation using the specified backend's AttentionImpl.""" @@ -291,13 +314,16 @@ def run_attention_backend( alibi_slopes=None, sliding_window=sliding_window, attn_type=attn_type, - kv_cache_dtype="auto", + kv_cache_dtype=kv_cache_dtype, ) # Create mock layer and output buffer mock_layer = MockAttentionLayer(device) output = torch.empty_like(query) + if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input: + query = query.to(current_platform.fp8_dtype()) + # Run forward pass # NOTE: The query, key, and value are already shaped correctly # in the calling test function. @@ -324,6 +350,7 @@ def _test_backend_correctness( atol: float = 1e-2, rtol: float = 1e-2, tensor_parallel_size: int = 1, + kv_cache_dtype: str = "auto", ): """ Test that all backends produce similar outputs to a reference implementation @@ -372,6 +399,7 @@ def _test_backend_correctness( num_gpu_blocks=8192, hf_config_override=hf_config_override, ) + vllm_config.cache_config.cache_dtype = kv_cache_dtype device = torch.device(f"{DEVICE_TYPE}:0") kv_cache_spec = create_standard_kv_cache_spec(vllm_config, attn_type) @@ -392,6 +420,13 @@ def _test_backend_correctness( block_size = vllm_config.cache_config.block_size scale = 1.0 / (head_size**0.5) + fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype) + if fp8_kv_cache: + query_fp8_dtype = current_platform.fp8_dtype() + kv_fp8_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] + atol = max(atol, 6e-2) + rtol = max(rtol, 1e-1) + # 2. Generate data and compute SDPA reference output all_q_vllm, all_k_vllm, all_v_vllm = [], [], [] all_sdpa_outputs = [] @@ -407,10 +442,17 @@ def _test_backend_correctness( k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device) v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device) + if fp8_kv_cache: + q_ref = q.to(query_fp8_dtype).to(dtype) + k_ref = k_full.to(kv_fp8_dtype).to(dtype) + v_ref = v_full.to(kv_fp8_dtype).to(dtype) + else: + q_ref, k_ref, v_ref = q, k_full, v_full + # SDPA expects (N, H, L, D), so unsqueeze batch and permute - q_sdpa_in = q.unsqueeze(0).transpose(1, 2) - k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2) - v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2) + q_sdpa_in = q_ref.unsqueeze(0).transpose(1, 2) + k_sdpa_in = k_ref.unsqueeze(0).transpose(1, 2) + v_sdpa_in = v_ref.unsqueeze(0).transpose(1, 2) if num_q_heads != num_kv_heads: assert num_q_heads % num_kv_heads == 0, ( @@ -471,6 +513,7 @@ def _test_backend_correctness( num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000, common_attn_metadata=common_attn_metadata, randomize_blocks=True, + kv_cache_dtype=kv_cache_dtype, ) # 4. Run vLLM backends and compare @@ -488,6 +531,12 @@ def _test_backend_correctness( else: backend_cls = None + if is_quantized_kv_cache(kv_cache_dtype) and ( + backend_cls is None + or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype) + ): + continue + if backend_name == AttentionBackendEnum.FLASHINFER: set_kv_cache_layout("HND") reset_kv_cache_layout = True @@ -521,6 +570,7 @@ def _test_backend_correctness( kv_cache_for_backend, sliding_window=sliding_window, attn_type=attn_type, + kv_cache_dtype=kv_cache_dtype, ) finally: if reset_kv_cache_layout: @@ -570,8 +620,13 @@ def _test_backend_correctness( ) @pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"]) @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"]) def test_causal_backend_correctness( - default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int + default_vllm_config, + batch_spec_name: str, + model: str, + tensor_parallel_size: int, + kv_cache_dtype: str, ): """Test backend's correctness with causal attention.""" @@ -612,6 +667,7 @@ def test_causal_backend_correctness( SMALL_BLOCK_BACKENDS, causal_mask_mod, tensor_parallel_size=tensor_parallel_size, + kv_cache_dtype=kv_cache_dtype, ) # Fast FlexAttention needs to run with block_size=128 @@ -623,6 +679,7 @@ def test_causal_backend_correctness( causal_mask_mod, block_size=128, tensor_parallel_size=tensor_parallel_size, + kv_cache_dtype=kv_cache_dtype, ) diff --git a/tests/v1/attention/utils.py b/tests/v1/attention/utils.py index 3daafcdc6f..20562cfe85 100644 --- a/tests/v1/attention/utils.py +++ b/tests/v1/attention/utils.py @@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import ( EncoderOnlyAttentionSpec, FullAttentionSpec, MambaSpec, + get_kv_quant_mode, ) @@ -178,6 +179,7 @@ def create_standard_kv_cache_spec( head_size=vllm_config.model_config.get_head_size(), dtype=vllm_config.model_config.dtype, sliding_window=vllm_config.model_config.get_sliding_window(), + kv_quant_mode=get_kv_quant_mode(vllm_config.cache_config.cache_dtype), ) diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index d44456530d..875d19a0f5 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -1083,7 +1083,8 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]: return {} # Analyze the functions to determine FA3/FA4-specific features - fa3_supports_fp8 = True + fa3_supports_fp8 = False + fa4_supports_fp8 = False fa3_supports_sinks = False fa4_supports_sinks = False fa3_compute_cap: str | None = None @@ -1093,6 +1094,44 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]: if not isinstance(node, ast.FunctionDef): continue + # Check flash_attn_supports_kv_cache_dtype for fp8 support per FA version. + # Accept both equality checks and membership checks such as `in (3, 4)`. + if node.name == "flash_attn_supports_kv_cache_dtype": + for n in ast.walk(node): + if not ( + isinstance(n, ast.Compare) + and len(n.ops) == 1 + and len(n.comparators) == 1 + ): + continue + is_version_compare = ( + isinstance(n.left, ast.Name) and n.left.id == "fa_version" + ) or ( + isinstance(n.left, ast.Call) + and isinstance(n.left.func, ast.Name) + and n.left.func.id == "get_flash_attn_version" + ) + if not is_version_compare: + continue + + versions: list[Any] = [] + comparator = n.comparators[0] + if isinstance(n.ops[0], ast.Eq) and isinstance( + comparator, ast.Constant + ): + versions = [comparator.value] + elif isinstance(n.ops[0], ast.In) and isinstance( + comparator, (ast.Tuple, ast.List, ast.Set) + ): + versions = [ + elt.value + for elt in comparator.elts + if isinstance(elt, ast.Constant) + ] + + fa3_supports_fp8 |= 3 in versions + fa4_supports_fp8 |= 4 in versions + # Check flash_attn_supports_sinks - looks for `fa_version == 3/4` # or `get_flash_attn_version() == 3/4` (also accepts `in (3, 4)`) if node.name == "flash_attn_supports_sinks": @@ -1199,7 +1238,7 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]: }, "fa4": { "compute_capability": fa4_compute_cap, - "supports_fp8": False, + "supports_fp8": fa4_supports_fp8, "supports_sink": fa4_supports_sinks, }, } @@ -1255,6 +1294,22 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]: # --------------------------------------------------------------------------- +def _apply_fp8_support(kv_cache_dtypes: str, supports_fp8: bool) -> str: + """Add or remove fp8 dtypes from a comma-separated kv_cache_dtypes string. + + The base FLASH_ATTN backend lists fp8 dtypes in ``supported_kv_cache_dtypes``, + but actual support varies by FA version (e.g. FA2 has no fp8 support), so each + variant's row is normalized to match its own capability. + """ + fp8_dtypes = {"fp8", "fp8_e4m3", "fp8_e5m2"} + dtypes = kv_cache_dtypes.split(", ") + base_dtypes = [dtype for dtype in dtypes if dtype not in fp8_dtypes] + if supports_fp8: + advertised_fp8_dtypes = [dtype for dtype in dtypes if dtype in fp8_dtypes] + return ", ".join(base_dtypes + advertised_fp8_dtypes) + return ", ".join(base_dtypes) + + def _expand_flash_attn_variants( all_backends: list[dict[str, Any]], fa_features: dict[str, dict[str, Any]], @@ -1275,6 +1330,9 @@ def _expand_flash_attn_variants( fa2["_sort_key"] = "FLASH_ATTN" fa2["_sort_order"] = 0 fa2["supports_sink"] = fa_features["fa2"]["supports_sink"] + fa2["kv_cache_dtypes"] = _apply_fp8_support( + backend["kv_cache_dtypes"], fa_features["fa2"]["supports_fp8"] + ) # Create FA3 entry (uses parsed compute_capability from fa_utils) fa3 = backend.copy() @@ -1284,11 +1342,9 @@ def _expand_flash_attn_variants( if fa_features["fa3"]["compute_capability"]: fa3["compute_capability"] = fa_features["fa3"]["compute_capability"] fa3["supports_sink"] = fa_features["fa3"]["supports_sink"] - if fa_features["fa3"]["supports_fp8"]: - base_dtypes = backend["kv_cache_dtypes"].split(", ") - fp8_dtypes = ["fp8", "fp8_e4m3", "fp8_e5m2"] - new_dtypes = [d for d in fp8_dtypes if d not in base_dtypes] - fa3["kv_cache_dtypes"] = ", ".join(base_dtypes + new_dtypes) + fa3["kv_cache_dtypes"] = _apply_fp8_support( + backend["kv_cache_dtypes"], fa_features["fa3"]["supports_fp8"] + ) expanded.append(fa2) expanded.append(fa3) @@ -1302,6 +1358,9 @@ def _expand_flash_attn_variants( if fa_features["fa4"].get("compute_capability"): fa4["compute_capability"] = fa_features["fa4"]["compute_capability"] fa4["supports_sink"] = fa_features["fa4"]["supports_sink"] + fa4["kv_cache_dtypes"] = _apply_fp8_support( + backend["kv_cache_dtypes"], fa_features["fa4"]["supports_fp8"] + ) expanded.append(fa4) return expanded diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 6c0debab99..718580dcf6 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -273,6 +273,29 @@ def is_fa_version_supported(fa_version: int) -> bool: return False +def flash_attn_supports_kv_cache_dtype( + kv_cache_dtype: str = "fp8_e4m3", + *, + requires_alibi: bool = False, + head_size: int | None = None, + head_size_v: int | None = None, + has_sinks: bool = False, +) -> bool: + if kv_cache_dtype == "fp8_e5m2": + return False + if current_platform.is_xpu(): + return True + fa_version = get_flash_attn_version( + requires_alibi=requires_alibi, + head_size=head_size, + head_size_v=head_size_v, + has_sinks=has_sinks, + ) + return (fa_version == 3 and current_platform.is_device_capability_family(90)) or ( + fa_version == 4 and current_platform.is_device_capability_family(100) + ) + + def flash_attn_supports_quant_query_input() -> bool: return not current_platform.is_xpu() diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 86198dda61..13bf435a54 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -22,6 +22,7 @@ from vllm.v1.attention.backend import ( MultipleOf, ) from vllm.v1.attention.backends.fa_utils import ( + flash_attn_supports_kv_cache_dtype, flash_attn_supports_quant_query_input, get_flash_attn_version, is_fa_version_supported, @@ -74,6 +75,8 @@ class FlashAttentionBackend(AttentionBackend): "auto", "float16", "bfloat16", + "fp8", + "fp8_e4m3", ] @staticmethod @@ -178,14 +181,11 @@ class FlashAttentionBackend(AttentionBackend): def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: if kv_cache_dtype is None: return True - if kv_cache_dtype in ("fp8", "fp8_e4m3"): - if current_platform.is_xpu(): - return True - return ( - get_flash_attn_version() == 3 - and current_platform.is_device_capability_family(90) - ) - return kv_cache_dtype in ["auto", "float16", "bfloat16"] + if kv_cache_dtype not in cls.supported_kv_cache_dtypes: + return False + if is_quantized_kv_cache(kv_cache_dtype): + return flash_attn_supports_kv_cache_dtype(kv_cache_dtype) + return True @classmethod def supports_mm_prefix(cls) -> bool: @@ -216,6 +216,17 @@ class FlashAttentionBackend(AttentionBackend): ) -> str | None: if has_sink and device_capability < DeviceCapability(9, 0): return "sink not supported on compute capability < 9.0" + if ( + kv_cache_dtype is not None + and is_quantized_kv_cache(kv_cache_dtype) + and not flash_attn_supports_kv_cache_dtype( + kv_cache_dtype, + head_size=head_size, + head_size_v=head_size, + has_sinks=has_sink, + ) + ): + return "FP8 KV cache requires FA3 on SM90 or FA4 on SM100" if ( use_mm_prefix and get_flash_attn_version(head_size=head_size, has_sinks=has_sink) != 4 @@ -772,6 +783,7 @@ class FlashAttentionImpl(AttentionImpl): self.vllm_flash_attn_version = get_flash_attn_version( requires_alibi=alibi_slopes is not None, head_size=head_size, + has_sinks=sinks is not None, ) logger.info_once( "Using FlashAttention version %s", @@ -780,6 +792,20 @@ class FlashAttentionImpl(AttentionImpl): # Cache the batch invariant result for use in forward passes self.batch_invariant_enabled = envs.VLLM_BATCH_INVARIANT + if is_quantized_kv_cache( + self.kv_cache_dtype + ) and not flash_attn_supports_kv_cache_dtype( + self.kv_cache_dtype, + requires_alibi=alibi_slopes is not None, + head_size=head_size, + head_size_v=head_size, + has_sinks=sinks is not None, + ): + raise NotImplementedError( + f"FlashAttention does not support {self.kv_cache_dtype}" + " kv-cache on this device." + ) + self.sinks = sinks if self.sinks is not None: assert flash_attn_supports_sinks(), (