A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 09:45 UTC
[Misc] Remove redundant attention var constants (#29650)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Nov 28, 2025 at 20:35 UTC
33b06a6f24be632e99c19cadb5004a27cf9605a0
7 files changed
+19
-63
tests/kernels/attention/test_attention_selector.py
+9
-10
@@ -11,7 +11,6 @@ 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.utils import STR_BACKEND_ENV_VAR, STR_FLASH_ATTN_VAL, STR_INVALID_VAL
14
15
16
@pytest.fixture(autouse=True)
@@ -83,7 +82,7 @@ def test_env(
82
):
83
"""Test attention backend selection with valid device-backend pairs."""
84
with monkeypatch.context() as m:
86
- m.setenv(STR_BACKEND_ENV_VAR, name)
85
+ m.setenv("VLLM_ATTENTION_BACKEND", name)
86
m.setenv("VLLM_MLA_DISABLE", "1" if use_mla else "0")
87
88
if device == "cpu":
@@ -237,27 +236,27 @@ def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
236
)
237
238
with monkeypatch.context() as m:
240
- m.setenv(STR_BACKEND_ENV_VAR, STR_FLASH_ATTN_VAL)
239
+ m.setenv("VLLM_ATTENTION_BACKEND", "FLASH_ATTN")
240
241
# Unsupported CUDA arch
242
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _=None: (7, 5))
243
backend = get_attn_backend(16, torch.float16, None, 16)
245
- assert backend.get_name() != STR_FLASH_ATTN_VAL
244
+ assert backend.get_name() != "FLASH_ATTN"
245
246
# Reset the monkeypatch for subsequent tests
247
monkeypatch.undo()
248
249
# Unsupported data type
250
backend = get_attn_backend(16, torch.float8_e4m3fn, None, 16)
252
- assert backend.get_name() != STR_FLASH_ATTN_VAL
251
+ assert backend.get_name() != "FLASH_ATTN"
252
253
# Unsupported kv cache data type
254
backend = get_attn_backend(16, torch.float16, "fp8", 16)
256
- assert backend.get_name() != STR_FLASH_ATTN_VAL
255
+ assert backend.get_name() != "FLASH_ATTN"
256
257
# Unsupported block size
258
backend = get_attn_backend(16, torch.float16, None, 8)
260
- assert backend.get_name() != STR_FLASH_ATTN_VAL
259
+ assert backend.get_name() != "FLASH_ATTN"
260
261
# flash-attn is not installed
262
import sys
@@ -265,7 +264,7 @@ def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
264
original_module = sys.modules.get("vllm_flash_attn")
265
monkeypatch.setitem(sys.modules, "vllm_flash_attn", None)
266
backend = get_attn_backend(16, torch.float16, None, 16)
268
- assert backend.get_name() != STR_FLASH_ATTN_VAL
267
+ assert backend.get_name() != "FLASH_ATTN"
268
269
# Restore the original module if it existed
270
if original_module is not None:
@@ -275,7 +274,7 @@ def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
274
275
# Unsupported head size
276
backend = get_attn_backend(17, torch.float16, None, 16)
278
- assert backend.get_name() != STR_FLASH_ATTN_VAL
277
+ assert backend.get_name() != "FLASH_ATTN"
278
279
280
def test_invalid_env(monkeypatch: pytest.MonkeyPatch):
@@ -284,7 +283,7 @@ def test_invalid_env(monkeypatch: pytest.MonkeyPatch):
283
monkeypatch.context() as m,
284
patch("vllm.platforms.current_platform", CudaPlatform()),
285
):
287
- m.setenv(STR_BACKEND_ENV_VAR, STR_INVALID_VAL)
286
+ m.setenv("VLLM_ATTENTION_BACKEND", "INVALID")
287
288
# Should raise ValueError for invalid backend
289
with pytest.raises(ValueError) as exc_info:
tests/kernels/attention/test_rocm_attention_selector.py
+5
-6
@@ -6,7 +6,6 @@ import torch
6
7
from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
8
from vllm.platforms.rocm import RocmPlatform
9
-from vllm.utils import STR_BACKEND_ENV_VAR
9
10
11
@pytest.fixture(autouse=True)
@@ -18,7 +17,7 @@ def clear_cache():
17
@pytest.mark.skip(reason="Skipped for now. Should be revisited.")
18
def test_selector(monkeypatch: pytest.MonkeyPatch):
19
with monkeypatch.context() as m:
21
- m.setenv(STR_BACKEND_ENV_VAR, "ROCM_ATTN")
20
+ m.setenv("VLLM_ATTENTION_BACKEND", "ROCM_ATTN")
21
22
# Set the current platform to ROCm using monkeypatch
23
monkeypatch.setattr("vllm.attention.selector.current_platform", RocmPlatform())
@@ -30,19 +29,19 @@ def test_selector(monkeypatch: pytest.MonkeyPatch):
29
# MLA test for deepseek related
30
31
# change the attention backend to triton MLA
33
- m.setenv(STR_BACKEND_ENV_VAR, "TRITON_MLA")
32
+ m.setenv("VLLM_ATTENTION_BACKEND", "TRITON_MLA")
33
backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False, use_mla=True)
34
assert backend.get_name() == "TRITON_MLA"
35
36
# If attention backend is None
37
# If use_mla is true
38
# The selected backend is triton MLA
40
- m.setenv(STR_BACKEND_ENV_VAR, None)
39
+ m.setenv("VLLM_ATTENTION_BACKEND", "")
40
backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False, use_mla=True)
41
assert backend.get_name() == "TRITON_MLA"
42
43
# change the attention backend to AITER MLA
45
- m.setenv(STR_BACKEND_ENV_VAR, "ROCM_AITER_MLA")
44
+ m.setenv("VLLM_ATTENTION_BACKEND", "ROCM_AITER_MLA")
45
backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False, use_mla=True)
46
assert backend.get_name() == "ROCM_AITER_MLA"
47
@@ -50,7 +49,7 @@ def test_selector(monkeypatch: pytest.MonkeyPatch):
49
# If use_mla is true
50
# If VLLM_ROCM_USE_AITER is enabled
51
# The selected backend is ROCM_AITER_MLA
53
- m.setenv(STR_BACKEND_ENV_VAR, None)
52
+ m.setenv("VLLM_ATTENTION_BACKEND", "")
53
m.setenv("VLLM_ROCM_USE_AITER", "1")
54
backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False, use_mla=True)
55
assert backend.get_name() == "ROCM_AITER_MLA"
tests/kernels/utils.py
-20
@@ -9,7 +9,6 @@ from numbers import Number
9
from typing import Any, NamedTuple
10
from unittest.mock import patch
11
12
-import pytest
12
import torch
13
from torch._prims_common import TensorLikeType
14
@@ -17,9 +16,6 @@ from tests.kernels.quant_utils import native_w8a8_block_matmul
16
from vllm.attention.backends.abstract import AttentionType
17
from vllm.model_executor.layers.activation import SiluAndMul
18
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
20
-from vllm.utils import (
21
- STR_BACKEND_ENV_VAR,
22
-)
19
from vllm.utils.torch_utils import make_tensor_with_pad
20
21
# For now, disable "test_aot_dispatch_dynamic" since there are some
@@ -217,22 +213,6 @@ def make_causal_mask(
213
return mask
214
215
220
-def override_backend_env_variable(
221
- mpatch: pytest.MonkeyPatch, backend_name: str
222
-) -> None:
223
- """
224
- Override the environment variable indicating the vLLM backend temporarily,
225
- using pytest monkeypatch to ensure that the env vars get
226
- reset once the test context exits.
227
-
228
- Arguments:
229
-
230
- * mpatch: pytest monkeypatch instance
231
- * backend_name: attention backend name to force
232
- """
233
- mpatch.setenv(STR_BACKEND_ENV_VAR, backend_name)
234
-
235
-
216
def ref_masked_attention(
217
query: torch.Tensor,
218
key: torch.Tensor,
tests/models/quantization/test_fp8.py
+1
-2
@@ -11,7 +11,6 @@ import pytest
11
from tests.quantization.utils import is_quant_method_supported
12
from vllm.attention.utils.fa_utils import flash_attn_supports_fp8
13
from vllm.platforms import current_platform
14
-from vllm.utils import STR_BACKEND_ENV_VAR
14
from ..utils import check_logprobs_close
15
16
@@ -76,7 +75,7 @@ def test_models(
75
76
with monkeypatch.context() as m:
77
m.setenv("TOKENIZERS_PARALLELISM", "true")
79
- m.setenv(STR_BACKEND_ENV_VAR, backend)
78
+ m.setenv("VLLM_ATTENTION_BACKEND", backend)
79
80
MAX_MODEL_LEN = 1024
81
NUM_LOG_PROBS = 8
vllm/attention/selector.py
+4
-5
@@ -19,7 +19,6 @@ from vllm.attention.backends.registry import (
19
)
20
from vllm.config.cache import CacheDType
21
from vllm.logger import init_logger
22
-from vllm.utils import STR_BACKEND_ENV_VAR
22
from vllm.utils.import_utils import resolve_obj_by_qualname
23
24
logger = init_logger(__name__)
@@ -35,7 +34,7 @@ def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
34
* AttentionBackendEnum value if an override is specified
35
* None otherwise
36
"""
38
- backend_name = os.environ.get(STR_BACKEND_ENV_VAR)
37
+ backend_name = os.environ.get("VLLM_ATTENTION_BACKEND")
38
if backend_name is None:
39
return None
40
if backend_name == "XFORMERS":
@@ -139,10 +138,10 @@ def _cached_get_attn_backend(
138
if backend_by_env_var.endswith("_VLLM_V1"):
139
logger.warning(
140
"The suffix '_VLLM_V1' in the environment variable "
142
- "%s is no longer necessary as V0 backends have been "
143
- "deprecated. Please remove this suffix from your "
141
+ "VLLM_ATTENTION_BACKEND is no longer necessary as "
142
+ "V0 backends have been deprecated. "
143
+ "Please remove this suffix from your "
144
"environment variable setting.",
145
- STR_BACKEND_ENV_VAR,
145
)
146
backend_by_env_var = backend_by_env_var.removesuffix("_VLLM_V1")
147
try:
vllm/model_executor/models/deepseek_eagle.py
-3
@@ -23,12 +23,9 @@ from vllm.model_executor.models.deepseek_v2 import (
23
DeepseekV2DecoderLayer,
24
DeepseekV3ForCausalLM,
25
)
26
-from vllm.utils import init_logger
26
27
from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight
28
30
-logger = init_logger(__name__)
31
-
29
30
@support_torch_compile
31
class DeepseekV2Model(nn.Module):
vllm/utils/__init__.py
-17
@@ -7,8 +7,6 @@ from typing import Any
7
8
import torch
9
10
-from vllm.logger import init_logger
11
-
10
_DEPRECATED_MAPPINGS = {
11
"cprofile": "profiling",
12
"cprofile_context": "profiling",
@@ -37,21 +35,6 @@ def __dir__() -> list[str]:
35
return sorted(list(globals().keys()) + list(_DEPRECATED_MAPPINGS.keys()))
36
37
40
-logger = init_logger(__name__)
41
-
42
-# Constants related to forcing the attention backend selection
43
-
44
-# String name of register which may be set in order to
45
-# force auto-selection of attention backend by Attention
46
-# wrapper
47
-STR_BACKEND_ENV_VAR: str = "VLLM_ATTENTION_BACKEND"
48
-
49
-# Possible string values of STR_BACKEND_ENV_VAR
50
-# register, corresponding to possible backends
51
-STR_FLASHINFER_ATTN_VAL: str = "FLASHINFER"
52
-STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
53
-STR_INVALID_VAL: str = "INVALID"
54
-
38
MASK_64_BITS = (1 << 64) - 1
39
40