diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 2a8e9cc781..2170c5c33c 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 f3e1a4f74c99145c0717709860bf765de1703779 + GIT_TAG 4eb29359a51b4fe67c8c0f0c5e931236971aa492 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 c4cedd9e4e..f0f363223f 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -127,6 +127,14 @@ Priority is **1 = highest** (tried first). > **†** 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`. > > **\*** 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. +> +> On Blackwell, when the FlashAttention backend is selected, `head_size=256` is served +> by a dedicated FA4 kernel that requires a KV cache block size of 128 (advertised +> automatically, so it is picked as long as `--block-size` is not pinned) and does not +> support logit soft capping, attention sinks, mm_prefix/R-SWA masking, DCP, or windowed +> encoder attention. Those configurations transparently fall back to FA2. A pinned +> `--block-size` that is not a multiple of 128 instead makes FlashAttention ineligible +> for such models, which is an error if the backend was requested explicitly. ## MiniMax M3 Sparse Attention Backends diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index e2f5c7f5d1..958b04307c 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/tests/kernels/attention/test_attention_selector.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from unittest.mock import patch +from contextlib import contextmanager +from unittest.mock import MagicMock, patch import pytest import torch @@ -14,6 +15,7 @@ from vllm.config import ( ) from vllm.platforms import current_platform from vllm.platforms.cpu import CpuPlatform +from vllm.platforms.interface import DeviceCapability if current_platform.is_cuda(): from vllm.platforms.cuda import CudaPlatform @@ -25,6 +27,7 @@ if current_platform.is_rocm(): else: RocmPlatform = None +from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend @@ -555,3 +558,175 @@ def test_flash_attn_accepts_handled_fp8_variants( # import order across earlier tests that patch vllm.platforms.current_platform. monkeypatch.setattr(fa_utils_mod.current_platform, "is_xpu", lambda: True) assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype) + + +blackwell_only = pytest.mark.skipif( + not current_platform.is_cuda(), reason="FA4 is CUDA-only" +) + + +@contextmanager +def _blackwell(vllm_config=None): + platform = MagicMock() + platform.is_xpu.return_value = False + platform.is_rocm.return_value = False + platform.get_device_capability.return_value = DeviceCapability(10, 0) + with ( + patch("vllm.v1.attention.backends.fa_utils.current_platform", platform), + patch( + "vllm.vllm_flash_attn.flash_attn_interface.is_fa_version_supported", + return_value=True, + ), + patch("vllm.config.get_current_vllm_config_or_none", return_value=vllm_config), + patch( + "vllm.v1.attention.backends.flash_attn.get_current_vllm_config_or_none", + return_value=vllm_config, + ), + ): + yield + + +def _hd256_config( + *, + is_mm_prefix_lm=False, + rswa_window=None, + dcp_size=1, + softcap=None, + head_size=256, + cache_dtype="auto", +): + vllm_config = MagicMock() + vllm_config.attention_config.flash_attn_version = None + vllm_config.model_config.is_mm_prefix_lm = is_mm_prefix_lm + vllm_config.model_config.rswa_window = rswa_window + vllm_config.model_config.hf_text_config.attn_logit_softcapping = softcap + vllm_config.model_config.get_head_size.return_value = head_size + vllm_config.cache_config.cache_dtype = cache_dtype + vllm_config.parallel_config.decode_context_parallel_size = dcp_size + return vllm_config + + +@blackwell_only +@pytest.mark.parametrize( + "kwargs,config_kwargs,expected", + [ + ({}, {}, 4), + ({"supports_fa4_hd256": False}, {}, 2), + ({"has_sinks": True}, {}, 2), + ({"requires_softcap": True}, {}, 2), + # Larger block sizes normalize to 128. + ({"kv_cache_block_size": 16}, {}, 2), + ({"kv_cache_block_size": 64}, {}, 2), + ({"kv_cache_block_size": 256}, {}, 4), + ({}, {"is_mm_prefix_lm": True}, 2), + ({}, {"rswa_window": 512}, 2), + ({}, {"dcp_size": 2}, 2), + ({}, {"softcap": 50.0}, 2), + ({}, {"cache_dtype": "fp8"}, 2), + ({"head_size": 128, "kv_cache_block_size": 16}, {}, 4), + ({"head_size": 192, "head_size_v": 128, "kv_cache_block_size": 16}, {}, 4), + ({"head_size": 256, "head_size_v": 128, "kv_cache_block_size": 16}, {}, 2), + ({"head_size": 256, "head_size_v": 64}, {}, 2), + ], +) +def test_fa4_hd256_fallback_matrix(kwargs, config_kwargs, expected): + from vllm.v1.attention.backends.fa_utils import get_flash_attn_version + + kwargs = { + "head_size": 256, + "kv_cache_block_size": 128, + "supports_fa4_hd256": True, + **kwargs, + } + with _blackwell(_hd256_config(head_size=kwargs["head_size"], **config_kwargs)): + assert get_flash_attn_version(**kwargs) == expected + + +@blackwell_only +@pytest.mark.parametrize( + "backend_name,config_kwargs,expected", + [ + ("FLASH_ATTN", {}, 128), + ("FLASH_ATTN", {"softcap": 50.0}, 16), + ("FLASH_ATTN", {"head_size": 128}, 16), + ("FLASH_ATTN", {"cache_dtype": "fp8"}, 16), + ("FLASH_ATTN_DIFFKV", {}, 16), + ], +) +def test_fa4_hd256_block_size_advertisement( + backend_name: str, config_kwargs: dict, expected: int +): + from vllm.v1.attention.backend import MultipleOf + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend + from vllm.v1.attention.backends.flash_attn_diffkv import ( + FlashAttentionDiffKVBackend, + ) + + backend = { + "FLASH_ATTN": FlashAttentionBackend, + "FLASH_ATTN_DIFFKV": FlashAttentionDiffKVBackend, + }[backend_name] + vllm_config = _hd256_config(**config_kwargs) + with _blackwell(vllm_config): + (size,) = backend.get_supported_kernel_block_sizes() + preferred = backend.get_preferred_block_size(16) + assert preferred == expected + if expected == 128: + assert size == 128 + else: + assert isinstance(size, MultipleOf) and size.base == expected + + +@blackwell_only +def test_fa4_hd256_mm_prefix_deselects_flash_attn(): + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend + + with _blackwell(_hd256_config(is_mm_prefix_lm=True)): + assert ( + FlashAttentionBackend.supports_combination( + head_size=256, + dtype=torch.bfloat16, + kv_cache_dtype=None, + block_size=128, + use_mla=False, + has_sink=False, + use_sparse=False, + use_mm_prefix=True, + device_capability=DeviceCapability(10, 0), + ) + is not None + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda() + or not current_platform.is_device_capability_family(100), + reason="requires a Blackwell GPU", +) +@pytest.mark.parametrize( + "attn_type,sliding_window,expected", + [ + (AttentionType.DECODER, None, 4), + (AttentionType.DECODER, 512, 4), + # Local encoder attention lacks the required seqused tensors. + (AttentionType.ENCODER_ONLY, None, 4), + (AttentionType.ENCODER_ONLY, 512, 2), + ], +) +def test_fa4_hd256_impl_selection(attn_type, sliding_window, expected): + from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl + + # Implementations are constructed before the platform updates the block size. + with set_current_vllm_config(VllmConfig()): + impl = FlashAttentionImpl( + num_heads=8, + head_size=256, + scale=0.0625, + num_kv_heads=8, + alibi_slopes=None, + sliding_window=sliding_window, + kv_cache_dtype="auto", + attn_type=attn_type, + ) + assert impl.vllm_flash_attn_version == expected + assert impl.fa4_hd256 == (expected == 4) diff --git a/tests/kernels/attention/test_flash_attn.py b/tests/kernels/attention/test_flash_attn.py index 2714cd8181..c58156c685 100644 --- a/tests/kernels/attention/test_flash_attn.py +++ b/tests/kernels/attention/test_flash_attn.py @@ -215,3 +215,78 @@ def test_varlen_with_paged_kv( torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), f"{torch.max(torch.abs(output - ref_output))}", ) + + +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="FA4 only has a dedicated head_size=256 kernel on Blackwell", +) +@pytest.mark.parametrize( + "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]] +) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS) +@torch.inference_mode() +def test_fa4_hd256_paged_call_shape( + seq_lens: list[tuple[int, int]], + num_heads: tuple[int, int], + sliding_window: int | None, +) -> None: + """Handles excess block-table columns with page-aligned ``max_seqlen_k``.""" + from vllm.v1.attention.backends.fa_utils import FA4_HD256_PAGE_SIZE + + if not is_fa_version_supported(4): + pytest.skip(f'FA4 unsupported: "{fa_version_unsupported_reason(4)}"') + torch.set_default_device("cuda") + set_random_seed(0) + + head_size = 256 + block_size = FA4_HD256_PAGE_SIZE + num_blocks = 2048 + dtype = torch.bfloat16 + num_query_heads, num_kv_heads = num_heads + query_lens = [x[0] for x in seq_lens] + kv_lens = [x[1] for x in seq_lens] + max_kv_len = max(kv_lens) + window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1) + scale = head_size**-0.5 + + query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype) + key_cache = torch.randn( + num_blocks, block_size, num_kv_heads, head_size, dtype=dtype + ) + value_cache = torch.randn_like(key_cache) + cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum( + dim=0, dtype=torch.int32 + ) + num_pages = (max_kv_len + block_size - 1) // block_size + block_tables = torch.randint( + 0, num_blocks, (len(seq_lens), num_pages + 3), dtype=torch.int32 + ) + + output = flash_attn_varlen_func( + q=query, + k=key_cache, + v=value_cache, + cu_seqlens_q=cu_query_lens, + seqused_k=torch.tensor(kv_lens, dtype=torch.int32), + max_seqlen_q=max(query_lens), + max_seqlen_k=num_pages * block_size, + softmax_scale=scale, + causal=True, + window_size=window_size, + block_table=block_tables[:, :num_pages], + num_splits=1, + fa_version=4, + ) + ref_output = ref_paged_attn( + query=query, + key_cache=key_cache, + value_cache=value_cache, + query_lens=query_lens, + kv_lens=kv_lens, + block_tables=block_tables, + scale=scale, + sliding_window=sliding_window, + ) + torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2) diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 4ebb1d05b9..0de6b85217 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -6,9 +6,13 @@ from typing import Any import vllm.envs as envs from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.utils.torch_utils import is_quantized_kv_cache logger = init_logger(__name__) +# The dedicated hd256 kernel requires one 128-token page per tile. +FA4_HD256_PAGE_SIZE = 128 + # Track whether upstream flash-attn is available on ROCm. # Set during module initialization and never modified afterwards. # This module-level flag avoids repeated import attempts and ensures @@ -72,6 +76,9 @@ def get_flash_attn_version( head_size: int | None = None, head_size_v: int | None = None, has_sinks: bool = False, + requires_softcap: bool = False, + kv_cache_block_size: int | None = None, + supports_fa4_hd256: bool = False, ) -> int | None: if current_platform.is_xpu(): return 2 @@ -168,26 +175,31 @@ def get_flash_attn_version( ) fa_version = 2 - # TODO: Restore the `requires_local_attention` restriction when FA4 - # head-dim 256 is re-enabled. - if fa_version == 4 and device_capability.major >= 10 and head_size == 256: - logger.warning_once( - "FA4 on Blackwell is temporarily disabled for head_size=256, " - "defaulting to FA version 2." - ) - fa_version = 2 + if fa_version == 4 and uses_fa4_hd256_kernel(head_size, head_size_v): + if not supports_fa4_hd256: + fa_version = 2 + elif ( + reason := _fa4_hd256_fallback_reason( + has_sinks, requires_softcap, kv_cache_block_size, vllm_config + ) + ) is not None: + logger.warning_once( + "FA4's Blackwell head_size=256 kernel does not support %s, " + "defaulting to FA version 2.", + reason, + ) + fa_version = 2 - # FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict - # supported head dimensions to ≤128. The 192/128 MLA prefill case is - # supported; 256 is temporarily disabled until upstream supports the - # required features. Development of symmetric 192, 384, and 512 support - # is tracked in https://github.com/Dao-AILab/flash-attention/issues/2456 + # FA4 head dimensions on Blackwell are limited by TMEM capacity. if ( fa_version == 4 and device_capability.major >= 10 and head_size is not None and head_size > 128 - and not (head_size == 192 and head_size_v == 128) + and not ( + (head_size == 256 and head_size_v in (None, 256)) + or (head_size == 192 and head_size_v == 128) + ) ): logger.warning_once( "FA4 on Blackwell does not support head_size=%d due to TMEM " @@ -209,6 +221,53 @@ def get_flash_attn_version( return None +def uses_fa4_hd256_kernel( + head_size: int | None, head_size_v: int | None = None +) -> bool: + """Return whether FA4 uses its dedicated hd256 kernel.""" + if head_size != 256: + return False + if head_size_v is not None and head_size_v != 256: + return False + capability = current_platform.get_device_capability() + return capability is not None and capability.major in (10, 11) + + +def _fa4_hd256_fallback_reason( + has_sinks: bool, + requires_softcap: bool, + kv_cache_block_size: int | None, + vllm_config: Any, +) -> str | None: + model_config = vllm_config.model_config if vllm_config is not None else None + cache_config = vllm_config.cache_config if vllm_config is not None else None + if has_sinks: + return "attention sinks" + if requires_softcap or ( + # Keep model-level and per-layer version selection consistent. + model_config is not None + and getattr(model_config.hf_text_config, "attn_logit_softcapping", None) + ): + return "logits soft capping" + if cache_config is not None and is_quantized_kv_cache(cache_config.cache_dtype): + return f"quantized KV cache dtype {cache_config.cache_dtype}" + if kv_cache_block_size is not None and kv_cache_block_size % FA4_HD256_PAGE_SIZE: + # Larger blocks are split into 128-token kernel pages. + return f"a KV cache block size of {kv_cache_block_size}" + if model_config is not None: + if model_config.is_mm_prefix_lm: + return "mm_prefix bidirectional attention" + if model_config.rswa_window is not None: + return "R-SWA" + if ( + vllm_config is not None + and vllm_config.parallel_config.decode_context_parallel_size > 1 + ): + # DCP reuses an unsliced block table with a non-page-aligned length. + return "decode context parallelism" + return None + + def is_fa_version_supported(fa_version: int) -> bool: try: from vllm.vllm_flash_attn.flash_attn_interface import ( @@ -227,6 +286,9 @@ def flash_attn_supports_kv_cache_dtype( head_size: int | None = None, head_size_v: int | None = None, has_sinks: bool = False, + requires_softcap: bool = False, + kv_cache_block_size: int | None = None, + supports_fa4_hd256: bool = False, ) -> bool: if kv_cache_dtype == "fp8_e5m2": return False @@ -237,6 +299,9 @@ def flash_attn_supports_kv_cache_dtype( head_size=head_size, head_size_v=head_size_v, has_sinks=has_sinks, + requires_softcap=requires_softcap, + kv_cache_block_size=kv_cache_block_size, + supports_fa4_hd256=supports_fa4_hd256, ) return (fa_version == 3 and current_platform.is_device_capability_family(90)) or ( fa_version == 4 and current_platform.is_device_capability_family(100) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 3d0244520c..9a6c66f5ec 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -24,11 +24,13 @@ from vllm.v1.attention.backend import ( MultipleOf, ) from vllm.v1.attention.backends.fa_utils import ( + FA4_HD256_PAGE_SIZE, flash_attn_supports_kv_cache_dtype, flash_attn_supports_quant_query_input, get_flash_attn_version, is_fa_version_supported, is_flash_attn_varlen_func_available, + uses_fa4_hd256_kernel, ) from vllm.v1.attention.backends.utils import ( fill_mm_prefix_query_ranges, @@ -85,15 +87,40 @@ class FlashAttentionBackend(AttentionBackend): "fp8", "fp8_e4m3", ] + head_size_v: int | None = None - @staticmethod - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + @classmethod + def _get_fa4_hd256_block_size(cls) -> int | None: + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None or vllm_config.model_config is None: + return None + + head_size = vllm_config.model_config.get_head_size() + if ( + uses_fa4_hd256_kernel(head_size, cls.head_size_v) + and get_flash_attn_version( + head_size=head_size, + head_size_v=cls.head_size_v, + supports_fa4_hd256=True, + ) + == 4 + ): + return FA4_HD256_PAGE_SIZE + return None + + @classmethod + def get_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]: + if block_size := cls._get_fa4_hd256_block_size(): + # Sliding-window specs select the smallest advertised size. + return [block_size] return [MultipleOf(16)] forward_includes_kv_cache_update: bool = False @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: + if block_size := cls._get_fa4_hd256_block_size(): + return max(default_block_size, block_size) if current_platform.is_xpu(): return max(default_block_size, 64) return super().get_preferred_block_size(default_block_size) @@ -194,12 +221,20 @@ class FlashAttentionBackend(AttentionBackend): head_size=head_size, head_size_v=head_size, has_sinks=has_sink, + kv_cache_block_size=block_size, + supports_fa4_hd256=True, ) ): 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 + and get_flash_attn_version( + head_size=head_size, + has_sinks=has_sink, + kv_cache_block_size=block_size, + supports_fa4_hd256=True, + ) + != 4 ): return ( "mm_prefix (PrefixLM bidirectional attention) requires " @@ -409,6 +444,15 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad self.max_num_splits = 0 # No upper bound on the number of splits. self.aot_schedule = get_flash_attn_version() == 3 + self.fa4_hd256 = uses_fa4_hd256_kernel(self.headdim) and ( + get_flash_attn_version( + head_size=self.headdim, + kv_cache_block_size=self.block_size, + supports_fa4_hd256=True, + ) + == 4 + ) + try: from vllm.distributed.parallel_state import get_dcp_group @@ -780,6 +824,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad metadata.scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata) def use_cascade_attention(self, *args, **kwargs) -> bool: + if self.fa4_hd256: + # Cascade may use a non-page-aligned prefix length. + return False return use_cascade_attention(*args, **kwargs) @@ -823,11 +870,30 @@ class FlashAttentionImpl(AttentionImpl): self.num_queries_per_kv = self.num_heads // self.num_kv_heads self.attn_type = attn_type + vllm_config = get_current_vllm_config_or_none() + uses_kv_cache = attn_type not in ( + AttentionType.ENCODER, + AttentionType.ENCODER_ONLY, + ) + # The final KV cache block size is unavailable during construction. 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, + requires_softcap=bool(self.logits_soft_cap), + supports_fa4_hd256=True, ) + self.fa4_hd256 = self.vllm_flash_attn_version == 4 and uses_fa4_hd256_kernel( + head_size + ) + if self.fa4_hd256 and not uses_kv_cache and sliding_window is not None: + # The hd256 kernel requires seqused_k for local attention. + logger.warning_once( + "FA4's Blackwell head_size=256 kernel does not support local " + "attention on encoder inputs, defaulting to FA version 2." + ) + self.vllm_flash_attn_version = 2 + self.fa4_hd256 = False logger.info_once( "Using FlashAttention version %s", self.vllm_flash_attn_version, @@ -843,6 +909,8 @@ class FlashAttentionImpl(AttentionImpl): head_size=head_size, head_size_v=head_size, has_sinks=sinks is not None, + requires_softcap=bool(self.logits_soft_cap), + supports_fa4_hd256=True, ): raise NotImplementedError( f"FlashAttention does not support {self.kv_cache_dtype}" @@ -861,7 +929,6 @@ class FlashAttentionImpl(AttentionImpl): self.supports_quant_query_input = flash_attn_supports_quant_query_input() - vllm_config = get_current_vllm_config_or_none() dcp_a2a = ( vllm_config is not None and vllm_config.parallel_config.decode_context_parallel_size > 1 @@ -1083,6 +1150,15 @@ class FlashAttentionImpl(AttentionImpl): ) causal = not has_window + num_splits = attn_metadata.max_num_splits + if self.fa4_hd256: + # hd256 requires page-aligned lengths, exact-width block + # tables, and no SplitKV. + num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE) + max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE + block_table = block_table[:, :num_pages] + num_splits = 1 + flash_attn_varlen_func( q=query[:num_actual_tokens], k=key_cache, @@ -1104,7 +1180,7 @@ class FlashAttentionImpl(AttentionImpl): k_descale=k_descale, v_descale=v_descale, dynamic_causal=dynamic_causal, - num_splits=attn_metadata.max_num_splits, + num_splits=num_splits, s_aux=self.sinks, mask_mod=rswa_mask_mod_fn or mm_mask_mod, aux_tensors=rswa_aux or mm_aux, @@ -1416,7 +1492,8 @@ class FlashAttentionImpl(AttentionImpl): else None, k_descale=layer._k_scale.expand(descale_shape), # type: ignore[operator] v_descale=layer._v_scale.expand(descale_shape), # type: ignore[operator] - num_splits=1 if self.batch_invariant_enabled else 0, + # The hd256 kernel does not support SplitKV. + num_splits=1 if self.batch_invariant_enabled or self.fa4_hd256 else 0, s_aux=self.sinks, ) diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index 326b21c2c1..40846758d8 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -87,6 +87,7 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl): head_size_v=FlashAttentionDiffKVBackend.head_size_v, has_sinks=self.sinks is not None, ) + self.fa4_hd256 = False def do_kv_cache_update( self,