When VLLM_USE_FLASH_ATTN_PA=0, support the use of default block_size

vllm/attention/layer.py

@@ -76,7 +76,7 @@ class Attention(nn.Module):
             calculate_kv_scales = cache_config.calculate_kv_scales
         else:
             kv_cache_dtype = "auto"
-            block_size = 64 if envs.VLLM_USE_FLASH_ATTN_PA or envs.VLLM_USE_FLASH_MLA else 16
+            block_size = 64 if envs.VLLM_USE_FLASH_ATTN_PA and envs.VLLM_USE_FLASH_MLA else 16
             is_attention_free = False
             calculate_kv_scales = False
         if num_kv_heads is None:
@@ -305,7 +305,7 @@ class MultiHeadAttention(nn.Module):
         attn_backend = get_attn_backend(head_size,
                                         dtype,
                                         kv_cache_dtype=None,
-                                        block_size=64 if envs.VLLM_USE_FLASH_ATTN_PA or envs.VLLM_USE_FLASH_MLA else 16,
+                                        block_size=64 if envs.VLLM_USE_FLASH_ATTN_PA and envs.VLLM_USE_FLASH_MLA else 16,
                                         is_attention_free=False)
         backend = backend_name_to_enum(attn_backend.get_name())
         if current_platform.is_rocm():

vllm/config.py

@@ -1499,7 +1499,7 @@ PrefixCachingHashAlgo = Literal["builtin", "sha256"]
 class CacheConfig:
     """Configuration for the KV cache."""
 
-    block_size: BlockSize = 64 if envs.VLLM_USE_FLASH_ATTN_PA or envs.VLLM_USE_FLASH_MLA else 16  # type: ignore
+    block_size: BlockSize = 64 if envs.VLLM_USE_FLASH_ATTN_PA and envs.VLLM_USE_FLASH_MLA else 16  # type: ignore
     """Size of a contiguous cache block in number of tokens. This is ignored on
     neuron devices and set to `--max-model-len`. On CUDA devices, only block
     sizes up to 32 are supported. On HPU devices, block size defaults to 128.