diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 2fdd5ce88d..a4941b4a18 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -1358,6 +1358,10 @@ def main(): profile_memory=args.profile_memory, warmup_ms=args.warmup_ms, prefill_backend=pb, + kv_lora_rank=args.kv_lora_rank, + qk_nope_head_dim=args.qk_nope_head_dim, + qk_rope_head_dim=args.qk_rope_head_dim, + v_head_dim=args.v_head_dim, ) result = run_benchmark(config) diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 755ce8961e..c4cedd9e4e 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -153,7 +153,9 @@ hardware and configuration. > **‡** Automatic selection tries FlashAttention first. On Blackwell > (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then -> TokenSpeed MLA. On other GPUs, only FlashAttention is considered. +> TokenSpeed MLA; for (qk_nope_head_dim=192, qk_rope_head_dim=64, +> v_head_dim=256) TRT-LLM Ragged is tried before FlashAttention. On other +> GPUs, only FlashAttention is considered. ### Decode Backends diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index d82397591f..f5985e7bc8 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -163,7 +163,7 @@ class TestGetMLAPrefillBackend: class TestAutoSelectMLAPrefillBackend: """Tests for fallback and error paths in auto-selection.""" - def test_blackwell_glm_dimensions_fall_back_to_trtllm(self): + def test_blackwell_glm_dimensions_use_trtllm(self): capability = DeviceCapability(major=10, minor=0) selector_config = MLAPrefillSelectorConfig( dtype=torch.bfloat16, @@ -184,11 +184,6 @@ class TestAutoSelectMLAPrefillBackend: with ( patch("vllm.platforms.current_platform") as mock_platform, patch.object(flash_attn_cls, "is_available", return_value=True), - patch( - "vllm.v1.attention.backends.mla.prefill.flash_attn." - "get_flash_attn_version", - return_value=4, - ), patch.object(trtllm_cls, "validate_configuration", return_value=[]), ): # Force the non-ROCm priority on the Blackwell. @@ -300,7 +295,12 @@ class TestROCmAiterFAPrefillSelection: with patch("vllm.platforms.current_platform") as mock_platform: mock_platform.is_rocm.return_value = True priorities = _get_mla_prefill_backend_priorities( - DeviceCapability(major=9, minor=5) + DeviceCapability(major=9, minor=5), + MLADimensions( + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ), ) assert priorities == [ diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 718580dcf6..9299af09a0 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -134,6 +134,7 @@ def get_flash_attn_version( head_size: int | None = None, head_size_v: int | None = None, has_sinks: bool = False, + requires_local_attention: bool = False, ) -> int | None: if current_platform.is_xpu(): return 2 @@ -230,17 +231,28 @@ def get_flash_attn_version( ) fa_version = 2 + if ( + fa_version == 4 + and device_capability.major >= 10 + and head_size == 256 + and requires_local_attention + ): + logger.warning_once( + "FA4 on Blackwell does not support local attention with " + "head_size=256, defaulting to FA version 2." + ) + fa_version = 2 + # FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict - # supported head dimensions. - # See: https://github.com/Dao-AILab/flash-attention/issues/1959 - # Exception: hdim 192 is supported for MLA's diff-headdim case - # (qk=192, v=128), added upstream in commits 1a15733e/1b36ab19. + # supported head dimensions to ≤128, with exceptions for 256 and 192/128 (MLA + # prefill). Development of symmetric 192, 384, and 512 support is being tracked + # in https://github.com/Dao-AILab/flash-attention/issues/2456 if ( fa_version == 4 and device_capability.major >= 10 and head_size is not None and head_size > 128 - and head_size != 192 + and not (head_size == 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 " diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 13bf435a54..416377deab 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -782,6 +782,7 @@ class FlashAttentionImpl(AttentionImpl): self.attn_type = attn_type self.vllm_flash_attn_version = get_flash_attn_version( requires_alibi=alibi_slopes is not None, + requires_local_attention=sliding_window is not None, head_size=head_size, has_sinks=sinks is not None, ) diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index c5fdbaebc0..d073f9b828 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -288,26 +288,23 @@ class FlashAttnPrefillBackend(MLAPrefillBackend): @classmethod def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool: - dims_deepseek = MLADimensions( - qk_nope_head_dim=128, - qk_rope_head_dim=64, - v_head_dim=128, - ) - dims_glm = MLADimensions( - qk_nope_head_dim=192, - qk_rope_head_dim=64, - v_head_dim=256, - ) - dims_mistral_s4 = MLADimensions( - qk_nope_head_dim=64, - qk_rope_head_dim=64, - v_head_dim=128, - ) - fa_version = get_flash_attn_version() - if fa_version == 4: - return mla_dimensions in [dims_deepseek, dims_mistral_s4] - else: - return mla_dimensions in [dims_deepseek, dims_glm, dims_mistral_s4] + return mla_dimensions in [ + MLADimensions( + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ), + MLADimensions( + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ), + MLADimensions( + qk_nope_head_dim=64, + qk_rope_head_dim=64, + v_head_dim=128, + ), + ] def __init__( self, diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index e0d54eee10..c313c90177 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -47,11 +47,13 @@ class MLAPrefillSelectorConfig(NamedTuple): def _get_mla_prefill_backend_priorities( device_capability: DeviceCapability, + mla_dimensions: MLADimensions, ) -> list[MLAPrefillBackendEnum]: """Get MLA prefill backend priorities based on device capability. Args: device_capability: The device's compute capability. + mla_dimensions: The model's MLA head dimensions. Returns: List of backends in priority order (highest priority first). @@ -65,6 +67,17 @@ def _get_mla_prefill_backend_priorities( ] if device_capability.major == 10: # Blackwell + if mla_dimensions == MLADimensions( + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ): + return [ + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] return [ MLAPrefillBackendEnum.FLASH_ATTN, MLAPrefillBackendEnum.TRTLLM_RAGGED, @@ -157,7 +170,10 @@ def _auto_select_mla_prefill_backend( Returns: The selected prefill backend class. """ - priorities = _get_mla_prefill_backend_priorities(device_capability) + priorities = _get_mla_prefill_backend_priorities( + device_capability, + selector_config.mla_dimensions, + ) all_invalid_reasons: dict[str, list[str]] = {} for backend_enum in priorities: