add VLLM_USE_FLASH_MLA to use flashmla backend

vllm/envs.py

@@ -18,6 +18,7 @@ if TYPE_CHECKING:
     LD_LIBRARY_PATH: Optional[str] = None
     VLLM_USE_TRITON_FLASH_ATTN: bool = False
     VLLM_USE_TRITON_OPT_MLA: bool = False
+    VLLM_USE_FLASH_MLA: bool = False
     VLLM_USE_OPT_OP: bool = False
     VLLM_USE_TC_PAGED_ATTN: bool = False
     VLLM_USE_PA_PRINT_PARAM: bool = False 
@@ -654,6 +655,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_USE_TRITON_OPT_MLA":
     lambda: bool(int(os.getenv("VLLM_USE_TRITON_OPT_MLA", "0"))),
     
+    # If set, vLLM will use FLASH MLA attention optimizations.
+    "VLLM_USE_FLASH_MLA":
+    lambda: bool(int(os.getenv("VLLM_USE_FLASH_MLA", "1"))),
+
     # If set, vLLM will use the Triton implementation of moe_align_block_size,
     # i.e. moe_align_block_size_triton in fused_moe.py.
     "VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON":

vllm/platforms/rocm.py

@@ -138,9 +138,6 @@ class RocmPlatform(Platform):
                              kv_cache_dtype, block_size, use_v1,
                              use_mla) -> str:
         if use_mla:
-            # logger.info("Using Triton MLA backend.")
-            # return "vllm.attention.backends.triton_mla.TritonMLABackend"
-            
             if selected_backend == _Backend.TRITON_MLA or block_size != 64:
                 if use_v1:
                     logger.info_once("Using Triton MLA backend on V1 engine.")
@@ -150,6 +147,7 @@ class RocmPlatform(Platform):
                     logger.info("Using Triton MLA backend.")
                     return "vllm.attention.backends.triton_mla.TritonMLABackend"  
             else:
+                if envs.VLLM_USE_FLASH_MLA:
                     from vllm.attention.backends.flashmla import (
                         is_flashmla_supported)
                     if not is_flashmla_supported()[0]:
@@ -171,6 +169,9 @@ class RocmPlatform(Platform):
                             logger.info("Using FlashMLA backend.")
                             return ("vllm.attention.backends."
                                     "flashmla.FlashMLABackend")
+                else:
+                    logger.info("Using Triton MLA backend (block size 64).")
+                    return "vllm.attention.backends.triton_mla.TritonMLABackend"
                         
         selected_backend = (_Backend.ROCM_FLASH if selected_backend
                             == _Backend.FLASH_ATTN else selected_backend)