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 14, 05:30 UTC
[Attention] Move Backend enum into registry (#25893)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Oct 2, 2025 at 23:32 UTC
2aaa42384280543078d0bdb6295128e541c22518
31 files changed
+99
-66
tests/compile/test_full_graph.py
+1
-1
@@ -11,8 +11,8 @@ import pytest
11
import torch
12
13
from tests.quantization.utils import is_quant_method_supported
14
-from tests.v1.attention.utils import _Backend
14
from vllm import LLM, SamplingParams
15
+from vllm.attention.backends.registry import _Backend
16
from vllm.attention.selector import global_force_attn_backend_context_manager
17
from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
18
PassConfig)
tests/compile/test_fusion_attn.py
+2
-2
@@ -8,11 +8,11 @@ import torch._dynamo
8
9
from tests.compile.backend import LazyInitPass, TestBackend
10
from tests.models.utils import check_outputs_equal
11
-from tests.v1.attention.utils import (BatchSpec, _Backend,
12
- create_common_attn_metadata)
11
+from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
12
from vllm import LLM, SamplingParams
13
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
14
from vllm.attention import Attention, AttentionMetadata
15
+from vllm.attention.backends.registry import _Backend
16
from vllm.attention.selector import global_force_attn_backend_context_manager
17
from vllm.compilation.fusion import QUANT_OPS
18
from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
tests/kernels/attention/test_mha_attn.py
+2
-1
@@ -10,8 +10,9 @@ from unittest.mock import patch
10
import pytest
11
import torch
12
13
+from vllm.attention.backends.registry import _Backend
14
from vllm.attention.layer import MultiHeadAttention
14
-from vllm.attention.selector import _Backend, _cached_get_attn_backend
15
+from vllm.attention.selector import _cached_get_attn_backend
16
from vllm.platforms import current_platform
17
from vllm.platforms.cpu import CpuPlatform
18
from vllm.platforms.cuda import CudaPlatform
tests/kernels/utils.py
+1
-1
@@ -15,10 +15,10 @@ from torch._prims_common import TensorLikeType
15
16
from tests.kernels.quant_utils import native_w8a8_block_matmul
17
from vllm.attention import AttentionBackend, AttentionMetadata, AttentionType
18
+from vllm.attention.backends.registry import _Backend
19
from vllm.model_executor.layers.activation import SiluAndMul
20
from vllm.model_executor.layers.fused_moe.utils import (
21
moe_kernel_quantize_input)
21
-from vllm.platforms.interface import _Backend
22
from vllm.utils import (STR_BACKEND_ENV_VAR, STR_FLASH_ATTN_VAL,
23
STR_XFORMERS_ATTN_VAL, make_tensor_with_pad)
24
tests/v1/attention/test_attention_backends.py
+2
-2
@@ -8,11 +8,11 @@ import pytest
8
import torch
9
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
10
11
-from tests.v1.attention.utils import (BatchSpec, _Backend,
12
- create_common_attn_metadata,
11
+from tests.v1.attention.utils import (BatchSpec, create_common_attn_metadata,
12
create_standard_kv_cache_spec,
13
create_vllm_config,
14
get_attention_backend)
15
+from vllm.attention.backends.registry import _Backend
16
from vllm.config import ModelConfig
17
from vllm.platforms import current_platform
18
from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv, is_torch_equal_or_newer
tests/v1/attention/test_mla_backends.py
+2
-2
@@ -6,12 +6,12 @@ from typing import Optional, Union
6
import pytest
7
import torch
8
9
-from tests.v1.attention.utils import (BatchSpec, _Backend,
10
- create_common_attn_metadata,
9
+from tests.v1.attention.utils import (BatchSpec, create_common_attn_metadata,
10
create_standard_kv_cache_spec,
11
create_vllm_config,
12
get_attention_backend)
13
from vllm import _custom_ops as ops
14
+from vllm.attention.backends.registry import _Backend
15
from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
16
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
17
from vllm.v1.kv_cache_interface import FullAttentionSpec
tests/v1/attention/utils.py
+2
-1
@@ -8,10 +8,11 @@ from typing import Optional, Union
8
import pytest
9
import torch
10
11
+from vllm.attention.backends.registry import _Backend
12
from vllm.config import (CacheConfig, CompilationConfig, DeviceConfig,
13
LoadConfig, ModelConfig, ModelDType, ParallelConfig,
14
SchedulerConfig, VllmConfig)
14
-from vllm.platforms import _Backend, current_platform
15
+from vllm.platforms import current_platform
16
from vllm.utils import resolve_obj_by_qualname
17
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
18
from vllm.v1.kv_cache_interface import FullAttentionSpec
tests/v1/spec_decode/test_eagle.py
+2
-2
@@ -8,10 +8,10 @@ import pytest
8
import torch
9
10
from tests.utils import get_attn_backend_list_based_on_platform
11
-from tests.v1.attention.utils import (BatchSpec, _Backend,
12
- create_common_attn_metadata,
11
+from tests.v1.attention.utils import (BatchSpec, create_common_attn_metadata,
12
create_standard_kv_cache_spec,
13
get_attention_backend)
14
+from vllm.attention.backends.registry import _Backend
15
from vllm.config import (CacheConfig, DeviceConfig, ModelConfig,
16
ParallelConfig, SchedulerConfig, SpeculativeConfig,
17
VllmConfig)
tests/v1/spec_decode/test_mtp.py
+2
-2
@@ -6,10 +6,10 @@ from unittest import mock
6
import pytest
7
import torch
8
9
-from tests.v1.attention.utils import (BatchSpec, _Backend,
10
- create_common_attn_metadata,
9
+from tests.v1.attention.utils import (BatchSpec, create_common_attn_metadata,
10
create_standard_kv_cache_spec,
11
get_attention_backend)
12
+from vllm.attention.backends.registry import _Backend
13
from vllm.config import (CacheConfig, DeviceConfig, ModelConfig,
14
ParallelConfig, SchedulerConfig, SpeculativeConfig,
15
VllmConfig)
tests/v1/spec_decode/test_tree_attention.py
+2
-1
@@ -6,9 +6,10 @@ from typing import Optional
6
7
import torch
8
9
-from tests.v1.attention.utils import (_Backend, create_standard_kv_cache_spec,
9
+from tests.v1.attention.utils import (create_standard_kv_cache_spec,
10
create_vllm_config,
11
get_attention_backend)
12
+from vllm.attention.backends.registry import _Backend
13
from vllm.config import ParallelConfig, SpeculativeConfig
14
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
15
vllm/attention/backends/registry.py
new
+27
@@ -0,0 +1,27 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Attention backend registry"""
4
+
5
+import enum
6
+
7
+
8
+class _Backend(enum.Enum):
9
+ FLASH_ATTN = enum.auto()
10
+ TRITON_ATTN = enum.auto()
11
+ XFORMERS = enum.auto()
12
+ ROCM_FLASH = enum.auto()
13
+ ROCM_AITER_MLA = enum.auto()
14
+ ROCM_AITER_FA = enum.auto() # used for ViT attn backend
15
+ TORCH_SDPA = enum.auto()
16
+ FLASHINFER = enum.auto()
17
+ FLASHINFER_MLA = enum.auto()
18
+ TRITON_MLA = enum.auto()
19
+ CUTLASS_MLA = enum.auto()
20
+ FLASHMLA = enum.auto()
21
+ FLASH_ATTN_MLA = enum.auto()
22
+ PALLAS = enum.auto()
23
+ IPEX = enum.auto()
24
+ NO_ATTENTION = enum.auto()
25
+ FLEX_ATTENTION = enum.auto()
26
+ TREE_ATTN = enum.auto()
27
+ ROCM_ATTN = enum.auto()
vllm/attention/layer.py
+2
-1
@@ -10,6 +10,7 @@ import torch.nn.functional as F
10
import vllm.envs as envs
11
from vllm.attention import AttentionType
12
from vllm.attention.backends.abstract import AttentionBackend
13
+from vllm.attention.backends.registry import _Backend
14
from vllm.attention.selector import backend_name_to_enum, get_attn_backend
15
from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
16
from vllm.config import CacheConfig, get_current_vllm_config
@@ -26,7 +27,7 @@ from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
27
from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
GroupShape)
29
from vllm.model_executor.models.vision import get_vit_attn_backend
29
-from vllm.platforms import _Backend, current_platform
30
+from vllm.platforms import current_platform
31
from vllm.utils import GiB_bytes, direct_register_custom_op
32
33
logger = init_logger(__name__)
vllm/attention/selector.py
+2
-1
@@ -11,8 +11,9 @@ import torch
11
12
import vllm.envs as envs
13
from vllm.attention.backends.abstract import AttentionBackend
14
+from vllm.attention.backends.registry import _Backend
15
from vllm.logger import init_logger
15
-from vllm.platforms import _Backend, current_platform
16
+from vllm.platforms import current_platform
17
from vllm.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
18
19
logger = init_logger(__name__)
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+2
-1
@@ -20,6 +20,7 @@ import torch
20
import zmq
21
22
from vllm import envs
23
+from vllm.attention.backends.registry import _Backend
24
from vllm.attention.selector import backend_name_to_enum, get_attn_backend
25
from vllm.config import VllmConfig
26
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@@ -32,7 +33,7 @@ from vllm.distributed.parallel_state import (
33
from vllm.distributed.utils import divide
34
from vllm.forward_context import ForwardContext
35
from vllm.logger import init_logger
35
-from vllm.platforms import _Backend, current_platform
36
+from vllm.platforms import current_platform
37
from vllm.utils import make_zmq_path, make_zmq_socket
38
from vllm.v1.attention.backends.utils import get_kv_cache_layout
39
from vllm.v1.core.sched.output import SchedulerOutput
vllm/envs.py
+3
-2
@@ -619,8 +619,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
619
# All possible options loaded dynamically from _Backend enum
620
"VLLM_ATTENTION_BACKEND":
621
env_with_choices("VLLM_ATTENTION_BACKEND", None,
622
- lambda: list(__import__('vllm.platforms.interface', \
623
- fromlist=['_Backend'])._Backend.__members__.keys())),
622
+ lambda: list(__import__(
623
+ 'vllm.attention.backends.registry',
624
+ fromlist=['_Backend'])._Backend.__members__.keys())),
625
626
# If set, vllm will use flashinfer sampler
627
"VLLM_USE_FLASHINFER_SAMPLER":
vllm/model_executor/models/dots_ocr.py
+1
-1
@@ -9,6 +9,7 @@ import torch.nn.functional as F
9
from torch.nn import LayerNorm
10
from transformers.models.qwen2_vl import Qwen2VLProcessor
11
12
+from vllm.attention.backends.registry import _Backend
13
from vllm.attention.layer import check_upstream_fa_availability
14
from vllm.config import VllmConfig
15
from vllm.distributed import utils as dist_utils
@@ -38,7 +39,6 @@ from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
39
from vllm.model_executor.models.vision import get_vit_attn_backend
40
from vllm.multimodal import MULTIMODAL_REGISTRY
41
from vllm.multimodal.inputs import MultiModalDataDict
41
-from vllm.platforms import _Backend
42
from vllm.sequence import IntermediateTensors
43
from vllm.transformers_utils.configs.dotsocr import (DotsOCRConfig,
44
DotsVisionConfig)
vllm/model_executor/models/ernie45_vl.py
+2
-1
@@ -34,6 +34,7 @@ import torch.nn.functional as F
34
from einops import rearrange, repeat
35
from transformers import BatchFeature
36
37
+from vllm.attention.backends.registry import _Backend
38
from vllm.attention.layer import check_upstream_fa_availability
39
from vllm.config import VllmConfig
40
from vllm.distributed import parallel_state
@@ -54,7 +55,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
55
BaseProcessingInfo, PromptReplacement,
56
PromptUpdate)
57
from vllm.multimodal.profiling import BaseDummyInputsBuilder
57
-from vllm.platforms import _Backend, current_platform
58
+from vllm.platforms import current_platform
59
from vllm.sequence import IntermediateTensors
60
from vllm.utils.tensor_schema import TensorSchema, TensorShape
61
vllm/model_executor/models/glm4_1v.py
+1
-1
@@ -46,6 +46,7 @@ from transformers.models.glm4v.video_processing_glm4v import (
46
Glm4vVideoProcessor)
47
from transformers.video_utils import VideoMetadata
48
49
+from vllm.attention.backends.registry import _Backend
50
from vllm.attention.layer import check_upstream_fa_availability
51
from vllm.config import VllmConfig
52
from vllm.distributed import (get_tensor_model_parallel_world_size,
@@ -69,7 +70,6 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
70
BaseProcessingInfo, PromptReplacement,
71
PromptUpdate, PromptUpdateDetails)
72
from vllm.multimodal.profiling import BaseDummyInputsBuilder
72
-from vllm.platforms import _Backend
73
from vllm.sequence import IntermediateTensors
74
from vllm.utils.tensor_schema import TensorSchema, TensorShape
75
vllm/model_executor/models/keye.py
+1
-1
@@ -17,6 +17,7 @@ from transformers.modeling_outputs import (BaseModelOutput,
17
BaseModelOutputWithPooling)
18
from transformers.utils import torch_int
19
20
+from vllm.attention.backends.registry import _Backend
21
from vllm.attention.layer import check_upstream_fa_availability
22
from vllm.config import VllmConfig
23
from vllm.distributed import get_tensor_model_parallel_world_size
@@ -39,7 +40,6 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
40
BaseProcessingInfo, PromptReplacement,
41
PromptUpdate)
42
from vllm.multimodal.profiling import BaseDummyInputsBuilder
42
-from vllm.platforms import _Backend
43
from vllm.sequence import IntermediateTensors
44
from vllm.utils import is_list_of
45
from vllm.utils.tensor_schema import TensorSchema, TensorShape
vllm/model_executor/models/qwen2_5_vl.py
+1
-1
@@ -38,6 +38,7 @@ from transformers.models.qwen2_5_vl import Qwen2_5_VLProcessor
38
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
39
Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig)
40
41
+from vllm.attention.backends.registry import _Backend
42
from vllm.attention.layer import check_upstream_fa_availability
43
from vllm.config import VllmConfig
44
from vllm.distributed import parallel_state
@@ -62,7 +63,6 @@ from vllm.multimodal.evs import (compute_mrope_for_media,
63
from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargs
64
from vllm.multimodal.parse import MultiModalDataItems
65
from vllm.multimodal.processing import PromptReplacement, PromptUpdate
65
-from vllm.platforms import _Backend
66
from vllm.sequence import IntermediateTensors
67
from vllm.utils import is_pin_memory_available
68
from vllm.utils.tensor_schema import TensorSchema, TensorShape
vllm/model_executor/models/qwen2_vl.py
+1
-1
@@ -41,6 +41,7 @@ from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
41
from transformers.models.qwen2_vl.video_processing_qwen2_vl import (
42
Qwen2VLVideoProcessor)
43
44
+from vllm.attention.backends.registry import _Backend
45
from vllm.attention.layer import check_upstream_fa_availability
46
from vllm.config import VllmConfig
47
from vllm.distributed import parallel_state, tensor_model_parallel_all_gather
@@ -65,7 +66,6 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
66
BaseProcessingInfo, PromptReplacement,
67
PromptUpdate)
68
from vllm.multimodal.profiling import BaseDummyInputsBuilder
68
-from vllm.platforms import _Backend
69
from vllm.sequence import IntermediateTensors
70
from vllm.transformers_utils.tokenizer import AnyTokenizer
71
from vllm.utils.tensor_schema import TensorSchema, TensorShape
vllm/model_executor/models/qwen3_vl.py
+1
-1
@@ -43,6 +43,7 @@ from transformers.models.qwen3_vl.video_processing_qwen3_vl import (
43
smart_resize as video_smart_resize)
44
from transformers.video_utils import VideoMetadata
45
46
+from vllm.attention.backends.registry import _Backend
47
from vllm.attention.layer import check_upstream_fa_availability
48
from vllm.compilation.decorators import support_torch_compile
49
from vllm.config import VllmConfig
@@ -66,7 +67,6 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
67
PromptReplacement, PromptUpdate,
68
PromptUpdateDetails)
69
from vllm.multimodal.profiling import BaseDummyInputsBuilder
69
-from vllm.platforms import _Backend
70
from vllm.sequence import IntermediateTensors
71
from vllm.utils import is_list_of
72
vllm/model_executor/models/siglip2navit.py
+1
-1
@@ -13,6 +13,7 @@ from torch.nn import functional as F
13
from transformers import Siglip2VisionConfig
14
from transformers.configuration_utils import PretrainedConfig
15
16
+from vllm.attention.backends.registry import _Backend
17
from vllm.attention.layer import check_upstream_fa_availability
18
from vllm.distributed import divide, get_tensor_model_parallel_world_size
19
from vllm.model_executor.layers.activation import get_act_fn
@@ -22,7 +23,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
23
RowParallelLinear)
24
from vllm.model_executor.layers.quantization import QuantizationConfig
25
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25
-from vllm.platforms import _Backend
26
27
from .vision import get_vit_attn_backend
28
vllm/model_executor/models/vision.py
+2
-1
@@ -10,11 +10,12 @@ from typing import (Callable, Final, Generic, Literal, Optional, Protocol,
10
import torch
11
from transformers import PretrainedConfig
12
13
+from vllm.attention.backends.registry import _Backend
14
from vllm.distributed import (get_tensor_model_parallel_rank,
15
get_tensor_model_parallel_world_size,
16
tensor_model_parallel_all_gather)
17
from vllm.logger import init_logger
17
-from vllm.platforms import _Backend, current_platform
18
+from vllm.platforms import current_platform
19
20
logger = init_logger(__name__)
21
vllm/platforms/__init__.py
-1
@@ -9,7 +9,6 @@ from vllm import envs
9
from vllm.plugins import load_plugins_by_group
10
from vllm.utils import resolve_obj_by_qualname, supports_xccl
11
12
-from .interface import _Backend # noqa: F401
12
from .interface import CpuArchEnum, Platform, PlatformEnum
13
14
logger = logging.getLogger(__name__)
vllm/platforms/cpu.py
+5
-2
@@ -15,13 +15,15 @@ import torch
15
from vllm.logger import init_logger
16
from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
17
18
-from .interface import CpuArchEnum, Platform, PlatformEnum, _Backend
18
+from .interface import CpuArchEnum, Platform, PlatformEnum
19
20
logger = init_logger(__name__)
21
22
if TYPE_CHECKING:
23
+ from vllm.attention.backends.registry import _Backend
24
from vllm.config import VllmConfig
25
else:
26
+ _Backend = None
27
VllmConfig = None
28
29
@@ -90,10 +92,11 @@ class CpuPlatform(Platform):
92
return "cpu"
93
94
@classmethod
93
- def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
95
+ def get_attn_backend_cls(cls, selected_backend: "_Backend", head_size: int,
96
dtype: torch.dtype, kv_cache_dtype: Optional[str],
97
block_size: int, use_v1: bool, use_mla: bool,
98
has_sink: bool, use_sparse: bool) -> str:
99
+ from vllm.attention.backends.registry import _Backend
100
if selected_backend and selected_backend != _Backend.TORCH_SDPA:
101
logger.info("Cannot use %s backend on CPU.", selected_backend)
102
if use_mla:
vllm/platforms/cuda.py
+7
-2
@@ -20,10 +20,13 @@ import vllm.envs as envs
20
from vllm.logger import init_logger
21
from vllm.utils import cuda_device_count_stateless, import_pynvml
22
23
-from .interface import DeviceCapability, Platform, PlatformEnum, _Backend
23
+from .interface import DeviceCapability, Platform, PlatformEnum
24
25
if TYPE_CHECKING:
26
+ from vllm.attention.backends.registry import _Backend
27
from vllm.config import ModelConfig, VllmConfig
28
+else:
29
+ _Backend = None
30
31
logger = init_logger(__name__)
32
@@ -202,7 +205,8 @@ class CudaPlatformBase(Platform):
205
206
@classmethod
207
def get_vit_attn_backend(cls, head_size: int,
205
- dtype: torch.dtype) -> _Backend:
208
+ dtype: torch.dtype) -> "_Backend":
209
+ from vllm.attention.backends.registry import _Backend
210
211
# For Blackwell GPUs, force TORCH_SDPA for now.
212
# See https://github.com/facebookresearch/xformers/issues/1317#issuecomment-3199392579 # noqa: E501
@@ -230,6 +234,7 @@ class CudaPlatformBase(Platform):
234
def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
235
kv_cache_dtype, block_size, use_v1, use_mla,
236
has_sink, use_sparse) -> str:
237
+ from vllm.attention.backends.registry import _Backend
238
if use_mla:
239
if not use_v1:
240
raise RuntimeError(
vllm/platforms/interface.py
+5
-26
@@ -17,12 +17,14 @@ from vllm.inputs import ProcessorInputs, PromptType
17
from vllm.logger import init_logger
18
19
if TYPE_CHECKING:
20
+ from vllm.attention.backends.registry import _Backend
21
from vllm.config import ModelConfig, VllmConfig
22
from vllm.lora.request import LoRARequest
23
from vllm.pooling_params import PoolingParams
24
from vllm.sampling_params import SamplingParams
25
from vllm.utils import FlexibleArgumentParser
26
else:
27
+ _Backend = None
28
ModelConfig = None
29
VllmConfig = None
30
LoRARequest = None
@@ -38,30 +40,6 @@ def in_wsl() -> bool:
40
return "microsoft" in " ".join(uname()).lower()
41
42
41
-class _Backend(enum.Enum):
42
- FLASH_ATTN = enum.auto()
43
- TRITON_ATTN = enum.auto()
44
- XFORMERS = enum.auto()
45
- ROCM_FLASH = enum.auto()
46
- ROCM_AITER_MLA = enum.auto() # Supported by V1
47
- ROCM_AITER_FA = enum.auto() # used for ViT attn backend
48
- TORCH_SDPA = enum.auto()
49
- FLASHINFER = enum.auto()
50
- FLASHINFER_MLA = enum.auto()
51
- TRITON_MLA = enum.auto() # Supported by V1
52
- CUTLASS_MLA = enum.auto()
53
- FLASHMLA = enum.auto() # Supported by V1
54
- FLASH_ATTN_MLA = enum.auto() # Supported by V1
55
- PALLAS = enum.auto()
56
- IPEX = enum.auto()
57
- DUAL_CHUNK_FLASH_ATTN = enum.auto()
58
- DIFFERENTIAL_FLASH_ATTN = enum.auto()
59
- NO_ATTENTION = enum.auto()
60
- FLEX_ATTENTION = enum.auto()
61
- TREE_ATTN = enum.auto()
62
- ROCM_ATTN = enum.auto()
63
-
64
-
43
class PlatformEnum(enum.Enum):
44
CUDA = enum.auto()
45
ROCM = enum.auto()
@@ -187,11 +165,12 @@ class Platform:
165
166
@classmethod
167
def get_vit_attn_backend(cls, head_size: int,
190
- dtype: torch.dtype) -> _Backend:
168
+ dtype: torch.dtype) -> "_Backend":
169
+ from vllm.attention.backends.registry import _Backend
170
return _Backend.TORCH_SDPA
171
172
@classmethod
194
- def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
173
+ def get_attn_backend_cls(cls, selected_backend: "_Backend", head_size: int,
174
dtype: torch.dtype, kv_cache_dtype: Optional[str],
175
block_size: int, use_v1: bool, use_mla: bool,
176
has_sink: bool, use_sparse: bool) -> str:
vllm/platforms/rocm.py
+7
-2
@@ -14,10 +14,13 @@ import vllm.envs as envs
14
from vllm.logger import init_logger
15
from vllm.utils import cuda_device_count_stateless
16
17
-from .interface import DeviceCapability, Platform, PlatformEnum, _Backend
17
+from .interface import DeviceCapability, Platform, PlatformEnum
18
19
if TYPE_CHECKING:
20
+ from vllm.attention.backends.registry import _Backend
21
from vllm.config import ModelConfig, VllmConfig
22
+else:
23
+ _Backend = None
24
25
logger = init_logger(__name__)
26
@@ -182,7 +185,8 @@ class RocmPlatform(Platform):
185
186
@classmethod
187
def get_vit_attn_backend(cls, head_size: int,
185
- dtype: torch.dtype) -> _Backend:
188
+ dtype: torch.dtype) -> "_Backend":
189
+ from vllm.attention.backends.registry import _Backend
190
if (envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA
191
and on_gfx9()):
192
# Note: AITER FA is only supported for Qwen-VL models.
@@ -196,6 +200,7 @@ class RocmPlatform(Platform):
200
def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
201
kv_cache_dtype, block_size, use_v1, use_mla,
202
has_sink, use_sparse) -> str:
203
+ from vllm.attention.backends.registry import _Backend
204
if use_sparse:
205
raise NotImplementedError(
206
"Sparse Attention is not supported on ROCm.")
vllm/platforms/tpu.py
+5
-2
@@ -11,9 +11,10 @@ from vllm.logger import init_logger
11
from vllm.sampling_params import SamplingParams, SamplingType
12
from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
13
14
-from .interface import Platform, PlatformEnum, _Backend
14
+from .interface import Platform, PlatformEnum
15
16
if TYPE_CHECKING:
17
+ from vllm.attention.backends.registry import _Backend
18
from vllm.config import BlockSize, ModelConfig, VllmConfig
19
from vllm.pooling_params import PoolingParams
20
else:
@@ -21,6 +22,7 @@ else:
22
ModelConfig = None
23
VllmConfig = None
24
PoolingParams = None
25
+ _Backend = None
26
27
logger = init_logger(__name__)
28
@@ -46,10 +48,11 @@ class TpuPlatform(Platform):
48
]
49
50
@classmethod
49
- def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
51
+ def get_attn_backend_cls(cls, selected_backend: "_Backend", head_size: int,
52
dtype: torch.dtype, kv_cache_dtype: Optional[str],
53
block_size: int, use_v1: bool, use_mla: bool,
54
has_sink, use_sparse) -> str:
55
+ from vllm.attention.backends.registry import _Backend
56
if use_sparse:
57
raise NotImplementedError(
58
"Sparse Attention is not supported on TPU.")
vllm/platforms/xpu.py
+5
-2
@@ -10,13 +10,15 @@ import vllm.envs as envs
10
from vllm.logger import init_logger
11
from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
12
13
-from .interface import DeviceCapability, Platform, PlatformEnum, _Backend
13
+from .interface import DeviceCapability, Platform, PlatformEnum
14
15
if TYPE_CHECKING:
16
+ from vllm.attention.backends.registry import _Backend
17
from vllm.config import ModelConfig, VllmConfig
18
else:
19
ModelConfig = None
20
VllmConfig = None
21
+ _Backend = None
22
23
logger = init_logger(__name__)
24
@@ -33,10 +35,11 @@ class XPUPlatform(Platform):
35
device_control_env_var: str = "ZE_AFFINITY_MASK"
36
37
@classmethod
36
- def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
38
+ def get_attn_backend_cls(cls, selected_backend: "_Backend", head_size: int,
39
dtype: torch.dtype, kv_cache_dtype: Optional[str],
40
block_size: int, use_v1: bool, use_mla: bool,
41
has_sink: bool, use_sparse) -> str:
42
+ from vllm.attention.backends.registry import _Backend
43
if use_sparse:
44
raise NotImplementedError(
45
"Sparse Attention is not supported on XPU.")