[Attention] FlashAttn MLA

.buildkite/check-wheel-size.py

@@ -5,11 +5,11 @@
 import sys
 import zipfile
 
-# Read the VLLM_MAX_SIZE_MB environment variable, defaulting to 400 MiB
-# Note that we have 400 MiB quota, please use it wisely.
-# See https://github.com/pypi/support/issues/3792 .
+# Read the VLLM_MAX_SIZE_MB environment variable, defaulting to 450 MiB
+# Note that we have 800 MiB quota, please use it wisely.
+# See https://github.com/pypi/support/issues/6326 .
 # Please also sync the value with the one in Dockerfile.
-VLLM_MAX_SIZE_MB = int(os.environ.get("VLLM_MAX_SIZE_MB", 400))
+VLLM_MAX_SIZE_MB = int(os.environ.get("VLLM_MAX_SIZE_MB", 450))
 
 
 def print_top_10_largest_files(zip_file):

cmake/external_projects/vllm_flash_attn.cmake

@@ -38,7 +38,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 57b4e68b9f9d94750b46de8f8dbd2bfcc86edd4f
+          GIT_TAG ee4d25bd84e0cbc7e0b9b9685085fd5db2dcb62a
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

docker/Dockerfile

@@ -237,7 +237,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
 # Check the size of the wheel if RUN_WHEEL_CHECK is true
 COPY .buildkite/check-wheel-size.py check-wheel-size.py
 # sync the default value with .buildkite/check-wheel-size.py
-ARG VLLM_MAX_SIZE_MB=400
+ARG VLLM_MAX_SIZE_MB=450
 ENV VLLM_MAX_SIZE_MB=$VLLM_MAX_SIZE_MB
 ARG RUN_WHEEL_CHECK=true
 RUN if [ "$RUN_WHEEL_CHECK" = "true" ]; then \

tests/kernels/attention/test_attention_selector.py

@@ -22,7 +22,7 @@ def clear_cache():
 
 # Define MLA and non-MLA backends separately
 DEVICE_MLA_BACKENDS = {
-    "cuda": ["TRITON_MLA", "FLASHMLA"],
+    "cuda": ["TRITON_MLA", "FLASHMLA", "FLASH_ATTN_MLA", "CUTLASS_MLA"],
     "hip": ["TRITON_MLA", "ROCM_AITER_MLA"],
     "cpu": [],
 }
@@ -98,21 +98,14 @@ def test_env(
             with patch("vllm.attention.selector.current_platform",
                        RocmPlatform()):
                 if use_mla:
-                    # Validate HIP MLA backend-block_size combinations
-                    valid_combination = (
-                        (name == "TRITON_MLA" and block_size != 1)
-                        or (name == "ROCM_AITER_MLA" and block_size == 1))
-
-                    if valid_combination:
-                        backend = get_attn_backend(16,
-                                                   torch.float16,
-                                                   torch.float16,
-                                                   block_size,
-                                                   False,
-                                                   use_mla=use_mla)
-                        expected = f"{name}_VLLM_V1" if use_v1 else name
-                        assert backend.get_name() == expected
-                    else:
+                    # ROCm MLA backend logic:
+                    # - TRITON_MLA: supported when block_size != 1
+                    # - ROCM_AITER_MLA: supported when block_size == 1
+                    # If backend is forced but doesn't match block_size,
+                    # should raise ValueError
+
+                    if name == "TRITON_MLA" and block_size == 1:
+                        # TRITON_MLA doesn't support block_size == 1
                         with pytest.raises(ValueError) as exc_info:
                             get_attn_backend(16,
                                              torch.float16,
@@ -122,6 +115,27 @@ def test_env(
                                              use_mla=use_mla)
                         assert f"The selected backend, {name}" in str(
                             exc_info.value)
+                    elif name == "ROCM_AITER_MLA" and block_size != 1:
+                        # ROCM_AITER_MLA only supports block_size == 1
+                        with pytest.raises(ValueError) as exc_info:
+                            get_attn_backend(16,
+                                             torch.float16,
+                                             torch.float16,
+                                             block_size,
+                                             False,
+                                             use_mla=use_mla)
+                        assert f"The selected backend, {name}" in str(
+                            exc_info.value)
+                    else:
+                        # Valid backend-block_size combination
+                        backend = get_attn_backend(16,
+                                                   torch.float16,
+                                                   torch.float16,
+                                                   block_size,
+                                                   False,
+                                                   use_mla=use_mla)
+                        expected = f"{name}_VLLM_V1" if use_v1 else name
+                        assert backend.get_name() == expected
                 else:
                     backend = get_attn_backend(16,
                                                torch.float16,
@@ -136,26 +150,68 @@ def test_env(
             with patch("vllm.attention.selector.current_platform",
                        CudaPlatform()):
                 if use_mla:
-                    if name == "FLASHMLA" and block_size == 64:
-                        from vllm.attention.backends.flashmla import (
-                            is_flashmla_supported)
-
-                        # only on cuda platforms with specific capability.
-                        is_supported, _ = is_flashmla_supported()
-
-                        if not is_supported:
-                            # if platform is not supported then skip this case.
-                            pytest.skip()
+                    # CUDA MLA backend logic:
+                    # - CUTLASS_MLA: only supported with block_size == 128
+                    #   and Blackwell GPUs (SM 10.0), V1 only
+                    # - FLASHMLA: only supported with block_size == 64
+                    # - FLASH_ATTN_MLA: V1 only
+                    # - TRITON_MLA: fallback for other cases
+
+                    if name == "CUTLASS_MLA":
+                        if not use_v1:
+                            # CUTLASS_MLA only supported on V1 engine
+                            pytest.skip(
+                                "CUTLASS_MLA only supported on V1 engine")
+                        elif block_size != 128:
+                            # CUTLASS_MLA only supports block_size == 128
+                            pytest.skip(
+                                "CUTLASS_MLA only supports block_size 128")
+                        else:
+                            backend = get_attn_backend(16,
+                                                       torch.float16,
+                                                       torch.float16,
+                                                       block_size,
+                                                       False,
+                                                       use_mla=use_mla)
+                            expected = "CUTLASS_MLA_VLLM_V1"
+                            assert backend.get_name() == expected
+                    elif name == "FLASHMLA":
+                        if block_size != 64:
+                            # FlashMLA only supports block_size == 64
+                            pytest.skip("FlashMLA only supports block_size 64")
+                        else:
+                            from vllm.attention.backends.flashmla import (
+                                is_flashmla_supported)
+                            is_supported, _ = is_flashmla_supported()
+                            if not is_supported:
+                                pytest.skip(
+                                    "FlashMLA not supported on this platform")
+                            else:
+                                backend = get_attn_backend(16,
+                                                           torch.float16,
+                                                           torch.float16,
+                                                           block_size,
+                                                           False,
+                                                           use_mla=use_mla)
+                                expected = f"{name}_VLLM_V1" if use_v1 else name
+                                assert backend.get_name() == expected
+                    elif name == "FLASH_ATTN_MLA":
+                        if not use_v1:
+                            # FlashAttention MLA only supported on V1 engine
+                            pytest.skip(
+                                "FlashAttention MLA only supported on V1 engine"
+                            )
                         else:
                             backend = get_attn_backend(16,
                                                        torch.float16,
                                                        torch.float16,
                                                        block_size,
                                                        False,
                                                        use_mla=use_mla)
-                            expected = f"{name}_VLLM_V1" if use_v1 else name
+                            expected = "FLASH_ATTN_MLA"
                             assert backend.get_name() == expected
                     else:
+                        # TRITON_MLA or other fallback
                         backend = get_attn_backend(16,
                                                    torch.float16,
                                                    torch.float16,

tests/v1/attention/test_attention_backends.py

@@ -70,22 +70,6 @@ def _convert_dtype_to_torch(dtype):
 }
 
 
-def create_dummy_kv_cache(kv_cache_spec: FullAttentionSpec,
-                          device: torch.device,
-                          num_blocks: int = 100) -> torch.Tensor:
-    """Create a dummy KV cache tensor for testing."""
-    kv_cache = torch.randn(
-        2,  # K and V
-        num_blocks,
-        kv_cache_spec.block_size,
-        kv_cache_spec.num_kv_heads,
-        kv_cache_spec.head_size,
-        dtype=_convert_dtype_to_torch(kv_cache_spec.dtype),
-        device=device,
-    )
-    return kv_cache
-
-
 def create_and_prepopulate_kv_cache(
         k_contexts: list[torch.Tensor],
         v_contexts: list[torch.Tensor],