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 10, 10:15 UTC
[1/N][Attention] Restructure attention: move files (#31916)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Jan 9, 2026 at 16:10 UTC
2612ba9285d828f3179801334b6d2db0046b9b6c
195 files changed
+425
-395
.buildkite/test-amd.yaml
+1
-1
@@ -1105,8 +1105,8 @@ steps:
1105
- vllm/v1/attention/backends/flashinfer.py
1106
- vllm/v1/attention/backends/mla/cutlass_mla.py
1107
- vllm/v1/attention/backends/mla/flashinfer_mla.py
1108
+ - vllm/v1/attention/selector.py
1109
- vllm/platforms/cuda.py
1109
- - vllm/attention/selector.py
1110
commands:
1111
- nvidia-smi
1112
- python3 examples/offline_inference/basic/chat.py
.buildkite/test-pipeline.yaml
+1
-1
@@ -954,8 +954,8 @@ steps:
954
- vllm/v1/attention/backends/flashinfer.py
955
- vllm/v1/attention/backends/mla/cutlass_mla.py
956
- vllm/v1/attention/backends/mla/flashinfer_mla.py
957
+ - vllm/v1/attention/selector.py
958
- vllm/platforms/cuda.py
958
- - vllm/attention/selector.py
959
commands:
960
- nvidia-smi
961
- python3 examples/offline_inference/basic/chat.py
.buildkite/test_areas/kernels.yaml
+1
-1
@@ -90,8 +90,8 @@ steps:
90
- vllm/v1/attention/backends/flashinfer.py
91
- vllm/v1/attention/backends/mla/cutlass_mla.py
92
- vllm/v1/attention/backends/mla/flashinfer_mla.py
93
+ - vllm/v1/attention/selector.py
94
- vllm/platforms/cuda.py
94
- - vllm/attention/selector.py
95
commands:
96
- nvidia-smi
97
- python3 examples/offline_inference/basic/chat.py
.github/CODEOWNERS
+4
-4
@@ -3,7 +3,6 @@
3
4
# This lists cover the "core" components of vLLM that require careful review
5
/vllm/attention @LucasWilkinson
6
-/vllm/attention/backends/abstract.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @njhill
6
/vllm/executor/executor_base.py @zhuohan123 @youkaichao @alexm-redhat @njhill @22quinn
7
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
8
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
@@ -27,6 +26,7 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
26
27
# vLLM V1
28
/vllm/v1/attention @LucasWilkinson
29
+/vllm/v1/attention/backend.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @njhill
30
/vllm/v1/attention/backends/mla @pavanimajety
31
/vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety
32
/vllm/v1/attention/backends/triton_attn.py @tdoublep
@@ -117,15 +117,15 @@ mkdocs.yaml @hmellor
117
/vllm/transformers_utils/tokenizers/mistral.py @patrickvonplaten
118
119
# Kernels
120
-/vllm/attention/ops/chunked_prefill_paged_decode.py @tdoublep
121
-/vllm/attention/ops/triton_unified_attention.py @tdoublep
120
+/vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep
121
+/vllm/v1/attention/ops/triton_unified_attention.py @tdoublep
122
123
# ROCm related: specify owner with write access to notify AMD folks for careful code review
124
/vllm/**/*rocm* @tjtanaa
125
/docker/Dockerfile.rocm* @gshtras @tjtanaa
126
/vllm/v1/attention/backends/rocm*.py @gshtras @tjtanaa
127
/vllm/v1/attention/backends/mla/rocm*.py @gshtras @tjtanaa
128
-/vllm/attention/ops/rocm*.py @gshtras @tjtanaa
128
+/vllm/v1/attention/ops/rocm*.py @gshtras @tjtanaa
129
/vllm/model_executor/layers/fused_moe/rocm*.py @gshtras @tjtanaa
130
/csrc/rocm @gshtras @tjtanaa
131
/requirements/*rocm* @tjtanaa
.github/mergify.yml
+2
-2
@@ -222,10 +222,10 @@ pull_request_rules:
222
- files~=^csrc/rocm/
223
- files~=^docker/Dockerfile.rocm
224
- files~=^requirements/rocm.*\.txt
225
- - files~=^vllm/attention/backends/rocm.*\.py
226
- - files~=^vllm/attention/ops/rocm.*\.py
225
- files~=^vllm/model_executor/layers/fused_moe/rocm.*\.py
226
+ - files~=^vllm/v1/attention/backends/rocm.*\.py
227
- files~=^vllm/v1/attention/backends/mla/rocm.*\.py
228
+ - files~=^vllm/v1/attention/ops/rocm.*\.py
229
- files~=^tests/kernels/.*_rocm.*\.py
230
- files=vllm/platforms/rocm.py
231
- title~=(?i)AMD
benchmarks/kernels/benchmark_reshape_and_cache_flash.py
+3
-3
@@ -7,9 +7,6 @@ import torch
7
from tabulate import tabulate
8
9
from vllm import _custom_ops as ops
10
-from vllm.attention.ops.triton_reshape_and_cache_flash import (
11
- triton_reshape_and_cache_flash,
12
-)
10
from vllm.logger import init_logger
11
from vllm.utils.argparse_utils import FlexibleArgumentParser
12
from vllm.utils.torch_utils import (
@@ -17,6 +14,9 @@ from vllm.utils.torch_utils import (
14
create_kv_caches_with_random_flash,
15
set_random_seed,
16
)
17
+from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
18
+ triton_reshape_and_cache_flash,
19
+)
20
21
logger = init_logger(__name__)
22
docs/contributing/model/basic.md
+1
-1
@@ -142,7 +142,7 @@ We use "mamba-like" to refer to layers that posses a state that is updated in-pl
142
For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`.
143
It is also necessary to implement the "attention meta-data" class which handles the meta-data that is common across all layers.
144
Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
145
-It is also worth noting that we should update `MAMBA_TYPE_TO_BACKEND_MAP` and `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/attention/backends/registry.py) when adding a new mamba backend.
145
+It is also worth noting that we should update `MAMBA_TYPE_TO_BACKEND_MAP` and `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend.
146
Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it.
147
Please see the calls to `direct_register_custom_op` in [vllm/model_executor/models/minimax_text_01.py](../../../vllm/model_executor/models/minimax_text_01.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
148
The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended.
docs/design/custom_op.md
+1
-1
@@ -60,7 +60,7 @@ For example:
60
**1. Attention:**
61
62
```python
63
---8<-- "vllm/attention/layers/mm_encoder_attention.py:mm_encoder_attn"
63
+--8<-- "vllm/model_executor/layers/attention/mm_encoder_attention.py:mm_encoder_attn"
64
65
--8<-- "vllm/model_executor/layers/mla.py:multi_head_latent_attention"
66
```
docs/design/plugin_system.md
+2
-2
@@ -124,7 +124,7 @@ Every plugin has three parts:
124
125
Please look at the worker base class [WorkerBase][vllm.v1.worker.worker_base.WorkerBase] for more functions that can be implemented.
126
127
-5. Implement the attention backend class `MyDummyAttention` in `my_dummy_attention.py`. The attention backend class should inherit from [AttentionBackend][vllm.attention.backends.abstract.AttentionBackend]. It's used to calculate attentions with your device. Take `vllm.v1.attention.backends` as examples, it contains many attention backend implementations.
127
+5. Implement the attention backend class `MyDummyAttention` in `my_dummy_attention.py`. The attention backend class should inherit from [AttentionBackend][vllm.v1.attention.backend.AttentionBackend]. It's used to calculate attentions with your device. Take `vllm.v1.attention.backends` as examples, it contains many attention backend implementations.
128
129
6. Implement custom ops for high performance. Most ops can be ran by pytorch native implementation, while the performance may not be good. In this case, you can implement specific custom ops for your plugins. Currently, there are kinds of custom ops vLLM supports:
130
@@ -153,5 +153,5 @@ The interface for the model/module may change during vLLM's development. If you
153
154
!!! warning "Deprecations"
155
- `use_v1` parameter in `Platform.get_attn_backend_cls` is deprecated. It has been removed in v0.13.0.
156
- - `_Backend` in `vllm.attention` is deprecated. It has been removed in v0.13.0. Please use `vllm.attention.backends.registry.register_backend` to add new attention backend to `AttentionBackendEnum` instead.
156
+ - `_Backend` in `vllm.attention` is deprecated. It has been removed in v0.13.0. Please use `vllm.v1.attention.backends.registry.register_backend` to add new attention backend to `AttentionBackendEnum` instead.
157
- `seed_everything` platform interface is deprecated. It will be removed in v0.15.0 or later. Please use `vllm.utils.torch_utils.set_random_seed` instead.
examples/offline_inference/basic/embed.py
+1
-1
@@ -4,10 +4,10 @@
4
from argparse import Namespace
5
6
from vllm import LLM, EngineArgs
7
-from vllm.attention.backends.registry import AttentionBackendEnum
7
from vllm.config import AttentionConfig
8
from vllm.platforms import current_platform
9
from vllm.utils.argparse_utils import FlexibleArgumentParser
10
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
11
12
13
def parse_args():
examples/offline_inference/basic/score.py
+1
-1
@@ -4,10 +4,10 @@
4
from argparse import Namespace
5
6
from vllm import LLM, EngineArgs
7
-from vllm.attention.backends.registry import AttentionBackendEnum
7
from vllm.config import AttentionConfig
8
from vllm.platforms import current_platform
9
from vllm.utils.argparse_utils import FlexibleArgumentParser
10
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
11
12
13
def parse_args():
tests/compile/fullgraph/test_full_cudagraph.py
+1
-1
@@ -9,10 +9,10 @@ import pytest
9
from tests.utils import wait_for_gpu_memory_to_clear
10
from tests.v1.attention.utils import full_cg_backend_configs as backend_configs
11
from vllm import LLM, SamplingParams
12
-from vllm.attention.backends.registry import AttentionBackendEnum
12
from vllm.config import CompilationConfig
13
from vllm.platforms import current_platform
14
from vllm.utils.torch_utils import is_torch_equal_or_newer
15
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
16
17
18
@contextlib.contextmanager
tests/compile/fullgraph/test_full_graph.py
+1
-1
@@ -10,10 +10,10 @@ import torch
10
11
from tests.quantization.utils import is_quant_method_supported
12
from vllm import LLM, SamplingParams
13
-from vllm.attention.backends.registry import AttentionBackendEnum
13
from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig
14
from vllm.platforms import current_platform
15
from vllm.utils.torch_utils import is_torch_equal_or_newer
16
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
17
18
from ...utils import create_new_process_for_each_test
19
tests/compile/test_fusion_attn.py
+2
-2
@@ -9,8 +9,6 @@ from tests.compile.backend import LazyInitPass, TestBackend
9
from tests.utils import flat_product
10
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
11
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
12
-from vllm.attention.backends.abstract import AttentionMetadata
13
-from vllm.attention.backends.registry import AttentionBackendEnum
12
from vllm.attention.layer import Attention
13
from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
14
from vllm.compilation.fx_utils import find_op_nodes
@@ -37,6 +35,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
35
from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp
36
from vllm.platforms import current_platform
37
from vllm.utils.flashinfer import has_flashinfer
38
+from vllm.v1.attention.backend import AttentionMetadata
39
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
40
from vllm.v1.kv_cache_interface import AttentionSpec
41
42
FP8_DTYPE = current_platform.fp8_dtype()
tests/compile/test_qk_norm_rope_fusion.py
+1
-1
@@ -5,7 +5,6 @@ import pytest
5
import torch
6
7
from tests.compile.backend import TestBackend
8
-from vllm.attention.backends.abstract import AttentionType
8
from vllm.attention.layer import Attention
9
from vllm.compilation.matcher_utils import FLASHINFER_ROTARY_OP, RMS_OP, ROTARY_OP
10
from vllm.compilation.noop_elimination import NoOpEliminationPass
@@ -25,6 +24,7 @@ from vllm.config import (
24
from vllm.model_executor.layers.layernorm import RMSNorm
25
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
26
from vllm.platforms import current_platform
27
+from vllm.v1.attention.backend import AttentionType
28
29
RSQRT_OP = torch.ops.aten.rsqrt.default
30
INDEX_SELECT_OP = torch.ops.aten.index.Tensor
tests/config/test_multimodal_config.py
+1
-1
@@ -3,8 +3,8 @@
3
4
import pytest
5
6
-from vllm.attention.backends.registry import AttentionBackendEnum
6
from vllm.config.multimodal import MultiModalConfig
7
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
8
9
10
def test_mm_encoder_attn_backend_str_conversion():
tests/engine/test_arg_utils.py
+1
-1
@@ -299,7 +299,7 @@ def test_compilation_config():
299
300
301
def test_attention_config():
302
- from vllm.attention.backends.registry import AttentionBackendEnum
302
+ from vllm.v1.attention.backends.registry import AttentionBackendEnum
303
304
parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
305
tests/kernels/attention/test_aiter_flash_attn.py
+1
-1
@@ -6,9 +6,9 @@ import pytest
6
import torch
7
8
import vllm.v1.attention.backends.rocm_aiter_fa # noqa: F401
9
-from vllm.attention.utils.fa_utils import is_flash_attn_varlen_func_available
9
from vllm.platforms import current_platform
10
from vllm.utils.torch_utils import set_random_seed
11
+from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
12
13
NUM_HEADS = [(4, 4), (8, 2)]
14
HEAD_SIZES = [128, 256]
tests/kernels/attention/test_attention.py
+2
-2
@@ -10,7 +10,7 @@ from tests.kernels.allclose_default import get_default_atol, get_default_rtol
10
from tests.kernels.utils import opcheck
11
from vllm import _custom_ops as ops
12
from vllm.attention.layer import Attention
13
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
13
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
14
from vllm.platforms import current_platform
15
from vllm.utils.mem_utils import get_max_shared_memory_bytes
16
from vllm.utils.torch_utils import set_random_seed
@@ -30,7 +30,7 @@ NUM_PREFILL_SEQS = [3] # Arbitrary values for testing
30
NUM_HEADS = [(40, 40), (64, 8)] # Arbitrary values for testing
31
32
# This should be sync with get_supported_head_sizes() in
33
-# vllm.attention.ops.paged_attn.PagedAttention
33
+# vllm.v1.attention.ops.paged_attn.PagedAttention
34
HEAD_SIZES = [32, 80, 128, 256]
35
36
BLOCK_SIZES = [16, 32]
tests/kernels/attention/test_attention_selector.py
+3
-3
@@ -6,13 +6,13 @@ from unittest.mock import patch
6
import pytest
7
import torch
8
9
-from vllm.attention.backends.registry import AttentionBackendEnum
10
-from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
9
from vllm.config import AttentionConfig, VllmConfig, set_current_vllm_config
10
from vllm.platforms import current_platform
11
from vllm.platforms.cpu import CpuPlatform
12
from vllm.platforms.cuda import CudaPlatform
13
from vllm.platforms.rocm import RocmPlatform
14
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
15
+from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend
16
17
18
@pytest.fixture(autouse=True)
@@ -182,7 +182,7 @@ def test_backend_selection(
182
expected = name
183
assert backend.get_name() == expected
184
elif name == "FLASH_ATTN_MLA":
185
- from vllm.attention.utils.fa_utils import (
185
+ from vllm.v1.attention.backends.fa_utils import (
186
flash_attn_supports_mla,
187
)
188
tests/kernels/attention/test_cache.py
+1
-1
@@ -270,7 +270,7 @@ def test_reshape_and_cache_flash(
270
v_scale,
271
)
272
elif implementation == "triton":
273
- from vllm.attention.ops.triton_reshape_and_cache_flash import (
273
+ from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
274
triton_reshape_and_cache_flash,
275
)
276
tests/kernels/attention/test_flashmla.py
+2
-2
@@ -7,12 +7,12 @@ import random
7
import pytest
8
import torch
9
10
-from vllm.attention.ops.flashmla import (
10
+from vllm.triton_utils import triton
11
+from vllm.v1.attention.ops.flashmla import (
12
flash_mla_with_kvcache,
13
get_mla_metadata,
14
is_flashmla_dense_supported,
15
)
15
-from vllm.triton_utils import triton
16
17
18
def cal_diff(
tests/kernels/attention/test_flashmla_sparse.py
+3
-3
@@ -5,7 +5,7 @@ import torch
5
6
7
def test_sparse_flashmla_metadata_smoke():
8
- import vllm.attention.ops.flashmla as fm
8
+ import vllm.v1.attention.ops.flashmla as fm
9
10
ok, reason = fm.is_flashmla_sparse_supported()
11
if not ok:
@@ -34,7 +34,7 @@ def test_sparse_flashmla_metadata_smoke():
34
35
36
def test_sparse_flashmla_decode_smoke():
37
- import vllm.attention.ops.flashmla as fm
37
+ import vllm.v1.attention.ops.flashmla as fm
38
39
ok, reason = fm.is_flashmla_sparse_supported()
40
if not ok:
@@ -97,7 +97,7 @@ def test_sparse_flashmla_decode_smoke():
97
98
99
def test_sparse_flashmla_prefill_smoke():
100
- import vllm.attention.ops.flashmla as fm
100
+ import vllm.v1.attention.ops.flashmla as fm
101
102
ok, reason = fm.is_flashmla_sparse_supported()
103
if not ok:
tests/kernels/attention/test_merge_attn_states.py
+2
-2
@@ -5,10 +5,10 @@ import pytest
5
import torch
6
7
from vllm._custom_ops import merge_attn_states as merge_attn_states_cuda
8
-from vllm.attention.ops.triton_merge_attn_states import (
8
+from vllm.platforms import current_platform
9
+from vllm.v1.attention.ops.triton_merge_attn_states import (
10
merge_attn_states as merge_attn_states_triton,
11
)
11
-from vllm.platforms import current_platform
12
13
14
# Naive PyTorch Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
tests/kernels/attention/test_mha_attn.py
+3
-3
@@ -12,14 +12,14 @@ from unittest.mock import patch
12
import pytest
13
import torch
14
15
-from vllm.attention.backends.registry import AttentionBackendEnum
16
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
17
-from vllm.attention.selector import _cached_get_attn_backend
15
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
16
from vllm.platforms import current_platform
17
from vllm.platforms.cpu import CpuPlatform
18
from vllm.platforms.cuda import CudaPlatform
19
from vllm.platforms.rocm import RocmPlatform
20
from vllm.utils.torch_utils import set_random_seed
21
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
22
+from vllm.v1.attention.selector import _cached_get_attn_backend
23
24
25
@pytest.fixture(autouse=True)
tests/kernels/attention/test_pack_unpack_triton.py
+1
-1
@@ -4,7 +4,7 @@
4
import torch
5
from torch.testing import assert_close
6
7
-from vllm.attention.ops.common import pack_seq_triton, unpack_seq_triton
7
+from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
8
9
10
def test_pack_seq_basic_fp8():
tests/kernels/attention/test_prefix_prefill.py
+4
-2
@@ -10,10 +10,12 @@ import pytest
10
import torch
11
import torch.nn.functional as F
12
13
-from vllm.attention.ops.chunked_prefill_paged_decode import chunked_prefill_paged_decode
14
-from vllm.attention.ops.prefix_prefill import context_attention_fwd
13
from vllm.platforms import current_platform
14
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
15
+from vllm.v1.attention.ops.chunked_prefill_paged_decode import (
16
+ chunked_prefill_paged_decode,
17
+)
18
+from vllm.v1.attention.ops.prefix_prefill import context_attention_fwd
19
20
NUM_HEADS = [64]
21
NUM_QUERIES_PER_KV = [1, 64]
tests/kernels/attention/test_rocm_attention_selector.py
+3
-3
@@ -4,10 +4,10 @@
4
import pytest
5
import torch
6
7
-from vllm.attention.backends.registry import AttentionBackendEnum
8
-from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
7
from vllm.config import AttentionConfig, VllmConfig, set_current_vllm_config
8
from vllm.platforms.rocm import RocmPlatform
9
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
10
+from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend
11
12
13
@pytest.fixture(autouse=True)
@@ -19,7 +19,7 @@ def clear_cache():
19
@pytest.mark.skip(reason="Skipped for now. Should be revisited.")
20
def test_selector(monkeypatch: pytest.MonkeyPatch):
21
# Set the current platform to ROCm using monkeypatch
22
- monkeypatch.setattr("vllm.attention.selector.current_platform", RocmPlatform())
22
+ monkeypatch.setattr("vllm.v1.attention.selector.current_platform", RocmPlatform())
23
24
# Test standard ROCm attention
25
attention_config = AttentionConfig(backend=AttentionBackendEnum.ROCM_ATTN)
tests/kernels/attention/test_triton_decode_attention.py
+1
-1
@@ -4,8 +4,8 @@
4
import pytest
5
import torch
6
7
-from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
7
from vllm.utils.math_utils import cdiv
8
+from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
9
10
11
@pytest.mark.parametrize("B", [3, 5])
tests/kernels/attention/test_triton_prefill_attention.py
+1
-1
@@ -5,7 +5,7 @@ import pytest
5
import torch
6
import torch.nn.functional as F
7
8
-from vllm.attention.ops.triton_prefill_attention import context_attention_fwd
8
+from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
9
10
11
def ref_masked_attention(
tests/kernels/attention/test_triton_unified_attention.py
+1
-1
@@ -5,10 +5,10 @@
5
import pytest
6
import torch
7
8
-from vllm.attention.ops.triton_unified_attention import unified_attention
8
from vllm.platforms import current_platform
9
from vllm.utils.math_utils import next_power_of_2
10
from vllm.utils.torch_utils import set_random_seed
11
+from vllm.v1.attention.ops.triton_unified_attention import unified_attention
12
13
NUM_HEADS = [(4, 4), (8, 2)]
14
HEAD_SIZES = [128, 256]
tests/kernels/utils.py
+1
-1
@@ -13,11 +13,11 @@ import torch
13
from torch._prims_common import TensorLikeType
14
15
from tests.kernels.quant_utils import native_w8a8_block_matmul
16
-from vllm.attention.backends.abstract import AttentionType
16
from vllm.model_executor.custom_op import CustomOp
17
from vllm.model_executor.layers.activation import SiluAndMul
18
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
19
from vllm.utils.torch_utils import make_tensor_with_pad
20
+from vllm.v1.attention.backend import AttentionType
21
22
# For now, disable "test_aot_dispatch_dynamic" since there are some
23
# bugs related to this test in PyTorch 2.4.
tests/models/multimodal/generation/test_vit_backend_functionality.py
+1
-1
@@ -14,10 +14,10 @@ import pytest
14
from transformers import AutoProcessor
15
16
from vllm import LLM, EngineArgs, SamplingParams
17
-from vllm.attention.backends.registry import AttentionBackendEnum
17
from vllm.multimodal.utils import encode_image_url
18
from vllm.multimodal.video import sample_frames_from_video
19
from vllm.platforms import current_platform
20
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
21
22
from ....utils import create_new_process_for_each_test
23
from ...utils import dummy_hf_overrides
tests/models/quantization/test_fp8.py
+1
-1
@@ -9,7 +9,7 @@ Note: these tests will only pass on L4 GPU.
9
import pytest
10
11
from tests.quantization.utils import is_quant_method_supported
12
-from vllm.attention.utils.fa_utils import flash_attn_supports_fp8
12
+from vllm.v1.attention.backends.fa_utils import flash_attn_supports_fp8
13
from vllm.platforms import current_platform
14
from ..utils import check_logprobs_close
15
tests/test_attention_backend_registry.py
+2
-2
@@ -1,10 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from vllm.attention.backends.abstract import (
3
+from vllm.v1.attention.backend import (
4
AttentionBackend,
5
AttentionImpl,
6
)
7
-from vllm.attention.backends.registry import (
7
+from vllm.v1.attention.backends.registry import (
8
AttentionBackendEnum,
9
MambaAttentionBackendEnum,
10
register_backend,
tests/v1/attention/test_attention_backends.py
+2
-2
@@ -15,8 +15,6 @@ from tests.v1.attention.utils import (
15
create_vllm_config,
16
try_get_attention_backend,
17
)
18
-from vllm.attention.backends.abstract import AttentionType
19
-from vllm.attention.backends.registry import AttentionBackendEnum
18
from vllm.config import ModelConfig
19
from vllm.platforms import current_platform
20
from vllm.utils.math_utils import cdiv
@@ -25,6 +23,8 @@ from vllm.utils.torch_utils import (
23
is_torch_equal_or_newer,
24
set_random_seed,
25
)
26
+from vllm.v1.attention.backend import AttentionType
27
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
28
from vllm.v1.attention.backends.utils import (
29
CommonAttentionMetadata,
30
set_kv_cache_layout,
tests/v1/attention/test_mla_backends.py
+3
-3
@@ -18,15 +18,15 @@ from tests.v1.attention.utils import (
18
try_get_attention_backend,
19
)
20
from vllm import _custom_ops as ops
21
-from vllm.attention.backends.registry import AttentionBackendEnum
22
-from vllm.attention.ops.flashmla import is_flashmla_dense_supported
23
-from vllm.attention.utils.fa_utils import flash_attn_supports_mla
21
from vllm.config.vllm import set_current_vllm_config
22
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
23
from vllm.utils.math_utils import cdiv
24
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
25
+from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
26
from vllm.v1.attention.backends.mla.common import QueryLenSupport
27
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
28
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
29
+from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
30
from vllm.v1.kv_cache_interface import FullAttentionSpec
31
32
BACKENDS_TO_TEST = [
tests/v1/attention/test_rocm_attention_backends_selection.py
+2
-2
@@ -7,9 +7,9 @@ from unittest.mock import MagicMock, patch
7
import pytest
8
import torch
9
10
-from vllm.attention.backends.registry import AttentionBackendEnum
11
-from vllm.attention.selector import AttentionSelectorConfig
10
from vllm.platforms import current_platform
11
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
12
+from vllm.v1.attention.selector import AttentionSelectorConfig
13
14
# ROCm-specific attention backend selection tests
15
pytestmark = pytest.mark.skipif(
tests/v1/attention/test_sparse_mla_backends.py
+1
-1
@@ -21,7 +21,6 @@ from tests.v1.attention.utils import (
21
create_vllm_config,
22
)
23
from vllm import _custom_ops as ops
24
-from vllm.attention.ops import flashmla
24
from vllm.config import set_current_vllm_config
25
from vllm.model_executor.layers.linear import ColumnParallelLinear
26
from vllm.platforms import current_platform
@@ -31,6 +30,7 @@ from vllm.v1.attention.backends.mla.flashmla_sparse import (
30
triton_convert_req_index_to_global_index,
31
)
32
from vllm.v1.attention.backends.utils import split_prefill_chunks
33
+from vllm.v1.attention.ops import flashmla
34
35
SPARSE_BACKEND_BATCH_SPECS = {
36
name: BATCH_SPECS[name]
tests/v1/attention/utils.py
+2
-2
@@ -7,8 +7,6 @@ from dataclasses import dataclass
7
import pytest
8
import torch
9
10
-from vllm.attention.backends.abstract import AttentionImpl
11
-from vllm.attention.backends.registry import AttentionBackendEnum
10
from vllm.config import (
11
CacheConfig,
12
CompilationConfig,
@@ -20,6 +18,8 @@ from vllm.config import (
18
VllmConfig,
19
)
20
from vllm.config.model import ModelDType
21
+from vllm.v1.attention.backend import AttentionImpl
22
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
23
from vllm.v1.attention.backends.utils import (
24
AttentionMetadataBuilder,
25
CommonAttentionMetadata,
tests/v1/determinism/utils.py
+1
-1
@@ -6,9 +6,9 @@ import random
6
import pytest
7
import torch
8
9
-from vllm.attention.utils.fa_utils import flash_attn_supports_mla
9
from vllm.platforms import current_platform
10
from vllm.utils.flashinfer import has_flashinfer
11
+from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
12
13
skip_unsupported = pytest.mark.skipif(
14
not (current_platform.is_cuda() and current_platform.has_device_capability(80)),
tests/v1/kv_connector/unit/test_backwards_compatibility.py
+1
-1
@@ -14,12 +14,12 @@ from unittest.mock import patch
14
15
import pytest
16
17
-from vllm.attention.backends.abstract import AttentionMetadata
17
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
18
from vllm.distributed.kv_transfer.kv_connector.v1 import (
19
KVConnectorBase_V1,
20
KVConnectorRole,
21
)
22
+from vllm.v1.attention.backend import AttentionMetadata
23
from vllm.v1.core.sched.output import SchedulerOutput
24
25
from .utils import create_scheduler, create_vllm_config
tests/v1/spec_decode/test_eagle.py
+1
-1
@@ -13,7 +13,6 @@ from tests.v1.attention.utils import (
13
create_standard_kv_cache_spec,
14
try_get_attention_backend,
15
)
16
-from vllm.attention.backends.registry import AttentionBackendEnum
16
from vllm.config import (
17
AttentionConfig,
18
CacheConfig,
@@ -27,6 +26,7 @@ from vllm.config import (
26
from vllm.config.load import LoadConfig
27
from vllm.model_executor.models.llama import LlamaForCausalLM
28
from vllm.platforms import current_platform
29
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
30
from vllm.v1.spec_decode.eagle import EagleProposer
31
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
32
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
tests/v1/spec_decode/test_mtp.py
+1
-1
@@ -12,7 +12,6 @@ from tests.v1.attention.utils import (
12
create_standard_kv_cache_spec,
13
try_get_attention_backend,
14
)
15
-from vllm.attention.backends.registry import AttentionBackendEnum
15
from vllm.config import (
16
CacheConfig,
17
DeviceConfig,
@@ -25,6 +24,7 @@ from vllm.config import (
24
from vllm.config.load import LoadConfig
25
from vllm.model_executor.models.llama import LlamaForCausalLM
26
from vllm.platforms import current_platform
27
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
28
from vllm.v1.spec_decode.eagle import EagleProposer
29
30
mimo_7b_dir = "XiaomiMiMo/MiMo-7B-Base"
tests/v1/spec_decode/test_tree_attention.py
+2
-2
@@ -11,9 +11,9 @@ from tests.v1.attention.utils import (
11
create_vllm_config,
12
try_get_attention_backend,
13
)
14
-from vllm.attention.backends.registry import AttentionBackendEnum
15
-from vllm.attention.utils.fa_utils import is_flash_attn_varlen_func_available
14
from vllm.config import ParallelConfig, SpeculativeConfig
15
+from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
16
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
17
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
18
19
if not is_flash_attn_varlen_func_available():
tests/v1/worker/test_gpu_model_runner.py
+2
-2
@@ -5,8 +5,6 @@ import numpy as np
5
import pytest
6
import torch
7
8
-from vllm.attention.backends.abstract import MultipleOf
9
-from vllm.attention.backends.registry import AttentionBackendEnum
8
from vllm.attention.layer import Attention
9
from vllm.config import (
10
AttentionConfig,
@@ -27,6 +25,8 @@ from vllm.sampling_params import SamplingParams
25
from vllm.utils.mem_constants import GiB_bytes
26
from vllm.utils.system_utils import update_environment_variables
27
from vllm.utils.torch_utils import set_random_seed
28
+from vllm.v1.attention.backend import MultipleOf
29
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
30
from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs
31
from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
32
from vllm.v1.kv_cache_interface import (
tools/pre_commit/mypy.py
+3
-1
@@ -73,7 +73,9 @@ EXCLUDE = [
73
"vllm/model_executor/models",
74
"vllm/model_executor/layers/fla/ops",
75
# Ignore triton kernels in ops.
76
- "vllm/attention/ops",
76
+ "vllm/v1/attention/ops",
77
+ # TODO(matt): remove.
78
+ "vllm/v1/attention/backends/fa_utils.py",
79
]
80
81
vllm/attention/layer.py
+7
-7
@@ -8,13 +8,6 @@ import torch
8
import torch.nn as nn
9
10
import vllm.envs as envs
11
-from vllm.attention.backends.abstract import (
12
- AttentionBackend,
13
- AttentionType,
14
- MLAAttentionImpl,
15
-)
16
-from vllm.attention.backends.registry import AttentionBackendEnum
17
-from vllm.attention.selector import get_attn_backend
11
from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
12
from vllm.attention.utils.kv_transfer_utils import maybe_transfer_kv_layer
13
from vllm.config import CacheConfig, get_current_vllm_config
@@ -37,6 +30,13 @@ from vllm.utils.torch_utils import (
30
direct_register_custom_op,
31
kv_cache_dtype_str_to_dtype,
32
)
33
+from vllm.v1.attention.backend import (
34
+ AttentionBackend,
35
+ AttentionType,
36
+ MLAAttentionImpl,
37
+)
38
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
39
+from vllm.v1.attention.selector import get_attn_backend
40
from vllm.v1.kv_cache_interface import (
41
FullAttentionSpec,
42
KVCacheSpec,
vllm/attention/ops/__init__.py
vllm/config/attention.py
+1
-1
@@ -6,9 +6,9 @@ from typing import Any, Literal
6
from pydantic import field_validator
7
from pydantic.dataclasses import dataclass
8
9
-from vllm.attention.backends.registry import AttentionBackendEnum
9
from vllm.config.utils import config
10
from vllm.logger import init_logger
11
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
12
13
logger = init_logger(__name__)
14
vllm/config/model.py
+1
-1
@@ -12,7 +12,6 @@ from pydantic import ConfigDict, Field, field_validator, model_validator
12
from pydantic.dataclasses import dataclass
13
14
import vllm.envs as envs
15
-from vllm.attention.backends.registry import AttentionBackendEnum
15
from vllm.config.model_arch import (
16
ModelArchitectureConfig,
17
)
@@ -50,6 +49,7 @@ from vllm.transformers_utils.model_arch_config_convertor import (
49
from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
50
from vllm.transformers_utils.utils import maybe_model_redirect
51
from vllm.utils.import_utils import LazyLoader
52
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
53
54
if TYPE_CHECKING:
55
from transformers import PretrainedConfig
vllm/config/multimodal.py
+2
-2
@@ -7,9 +7,9 @@ from typing import Any, Literal, TypeAlias
7
from pydantic import ConfigDict, Field, field_validator, model_validator
8
from pydantic.dataclasses import dataclass
9
10
-from vllm.attention.backends.registry import AttentionBackendEnum
10
from vllm.config.utils import config
11
from vllm.utils.hashing import safe_hash
12
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
13
14
15
@dataclass
@@ -124,7 +124,7 @@ class MultiModalConfig:
124
mm_encoder_attn_backend: AttentionBackendEnum | None = None
125
"""Optional override for the multi-modal encoder attention backend when
126
using vision transformers. Accepts any value from
127
- `vllm.attention.backends.registry.AttentionBackendEnum` (e.g. `FLASH_ATTN`)."""
127
+ `vllm.v1.attention.backends.registry.AttentionBackendEnum` (e.g. `FLASH_ATTN`)."""
128
interleave_mm_strings: bool = False
129
"""Enable fully interleaved support for multimodal prompts, while using
130
--chat-template-content-format=string."""
vllm/distributed/kv_transfer/kv_connector/utils.py
+1
-1
@@ -10,10 +10,10 @@ from typing import TYPE_CHECKING, Literal
10
11
import torch
12
13
-from vllm.attention.backends.abstract import AttentionBackend
13
from vllm.config import get_current_vllm_config
14
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
15
from vllm.logger import init_logger
16
+from vllm.v1.attention.backend import AttentionBackend
17
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
18
19
if TYPE_CHECKING:
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+1
-1
@@ -42,8 +42,8 @@ from typing import TYPE_CHECKING, Any, Literal, Optional
42
43
import torch
44
45
-from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata
45
from vllm.logger import init_logger
46
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
47
from vllm.v1.core.sched.output import SchedulerOutput
48
from vllm.v1.outputs import KVConnectorOutput
49
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+1
-1
@@ -36,7 +36,6 @@ from typing import TYPE_CHECKING, Any, Optional
36
37
import torch
38
39
-from vllm.attention.backends.abstract import AttentionMetadata
39
from vllm.distributed.kv_transfer.kv_connector.v1 import (
40
KVConnectorBase_V1,
41
KVConnectorRole,
@@ -44,6 +43,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
43
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
44
from vllm.logger import init_logger
45
from vllm.utils.math_utils import cdiv
46
+from vllm.v1.attention.backend import AttentionMetadata
47
48
if TYPE_CHECKING:
49
from vllm.config import VllmConfig
vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
+1
-1
@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Any, Optional
7
import safetensors
8
import torch
9
10
-from vllm.attention.backends.abstract import AttentionMetadata
10
from vllm.config import VllmConfig
11
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
12
KVConnectorBase_V1,
@@ -16,6 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
15
)
16
from vllm.logger import init_logger
17
from vllm.utils.hashing import safe_hash
18
+from vllm.v1.attention.backend import AttentionMetadata
19
from vllm.v1.attention.backends.mla.common import MLACommonMetadata
20
from vllm.v1.core.sched.output import SchedulerOutput
21
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
+1
-1
@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any
5
6
import torch
7
8
-from vllm.attention.backends.abstract import AttentionMetadata
8
from vllm.config import VllmConfig
9
from vllm.distributed.kv_events import (
10
BlockStored,
@@ -19,6 +18,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
18
KVConnectorRole,
19
)
20
from vllm.logger import init_logger
21
+from vllm.v1.attention.backend import AttentionMetadata
22
from vllm.v1.core.sched.output import SchedulerOutput
23
from vllm.v1.outputs import KVConnectorOutput
24
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
+1
-1
@@ -36,7 +36,6 @@ except ImportError:
36
PluginLauncher as RuntimePluginLauncher,
37
)
38
39
-from vllm.attention.backends.abstract import AttentionMetadata
39
from vllm.config import VllmConfig
40
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
41
KVConnectorBase_V1,
@@ -54,6 +53,7 @@ from vllm.distributed.parallel_state import get_tensor_model_parallel_rank, get_
53
from vllm.sampling_params import SamplingParams
54
from vllm.utils.math_utils import cdiv
55
from vllm.utils.torch_utils import get_kv_cache_torch_dtype
56
+from vllm.v1.attention.backend import AttentionMetadata
57
from vllm.v1.core.sched.output import SchedulerOutput
58
from vllm.version import __version__ as VLLM_VERSION
59
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+1
-1
@@ -10,13 +10,13 @@ import zmq
10
from lmcache.integration.vllm.utils import mla_enabled
11
from lmcache.utils import init_logger as lmcache_init_logger
12
13
-from vllm.attention.backends.abstract import AttentionMetadata
13
from vllm.config import VllmConfig
14
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
15
KVConnectorBase_V1,
16
KVConnectorMetadata,
17
KVConnectorRole,
18
)
19
+from vllm.v1.attention.backend import AttentionMetadata
20
from vllm.v1.core.sched.output import SchedulerOutput
21
from vllm.v1.outputs import KVConnectorOutput
22
from vllm.v1.request import RequestStatus
vllm/distributed/kv_transfer/kv_connector/v1/mooncake_connector.py
+2
-2
@@ -16,8 +16,6 @@ import zmq
16
import zmq.asyncio
17
18
from vllm import envs
19
-from vllm.attention.backends.abstract import AttentionMetadata
20
-from vllm.attention.selector import get_attn_backend
19
from vllm.config import VllmConfig
20
from vllm.distributed.kv_transfer.kv_connector.utils import TpKVTopology
21
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@@ -33,7 +31,9 @@ from vllm.distributed.parallel_state import (
31
from vllm.forward_context import ForwardContext
32
from vllm.logger import init_logger
33
from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket
34
+from vllm.v1.attention.backend import AttentionMetadata
35
from vllm.v1.attention.backends.utils import get_kv_cache_layout
36
+from vllm.v1.attention.selector import get_attn_backend
37
from vllm.v1.core.sched.output import SchedulerOutput
38
from vllm.v1.request import RequestStatus
39
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+2
-2
@@ -15,7 +15,6 @@ import numpy as np
15
import torch
16
import zmq
17
18
-from vllm.attention.selector import get_attn_backend
18
from vllm.config import VllmConfig
19
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
20
KVConnectorBase_V1,
@@ -56,11 +55,12 @@ from vllm.utils.network_utils import (
55
make_zmq_path,
56
make_zmq_socket,
57
)
58
+from vllm.v1.attention.selector import get_attn_backend
59
from vllm.v1.core.sched.output import SchedulerOutput
60
from vllm.v1.request import RequestStatus
61
62
if TYPE_CHECKING:
63
- from vllm.attention.backends.abstract import AttentionMetadata
63
+ from vllm.v1.attention.backend import AttentionMetadata
64
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
65
from vllm.v1.kv_cache_interface import KVCacheConfig
66
from vllm.v1.request import Request
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+1
-1
@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Any
7
8
import torch
9
10
-from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata
10
from vllm.config import VllmConfig
11
from vllm.config.kv_transfer import KVTransferConfig
12
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType
@@ -24,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
23
PromMetricT,
24
)
25
from vllm.logger import init_logger
26
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
27
from vllm.v1.core.sched.output import SchedulerOutput
28
from vllm.v1.outputs import KVConnectorOutput
29
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+2
-2
@@ -20,8 +20,6 @@ import torch
20
import zmq
21
22
from vllm import envs
23
-from vllm.attention.backends.abstract import AttentionMetadata
24
-from vllm.attention.selector import get_attn_backend
23
from vllm.config import VllmConfig
24
from vllm.distributed.kv_transfer.kv_connector.utils import (
25
EngineId,
@@ -50,7 +48,9 @@ from vllm.forward_context import ForwardContext
48
from vllm.logger import init_logger
49
from vllm.platforms import current_platform
50
from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
51
+from vllm.v1.attention.backend import AttentionMetadata
52
from vllm.v1.attention.backends.utils import get_kv_cache_layout
53
+from vllm.v1.attention.selector import get_attn_backend
54
from vllm.v1.core.sched.output import SchedulerOutput
55
from vllm.v1.worker.block_table import BlockTable
56
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+1
-1
@@ -8,7 +8,6 @@ from typing import Any
8
9
import torch
10
11
-from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata
11
from vllm.attention.layer import Attention
12
from vllm.config import VllmConfig, get_layers_from_vllm_config
13
from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
@@ -20,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
19
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
20
from vllm.forward_context import ForwardContext
21
from vllm.logger import init_logger
22
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
23
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
24
from vllm.v1.core.kv_cache_utils import BlockHash
25
from vllm.v1.core.sched.output import SchedulerOutput
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py
+1
-1
@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Any, Optional
7
import regex as re
8
import torch
9
10
-from vllm.attention.backends.abstract import AttentionMetadata
10
from vllm.config import VllmConfig
11
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
12
KVConnectorBase_V1,
@@ -19,6 +18,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_engine import (
18
)
19
from vllm.distributed.parallel_state import get_world_group
20
from vllm.logger import init_logger
21
+from vllm.v1.attention.backend import AttentionMetadata
22
from vllm.v1.attention.backends.mla.common import MLACommonMetadata
23
from vllm.v1.core.sched.output import SchedulerOutput
24
vllm/engine/arg_utils.py
+1
-1
@@ -32,7 +32,6 @@ from pydantic.fields import FieldInfo
32
from typing_extensions import TypeIs
33
34
import vllm.envs as envs
35
-from vllm.attention.backends.registry import AttentionBackendEnum
35
from vllm.config import (
36
AttentionConfig,
37
CacheConfig,
@@ -94,6 +93,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
93
from vllm.utils.mem_constants import GiB_bytes
94
from vllm.utils.network_utils import get_ip
95
from vllm.utils.torch_utils import resolve_kv_cache_dtype_string
96
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
97
from vllm.v1.sample.logits_processor import LogitsProcessor
98
99
if TYPE_CHECKING:
vllm/envs.py
+1
-1
@@ -684,7 +684,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
684
None,
685
lambda: list(
686
__import__(
687
- "vllm.attention.backends.registry", fromlist=["AttentionBackendEnum"]
687
+ "vllm.v1.attention.backends.registry", fromlist=["AttentionBackendEnum"]
688
).AttentionBackendEnum.__members__.keys()
689
),
690
),
vllm/forward_context.py
+1
-1
@@ -10,10 +10,10 @@ from typing import Any, NamedTuple
10
import torch
11
12
import vllm.envs as envs
13
-from vllm.attention.backends.abstract import AttentionMetadata
13
from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
14
from vllm.logger import init_logger
15
from vllm.platforms import current_platform
16
+from vllm.v1.attention.backend import AttentionMetadata
17
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
18
from vllm.v1.worker.ubatch_utils import UBatchSlices
19
vllm/model_executor/layers/attention/__init__.py
renamed
vllm/model_executor/layers/attention/chunked_local_attention.py
renamed
+2
-2
@@ -4,12 +4,11 @@ import functools
4
5
import torch
6
7
-from vllm.attention.backends.abstract import AttentionBackend
7
from vllm.attention.layer import Attention
9
-from vllm.attention.selector import get_attn_backend
8
from vllm.config import CacheConfig
9
from vllm.config.vllm import VllmConfig
10
from vllm.model_executor.layers.quantization import QuantizationConfig
11
+from vllm.v1.attention.backend import AttentionBackend
12
from vllm.v1.attention.backends.utils import (
13
AttentionCGSupport,
14
AttentionMetadataBuilder,
@@ -17,6 +16,7 @@ from vllm.v1.attention.backends.utils import (
16
make_local_attention_virtual_batches,
17
subclass_attention_backend,
18
)
19
+from vllm.v1.attention.selector import get_attn_backend
20
from vllm.v1.kv_cache_interface import (
21
AttentionSpec,
22
ChunkedLocalAttentionSpec,
vllm/model_executor/layers/attention/cross_attention.py
renamed
+6
-6
@@ -6,20 +6,20 @@ from copy import copy
6
import numpy as np
7
import torch
8
9
-from vllm.attention.backends.abstract import (
10
- AttentionBackend,
11
- AttentionMetadata,
12
- AttentionType,
13
-)
9
from vllm.attention.layer import Attention
15
-from vllm.attention.selector import get_attn_backend
10
from vllm.config import CacheConfig, VllmConfig
11
from vllm.logger import init_logger
12
from vllm.utils.math_utils import cdiv
13
+from vllm.v1.attention.backend import (
14
+ AttentionBackend,
15
+ AttentionMetadata,
16
+ AttentionType,
17
+)
18
from vllm.v1.attention.backends.utils import (
19
CommonAttentionMetadata,
20
subclass_attention_backend,
21
)
22
+from vllm.v1.attention.selector import get_attn_backend
23
from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec
24
25
logger = init_logger(__name__)
vllm/model_executor/layers/attention/encoder_only_attention.py
renamed
+5
-5
@@ -5,19 +5,19 @@ from copy import copy
5
6
import torch
7
8
-from vllm.attention.backends.abstract import (
8
+from vllm.attention.layer import Attention
9
+from vllm.config import CacheConfig
10
+from vllm.config.vllm import VllmConfig
11
+from vllm.v1.attention.backend import (
12
AttentionBackend,
13
AttentionMetadata,
14
AttentionType,
15
)
13
-from vllm.attention.layer import Attention
14
-from vllm.attention.selector import get_attn_backend
15
-from vllm.config import CacheConfig
16
-from vllm.config.vllm import VllmConfig
16
from vllm.v1.attention.backends.utils import (
17
CommonAttentionMetadata,
18
subclass_attention_backend,
19
)
20
+from vllm.v1.attention.selector import get_attn_backend
21
from vllm.v1.kv_cache_interface import KVCacheSpec
22
23
vllm/model_executor/layers/attention/mm_encoder_attention.py
renamed
+6
-6
@@ -4,16 +4,16 @@
4
5
import torch
6
7
-from vllm.attention.backends.registry import AttentionBackendEnum
8
-from vllm.attention.ops.vit_attn_wrappers import (
9
- vit_flash_attn_wrapper,
10
- vit_torch_sdpa_wrapper,
11
-)
12
-from vllm.attention.utils.fa_utils import get_flash_attn_version
7
from vllm.config import MultiModalConfig
8
from vllm.logger import init_logger
9
from vllm.model_executor.custom_op import CustomOp
10
from vllm.model_executor.models.vision import get_vit_attn_backend
11
+from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
12
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
13
+from vllm.v1.attention.ops.vit_attn_wrappers import (
14
+ vit_flash_attn_wrapper,
15
+ vit_torch_sdpa_wrapper,
16
+)
17
18
logger = init_logger(__name__)
19
vllm/model_executor/layers/attention/static_sink_attention.py
renamed
+9
-9
@@ -4,26 +4,26 @@ import functools
4
5
import torch
6
7
-from vllm.attention.backends.abstract import (
8
- AttentionBackend,
9
- AttentionMetadata,
10
- AttentionType,
11
-)
7
from vllm.attention.layer import Attention
13
-from vllm.attention.ops.triton_reshape_and_cache_flash import (
14
- triton_reshape_and_cache_flash_diffkv,
15
-)
16
-from vllm.attention.selector import get_attn_backend
8
from vllm.config import CacheConfig, VllmConfig
9
from vllm.forward_context import ForwardContext, get_forward_context
10
from vllm.logger import init_logger
11
from vllm.model_executor.custom_op import CustomOp
12
from vllm.utils.math_utils import cdiv
13
from vllm.utils.torch_utils import direct_register_custom_op
14
+from vllm.v1.attention.backend import (
15
+ AttentionBackend,
16
+ AttentionMetadata,
17
+ AttentionType,
18
+)
19
from vllm.v1.attention.backends.utils import (
20
CommonAttentionMetadata,
21
subclass_attention_backend,
22
)
23
+from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
24
+ triton_reshape_and_cache_flash_diffkv,
25
+)
26
+from vllm.v1.attention.selector import get_attn_backend
27
from vllm.v1.kv_cache_interface import (
28
AttentionSpec,
29
KVCacheSpec,
vllm/model_executor/layers/attention_layer_base.py
+1
-1
@@ -4,8 +4,8 @@
4
5
from abc import ABC, abstractmethod
6
7
-from vllm.attention.backends.abstract import AttentionBackend, AttentionImpl
7
from vllm.config import VllmConfig
8
+from vllm.v1.attention.backend import AttentionBackend, AttentionImpl
9
from vllm.v1.kv_cache_interface import KVCacheSpec
10
11
vllm/model_executor/layers/batch_invariant.py
+1
-1
@@ -6,11 +6,11 @@ from typing import Any
6
7
import torch
8
9
-from vllm.attention.backends.registry import AttentionBackendEnum
9
from vllm.logger import init_logger
10
from vllm.platforms import current_platform
11
from vllm.triton_utils import tl, triton
12
from vllm.utils.torch_utils import is_torch_equal_or_newer
13
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
14
15
logger = init_logger(__name__)
16
vllm/model_executor/layers/kda.py
+1
-1
@@ -5,7 +5,6 @@ import torch
5
from einops import rearrange
6
from torch import nn
7
8
-from vllm.attention.backends.abstract import AttentionMetadata
8
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
9
from vllm.distributed import (
10
divide,
@@ -17,6 +16,7 @@ from vllm.logger import init_logger
16
from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
17
from vllm.model_executor.utils import set_weight_attrs
18
from vllm.utils.torch_utils import direct_register_custom_op
19
+from vllm.v1.attention.backend import AttentionMetadata
20
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
21
22
from .fla.ops.kda import (
vllm/model_executor/layers/mamba/abstract.py
+2
-2
@@ -5,10 +5,10 @@ from collections.abc import Iterable
5
6
import torch
7
8
-from vllm.attention.backends.abstract import AttentionBackend
9
-from vllm.attention.selector import get_mamba_attn_backend
8
from vllm.config import VllmConfig
9
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10
+from vllm.v1.attention.backend import AttentionBackend
11
+from vllm.v1.attention.selector import get_mamba_attn_backend
12
from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
13
14
vllm/model_executor/layers/mamba/linear_attn.py
+1
-1
@@ -8,7 +8,6 @@ import torch.nn.functional as F
8
from einops import rearrange
9
from torch import nn
10
11
-from vllm.attention.backends.abstract import AttentionMetadata
11
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
12
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
13
from vllm.distributed.parallel_state import (
@@ -29,6 +28,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
28
)
29
from vllm.model_executor.layers.quantization import QuantizationConfig
30
from vllm.utils.torch_utils import direct_register_custom_op
31
+from vllm.v1.attention.backend import AttentionMetadata
32
from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
33
34
vllm/model_executor/layers/mamba/mamba_mixer2.py
+1
-1
@@ -5,7 +5,6 @@
5
import torch
6
from torch import nn
7
8
-from vllm.attention.backends.abstract import AttentionMetadata
8
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
9
from vllm.distributed import (
10
divide,
@@ -43,6 +42,7 @@ from vllm.model_executor.model_loader.weight_utils import (
42
)
43
from vllm.model_executor.utils import set_weight_attrs
44
from vllm.utils.torch_utils import direct_register_custom_op
45
+from vllm.v1.attention.backend import AttentionMetadata
46
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
47
48
# Added by the IBM Team, 2024
vllm/model_executor/layers/mamba/short_conv.py
+1
-1
@@ -4,7 +4,6 @@
4
5
import torch
6
7
-from vllm.attention.backends.abstract import AttentionMetadata
7
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
8
from vllm.distributed import get_tensor_model_parallel_world_size
9
from vllm.forward_context import ForwardContext, get_forward_context
@@ -24,6 +23,7 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
23
causal_conv1d_update,
24
)
25
from vllm.utils.torch_utils import direct_register_custom_op
26
+from vllm.v1.attention.backend import AttentionMetadata
27
from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadata
28
29
vllm/model_executor/models/afmoe.py
+1
-1
@@ -9,7 +9,6 @@ from itertools import islice
9
import torch
10
from torch import nn
11
12
-from vllm.attention.backends.abstract import AttentionType
12
from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
@@ -50,6 +49,7 @@ from vllm.model_executor.models.utils import (
49
maybe_prefix,
50
)
51
from vllm.sequence import IntermediateTensors
52
+from vllm.v1.attention.backend import AttentionType
53
54
logger = init_logger(__name__)
55
vllm/model_executor/models/aimv2.py
+1
-1
@@ -8,10 +8,10 @@ from collections.abc import Iterable
8
import torch
9
import torch.nn as nn
10
11
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
11
from vllm.distributed import get_tensor_model_parallel_world_size
12
from vllm.distributed.utils import divide
13
from vllm.model_executor.layers.activation import SiluAndMul
14
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
15
from vllm.model_executor.layers.conv import Conv2dLayer
16
from vllm.model_executor.layers.layernorm import RMSNorm
17
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/apertus.py
+4
-2
@@ -32,13 +32,14 @@ import torch
32
from torch import nn
33
from transformers import ApertusConfig
34
35
-from vllm.attention.backends.abstract import AttentionType
35
from vllm.attention.layer import Attention
37
-from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
36
from vllm.compilation.decorators import support_torch_compile
37
from vllm.config import CacheConfig, VllmConfig
38
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
39
from vllm.model_executor.layers.activation import XIELU
40
+from vllm.model_executor.layers.attention.encoder_only_attention import (
41
+ EncoderOnlyAttention,
42
+)
43
from vllm.model_executor.layers.layernorm import RMSNorm
44
from vllm.model_executor.layers.linear import (
45
ColumnParallelLinear,
@@ -57,6 +58,7 @@ from vllm.model_executor.model_loader.weight_utils import (
58
maybe_remap_kv_scale_name,
59
)
60
from vllm.sequence import IntermediateTensors
61
+from vllm.v1.attention.backend import AttentionType
62
63
from .interfaces import SupportsLoRA, SupportsPP
64
from .utils import (
vllm/model_executor/models/bert.py
+3
-1
@@ -7,11 +7,13 @@ import torch
7
from torch import nn
8
from transformers import BertConfig
9
10
-from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
10
from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import CacheConfig, PoolerConfig, VllmConfig
12
from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.model_executor.layers.activation import get_act_fn
14
+from vllm.model_executor.layers.attention.encoder_only_attention import (
15
+ EncoderOnlyAttention,
16
+)
17
from vllm.model_executor.layers.linear import (
18
ColumnParallelLinear,
19
QKVParallelLinear,
vllm/model_executor/models/bert_with_rope.py
+3
-1
@@ -6,7 +6,6 @@ import torch
6
from torch import nn
7
from transformers import PretrainedConfig
8
9
-from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
9
from vllm.compilation.decorators import support_torch_compile
10
from vllm.config import CacheConfig, VllmConfig
11
from vllm.distributed import (
@@ -16,6 +15,9 @@ from vllm.distributed import (
15
tensor_model_parallel_all_reduce,
16
)
17
from vllm.model_executor.layers.activation import get_act_and_mul_fn, get_act_fn
18
+from vllm.model_executor.layers.attention.encoder_only_attention import (
19
+ EncoderOnlyAttention,
20
+)
21
from vllm.model_executor.layers.fused_moe import activation_without_mul, fused_topk
22
from vllm.model_executor.layers.linear import (
23
ColumnParallelLinear,
vllm/model_executor/models/blip.py
+1
-1
@@ -9,9 +9,9 @@ import torch
9
import torch.nn as nn
10
from transformers import Blip2VisionConfig, BlipVisionConfig
11
12
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
12
from vllm.distributed import divide, get_tensor_model_parallel_world_size
13
from vllm.model_executor.layers.activation import get_act_fn
14
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
15
from vllm.model_executor.layers.conv import Conv2dLayer
16
from vllm.model_executor.layers.linear import (
17
ColumnParallelLinear,
vllm/model_executor/models/clip.py
+1
-1
@@ -15,11 +15,11 @@ from transformers import (
15
)
16
17
from vllm.attention.layer import Attention
18
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
18
from vllm.config import VllmConfig
19
from vllm.config.multimodal import BaseDummyOptions, MultiModalConfig
20
from vllm.distributed import divide, get_tensor_model_parallel_world_size
21
from vllm.model_executor.layers.activation import get_act_fn
22
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
23
from vllm.model_executor.layers.conv import Conv2dLayer
24
from vllm.model_executor.layers.linear import (
25
ColumnParallelLinear,
vllm/model_executor/models/config.py
+1
-1
@@ -4,12 +4,12 @@ from copy import deepcopy
4
from math import lcm
5
from typing import TYPE_CHECKING
6
7
-from vllm.attention.backends.registry import AttentionBackendEnum
7
from vllm.logger import init_logger
8
from vllm.model_executor.models import ModelRegistry
9
from vllm.platforms import current_platform
10
from vllm.utils.math_utils import cdiv, round_up
11
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
12
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
13
from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, MLAAttentionSpec
14
15
if TYPE_CHECKING:
vllm/model_executor/models/deepencoder.py
+1
-1
@@ -18,8 +18,8 @@ import torch.nn as nn
18
import torch.nn.functional as F
19
from transformers import CLIPVisionConfig
20
21
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
21
from vllm.config import MultiModalConfig
22
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
23
from vllm.model_executor.layers.conv import Conv2dLayer
24
from vllm.model_executor.layers.quantization import QuantizationConfig
25
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
vllm/model_executor/models/deepseek_v2.py
+6
-4
@@ -33,9 +33,7 @@ from torch import nn
33
from transformers import DeepseekV2Config, DeepseekV3Config
34
35
from vllm._aiter_ops import rocm_aiter_ops
36
-from vllm.attention.backends.abstract import AttentionBackend
36
from vllm.attention.layer import Attention
38
-from vllm.attention.ops.common import pack_seq_triton, unpack_seq_triton
37
from vllm.compilation.decorators import support_torch_compile
38
from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config
39
from vllm.distributed import (
@@ -78,10 +76,12 @@ from vllm.platforms import current_platform
76
from vllm.sequence import IntermediateTensors
77
from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits
78
from vllm.utils.torch_utils import direct_register_custom_op
79
+from vllm.v1.attention.backend import AttentionBackend
80
from vllm.v1.attention.backends.mla.indexer import (
81
DeepseekV32IndexerBackend,
82
DeepseekV32IndexerMetadata,
83
)
84
+from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
85
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
86
from vllm.v1.worker.workspace import current_workspace_manager
87
@@ -679,7 +679,9 @@ def sparse_attn_indexer(
679
)
680
fp8_mqa_logits_func = fp8_mqa_logits
681
if current_platform.is_rocm():
682
- from vllm.attention.ops.rocm_aiter_mla_sparse import rocm_fp8_mqa_logits
682
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
683
+ rocm_fp8_mqa_logits,
684
+ )
685
686
fp8_mqa_logits_func = rocm_fp8_mqa_logits
687
logits = fp8_mqa_logits_func(
@@ -729,7 +731,7 @@ def sparse_attn_indexer(
731
num_padded_tokens = batch_size * next_n
732
fp8_paged_mqa_logits_func = fp8_paged_mqa_logits
733
if current_platform.is_rocm():
732
- from vllm.attention.ops.rocm_aiter_mla_sparse import (
734
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
735
rocm_fp8_paged_mqa_logits,
736
)
737
vllm/model_executor/models/dots_ocr.py
+4
-4
@@ -8,10 +8,6 @@ import torch.nn as nn
8
from torch.nn import LayerNorm
9
from transformers.models.qwen2_vl import Qwen2VLProcessor
10
11
-from vllm.attention.backends.registry import AttentionBackendEnum
12
-from vllm.attention.layers.mm_encoder_attention import (
13
- MMEncoderAttention,
14
-)
11
from vllm.config import MultiModalConfig, VllmConfig
12
from vllm.config.multimodal import BaseDummyOptions
13
from vllm.distributed import utils as dist_utils
@@ -20,6 +16,9 @@ from vllm.distributed.parallel_state import (
16
get_tensor_model_parallel_world_size,
17
)
18
from vllm.model_executor.layers.activation import SiluAndMul
19
+from vllm.model_executor.layers.attention.mm_encoder_attention import (
20
+ MMEncoderAttention,
21
+)
22
from vllm.model_executor.layers.conv import Conv2dLayer
23
from vllm.model_executor.layers.layernorm import RMSNorm
24
from vllm.model_executor.layers.linear import (
@@ -59,6 +58,7 @@ from vllm.multimodal.inputs import MultiModalDataDict
58
from vllm.sequence import IntermediateTensors
59
from vllm.transformers_utils.configs.dotsocr import DotsOCRConfig, DotsVisionConfig
60
from vllm.utils.tensor_schema import TensorSchema, TensorShape
61
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
62
63
from .vision import run_dp_sharded_mrope_vision_model
64
vllm/model_executor/models/ernie45_vl.py
+4
-4
@@ -36,16 +36,15 @@ import torch.nn.functional as F
36
from einops import rearrange
37
from transformers import BatchFeature
38
39
-from vllm.attention.backends.registry import AttentionBackendEnum
40
-from vllm.attention.layers.mm_encoder_attention import (
41
- MMEncoderAttention,
42
-)
39
from vllm.config import MultiModalConfig, VllmConfig
40
from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
41
from vllm.distributed import parallel_state
42
from vllm.distributed import utils as dist_utils
43
from vllm.logger import init_logger
44
from vllm.model_executor.layers.activation import QuickGELU
45
+from vllm.model_executor.layers.attention.mm_encoder_attention import (
46
+ MMEncoderAttention,
47
+)
48
from vllm.model_executor.layers.layernorm import RMSNorm
49
from vllm.model_executor.layers.linear import (
50
ColumnParallelLinear,
@@ -74,6 +73,7 @@ from vllm.multimodal.processing import (
73
from vllm.multimodal.profiling import BaseDummyInputsBuilder
74
from vllm.sequence import IntermediateTensors
75
from vllm.utils.tensor_schema import TensorSchema, TensorShape
76
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
77
78
from .ernie45_vl_moe import Ernie4_5_VLMoeForCausalLM
79
from .interfaces import (
vllm/model_executor/models/gemma3.py
+4
-2
@@ -22,13 +22,15 @@ import torch
22
from torch import nn
23
from transformers import Gemma3TextConfig
24
25
-from vllm.attention.backends.abstract import AttentionType
25
from vllm.attention.layer import Attention
26
from vllm.compilation.decorators import support_torch_compile
27
from vllm.config import CacheConfig, VllmConfig
28
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
29
from vllm.logger import init_logger
30
from vllm.model_executor.layers.activation import GeluAndMul
31
+from vllm.model_executor.layers.attention.encoder_only_attention import (
32
+ EncoderOnlyAttention,
33
+)
34
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
35
from vllm.model_executor.layers.linear import (
36
MergedColumnParallelLinear,
@@ -47,8 +49,8 @@ from vllm.model_executor.model_loader.weight_utils import (
49
maybe_remap_kv_scale_name,
50
)
51
from vllm.sequence import IntermediateTensors
52
+from vllm.v1.attention.backend import AttentionType
53
51
-from ...attention.layers.encoder_only_attention import EncoderOnlyAttention
54
from .interfaces import SupportsLoRA, SupportsPP
55
from .utils import (
56
AutoWeightsLoader,
vllm/model_executor/models/glm4.py
+1
-1
@@ -29,7 +29,6 @@ import torch
29
from torch import nn
30
from transformers import Glm4Config
31
32
-from vllm.attention.backends.abstract import AttentionType
32
from vllm.attention.layer import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
@@ -41,6 +40,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
40
from vllm.model_executor.layers.rotary_embedding import get_rope
41
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
42
from vllm.sequence import IntermediateTensors
43
+from vllm.v1.attention.backend import AttentionType
44
45
from .interfaces import SupportsLoRA, SupportsPP
46
from .llama import LlamaMLP as Glm4MLP
vllm/model_executor/models/glm4_1v.py
+4
-4
@@ -46,15 +46,14 @@ from transformers.models.glm4v.image_processing_glm4v import (
46
from transformers.models.glm4v.video_processing_glm4v import Glm4vVideoProcessor
47
from transformers.video_utils import VideoMetadata
48
49
-from vllm.attention.backends.registry import AttentionBackendEnum
50
-from vllm.attention.layers.mm_encoder_attention import (
51
- MMEncoderAttention,
52
-)
49
from vllm.config import MultiModalConfig, VllmConfig
50
from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
51
from vllm.distributed import get_tensor_model_parallel_world_size, parallel_state
52
from vllm.distributed import utils as dist_utils
53
from vllm.logger import init_logger
54
+from vllm.model_executor.layers.attention.mm_encoder_attention import (
55
+ MMEncoderAttention,
56
+)
57
from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
58
from vllm.model_executor.layers.layernorm import RMSNorm
59
from vllm.model_executor.layers.linear import (
@@ -89,6 +88,7 @@ from vllm.multimodal.processing import (
88
from vllm.multimodal.profiling import BaseDummyInputsBuilder
89
from vllm.sequence import IntermediateTensors
90
from vllm.utils.tensor_schema import TensorSchema, TensorShape
91
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
92
93
from ..layers.activation import SiluAndMul
94
from .interfaces import (
vllm/model_executor/models/glm4v.py
+1
-1
@@ -19,11 +19,11 @@ from transformers import BatchFeature, PreTrainedTokenizer, TensorType
19
from transformers.image_utils import ImageInput
20
from transformers.tokenization_utils_base import TextInput
21
22
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
22
from vllm.config import VllmConfig
23
from vllm.config.multimodal import BaseDummyOptions
24
from vllm.distributed import get_tensor_model_parallel_world_size
25
from vllm.model_executor.layers.activation import SiluAndMul, get_act_fn
26
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
27
from vllm.model_executor.layers.conv import Conv2dLayer
28
from vllm.model_executor.layers.linear import (
29
ColumnParallelLinear,
vllm/model_executor/models/glmasr.py
+1
-1
@@ -11,12 +11,12 @@ from transformers import BatchFeature
11
from transformers.models.glmasr import GlmAsrConfig, GlmAsrProcessor
12
from transformers.models.whisper import WhisperFeatureExtractor
13
14
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
14
from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
15
from vllm.config.multimodal import BaseDummyOptions
16
from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
17
from vllm.inputs.data import PromptType
18
from vllm.model_executor.layers.activation import get_act_fn
19
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
20
from vllm.model_executor.layers.linear import (
21
ColumnParallelLinear,
22
QKVParallelLinear,
vllm/model_executor/models/gpt_oss.py
+1
-1
@@ -7,7 +7,6 @@ import torch.distributed as dist
7
from torch import nn
8
from transformers import GptOssConfig
9
10
-from vllm.attention.backends.abstract import AttentionType
10
from vllm.attention.layer import Attention
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import CacheConfig, VllmConfig
@@ -37,6 +36,7 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk
36
from vllm.platforms import current_platform
37
from vllm.sequence import IntermediateTensors
38
from vllm.utils.math_utils import cdiv
39
+from vllm.v1.attention.backend import AttentionType
40
41
from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
42
from .utils import (
vllm/model_executor/models/hunyuan_v1.py
+1
-1
@@ -33,7 +33,6 @@ import torch
33
from torch import nn
34
from transformers import PretrainedConfig
35
36
-from vllm.attention.backends.abstract import AttentionType
36
from vllm.attention.layer import Attention
37
from vllm.compilation.decorators import support_torch_compile
38
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
@@ -65,6 +64,7 @@ from vllm.model_executor.model_loader.weight_utils import (
64
maybe_remap_kv_scale_name,
65
)
66
from vllm.sequence import IntermediateTensors
67
+from vllm.v1.attention.backend import AttentionType
68
69
from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
70
from .utils import (
vllm/model_executor/models/hunyuan_vision.py
+2
-2
@@ -33,14 +33,13 @@ import torch.nn as nn
33
import torch.nn.functional as F
34
from transformers import BatchFeature
35
36
-from vllm.attention.backends.registry import AttentionBackendEnum
37
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
36
from vllm.config import MultiModalConfig, VllmConfig
37
from vllm.config.multimodal import BaseDummyOptions
38
from vllm.distributed import parallel_state
39
from vllm.distributed import utils as dist_utils
40
from vllm.logger import init_logger
41
from vllm.model_executor.layers.activation import get_act_fn
42
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
43
from vllm.model_executor.layers.layernorm import RMSNorm
44
from vllm.model_executor.layers.linear import (
45
ColumnParallelLinear,
@@ -81,6 +80,7 @@ from vllm.transformers_utils.configs.hunyuan_vl import (
80
from vllm.transformers_utils.processors.hunyuan_vl import HunYuanVLProcessor
81
from vllm.transformers_utils.processors.hunyuan_vl_image import smart_resize
82
from vllm.utils.tensor_schema import TensorSchema, TensorShape
83
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
84
85
from .interfaces import (
86
MultiModalEmbeddings,
vllm/model_executor/models/idefics2_vision_model.py
+1
-1
@@ -27,9 +27,9 @@ from transformers.models.idefics2.configuration_idefics2 import (
27
Idefics2VisionConfig,
28
)
29
30
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
30
from vllm.distributed import get_tensor_model_parallel_world_size
31
from vllm.model_executor.layers.activation import get_act_fn
32
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
33
from vllm.model_executor.layers.conv import Conv2dLayer
34
from vllm.model_executor.layers.linear import (
35
ColumnParallelLinear,
vllm/model_executor/models/intern_vit.py
+1
-1
@@ -15,7 +15,6 @@ import torch.nn as nn
15
import torch.nn.functional as F
16
from transformers import PretrainedConfig
17
18
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
18
from vllm.distributed import (
19
divide,
20
get_tensor_model_parallel_rank,
@@ -24,6 +23,7 @@ from vllm.distributed import (
23
tensor_model_parallel_all_gather,
24
)
25
from vllm.model_executor.layers.activation import get_act_fn
26
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
27
from vllm.model_executor.layers.conv import Conv2dLayer
28
from vllm.model_executor.layers.layernorm import RMSNorm
29
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/interns1_vit.py
+1
-1
@@ -14,8 +14,8 @@ import torch.nn as nn
14
from transformers import PretrainedConfig
15
from transformers.utils import torch_int
16
17
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
17
from vllm.model_executor.layers.activation import get_act_fn
18
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
19
from vllm.model_executor.layers.conv import Conv2dLayer
20
from vllm.model_executor.layers.layernorm import RMSNorm
21
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
vllm/model_executor/models/iquest_loopcoder.py
+1
-1
@@ -24,7 +24,6 @@ import torch
24
from torch import nn
25
from transformers import PretrainedConfig
26
27
-from vllm.attention.backends.abstract import AttentionType
27
from vllm.attention.layer import Attention
28
from vllm.compilation.decorators import support_torch_compile
29
from vllm.config import CacheConfig, VllmConfig
@@ -48,6 +47,7 @@ from vllm.model_executor.model_loader.weight_utils import (
47
)
48
from vllm.model_executor.models.llama import LlamaMLP
49
from vllm.sequence import IntermediateTensors
50
+from vllm.v1.attention.backend import AttentionType
51
52
from .utils import (
53
AutoWeightsLoader,
vllm/model_executor/models/isaac.py
+1
-1
@@ -16,11 +16,11 @@ from transformers.image_processing_utils import BatchFeature
16
from transformers.tokenization_utils import TensorType
17
from typing_extensions import TypedDict, Unpack
18
19
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
19
from vllm.config import MultiModalConfig, VllmConfig
20
from vllm.config.model import ModelConfig
21
from vllm.distributed import parallel_state
22
from vllm.distributed import utils as dist_utils
23
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
24
from vllm.model_executor.layers.linear import (
25
ColumnParallelLinear,
26
QKVParallelLinear,
vllm/model_executor/models/keye.py
+3
-3
@@ -16,13 +16,13 @@ from transformers.feature_extraction_utils import BatchFeature
16
from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
17
from transformers.utils import torch_int
18
19
-from vllm.attention.layers.mm_encoder_attention import (
20
- MMEncoderAttention,
21
-)
19
from vllm.config import MultiModalConfig, VllmConfig
20
from vllm.config.multimodal import BaseDummyOptions
21
from vllm.distributed import get_tensor_model_parallel_world_size
22
from vllm.logger import init_logger
23
+from vllm.model_executor.layers.attention.mm_encoder_attention import (
24
+ MMEncoderAttention,
25
+)
26
from vllm.model_executor.layers.conv import Conv2dLayer
27
from vllm.model_executor.layers.linear import (
28
ColumnParallelLinear,
vllm/model_executor/models/llama.py
+4
-2
@@ -31,13 +31,14 @@ import torch
31
from torch import nn
32
from transformers import LlamaConfig
33
34
-from vllm.attention.backends.abstract import AttentionType
34
from vllm.attention.layer import Attention
36
-from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig
37
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
38
from vllm.model_executor.layers.activation import SiluAndMul
39
+from vllm.model_executor.layers.attention.encoder_only_attention import (
40
+ EncoderOnlyAttention,
41
+)
42
from vllm.model_executor.layers.layernorm import RMSNorm
43
from vllm.model_executor.layers.linear import (
44
MergedColumnParallelLinear,
@@ -56,6 +57,7 @@ from vllm.model_executor.model_loader.weight_utils import (
57
maybe_remap_kv_scale_name,
58
)
59
from vllm.sequence import IntermediateTensors
60
+from vllm.v1.attention.backend import AttentionType
61
62
from .adapters import as_embedding_model, as_seq_cls_model
63
from .interfaces import (
vllm/model_executor/models/llama4.py
+3
-1
@@ -25,7 +25,6 @@ from torch import nn
25
from transformers import Llama4TextConfig
26
27
from vllm.attention.layer import Attention
28
-from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
28
from vllm.compilation.decorators import support_torch_compile
29
from vllm.config import CacheConfig, VllmConfig
30
from vllm.distributed import (
@@ -34,6 +33,9 @@ from vllm.distributed import (
33
tensor_model_parallel_all_gather,
34
)
35
from vllm.logger import init_logger
36
+from vllm.model_executor.layers.attention.chunked_local_attention import (
37
+ ChunkedLocalAttention,
38
+)
39
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
40
from vllm.model_executor.layers.layernorm import RMSNorm
41
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/mimo_v2_flash.py
+1
-1
@@ -6,7 +6,6 @@ from itertools import islice
6
import torch
7
from torch import nn
8
9
-from vllm.attention.backends.abstract import AttentionType
9
from vllm.attention.layer import Attention
10
from vllm.config import (
11
CacheConfig,
@@ -43,6 +42,7 @@ from vllm.model_executor.model_loader.weight_utils import (
42
)
43
from vllm.model_executor.models.utils import sequence_parallel_chunk
44
from vllm.sequence import IntermediateTensors
45
+from vllm.v1.attention.backend import AttentionType
46
47
from .interfaces import MixtureOfExperts, SupportsPP
48
from .utils import (
vllm/model_executor/models/minimax_text_01.py
+1
-1
@@ -14,7 +14,6 @@ import torch
14
from torch import nn
15
from transformers import MiniMaxConfig
16
17
-from vllm.attention.backends.abstract import AttentionMetadata
17
from vllm.attention.layer import Attention
18
from vllm.compilation.decorators import support_torch_compile
19
from vllm.config import CacheConfig, ModelConfig, VllmConfig
@@ -48,6 +47,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
47
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
48
from vllm.model_executor.models.utils import maybe_prefix
49
from vllm.sequence import IntermediateTensors
50
+from vllm.v1.attention.backend import AttentionMetadata
51
52
from .interfaces import HasInnerState, IsHybrid
53
from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers
vllm/model_executor/models/mllama4.py
+1
-1
@@ -31,10 +31,10 @@ from transformers.models.llama4.image_processing_llama4_fast import (
31
get_best_fit,
32
)
33
34
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
34
from vllm.config import VllmConfig
35
from vllm.config.multimodal import BaseDummyOptions
36
from vllm.distributed import get_tensor_model_parallel_world_size
37
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
38
from vllm.model_executor.layers.fused_moe import FusedMoE
39
from vllm.model_executor.layers.linear import (
40
ColumnParallelLinear,
vllm/model_executor/models/modernbert.py
+3
-1
@@ -7,10 +7,12 @@ from torch import nn
7
from transformers import ModernBertConfig
8
from transformers.activations import ACT2FN
9
10
-from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
10
from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import VllmConfig
12
from vllm.distributed import get_tensor_model_parallel_world_size
13
+from vllm.model_executor.layers.attention.encoder_only_attention import (
14
+ EncoderOnlyAttention,
15
+)
16
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
17
from vllm.model_executor.layers.pooler import DispatchPooler
18
from vllm.model_executor.layers.pooler.seqwise import (
vllm/model_executor/models/molmo.py
+1
-1
@@ -18,7 +18,6 @@ from transformers.image_utils import ImageInput
18
from transformers.tokenization_utils_base import TextInput
19
20
from vllm.attention.layer import Attention
21
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
21
from vllm.compilation.decorators import support_torch_compile
22
from vllm.config import CacheConfig, VllmConfig
23
from vllm.config.multimodal import BaseDummyOptions
@@ -30,6 +29,7 @@ from vllm.distributed import (
29
tensor_model_parallel_all_gather,
30
)
31
from vllm.model_executor.layers.activation import MulAndSilu, QuickGELU, SiluAndMul
32
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
33
from vllm.model_executor.layers.layernorm import RMSNorm
34
from vllm.model_executor.layers.linear import (
35
ColumnParallelLinear,
vllm/model_executor/models/moonvit.py
+1
-1
@@ -52,9 +52,9 @@ import torch.nn.functional as F
52
from transformers.activations import ACT2FN
53
from transformers.modeling_utils import PreTrainedModel
54
55
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
55
from vllm.config import MultiModalConfig
56
from vllm.distributed import divide, get_tensor_model_parallel_world_size
57
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
58
from vllm.model_executor.layers.conv import Conv2dLayer
59
from vllm.model_executor.layers.linear import (
60
ColumnParallelLinear,
vllm/model_executor/models/nemotron_nas.py
+1
-1
@@ -31,7 +31,6 @@ import torch
31
from torch import nn
32
from transformers import LlamaConfig
33
34
-from vllm.attention.backends.abstract import AttentionType
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
36
from vllm.distributed import get_pp_group
@@ -49,6 +48,7 @@ from vllm.model_executor.model_loader.weight_utils import (
48
)
49
from vllm.model_executor.models.llama import LlamaAttention, LlamaMLP
50
from vllm.sequence import IntermediateTensors
51
+from vllm.v1.attention.backend import AttentionType
52
53
from .interfaces import HasNoOps, SupportsLoRA, SupportsPP
54
from .utils import (
vllm/model_executor/models/nemotron_parse.py
+1
-1
@@ -26,7 +26,6 @@ from transformers import (
26
TensorType,
27
)
28
29
-from vllm.attention.backends.abstract import AttentionType
29
from vllm.config import CacheConfig, VllmConfig
30
from vllm.config.lora import LoRAConfig
31
from vllm.config.multimodal import BaseDummyOptions
@@ -63,6 +62,7 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder
62
from vllm.transformers_utils.configs.radio import RadioConfig
63
from vllm.transformers_utils.tokenizer import AnyTokenizer
64
from vllm.utils.tensor_schema import TensorSchema, TensorShape
65
+from vllm.v1.attention.backend import AttentionType
66
67
logger = init_logger(__name__)
68
DEFAULT_FINAL_IMAGE_SIZE = (2048, 1648)
vllm/model_executor/models/openpangu.py
+3
-1
@@ -30,7 +30,6 @@ from torch import nn
30
from transformers import PretrainedConfig
31
32
from vllm.attention.layer import Attention, AttentionType
33
-from vllm.attention.layers.static_sink_attention import StaticSinkAttention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, ParallelConfig, VllmConfig
35
from vllm.distributed import (
@@ -42,6 +41,9 @@ from vllm.distributed import (
41
tensor_model_parallel_all_gather,
42
)
43
from vllm.model_executor.layers.activation import SiluAndMul
44
+from vllm.model_executor.layers.attention.static_sink_attention import (
45
+ StaticSinkAttention,
46
+)
47
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
48
from vllm.model_executor.layers.layernorm import RMSNorm
49
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/ouro.py
+1
-1
@@ -33,7 +33,6 @@ import torch
33
from torch import nn
34
from transformers import PretrainedConfig
35
36
-from vllm.attention.backends.abstract import AttentionType
36
from vllm.attention.layer import Attention
37
from vllm.compilation.decorators import support_torch_compile
38
from vllm.config import CacheConfig, VllmConfig
@@ -57,6 +56,7 @@ from vllm.model_executor.model_loader.weight_utils import (
56
maybe_remap_kv_scale_name,
57
)
58
from vllm.sequence import IntermediateTensors
59
+from vllm.v1.attention.backend import AttentionType
60
61
from .interfaces import SupportsLoRA
62
from .utils import (
vllm/model_executor/models/paddleocr_vl.py
+4
-4
@@ -30,14 +30,13 @@ from transformers.modeling_outputs import (
30
)
31
from transformers.utils import torch_int
32
33
-from vllm.attention.backends.registry import AttentionBackendEnum
34
-from vllm.attention.layers.mm_encoder_attention import (
35
- MMEncoderAttention,
36
-)
33
from vllm.config import MultiModalConfig, VllmConfig
34
from vllm.config.multimodal import BaseDummyOptions
35
from vllm.distributed import parallel_state
36
from vllm.distributed import utils as dist_utils
37
+from vllm.model_executor.layers.attention.mm_encoder_attention import (
38
+ MMEncoderAttention,
39
+)
40
from vllm.model_executor.layers.conv import Conv2dLayer
41
from vllm.model_executor.layers.linear import (
42
QKVParallelLinear,
@@ -72,6 +71,7 @@ from vllm.multimodal.processing import (
71
from vllm.multimodal.profiling import BaseDummyInputsBuilder
72
from vllm.sequence import IntermediateTensors
73
from vllm.utils.tensor_schema import TensorSchema, TensorShape
74
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
75
76
from .ernie45 import Ernie4_5ForCausalLM
77
from .interfaces import MultiModalEmbeddings, SupportsMRoPE, SupportsMultiModal
vllm/model_executor/models/plamo2.py
+1
-1
@@ -9,7 +9,6 @@ import torch
9
from torch import nn
10
from transformers import PretrainedConfig
11
12
-from vllm.attention.backends.abstract import AttentionMetadata
12
from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import VllmConfig, get_current_vllm_config
@@ -66,6 +65,7 @@ from vllm.model_executor.models.utils import (
65
from vllm.model_executor.utils import set_weight_attrs
66
from vllm.sequence import IntermediateTensors
67
from vllm.utils.torch_utils import direct_register_custom_op
68
+from vllm.v1.attention.backend import AttentionMetadata
69
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
70
71
vllm/model_executor/models/qwen2.py
+4
-2
@@ -33,13 +33,14 @@ import torch
33
from torch import nn
34
from transformers import Qwen2Config
35
36
-from vllm.attention.backends.abstract import AttentionType
36
from vllm.attention.layer import Attention
38
-from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
37
from vllm.compilation.decorators import support_torch_compile
38
from vllm.config import CacheConfig, VllmConfig
39
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
40
from vllm.model_executor.layers.activation import SiluAndMul
41
+from vllm.model_executor.layers.attention.encoder_only_attention import (
42
+ EncoderOnlyAttention,
43
+)
44
from vllm.model_executor.layers.layernorm import RMSNorm
45
from vllm.model_executor.layers.linear import (
46
MergedColumnParallelLinear,
@@ -59,6 +60,7 @@ from vllm.model_executor.model_loader.weight_utils import (
60
)
61
from vllm.sequence import IntermediateTensors
62
from vllm.transformers_utils.config import is_interleaved, set_default_rope_theta
63
+from vllm.v1.attention.backend import AttentionType
64
65
from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
66
from .utils import (
vllm/model_executor/models/qwen2_5_vl.py
+2
-2
@@ -41,8 +41,6 @@ from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
41
Qwen2_5_VLVisionConfig,
42
)
43
44
-from vllm.attention.backends.registry import AttentionBackendEnum
45
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
44
from vllm.compilation.decorators import support_torch_compile
45
from vllm.config import MultiModalConfig, VllmConfig
46
from vllm.distributed import parallel_state
@@ -50,6 +48,7 @@ from vllm.distributed import utils as dist_utils
48
from vllm.forward_context import set_forward_context
49
from vllm.logger import init_logger
50
from vllm.model_executor.layers.activation import get_act_and_mul_fn
51
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
52
from vllm.model_executor.layers.conv import Conv3dLayer
53
from vllm.model_executor.layers.layernorm import RMSNorm
54
from vllm.model_executor.layers.linear import (
@@ -83,6 +82,7 @@ from vllm.multimodal.processing import PromptReplacement, PromptUpdate
82
from vllm.sequence import IntermediateTensors
83
from vllm.utils.platform_utils import is_pin_memory_available
84
from vllm.utils.tensor_schema import TensorSchema, TensorShape
85
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
86
87
from .interfaces import (
88
MultiModalEmbeddings,
vllm/model_executor/models/qwen2_vl.py
+2
-2
@@ -43,14 +43,13 @@ from transformers.models.qwen2_vl.configuration_qwen2_vl import (
43
from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
44
from transformers.models.qwen2_vl.video_processing_qwen2_vl import Qwen2VLVideoProcessor
45
46
-from vllm.attention.backends.registry import AttentionBackendEnum
47
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
46
from vllm.config import MultiModalConfig, VllmConfig
47
from vllm.config.multimodal import BaseDummyOptions
48
from vllm.distributed import parallel_state, tensor_model_parallel_all_gather
49
from vllm.distributed import utils as dist_utils
50
from vllm.logger import init_logger
51
from vllm.model_executor.layers.activation import QuickGELU
52
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
53
from vllm.model_executor.layers.conv import Conv3dLayer
54
from vllm.model_executor.layers.linear import (
55
ColumnParallelLinear,
@@ -90,6 +89,7 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder
89
from vllm.sequence import IntermediateTensors
90
from vllm.tokenizers import TokenizerLike
91
from vllm.utils.tensor_schema import TensorSchema, TensorShape
92
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
93
94
from .interfaces import (
95
MultiModalEmbeddings,
vllm/model_executor/models/qwen3.py
+1
-1
@@ -30,7 +30,6 @@ import torch
30
from torch import nn
31
from transformers import Qwen3Config
32
33
-from vllm.attention.backends.abstract import AttentionType
33
from vllm.attention.layer import Attention
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
@@ -44,6 +43,7 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
43
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
44
from vllm.sequence import IntermediateTensors
45
from vllm.transformers_utils.config import set_default_rope_theta
46
+from vllm.v1.attention.backend import AttentionType
47
48
from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
49
from .qwen2 import Qwen2MLP as Qwen3MLP
vllm/model_executor/models/qwen3_next.py
+1
-1
@@ -10,7 +10,6 @@ from einops import rearrange
10
from torch import nn
11
from transformers.activations import ACT2FN
12
13
-from vllm.attention.backends.abstract import AttentionMetadata
13
from vllm.attention.layer import Attention
14
from vllm.compilation.decorators import support_torch_compile
15
from vllm.config import (
@@ -75,6 +74,7 @@ from vllm.sequence import IntermediateTensors
74
from vllm.transformers_utils.configs import Qwen3NextConfig
75
from vllm.triton_utils import tl, triton
76
from vllm.utils.torch_utils import direct_register_custom_op
77
+from vllm.v1.attention.backend import AttentionMetadata
78
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
79
80
from .interfaces import (
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+1
-1
@@ -46,7 +46,6 @@ from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import (
46
)
47
from transformers.models.whisper import WhisperFeatureExtractor
48
49
-from vllm.attention.backends.registry import AttentionBackendEnum
49
from vllm.compilation.decorators import support_torch_compile
50
from vllm.config import MultiModalConfig, VllmConfig
51
from vllm.distributed import get_pp_group
@@ -75,6 +74,7 @@ from vllm.multimodal.processing import (
74
PromptUpdateDetails,
75
)
76
from vllm.sequence import IntermediateTensors
77
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
78
79
from .interfaces import (
80
MultiModalEmbeddings,
vllm/model_executor/models/qwen3_vl.py
+1
-1
@@ -48,7 +48,6 @@ from transformers.models.qwen3_vl.video_processing_qwen3_vl import (
48
)
49
from transformers.video_utils import VideoMetadata
50
51
-from vllm.attention.backends.registry import AttentionBackendEnum
51
from vllm.compilation.decorators import support_torch_compile
52
from vllm.config import MultiModalConfig, VllmConfig
53
from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
@@ -92,6 +91,7 @@ from vllm.multimodal.processing import (
91
from vllm.multimodal.profiling import BaseDummyInputsBuilder
92
from vllm.sequence import IntermediateTensors
93
from vllm.utils.collection_utils import is_list_of
94
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
95
96
from .interfaces import (
97
MultiModalEmbeddings,
vllm/model_executor/models/seed_oss.py
+1
-1
@@ -30,7 +30,6 @@ import torch
30
from torch import nn
31
from transformers import PretrainedConfig as SeedOssConfig
32
33
-from vllm.attention.backends.abstract import AttentionType
33
from vllm.attention.layer import Attention
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
@@ -56,6 +55,7 @@ from vllm.model_executor.model_loader.weight_utils import (
55
)
56
from vllm.sequence import IntermediateTensors
57
from vllm.transformers_utils.config import set_default_rope_theta
58
+from vllm.v1.attention.backend import AttentionType
59
60
from .interfaces import SupportsLoRA, SupportsPP
61
from .utils import (
vllm/model_executor/models/siglip.py
+4
-2
@@ -15,12 +15,14 @@ from transformers import (
15
SiglipVisionConfig,
16
)
17
18
-from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
19
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
18
from vllm.config import VllmConfig
19
from vllm.config.multimodal import BaseDummyOptions, MultiModalConfig
20
from vllm.distributed import divide, get_tensor_model_parallel_world_size
21
from vllm.model_executor.layers.activation import get_act_fn
22
+from vllm.model_executor.layers.attention.encoder_only_attention import (
23
+ EncoderOnlyAttention,
24
+)
25
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
26
from vllm.model_executor.layers.conv import Conv2dLayer
27
from vllm.model_executor.layers.linear import (
28
ColumnParallelLinear,
vllm/model_executor/models/siglip2.py
+1
-1
@@ -10,11 +10,11 @@ from torch import nn
10
from torch.nn import functional as F
11
from transformers import Siglip2VisionConfig
12
13
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import MultiModalConfig
15
from vllm.distributed import get_tensor_model_parallel_world_size
16
from vllm.model_executor.layers.activation import get_act_fn
17
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
18
from vllm.model_executor.layers.linear import (
19
ColumnParallelLinear,
20
QKVParallelLinear,
vllm/model_executor/models/siglip2navit.py
+1
-1
@@ -11,10 +11,10 @@ from torch.nn import functional as F
11
from transformers import Siglip2VisionConfig
12
from transformers.configuration_utils import PretrainedConfig
13
14
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
14
from vllm.config import MultiModalConfig
15
from vllm.distributed import divide, get_tensor_model_parallel_world_size
16
from vllm.model_executor.layers.activation import get_act_fn
17
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
18
from vllm.model_executor.layers.conv import Conv2dLayer
19
from vllm.model_executor.layers.linear import (
20
ColumnParallelLinear,
vllm/model_executor/models/step3_vl.py
+1
-1
@@ -15,11 +15,11 @@ from torchvision import transforms
15
from torchvision.transforms.functional import InterpolationMode
16
from transformers import BatchFeature, PretrainedConfig, TensorType
17
18
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
18
from vllm.config import VllmConfig
19
from vllm.config.multimodal import BaseDummyOptions
20
from vllm.distributed import get_tensor_model_parallel_world_size
21
from vllm.model_executor.layers.activation import get_act_fn
22
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
23
from vllm.model_executor.layers.conv import Conv2dLayer
24
from vllm.model_executor.layers.linear import (
25
ColumnParallelLinear,
vllm/model_executor/models/transformers/base.py
+4
-2
@@ -27,13 +27,14 @@ from torch import nn
27
from transformers import AutoModel
28
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
29
30
-from vllm.attention.backends.abstract import AttentionType
30
from vllm.attention.layer import Attention
32
-from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
31
from vllm.config.utils import getattr_iter
32
from vllm.distributed import get_pp_group, get_tp_group
33
from vllm.distributed.utils import get_pp_indices
34
from vllm.logger import init_logger
35
+from vllm.model_executor.layers.attention.encoder_only_attention import (
36
+ EncoderOnlyAttention,
37
+)
38
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
39
from vllm.model_executor.models.interfaces import (
40
SupportsEagle,
@@ -59,6 +60,7 @@ from vllm.model_executor.models.utils import (
60
maybe_prefix,
61
)
62
from vllm.sequence import IntermediateTensors
63
+from vllm.v1.attention.backend import AttentionType
64
65
if TYPE_CHECKING:
66
from transformers import PreTrainedModel
vllm/model_executor/models/vision.py
+1
-1
@@ -10,7 +10,6 @@ from typing import Final, Generic, Literal, Protocol, TypeAlias, TypeVar
10
import torch
11
from transformers import PretrainedConfig
12
13
-from vllm.attention.backends.registry import AttentionBackendEnum
13
from vllm.config import VllmConfig
14
from vllm.distributed import (
15
get_tensor_model_parallel_rank,
@@ -19,6 +18,7 @@ from vllm.distributed import (
18
)
19
from vllm.logger import init_logger
20
from vllm.platforms import current_platform
21
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
22
23
logger = init_logger(__name__)
24
vllm/model_executor/models/whisper.py
+5
-5
@@ -18,18 +18,15 @@ from transformers import (
18
)
19
from transformers.models.whisper.modeling_whisper import sinusoids
20
21
-from vllm.attention.backends.abstract import (
22
- AttentionType,
23
-)
21
from vllm.attention.layer import Attention
25
-from vllm.attention.layers.cross_attention import CrossAttention
26
-from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
22
from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig
23
from vllm.config.multimodal import BaseDummyOptions
24
from vllm.distributed import get_tensor_model_parallel_world_size
25
from vllm.inputs.data import PromptType
26
from vllm.logger import init_logger
27
from vllm.model_executor.layers.activation import get_act_fn
28
+from vllm.model_executor.layers.attention.cross_attention import CrossAttention
29
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
30
from vllm.model_executor.layers.linear import (
31
ColumnParallelLinear,
32
QKVParallelLinear,
@@ -62,6 +59,9 @@ from vllm.transformers_utils.processor import cached_processor_from_config
59
from vllm.utils.jsontree import json_map_leaves
60
from vllm.utils.tensor_schema import TensorSchema, TensorShape
61
from vllm.utils.torch_utils import set_default_torch_dtype
62
+from vllm.v1.attention.backend import (
63
+ AttentionType,
64
+)
65
66
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsTranscription
67
from .utils import (
vllm/model_executor/models/whisper_utils.py
+5
-5
@@ -9,20 +9,20 @@ import torch
9
import torch.nn.functional as F
10
from torch import nn
11
12
-from vllm.attention.backends.abstract import (
12
+from vllm.attention.layer import Attention
13
+from vllm.config import CacheConfig, VllmConfig
14
+from vllm.model_executor.layers.quantization import QuantizationConfig
15
+from vllm.v1.attention.backend import (
16
AttentionBackend,
17
AttentionMetadata,
18
AttentionType,
19
)
17
-from vllm.attention.layer import Attention
18
-from vllm.attention.selector import get_attn_backend
19
-from vllm.config import CacheConfig, VllmConfig
20
-from vllm.model_executor.layers.quantization import QuantizationConfig
20
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
21
from vllm.v1.attention.backends.utils import (
22
CommonAttentionMetadata,
23
subclass_attention_backend_with_overrides,
24
)
25
+from vllm.v1.attention.selector import get_attn_backend
26
from vllm.v1.kv_cache_interface import AttentionSpec
27
28
# From https://platform.openai.com/docs/guides/speech-to-text/supported-languages
vllm/platforms/cpu.py
+2
-2
@@ -15,16 +15,16 @@ import regex as re
15
import torch
16
17
from vllm import envs
18
-from vllm.attention.backends.registry import AttentionBackendEnum
18
from vllm.logger import init_logger
19
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
20
21
from .interface import CpuArchEnum, Platform, PlatformEnum
22
23
logger = init_logger(__name__)
24
25
if TYPE_CHECKING:
26
- from vllm.attention.selector import AttentionSelectorConfig
26
from vllm.config import VllmConfig
27
+ from vllm.v1.attention.selector import AttentionSelectorConfig
28
else:
29
VllmConfig = None
30
vllm/platforms/cuda.py
+4
-4
@@ -14,17 +14,17 @@ from typing_extensions import ParamSpec
14
15
# import custom ops, trigger op registration
16
import vllm._C # noqa
17
-from vllm.attention.backends.registry import AttentionBackendEnum
17
from vllm.logger import init_logger
18
from vllm.utils.import_utils import import_pynvml
19
from vllm.utils.torch_utils import cuda_device_count_stateless
20
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
21
22
from .interface import DeviceCapability, Platform, PlatformEnum
23
24
if TYPE_CHECKING:
25
- from vllm.attention.selector import AttentionSelectorConfig
25
from vllm.config import VllmConfig
26
from vllm.config.cache import CacheDType
27
+ from vllm.v1.attention.selector import AttentionSelectorConfig
28
else:
29
VllmConfig = None
30
CacheDType = None
@@ -148,7 +148,7 @@ class CudaPlatformBase(Platform):
148
149
@classmethod
150
def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
151
- from vllm.attention.backends.registry import AttentionBackendEnum
151
+ from vllm.v1.attention.backends.registry import AttentionBackendEnum
152
153
parallel_config = vllm_config.parallel_config
154
model_config = vllm_config.model_config
@@ -200,7 +200,7 @@ class CudaPlatformBase(Platform):
200
use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA
201
use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA
202
203
- from vllm.attention.ops.flashmla import is_flashmla_dense_supported
203
+ from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
204
205
if (
206
use_flashmla
vllm/platforms/interface.py
+2
-2
@@ -13,18 +13,18 @@ import numpy as np
13
import torch
14
from typing_extensions import deprecated
15
16
-from vllm.attention.backends.registry import AttentionBackendEnum
16
from vllm.logger import init_logger
17
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
18
19
if TYPE_CHECKING:
20
from torch.distributed import PrefixStore, ProcessGroup
21
22
- from vllm.attention.selector import AttentionSelectorConfig
22
from vllm.config import VllmConfig
23
from vllm.inputs import ProcessorInputs, PromptType
24
from vllm.pooling_params import PoolingParams
25
from vllm.sampling_params import SamplingParams
26
from vllm.utils.argparse_utils import FlexibleArgumentParser
27
+ from vllm.v1.attention.selector import AttentionSelectorConfig
28
else:
29
FlexibleArgumentParser = object
30
vllm/platforms/rocm.py
+2
-2
@@ -8,15 +8,15 @@ from typing import TYPE_CHECKING, Optional
8
import torch
9
10
import vllm.envs as envs
11
-from vllm.attention.backends.registry import AttentionBackendEnum
11
from vllm.logger import init_logger
12
from vllm.utils.torch_utils import cuda_device_count_stateless
13
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
14
15
from .interface import DeviceCapability, Platform, PlatformEnum
16
17
if TYPE_CHECKING:
18
- from vllm.attention.selector import AttentionSelectorConfig
18
from vllm.config import VllmConfig
19
+ from vllm.v1.attention.selector import AttentionSelectorConfig
20
21
logger = init_logger(__name__)
22
vllm/platforms/xpu.py
+2
-2
@@ -7,14 +7,14 @@ from typing import TYPE_CHECKING, Optional
7
8
import torch
9
10
-from vllm.attention.backends.registry import AttentionBackendEnum
10
from vllm.logger import init_logger
11
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
12
13
from .interface import DeviceCapability, Platform, PlatformEnum
14
15
if TYPE_CHECKING:
16
- from vllm.attention.selector import AttentionSelectorConfig
16
from vllm.config import VllmConfig
17
+ from vllm.v1.attention.selector import AttentionSelectorConfig
18
else:
19
VllmConfig = None
20
vllm/v1/attention/backend.py
renamed
vllm/v1/attention/backends/cpu_attn.py
+4
-4
@@ -6,16 +6,16 @@ from typing import ClassVar
6
import torch
7
8
from vllm import _custom_ops as ops
9
-from vllm.attention.backends.abstract import (
9
+from vllm.config import VllmConfig
10
+from vllm.logger import init_logger
11
+from vllm.platforms import CpuArchEnum, current_platform
12
+from vllm.v1.attention.backend import (
13
AttentionBackend,
14
AttentionImpl,
15
AttentionLayer,
16
AttentionType,
17
is_quantized_kv_cache,
18
)
16
-from vllm.config import VllmConfig
17
-from vllm.logger import init_logger
18
-from vllm.platforms import CpuArchEnum, current_platform
19
from vllm.v1.attention.backends.utils import (
20
AttentionMetadataBuilder,
21
CommonAttentionMetadata,
vllm/v1/attention/backends/fa_utils.py
renamed
vllm/v1/attention/backends/flash_attn.py
+6
-6
@@ -9,24 +9,24 @@ from typing import ClassVar
9
import numpy as np
10
import torch
11
12
-from vllm.attention.backends.abstract import (
12
+from vllm.attention.layer import Attention
13
+from vllm.v1.attention.backend import (
14
AttentionBackend,
15
AttentionImpl,
16
AttentionType,
17
MultipleOf,
18
is_quantized_kv_cache,
19
)
19
-from vllm.attention.layer import Attention
20
-from vllm.attention.ops.common import cp_lse_ag_out_rs
21
-from vllm.attention.ops.merge_attn_states import merge_attn_states
22
-from vllm.attention.utils.fa_utils import (
20
+from vllm.v1.attention.backends.fa_utils import (
21
flash_attn_supports_fp8,
22
get_flash_attn_version,
23
is_flash_attn_varlen_func_available,
24
)
25
+from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
26
+from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
27
28
if is_flash_attn_varlen_func_available():
29
- from vllm.attention.utils.fa_utils import (
29
+ from vllm.v1.attention.backends.fa_utils import (
30
flash_attn_supports_sinks,
31
flash_attn_varlen_func,
32
get_scheduler_metadata,
vllm/v1/attention/backends/flash_attn_diffkv.py
+4
-4
@@ -4,14 +4,14 @@
4
5
import torch
6
7
-from vllm.attention.backends.abstract import AttentionType
8
-from vllm.attention.ops.triton_reshape_and_cache_flash import (
7
+from vllm.v1.attention.backend import AttentionType
8
+from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
9
+from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
10
triton_reshape_and_cache_flash_diffkv,
11
)
11
-from vllm.attention.utils.fa_utils import is_flash_attn_varlen_func_available
12
13
if is_flash_attn_varlen_func_available():
14
- from vllm.attention.utils.fa_utils import flash_attn_varlen_func
14
+ from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
15
from vllm.logger import init_logger
16
from vllm.v1.attention.backends.utils import get_kv_cache_layout
17
vllm/v1/attention/backends/flashinfer.py
+8
-8
@@ -19,14 +19,6 @@ from flashinfer.utils import FP4Tensor
19
from typing_extensions import override
20
21
from vllm import envs
22
-from vllm.attention.backends.abstract import (
23
- AttentionBackend,
24
- AttentionImpl,
25
- AttentionType,
26
- MultipleOf,
27
-)
28
-from vllm.attention.ops.common import cp_lse_ag_out_rs
29
-from vllm.attention.ops.merge_attn_states import merge_attn_states
22
from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config
23
from vllm.config.cache import CacheDType
24
from vllm.distributed.parallel_state import get_dcp_group
@@ -48,6 +40,12 @@ from vllm.utils.flashinfer import (
40
)
41
from vllm.utils.math_utils import cdiv
42
from vllm.utils.platform_utils import is_pin_memory_available
43
+from vllm.v1.attention.backend import (
44
+ AttentionBackend,
45
+ AttentionImpl,
46
+ AttentionType,
47
+ MultipleOf,
48
+)
49
from vllm.v1.attention.backends.utils import (
50
AttentionCGSupport,
51
AttentionMetadataBuilder,
@@ -59,6 +57,8 @@ from vllm.v1.attention.backends.utils import (
57
infer_global_hyperparameters,
58
split_decodes_and_prefills,
59
)
60
+from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
61
+from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
62
from vllm.v1.kv_cache_interface import AttentionSpec
63
from vllm.v1.utils import CpuGpuBuffer
64
vllm/v1/attention/backends/flex_attention.py
+6
-6
@@ -20,12 +20,6 @@ from torch.nn.attention.flex_attention import (
20
or_masks,
21
)
22
23
-from vllm.attention.backends.abstract import (
24
- AttentionBackend,
25
- AttentionImpl,
26
- AttentionType,
27
- is_quantized_kv_cache,
28
-)
23
from vllm.config import VllmConfig
24
from vllm.config.cache import CacheDType
25
from vllm.logger import init_logger
@@ -35,6 +29,12 @@ from vllm.model_executor.layers.batch_invariant import (
29
from vllm.platforms import current_platform
30
from vllm.utils.math_utils import cdiv
31
from vllm.utils.torch_utils import is_torch_equal_or_newer
32
+from vllm.v1.attention.backend import (
33
+ AttentionBackend,
34
+ AttentionImpl,
35
+ AttentionType,
36
+ is_quantized_kv_cache,
37
+)
38
from vllm.v1.attention.backends.utils import (
39
AttentionMetadataBuilder,
40
CommonAttentionMetadata,
vllm/v1/attention/backends/gdn_attn.py
+1
-1
@@ -6,8 +6,8 @@ from dataclasses import dataclass
6
7
import torch
8
9
-from vllm.attention.backends.abstract import AttentionBackend
9
from vllm.config import VllmConfig
10
+from vllm.v1.attention.backend import AttentionBackend
11
from vllm.v1.attention.backends.utils import (
12
PAD_SLOT_ID,
13
AttentionCGSupport,
vllm/v1/attention/backends/linear_attn.py
+1
-1
@@ -4,8 +4,8 @@ from dataclasses import dataclass
4
5
import torch
6
7
-from vllm.attention.backends.abstract import AttentionBackend
7
from vllm.config import VllmConfig
8
+from vllm.v1.attention.backend import AttentionBackend
9
from vllm.v1.attention.backends.utils import (
10
AttentionCGSupport,
11
AttentionMetadataBuilder,
vllm/v1/attention/backends/mamba1_attn.py
+1
-1
@@ -3,7 +3,7 @@
3
4
from dataclasses import dataclass
5
6
-from vllm.attention.backends.abstract import AttentionBackend
6
+from vllm.v1.attention.backend import AttentionBackend
7
from vllm.v1.attention.backends.mamba_attn import (
8
BaseMambaAttentionMetadata,
9
BaseMambaAttentionMetadataBuilder,
vllm/v1/attention/backends/mamba2_attn.py
+1
-1
@@ -5,9 +5,9 @@ from dataclasses import dataclass, replace
5
6
import torch
7
8
-from vllm.attention.backends.abstract import AttentionBackend
8
from vllm.config import VllmConfig
9
from vllm.utils.math_utils import cdiv
10
+from vllm.v1.attention.backend import AttentionBackend
11
from vllm.v1.attention.backends.mamba_attn import (
12
BaseMambaAttentionMetadata,
13
BaseMambaAttentionMetadataBuilder,
vllm/v1/attention/backends/mla/common.py
+9
-9
@@ -199,15 +199,6 @@ from tqdm import tqdm
199
from vllm import _custom_ops as ops
200
from vllm import envs
201
from vllm._aiter_ops import rocm_aiter_ops
202
-from vllm.attention.backends.abstract import (
203
- AttentionBackend,
204
- AttentionLayer,
205
- AttentionMetadata,
206
- MLAAttentionImpl,
207
-)
208
-from vllm.attention.ops.common import cp_lse_ag_out_rs
209
-from vllm.attention.ops.merge_attn_states import merge_attn_states
210
-from vllm.attention.utils.fa_utils import get_flash_attn_version
202
from vllm.config import ModelConfig, VllmConfig, get_current_vllm_config
203
from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
204
from vllm.logger import init_logger
@@ -222,6 +213,13 @@ from vllm.model_executor.layers.linear import (
213
from vllm.platforms import current_platform
214
from vllm.utils.flashinfer import has_nvidia_artifactory
215
from vllm.utils.math_utils import cdiv, round_down
216
+from vllm.v1.attention.backend import (
217
+ AttentionBackend,
218
+ AttentionLayer,
219
+ AttentionMetadata,
220
+ MLAAttentionImpl,
221
+)
222
+from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
223
from vllm.v1.attention.backends.utils import (
224
AttentionMetadataBuilder,
225
CommonAttentionMetadata,
@@ -230,6 +228,8 @@ from vllm.v1.attention.backends.utils import (
228
infer_global_hyperparameters,
229
split_decodes_and_prefills,
230
)
231
+from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
232
+from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
233
from vllm.v1.kv_cache_interface import AttentionSpec
234
235
vllm/v1/attention/backends/mla/cutlass_mla.py
+4
-4
@@ -7,15 +7,15 @@ from typing import ClassVar
7
import torch
8
9
import vllm._custom_ops as ops
10
-from vllm.attention.backends.abstract import (
10
+from vllm.config.cache import CacheDType
11
+from vllm.logger import init_logger
12
+from vllm.platforms.interface import DeviceCapability
13
+from vllm.v1.attention.backend import (
14
AttentionLayer,
15
AttentionType,
16
MultipleOf,
17
is_quantized_kv_cache,
18
)
16
-from vllm.config.cache import CacheDType
17
-from vllm.logger import init_logger
18
-from vllm.platforms.interface import DeviceCapability
19
from vllm.v1.attention.backends.mla.common import (
20
MLACommonBackend,
21
MLACommonImpl,
vllm/v1/attention/backends/mla/flashattn_mla.py
+9
-9
@@ -6,23 +6,23 @@ from typing import ClassVar
6
7
import torch
8
9
-from vllm.attention.backends.abstract import (
9
+from vllm.config import VllmConfig
10
+from vllm.config.cache import CacheDType
11
+from vllm.logger import init_logger
12
+from vllm.model_executor.layers.batch_invariant import (
13
+ vllm_is_batch_invariant,
14
+)
15
+from vllm.platforms.interface import DeviceCapability
16
+from vllm.v1.attention.backend import (
17
AttentionLayer,
18
AttentionType,
19
MultipleOf,
20
is_quantized_kv_cache,
21
)
15
-from vllm.attention.utils.fa_utils import (
22
+from vllm.v1.attention.backends.fa_utils import (
23
flash_attn_supports_mla,
24
get_flash_attn_version,
25
)
19
-from vllm.config import VllmConfig
20
-from vllm.config.cache import CacheDType
21
-from vllm.logger import init_logger
22
-from vllm.model_executor.layers.batch_invariant import (
23
- vllm_is_batch_invariant,
24
-)
25
-from vllm.platforms.interface import DeviceCapability
26
from vllm.v1.attention.backends.mla.common import (
27
MLACommonBackend,
28
MLACommonDecodeMetadata,
vllm/v1/attention/backends/mla/flashinfer_mla.py
+4
-4
@@ -6,14 +6,14 @@ from typing import ClassVar
6
import torch
7
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
8
9
-from vllm.attention.backends.abstract import (
9
+from vllm.config.cache import CacheDType
10
+from vllm.logger import init_logger
11
+from vllm.platforms.interface import DeviceCapability
12
+from vllm.v1.attention.backend import (
13
AttentionLayer,
14
AttentionType,
15
MultipleOf,
16
)
14
-from vllm.config.cache import CacheDType
15
-from vllm.logger import init_logger
16
-from vllm.platforms.interface import DeviceCapability
17
from vllm.v1.attention.backends.mla.common import (
18
MLACommonBackend,
19
MLACommonImpl,
vllm/v1/attention/backends/mla/flashmla.py
+8
-8
@@ -6,12 +6,6 @@ from typing import ClassVar
6
7
import torch
8
9
-from vllm.attention.backends.abstract import AttentionLayer, AttentionType, MultipleOf
10
-from vllm.attention.ops.flashmla import (
11
- flash_mla_with_kvcache,
12
- get_mla_metadata,
13
- is_flashmla_dense_supported,
14
-)
9
from vllm.config import VllmConfig
10
from vllm.config.cache import CacheDType
11
from vllm.logger import init_logger
@@ -19,6 +13,7 @@ from vllm.model_executor.layers.batch_invariant import (
13
vllm_is_batch_invariant,
14
)
15
from vllm.platforms.interface import DeviceCapability
16
+from vllm.v1.attention.backend import AttentionLayer, AttentionType, MultipleOf
17
from vllm.v1.attention.backends.mla.common import (
18
MLACommonBackend,
19
MLACommonDecodeMetadata,
@@ -32,6 +27,11 @@ from vllm.v1.attention.backends.utils import (
27
reshape_attn_output_for_spec_decode,
28
reshape_query_for_spec_decode,
29
)
30
+from vllm.v1.attention.ops.flashmla import (
31
+ flash_mla_with_kvcache,
32
+ get_mla_metadata,
33
+ is_flashmla_dense_supported,
34
+)
35
from vllm.v1.kv_cache_interface import AttentionSpec
36
37
logger = init_logger(__name__)
@@ -78,11 +78,11 @@ class FlashMLABackend(MLACommonBackend):
78
device_capability: DeviceCapability,
79
) -> str | None:
80
if use_sparse:
81
- from vllm.attention.ops.flashmla import is_flashmla_sparse_supported
81
+ from vllm.v1.attention.ops.flashmla import is_flashmla_sparse_supported
82
83
return is_flashmla_sparse_supported()[1]
84
else:
85
- from vllm.attention.ops.flashmla import is_flashmla_dense_supported
85
+ from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
86
87
return is_flashmla_dense_supported()[1]
88
vllm/v1/attention/backends/mla/flashmla_sparse.py
+11
-11
@@ -7,17 +7,6 @@ import numpy as np
7
import torch
8
9
from vllm import _custom_ops as ops
10
-from vllm.attention.backends.abstract import (
11
- AttentionBackend,
12
- AttentionLayer,
13
- AttentionMetadata,
14
- MultipleOf,
15
-)
16
-from vllm.attention.ops.flashmla import (
17
- flash_mla_sparse_prefill,
18
- flash_mla_with_kvcache,
19
- get_mla_metadata,
20
-)
10
from vllm.config import VllmConfig, get_current_vllm_config
11
from vllm.config.cache import CacheDType
12
from vllm.logger import init_logger
@@ -25,6 +14,12 @@ from vllm.platforms import current_platform
14
from vllm.platforms.interface import DeviceCapability
15
from vllm.triton_utils import tl, triton
16
from vllm.utils.math_utils import cdiv
17
+from vllm.v1.attention.backend import (
18
+ AttentionBackend,
19
+ AttentionLayer,
20
+ AttentionMetadata,
21
+ MultipleOf,
22
+)
23
from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl, get_mla_dims
24
from vllm.v1.attention.backends.utils import (
25
AttentionCGSupport,
@@ -35,6 +30,11 @@ from vllm.v1.attention.backends.utils import (
30
split_decodes_and_prefills,
31
split_prefill_chunks,
32
)
33
+from vllm.v1.attention.ops.flashmla import (
34
+ flash_mla_sparse_prefill,
35
+ flash_mla_with_kvcache,
36
+ get_mla_metadata,
37
+)
38
from vllm.v1.kv_cache_interface import AttentionSpec
39
from vllm.v1.worker.workspace import current_workspace_manager
40
vllm/v1/attention/backends/mla/indexer.py
+4
-4
@@ -5,14 +5,14 @@ from typing import ClassVar
5
6
import torch
7
8
-from vllm.attention.backends.abstract import (
9
- AttentionBackend,
10
- MultipleOf,
11
-)
8
from vllm.config import VllmConfig
9
from vllm.logger import init_logger
10
from vllm.platforms import current_platform
11
from vllm.utils.deep_gemm import get_paged_mqa_logits_metadata, is_deep_gemm_supported
12
+from vllm.v1.attention.backend import (
13
+ AttentionBackend,
14
+ MultipleOf,
15
+)
16
from vllm.v1.attention.backends.utils import (
17
AttentionCGSupport,
18
AttentionMetadataBuilder,
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+1
-1
@@ -7,8 +7,8 @@ from typing import ClassVar
7
import torch
8
9
from vllm._aiter_ops import rocm_aiter_ops
10
-from vllm.attention.backends.abstract import AttentionLayer, MultipleOf
10
from vllm.config import VllmConfig
11
+from vllm.v1.attention.backend import AttentionLayer, MultipleOf
12
from vllm.v1.attention.backends.mla.common import (
13
MLACommonBackend,
14
MLACommonDecodeMetadata,
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+3
-3
@@ -9,13 +9,13 @@ import torch
9
10
from vllm import _custom_ops as ops
11
from vllm._aiter_ops import rocm_aiter_ops
12
-from vllm.attention.backends.abstract import (
12
+from vllm.config import VllmConfig
13
+from vllm.logger import init_logger
14
+from vllm.v1.attention.backend import (
15
AttentionBackend,
16
AttentionLayer,
17
AttentionMetadata,
18
)
17
-from vllm.config import VllmConfig
18
-from vllm.logger import init_logger
19
from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl, get_mla_dims
20
from vllm.v1.attention.backends.mla.flashmla_sparse import (
21
triton_convert_req_index_to_global_index,
vllm/v1/attention/backends/mla/triton_mla.py
+6
-6
@@ -5,23 +5,23 @@ from typing import ClassVar
5
6
import torch
7
8
-from vllm.attention.backends.abstract import (
9
- AttentionLayer,
10
- AttentionType,
11
- is_quantized_kv_cache,
12
-)
13
-from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
8
from vllm.config.cache import CacheDType
9
from vllm.logger import init_logger
10
from vllm.model_executor.layers.batch_invariant import (
11
vllm_is_batch_invariant,
12
)
13
from vllm.platforms.interface import DeviceCapability
14
+from vllm.v1.attention.backend import (
15
+ AttentionLayer,
16
+ AttentionType,
17
+ is_quantized_kv_cache,
18
+)
19
from vllm.v1.attention.backends.mla.common import (
20
MLACommonBackend,
21
MLACommonImpl,
22
MLACommonMetadata,
23
)
24
+from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
25
26
logger = init_logger(__name__)
27
vllm/v1/attention/backends/registry.py
renamed
+1
-1
@@ -10,7 +10,7 @@ from vllm.logger import init_logger
10
from vllm.utils.import_utils import resolve_obj_by_qualname
11
12
if TYPE_CHECKING:
13
- from vllm.attention.backends.abstract import AttentionBackend
13
+ from vllm.v1.attention.backend import AttentionBackend
14
15
logger = init_logger(__name__)
16
vllm/v1/attention/backends/rocm_aiter_fa.py
+7
-7
@@ -7,25 +7,25 @@ from typing import ClassVar
7
8
import torch
9
10
-from vllm.attention.backends.abstract import (
11
- AttentionBackend,
12
- AttentionImpl,
13
- AttentionType,
14
- MultipleOf,
15
-)
10
from vllm.attention.layer import Attention
17
-from vllm.attention.ops.merge_attn_states import merge_attn_states
11
from vllm.config import VllmConfig, get_layers_from_vllm_config
12
from vllm.logger import init_logger
13
from vllm.platforms import current_platform
14
from vllm.utils.math_utils import cdiv
15
from vllm.utils.platform_utils import get_cu_count
16
+from vllm.v1.attention.backend import (
17
+ AttentionBackend,
18
+ AttentionImpl,
19
+ AttentionType,
20
+ MultipleOf,
21
+)
22
from vllm.v1.attention.backends.utils import (
23
AttentionCGSupport,
24
AttentionMetadataBuilder,
25
CommonAttentionMetadata,
26
split_decodes_prefills_and_extends,
27
)
28
+from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
29
from vllm.v1.kv_cache_interface import AttentionSpec
30
31
_PARTITION_SIZE_ROCM = 256
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+1
-1
@@ -5,12 +5,12 @@
5
import torch
6
7
from vllm import _custom_ops as ops
8
-from vllm.attention.backends.abstract import AttentionType
8
from vllm.logger import init_logger
9
from vllm.model_executor.layers.quantization.utils.quant_utils import (
10
QuantKey,
11
kFp8StaticTensorSym,
12
)
13
+from vllm.v1.attention.backend import AttentionType
14
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
15
from vllm.v1.attention.backends.rocm_attn import (
16
RocmAttentionBackend,
vllm/v1/attention/backends/rocm_attn.py
+13
-11
@@ -7,17 +7,6 @@ from typing import ClassVar
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.chunked_prefill_paged_decode import chunked_prefill_paged_decode
17
-from vllm.attention.ops.paged_attn import PagedAttention
18
-from vllm.attention.ops.triton_reshape_and_cache_flash import (
19
- triton_reshape_and_cache_flash,
20
-)
10
from vllm.config import VllmConfig
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -25,12 +14,25 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
14
kFp8StaticTensorSym,
15
)
16
from vllm.platforms import current_platform
17
+from vllm.v1.attention.backend import (
18
+ AttentionBackend,
19
+ AttentionImpl,
20
+ AttentionType,
21
+ MultipleOf,
22
+)
23
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
24
from vllm.v1.attention.backends.utils import (
25
AttentionCGSupport,
26
AttentionMetadataBuilder,
27
CommonAttentionMetadata,
28
)
29
+from vllm.v1.attention.ops.chunked_prefill_paged_decode import (
30
+ chunked_prefill_paged_decode,
31
+)
32
+from vllm.v1.attention.ops.paged_attn import PagedAttention
33
+from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
34
+ triton_reshape_and_cache_flash,
35
+)
36
from vllm.v1.kv_cache_interface import AttentionSpec
37
38
logger = init_logger(__name__)
vllm/v1/attention/backends/short_conv_attn.py
+1
-1
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
5
-from vllm.attention.backends.abstract import AttentionBackend
5
+from vllm.v1.attention.backend import AttentionBackend
6
from vllm.v1.attention.backends.mamba_attn import (
7
BaseMambaAttentionMetadata,
8
BaseMambaAttentionMetadataBuilder,
vllm/v1/attention/backends/tree_attn.py
+4
-4
@@ -9,20 +9,20 @@ from typing import ClassVar, Optional
9
import torch
10
11
from vllm import _custom_ops as ops
12
-from vllm.attention.backends.abstract import (
12
+from vllm.config import VllmConfig
13
+from vllm.logger import init_logger
14
+from vllm.v1.attention.backend import (
15
AttentionBackend,
16
AttentionImpl,
17
AttentionType,
18
MultipleOf,
19
)
18
-from vllm.attention.ops.triton_unified_attention import unified_attention
19
-from vllm.config import VllmConfig
20
-from vllm.logger import init_logger
20
from vllm.v1.attention.backends.utils import (
21
AttentionMetadataBuilder,
22
CommonAttentionMetadata,
23
split_decodes_and_prefills,
24
)
25
+from vllm.v1.attention.ops.triton_unified_attention import unified_attention
26
from vllm.v1.kv_cache_interface import AttentionSpec
27
28
logger = init_logger(__name__)
vllm/v1/attention/backends/triton_attn.py
+11
-11
@@ -7,17 +7,6 @@ from typing import ClassVar
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_prefill_attention import context_attention_fwd
17
-from vllm.attention.ops.triton_reshape_and_cache_flash import (
18
- triton_reshape_and_cache_flash,
19
-)
20
-from vllm.attention.ops.triton_unified_attention import unified_attention
10
from vllm.config import CUDAGraphMode, VllmConfig
11
from vllm.config.cache import CacheDType
12
from vllm.logger import init_logger
@@ -28,11 +17,22 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
from vllm.platforms import current_platform
18
from vllm.platforms.interface import DeviceCapability
19
from vllm.utils.math_utils import next_power_of_2
20
+from vllm.v1.attention.backend import (
21
+ AttentionBackend,
22
+ AttentionImpl,
23
+ AttentionType,
24
+ MultipleOf,
25
+)
26
from vllm.v1.attention.backends.utils import (
27
AttentionCGSupport,
28
AttentionMetadataBuilder,
29
CommonAttentionMetadata,
30
)
31
+from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
32
+from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
33
+ triton_reshape_and_cache_flash,
34
+)
35
+from vllm.v1.attention.ops.triton_unified_attention import unified_attention
36
from vllm.v1.kv_cache_interface import AttentionSpec
37
38
logger = init_logger(__name__)
vllm/v1/attention/backends/utils.py
+5
-5
@@ -29,16 +29,16 @@ if TYPE_CHECKING:
29
from vllm.v1.worker.gpu_input_batch import InputBatch
30
31
import vllm.envs as envs
32
-from vllm.attention.backends.abstract import (
33
- AttentionBackend,
34
- AttentionImpl,
35
- AttentionMetadata,
36
-)
32
from vllm.distributed.kv_transfer.kv_connector.utils import (
33
get_kv_connector_cache_layout,
34
)
35
from vllm.logger import init_logger
36
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
37
+from vllm.v1.attention.backend import (
38
+ AttentionBackend,
39
+ AttentionImpl,
40
+ AttentionMetadata,
41
+)
42
from vllm.v1.kv_cache_interface import AttentionSpec
43
from vllm.v1.worker.ubatch_utils import UBatchSlice
44
vllm/v1/attention/ops/__init__.py
renamed
vllm/v1/attention/ops/chunked_prefill_paged_decode.py
renamed
vllm/v1/attention/ops/common.py
renamed
vllm/v1/attention/ops/flashmla.py
renamed
vllm/v1/attention/ops/merge_attn_states.py
renamed
+1
-1
@@ -40,7 +40,7 @@ def merge_attn_states(
40
output, prefix_output, prefix_lse, suffix_output, suffix_lse, output_lse
41
)
42
else:
43
- from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
43
+ from vllm.v1.attention.ops.triton_merge_attn_states import merge_attn_states
44
45
return merge_attn_states(
46
output, prefix_output, prefix_lse, suffix_output, suffix_lse, output_lse
vllm/v1/attention/ops/paged_attn.py
renamed
vllm/v1/attention/ops/pallas_kv_cache_update.py
renamed
vllm/v1/attention/ops/prefix_prefill.py
renamed
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
renamed
vllm/v1/attention/ops/triton_decode_attention.py
renamed
vllm/v1/attention/ops/triton_merge_attn_states.py
renamed
vllm/v1/attention/ops/triton_prefill_attention.py
renamed
vllm/v1/attention/ops/triton_reshape_and_cache_flash.py
renamed
vllm/v1/attention/ops/triton_unified_attention.py
renamed
vllm/v1/attention/ops/vit_attn_wrappers.py
renamed
+1
-1
@@ -35,7 +35,7 @@ def flash_attn_maxseqlen_wrapper(
35
if is_rocm_aiter:
36
from aiter import flash_attn_varlen_func
37
else:
38
- from vllm.attention.utils.fa_utils import flash_attn_varlen_func
38
+ from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
39
40
if not current_platform.is_rocm() and fa_version is not None:
41
kwargs["fa_version"] = fa_version
vllm/v1/attention/selector.py
renamed
+5
-5
@@ -6,14 +6,14 @@ from typing import NamedTuple, cast, get_args
6
7
import torch
8
9
-from vllm.attention.backends.abstract import AttentionBackend, AttentionType
10
-from vllm.attention.backends.registry import (
11
- MAMBA_TYPE_TO_BACKEND_MAP,
12
- MambaAttentionBackendEnum,
13
-)
9
from vllm.config.cache import CacheDType
10
from vllm.logger import init_logger
11
from vllm.utils.import_utils import resolve_obj_by_qualname
12
+from vllm.v1.attention.backend import AttentionBackend, AttentionType
13
+from vllm.v1.attention.backends.registry import (
14
+ MAMBA_TYPE_TO_BACKEND_MAP,
15
+ MambaAttentionBackendEnum,
16
+)
17
18
logger = init_logger(__name__)
19
vllm/v1/kv_offload/cpu.py
+1
-1
@@ -4,9 +4,9 @@ from collections.abc import Iterator
4
5
import torch
6
7
-from vllm.attention.backends.abstract import AttentionBackend
7
from vllm.config import VllmConfig
8
from vllm.platforms import current_platform
9
+from vllm.v1.attention.backend import AttentionBackend
10
from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
11
from vllm.v1.kv_offload.arc_manager import ARCOffloadingManager
12
from vllm.v1.kv_offload.backends.cpu import CPUBackend
vllm/v1/kv_offload/spec.py
+1
-1
@@ -6,8 +6,8 @@ from typing import TYPE_CHECKING
6
7
import torch
8
9
-from vllm.attention.backends.abstract import AttentionBackend
9
from vllm.logger import init_logger
10
+from vllm.v1.attention.backend import AttentionBackend
11
from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
12
from vllm.v1.kv_offload.worker.worker import OffloadingHandler
13
vllm/v1/kv_offload/worker/cpu_gpu.py
+1
-1
@@ -6,9 +6,9 @@ import numpy as np
6
import torch
7
8
from vllm import _custom_ops as ops
9
-from vllm.attention.backends.abstract import AttentionBackend
9
from vllm.logger import init_logger
10
from vllm.utils.platform_utils import is_pin_memory_available
11
+from vllm.v1.attention.backend import AttentionBackend
12
from vllm.v1.kv_offload.mediums import BlockIDsLoadStoreSpec
13
from vllm.v1.kv_offload.worker.worker import (
14
OffloadingHandler,
vllm/v1/spec_decode/eagle.py
+1
-1
@@ -8,7 +8,6 @@ import numpy as np
8
import torch
9
import torch.nn as nn
10
11
-from vllm.attention.backends.registry import AttentionBackendEnum
11
from vllm.config import (
12
CompilationMode,
13
CUDAGraphMode,
@@ -27,6 +26,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
26
from vllm.platforms import current_platform
27
from vllm.triton_utils import triton
28
from vllm.utils.platform_utils import is_pin_memory_available
29
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
30
from vllm.v1.attention.backends.tree_attn import (
31
TreeAttentionMetadata,
32
TreeAttentionMetadataBuilder,
vllm/v1/worker/gpu/attn_utils.py
+1
-1
@@ -6,9 +6,9 @@ from typing import Any, cast
6
import numpy as np
7
import torch
8
9
-from vllm.attention.backends.abstract import AttentionBackend
9
from vllm.config import VllmConfig, get_layers_from_vllm_config
10
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
11
+from vllm.v1.attention.backend import AttentionBackend
12
from vllm.v1.attention.backends.utils import (
13
AttentionMetadataBuilder,
14
CommonAttentionMetadata,
vllm/v1/worker/gpu_model_runner.py
+6
-6
@@ -20,12 +20,6 @@ import torch.nn as nn
20
from tqdm import tqdm
21
22
import vllm.envs as envs
23
-from vllm.attention.backends.abstract import (
24
- AttentionBackend,
25
- AttentionMetadata,
26
- AttentionType,
27
- MultipleOf,
28
-)
23
from vllm.attention.layer import Attention, MLAAttention
24
from vllm.compilation.counter import compilation_counter
25
from vllm.compilation.cuda_graph import CUDAGraphStat, CUDAGraphWrapper
@@ -101,6 +95,12 @@ from vllm.utils.torch_utils import (
95
kv_cache_dtype_str_to_dtype,
96
supports_dynamo,
97
)
98
+from vllm.v1.attention.backend import (
99
+ AttentionBackend,
100
+ AttentionMetadata,
101
+ AttentionType,
102
+ MultipleOf,
103
+)
104
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
105
from vllm.v1.attention.backends.utils import (
106
AttentionCGSupport,
vllm/v1/worker/kv_connector_model_runner_mixin.py
+1
-1
@@ -11,7 +11,6 @@ from typing import TYPE_CHECKING
11
12
import torch
13
14
-from vllm.attention.backends.abstract import AttentionBackend
14
from vllm.config import VllmConfig
15
from vllm.config.cache import CacheDType
16
from vllm.distributed.kv_transfer import (
@@ -22,6 +21,7 @@ from vllm.distributed.kv_transfer import (
21
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
22
from vllm.forward_context import get_forward_context, set_forward_context
23
from vllm.logger import init_logger
24
+from vllm.v1.attention.backend import AttentionBackend
25
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
26
from vllm.v1.outputs import (
27
EMPTY_MODEL_RUNNER_OUTPUT,
vllm/v1/worker/utils.py
+1
-1
@@ -7,7 +7,6 @@ from dataclasses import dataclass, field
7
import torch
8
from typing_extensions import deprecated
9
10
-from vllm.attention.backends.abstract import AttentionBackend
10
from vllm.attention.layer import Attention
11
from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
12
from vllm.logger import init_logger
@@ -17,6 +16,7 @@ from vllm.multimodal.cache import processor_only_cache_from_config
16
from vllm.multimodal.registry import MultiModalRegistry
17
from vllm.platforms import current_platform
18
from vllm.utils.mem_utils import MemorySnapshot, format_gib
19
+from vllm.v1.attention.backend import AttentionBackend
20
from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
21
from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
22
from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec