fc1d8be3dc · tduncan/Nvidia-Optimization-Stack-VLLM

tests/v1/attention/test_rocm_attention_backends_selection.py

+3 -6

@@ -139,14 +139,13 @@ def test_standard_attention_backend_selection(

139 import importlib

140

141 import vllm.envs as envs

142 - from vllm.attention.backends.registry import _Backend

142

143 importlib.reload(envs)

144

145 # Convert string backend to enum if provided

146 backend_enum = None

147 if selected_backend:

149 - backend_enum = getattr(_Backend, selected_backend)

148 + backend_enum = getattr(AttentionBackendEnum, selected_backend)

149

150 # Get the backend class path

151 from vllm.platforms.rocm import RocmPlatform

@@ -253,7 +252,6 @@ def test_mla_backend_selection(

252 import importlib

253

254 import vllm.envs as envs

256 - from vllm.attention.backends.registry import _Backend

255

256 importlib.reload(envs)

257

@@ -269,7 +267,7 @@ def test_mla_backend_selection(

267 # Convert string backend to enum if provided

268 backend_enum = None

269 if selected_backend:

272 - backend_enum = getattr(_Backend, selected_backend)

270 + backend_enum = getattr(AttentionBackendEnum, selected_backend)

271

272 from vllm.platforms.rocm import RocmPlatform

273

@@ -301,7 +299,6 @@ def test_mla_backend_selection(

299

300 def test_aiter_fa_requires_gfx9(mock_vllm_config):

301 """Test that ROCM_AITER_FA requires gfx9 architecture."""

304 - from vllm.attention.backends.registry import _Backend

302 from vllm.platforms.rocm import RocmPlatform

303

304 # Mock on_gfx9 to return False

@@ -313,7 +310,7 @@ def test_aiter_fa_requires_gfx9(mock_vllm_config):

310 ),

311 ):

312 RocmPlatform.get_attn_backend_cls(

316 - selected_backend=_Backend.ROCM_AITER_FA,

313 + selected_backend=AttentionBackendEnum.ROCM_AITER_FA,

314 head_size=128,

315 dtype=torch.float16,

316 kv_cache_dtype="auto",

tests/v1/kv_connector/unit/test_backwards_compatibility.py

+3 -3

@@ -14,6 +14,7 @@ from unittest.mock import patch

14

15 import pytest

16

17 +from vllm.attention.backends.abstract import AttentionMetadata

18 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory

19 from vllm.distributed.kv_transfer.kv_connector.v1 import (

20 KVConnectorBase_V1,

@@ -24,7 +25,6 @@ from vllm.v1.core.sched.output import SchedulerOutput

25 from .utils import create_scheduler, create_vllm_config

26

27 if TYPE_CHECKING:

27 - from vllm.attention.backends.abstract import AttentionMetadata

28 from vllm.config import VllmConfig

29 from vllm.forward_context import ForwardContext

30 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

@@ -68,7 +68,7 @@ class OldStyleTestConnector(KVConnectorBase_V1):

68 self,

69 layer_name: str,

70 kv_layer,

71 - attn_metadata: "AttentionMetadata",

71 + attn_metadata: AttentionMetadata,

72 **kwargs,

73 ) -> None:

74 pass

@@ -119,7 +119,7 @@ class NewStyleTestConnector(KVConnectorBase_V1):

119 self,

120 layer_name: str,

121 kv_layer,

122 - attn_metadata: "AttentionMetadata",

122 + attn_metadata: AttentionMetadata,

123 **kwargs,

124 ) -> None:

125 pass

vllm/attention/backends/abstract.py

+4 -7

@@ -6,11 +6,10 @@ from typing import TYPE_CHECKING, ClassVar, Generic, Protocol, TypeVar, get_args

6

7 import torch

8

9 -from vllm.model_executor.layers.linear import ColumnParallelLinear

10 -from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey

11 -

9 if TYPE_CHECKING:

10 from vllm.config.cache import CacheDType

11 + from vllm.model_executor.layers.linear import ColumnParallelLinear

12 + from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey

13 from vllm.platforms.interface import DeviceCapability

14 from vllm.v1.attention.backends.utils import KVCacheLayoutType

15

@@ -178,8 +177,6 @@ class AttentionBackend(ABC):

177 By default, only supports decoder attention.

178 Backends should override this to support other attention types.

179 """

181 - from vllm.attention.backends.abstract import AttentionType

182 -

180 return attn_type == AttentionType.DECODER

181

182 @classmethod

@@ -360,7 +357,7 @@ class AttentionImpl(ABC, Generic[T]):

357 ) -> torch.Tensor:

358 raise NotImplementedError

359

363 - def fused_output_quant_supported(self, quant_key: QuantKey):

360 + def fused_output_quant_supported(self, quant_key: "QuantKey"):

361 """

362 Does this attention implementation support fused output quantization.

363 This is used by the AttnFusionPass to only fuse output quantization

@@ -412,7 +409,7 @@ class MLAAttentionImpl(AttentionImpl[T], Generic[T]):

409 qk_rope_head_dim: int,

410 qk_head_dim: int,

411 v_head_dim: int,

415 - kv_b_proj: ColumnParallelLinear,

412 + kv_b_proj: "ColumnParallelLinear",

413 indexer: object | None = None,

414 ) -> None:

415 raise NotImplementedError

vllm/attention/layers/chunked_local_attention.py

+1 -2

@@ -5,6 +5,7 @@ import functools

5 import torch

6

7 from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata

8 +from vllm.attention.layer import Attention

9 from vllm.attention.selector import get_attn_backend

10 from vllm.config import CacheConfig

11 from vllm.config.vllm import VllmConfig

@@ -22,8 +23,6 @@ from vllm.v1.kv_cache_interface import (

23 KVCacheSpec,

24 )

25

25 -from ..layer import Attention

26 -

26

27 @functools.lru_cache

28 def create_chunked_local_attention_backend(

vllm/config/model.py

+1 -2

@@ -14,6 +14,7 @@ from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE

14 from transformers.configuration_utils import ALLOWED_LAYER_TYPES

15

16 import vllm.envs as envs

17 +from vllm.attention.backends.registry import AttentionBackendEnum

18 from vllm.config.multimodal import MMCacheType, MMEncoderTPMode, MultiModalConfig

19 from vllm.config.pooler import PoolerConfig

20 from vllm.config.scheduler import RunnerType

@@ -53,7 +54,6 @@ if TYPE_CHECKING:

54

55 import vllm.model_executor.layers.quantization as me_quant

56 import vllm.model_executor.models as me_models

56 - from vllm.attention.backends.registry import AttentionBackendEnum

57 from vllm.config.load import LoadConfig

58 from vllm.config.parallel import ParallelConfig

59 from vllm.model_executor.layers.quantization import QuantizationMethods

@@ -61,7 +61,6 @@ if TYPE_CHECKING:

61 else:

62 PretrainedConfig = Any

63

64 - AttentionBackendEnum = Any

64 me_quant = LazyLoader(

65 "model_executor", globals(), "vllm.model_executor.layers.quantization"

66 )

vllm/config/multimodal.py

+2 -9

@@ -2,19 +2,15 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Mapping

5 -from typing import TYPE_CHECKING, Any, Literal, TypeAlias

5 +from typing import Any, Literal, TypeAlias

6

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

11 from vllm.config.utils import config

12 from vllm.utils.hashing import safe_hash

13

13 -if TYPE_CHECKING:

14 - from vllm.attention.backends.registry import AttentionBackendEnum

15 -else:

16 - AttentionBackendEnum = Any

17 -

14

15 @dataclass

16 class BaseDummyOptions:

@@ -170,9 +166,6 @@ class MultiModalConfig:

166 def _validate_mm_encoder_attn_backend(

167 cls, value: str | AttentionBackendEnum | None

168 ) -> AttentionBackendEnum | None:

173 - # We need to import the real type here (deferred to avoid circular import).

174 - from vllm.attention.backends.registry import AttentionBackendEnum

175 -

169 if isinstance(value, str) and value.upper() == "XFORMERS":

170 raise ValueError(

171 "Attention backend 'XFORMERS' has been removed (See PR #29262 for "

vllm/distributed/kv_transfer/kv_connector/v1/base.py

+2 -2

@@ -42,12 +42,12 @@ from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional

42

43 import torch

44

45 +from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata

46 from vllm.logger import init_logger

47 from vllm.v1.core.sched.output import SchedulerOutput

48 from vllm.v1.outputs import KVConnectorOutput

49

50 if TYPE_CHECKING:

50 - from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata

51 from vllm.config import VllmConfig

52 from vllm.distributed.kv_events import KVCacheEvent

53 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (

@@ -239,7 +239,7 @@ class KVConnectorBase_V1(ABC):

239 return

240

241 def register_cross_layers_kv_cache(

242 - self, kv_cache: torch.Tensor, attn_backend: type["AttentionBackend"]

242 + self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]

243 ):

244 """

245 Initialize with a single KV cache tensor used by all layers.

vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py

+2 -2

@@ -36,6 +36,7 @@ from typing import TYPE_CHECKING, Any, Optional

36

37 import torch

38

39 +from vllm.attention.backends.abstract import AttentionMetadata

40 from vllm.distributed.kv_transfer.kv_connector.v1 import (

41 KVConnectorBase_V1,

42 KVConnectorRole,

@@ -45,7 +46,6 @@ from vllm.logger import init_logger

46 from vllm.utils.math_utils import cdiv

47

48 if TYPE_CHECKING:

48 - from vllm.attention.backends.abstract import AttentionMetadata

49 from vllm.config import VllmConfig

50 from vllm.forward_context import ForwardContext

51 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

@@ -117,7 +117,7 @@ class DecodeBenchConnector(KVConnectorBase_V1):

117 self,

118 layer_name: str,

119 kv_layer: torch.Tensor,

120 - attn_metadata: "AttentionMetadata",

120 + attn_metadata: AttentionMetadata,

121 **kwargs: Any,

122 ) -> None:

123 # This connector doesn't save KV cache (benchmarking only)

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py

+2 -2

@@ -7,6 +7,7 @@ from lmcache.integration.vllm.vllm_v1_adapter import (

7 LMCacheConnectorV1Impl as LMCacheConnectorLatestImpl,

8 )

9

10 +from vllm.attention.backends.abstract import AttentionMetadata

11 from vllm.config import VllmConfig

12 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

13 KVConnectorBase_V1,

@@ -17,7 +18,6 @@ from vllm.logger import init_logger

18 from vllm.v1.core.sched.output import SchedulerOutput

19

20 if TYPE_CHECKING:

20 - from vllm.attention.backends.abstract import AttentionMetadata

21 from vllm.forward_context import ForwardContext

22 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

23 from vllm.v1.kv_cache_interface import KVCacheConfig

@@ -91,7 +91,7 @@ class LMCacheConnectorV1(KVConnectorBase_V1):

91 self,

92 layer_name: str,

93 kv_layer: torch.Tensor,

94 - attn_metadata: "AttentionMetadata",

94 + attn_metadata: AttentionMetadata,

95 **kwargs: Any,

96 ) -> None:

97 """

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py

+2 -2

@@ -29,6 +29,7 @@ from lmcache.v1.lookup_client.lmcache_async_lookup_client import (

29 from lmcache.v1.offload_server.zmq_server import ZMQOffloadServer

30 from lmcache.v1.plugin.plugin_launcher import PluginLauncher

31

32 +from vllm.attention.backends.abstract import AttentionMetadata

33 from vllm.config import VllmConfig

34 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

35 KVConnectorBase_V1,

@@ -50,7 +51,6 @@ from vllm.v1.core.sched.output import SchedulerOutput

51 from vllm.version import __version__ as VLLM_VERSION

52

53 if TYPE_CHECKING:

53 - from vllm.attention.backends.abstract import AttentionMetadata

54 from vllm.forward_context import ForwardContext

55 from vllm.multimodal.inputs import PlaceholderRange

56 from vllm.v1.core.kv_cache_manager import KVCacheManager

@@ -915,7 +915,7 @@ class LMCacheConnectorV1Impl:

915 self,

916 layer_name: str,

917 kv_layer: torch.Tensor,

918 - attn_metadata: "AttentionMetadata",

918 + attn_metadata: AttentionMetadata,

919 **kwargs,

920 ) -> None:

921 """Start saving the a layer of KV cache from vLLM's paged buffer

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py

+2 -2

@@ -10,6 +10,7 @@ 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

14 from vllm.config import VllmConfig

15 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

16 KVConnectorBase_V1,

@@ -26,7 +27,6 @@ from vllm.v1.outputs import KVConnectorOutput

27 from vllm.v1.utils import ConstantList

28

29 if TYPE_CHECKING:

29 - from vllm.attention.backends.abstract import AttentionMetadata

30 from vllm.config import VllmConfig

31 from vllm.distributed.kv_events import KVCacheEvent

32 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (

@@ -490,7 +490,7 @@ class LMCacheMPConnector(KVConnectorBase_V1):

490 self,

491 layer_name: str,

492 kv_layer: torch.Tensor,

493 - attn_metadata: "AttentionMetadata",

493 + attn_metadata: AttentionMetadata,

494 **kwargs: Any,

495 ) -> None:

496 """

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

+2 -2

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any

7

8 import torch

9

10 +from vllm.attention.backends.abstract import AttentionMetadata

11 from vllm.config import VllmConfig

12 from vllm.config.kv_transfer import KVTransferConfig

13 from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType

@@ -27,7 +28,6 @@ from vllm.v1.core.sched.output import SchedulerOutput

28 from vllm.v1.outputs import KVConnectorOutput

29

30 if TYPE_CHECKING:

30 - from vllm.attention.backends.abstract import AttentionMetadata

31 from vllm.distributed.kv_events import KVCacheEvent

32 from vllm.forward_context import ForwardContext

33 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

@@ -216,7 +216,7 @@ class MultiConnector(KVConnectorBase_V1):

216 self,

217 layer_name: str,

218 kv_layer: torch.Tensor,

219 - attn_metadata: "AttentionMetadata",

219 + attn_metadata: AttentionMetadata,

220 **kwargs,

221 ) -> None:

222 for c in self._connectors:

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+2 -3

@@ -20,7 +20,7 @@ import torch

20 import zmq

21

22 from vllm import envs

23 -from vllm.attention.backends.abstract import AttentionBackend

23 +from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata

24 from vllm.attention.backends.registry import AttentionBackendEnum

25 from vllm.attention.selector import get_attn_backend

26 from vllm.config import VllmConfig

@@ -51,7 +51,6 @@ from vllm.v1.core.sched.output import SchedulerOutput

51 from vllm.v1.worker.block_table import BlockTable

52

53 if TYPE_CHECKING:

54 - from vllm.attention.backends.abstract import AttentionMetadata

54 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

55 from vllm.v1.kv_cache_interface import KVCacheConfig

56 from vllm.v1.request import Request

@@ -308,7 +307,7 @@ class NixlConnector(KVConnectorBase_V1):

307 self,

308 layer_name: str,

309 kv_layer: torch.Tensor,

311 - attn_metadata: "AttentionMetadata",

310 + attn_metadata: AttentionMetadata,

311 **kwargs,

312 ) -> None:

313 """NixlConnector does not save explicitly."""

vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py

+2 -2

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Optional

7 import regex as re

8 import torch

9

10 +from vllm.attention.backends.abstract import AttentionMetadata

11 from vllm.config import VllmConfig

12 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

13 KVConnectorBase_V1,

@@ -22,7 +23,6 @@ from vllm.v1.attention.backends.mla.common import MLACommonMetadata

23 from vllm.v1.core.sched.output import SchedulerOutput

24

25 if TYPE_CHECKING:

25 - from vllm.attention.backends.abstract import AttentionMetadata

26 from vllm.forward_context import ForwardContext

27 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

28 from vllm.v1.kv_cache_interface import KVCacheConfig

@@ -243,7 +243,7 @@ class P2pNcclConnector(KVConnectorBase_V1):

243 self,

244 layer_name: str,

245 kv_layer: torch.Tensor,

246 - attn_metadata: "AttentionMetadata",

246 + attn_metadata: AttentionMetadata,

247 **kwargs: Any,

248 ) -> None:

249 """Start saving the KV cache of the layer from vLLM's paged buffer

vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py

+2 -2

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Optional

7 import safetensors

8 import torch

9

10 +from vllm.attention.backends.abstract import AttentionMetadata

11 from vllm.config import VllmConfig

12 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

13 KVConnectorBase_V1,

@@ -19,7 +20,6 @@ from vllm.v1.attention.backends.mla.common import MLACommonMetadata

20 from vllm.v1.core.sched.output import SchedulerOutput

21

22 if TYPE_CHECKING:

22 - from vllm.attention.backends.abstract import AttentionMetadata

23 from vllm.forward_context import ForwardContext

24 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

25 from vllm.v1.kv_cache_interface import KVCacheConfig

@@ -211,7 +211,7 @@ class SharedStorageConnector(KVConnectorBase_V1):

211 self,

212 layer_name: str,

213 kv_layer: torch.Tensor,

214 - attn_metadata: "AttentionMetadata",

214 + attn_metadata: AttentionMetadata,

215 **kwargs: Any,

216 ) -> None:

217 """Start saving the KV cache of the layer from vLLM's paged buffer

vllm/forward_context.py

+3 -5

@@ -5,19 +5,17 @@ import time

5 from collections import defaultdict

6 from contextlib import contextmanager

7 from dataclasses import dataclass

8 -from typing import TYPE_CHECKING, Any, NamedTuple

8 +from typing import Any, NamedTuple

9

10 import torch

11

12 import vllm.envs as envs

13 +from vllm.attention.backends.abstract import AttentionMetadata

14 from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig

15 from vllm.logger import init_logger

16 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp

17 from vllm.v1.worker.ubatch_utils import UBatchSlices

18

18 -if TYPE_CHECKING:

19 - from vllm.attention.backends.abstract import AttentionMetadata

20 -

19 logger = init_logger(__name__)

20

21 track_batchsize: bool = envs.VLLM_LOG_BATCHSIZE_INTERVAL >= 0

@@ -195,7 +193,7 @@ class ForwardContext:

193 for each microbatch.

194 Set dynamically for each forward pass

195 """

198 - attn_metadata: dict[str, "AttentionMetadata"] | list[dict[str, "AttentionMetadata"]]

196 + attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]]

197 # TODO: remove after making all virtual_engines share the same kv cache

198 virtual_engine: int # set dynamically for each forward pass

199 # set dynamically for each forward pass

vllm/model_executor/layers/attention_layer_base.py

+2 -5

@@ -3,14 +3,11 @@

3 """Base class for attention-like layers."""

4

5 from abc import ABC, abstractmethod

6 -from typing import TYPE_CHECKING

6

7 +from vllm.attention.backends.abstract import AttentionBackend

8 from vllm.config import VllmConfig

9 from vllm.v1.kv_cache_interface import KVCacheSpec

10

11 -if TYPE_CHECKING:

12 - from vllm.attention.backends.abstract import AttentionBackend

13 -

11

12 class AttentionLayerBase(ABC):

13 """

@@ -22,7 +19,7 @@ class AttentionLayerBase(ABC):

19 """

20

21 @abstractmethod

25 - def get_attn_backend(self) -> type["AttentionBackend"]:

22 + def get_attn_backend(self) -> type[AttentionBackend]:

23 """Get the attention backend class for this layer."""

24 pass

25

vllm/model_executor/layers/mamba/abstract.py

+2 -5

@@ -2,18 +2,15 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from abc import abstractmethod

4 from collections.abc import Iterable

5 -from typing import TYPE_CHECKING

5

6 import torch

7

8 +from vllm.attention.backends.abstract import AttentionBackend

9 from vllm.attention.selector import get_mamba_attn_backend

10 from vllm.config import VllmConfig

11 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

12 from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec

13

14 -if TYPE_CHECKING:

15 - from vllm.attention.backends.abstract import AttentionBackend

16 -

14

15 class MambaBase(AttentionLayerBase):

16 """

@@ -66,6 +63,6 @@ class MambaBase(AttentionLayerBase):

63 ),

64 )

65

69 - def get_attn_backend(self) -> type["AttentionBackend"]:

66 + def get_attn_backend(self) -> type[AttentionBackend]:

67 """Get the attention backend class for this Mamba layer."""

68 return get_mamba_attn_backend(self.mamba_type)

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py

+1 -2

@@ -18,6 +18,7 @@ from compressed_tensors.quantization import (

18 from compressed_tensors.transform import TransformConfig

19

20 import vllm.envs as envs

21 +from vllm.attention.layer import Attention

22 from vllm.logger import init_logger

23 from vllm.model_executor.layers.fused_moe import FusedMoE

24 from vllm.model_executor.layers.linear import (

@@ -131,8 +132,6 @@ class CompressedTensorsConfig(QuantizationConfig):

132 layer: torch.nn.Module,

133 prefix: str,

134 ) -> Optional["QuantizeMethodBase"]:

134 - from vllm.attention.layer import Attention # Avoid circular import

135 -

135 if isinstance(layer, LinearBase):

136 # collect schemes

137 quant_scheme = self.get_scheme(layer=layer, layer_name=prefix)

vllm/model_executor/layers/quantization/fp8.py

+1 -3

@@ -14,6 +14,7 @@ import vllm.envs as envs

14 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

15 from vllm import _custom_ops as ops

16 from vllm._aiter_ops import rocm_aiter_ops

17 +from vllm.attention.layer import Attention

18 from vllm.distributed import get_tensor_model_parallel_world_size

19 from vllm.logger import init_logger

20 from vllm.model_executor.layers.batch_invariant import (

@@ -277,7 +278,6 @@ class Fp8Config(QuantizationConfig):

278 def get_xpu_quant_method(

279 self, layer: torch.nn.Module, prefix: str

280 ) -> Optional["QuantizeMethodBase"]:

280 - from vllm.attention.layer import Attention

281 from vllm.model_executor.layers.quantization.ipex_quant import (

282 XPUFp8LinearMethod,

283 XPUFp8MoEMethod,

@@ -307,8 +307,6 @@ class Fp8Config(QuantizationConfig):

307 def get_quant_method(

308 self, layer: torch.nn.Module, prefix: str

309 ) -> Optional["QuantizeMethodBase"]:

310 - from vllm.attention.layer import Attention # Avoid circular import

311 -

310 if current_platform.is_xpu():

311 return self.get_xpu_quant_method(layer, prefix)

312 if isinstance(layer, LinearBase):

vllm/model_executor/layers/quantization/modelopt.py

+1 -2

@@ -12,6 +12,7 @@ from torch.nn.parameter import Parameter

12 import vllm.envs as envs

13 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

14 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant

15 +from vllm.attention.layer import Attention

16 from vllm.logger import init_logger

17 from vllm.model_executor.layers.fused_moe.config import (

18 FusedMoEQuantConfig,

@@ -149,8 +150,6 @@ class ModelOptQuantConfigBase(QuantizationConfig):

150 def get_quant_method(

151 self, layer: torch.nn.Module, prefix: str

152 ) -> Optional["QuantizeMethodBase"]:

152 - from vllm.attention.layer import Attention # Avoid circular import

153 -

153 # handle kv-cache first so we can focus only on weight quantization thereafter

154 if isinstance(layer, Attention):

155 return self.KVCacheMethodCls(self)

vllm/model_executor/layers/quantization/mxfp4.py

+1 -2

@@ -8,6 +8,7 @@ import torch

8 from torch.nn.parameter import Parameter

9

10 from vllm import envs

11 +from vllm.attention.layer import Attention

12 from vllm.config import get_current_vllm_config

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.fused_moe import (

@@ -184,8 +185,6 @@ class Mxfp4Config(QuantizationConfig):

185 def get_quant_method(

186 self, layer: torch.nn.Module, prefix: str

187 ) -> Optional["QuantizeMethodBase"]:

187 - from vllm.attention.layer import Attention # Avoid circular import

188 -

188 if isinstance(layer, LinearBase):

189 if self.ignored_layers and is_layer_skipped(

190 prefix=prefix,

vllm/model_executor/layers/quantization/petit.py

+1 -2

@@ -8,6 +8,7 @@ import regex as re

8 import torch

9 from torch.nn.parameter import Parameter

10

11 +from vllm.attention.layer import Attention

12 from vllm.logger import init_logger

13 from vllm.model_executor.layers.linear import (

14 LinearBase,

@@ -159,8 +160,6 @@ class PetitNvFp4Config(QuantizationConfig):

160 def get_quant_method(

161 self, layer: torch.nn.Module, prefix: str

162 ) -> Optional["QuantizeMethodBase"]:

162 - from vllm.attention.layer import Attention # Avoid circular import

163 -

163 exclude = self.require_exclude_modules()

164

165 if isinstance(layer, LinearBase):

vllm/model_executor/layers/quantization/ptpc_fp8.py

+1 -2

@@ -7,6 +7,7 @@ import torch

7 from torch.nn.parameter import Parameter

8

9 from vllm import _custom_ops as ops

10 +from vllm.attention.layer import Attention

11 from vllm.logger import init_logger

12 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod

13 from vllm.model_executor.layers.quantization import QuantizationMethods

@@ -65,8 +66,6 @@ class PTPCFp8Config(Fp8Config):

66 def get_quant_method(

67 self, layer: torch.nn.Module, prefix: str

68 ) -> Optional["QuantizeMethodBase"]:

68 - from vllm.attention.layer import Attention # Avoid circular import

69 -

69 if isinstance(layer, LinearBase):

70 if is_layer_skipped(prefix, self.ignored_layers):

71 return UnquantizedLinearMethod()

vllm/model_executor/layers/quantization/quark/quark.py

+1 -2

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Optional, cast

6

7 import torch

8

9 +from vllm.attention.layer import Attention

10 from vllm.logger import init_logger

11 from vllm.model_executor.layers.fused_moe import FusedMoE

12 from vllm.model_executor.layers.linear import (

@@ -102,8 +103,6 @@ class QuarkConfig(QuantizationConfig):

103 def get_quant_method(

104 self, layer: torch.nn.Module, prefix: str

105 ) -> Optional["QuantizeMethodBase"]:

105 - from vllm.attention.layer import Attention # Avoid circular import

106 -

106 # Check if the layer is skipped for quantization.

107 exclude_layers = cast(list[str], self.quant_config.get("exclude"))

108 if should_ignore_layer(

vllm/platforms/cpu.py

+1 -4

@@ -14,6 +14,7 @@ import regex as re

14 import torch

15

16 from vllm import envs

17 +from vllm.attention.backends.registry import AttentionBackendEnum

18 from vllm.logger import init_logger

19

20 from .interface import CpuArchEnum, Platform, PlatformEnum

@@ -21,10 +22,8 @@ from .interface import CpuArchEnum, Platform, PlatformEnum

22 logger = init_logger(__name__)

23

24 if TYPE_CHECKING:

24 - from vllm.attention.backends.registry import AttentionBackendEnum

25 from vllm.config import VllmConfig

26 else:

27 - AttentionBackendEnum = None

27 VllmConfig = None

28

29

@@ -135,8 +134,6 @@ class CpuPlatform(Platform):

134 use_sparse: bool,

135 attn_type: str | None = None,

136 ) -> str:

138 - from vllm.attention.backends.registry import AttentionBackendEnum

139 -

137 if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN:

138 logger.info("Cannot use %s backend on CPU.", selected_backend)

139 if use_mla:

vllm/platforms/cuda.py

+2 -8

@@ -15,6 +15,8 @@ from typing_extensions import ParamSpec

15 # import custom ops, trigger op registration

16 import vllm._C # noqa

17 import vllm.envs as envs

18 +from vllm.attention.backends.abstract import AttentionType

19 +from vllm.attention.backends.registry import AttentionBackendEnum

20 from vllm.logger import init_logger

21 from vllm.utils.import_utils import import_pynvml

22 from vllm.utils.torch_utils import cuda_device_count_stateless

@@ -22,11 +24,9 @@ from vllm.utils.torch_utils import cuda_device_count_stateless

24 from .interface import DeviceCapability, Platform, PlatformEnum

25

26 if TYPE_CHECKING:

25 - from vllm.attention.backends.registry import AttentionBackendEnum

27 from vllm.config import VllmConfig

28 from vllm.config.cache import CacheDType

29 else:

29 - AttentionBackendEnum = None

30 VllmConfig = None

31 CacheDType = None

32

@@ -48,8 +48,6 @@ def _get_backend_priorities(

48 device_capability: DeviceCapability,

49 ) -> list[AttentionBackendEnum]:

50 """Get backend priorities with lazy import to avoid circular dependency."""

51 - from vllm.attention.backends.registry import AttentionBackendEnum

52 -

51 if use_mla:

52 if device_capability.major == 10:

53 return [

@@ -265,8 +263,6 @@ class CudaPlatformBase(Platform):

263 def get_vit_attn_backend(

264 cls, head_size: int, dtype: torch.dtype

265 ) -> "AttentionBackendEnum":

268 - from vllm.attention.backends.registry import AttentionBackendEnum

269 -

266 # Try FlashAttention first

267 try:

268 backend_class = AttentionBackendEnum.FLASH_ATTN.get_class()

@@ -335,8 +331,6 @@ class CudaPlatformBase(Platform):

331 use_sparse: bool,

332 attn_type: str | None = None,

333 ) -> str:

338 - from vllm.attention.backends.abstract import AttentionType

339 -

334 if attn_type is None:

335 attn_type = AttentionType.DECODER

336

vllm/platforms/interface.py

+1 -4

@@ -12,12 +12,12 @@ from typing import TYPE_CHECKING, Any, NamedTuple

12 import numpy as np

13 import torch

14

15 +from vllm.attention.backends.registry import AttentionBackendEnum

16 from vllm.logger import init_logger

17

18 if TYPE_CHECKING:

19 from torch.distributed import PrefixStore, ProcessGroup

20

20 - from vllm.attention.backends.registry import AttentionBackendEnum

21 from vllm.config import VllmConfig

22 from vllm.config.cache import CacheDType

23 from vllm.inputs import ProcessorInputs, PromptType

@@ -226,9 +226,6 @@ class Platform:

226 def get_vit_attn_backend(

227 cls, head_size: int, dtype: torch.dtype

228 ) -> "AttentionBackendEnum":

229 - # Import AttentionBackendEnum here to avoid circular import.

230 - from vllm.attention.backends.registry import AttentionBackendEnum

231 -

229 return AttentionBackendEnum.TORCH_SDPA

230

231 @classmethod

vllm/platforms/rocm.py

+1 -5

@@ -8,16 +8,14 @@ from typing import TYPE_CHECKING

8 import torch

9

10 import vllm.envs as envs

11 +from vllm.attention.backends.registry import AttentionBackendEnum

12 from vllm.logger import init_logger

13 from vllm.utils.torch_utils import cuda_device_count_stateless

14

15 from .interface import DeviceCapability, Platform, PlatformEnum

16

17 if TYPE_CHECKING:

17 - from vllm.attention.backends.registry import AttentionBackendEnum

18 from vllm.config import VllmConfig

19 -else:

20 - AttentionBackendEnum = None

19

20 logger = init_logger(__name__)

21

@@ -196,7 +194,6 @@ class RocmPlatform(Platform):

194 from importlib.util import find_spec

195

196 from vllm._aiter_ops import rocm_aiter_ops

199 - from vllm.attention.backends.registry import AttentionBackendEnum

197

198 if rocm_aiter_ops.is_mha_enabled():

199 # Note: AITER FA is only supported for Qwen-VL models.

@@ -222,7 +219,6 @@ class RocmPlatform(Platform):

219 attn_type: str | None = None,

220 ) -> str:

221 from vllm._aiter_ops import rocm_aiter_ops

225 - from vllm.attention.backends.registry import AttentionBackendEnum

222

223 if use_sparse:

224 if kv_cache_dtype.startswith("fp8"):

vllm/platforms/tpu.py

+1 -4

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, cast

7 import torch

8 from tpu_info import device

9

10 +from vllm.attention.backends.registry import AttentionBackendEnum

11 from vllm.inputs import ProcessorInputs, PromptType

12 from vllm.logger import init_logger

13

@@ -15,7 +16,6 @@ from .interface import Platform, PlatformEnum

16 if TYPE_CHECKING:

17 from typing import TypeAlias

18

18 - from vllm.attention.backends.registry import AttentionBackendEnum

19 from vllm.config import VllmConfig

20 from vllm.config.cache import BlockSize

21 from vllm.pooling_params import PoolingParams

@@ -26,7 +26,6 @@ else:

26 BlockSize = None

27 VllmConfig = None

28 PoolingParams = None

29 - AttentionBackendEnum = None

29 ParamsType = None

30

31 logger = init_logger(__name__)

@@ -67,8 +66,6 @@ class TpuPlatform(Platform):

66 use_sparse,

67 attn_type: str | None = None,

68 ) -> str:

70 - from vllm.attention.backends.registry import AttentionBackendEnum

71 -

69 if use_sparse:

70 raise NotImplementedError("Sparse Attention is not supported on TPU.")

71 if selected_backend != AttentionBackendEnum.PALLAS:

vllm/platforms/xpu.py

+1 -6

@@ -8,16 +8,15 @@ from typing import TYPE_CHECKING

8 import torch

9

10 import vllm.envs as envs

11 +from vllm.attention.backends.registry import AttentionBackendEnum

12 from vllm.logger import init_logger

13

14 from .interface import DeviceCapability, Platform, PlatformEnum

15

16 if TYPE_CHECKING:

16 - from vllm.attention.backends.registry import AttentionBackendEnum

17 from vllm.config import VllmConfig

18 else:

19 VllmConfig = None

20 - AttentionBackendEnum = None

20

21 logger = init_logger(__name__)

22

@@ -60,8 +59,6 @@ class XPUPlatform(Platform):

59 "only NHD layout is supported by XPU attention kernels."

60 )

61

63 - from vllm.attention.backends.registry import AttentionBackendEnum

64 -

62 if use_sparse:

63 raise NotImplementedError("Sparse Attention is not supported on XPU.")

64 if selected_backend == AttentionBackendEnum.TRITON_ATTN:

@@ -116,8 +113,6 @@ class XPUPlatform(Platform):

113 def get_vit_attn_backend(

114 cls, head_size: int, dtype: torch.dtype

115 ) -> "AttentionBackendEnum":

119 - from vllm.attention.backends.registry import AttentionBackendEnum

120 -

116 return AttentionBackendEnum.FLASH_ATTN

117

118 @classmethod

vllm/v1/attention/backends/cpu_attn.py

-2

@@ -51,8 +51,6 @@ class CPUAttentionBackend(AttentionBackend):

51 @classmethod

52 def supports_attn_type(cls, attn_type: str) -> bool:

53 """CPU attention supports decoder and encoder-only attention."""

54 - from vllm.attention.backends.abstract import AttentionType

55 -

54 return attn_type in (

55 AttentionType.DECODER,

56 AttentionType.ENCODER,

vllm/v1/attention/backends/flash_attn.py

-2

@@ -84,8 +84,6 @@ class FlashAttentionBackend(AttentionBackend):

84 @classmethod

85 def supports_attn_type(cls, attn_type: str) -> bool:

86 """FlashAttention supports all attention types."""

87 - from vllm.attention.backends.abstract import AttentionType

88 -

87 return attn_type in (

88 AttentionType.DECODER,

89 AttentionType.ENCODER,

vllm/v1/attention/backends/flex_attention.py

-2

@@ -87,8 +87,6 @@ class FlexAttentionBackend(AttentionBackend):

87 @classmethod

88 def supports_attn_type(cls, attn_type: str) -> bool:

89 """FlexAttention supports both decoder and encoder-only attention."""

90 - from vllm.attention.backends.abstract import AttentionType

91 -

90 return attn_type in (AttentionType.DECODER, AttentionType.ENCODER_ONLY)

91

92 @staticmethod

vllm/v1/attention/backends/utils.py

+5 -2

@@ -24,12 +24,15 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config

24 from vllm.utils.math_utils import cdiv

25

26 if TYPE_CHECKING:

27 - from vllm.attention.backends.abstract import AttentionImpl

27 from vllm.v1.core.sched.output import SchedulerOutput

28 from vllm.v1.worker.gpu_input_batch import InputBatch

29

30 import vllm.envs as envs

32 -from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata

31 +from vllm.attention.backends.abstract import (

32 + AttentionBackend,

33 + AttentionImpl,

34 + AttentionMetadata,

35 +)

36 from vllm.distributed.kv_transfer.kv_connector.utils import (

37 get_kv_connector_cache_layout,

38 )

vllm/v1/kv_offload/spec.py

+2 -2

@@ -6,12 +6,12 @@ from typing import TYPE_CHECKING

6

7 import torch

8

9 +from vllm.attention.backends.abstract import AttentionBackend

10 from vllm.logger import init_logger

11 from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager

12 from vllm.v1.kv_offload.worker.worker import OffloadingHandler

13

14 if TYPE_CHECKING:

14 - from vllm.attention.backends.abstract import AttentionBackend

15 from vllm.config import VllmConfig

16

17 logger = init_logger(__name__)

@@ -51,7 +51,7 @@ class OffloadingSpec(ABC):

51 def get_handlers(

52 self,

53 kv_caches: dict[str, torch.Tensor],

54 - attn_backends: dict[str, type["AttentionBackend"]],

54 + attn_backends: dict[str, type[AttentionBackend]],

55 ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:

56 """

57 Get offloading handlers along with their respective src and dst types.

vllm/v1/spec_decode/eagle.py

+1 -2

@@ -8,6 +8,7 @@ import numpy as np

8 import torch

9 import torch.nn as nn

10

11 +from vllm.attention.backends.registry import AttentionBackendEnum

12 from vllm.config import (

13 CompilationMode,

14 CUDAGraphMode,

@@ -157,8 +158,6 @@ class EagleProposer:

158 )

159

160 # Determine allowed attention backends once during initialization.

160 - from vllm.attention.backends.registry import AttentionBackendEnum

161 -

161 self.allowed_attn_types: tuple | None = None

162 if current_platform.is_rocm():

163 rocm_types = [TritonAttentionMetadata, FlashAttentionMetadata]

vllm/v1/worker/utils.py

+2 -5

@@ -2,11 +2,11 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections import defaultdict

4 from dataclasses import dataclass, field

5 -from typing import TYPE_CHECKING

5

6 import torch

7

8 from vllm.attention.backends.abstract import AttentionBackend

9 +from vllm.attention.layer import Attention

10 from vllm.config import ModelConfig, SchedulerConfig, VllmConfig

11 from vllm.model_executor.models.interfaces import MultiModalEmbeddings

12 from vllm.model_executor.models.utils import extract_layer_index

@@ -17,9 +17,6 @@ from vllm.v1.attention.backends.utils import AttentionMetadataBuilder

17 from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget

18 from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec

19

20 -if TYPE_CHECKING:

21 - from vllm.attention.layer import Attention

22 -

20

21 class MultiModalBudget:

22 """Helper class to calculate budget information for multi-modal models."""

@@ -278,7 +275,7 @@ def add_kv_sharing_layers_to_kv_cache_groups(

275

276 def bind_kv_cache(

277 kv_caches: dict[str, torch.Tensor],

281 - forward_context: dict[str, "Attention"],

278 + forward_context: dict[str, Attention],

279 runner_kv_caches: list[torch.Tensor],

280 num_attn_module: int = 1,

281 ) -> None:

[Attention] Update attention imports (#29540)