iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Attention] FlashAttention 4 SM100 FP8 kv cache support (#42569)

Browse Source 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Matthew Bonanni
2026-07-20 10:53:27 -07:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent b23bd73f54
commit bd091079cb
8 changed files with 194 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/external_projects/vllm_flash_attn.cmake
+1 -1
View File
@@ -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
docs/design/attention_backends.md
+2 -2
View File
@@ -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 |
tests/models/quantization/test_fp8.py
+2 -8
View File
@@ -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."
)
tests/v1/attention/test_attention_backends.py
+64 -7
View File
@@ -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,
)
tests/v1/attention/utils.py
+2
View File
@@ -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),
)
tools/pre_commit/generate_attention_backend_docs.py
+66 -7
View File
@@ -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
vllm/v1/attention/backends/fa_utils.py
+23
View File
@@ -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()
vllm/v1/attention/backends/flash_attn.py
+34 -8
View File
@@ -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(), (