[SM100] Hdim 256 optimized (#52980)
Signed-off-by: Simon Veitner <sveitner@redhat.com>
simon-veitner-redhat committed
Aug 26, 2026 at 05:28 UTC
7156c63bef1f80a2edcc67901fa7ebb3892b4f3d
7 files changed
+423
-22
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 f3e1a4f74c99145c0717709860bf765de1703779
42
+ GIT_TAG 4eb29359a51b4fe67c8c0f0c5e931236971aa492
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
+8
@@ -127,6 +127,14 @@ Priority is **1 = highest** (tried first).
127
> **†** FlashInfer Native is the regular FlashInfer path. XQA is the SM90 decode path exposed through FlashInfer's TRTLLM decode API. trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen via `--attention-config.use_trtllm_attention=0`.
128
>
129
> **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
130
+>
131
+> On Blackwell, when the FlashAttention backend is selected, `head_size=256` is served
132
+> by a dedicated FA4 kernel that requires a KV cache block size of 128 (advertised
133
+> automatically, so it is picked as long as `--block-size` is not pinned) and does not
134
+> support logit soft capping, attention sinks, mm_prefix/R-SWA masking, DCP, or windowed
135
+> encoder attention. Those configurations transparently fall back to FA2. A pinned
136
+> `--block-size` that is not a multiple of 128 instead makes FlashAttention ineligible
137
+> for such models, which is an error if the backend was requested explicitly.
138
139
## MiniMax M3 Sparse Attention Backends
140
tests/kernels/attention/test_attention_selector.py
+176
-1
@@ -1,7 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from unittest.mock import patch
4
+from contextlib import contextmanager
5
+from unittest.mock import MagicMock, patch
6
7
import pytest
8
import torch
@@ -14,6 +15,7 @@ from vllm.config import (
15
)
16
from vllm.platforms import current_platform
17
from vllm.platforms.cpu import CpuPlatform
18
+from vllm.platforms.interface import DeviceCapability
19
20
if current_platform.is_cuda():
21
from vllm.platforms.cuda import CudaPlatform
@@ -25,6 +27,7 @@ if current_platform.is_rocm():
27
else:
28
RocmPlatform = None
29
30
+from vllm.v1.attention.backend import AttentionType
31
from vllm.v1.attention.backends.registry import AttentionBackendEnum
32
from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend
33
@@ -555,3 +558,175 @@ def test_flash_attn_accepts_handled_fp8_variants(
558
# import order across earlier tests that patch vllm.platforms.current_platform.
559
monkeypatch.setattr(fa_utils_mod.current_platform, "is_xpu", lambda: True)
560
assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
561
+
562
+
563
+blackwell_only = pytest.mark.skipif(
564
+ not current_platform.is_cuda(), reason="FA4 is CUDA-only"
565
+)
566
+
567
+
568
+@contextmanager
569
+def _blackwell(vllm_config=None):
570
+ platform = MagicMock()
571
+ platform.is_xpu.return_value = False
572
+ platform.is_rocm.return_value = False
573
+ platform.get_device_capability.return_value = DeviceCapability(10, 0)
574
+ with (
575
+ patch("vllm.v1.attention.backends.fa_utils.current_platform", platform),
576
+ patch(
577
+ "vllm.vllm_flash_attn.flash_attn_interface.is_fa_version_supported",
578
+ return_value=True,
579
+ ),
580
+ patch("vllm.config.get_current_vllm_config_or_none", return_value=vllm_config),
581
+ patch(
582
+ "vllm.v1.attention.backends.flash_attn.get_current_vllm_config_or_none",
583
+ return_value=vllm_config,
584
+ ),
585
+ ):
586
+ yield
587
+
588
+
589
+def _hd256_config(
590
+ *,
591
+ is_mm_prefix_lm=False,
592
+ rswa_window=None,
593
+ dcp_size=1,
594
+ softcap=None,
595
+ head_size=256,
596
+ cache_dtype="auto",
597
+):
598
+ vllm_config = MagicMock()
599
+ vllm_config.attention_config.flash_attn_version = None
600
+ vllm_config.model_config.is_mm_prefix_lm = is_mm_prefix_lm
601
+ vllm_config.model_config.rswa_window = rswa_window
602
+ vllm_config.model_config.hf_text_config.attn_logit_softcapping = softcap
603
+ vllm_config.model_config.get_head_size.return_value = head_size
604
+ vllm_config.cache_config.cache_dtype = cache_dtype
605
+ vllm_config.parallel_config.decode_context_parallel_size = dcp_size
606
+ return vllm_config
607
+
608
+
609
+@blackwell_only
610
+@pytest.mark.parametrize(
611
+ "kwargs,config_kwargs,expected",
612
+ [
613
+ ({}, {}, 4),
614
+ ({"supports_fa4_hd256": False}, {}, 2),
615
+ ({"has_sinks": True}, {}, 2),
616
+ ({"requires_softcap": True}, {}, 2),
617
+ # Larger block sizes normalize to 128.
618
+ ({"kv_cache_block_size": 16}, {}, 2),
619
+ ({"kv_cache_block_size": 64}, {}, 2),
620
+ ({"kv_cache_block_size": 256}, {}, 4),
621
+ ({}, {"is_mm_prefix_lm": True}, 2),
622
+ ({}, {"rswa_window": 512}, 2),
623
+ ({}, {"dcp_size": 2}, 2),
624
+ ({}, {"softcap": 50.0}, 2),
625
+ ({}, {"cache_dtype": "fp8"}, 2),
626
+ ({"head_size": 128, "kv_cache_block_size": 16}, {}, 4),
627
+ ({"head_size": 192, "head_size_v": 128, "kv_cache_block_size": 16}, {}, 4),
628
+ ({"head_size": 256, "head_size_v": 128, "kv_cache_block_size": 16}, {}, 2),
629
+ ({"head_size": 256, "head_size_v": 64}, {}, 2),
630
+ ],
631
+)
632
+def test_fa4_hd256_fallback_matrix(kwargs, config_kwargs, expected):
633
+ from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
634
+
635
+ kwargs = {
636
+ "head_size": 256,
637
+ "kv_cache_block_size": 128,
638
+ "supports_fa4_hd256": True,
639
+ **kwargs,
640
+ }
641
+ with _blackwell(_hd256_config(head_size=kwargs["head_size"], **config_kwargs)):
642
+ assert get_flash_attn_version(**kwargs) == expected
643
+
644
+
645
+@blackwell_only
646
+@pytest.mark.parametrize(
647
+ "backend_name,config_kwargs,expected",
648
+ [
649
+ ("FLASH_ATTN", {}, 128),
650
+ ("FLASH_ATTN", {"softcap": 50.0}, 16),
651
+ ("FLASH_ATTN", {"head_size": 128}, 16),
652
+ ("FLASH_ATTN", {"cache_dtype": "fp8"}, 16),
653
+ ("FLASH_ATTN_DIFFKV", {}, 16),
654
+ ],
655
+)
656
+def test_fa4_hd256_block_size_advertisement(
657
+ backend_name: str, config_kwargs: dict, expected: int
658
+):
659
+ from vllm.v1.attention.backend import MultipleOf
660
+ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
661
+ from vllm.v1.attention.backends.flash_attn_diffkv import (
662
+ FlashAttentionDiffKVBackend,
663
+ )
664
+
665
+ backend = {
666
+ "FLASH_ATTN": FlashAttentionBackend,
667
+ "FLASH_ATTN_DIFFKV": FlashAttentionDiffKVBackend,
668
+ }[backend_name]
669
+ vllm_config = _hd256_config(**config_kwargs)
670
+ with _blackwell(vllm_config):
671
+ (size,) = backend.get_supported_kernel_block_sizes()
672
+ preferred = backend.get_preferred_block_size(16)
673
+ assert preferred == expected
674
+ if expected == 128:
675
+ assert size == 128
676
+ else:
677
+ assert isinstance(size, MultipleOf) and size.base == expected
678
+
679
+
680
+@blackwell_only
681
+def test_fa4_hd256_mm_prefix_deselects_flash_attn():
682
+ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
683
+
684
+ with _blackwell(_hd256_config(is_mm_prefix_lm=True)):
685
+ assert (
686
+ FlashAttentionBackend.supports_combination(
687
+ head_size=256,
688
+ dtype=torch.bfloat16,
689
+ kv_cache_dtype=None,
690
+ block_size=128,
691
+ use_mla=False,
692
+ has_sink=False,
693
+ use_sparse=False,
694
+ use_mm_prefix=True,
695
+ device_capability=DeviceCapability(10, 0),
696
+ )
697
+ is not None
698
+ )
699
+
700
+
701
+@pytest.mark.skipif(
702
+ not current_platform.is_cuda()
703
+ or not current_platform.is_device_capability_family(100),
704
+ reason="requires a Blackwell GPU",
705
+)
706
+@pytest.mark.parametrize(
707
+ "attn_type,sliding_window,expected",
708
+ [
709
+ (AttentionType.DECODER, None, 4),
710
+ (AttentionType.DECODER, 512, 4),
711
+ # Local encoder attention lacks the required seqused tensors.
712
+ (AttentionType.ENCODER_ONLY, None, 4),
713
+ (AttentionType.ENCODER_ONLY, 512, 2),
714
+ ],
715
+)
716
+def test_fa4_hd256_impl_selection(attn_type, sliding_window, expected):
717
+ from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl
718
+
719
+ # Implementations are constructed before the platform updates the block size.
720
+ with set_current_vllm_config(VllmConfig()):
721
+ impl = FlashAttentionImpl(
722
+ num_heads=8,
723
+ head_size=256,
724
+ scale=0.0625,
725
+ num_kv_heads=8,
726
+ alibi_slopes=None,
727
+ sliding_window=sliding_window,
728
+ kv_cache_dtype="auto",
729
+ attn_type=attn_type,
730
+ )
731
+ assert impl.vllm_flash_attn_version == expected
732
+ assert impl.fa4_hd256 == (expected == 4)
tests/kernels/attention/test_flash_attn.py
+75
@@ -215,3 +215,78 @@ def test_varlen_with_paged_kv(
215
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
216
f"{torch.max(torch.abs(output - ref_output))}",
217
)
218
+
219
+
220
+@pytest.mark.skipif(
221
+ not current_platform.is_device_capability_family(100),
222
+ reason="FA4 only has a dedicated head_size=256 kernel on Blackwell",
223
+)
224
+@pytest.mark.parametrize(
225
+ "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
226
+)
227
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
228
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
229
+@torch.inference_mode()
230
+def test_fa4_hd256_paged_call_shape(
231
+ seq_lens: list[tuple[int, int]],
232
+ num_heads: tuple[int, int],
233
+ sliding_window: int | None,
234
+) -> None:
235
+ """Handles excess block-table columns with page-aligned ``max_seqlen_k``."""
236
+ from vllm.v1.attention.backends.fa_utils import FA4_HD256_PAGE_SIZE
237
+
238
+ if not is_fa_version_supported(4):
239
+ pytest.skip(f'FA4 unsupported: "{fa_version_unsupported_reason(4)}"')
240
+ torch.set_default_device("cuda")
241
+ set_random_seed(0)
242
+
243
+ head_size = 256
244
+ block_size = FA4_HD256_PAGE_SIZE
245
+ num_blocks = 2048
246
+ dtype = torch.bfloat16
247
+ num_query_heads, num_kv_heads = num_heads
248
+ query_lens = [x[0] for x in seq_lens]
249
+ kv_lens = [x[1] for x in seq_lens]
250
+ max_kv_len = max(kv_lens)
251
+ window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
252
+ scale = head_size**-0.5
253
+
254
+ query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
255
+ key_cache = torch.randn(
256
+ num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
257
+ )
258
+ value_cache = torch.randn_like(key_cache)
259
+ cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
260
+ dim=0, dtype=torch.int32
261
+ )
262
+ num_pages = (max_kv_len + block_size - 1) // block_size
263
+ block_tables = torch.randint(
264
+ 0, num_blocks, (len(seq_lens), num_pages + 3), dtype=torch.int32
265
+ )
266
+
267
+ output = flash_attn_varlen_func(
268
+ q=query,
269
+ k=key_cache,
270
+ v=value_cache,
271
+ cu_seqlens_q=cu_query_lens,
272
+ seqused_k=torch.tensor(kv_lens, dtype=torch.int32),
273
+ max_seqlen_q=max(query_lens),
274
+ max_seqlen_k=num_pages * block_size,
275
+ softmax_scale=scale,
276
+ causal=True,
277
+ window_size=window_size,
278
+ block_table=block_tables[:, :num_pages],
279
+ num_splits=1,
280
+ fa_version=4,
281
+ )
282
+ ref_output = ref_paged_attn(
283
+ query=query,
284
+ key_cache=key_cache,
285
+ value_cache=value_cache,
286
+ query_lens=query_lens,
287
+ kv_lens=kv_lens,
288
+ block_tables=block_tables,
289
+ scale=scale,
290
+ sliding_window=sliding_window,
291
+ )
292
+ torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
vllm/v1/attention/backends/fa_utils.py
+79
-14
@@ -6,9 +6,13 @@ from typing import Any
6
import vllm.envs as envs
7
from vllm.logger import init_logger
8
from vllm.platforms import current_platform
9
+from vllm.utils.torch_utils import is_quantized_kv_cache
10
11
logger = init_logger(__name__)
12
13
+# The dedicated hd256 kernel requires one 128-token page per tile.
14
+FA4_HD256_PAGE_SIZE = 128
15
+
16
# Track whether upstream flash-attn is available on ROCm.
17
# Set during module initialization and never modified afterwards.
18
# This module-level flag avoids repeated import attempts and ensures
@@ -72,6 +76,9 @@ def get_flash_attn_version(
76
head_size: int | None = None,
77
head_size_v: int | None = None,
78
has_sinks: bool = False,
79
+ requires_softcap: bool = False,
80
+ kv_cache_block_size: int | None = None,
81
+ supports_fa4_hd256: bool = False,
82
) -> int | None:
83
if current_platform.is_xpu():
84
return 2
@@ -168,26 +175,31 @@ def get_flash_attn_version(
175
)
176
fa_version = 2
177
171
- # TODO: Restore the `requires_local_attention` restriction when FA4
172
- # head-dim 256 is re-enabled.
173
- if fa_version == 4 and device_capability.major >= 10 and head_size == 256:
174
- logger.warning_once(
175
- "FA4 on Blackwell is temporarily disabled for head_size=256, "
176
- "defaulting to FA version 2."
177
- )
178
- fa_version = 2
178
+ if fa_version == 4 and uses_fa4_hd256_kernel(head_size, head_size_v):
179
+ if not supports_fa4_hd256:
180
+ fa_version = 2
181
+ elif (
182
+ reason := _fa4_hd256_fallback_reason(
183
+ has_sinks, requires_softcap, kv_cache_block_size, vllm_config
184
+ )
185
+ ) is not None:
186
+ logger.warning_once(
187
+ "FA4's Blackwell head_size=256 kernel does not support %s, "
188
+ "defaulting to FA version 2.",
189
+ reason,
190
+ )
191
+ fa_version = 2
192
180
- # FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict
181
- # supported head dimensions to ≤128. The 192/128 MLA prefill case is
182
- # supported; 256 is temporarily disabled until upstream supports the
183
- # required features. Development of symmetric 192, 384, and 512 support
184
- # is tracked in https://github.com/Dao-AILab/flash-attention/issues/2456
193
+ # FA4 head dimensions on Blackwell are limited by TMEM capacity.
194
if (
195
fa_version == 4
196
and device_capability.major >= 10
197
and head_size is not None
198
and head_size > 128
190
- and not (head_size == 192 and head_size_v == 128)
199
+ and not (
200
+ (head_size == 256 and head_size_v in (None, 256))
201
+ or (head_size == 192 and head_size_v == 128)
202
+ )
203
):
204
logger.warning_once(
205
"FA4 on Blackwell does not support head_size=%d due to TMEM "
@@ -209,6 +221,53 @@ def get_flash_attn_version(
221
return None
222
223
224
+def uses_fa4_hd256_kernel(
225
+ head_size: int | None, head_size_v: int | None = None
226
+) -> bool:
227
+ """Return whether FA4 uses its dedicated hd256 kernel."""
228
+ if head_size != 256:
229
+ return False
230
+ if head_size_v is not None and head_size_v != 256:
231
+ return False
232
+ capability = current_platform.get_device_capability()
233
+ return capability is not None and capability.major in (10, 11)
234
+
235
+
236
+def _fa4_hd256_fallback_reason(
237
+ has_sinks: bool,
238
+ requires_softcap: bool,
239
+ kv_cache_block_size: int | None,
240
+ vllm_config: Any,
241
+) -> str | None:
242
+ model_config = vllm_config.model_config if vllm_config is not None else None
243
+ cache_config = vllm_config.cache_config if vllm_config is not None else None
244
+ if has_sinks:
245
+ return "attention sinks"
246
+ if requires_softcap or (
247
+ # Keep model-level and per-layer version selection consistent.
248
+ model_config is not None
249
+ and getattr(model_config.hf_text_config, "attn_logit_softcapping", None)
250
+ ):
251
+ return "logits soft capping"
252
+ if cache_config is not None and is_quantized_kv_cache(cache_config.cache_dtype):
253
+ return f"quantized KV cache dtype {cache_config.cache_dtype}"
254
+ if kv_cache_block_size is not None and kv_cache_block_size % FA4_HD256_PAGE_SIZE:
255
+ # Larger blocks are split into 128-token kernel pages.
256
+ return f"a KV cache block size of {kv_cache_block_size}"
257
+ if model_config is not None:
258
+ if model_config.is_mm_prefix_lm:
259
+ return "mm_prefix bidirectional attention"
260
+ if model_config.rswa_window is not None:
261
+ return "R-SWA"
262
+ if (
263
+ vllm_config is not None
264
+ and vllm_config.parallel_config.decode_context_parallel_size > 1
265
+ ):
266
+ # DCP reuses an unsliced block table with a non-page-aligned length.
267
+ return "decode context parallelism"
268
+ return None
269
+
270
+
271
def is_fa_version_supported(fa_version: int) -> bool:
272
try:
273
from vllm.vllm_flash_attn.flash_attn_interface import (
@@ -227,6 +286,9 @@ def flash_attn_supports_kv_cache_dtype(
286
head_size: int | None = None,
287
head_size_v: int | None = None,
288
has_sinks: bool = False,
289
+ requires_softcap: bool = False,
290
+ kv_cache_block_size: int | None = None,
291
+ supports_fa4_hd256: bool = False,
292
) -> bool:
293
if kv_cache_dtype == "fp8_e5m2":
294
return False
@@ -237,6 +299,9 @@ def flash_attn_supports_kv_cache_dtype(
299
head_size=head_size,
300
head_size_v=head_size_v,
301
has_sinks=has_sinks,
302
+ requires_softcap=requires_softcap,
303
+ kv_cache_block_size=kv_cache_block_size,
304
+ supports_fa4_hd256=supports_fa4_hd256,
305
)
306
return (fa_version == 3 and current_platform.is_device_capability_family(90)) or (
307
fa_version == 4 and current_platform.is_device_capability_family(100)
vllm/v1/attention/backends/flash_attn.py
+83
-6
@@ -24,11 +24,13 @@ from vllm.v1.attention.backend import (
24
MultipleOf,
25
)
26
from vllm.v1.attention.backends.fa_utils import (
27
+ FA4_HD256_PAGE_SIZE,
28
flash_attn_supports_kv_cache_dtype,
29
flash_attn_supports_quant_query_input,
30
get_flash_attn_version,
31
is_fa_version_supported,
32
is_flash_attn_varlen_func_available,
33
+ uses_fa4_hd256_kernel,
34
)
35
from vllm.v1.attention.backends.utils import (
36
fill_mm_prefix_query_ranges,
@@ -85,15 +87,40 @@ class FlashAttentionBackend(AttentionBackend):
87
"fp8",
88
"fp8_e4m3",
89
]
90
+ head_size_v: int | None = None
91
89
- @staticmethod
90
- def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
92
+ @classmethod
93
+ def _get_fa4_hd256_block_size(cls) -> int | None:
94
+ vllm_config = get_current_vllm_config_or_none()
95
+ if vllm_config is None or vllm_config.model_config is None:
96
+ return None
97
+
98
+ head_size = vllm_config.model_config.get_head_size()
99
+ if (
100
+ uses_fa4_hd256_kernel(head_size, cls.head_size_v)
101
+ and get_flash_attn_version(
102
+ head_size=head_size,
103
+ head_size_v=cls.head_size_v,
104
+ supports_fa4_hd256=True,
105
+ )
106
+ == 4
107
+ ):
108
+ return FA4_HD256_PAGE_SIZE
109
+ return None
110
+
111
+ @classmethod
112
+ def get_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]:
113
+ if block_size := cls._get_fa4_hd256_block_size():
114
+ # Sliding-window specs select the smallest advertised size.
115
+ return [block_size]
116
return [MultipleOf(16)]
117
118
forward_includes_kv_cache_update: bool = False
119
120
@classmethod
121
def get_preferred_block_size(cls, default_block_size: int) -> int:
122
+ if block_size := cls._get_fa4_hd256_block_size():
123
+ return max(default_block_size, block_size)
124
if current_platform.is_xpu():
125
return max(default_block_size, 64)
126
return super().get_preferred_block_size(default_block_size)
@@ -194,12 +221,20 @@ class FlashAttentionBackend(AttentionBackend):
221
head_size=head_size,
222
head_size_v=head_size,
223
has_sinks=has_sink,
224
+ kv_cache_block_size=block_size,
225
+ supports_fa4_hd256=True,
226
)
227
):
228
return "FP8 KV cache requires FA3 on SM90 or FA4 on SM100"
229
if (
230
use_mm_prefix
202
- and get_flash_attn_version(head_size=head_size, has_sinks=has_sink) != 4
231
+ and get_flash_attn_version(
232
+ head_size=head_size,
233
+ has_sinks=has_sink,
234
+ kv_cache_block_size=block_size,
235
+ supports_fa4_hd256=True,
236
+ )
237
+ != 4
238
):
239
return (
240
"mm_prefix (PrefixLM bidirectional attention) requires "
@@ -409,6 +444,15 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
444
self.max_num_splits = 0 # No upper bound on the number of splits.
445
self.aot_schedule = get_flash_attn_version() == 3
446
447
+ self.fa4_hd256 = uses_fa4_hd256_kernel(self.headdim) and (
448
+ get_flash_attn_version(
449
+ head_size=self.headdim,
450
+ kv_cache_block_size=self.block_size,
451
+ supports_fa4_hd256=True,
452
+ )
453
+ == 4
454
+ )
455
+
456
try:
457
from vllm.distributed.parallel_state import get_dcp_group
458
@@ -780,6 +824,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
824
metadata.scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)
825
826
def use_cascade_attention(self, *args, **kwargs) -> bool:
827
+ if self.fa4_hd256:
828
+ # Cascade may use a non-page-aligned prefix length.
829
+ return False
830
return use_cascade_attention(*args, **kwargs)
831
832
@@ -823,11 +870,30 @@ class FlashAttentionImpl(AttentionImpl):
870
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
871
872
self.attn_type = attn_type
873
+ vllm_config = get_current_vllm_config_or_none()
874
+ uses_kv_cache = attn_type not in (
875
+ AttentionType.ENCODER,
876
+ AttentionType.ENCODER_ONLY,
877
+ )
878
+ # The final KV cache block size is unavailable during construction.
879
self.vllm_flash_attn_version = get_flash_attn_version(
880
requires_alibi=alibi_slopes is not None,
881
head_size=head_size,
882
has_sinks=sinks is not None,
883
+ requires_softcap=bool(self.logits_soft_cap),
884
+ supports_fa4_hd256=True,
885
+ )
886
+ self.fa4_hd256 = self.vllm_flash_attn_version == 4 and uses_fa4_hd256_kernel(
887
+ head_size
888
)
889
+ if self.fa4_hd256 and not uses_kv_cache and sliding_window is not None:
890
+ # The hd256 kernel requires seqused_k for local attention.
891
+ logger.warning_once(
892
+ "FA4's Blackwell head_size=256 kernel does not support local "
893
+ "attention on encoder inputs, defaulting to FA version 2."
894
+ )
895
+ self.vllm_flash_attn_version = 2
896
+ self.fa4_hd256 = False
897
logger.info_once(
898
"Using FlashAttention version %s",
899
self.vllm_flash_attn_version,
@@ -843,6 +909,8 @@ class FlashAttentionImpl(AttentionImpl):
909
head_size=head_size,
910
head_size_v=head_size,
911
has_sinks=sinks is not None,
912
+ requires_softcap=bool(self.logits_soft_cap),
913
+ supports_fa4_hd256=True,
914
):
915
raise NotImplementedError(
916
f"FlashAttention does not support {self.kv_cache_dtype}"
@@ -861,7 +929,6 @@ class FlashAttentionImpl(AttentionImpl):
929
930
self.supports_quant_query_input = flash_attn_supports_quant_query_input()
931
864
- vllm_config = get_current_vllm_config_or_none()
932
dcp_a2a = (
933
vllm_config is not None
934
and vllm_config.parallel_config.decode_context_parallel_size > 1
@@ -1083,6 +1150,15 @@ class FlashAttentionImpl(AttentionImpl):
1150
)
1151
causal = not has_window
1152
1153
+ num_splits = attn_metadata.max_num_splits
1154
+ if self.fa4_hd256:
1155
+ # hd256 requires page-aligned lengths, exact-width block
1156
+ # tables, and no SplitKV.
1157
+ num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE)
1158
+ max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE
1159
+ block_table = block_table[:, :num_pages]
1160
+ num_splits = 1
1161
+
1162
flash_attn_varlen_func(
1163
q=query[:num_actual_tokens],
1164
k=key_cache,
@@ -1104,7 +1180,7 @@ class FlashAttentionImpl(AttentionImpl):
1180
k_descale=k_descale,
1181
v_descale=v_descale,
1182
dynamic_causal=dynamic_causal,
1107
- num_splits=attn_metadata.max_num_splits,
1183
+ num_splits=num_splits,
1184
s_aux=self.sinks,
1185
mask_mod=rswa_mask_mod_fn or mm_mask_mod,
1186
aux_tensors=rswa_aux or mm_aux,
@@ -1416,7 +1492,8 @@ class FlashAttentionImpl(AttentionImpl):
1492
else None,
1493
k_descale=layer._k_scale.expand(descale_shape), # type: ignore[operator]
1494
v_descale=layer._v_scale.expand(descale_shape), # type: ignore[operator]
1419
- num_splits=1 if self.batch_invariant_enabled else 0,
1495
+ # The hd256 kernel does not support SplitKV.
1496
+ num_splits=1 if self.batch_invariant_enabled or self.fa4_hd256 else 0,
1497
s_aux=self.sinks,
1498
)
1499
vllm/v1/attention/backends/flash_attn_diffkv.py
+1
@@ -87,6 +87,7 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):
87
head_size_v=FlashAttentionDiffKVBackend.head_size_v,
88
has_sinks=self.sinks is not None,
89
)
90
+ self.fa4_hd256 = False
91
92
def do_kv_cache_update(
93
self,