[Core] Remove FlashAttention block size restriction for hybrid models (#36701)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Thomas Parnell committed
Jun 26, 2026 at 23:16 UTC
c40d307731b82a9d472001c5feff18c24797d9df
1 file changed
-17
vllm/v1/attention/backends/flash_attn.py
-17
@@ -43,7 +43,6 @@ if is_flash_attn_varlen_func_available():
43
import vllm.envs as envs
44
from vllm.config import (
45
VllmConfig,
46
- get_current_vllm_config,
46
get_current_vllm_config_or_none,
47
get_layers_from_vllm_config,
48
)
@@ -75,22 +74,6 @@ class FlashAttentionBackend(AttentionBackend):
74
75
@staticmethod
76
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
78
- vllm_config = get_current_vllm_config()
79
- model_config = vllm_config.model_config
80
- cache_config = vllm_config.cache_config
81
- if (
82
- model_config
83
- and model_config.is_hybrid
84
- and (
85
- cache_config.mamba_ssm_cache_dtype == "float32"
86
- or cache_config.mamba_cache_dtype == "float32"
87
- )
88
- ):
89
- # NOTE(tdoublep): while in principle, FA supports
90
- # MultipleOf(16), these are the block sizes that do not
91
- # suffer from the NaN propagation problem described here:
92
- # https://github.com/Dao-AILab/flash-attention/issues/1974
93
- return [16, 32, 64]
77
return [MultipleOf(16)]
78
79
forward_includes_kv_cache_update: bool = False