refactoring backends(pt3)

Signed-off-by: vnadathur <[email protected]>

Author: vnadathur

tests/v1/attention/test_mla_backends.py

@@ -829,4 +829,4 @@ def test_backend_correctness(
 
         summary = f"{len(failures)} backend(s) failed: {', '.join(backend_names)}"
         detailed_msg = "\n".join(failures)
-        pytest.fail(f"{summary}\n{detailed_msg}")
+        pytest.fail(f"{summary}\n{detailed_msg}")

tests/v1/attention/test_sparse_mla_backends.py

@@ -23,7 +23,7 @@
 from vllm import _custom_ops as ops
 from vllm.attention.ops import flashmla
 from vllm.model_executor.layers.linear import ColumnParallelLinear
-from vllm.utils import cdiv
+from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend
 from vllm.v1.attention.backends.mla.indexer import split_prefill_chunks
 
@@ -389,4 +389,4 @@ def test_sparse_backend_decode_correctness(
 )
 def test_split_prefill_chunks(seq_lens, max_buf, start, expected):
     out = split_prefill_chunks(seq_lens, max_buf, start)
-    assert out == expected
+    assert out == expected

vllm/attention/backends/abstract.py

@@ -255,4 +255,4 @@ def forward(
 
 
 def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
-    return kv_cache_dtype != "auto"
+    return kv_cache_dtype != "auto"

vllm/attention/layer.py

@@ -104,19 +104,6 @@ def maybe_get_vit_flash_attn_backend(
 
     if current_platform.is_rocm() and attn_backend == _Backend.FLASH_ATTN:
         use_upstream_fa = True
-    elif current_platform.is_cuda():
-        if attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(
-            torch.get_default_dtype()
-        ):
-            attn_backend = _Backend.FLASH_ATTN
-            use_upstream_fa = True
-    elif current_platform.is_xpu():
-        assert attn_backend == _Backend.FLASH_ATTN, (
-            "XPU platform only supports FLASH_ATTN as vision attention backend."
-        )
-        use_upstream_fa = False
-    else:
-        return _Backend.TORCH_SDPA, None
 
     if attn_backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}:
         if attn_backend == _Backend.ROCM_AITER_FA:
@@ -125,7 +112,7 @@ def maybe_get_vit_flash_attn_backend(
             if use_upstream_fa:
                 from flash_attn import flash_attn_varlen_func
             else:
-                from vllm.attention.utils.fa_utils import flash_attn_varlen_func
+                from vllm.vllm_flash_attn import flash_attn_varlen_func
     else:
         flash_attn_varlen_func = None
 
@@ -482,18 +469,22 @@ def __init__(
         # If vllm native fa is selected, we use it directly.
         use_upstream_fa = False
 
-        self.attn_backend = (
-            backend
-            if backend
-            in {
-                _Backend.TORCH_SDPA,
-                _Backend.XFORMERS,
-                _Backend.PALLAS,
-                _Backend.ROCM_AITER_FA,
-                _Backend.FLASH_ATTN,
-            }
-            else _Backend.TORCH_SDPA
-        )
+        if current_platform.is_xpu():
+            # currently, only torch_sdpa is supported on xpu
+            self.attn_backend = _Backend.TORCH_SDPA
+        else:
+            self.attn_backend = (
+                backend
+                if backend
+                in {
+                    _Backend.TORCH_SDPA,
+                    _Backend.XFORMERS,
+                    _Backend.PALLAS,
+                    _Backend.ROCM_AITER_FA,
+                    _Backend.FLASH_ATTN,
+                }
+                else _Backend.TORCH_SDPA
+            )
 
         self.attn_backend, self._flash_attn_varlen_func = (
             maybe_get_vit_flash_attn_backend(