@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0ff70821c9
Fork 0

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 12, 16:00 UTC

[Core] Deprecate `xformers` (#29262)

Signed-off-by: Roger Wang <hey@rogerw.io>

Roger Wang committed Nov 23, 2025 at 20:18 UTC 0ff70821c9b0b991197fa7f3264bf9dd78b8d4b3
31 files changed +77 -963
docker/Dockerfile.nightly_torch
+1 -34
@@ -76,34 +76,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
76 RUN --mount=type=cache,target=/root/.cache/uv \
77 uv pip install --system -r requirements/common.txt
78
79 -# must put before installing xformers, so it can install the correct version of xfomrers.
80 -ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
81 -ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
82 -
83 -# Build xformers with cuda and torch nightly
84 -# following official xformers guidance: https://github.com/facebookresearch/xformers#build
85 -# todo(elainewy): cache xformers build result for faster build
86 -ARG max_jobs=16
87 -ENV MAX_JOBS=${max_jobs}
88 -ARG XFORMERS_COMMIT=f2de641ef670510cadab099ce6954031f52f191c
89 -
90 -ENV CCACHE_DIR=/root/.cache/ccache
91 -RUN --mount=type=cache,target=/root/.cache/ccache \
92 - --mount=type=cache,target=/root/.cache/uv \
93 - echo 'git clone xformers...' \
94 - && git clone https://github.com/facebookresearch/xformers.git --recursive \
95 - && cd xformers \
96 - && git checkout ${XFORMERS_COMMIT} \
97 - && git submodule update --init --recursive \
98 - && echo 'finish git clone xformers...' \
99 - && rm -rf build \
100 - && python3 setup.py bdist_wheel --dist-dir=../xformers-dist --verbose \
101 - && cd .. \
102 - && rm -rf xformers
103 -
104 -RUN --mount=type=cache,target=/root/.cache/uv \
105 - uv pip install --system xformers-dist/*.whl --verbose
106 -
79 # build can take a long time, and the torch nightly version fetched from url can be different in next docker stage.
80 # track the nightly torch version used in the build, when we set up runtime environment we can make sure the version is the same
81 RUN uv pip freeze | grep -i '^torch\|^torchvision\|^torchaudio' > torch_build_versions.txt
@@ -233,11 +205,6 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/vllm
205 --mount=type=cache,target=/root/.cache/uv \
206 uv pip install --system vllm-dist/*.whl --verbose
207
236 -# install xformers again for the new environment
237 -RUN --mount=type=bind,from=base,src=/workspace/xformers-dist,target=/vllm-workspace/xformers-dist \
238 - --mount=type=cache,target=/root/.cache/uv \
239 - uv pip install --system /vllm-workspace/xformers-dist/*.whl --verbose
240 -
208 ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
209
210 # install package for build flashinfer
@@ -307,7 +274,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
274 uv pip install --system -r requirements/nightly_torch_test.txt
275
276 # Logging to confirm the torch versions
310 -RUN pip freeze | grep -E 'torch|xformers|vllm|flashinfer'
277 +RUN pip freeze | grep -E 'torch|vllm|flashinfer'
278
279 # Logging to confirm all the packages are installed
280 RUN pip freeze
docs/contributing/ci/update_pytorch_version.md
-15
@@ -98,21 +98,6 @@ to warm it up so that future builds are faster.
98 <img width="60%" alt="Buildkite new build popup" src="https://github.com/user-attachments/assets/a8ff0fcd-76e0-4e91-b72f-014e3fdb6b94">
99 </p>
100
101 -## Update dependencies
102 -
103 -Several vLLM dependencies like xFormers depend on PyTorch and need
104 -to be updated accordingly. Rather than waiting for all of them to publish new
105 -releases (which would take too much time), they can be built from
106 -source to unblock the update process.
107 -
108 -### xFormers
109 -
110 -```bash
111 -export TORCH_CUDA_ARCH_LIST='7.5 8.0+PTX 9.0a'
112 -MAX_JOBS=16 uv pip install --system \
113 - --no-build-isolation "git+https://github.com/facebookresearch/xformers@v0.0.32.post2"
114 -```
115 -
101 ## Update all the different vLLM platforms
102
103 Rather than attempting to update all vLLM platforms in a single pull request, it's more manageable
docs/getting_started/quickstart.md
+1 -1
@@ -283,7 +283,7 @@ Currently, vLLM supports multiple backends for efficient Attention computation a
283
284 If desired, you can also manually set the backend of your choice by configuring the environment variable `VLLM_ATTENTION_BACKEND` to one of the following options:
285
286 -- On NVIDIA CUDA: `FLASH_ATTN`, `FLASHINFER` or `XFORMERS`.
286 +- On NVIDIA CUDA: `FLASH_ATTN` or `FLASHINFER`.
287 - On AMD ROCm: `TRITON_ATTN`, `ROCM_ATTN`, `ROCM_AITER_FA` or `ROCM_AITER_UNIFIED_ATTN`.
288
289 For AMD ROCm, you can further control the specific Attention implementation using the following variables:
examples/online_serving/openai_embedding_long_text/service.sh
-1
@@ -22,7 +22,6 @@ API_KEY=${API_KEY:-"your-api-key"}
22 POOLING_TYPE=${POOLING_TYPE:-"auto"} # auto, MEAN, CLS, LAST
23 export VLLM_ENABLE_CHUNKED_PROCESSING=true
24 export CUDA_VISIBLE_DEVICES=2,3,4,5
25 -# export VLLM_ATTENTION_BACKEND=XFORMERS
25
26 echo "🚀 Starting vLLM Embedding Server with Enhanced Chunked Processing"
27 echo "=================================================================="
requirements/cuda.txt
-1
@@ -9,6 +9,5 @@ torch==2.9.0
9 torchaudio==2.9.0
10 # These must be updated alongside torch
11 torchvision==0.24.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
12 -xformers==0.0.33.post1; platform_system == 'Linux' and platform_machine == 'x86_64' # Requires PyTorch >= 2.9
12 # FlashInfer should be updated together with the Dockerfile
13 flashinfer-python==0.5.2
tests/basic_correctness/test_basic_correctness.py
-3
@@ -74,9 +74,6 @@ def test_models(
74 model_executor: str,
75 enable_prompt_embeds: bool,
76 ) -> None:
77 - if backend == "XFORMERS" and model == "google/gemma-2-2b-it":
78 - pytest.skip(f"{backend} does not support gemma2 with full context length.")
79 -
77 with monkeypatch.context() as m:
78 m.setenv("VLLM_ATTENTION_BACKEND", backend)
79
tests/kernels/attention/test_attention.py
-129
@@ -13,12 +13,6 @@ from vllm.attention.layer import Attention, MultiHeadAttention
13 from vllm.platforms import current_platform
14 from vllm.utils.mem_utils import get_max_shared_memory_bytes
15
16 -if not current_platform.is_rocm():
17 - from xformers import ops as xops
18 - from xformers.ops.fmha.attn_bias import BlockDiagonalCausalMask
19 -
20 - from tests.kernels.utils import make_alibi_bias
21 -
16 FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
17 # This will change depending on the compute capability.
18 # - 512 as a buffer
@@ -448,129 +442,6 @@ def ref_multi_query_kv_attention(
442 return torch.cat(ref_outputs, dim=0)
443
444
451 -@pytest.mark.parametrize("num_seqs", NUM_PREFILL_SEQS)
452 -@pytest.mark.parametrize("num_heads", NUM_HEADS)
453 -@pytest.mark.parametrize("head_size", HEAD_SIZES)
454 -@pytest.mark.parametrize("dtype", DTYPES)
455 -@pytest.mark.parametrize("seed", SEEDS)
456 -@pytest.mark.parametrize("device", CUDA_DEVICES)
457 -@pytest.mark.skipif(
458 - current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
459 -)
460 -@torch.inference_mode()
461 -def test_multi_query_kv_attention(
462 - num_seqs: int,
463 - num_heads: tuple[int, int],
464 - head_size: int,
465 - dtype: torch.dtype,
466 - seed: int,
467 - device: str,
468 - use_alibi: bool = False,
469 -) -> None:
470 - current_platform.seed_everything(seed)
471 - torch.set_default_device(device)
472 - # MAX_SEQ_LEN sometimes causes OOM in the reference implementation.
473 - # As the xformers library is already tested with its own tests, we can use
474 - # a smaller MAX_SEQ_LEN here.
475 - max_len = min(MAX_SEQ_LEN, 4096)
476 - seq_lens = random.sample(range(1, max_len), num_seqs)
477 - num_tokens = sum(seq_lens)
478 -
479 - scale = float(1.0 / (head_size**0.5))
480 - num_query_heads, num_kv_heads = num_heads
481 - qkv = torch.empty(
482 - num_tokens, num_query_heads + 2 * num_kv_heads, head_size, dtype=dtype
483 - )
484 - qkv.uniform_(-scale, scale)
485 - query, key, value = qkv.split([num_query_heads, num_kv_heads, num_kv_heads], dim=1)
486 -
487 - num_queries_per_kv = num_query_heads // num_kv_heads
488 - if num_queries_per_kv > 1:
489 - # Handle MQA and GQA
490 - key = torch.repeat_interleave(key, num_queries_per_kv, dim=1)
491 - value = torch.repeat_interleave(value, num_queries_per_kv, dim=1)
492 - alibi_bias = None
493 - if use_alibi:
494 - alibi_slopes = torch.randn(num_query_heads, dtype=torch.float)
495 - attn_bias = make_alibi_bias(alibi_slopes, num_kv_heads, dtype, seq_lens)
496 - output = torch.empty_like(query)
497 - start = 0
498 - # Dynamic sequence length not supported with custom attn_bias.
499 - for i, seq_len in enumerate(seq_lens):
500 - end = start + seq_len
501 - out = xops.memory_efficient_attention_forward(
502 - query[None, start:end],
503 - key[None, start:end],
504 - value[None, start:end],
505 - attn_bias=attn_bias[i],
506 - p=0.0,
507 - scale=scale,
508 - )
509 - output[start:end].copy_(out.view_as(query[start:end]))
510 - start += seq_len
511 - # xformers.AttentionBias to Tensor for use in reference impl.
512 - alibi_bias = [
513 - b.materialize((1, num_query_heads, i, i), device=device).squeeze()
514 - for b, i in zip(attn_bias, seq_lens)
515 - ]
516 - else:
517 - attn_bias = BlockDiagonalCausalMask.from_seqlens(seq_lens)
518 - output = xops.memory_efficient_attention_forward(
519 - query.unsqueeze(0),
520 - key.unsqueeze(0),
521 - value.unsqueeze(0),
522 - attn_bias=attn_bias,
523 - p=0.0,
524 - scale=scale,
525 - )
526 - output = output.squeeze(0)
527 -
528 - cu_seq_lens = [0]
529 - for seq_len in seq_lens:
530 - cu_seq_lens.append(cu_seq_lens[-1] + seq_len)
531 - ref_output = ref_multi_query_kv_attention(
532 - cu_seq_lens,
533 - query,
534 - key,
535 - value,
536 - scale,
537 - alibi_bias,
538 - dtype,
539 - )
540 - atol = get_default_atol(output) if current_platform.is_rocm() else 1e-3
541 - rtol = get_default_rtol(output) if current_platform.is_rocm() else 1e-5
542 - torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
543 -
544 -
545 -@pytest.mark.parametrize("num_seqs", NUM_PREFILL_SEQS)
546 -@pytest.mark.parametrize("num_heads", NUM_HEADS)
547 -@pytest.mark.parametrize("head_size", [64])
548 -@pytest.mark.parametrize("dtype", DTYPES)
549 -@pytest.mark.parametrize("seed", SEEDS)
550 -@pytest.mark.parametrize("device", CUDA_DEVICES)
551 -@pytest.mark.skipif(
552 - current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
553 -)
554 -@torch.inference_mode()
555 -def test_multi_query_kv_attention_with_alibi(
556 - num_seqs: int,
557 - num_heads: tuple[int, int],
558 - head_size: int,
559 - dtype: torch.dtype,
560 - seed: int,
561 - device: str,
562 -) -> None:
563 - return test_multi_query_kv_attention(
564 - num_seqs,
565 - num_heads,
566 - head_size,
567 - dtype,
568 - seed,
569 - device,
570 - use_alibi=True,
571 - )
572 -
573 -
445 @pytest.mark.parametrize("attention_cls", [Attention, MultiHeadAttention])
446 def test_num_heads_not_divisble_by_num_kv_heads(attention_cls: type) -> None:
447 head_size = 64
tests/kernels/attention/test_attention_selector.py
+1 -7
@@ -34,7 +34,7 @@ DEVICE_MLA_BACKENDS = {
34 }
35
36 DEVICE_REGULAR_ATTN_BACKENDS = {
37 - "cuda": ["XFORMERS", "FLASHINFER", "FLASH_ATTN"],
37 + "cuda": ["FLASHINFER", "FLASH_ATTN"],
38 "hip": ["ROCM_ATTN"],
39 "cpu": ["CPU_ATTN"],
40 }
@@ -207,12 +207,6 @@ def test_env(
207 )
208 expected = "FLASHINFER"
209 assert backend.get_name() == expected
210 - elif name == "XFORMERS":
211 - backend = get_attn_backend(
212 - 32, torch.float16, None, block_size, use_mla=use_mla
213 - )
214 - expected = "XFORMERS"
215 - assert backend.get_name() == expected
210 elif name == "FLASH_ATTN":
211 backend = get_attn_backend(
212 32, torch.float16, None, block_size, use_mla=use_mla
tests/kernels/attention/test_mha_attn.py
-4
@@ -24,10 +24,6 @@ from vllm.platforms.rocm import RocmPlatform
24 def clear_cache():
25 """Clear lru cache to ensure each test case runs without caching."""
26 _cached_get_attn_backend.cache_clear()
27 - # Clear xformers availability cache
28 - import vllm.attention.layer as layer_module
29 -
30 - layer_module.USE_XFORMERS_OPS = None
27
28
29 @pytest.mark.parametrize("device", ["cpu", "hip", "cuda"])
tests/kernels/utils.py
+11 -67
@@ -509,43 +509,6 @@ def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs:
509 )
510
511
512 -def make_alibi_bias(
513 - alibi_slopes: torch.Tensor,
514 - num_kv_heads: int,
515 - dtype: torch.dtype,
516 - seq_lens: list[int],
517 -) -> list[Any]:
518 - """Create ALiBi biases compatible with xFormers attention tests."""
519 - from xformers.ops.fmha.attn_bias import LowerTriangularMaskWithTensorBias
520 -
521 - if alibi_slopes is None:
522 - return [None for _ in seq_lens]
523 -
524 - attn_biases: list[Any] = []
525 - num_heads = alibi_slopes.shape[0]
526 - assert num_heads >= num_kv_heads, (
527 - "ALiBi slopes expect at least as many heads as KV heads"
528 - )
529 -
530 - for seq_len in seq_lens:
531 - bias = torch.arange(seq_len, dtype=dtype, device=alibi_slopes.device)
532 - bias = bias[None, :] - bias[:, None]
533 -
534 - padded_len = (seq_len + 7) // 8 * 8
535 - bias_tensor = torch.empty(
536 - 1,
537 - num_heads,
538 - seq_len,
539 - padded_len,
540 - device=alibi_slopes.device,
541 - dtype=dtype,
542 - )[:, :, :, :seq_len].copy_(bias)
543 - bias_tensor.mul_(alibi_slopes[:, None, None])
544 - attn_biases.append(LowerTriangularMaskWithTensorBias(bias_tensor))
545 -
546 - return attn_biases
547 -
548 -
512 def _make_metadata_tensors(
513 seq_lens: list[int] | None,
514 context_lens: list[int] | None,
@@ -649,23 +612,12 @@ def make_kv_cache(
612
613 Returns:
614
652 - * kv_cache: 2 x num_blocks x (block_size * num_heads * head_size)
653 - * for backend 'XFORMERS'
615 * kv_cache: 2 x num_blocks x block_size x num_heads x head_size
616 * for backend 'FLASH_ATTN'
617 """
657 - if backend == "XFORMERS":
658 - kv_cache = torch.rand((2, num_blocks, block_size * num_heads * head_size)).to(
659 - device
660 - )
661 - elif backend == "FLASH_ATTN":
662 - kv_cache = torch.rand((2, num_blocks, block_size, num_heads, head_size)).to(
663 - device
664 - )
665 - else:
666 - raise ValueError(
667 - f"Unknown backend value: '{backend}'. Expected 'XFORMERS' or 'FLASH_ATTN'."
668 - )
618 + if backend != "FLASH_ATTN":
619 + raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.")
620 + kv_cache = torch.rand((2, num_blocks, block_size, num_heads, head_size)).to(device)
621 if default_val is not None:
622 kv_cache[:, :, :] = default_val
623 return kv_cache
@@ -843,22 +795,14 @@ def assert_actual_matches_ideal(
795 * output_under_test: actually observed output value
796 """
797 ideal_output = test_params.packed_qkvo.ideal_output
846 - if backend == "XFORMERS":
847 - torch.testing.assert_close(
848 - ideal_output, output_under_test.view_as(ideal_output)
849 - )
850 -
851 - elif backend == "FLASH_ATTN":
852 - # For FlashAttention override the accuracy thresholds to non default
853 - # values since we notice a higher difference between the ideal and
854 - # actual output.
855 - torch.testing.assert_close(
856 - ideal_output, output_under_test.view_as(ideal_output), atol=0.01, rtol=0.016
857 - )
858 - else:
859 - raise ValueError(
860 - f"Unknown backend value: '{backend}'. Expected 'XFORMERS' or 'FLASH_ATTN'."
861 - )
798 + if backend != "FLASH_ATTN":
799 + raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.")
800 + # For FlashAttention override the accuracy thresholds to non default
801 + # values since we notice a higher difference between the ideal and
802 + # actual output.
803 + torch.testing.assert_close(
804 + ideal_output, output_under_test.view_as(ideal_output), atol=0.01, rtol=0.016
805 + )
806
807
808 # Copied/modified from torch._refs.__init__.py
tests/lora/test_minicpmv_tp.py
-12
@@ -57,10 +57,6 @@ def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
57 return generated_texts
58
59
60 -@pytest.mark.xfail(
61 - current_platform.is_rocm(),
62 - reason="MiniCPM-V dependency xformers incompatible with ROCm",
63 -)
60 def test_minicpmv_lora(minicpmv_lora_files):
61 llm = vllm.LLM(
62 MODEL_PATH,
@@ -84,10 +80,6 @@ def test_minicpmv_lora(minicpmv_lora_files):
80 @pytest.mark.skipif(
81 current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
82 )
87 -@pytest.mark.xfail(
88 - current_platform.is_rocm(),
89 - reason="MiniCPM-V dependency xformers incompatible with ROCm",
90 -)
83 @multi_gpu_test(num_gpus=4)
84 def test_minicpmv_tp4_wo_fully_sharded_loras(minicpmv_lora_files):
85 llm = vllm.LLM(
@@ -108,10 +100,6 @@ def test_minicpmv_tp4_wo_fully_sharded_loras(minicpmv_lora_files):
100 @pytest.mark.skipif(
101 current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
102 )
111 -@pytest.mark.xfail(
112 - current_platform.is_rocm(),
113 - reason="MiniCPM-V dependency xformers incompatible with ROCm",
114 -)
103 @multi_gpu_test(num_gpus=4)
104 def test_minicpmv_tp4_fully_sharded_loras(minicpmv_lora_files):
105 llm = vllm.LLM(
tests/lora/test_qwen2vl.py
-15
@@ -2,12 +2,9 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass
4
5 -import pytest
6 -
5 import vllm
6 from vllm.assets.image import ImageAsset
7 from vllm.lora.request import LoRARequest
10 -from vllm.platforms import current_platform
8 from vllm.sampling_params import BeamSearchParams
9
10
@@ -142,10 +139,6 @@ QWEN2VL_MODEL_PATH = "Qwen/Qwen2-VL-2B-Instruct"
139 QWEN25VL_MODEL_PATH = "Qwen/Qwen2.5-VL-3B-Instruct"
140
141
145 -@pytest.mark.xfail(
146 - current_platform.is_rocm(),
147 - reason="Qwen2-VL dependency xformers incompatible with ROCm",
148 -)
142 def test_qwen2vl_lora(qwen2vl_lora_files):
143 """Test Qwen 2.0 VL model with LoRA"""
144 config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files)
@@ -156,10 +149,6 @@ def test_qwen2vl_lora(qwen2vl_lora_files):
149 tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id)
150
151
159 -@pytest.mark.xfail(
160 - current_platform.is_rocm(),
161 - reason="Qwen2-VL dependency xformers incompatible with ROCm",
162 -)
152 def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):
153 """Test Qwen 2.0 VL model with LoRA through beam search."""
154 config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files)
@@ -178,10 +167,6 @@ def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):
167 )
168
169
181 -@pytest.mark.xfail(
182 - current_platform.is_rocm(),
183 - reason="Qwen2.5-VL dependency xformers incompatible with ROCm",
184 -)
170 def test_qwen25vl_lora(qwen25vl_lora_files):
171 """Test Qwen 2.5 VL model with LoRA"""
172 config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files)
vllm/attention/backends/registry.py
-1
@@ -43,7 +43,6 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
43
44 FLASH_ATTN = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
45 TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"
46 - XFORMERS = "vllm.v1.attention.backends.xformers.XFormersAttentionBackend"
46 ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
47 ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"
48 ROCM_AITER_TRITON_MLA = (
vllm/attention/layer.py
-38
@@ -51,31 +51,6 @@ else:
51
52 FP8_DTYPE = current_platform.fp8_dtype()
53 logger = init_logger(__name__)
54 -USE_XFORMERS_OPS = None
55 -
56 -
57 -def check_xformers_availability():
58 - global USE_XFORMERS_OPS
59 - if USE_XFORMERS_OPS is not None:
60 - return USE_XFORMERS_OPS
61 -
62 - if current_platform.is_cuda() and current_platform.has_device_capability(100):
63 - # Xformers FA is not compatible with B200
64 - USE_XFORMERS_OPS = False
65 - else:
66 - try:
67 - from importlib.util import find_spec
68 -
69 - find_spec("xformers.ops")
70 - USE_XFORMERS_OPS = True
71 - except ImportError:
72 - USE_XFORMERS_OPS = False
73 -
74 - # the warning only needs to be shown once
75 - if not USE_XFORMERS_OPS:
76 - logger.warning("Xformers is not available, falling back.")
77 -
78 - return USE_XFORMERS_OPS
54
55
56 def check_upstream_fa_availability(dtype: torch.dtype):
@@ -533,7 +508,6 @@ class MultiHeadAttention(nn.Module):
508 if backend
509 in {
510 AttentionBackendEnum.TORCH_SDPA,
536 - AttentionBackendEnum.XFORMERS,
511 AttentionBackendEnum.PALLAS,
512 AttentionBackendEnum.ROCM_AITER_FA,
513 AttentionBackendEnum.FLASH_ATTN,
@@ -549,12 +523,6 @@ class MultiHeadAttention(nn.Module):
523 )
524 )
525
552 - if (
553 - self.attn_backend == AttentionBackendEnum.XFORMERS
554 - and not check_xformers_availability()
555 - ):
556 - self.attn_backend = AttentionBackendEnum.TORCH_SDPA
557 -
526 self.is_flash_attn_backend = self.attn_backend in {
527 AttentionBackendEnum.FLASH_ATTN,
528 AttentionBackendEnum.ROCM_AITER_FA,
@@ -614,12 +582,6 @@ class MultiHeadAttention(nn.Module):
582 max_seqlen_k=kv_len,
583 softmax_scale=self.scale,
584 )
617 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
618 - from xformers import ops as xops
619 -
620 - out = xops.memory_efficient_attention_forward(
621 - query, key, value, scale=self.scale
622 - )
585 elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
586 query, key, value = (x.transpose(1, 2) for x in (query, key, value))
587 out = F.scaled_dot_product_attention(query, key, value, scale=self.scale)
vllm/attention/ops/vit_attn_wrappers.py
+1 -37
@@ -3,7 +3,7 @@
3 """
4 This file contains ops for ViT attention to be compatible with torch.compile
5 as there are operations here not supported by torch.compile (for instance,
6 -`to_list` in xformers attn, or `.item()` in flash attention)
6 +`.item()` in flash attention)
7
8 Using these ops and wrapping vision blocks with `torch.compile` can speed up
9 throughput in vision models by ~5% relative on H100, and improve token
@@ -19,42 +19,6 @@ import torch.nn.functional as F
19 from vllm.utils.torch_utils import direct_register_custom_op
20
21
22 -def xformers_attn_seqlens_wrapper(
23 - q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, seqlens: torch.Tensor
24 -) -> torch.Tensor:
25 - from xformers import ops as xops
26 - from xformers.ops.fmha.attn_bias import BlockDiagonalMask
27 -
28 - attn_bias = BlockDiagonalMask.from_seqlens(
29 - q_seqlen=seqlens.tolist(), kv_seqlen=None, device=q.device
30 - )
31 - context_layer = xops.memory_efficient_attention_forward(
32 - q, k, v, attn_bias=attn_bias, p=0, scale=None
33 - )
34 - context_layer = einops.rearrange(context_layer, "b s h d -> s b (h d)").contiguous()
35 - return context_layer
36 -
37 -
38 -def xformers_attn_seqlens_wrapper_fake(
39 - q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, seqlens: torch.Tensor
40 -) -> torch.Tensor:
41 - b, s, h, d = q.shape
42 - return torch.empty((s, b, h * d), dtype=q.dtype, device=q.device)
43 -
44 -
45 -direct_register_custom_op(
46 - op_name="xformers_attn_seqlens_wrapper",
47 - op_func=xformers_attn_seqlens_wrapper,
48 - fake_impl=xformers_attn_seqlens_wrapper_fake,
49 -)
50 -
51 -
52 -def vit_xformers_attn_wrapper(
53 - q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, seqlens: torch.Tensor
54 -) -> torch.Tensor:
55 - return torch.ops.vllm.xformers_attn_seqlens_wrapper(q, k, v, seqlens)
56 -
57 -
22 def flash_attn_maxseqlen_wrapper(
23 q: torch.Tensor,
24 k: torch.Tensor,
vllm/attention/selector.py
+8 -1
@@ -36,7 +36,14 @@ def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
36 * None otherwise
37 """
38 backend_name = os.environ.get(STR_BACKEND_ENV_VAR)
39 - return None if backend_name is None else AttentionBackendEnum[backend_name]
39 + if backend_name is None:
40 + return None
41 + if backend_name == "XFORMERS":
42 + raise ValueError(
43 + "Attention backend 'XFORMERS' has been removed (See PR #29262 for "
44 + "details). Please select a supported attention backend."
45 + )
46 + return AttentionBackendEnum[backend_name]
47
48
49 # Global state allows a particular choice of backend
vllm/config/multimodal.py
+6
@@ -173,6 +173,12 @@ class MultiModalConfig:
173 # We need to import the real type here (deferred to avoid circular import).
174 from vllm.attention.backends.registry import AttentionBackendEnum
175
176 + if isinstance(value, str) and value.upper() == "XFORMERS":
177 + raise ValueError(
178 + "Attention backend 'XFORMERS' has been removed (See PR #29262 for "
179 + "details). Please select a supported attention backend."
180 + )
181 +
182 if value is None or isinstance(value, AttentionBackendEnum):
183 return value
184
vllm/envs.py
-1
@@ -640,7 +640,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
640 # Example options:
641 # - "TORCH_SDPA": use torch.nn.MultiheadAttention
642 # - "FLASH_ATTN": use FlashAttention
643 - # - "XFORMERS": use XFormers
643 # - "FLASHINFER": use flashinfer
644 # - "FLASHMLA": use FlashMLA
645 # - "FLASH_ATTN_MLA": use FlashAttention for MLA
vllm/model_executor/models/dots_ocr.py
+4 -23
@@ -306,7 +306,6 @@ class DotsVisionAttention(nn.Module):
306 if self.attn_backend not in {
307 AttentionBackendEnum.FLASH_ATTN,
308 AttentionBackendEnum.TORCH_SDPA,
309 - AttentionBackendEnum.XFORMERS,
309 AttentionBackendEnum.ROCM_AITER_FA,
310 }:
311 raise RuntimeError(
@@ -324,7 +323,6 @@ class DotsVisionAttention(nn.Module):
323 rotary_pos_emb: torch.Tensor | None = None,
324 *,
325 max_seqlen: int | None = None,
327 - seqlens: list[int] | None = None,
326 ) -> torch.Tensor:
327 # [S, C] -> [S, B=1, C]
328 x = hidden_states.unsqueeze(1)
@@ -374,16 +372,6 @@ class DotsVisionAttention(nn.Module):
372 out_i = out_i.permute(0, 2, 1, 3)
373 outputs.append(out_i)
374 context_layer = torch.cat(outputs, dim=1) if outputs else q[:, :0]
377 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
378 - from xformers import ops as xops
379 - from xformers.ops.fmha.attn_bias import BlockDiagonalMask
380 -
381 - attn_bias = BlockDiagonalMask.from_seqlens(
382 - q_seqlen=seqlens, kv_seqlen=None, device=q.device
383 - )
384 - context_layer = xops.memory_efficient_attention_forward(
385 - q, k, v, attn_bias=attn_bias, p=0, scale=None
386 - )
375 else:
376 raise RuntimeError("Unsupported attention backend")
377
@@ -545,14 +533,12 @@ class DotsVisionBlock(nn.Module):
533 cu_seqlens: torch.Tensor,
534 rotary_pos_emb: torch.Tensor,
535 max_seqlen: int | None = None,
548 - seqlens: list[int] | None = None,
536 ) -> torch.Tensor:
537 hidden_states = hidden_states + self.attn(
538 self.norm1(hidden_states),
539 cu_seqlens=cu_seqlens,
540 rotary_pos_emb=rotary_pos_emb,
541 max_seqlen=max_seqlen,
555 - seqlens=seqlens,
542 )
543 hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
544 return hidden_states
@@ -663,18 +649,14 @@ class DotsVisionTransformer(nn.Module):
649 rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
650 return rotary_pos_emb
651
666 - def compute_attn_mask_seqlen(
667 - self, cu_seqlens: torch.Tensor
668 - ) -> tuple[int | None, list[int] | None]:
669 - max_seqlen, seqlens = None, None
652 + def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> int | None:
653 + max_seqlen = None
654 if (
655 self.attn_backend == AttentionBackendEnum.FLASH_ATTN
656 or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
657 ):
658 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
675 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
676 - seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
677 - return max_seqlen, seqlens
659 + return max_seqlen
660
661 def forward(
662 self, hidden_states: torch.Tensor, grid_thw: list[list[int]]
@@ -694,14 +676,13 @@ class DotsVisionTransformer(nn.Module):
676 )
677 cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
678
697 - max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
679 + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
680 for blk in self.blocks:
681 hidden_states = blk(
682 hidden_states,
683 cu_seqlens=cu_seqlens,
684 rotary_pos_emb=rotary_pos_emb,
685 max_seqlen=max_seqlen,
704 - seqlens=seqlens,
686 )
687
688 if self.post_trunk_norm is not None:
vllm/model_executor/models/ernie45_vl.py
+5 -28
@@ -214,7 +214,6 @@ class Ernie4_5_VisionAttention(nn.Module):
214 if self.attn_backend not in {
215 AttentionBackendEnum.FLASH_ATTN,
216 AttentionBackendEnum.TORCH_SDPA,
217 - AttentionBackendEnum.XFORMERS,
217 AttentionBackendEnum.ROCM_AITER_FA,
218 }:
219 raise RuntimeError(
@@ -259,7 +258,6 @@ class Ernie4_5_VisionAttention(nn.Module):
258 cu_seqlens: torch.Tensor,
259 rotary_pos_emb: torch.Tensor,
260 max_seqlen: int | None = None, # Only used for Flash Attention
262 - seqlens: list[int] | None = None, # Only used for xFormers
261 ) -> torch.Tensor:
262 # [s, b, c] --> [s, b, head * 3 * head_dim]
263 x, _ = self.qkv(x)
@@ -311,20 +309,6 @@ class Ernie4_5_VisionAttention(nn.Module):
309 context_layer = rearrange(
310 context_layer, "b s h d -> s b (h d)"
311 ).contiguous()
314 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
315 - from xformers import ops as xops
316 - from xformers.ops.fmha.attn_bias import BlockDiagonalMask
317 -
318 - attn_bias = BlockDiagonalMask.from_seqlens(
319 - q_seqlen=seqlens, kv_seqlen=None, device=q.device
320 - )
321 -
322 - context_layer = xops.memory_efficient_attention_forward(
323 - q, k, v, attn_bias=attn_bias, p=0, scale=None
324 - )
325 - context_layer = rearrange(
326 - context_layer, "b s h d -> s b (h d)"
327 - ).contiguous()
312
313 output, _ = self.proj(context_layer)
314 return output
@@ -404,14 +388,12 @@ class Ernie4_5_VisionBlock(nn.Module):
388 cu_seqlens: torch.Tensor,
389 rotary_pos_emb: torch.Tensor,
390 max_seqlen: int | None = None, # Only used for Flash Attention
407 - seqlens: list[int] | None = None, # Only used for xFormers
391 ) -> torch.Tensor:
392 hidden_states = hidden_states + self.attn(
393 self.norm1(hidden_states),
394 cu_seqlens=cu_seqlens,
395 rotary_pos_emb=rotary_pos_emb,
396 max_seqlen=max_seqlen,
414 - seqlens=seqlens,
397 )
398 hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
399 return hidden_states
@@ -562,18 +544,14 @@ class Ernie4_5_VisionTransformer(nn.Module):
544 rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
545 return rotary_pos_emb
546
565 - def compute_attn_mask_seqlen(
566 - self, cu_seqlens: torch.Tensor
567 - ) -> tuple[int | None, list[int] | None]:
568 - max_seqlen, seqlens = None, None
547 + def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> int | None:
548 + max_seqlen = None
549 if (
550 self.attn_backend == AttentionBackendEnum.FLASH_ATTN
551 or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
552 ):
553 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
574 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
575 - seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
576 - return max_seqlen, seqlens
554 + return max_seqlen
555
556 def forward(
557 self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, num_pad=0
@@ -598,8 +576,8 @@ class Ernie4_5_VisionTransformer(nn.Module):
576 if hidden_states.ndim == 2:
577 hidden_states = hidden_states.unsqueeze(dim=1)
578
601 - # pre-compute seqlens for attn mask to reduce cuMemcpy operations
602 - max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
579 + # pre-compute max_seqlen for attn mask to reduce cuMemcpy operations
580 + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
581
582 for i, blk in enumerate(self.blocks):
583 hidden_states = blk(
@@ -607,7 +585,6 @@ class Ernie4_5_VisionTransformer(nn.Module):
585 cu_seqlens=cu_seqlens,
586 rotary_pos_emb=rotary_pos_emb,
587 max_seqlen=max_seqlen,
610 - seqlens=seqlens,
588 )
589
590 final_output = self.ln(hidden_states)
vllm/model_executor/models/glm4_1v.py
+6 -25
@@ -309,7 +309,6 @@ class Glm4vVisionAttention(nn.Module):
309 if self.attn_backend not in {
310 AttentionBackendEnum.FLASH_ATTN,
311 AttentionBackendEnum.TORCH_SDPA,
312 - AttentionBackendEnum.XFORMERS,
312 AttentionBackendEnum.ROCM_AITER_FA,
313 }:
314 raise RuntimeError(
@@ -345,7 +344,6 @@ class Glm4vVisionAttention(nn.Module):
344 rotary_pos_emb_cos: torch.Tensor,
345 rotary_pos_emb_sin: torch.Tensor,
346 max_seqlen: int | None = None, # Only used for Flash Attention
348 - seqlens: list[int] | None = None, # Only used for xFormers
347 ) -> torch.Tensor:
348 # [s, b, c] --> [s, b, head * 3 * head_dim]
349 x, _ = self.qkv(x)
@@ -400,20 +398,6 @@ class Glm4vVisionAttention(nn.Module):
398 context_layer = rearrange(
399 context_layer, "b s h d -> s b (h d)"
400 ).contiguous()
403 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
404 - from xformers import ops as xops
405 - from xformers.ops.fmha.attn_bias import BlockDiagonalMask
406 -
407 - attn_bias = BlockDiagonalMask.from_seqlens(
408 - q_seqlen=seqlens, kv_seqlen=None, device=q.device
409 - )
410 -
411 - context_layer = xops.memory_efficient_attention_forward(
412 - q, k, v, attn_bias=attn_bias, p=0, scale=None
413 - )
414 - context_layer = rearrange(
415 - context_layer, "b s h d -> s b (h d)"
416 - ).contiguous()
401
402 output, _ = self.proj(context_layer)
403 return output
@@ -461,7 +445,6 @@ class Glm4vVisionBlock(nn.Module):
445 rotary_pos_emb_cos: torch.Tensor,
446 rotary_pos_emb_sin: torch.Tensor,
447 max_seqlen: int | None = None, # Only used for Flash Attention
464 - seqlens: list[int] | None = None, # Only used for xFormers
448 ) -> torch.Tensor:
449 x_attn = self.attn(
450 self.norm1(x),
@@ -469,7 +452,6 @@ class Glm4vVisionBlock(nn.Module):
452 rotary_pos_emb_cos=rotary_pos_emb_cos,
453 rotary_pos_emb_sin=rotary_pos_emb_sin,
454 max_seqlen=max_seqlen,
472 - seqlens=seqlens,
455 )
456 x_fused_norm, residual = self.norm2(x, residual=x_attn)
457 x = residual + self.mlp(x_fused_norm)
@@ -803,15 +785,14 @@ class Glm4vVisionTransformer(nn.Module):
785 def compute_attn_mask_seqlen(
786 self,
787 cu_seqlens: torch.Tensor,
806 - ) -> tuple[int | None, list[int] | None]:
807 - max_seqlen, seqlens = None, None
808 - seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
788 + ) -> int | None:
789 + max_seqlen = None
790 if (
791 self.attn_backend == AttentionBackendEnum.FLASH_ATTN
792 or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
793 ):
794 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
814 - return max_seqlen, seqlens
795 + return max_seqlen
796
797 def forward(
798 self,
@@ -836,8 +817,9 @@ class Glm4vVisionTransformer(nn.Module):
817 ).cumsum(dim=0, dtype=torch.int32)
818 cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0)
819
839 - # pre-compute seqlens for attn mask to reduce cuMemcpy operations
840 - max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
820 + # pre-compute max_seqlen for attn mask to reduce cuMemcpy operations
821 + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
822 + seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
823 x = self.embeddings(
824 x, seqlens, grid_thw, image_type_ids[:, 0], image_type_ids[:, 1]
825 )
@@ -851,7 +833,6 @@ class Glm4vVisionTransformer(nn.Module):
833 rotary_pos_emb_cos=rotary_pos_emb_cos,
834 rotary_pos_emb_sin=rotary_pos_emb_sin,
835 max_seqlen=max_seqlen,
854 - seqlens=seqlens,
836 )
837
838 # adapter
vllm/model_executor/models/keye.py
+17 -13
@@ -9,6 +9,7 @@ from typing import Annotated, Any, Literal, TypeAlias, TypeVar
9 import numpy as np
10 import torch
11 import torch.nn as nn
12 +import torch.nn.functional as F
13 from einops import rearrange
14 from transformers import PretrainedConfig
15 from transformers.activations import GELUActivation
@@ -424,7 +425,7 @@ class KeyeSiglipAttention(nn.Module):
425
426 if self.attn_backend not in {
427 AttentionBackendEnum.FLASH_ATTN,
427 - AttentionBackendEnum.XFORMERS,
428 + AttentionBackendEnum.TORCH_SDPA,
429 AttentionBackendEnum.ROCM_AITER_FA,
430 }:
431 raise RuntimeError(
@@ -451,7 +452,6 @@ class KeyeSiglipAttention(nn.Module):
452 )
453
454 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
454 - seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
455 batch_size = q.shape[0]
456
457 if rope_emb is None:
@@ -498,17 +498,21 @@ class KeyeSiglipAttention(nn.Module):
498 softmax_scale=self.scale,
499 )
500 context_layer = rearrange(output, "(b s) ... -> b s ...", b=batch_size)
501 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
502 - from xformers import ops as xops
503 - from xformers.ops.fmha.attn_bias import BlockDiagonalMask
504 -
505 - attn_bias = BlockDiagonalMask.from_seqlens(
506 - q_seqlen=seqlens, kv_seqlen=None, device=q.device
507 - )
508 -
509 - context_layer = xops.memory_efficient_attention_forward(
510 - q, k, v, attn_bias=attn_bias, p=0, scale=None
511 - )
501 + elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
502 + outputs = []
503 + for i in range(1, len(cu_seqlens)):
504 + start_idx = cu_seqlens[i - 1]
505 + end_idx = cu_seqlens[i]
506 + q_i = q[:, start_idx:end_idx]
507 + k_i = k[:, start_idx:end_idx]
508 + v_i = v[:, start_idx:end_idx]
509 + q_i, k_i, v_i = (
510 + rearrange(x, "b s h d -> b h s d") for x in (q_i, k_i, v_i)
511 + )
512 + output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
513 + output_i = rearrange(output_i, "b h s d -> b s h d ")
514 + outputs.append(output_i)
515 + context_layer = torch.cat(outputs, dim=1) if outputs else q[:, :0]
516
517 context_layer = rearrange(context_layer, "b s h d -> b s (h d)").contiguous()
518
vllm/model_executor/models/paddleocr_vl.py
-13
@@ -38,7 +38,6 @@ from vllm.attention.layer import (
38 )
39 from vllm.attention.ops.vit_attn_wrappers import (
40 vit_flash_attn_wrapper,
41 - vit_xformers_attn_wrapper,
41 )
42 from vllm.config import VllmConfig
43 from vllm.config.multimodal import BaseDummyOptions
@@ -657,7 +656,6 @@ class SiglipAttention(nn.Module):
656 cu_seqlens: torch.Tensor,
657 rotary_pos_emb: torch.Tensor | None,
658 max_seqlen: torch.Tensor | None,
660 - seqlens: torch.Tensor | None,
659 ) -> torch.Tensor:
660 batch_size, _, _ = hidden_states.shape
661
@@ -703,10 +701,6 @@ class SiglipAttention(nn.Module):
701 context_layer = rearrange(
702 context_layer, "b s h d -> s b (h d)"
703 ).contiguous()
706 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
707 - if seqlens is None:
708 - raise ValueError("xFormers attention backend requires seqlens tensor.")
709 - context_layer = vit_xformers_attn_wrapper(q, k, v, seqlens)
704 else:
705 raise RuntimeError(
706 f"PaddleOCR-VL does not support {self.attn_backend} backend now."
@@ -818,7 +812,6 @@ class SiglipEncoderLayer(nn.Module):
812 cu_seqlens: torch.Tensor,
813 rotary_pos_emb: torch.Tensor | None,
814 max_seqlen: torch.Tensor | None,
821 - seqlens: torch.Tensor | None,
815 ) -> torch.Tensor:
816 residual = hidden_states
817
@@ -828,7 +821,6 @@ class SiglipEncoderLayer(nn.Module):
821 cu_seqlens=cu_seqlens,
822 rotary_pos_emb=rotary_pos_emb,
823 max_seqlen=max_seqlen,
831 - seqlens=seqlens,
824 )
825
826 hidden_states = residual + hidden_states
@@ -870,7 +862,6 @@ class SiglipEncoder(nn.Module):
862 if self.attn_backend not in {
863 AttentionBackendEnum.FLASH_ATTN,
864 AttentionBackendEnum.TORCH_SDPA,
873 - AttentionBackendEnum.XFORMERS,
865 AttentionBackendEnum.ROCM_AITER_FA,
866 }:
867 raise RuntimeError(
@@ -943,14 +934,11 @@ class SiglipEncoder(nn.Module):
934 cu_seqlens = cu_seqlens.to(device=device)
935
936 max_seqlen = None
946 - seqlens = None
937 if self.attn_backend in {
938 AttentionBackendEnum.FLASH_ATTN,
939 AttentionBackendEnum.ROCM_AITER_FA,
940 }:
941 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
952 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
953 - seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
942
943 hidden_states = inputs_embeds
944 for encoder_layer in self.layers:
@@ -959,7 +947,6 @@ class SiglipEncoder(nn.Module):
947 cu_seqlens=cu_seqlens,
948 rotary_pos_emb=rotary_pos_emb,
949 max_seqlen=max_seqlen,
962 - seqlens=seqlens,
950 )
951 return hidden_states
952
vllm/model_executor/models/pixtral.py
+1
@@ -74,6 +74,7 @@ from .vision import (
74 )
75
76 try:
77 + # Note: vLLM does not install xformers by default.
78 from xformers import ops as xops
79
80 if current_platform.is_cuda() and current_platform.has_device_capability(100):
vllm/model_executor/models/qwen2_5_vl.py
+4 -21
@@ -46,7 +46,6 @@ from vllm.attention.layer import maybe_get_vit_flash_attn_backend
46 from vllm.attention.ops.vit_attn_wrappers import (
47 vit_flash_attn_wrapper,
48 vit_torch_sdpa_wrapper,
49 - vit_xformers_attn_wrapper,
49 )
50 from vllm.compilation.decorators import support_torch_compile
51 from vllm.config import VllmConfig
@@ -375,7 +374,6 @@ class Qwen2_5_VisionAttention(nn.Module):
374 rotary_pos_emb_cos: torch.Tensor,
375 rotary_pos_emb_sin: torch.Tensor,
376 max_seqlen: torch.Tensor, # Only used for Flash Attention
378 - seqlens: torch.Tensor, # Only used for xFormers
377 ) -> torch.Tensor:
378 # [s, b, c] --> [s, b, head * 3 * head_dim]
379 x, _ = self.qkv(x)
@@ -435,8 +433,6 @@ class Qwen2_5_VisionAttention(nn.Module):
433 v,
434 cu_seqlens,
435 )
438 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
439 - context_layer = vit_xformers_attn_wrapper(q, k, v, seqlens)
436
437 output, _ = self.proj(context_layer)
438 return output
@@ -448,9 +444,7 @@ class Qwen2_5_VisionAttention(nn.Module):
444 "cu_seqlens": 0,
445 "rotary_pos_emb_cos": 0,
446 "rotary_pos_emb_sin": 0,
451 - "seqlens": 0,
447 },
453 - mark_unbacked_dims={"seqlens": 0},
448 enable_if=should_torch_compile_mm_vit,
449 )
450 class Qwen2_5_VisionBlock(nn.Module):
@@ -501,7 +495,6 @@ class Qwen2_5_VisionBlock(nn.Module):
495 rotary_pos_emb_cos: torch.Tensor,
496 rotary_pos_emb_sin: torch.Tensor,
497 max_seqlen: torch.Tensor, # Only used for Flash Attention
504 - seqlens: torch.Tensor, # Only used for xFormers
498 ) -> torch.Tensor:
499 x_attn = self.attn(
500 self.norm1(x),
@@ -509,7 +502,6 @@ class Qwen2_5_VisionBlock(nn.Module):
502 rotary_pos_emb_cos=rotary_pos_emb_cos,
503 rotary_pos_emb_sin=rotary_pos_emb_sin,
504 max_seqlen=max_seqlen,
512 - seqlens=seqlens,
505 )
506 x_fused_norm, residual = self.norm2(x, residual=x_attn)
507 x = residual + self.mlp(x_fused_norm)
@@ -670,7 +662,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
662 if self.attn_backend not in {
663 AttentionBackendEnum.FLASH_ATTN,
664 AttentionBackendEnum.TORCH_SDPA,
673 - AttentionBackendEnum.XFORMERS,
665 AttentionBackendEnum.ROCM_AITER_FA,
666 }:
667 raise RuntimeError(
@@ -822,17 +813,14 @@ class Qwen2_5_VisionTransformer(nn.Module):
813 def compute_attn_mask_seqlen(
814 self,
815 cu_seqlens: torch.Tensor,
825 - ) -> tuple[torch.Tensor, torch.Tensor]:
816 + ) -> torch.Tensor:
817 max_seqlen = torch.zeros([], device=cu_seqlens.device)
827 - seqlens = torch.zeros(1, device=cu_seqlens.device)
818 if self.attn_backend in {
819 AttentionBackendEnum.FLASH_ATTN,
820 AttentionBackendEnum.ROCM_AITER_FA,
821 }:
822 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
833 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
834 - seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
835 - return max_seqlen, seqlens
823 + return max_seqlen
824
825 @staticmethod
826 def invert_permutation(perm: torch.Tensor) -> torch.Tensor:
@@ -897,10 +885,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
885
886 # transformers
887 # pre-compute seqlens for window/full attn to reduce cuMemcpy operations
900 - max_seqlen_full, seqlens_full = self.compute_attn_mask_seqlen(cu_seqlens)
901 - max_seqlen_window, seqlens_window = self.compute_attn_mask_seqlen(
902 - cu_window_seqlens
903 - )
888 + max_seqlen_full = self.compute_attn_mask_seqlen(cu_seqlens)
889 + max_seqlen_window = self.compute_attn_mask_seqlen(cu_window_seqlens)
890
891 cu_seqlens = cu_seqlens.to(device=self.device, non_blocking=True)
892 cu_window_seqlens = cu_window_seqlens.to(device=self.device, non_blocking=True)
@@ -927,11 +913,9 @@ class Qwen2_5_VisionTransformer(nn.Module):
913 if layer_num in self.fullatt_block_indexes:
914 cu_seqlens_now = cu_seqlens
915 max_seqlen_now = max_seqlen_full
930 - seqlens_now = seqlens_full
916 else:
917 cu_seqlens_now = cu_window_seqlens
918 max_seqlen_now = max_seqlen_window
934 - seqlens_now = seqlens_window
919
920 hidden_states = blk(
921 hidden_states,
@@ -939,7 +923,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
923 rotary_pos_emb_cos=rotary_pos_emb_cos,
924 rotary_pos_emb_sin=rotary_pos_emb_sin,
925 max_seqlen=max_seqlen_now,
942 - seqlens=seqlens_now,
926 )
927
928 # For Qwen2.5-VL-3B, float16 will overflow at last block
vllm/model_executor/models/qwen2_vl.py
+4 -27
@@ -348,7 +348,6 @@ class Qwen2VisionAttention(nn.Module):
348 if self.attn_backend not in {
349 AttentionBackendEnum.FLASH_ATTN,
350 AttentionBackendEnum.TORCH_SDPA,
351 - AttentionBackendEnum.XFORMERS,
351 AttentionBackendEnum.ROCM_AITER_FA,
352 }:
353 raise RuntimeError(
@@ -384,7 +383,6 @@ class Qwen2VisionAttention(nn.Module):
383 rotary_pos_emb_cos: torch.Tensor,
384 rotary_pos_emb_sin: torch.Tensor,
385 max_seqlen: int | None = None, # Only used for Flash Attention
387 - seqlens: list[int] | None = None, # Only used for xFormers
386 ) -> torch.Tensor:
387 # [s, b, c] --> [s, b, 3 * head * head_dim]
388 x, _ = self.qkv(x)
@@ -445,20 +443,6 @@ class Qwen2VisionAttention(nn.Module):
443 context_layer = rearrange(
444 context_layer, "b s h d -> s b (h d)"
445 ).contiguous()
448 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
449 - from xformers import ops as xops
450 - from xformers.ops.fmha.attn_bias import BlockDiagonalMask
451 -
452 - attn_bias = BlockDiagonalMask.from_seqlens(
453 - q_seqlen=seqlens, kv_seqlen=None, device=q.device
454 - )
455 -
456 - context_layer = xops.memory_efficient_attention_forward(
457 - q, k, v, attn_bias=attn_bias, p=0, scale=None
458 - )
459 - context_layer = rearrange(
460 - context_layer, "b s h d -> s b (h d)"
461 - ).contiguous()
446
447 output, _ = self.proj(context_layer)
448 return output
@@ -509,7 +493,6 @@ class Qwen2VisionBlock(nn.Module):
493 rotary_pos_emb_cos: torch.Tensor,
494 rotary_pos_emb_sin: torch.Tensor,
495 max_seqlen: int | None = None, # Only used for Flash Attention
512 - seqlens: list[int] | None = None, # Only used for xFormers
496 ) -> torch.Tensor:
497 x = x + self.attn(
498 self.norm1(x),
@@ -517,7 +500,6 @@ class Qwen2VisionBlock(nn.Module):
500 rotary_pos_emb_cos=rotary_pos_emb_cos,
501 rotary_pos_emb_sin=rotary_pos_emb_sin,
502 max_seqlen=max_seqlen,
520 - seqlens=seqlens,
503 )
504
505 x = x + self.mlp(self.norm2(x))
@@ -728,18 +710,14 @@ class Qwen2VisionTransformer(nn.Module):
710 sin_combined = sin[pos_ids].flatten(1)
711 return cos_combined, sin_combined
712
731 - def compute_attn_mask_seqlen(
732 - self, cu_seqlens: torch.Tensor
733 - ) -> tuple[int | None, list[int] | None]:
734 - max_seqlen, seqlens = None, None
713 + def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> int | None:
714 + max_seqlen = None
715 if self.attn_backend in {
716 AttentionBackendEnum.FLASH_ATTN,
717 AttentionBackendEnum.ROCM_AITER_FA,
718 }:
719 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
740 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
741 - seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
742 - return max_seqlen, seqlens
720 + return max_seqlen
721
722 def forward(
723 self,
@@ -771,7 +749,7 @@ class Qwen2VisionTransformer(nn.Module):
749 x = x.unsqueeze(1)
750
751 # pre-compute seqlens for attn mask to reduce cuMemcpy operations
774 - max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
752 + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
753 cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
754 for blk in self.blocks:
755 x = blk(
@@ -780,7 +758,6 @@ class Qwen2VisionTransformer(nn.Module):
758 rotary_pos_emb_cos=rotary_pos_emb_cos,
759 rotary_pos_emb_sin=rotary_pos_emb_sin,
760 max_seqlen=max_seqlen,
783 - seqlens=seqlens,
761 )
762
763 # adapter
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+3 -9
@@ -224,7 +224,6 @@ class Qwen3_VisionBlock(nn.Module):
224 rotary_pos_emb_cos: torch.Tensor,
225 rotary_pos_emb_sin: torch.Tensor,
226 max_seqlen: torch.Tensor, # Only used for Flash Attention
227 - seqlens: torch.Tensor, # Only used for xFormers
227 ) -> torch.Tensor:
228 x = x + self.attn(
229 self.norm1(x),
@@ -232,7 +231,6 @@ class Qwen3_VisionBlock(nn.Module):
231 rotary_pos_emb_cos=rotary_pos_emb_cos,
232 rotary_pos_emb_sin=rotary_pos_emb_sin,
233 max_seqlen=max_seqlen,
235 - seqlens=seqlens,
234 )
235
236 x = x + self.mlp(self.norm2(x))
@@ -500,14 +498,11 @@ class Qwen3Omni_VisionTransformer(nn.Module):
498 def compute_attn_mask_seqlen(
499 self,
500 cu_seqlens: torch.Tensor,
503 - ) -> tuple[torch.Tensor, torch.Tensor]:
501 + ) -> torch.Tensor:
502 max_seqlen = torch.zeros([], device=cu_seqlens.device)
505 - seqlens = torch.zeros(1, device=cu_seqlens.device)
503 if self.attn_backend == AttentionBackendEnum.FLASH_ATTN:
504 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
508 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
509 - seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
510 - return max_seqlen, seqlens
505 + return max_seqlen
506
507 def forward(
508 self,
@@ -533,7 +528,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):
528 hidden_states = hidden_states.unsqueeze(1)
529 rotary_pos_emb_cos = rotary_pos_emb_cos.to(hidden_states.device)
530 rotary_pos_emb_sin = rotary_pos_emb_sin.to(hidden_states.device)
536 - max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
531 + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
532
533 hidden_states_list = []
534 deepstack_visual_indexes = self.deepstack_visual_indexes
@@ -545,7 +540,6 @@ class Qwen3Omni_VisionTransformer(nn.Module):
540 rotary_pos_emb_cos=rotary_pos_emb_cos,
541 rotary_pos_emb_sin=rotary_pos_emb_sin,
542 max_seqlen=max_seqlen,
548 - seqlens=seqlens,
543 )
544 if (
545 deepstack_visual_indexes is not None
vllm/model_executor/models/qwen3_vl.py
+3 -10
@@ -235,7 +235,6 @@ class Qwen3_VisionBlock(nn.Module):
235 rotary_pos_emb_cos: torch.Tensor,
236 rotary_pos_emb_sin: torch.Tensor,
237 max_seqlen: torch.Tensor, # Only used for Flash Attention
238 - seqlens: torch.Tensor, # Only used for xFormers
238 ) -> torch.Tensor:
239 x = x + self.attn(
240 self.norm1(x),
@@ -243,7 +242,6 @@ class Qwen3_VisionBlock(nn.Module):
242 rotary_pos_emb_cos=rotary_pos_emb_cos,
243 rotary_pos_emb_sin=rotary_pos_emb_sin,
244 max_seqlen=max_seqlen,
246 - seqlens=seqlens,
245 )
246
247 x = x + self.mlp(self.norm2(x))
@@ -391,7 +389,6 @@ class Qwen3_VisionTransformer(nn.Module):
389 if self.attn_backend not in {
390 AttentionBackendEnum.FLASH_ATTN,
391 AttentionBackendEnum.TORCH_SDPA,
394 - AttentionBackendEnum.XFORMERS,
392 AttentionBackendEnum.ROCM_AITER_FA,
393 }:
394 raise RuntimeError(
@@ -531,17 +528,14 @@ class Qwen3_VisionTransformer(nn.Module):
528 def compute_attn_mask_seqlen(
529 self,
530 cu_seqlens: torch.Tensor,
534 - ) -> tuple[torch.Tensor, torch.Tensor]:
531 + ) -> torch.Tensor:
532 max_seqlen = torch.zeros([], device=cu_seqlens.device)
536 - seqlens = torch.zeros(1, device=cu_seqlens.device)
533 if (
534 self.attn_backend == AttentionBackendEnum.FLASH_ATTN
535 or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
536 ):
537 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
542 - elif self.attn_backend == AttentionBackendEnum.XFORMERS:
543 - seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
544 - return max_seqlen, seqlens
538 + return max_seqlen
539
540 def forward(
541 self,
@@ -569,7 +563,7 @@ class Qwen3_VisionTransformer(nn.Module):
563 cu_seqlens = torch.from_numpy(cu_seqlens)
564
565 hidden_states = hidden_states.unsqueeze(1)
572 - max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
566 + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
567 cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
568
569 deepstack_feature_lists = []
@@ -580,7 +574,6 @@ class Qwen3_VisionTransformer(nn.Module):
574 rotary_pos_emb_cos=rotary_pos_emb_cos,
575 rotary_pos_emb_sin=rotary_pos_emb_sin,
576 max_seqlen=max_seqlen,
583 - seqlens=seqlens,
577 )
578 if layer_num in self.deepstack_visual_indexes:
579 deepstack_merger_idx = self.deepstack_visual_indexes.index(layer_num)
vllm/platforms/cuda.py
+1 -6
@@ -277,12 +277,7 @@ class CudaPlatformBase(Platform):
277 except ImportError:
278 pass
279
280 - if cls.has_device_capability(100):
281 - # xFormers doesn't support Blackwell, fall back to SDPA
282 - # See https://github.com/facebookresearch/xformers/issues/1317#issuecomment-3199392579 # noqa: E501
283 - return AttentionBackendEnum.TORCH_SDPA
284 - else:
285 - return AttentionBackendEnum.XFORMERS
280 + return AttentionBackendEnum.TORCH_SDPA
281
282 @classmethod
283 def get_valid_backends(
vllm/utils/__init__.py
-1
@@ -49,7 +49,6 @@ STR_BACKEND_ENV_VAR: str = "VLLM_ATTENTION_BACKEND"
49 # Possible string values of STR_BACKEND_ENV_VAR
50 # register, corresponding to possible backends
51 STR_FLASHINFER_ATTN_VAL: str = "FLASHINFER"
52 -STR_XFORMERS_ATTN_VAL: str = "XFORMERS"
52 STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
53 STR_INVALID_VAL: str = "INVALID"
54
vllm/v1/attention/backends/xformers.py deleted
-420
@@ -1,420 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Attention layer with XFormersAttention."""
4 -
5 -from dataclasses import dataclass
6 -from typing import ClassVar, Optional
7 -
8 -import torch
9 -
10 -from vllm.attention.backends.abstract import (
11 - AttentionBackend,
12 - AttentionImpl,
13 - AttentionType,
14 - MultipleOf,
15 -)
16 -from vllm.attention.ops.triton_unified_attention import unified_attention
17 -from vllm.config import VllmConfig
18 -from vllm.logger import init_logger
19 -from vllm.v1.attention.backends.utils import (
20 - AttentionMetadataBuilder,
21 - CommonAttentionMetadata,
22 - split_decodes_and_prefills,
23 -)
24 -from vllm.v1.kv_cache_interface import AttentionSpec
25 -
26 -try:
27 - from xformers import ops as xops
28 - from xformers.ops.fmha.attn_bias import (
29 - AttentionBias,
30 - PagedBlockDiagonalCausalWithOffsetPaddedKeysMask,
31 - )
32 -
33 - XFORMERS_AVAILABLE = True
34 -except ImportError:
35 - XFORMERS_AVAILABLE = False
36 -
37 -from vllm import _custom_ops as ops
38 -
39 -logger = init_logger(__name__)
40 -
41 -
42 -class XFormersAttentionBackend(AttentionBackend):
43 - accept_output_buffer: bool = True
44 - supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
45 -
46 - @staticmethod
47 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
48 - return [MultipleOf(16)]
49 -
50 - @classmethod
51 - def get_supported_head_sizes(cls) -> list[int]:
52 - return [
53 - 32,
54 - 40,
55 - 48,
56 - 56,
57 - 64,
58 - 72,
59 - 80,
60 - 88,
61 - 96,
62 - 104,
63 - 112,
64 - 120,
65 - 128,
66 - 136,
67 - 144,
68 - 152,
69 - 160,
70 - 168,
71 - 176,
72 - 184,
73 - 192,
74 - 200,
75 - 208,
76 - 216,
77 - 224,
78 - 232,
79 - 240,
80 - 248,
81 - 256,
82 - ]
83 -
84 - @staticmethod
85 - def get_name() -> str:
86 - return "XFORMERS"
87 -
88 - @staticmethod
89 - def get_impl_cls() -> type["XFormersAttentionImpl"]:
90 - return XFormersAttentionImpl
91 -
92 - @staticmethod
93 - def get_kv_cache_shape(
94 - num_blocks: int,
95 - block_size: int,
96 - num_kv_heads: int,
97 - head_size: int,
98 - cache_dtype_str: str = "auto",
99 - ) -> tuple[int, ...]:
100 - if block_size % 16 != 0:
101 - raise ValueError("Block size must be a multiple of 16.")
102 - return (2, num_blocks, block_size, num_kv_heads, head_size)
103 -
104 - @staticmethod
105 - def get_builder_cls() -> type["XFormersAttentionMetadataBuilder"]:
106 - return XFormersAttentionMetadataBuilder
107 -
108 - @staticmethod
109 - def use_cascade_attention(*args, **kwargs) -> bool:
110 - return False
111 -
112 -
113 -@dataclass
114 -class XFormersAttentionMetadata:
115 - num_actual_tokens: int # Number of tokens excluding padding.
116 - max_query_len: int
117 - query_start_loc: torch.Tensor
118 - max_seq_len: int
119 - seq_lens: torch.Tensor
120 - block_table: torch.Tensor
121 - slot_mapping: torch.Tensor
122 -
123 - num_prefill_tokens: int = 0
124 - num_decode_tokens: int = 0
125 - num_prefills: int = 0
126 - num_decodes: int = 0
127 -
128 - # Biases for different attention types.
129 - attn_bias: Optional["AttentionBias"] = None
130 -
131 - # Self-attention prefill/decode metadata cache
132 - _cached_prefill_metadata: Optional["XFormersAttentionMetadata"] = None
133 - _cached_decode_metadata: Optional["XFormersAttentionMetadata"] = None
134 -
135 - @property
136 - def prefill_metadata(self) -> Optional["XFormersAttentionMetadata"]:
137 - if self.num_prefills == 0:
138 - return None
139 -
140 - if self._cached_prefill_metadata is not None:
141 - # Recover cached prefill-phase attention
142 - # metadata structure
143 - return self._cached_prefill_metadata
144 -
145 - q_start_loc = self.query_start_loc[self.num_decodes :]
146 - q_seqlens = torch.diff(q_start_loc)
147 - kv_seqlens = self.seq_lens[self.num_decodes :]
148 - # Construct & cache prefill-phase attention metadata structure
149 - self._cached_prefill_metadata = XFormersAttentionMetadata(
150 - num_actual_tokens=self.num_prefill_tokens,
151 - max_query_len=int(q_seqlens.max().item()),
152 - query_start_loc=q_start_loc - q_start_loc[0],
153 - max_seq_len=int(kv_seqlens.max().item()),
154 - seq_lens=kv_seqlens,
155 - block_table=self.block_table[self.num_decodes :],
156 - slot_mapping=self.slot_mapping[self.num_decode_tokens :],
157 - )
158 - return self._cached_prefill_metadata
159 -
160 - @property
161 - def decode_metadata(self) -> Optional["XFormersAttentionMetadata"]:
162 - if self.num_decode_tokens == 0:
163 - return None
164 -
165 - if self._cached_decode_metadata is not None:
166 - # Recover cached decode-phase attention
167 - # metadata structure
168 - return self._cached_decode_metadata
169 -
170 - q_start_loc = self.query_start_loc
171 - q_seqlens = torch.diff(q_start_loc)
172 - decode_kv_seqlens = self.seq_lens[: self.num_decodes]
173 - # Construct & cache decode-phase attention metadata structure
174 - self._cached_decode_metadata = XFormersAttentionMetadata(
175 - num_actual_tokens=self.num_decode_tokens,
176 - max_query_len=int(q_seqlens[: self.num_decodes].max().item()),
177 - query_start_loc=q_start_loc[: self.num_decodes + 1],
178 - max_seq_len=int(decode_kv_seqlens.max().item()),
179 - seq_lens=decode_kv_seqlens,
180 - block_table=self.block_table[: self.num_decodes],
181 - slot_mapping=self.slot_mapping[: self.num_decode_tokens],
182 - attn_bias=self.attn_bias,
183 - )
184 - return self._cached_decode_metadata
185 -
186 -
187 -class XFormersAttentionMetadataBuilder(
188 - AttentionMetadataBuilder[XFormersAttentionMetadata]
189 -):
190 - reorder_batch_threshold: int = 1
191 -
192 - def __init__(
193 - self,
194 - kv_cache_spec: AttentionSpec,
195 - layer_names: list[str],
196 - vllm_config: VllmConfig,
197 - device: torch.device,
198 - ):
199 - super().__init__(kv_cache_spec, layer_names, vllm_config, device)
200 -
201 - assert XFORMERS_AVAILABLE
202 - self.block_size = kv_cache_spec.block_size
203 - self._num_decodes = 0
204 - self._num_decode_tokens = 0
205 -
206 - def build(
207 - self,
208 - common_prefix_len: int,
209 - common_attn_metadata: CommonAttentionMetadata,
210 - fast_build: bool = False,
211 - ) -> XFormersAttentionMetadata:
212 - num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
213 - split_decodes_and_prefills(
214 - common_attn_metadata, decode_threshold=self.reorder_batch_threshold
215 - )
216 - )
217 -
218 - num_actual_tokens = common_attn_metadata.num_actual_tokens
219 - q_start_loc = common_attn_metadata.query_start_loc
220 - q_seqlens = torch.diff(q_start_loc)
221 - max_query_len = common_attn_metadata.max_query_len
222 - kv_seqlens = common_attn_metadata.seq_lens
223 - max_seq_len = common_attn_metadata.max_seq_len
224 - block_table = common_attn_metadata.block_table_tensor
225 - slot_mapping = common_attn_metadata.slot_mapping
226 -
227 - bias = None
228 - if num_decodes > 0:
229 - # Construct the decoder bias.
230 - decode_q_seqlens = q_seqlens[:num_decodes]
231 - decode_kv_seqlens = kv_seqlens[:num_decodes]
232 - bias = PagedBlockDiagonalCausalWithOffsetPaddedKeysMask.from_seqlens(
233 - q_seqlen=decode_q_seqlens.tolist(),
234 - kv_seqlen=decode_kv_seqlens.tolist(),
235 - page_size=self.block_size,
236 - block_tables=block_table[:num_decodes],
237 - device=block_table.device,
238 - )
239 -
240 - return XFormersAttentionMetadata(
241 - num_actual_tokens=num_actual_tokens,
242 - num_prefill_tokens=num_prefill_tokens,
243 - num_decode_tokens=num_decode_tokens,
244 - num_prefills=num_prefills,
245 - num_decodes=num_decodes,
246 - max_query_len=max_query_len,
247 - query_start_loc=q_start_loc,
248 - max_seq_len=max_seq_len,
249 - seq_lens=kv_seqlens,
250 - block_table=block_table,
251 - slot_mapping=slot_mapping,
252 - attn_bias=bias,
253 - )
254 -
255 -
256 -class XFormersAttentionImpl(AttentionImpl):
257 - def __init__(
258 - self,
259 - num_heads: int,
260 - head_size: int,
261 - scale: float,
262 - num_kv_heads: int,
263 - alibi_slopes: list[float] | None,
264 - sliding_window: int | None,
265 - kv_cache_dtype: str,
266 - logits_soft_cap: float | None = None,
267 - attn_type: AttentionType = AttentionType.DECODER,
268 - kv_sharing_target_layer_name: str | None = None,
269 - ) -> None:
270 - if kv_sharing_target_layer_name is not None:
271 - raise NotImplementedError("KV sharing is not supported in V0.")
272 - if alibi_slopes is not None:
273 - raise NotImplementedError("XFormers does not support alibi slopes yet.")
274 - self.num_heads = num_heads
275 - self.head_size = head_size
276 - self.scale = float(scale)
277 - self.num_kv_heads = num_kv_heads
278 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads
279 - self.kv_cache_dtype = kv_cache_dtype
280 - self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
281 - if alibi_slopes is not None:
282 - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
283 - self.alibi_slopes = alibi_slopes
284 - if sliding_window is None:
285 - self.sliding_window = (-1, -1)
286 - else:
287 - self.sliding_window = (sliding_window - 1, 0)
288 - if logits_soft_cap is None:
289 - # Setting logits_soft_cap to 0 means no soft cap.
290 - logits_soft_cap = 0
291 - self.logits_soft_cap = logits_soft_cap
292 -
293 - if attn_type != AttentionType.DECODER:
294 - raise NotImplementedError(
295 - "Encoder self-attention and "
296 - "encoder/decoder cross-attention "
297 - "are not implemented for "
298 - "XFormersAttentionImpl."
299 - )
300 -
301 - def forward(
302 - self,
303 - layer: torch.nn.Module,
304 - query: torch.Tensor,
305 - key: torch.Tensor,
306 - value: torch.Tensor,
307 - kv_cache: torch.Tensor,
308 - attn_metadata: XFormersAttentionMetadata,
309 - output: torch.Tensor | None = None,
310 - output_scale: torch.Tensor | None = None,
311 - output_block_scale: torch.Tensor | None = None,
312 - ) -> torch.Tensor:
313 - """Forward pass with XFormers.
314 -
315 - Args:
316 - query: shape = [num_tokens, num_heads, head_size]
317 - key: shape = [num_tokens, num_kv_heads, head_size]
318 - value: shape = [num_tokens, num_kv_heads, head_size]
319 - kv_cache: shape =
320 - [2, num_blocks, block_size, num_kv_heads, head_size]
321 - attn_metadata: Metadata for attention.
322 - Returns:
323 - shape = [num_tokens, num_heads * head_size]
324 - """
325 - assert output is not None, "Output tensor must be provided."
326 -
327 - if output_scale is not None or output_block_scale is not None:
328 - raise NotImplementedError(
329 - "fused output quantization is not yet supported"
330 - " for XFormersAttentionImpl"
331 - )
332 -
333 - if attn_metadata is None:
334 - # Profiling run.
335 - return output.fill_(0)
336 -
337 - # Cache the input KVs.
338 - key_cache, value_cache = kv_cache.unbind(0)
339 - if self.kv_sharing_target_layer_name is None:
340 - # Reshape the input keys and values and store them in the cache.
341 - # Skip this if sharing KV cache with an earlier attention layer.
342 - # NOTE(woosuk): Here, key and value are padded while slot_mapping is
343 - # not padded. However, we don't need to do key[:num_actual_tokens]
344 - # and value[:num_actual_tokens] because the reshape_and_cache_flash
345 - # op uses the slot_mapping's shape to determine the number of
346 - # actual tokens.
347 - ops.reshape_and_cache_flash(
348 - key,
349 - value,
350 - key_cache,
351 - value_cache,
352 - attn_metadata.slot_mapping,
353 - self.kv_cache_dtype,
354 - layer._k_scale,
355 - layer._v_scale,
356 - )
357 -
358 - num_actual_tokens = attn_metadata.num_actual_tokens
359 - num_decode_tokens = attn_metadata.num_decode_tokens
360 - if prefill_meta := attn_metadata.prefill_metadata:
361 - descale_shape = (prefill_meta.query_start_loc.shape[0] - 1, key.shape[1])
362 - unified_attention(
363 - q=query[num_decode_tokens:num_actual_tokens],
364 - k=key_cache,
365 - v=value_cache,
366 - out=output[num_decode_tokens:num_actual_tokens],
367 - cu_seqlens_q=prefill_meta.query_start_loc,
368 - max_seqlen_q=prefill_meta.max_query_len,
369 - seqused_k=prefill_meta.seq_lens,
370 - max_seqlen_k=prefill_meta.max_seq_len,
371 - softmax_scale=self.scale,
372 - causal=True,
373 - alibi_slopes=self.alibi_slopes,
374 - window_size=self.sliding_window,
375 - block_table=prefill_meta.block_table,
376 - softcap=self.logits_soft_cap,
377 - q_descale=None, # Not supported
378 - k_descale=layer._k_scale.expand(descale_shape),
379 - v_descale=layer._v_scale.expand(descale_shape),
380 - )
381 -
382 - if decode_meta := attn_metadata.decode_metadata:
383 - # Query for decode. KV is not needed because it is already cached.
384 - decode_query = query[:num_decode_tokens]
385 - # Reshape query to [1, B_T, G, H, D].
386 - q = decode_query.view(
387 - 1, -1, self.num_kv_heads, self.num_queries_per_kv, self.head_size
388 - )
389 - # Reshape the k and v caches to [1, Bkv_T, G, H, D]
390 - cache_k = key_cache.view(
391 - 1, -1, self.num_kv_heads, 1, self.head_size
392 - ).expand(
393 - 1,
394 - -1,
395 - self.num_kv_heads,
396 - self.num_queries_per_kv,
397 - self.head_size,
398 - )
399 - cache_v = value_cache.view(
400 - 1, -1, self.num_kv_heads, 1, self.head_size
401 - ).expand(
402 - 1,
403 - -1,
404 - self.num_kv_heads,
405 - self.num_queries_per_kv,
406 - self.head_size,
407 - )
408 -
409 - attn_bias = decode_meta.attn_bias
410 - output[:num_decode_tokens] = xops.memory_efficient_attention_forward(
411 - q,
412 - cache_k,
413 - cache_v,
414 - attn_bias=attn_bias,
415 - p=0.0,
416 - scale=self.scale,
417 - ).view(decode_query.shape)
418 -
419 - # Reshape the output tensor.
420 - return output