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 
@@ -653,6 +654,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # If set, vLLM will use optimized MLA attention optimizations.
     "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/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.")
@@ -148,29 +145,33 @@ class RocmPlatform(Platform):
                             "triton_mla.TritonMLABackend")
                 else:
                     logger.info("Using Triton MLA backend.")
-                    return "vllm.attention.backends.triton_mla.TritonMLABackend"
+                    return "vllm.attention.backends.triton_mla.TritonMLABackend"  
             else:
-                from vllm.attention.backends.flashmla import (
-                    is_flashmla_supported)
-                if not is_flashmla_supported()[0]:
-                    logger.warning(
-                        "FlashMLA backend is not supported due to %s",
-                        is_flashmla_supported()[1])
-                elif block_size != 64:
-                    logger.warning(
-                        "FlashMLA backend is not supported for block size %d"
-                        " (currently only supports block size 64).",
-                        block_size)
-                else:
-                    if use_v1:
-                        logger.info_once(
-                            "Using FlashMLA backend on V1 engine.")
-                        return ("vllm.v1.attention.backends.mla."
-                                "flashmla.FlashMLABackend")
+                if envs.VLLM_USE_FLASH_MLA:
+                    from vllm.attention.backends.flashmla import (
+                        is_flashmla_supported)
+                    if not is_flashmla_supported()[0]:
+                        logger.warning(
+                            "FlashMLA backend is not supported due to %s",
+                            is_flashmla_supported()[1])
+                    elif block_size != 64:
+                        logger.warning(
+                            "FlashMLA backend is not supported for block size %d"
+                            " (currently only supports block size 64).",
+                            block_size)
                     else:
-                        logger.info("Using FlashMLA backend.")
-                        return ("vllm.attention.backends."
-                                "flashmla.FlashMLABackend")
+                        if use_v1:
+                            logger.info_once(
+                                "Using FlashMLA backend on V1 engine.")
+                            return ("vllm.v1.attention.backends.mla."
+                                    "flashmla.FlashMLABackend")
+                        else:
+                            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)