quick fix for 37665 (#37722)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

vllm/model_executor/models/config.py

@@ -4,12 +4,14 @@ from copy import deepcopy
 from math import lcm
 from typing import TYPE_CHECKING
 
+from vllm.config.vllm import set_current_vllm_config
+from vllm.distributed.kv_transfer.kv_connector.utils import get_current_attn_backends
 from vllm.logger import init_logger
 from vllm.model_executor.models import ModelRegistry
 from vllm.utils.math_utils import cdiv, round_up
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, MLAAttentionSpec
+from vllm.v1.worker.utils import select_common_block_size
 
 if TYPE_CHECKING:
     from vllm.config import ModelConfig, VllmConfig
@@ -118,7 +120,6 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
         # Enable FULL_AND_PIECEWISE by default
         MambaModelConfig.verify_and_update_config(vllm_config)
 
-        attention_config = vllm_config.attention_config
         cache_config = vllm_config.cache_config
         model_config = vllm_config.model_config
         parallel_config = vllm_config.parallel_config
@@ -128,6 +129,12 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
         else:
             kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
 
+        # Get actual kernel block alignment from selected backend(s).
+        with set_current_vllm_config(vllm_config):
+            backends = get_current_attn_backends(vllm_config)
+            kernel_block_alignment_size = select_common_block_size(
+                cache_config.block_size, backends
+            )
         # get attention page size (for 1 token)
         # Attention backend constraints:
         # - FlashAttention (FA) requires block size to be multiple of 16
@@ -135,10 +142,6 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
         #   * CUTLASS_MLA backend: kernel_block_size 128 alignment
         #   * Other MLA backends: kernel_block_size 64 alignment
         if model_config.use_mla:
-            use_cutlass_mla = (
-                attention_config.backend == AttentionBackendEnum.CUTLASS_MLA
-            )
-            kernel_block_alignment_size = 128 if use_cutlass_mla else 64
             attn_page_size_1_token = MLAAttentionSpec(
                 block_size=1,
                 num_kv_heads=model_config.get_num_kv_heads(parallel_config),
@@ -146,7 +149,6 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
                 dtype=kv_cache_dtype,
             ).page_size_bytes
         else:
-            kernel_block_alignment_size = 16
             attn_page_size_1_token = FullAttentionSpec(
                 block_size=1,
                 num_kv_heads=model_config.get_num_kv_heads(parallel_config),