@tduncan / Nvidia-Optimization-Stack-VLLM / commits / bd091079cb

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

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Matthew Bonanni committed Jul 20, 2026 at 13:53 UTC bd091079cba0800d8b8ee8ed22feab5864d1b101

8 files changed +194 -33

cmake/external_projects/vllm_flash_attn.cmake

+1 -1

@@ -39,7 +39,7 @@ else()

39 FetchContent_Declare(

40 vllm-flash-attn

41 GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git

42 - GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8

42 + GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199

43 GIT_PROGRESS TRUE

44 # Don't share the vllm-flash-attn build between build types

45 BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

docs/design/attention_backends.md

+2 -2

@@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first).

164 | `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 |

165 | `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 |

166 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |

167 -| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |

168 -| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |

167 +| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |

168 +| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |

169 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |

170 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |

171 | `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |

tests/models/quantization/test_fp8.py

+2 -8

@@ -9,8 +9,8 @@ Note: these tests will only pass on L4 GPU.

9 import pytest

11 from tests.quantization.utils import is_quant_method_supported

12 +from vllm.v1.attention.backends.fa_utils import flash_attn_supports_kv_cache_dtype

13 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

@@ -70,13 +70,7 @@ def test_models(

70 if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda():

71 pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.")

73 - if not (

74 - current_platform.is_xpu()

75 - or (

76 - get_flash_attn_version() == 3

77 - and current_platform.is_device_capability_family(90)

78 - )

79 - ):

73 + if not flash_attn_supports_kv_cache_dtype(kv_cache_dtype):

74 pytest.skip(

75 f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention."

76 )

tests/v1/attention/test_attention_backends.py

+64 -7

@@ -21,6 +21,7 @@ from vllm.platforms import current_platform

21 from vllm.utils.math_utils import cdiv

22 from vllm.utils.torch_utils import (

23 STR_DTYPE_TO_TORCH_DTYPE,

24 + is_quantized_kv_cache,

25 is_torch_equal_or_newer,

26 set_random_seed,

27 )

@@ -45,6 +46,11 @@ BACKENDS_TO_TEST = [

47 DEVICE_TYPE = current_platform.device_type

49 +FP8_KV_CACHE_DTYPES = {

50 + "fp8": torch.float8_e4m3fn,

51 + "fp8_e4m3": torch.float8_e4m3fn,

52 +}

53 +

54 # Remove flashinfer from the list if it's not available

55 try:

56 import flashinfer # noqa: F401

@@ -110,6 +116,7 @@ def create_and_prepopulate_kv_cache(

116 num_blocks: int,

117 common_attn_metadata: CommonAttentionMetadata,

118 randomize_blocks: bool = True,

119 + kv_cache_dtype: str = "auto",

120 ) -> torch.Tensor:

121 """Create and prepopulate a KV cache with context data.

122

@@ -140,8 +147,18 @@ def create_and_prepopulate_kv_cache(

147 block_table = common_attn_metadata.block_table_tensor

148 slot_mapping = common_attn_metadata.slot_mapping

149

150 + # For an fp8 kv cache, store the cache in the fp8 dtype so that assigning

151 + # the higher-precision context tensors quantizes them, mirroring runtime.

152 + fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)

153 + storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype

154 +

155 kv_cache = torch.zeros(

144 - num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device

156 + num_blocks,

157 + block_size,

158 + num_kv_heads,

159 + 2 * head_size,

160 + dtype=storage_dtype,

161 + device=device,

162 )

163 kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)

164

@@ -195,7 +212,12 @@ def create_and_prepopulate_kv_cache(

212 ] * block_size + token_inter_block_offsets.to(device)

213

214 # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)

198 - return kv_cache.transpose(1, 2).contiguous()

215 + kv_cache = kv_cache.transpose(1, 2).contiguous()

216 +

217 + if fp8_kv_cache:

218 + kv_cache = kv_cache.view(torch.uint8)

219 +

220 + return kv_cache

221

222

223 class MockAttentionLayer:

@@ -224,6 +246,7 @@ def run_attention_backend(

246 kv_cache: torch.Tensor,

247 attn_type: AttentionType = AttentionType.DECODER,

248 sliding_window: int | None = None,

249 + kv_cache_dtype: str = "auto",

250 ) -> torch.Tensor:

251 """Run attention computation using the specified backend's AttentionImpl."""

252

@@ -291,13 +314,16 @@ def run_attention_backend(

314 alibi_slopes=None,

315 sliding_window=sliding_window,

316 attn_type=attn_type,

294 - kv_cache_dtype="auto",

317 + kv_cache_dtype=kv_cache_dtype,

318 )

319

320 # Create mock layer and output buffer

321 mock_layer = MockAttentionLayer(device)

322 output = torch.empty_like(query)

323

324 + if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input:

325 + query = query.to(current_platform.fp8_dtype())

326 +

327 # Run forward pass

328 # NOTE: The query, key, and value are already shaped correctly

329 # in the calling test function.

@@ -324,6 +350,7 @@ def _test_backend_correctness(

350 atol: float = 1e-2,

351 rtol: float = 1e-2,

352 tensor_parallel_size: int = 1,

353 + kv_cache_dtype: str = "auto",

354 ):

355 """

356 Test that all backends produce similar outputs to a reference implementation

@@ -372,6 +399,7 @@ def _test_backend_correctness(

399 num_gpu_blocks=8192,

400 hf_config_override=hf_config_override,

401 )

402 + vllm_config.cache_config.cache_dtype = kv_cache_dtype

403 device = torch.device(f"{DEVICE_TYPE}:0")

404

405 kv_cache_spec = create_standard_kv_cache_spec(vllm_config, attn_type)

@@ -392,6 +420,13 @@ def _test_backend_correctness(

420 block_size = vllm_config.cache_config.block_size

421 scale = 1.0 / (head_size**0.5)

422

423 + fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)

424 + if fp8_kv_cache:

425 + query_fp8_dtype = current_platform.fp8_dtype()

426 + kv_fp8_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype]

427 + atol = max(atol, 6e-2)

428 + rtol = max(rtol, 1e-1)

429 +

430 # 2. Generate data and compute SDPA reference output

431 all_q_vllm, all_k_vllm, all_v_vllm = [], [], []

432 all_sdpa_outputs = []

@@ -407,10 +442,17 @@ def _test_backend_correctness(

442 k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)

443 v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)

444

445 + if fp8_kv_cache:

446 + q_ref = q.to(query_fp8_dtype).to(dtype)

447 + k_ref = k_full.to(kv_fp8_dtype).to(dtype)

448 + v_ref = v_full.to(kv_fp8_dtype).to(dtype)

449 + else:

450 + q_ref, k_ref, v_ref = q, k_full, v_full

451 +

452 # SDPA expects (N, H, L, D), so unsqueeze batch and permute

411 - q_sdpa_in = q.unsqueeze(0).transpose(1, 2)

412 - k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)

413 - v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)

453 + q_sdpa_in = q_ref.unsqueeze(0).transpose(1, 2)

454 + k_sdpa_in = k_ref.unsqueeze(0).transpose(1, 2)

455 + v_sdpa_in = v_ref.unsqueeze(0).transpose(1, 2)

456

457 if num_q_heads != num_kv_heads:

458 assert num_q_heads % num_kv_heads == 0, (

@@ -471,6 +513,7 @@ def _test_backend_correctness(

513 num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,

514 common_attn_metadata=common_attn_metadata,

515 randomize_blocks=True,

516 + kv_cache_dtype=kv_cache_dtype,

517 )

518

519 # 4. Run vLLM backends and compare

@@ -488,6 +531,12 @@ def _test_backend_correctness(

531 else:

532 backend_cls = None

533

534 + if is_quantized_kv_cache(kv_cache_dtype) and (

535 + backend_cls is None

536 + or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)

537 + ):

538 + continue

539 +

540 if backend_name == AttentionBackendEnum.FLASHINFER:

541 set_kv_cache_layout("HND")

542 reset_kv_cache_layout = True

@@ -521,6 +570,7 @@ def _test_backend_correctness(

570 kv_cache_for_backend,

571 sliding_window=sliding_window,

572 attn_type=attn_type,

573 + kv_cache_dtype=kv_cache_dtype,

574 )

575 finally:

576 if reset_kv_cache_layout:

@@ -570,8 +620,13 @@ def _test_backend_correctness(

620 )

621 @pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])

622 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])

623 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])

624 def test_causal_backend_correctness(

574 - default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int

625 + default_vllm_config,

626 + batch_spec_name: str,

627 + model: str,

628 + tensor_parallel_size: int,

629 + kv_cache_dtype: str,

630 ):

631 """Test backend's correctness with causal attention."""

632

@@ -612,6 +667,7 @@ def test_causal_backend_correctness(

667 SMALL_BLOCK_BACKENDS,

668 causal_mask_mod,

669 tensor_parallel_size=tensor_parallel_size,

670 + kv_cache_dtype=kv_cache_dtype,

671 )

672

673 # Fast FlexAttention needs to run with block_size=128

@@ -623,6 +679,7 @@ def test_causal_backend_correctness(

679 causal_mask_mod,

680 block_size=128,

681 tensor_parallel_size=tensor_parallel_size,

682 + kv_cache_dtype=kv_cache_dtype,

683 )

684

685

tests/v1/attention/utils.py

@@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import (

33 EncoderOnlyAttentionSpec,

34 FullAttentionSpec,

35 MambaSpec,

36 + get_kv_quant_mode,

37 )

@@ -178,6 +179,7 @@ def create_standard_kv_cache_spec(

179 head_size=vllm_config.model_config.get_head_size(),

180 dtype=vllm_config.model_config.dtype,

181 sliding_window=vllm_config.model_config.get_sliding_window(),

182 + kv_quant_mode=get_kv_quant_mode(vllm_config.cache_config.cache_dtype),

183 )

184

185

tools/pre_commit/generate_attention_backend_docs.py

+66 -7

@@ -1083,7 +1083,8 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:

1083 return {}

1084

1085 # Analyze the functions to determine FA3/FA4-specific features

1086 - fa3_supports_fp8 = True

1086 + fa3_supports_fp8 = False

1087 + fa4_supports_fp8 = False

1088 fa3_supports_sinks = False

1089 fa4_supports_sinks = False

1090 fa3_compute_cap: str | None = None

@@ -1093,6 +1094,44 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:

1094 if not isinstance(node, ast.FunctionDef):

1095 continue

1096

1097 + # Check flash_attn_supports_kv_cache_dtype for fp8 support per FA version.

1098 + # Accept both equality checks and membership checks such as `in (3, 4)`.

1099 + if node.name == "flash_attn_supports_kv_cache_dtype":

1100 + for n in ast.walk(node):

1101 + if not (

1102 + isinstance(n, ast.Compare)

1103 + and len(n.ops) == 1

1104 + and len(n.comparators) == 1

1105 + ):

1106 + continue

1107 + is_version_compare = (

1108 + isinstance(n.left, ast.Name) and n.left.id == "fa_version"

1109 + ) or (

1110 + isinstance(n.left, ast.Call)

1111 + and isinstance(n.left.func, ast.Name)

1112 + and n.left.func.id == "get_flash_attn_version"

1113 + )

1114 + if not is_version_compare:

1115 + continue

1116 +

1117 + versions: list[Any] = []

1118 + comparator = n.comparators[0]

1119 + if isinstance(n.ops[0], ast.Eq) and isinstance(

1120 + comparator, ast.Constant

1121 + ):

1122 + versions = [comparator.value]

1123 + elif isinstance(n.ops[0], ast.In) and isinstance(

1124 + comparator, (ast.Tuple, ast.List, ast.Set)

1125 + ):

1126 + versions = [

1127 + elt.value

1128 + for elt in comparator.elts

1129 + if isinstance(elt, ast.Constant)

1130 + ]

1131 +

1132 + fa3_supports_fp8 |= 3 in versions

1133 + fa4_supports_fp8 |= 4 in versions

1134 +

1135 # Check flash_attn_supports_sinks - looks for `fa_version == 3/4`

1136 # or `get_flash_attn_version() == 3/4` (also accepts `in (3, 4)`)

1137 if node.name == "flash_attn_supports_sinks":

@@ -1199,7 +1238,7 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:

1238 },

1239 "fa4": {

1240 "compute_capability": fa4_compute_cap,

1202 - "supports_fp8": False,

1241 + "supports_fp8": fa4_supports_fp8,

1242 "supports_sink": fa4_supports_sinks,

1243 },

1244 }

@@ -1255,6 +1294,22 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:

1294 # ---------------------------------------------------------------------------

1295

1296

1297 +def _apply_fp8_support(kv_cache_dtypes: str, supports_fp8: bool) -> str:

1298 + """Add or remove fp8 dtypes from a comma-separated kv_cache_dtypes string.

1299 +

1300 + The base FLASH_ATTN backend lists fp8 dtypes in ``supported_kv_cache_dtypes``,

1301 + but actual support varies by FA version (e.g. FA2 has no fp8 support), so each

1302 + variant's row is normalized to match its own capability.

1303 + """

1304 + fp8_dtypes = {"fp8", "fp8_e4m3", "fp8_e5m2"}

1305 + dtypes = kv_cache_dtypes.split(", ")

1306 + base_dtypes = [dtype for dtype in dtypes if dtype not in fp8_dtypes]

1307 + if supports_fp8:

1308 + advertised_fp8_dtypes = [dtype for dtype in dtypes if dtype in fp8_dtypes]

1309 + return ", ".join(base_dtypes + advertised_fp8_dtypes)

1310 + return ", ".join(base_dtypes)

1311 +

1312 +

1313 def _expand_flash_attn_variants(

1314 all_backends: list[dict[str, Any]],

1315 fa_features: dict[str, dict[str, Any]],

@@ -1275,6 +1330,9 @@ def _expand_flash_attn_variants(

1330 fa2["_sort_key"] = "FLASH_ATTN"

1331 fa2["_sort_order"] = 0

1332 fa2["supports_sink"] = fa_features["fa2"]["supports_sink"]

1333 + fa2["kv_cache_dtypes"] = _apply_fp8_support(

1334 + backend["kv_cache_dtypes"], fa_features["fa2"]["supports_fp8"]

1335 + )

1336

1337 # Create FA3 entry (uses parsed compute_capability from fa_utils)

1338 fa3 = backend.copy()

@@ -1284,11 +1342,9 @@ def _expand_flash_attn_variants(

1342 if fa_features["fa3"]["compute_capability"]:

1343 fa3["compute_capability"] = fa_features["fa3"]["compute_capability"]

1344 fa3["supports_sink"] = fa_features["fa3"]["supports_sink"]

1287 - if fa_features["fa3"]["supports_fp8"]:

1288 - base_dtypes = backend["kv_cache_dtypes"].split(", ")

1289 - fp8_dtypes = ["fp8", "fp8_e4m3", "fp8_e5m2"]

1290 - new_dtypes = [d for d in fp8_dtypes if d not in base_dtypes]

1291 - fa3["kv_cache_dtypes"] = ", ".join(base_dtypes + new_dtypes)

1345 + fa3["kv_cache_dtypes"] = _apply_fp8_support(

1346 + backend["kv_cache_dtypes"], fa_features["fa3"]["supports_fp8"]

1347 + )

1348

1349 expanded.append(fa2)

1350 expanded.append(fa3)

@@ -1302,6 +1358,9 @@ def _expand_flash_attn_variants(

1358 if fa_features["fa4"].get("compute_capability"):

1359 fa4["compute_capability"] = fa_features["fa4"]["compute_capability"]

1360 fa4["supports_sink"] = fa_features["fa4"]["supports_sink"]

1361 + fa4["kv_cache_dtypes"] = _apply_fp8_support(

1362 + backend["kv_cache_dtypes"], fa_features["fa4"]["supports_fp8"]

1363 + )

1364 expanded.append(fa4)

1365

1366 return expanded

vllm/v1/attention/backends/fa_utils.py

+23

@@ -273,6 +273,29 @@ def is_fa_version_supported(fa_version: int) -> bool:

273 return False

274

275

276 +def flash_attn_supports_kv_cache_dtype(

277 + kv_cache_dtype: str = "fp8_e4m3",

278 + *,

279 + requires_alibi: bool = False,

280 + head_size: int | None = None,

281 + head_size_v: int | None = None,

282 + has_sinks: bool = False,

283 +) -> bool:

284 + if kv_cache_dtype == "fp8_e5m2":

285 + return False

286 + if current_platform.is_xpu():

287 + return True

288 + fa_version = get_flash_attn_version(

289 + requires_alibi=requires_alibi,

290 + head_size=head_size,

291 + head_size_v=head_size_v,

292 + has_sinks=has_sinks,

293 + )

294 + return (fa_version == 3 and current_platform.is_device_capability_family(90)) or (

295 + fa_version == 4 and current_platform.is_device_capability_family(100)

296 + )

297 +

298 +

299 def flash_attn_supports_quant_query_input() -> bool:

300 return not current_platform.is_xpu()

301

vllm/v1/attention/backends/flash_attn.py

+34 -8

@@ -22,6 +22,7 @@ from vllm.v1.attention.backend import (

22 MultipleOf,

23 )

24 from vllm.v1.attention.backends.fa_utils import (

25 + flash_attn_supports_kv_cache_dtype,

26 flash_attn_supports_quant_query_input,

27 get_flash_attn_version,

28 is_fa_version_supported,

@@ -74,6 +75,8 @@ class FlashAttentionBackend(AttentionBackend):

75 "auto",

76 "float16",

77 "bfloat16",

78 + "fp8",

79 + "fp8_e4m3",

80 ]

82 @staticmethod

@@ -178,14 +181,11 @@ class FlashAttentionBackend(AttentionBackend):

181 def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:

182 if kv_cache_dtype is None:

183 return True

181 - if kv_cache_dtype in ("fp8", "fp8_e4m3"):

182 - if current_platform.is_xpu():

183 - return True

184 - return (

185 - get_flash_attn_version() == 3

186 - and current_platform.is_device_capability_family(90)

187 - )

188 - return kv_cache_dtype in ["auto", "float16", "bfloat16"]

184 + if kv_cache_dtype not in cls.supported_kv_cache_dtypes:

185 + return False

186 + if is_quantized_kv_cache(kv_cache_dtype):

187 + return flash_attn_supports_kv_cache_dtype(kv_cache_dtype)

188 + return True

189

190 @classmethod

191 def supports_mm_prefix(cls) -> bool:

@@ -216,6 +216,17 @@ class FlashAttentionBackend(AttentionBackend):

216 ) -> str | None:

217 if has_sink and device_capability < DeviceCapability(9, 0):

218 return "sink not supported on compute capability < 9.0"

219 + if (

220 + kv_cache_dtype is not None

221 + and is_quantized_kv_cache(kv_cache_dtype)

222 + and not flash_attn_supports_kv_cache_dtype(

223 + kv_cache_dtype,

224 + head_size=head_size,

225 + head_size_v=head_size,

226 + has_sinks=has_sink,

227 + )

228 + ):

229 + return "FP8 KV cache requires FA3 on SM90 or FA4 on SM100"

230 if (

231 use_mm_prefix

232 and get_flash_attn_version(head_size=head_size, has_sinks=has_sink) != 4

@@ -772,6 +783,7 @@ class FlashAttentionImpl(AttentionImpl):

783 self.vllm_flash_attn_version = get_flash_attn_version(

784 requires_alibi=alibi_slopes is not None,

785 head_size=head_size,

786 + has_sinks=sinks is not None,

787 )

788 logger.info_once(

789 "Using FlashAttention version %s",

@@ -780,6 +792,20 @@ class FlashAttentionImpl(AttentionImpl):

792 # Cache the batch invariant result for use in forward passes

793 self.batch_invariant_enabled = envs.VLLM_BATCH_INVARIANT

794

795 + if is_quantized_kv_cache(

796 + self.kv_cache_dtype

797 + ) and not flash_attn_supports_kv_cache_dtype(

798 + self.kv_cache_dtype,

799 + requires_alibi=alibi_slopes is not None,

800 + head_size=head_size,

801 + head_size_v=head_size,

802 + has_sinks=sinks is not None,

803 + ):

804 + raise NotImplementedError(

805 + f"FlashAttention does not support {self.kv_cache_dtype}"

806 + " kv-cache on this device."

807 + )

808 +

809 self.sinks = sinks

810 if self.sinks is not None:

811 assert flash_attn_supports_sinks(), (