A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 19:00 UTC
fix(ROCm): Make flash_attn import optional in MLA attention (#33511)
Signed-off-by: rabi <ramishra@redhat.com>
Rabi Mishra committed
Feb 6, 2026 at 07:52 UTC
20d7454c9bb0c2de7f59863f5030e5f494cab178
1 file changed
+19
-3
vllm/model_executor/layers/attention/mla_attention.py
+19
-3
@@ -919,10 +919,20 @@ try:
919
920
is_vllm_fa = True
921
except ImportError:
922
- # For rocm use upstream flash attention
923
- if current_platform.is_rocm():
924
- from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]
922
is_vllm_fa = False
923
+ flash_attn_varlen_func = None # type: ignore[assignment]
924
+ # On ROCm, vllm_flash_attn is not available, try upstream flash_attn instead.
925
+ # On CUDA, vllm_flash_attn should always be available (built with vLLM),
926
+ # so we don't attempt the fallback there.
927
+ if current_platform.is_rocm():
928
+ try:
929
+ from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]
930
+ except ImportError:
931
+ logger.debug(
932
+ "flash_attn not available on ROCm; "
933
+ "MLA models using TRITON_MLA will require flash_attn. "
934
+ "AITER_MLA backends use aiter kernels instead."
935
+ )
936
937
938
def dynamic_per_batched_tensor_quant(
@@ -1917,6 +1927,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1927
self._run_prefill_new_tokens = self._run_prefill_new_tokens_cudnn
1928
self._pad_v = False
1929
else: # Use FlashAttention
1930
+ if flash_attn_varlen_func is None:
1931
+ raise RuntimeError(
1932
+ "MLA attention requires FlashAttention but it is not "
1933
+ "available. Please install flash_attn or use "
1934
+ "--attention-backend ROCM_AITER_MLA."
1935
+ )
1936
logger.info_once("Using FlashAttention prefill for MLA", scope="local")
1937
self._run_prefill_context_chunk = self._run_prefill_context_chunk_fa
1938
self._run_prefill_new_tokens = self._run_prefill_new_tokens_fa