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 9, 11:15 UTC
[ez] Remove checks for torch version <= 2.8 (#33209)
Signed-off-by: angelayi <yiangela7@gmail.com>
Angela Yi committed
Jan 28, 2026 at 13:03 UTC
4197168ea5877de51dc8aede9388c9be9ecc76ac
11 files changed
+30
-139
vllm/_aiter_ops.py
+1
-8
@@ -8,7 +8,7 @@ from torch._ops import OpOverload
8
9
import vllm.envs as envs
10
from vllm.platforms import current_platform
11
-from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
11
+from vllm.utils.torch_utils import direct_register_custom_op
12
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
13
rocm_aiter_sparse_attn_indexer,
14
rocm_aiter_sparse_attn_indexer_fake,
@@ -1015,12 +1015,6 @@ class rocm_aiter_ops:
1015
def register_ops_once() -> None:
1016
global _OPS_REGISTERED
1017
if not _OPS_REGISTERED:
1018
- tags = (
1019
- tuple()
1020
- if is_torch_equal_or_newer("2.7.0")
1021
- else (torch.Tag.needs_fixed_stride_order,)
1022
- )
1023
-
1018
# register all the custom ops here
1019
direct_register_custom_op(
1020
op_name="rocm_aiter_asm_moe_tkw1",
@@ -1075,7 +1069,6 @@ class rocm_aiter_ops:
1069
op_func=_rocm_aiter_mla_decode_fwd_impl,
1070
mutates_args=["o"],
1071
fake_impl=_rocm_aiter_mla_decode_fwd_fake,
1078
- tags=tags,
1072
)
1073
1074
direct_register_custom_op(
vllm/compilation/backends.py
+2
-5
@@ -33,7 +33,6 @@ from vllm.logger import init_logger
33
from vllm.logging_utils import lazy
34
from vllm.platforms import current_platform
35
from vllm.utils.import_utils import resolve_obj_by_qualname
36
-from vllm.utils.torch_utils import is_torch_equal_or_newer
36
37
from .compiler_interface import (
38
CompilerInterface,
@@ -94,10 +93,8 @@ def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface:
93
if compilation_config.backend == "inductor":
94
# Use standalone compile only if requested, version is new enough,
95
# and the symbol actually exists in this PyTorch build.
97
- if (
98
- envs.VLLM_USE_STANDALONE_COMPILE
99
- and is_torch_equal_or_newer("2.8.0.dev")
100
- and hasattr(torch._inductor, "standalone_compile")
96
+ if envs.VLLM_USE_STANDALONE_COMPILE and hasattr(
97
+ torch._inductor, "standalone_compile"
98
):
99
logger.debug("Using InductorStandaloneAdaptor")
100
return InductorStandaloneAdaptor(
vllm/compilation/compiler_interface.py
+13
-14
@@ -501,20 +501,19 @@ class InductorAdaptor(CompilerInterface):
501
# get hit.
502
# TODO(zou3519): we're going to replace this all with
503
# standalone_compile sometime.
504
- if is_torch_equal_or_newer("2.6"):
505
- stack.enter_context(
506
- torch._inductor.config.patch(fx_graph_remote_cache=False)
507
- )
508
- # InductorAdaptor (unfortunately) requires AOTAutogradCache
509
- # to be turned off to run. It will fail to acquire the hash_str
510
- # and error if not.
511
- # StandaloneInductorAdaptor (PyTorch 2.8+) fixes this problem.
512
- stack.enter_context(
513
- torch._functorch.config.patch(enable_autograd_cache=False)
514
- )
515
- stack.enter_context(
516
- torch._functorch.config.patch(enable_remote_autograd_cache=False)
517
- )
504
+ stack.enter_context(
505
+ torch._inductor.config.patch(fx_graph_remote_cache=False)
506
+ )
507
+ # InductorAdaptor (unfortunately) requires AOTAutogradCache
508
+ # to be turned off to run. It will fail to acquire the hash_str
509
+ # and error if not.
510
+ # StandaloneInductorAdaptor (PyTorch 2.8+) fixes this problem.
511
+ stack.enter_context(
512
+ torch._functorch.config.patch(enable_autograd_cache=False)
513
+ )
514
+ stack.enter_context(
515
+ torch._functorch.config.patch(enable_remote_autograd_cache=False)
516
+ )
517
518
compiled_graph = compile_fx(
519
graph,
vllm/compilation/decorators.py
+1
-42
@@ -7,12 +7,11 @@ import inspect
7
import os
8
import sys
9
from collections.abc import Callable, Generator
10
-from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload
10
+from typing import TYPE_CHECKING, Any, TypeVar, overload
11
from unittest.mock import patch
12
13
import torch
14
import torch.nn as nn
15
-from packaging import version
15
from torch._dynamo.symbolic_convert import InliningInstructionTranslator
16
17
import vllm.envs as envs
@@ -540,7 +539,6 @@ def _support_torch_compile(
539
torch._dynamo.config.patch(**dynamo_config_patches),
540
maybe_use_cudagraph_partition_wrapper(self.vllm_config),
541
torch.fx.experimental._config.patch(**fx_config_patches),
543
- _torch27_patch_tensor_subclasses(),
542
torch._inductor.config.patch(**inductor_config_patches),
543
):
544
use_aot_compile = envs.VLLM_USE_AOT_COMPILE
@@ -647,42 +645,3 @@ def maybe_use_cudagraph_partition_wrapper(
645
and compilation_config.use_inductor_graph_partition
646
):
647
torch._inductor.utils.set_customized_partition_wrappers(None)
650
-
651
-
652
-@contextlib.contextmanager
653
-def _torch27_patch_tensor_subclasses() -> Generator[None, None, None]:
654
- """
655
- Add support for using tensor subclasses (ie `BasevLLMParameter`, ect) when
656
- using torch 2.7.0. This enables using weight_loader_v2 and the use of
657
- `BasevLLMParameters` without having to replace them with regular tensors
658
- before `torch.compile`-time.
659
- """
660
- from vllm.model_executor.parameter import (
661
- BasevLLMParameter,
662
- ModelWeightParameter,
663
- RowvLLMParameter,
664
- _ColumnvLLMParameter,
665
- )
666
-
667
- def return_false(*args: Any, **kwargs: Any) -> Literal[False]:
668
- return False
669
-
670
- if version.parse("2.7") <= version.parse(torch.__version__) < version.parse("2.8"):
671
- yield
672
- return
673
-
674
- with (
675
- torch._dynamo.config.patch(
676
- "traceable_tensor_subclasses",
677
- [
678
- BasevLLMParameter,
679
- ModelWeightParameter,
680
- _ColumnvLLMParameter,
681
- RowvLLMParameter,
682
- ],
683
- ),
684
- patch(
685
- "torch._dynamo.variables.torch.can_dispatch_torch_function", return_false
686
- ),
687
- ):
688
- yield
vllm/compilation/inductor_pass.py
+1
-9
@@ -16,18 +16,10 @@ import torch
16
from torch import fx
17
from torch._subclasses.fake_tensor import FakeTensorMode, unset_fake_temporarily
18
19
-from vllm.utils.torch_utils import is_torch_equal_or_newer
20
-
19
if TYPE_CHECKING:
20
from vllm.config.utils import Range
21
24
-if is_torch_equal_or_newer("2.6"):
25
- from torch._inductor.custom_graph_pass import CustomGraphPass
26
-else:
27
- # CustomGraphPass is not present in 2.5 or lower, import our version
28
- from .torch25_custom_graph_pass import (
29
- Torch25CustomGraphPass as CustomGraphPass,
30
- )
22
+from torch._inductor.custom_graph_pass import CustomGraphPass
23
24
_pass_context = None
25
P = ParamSpec("P")
vllm/config/compilation.py
+3
-4
@@ -777,10 +777,9 @@ class CompilationConfig:
777
# and it is not yet a priority. RFC here:
778
# https://github.com/vllm-project/vllm/issues/14703
779
780
- if is_torch_equal_or_newer("2.6"):
781
- KEY = "enable_auto_functionalized_v2"
782
- if KEY not in self.inductor_compile_config:
783
- self.inductor_compile_config[KEY] = False
780
+ KEY = "enable_auto_functionalized_v2"
781
+ if KEY not in self.inductor_compile_config:
782
+ self.inductor_compile_config[KEY] = False
783
784
for k, v in self.inductor_passes.items():
785
if not isinstance(v, str):
vllm/distributed/utils.py
+7
-26
@@ -31,7 +31,6 @@ import vllm.envs as envs
31
from vllm.logger import init_logger
32
from vllm.utils.network_utils import get_tcp_uri
33
from vllm.utils.system_utils import suppress_stdout
34
-from vllm.utils.torch_utils import is_torch_equal_or_newer
34
35
logger = init_logger(__name__)
36
@@ -429,20 +428,11 @@ def init_gloo_process_group(
428
different torch versions.
429
"""
430
with suppress_stdout():
432
- if is_torch_equal_or_newer("2.6"):
433
- pg = ProcessGroup(
434
- prefix_store,
435
- group_rank,
436
- group_size,
437
- )
438
- else:
439
- options = ProcessGroup.Options(backend="gloo")
440
- pg = ProcessGroup(
441
- prefix_store,
442
- group_rank,
443
- group_size,
444
- options,
445
- )
431
+ pg = ProcessGroup(
432
+ prefix_store,
433
+ group_rank,
434
+ group_size,
435
+ )
436
from torch.distributed.distributed_c10d import ProcessGroupGloo
437
438
backend_class = ProcessGroupGloo(
@@ -450,9 +440,7 @@ def init_gloo_process_group(
440
)
441
backend_type = ProcessGroup.BackendType.GLOO
442
device = torch.device("cpu")
453
- if is_torch_equal_or_newer("2.6"):
454
- # _set_default_backend is supported in torch >= 2.6
455
- pg._set_default_backend(backend_type)
443
+ pg._set_default_backend(backend_type)
444
backend_class._set_sequence_number_for_group()
445
446
pg._register_backend(device, backend_type, backend_class)
@@ -534,12 +522,5 @@ def stateless_destroy_torch_distributed_process_group(pg: ProcessGroup) -> None:
522
Destroy ProcessGroup returned by
523
stateless_init_torch_distributed_process_group().
524
"""
537
- if is_torch_equal_or_newer("2.7"):
538
- pg.shutdown()
539
- else:
540
- # Lazy import for non-CUDA backends.
541
- from torch.distributed.distributed_c10d import _shutdown_backend
542
-
543
- _shutdown_backend(pg)
544
-
525
+ pg.shutdown()
526
_unregister_process_group(pg.group_name)
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
-11
@@ -52,7 +52,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
52
)
53
from vllm.platforms import current_platform
54
from vllm.triton_utils import tl, triton
55
-from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
55
+from vllm.utils.torch_utils import direct_register_custom_op
56
57
logger = init_logger(__name__)
58
@@ -1406,11 +1406,6 @@ direct_register_custom_op(
1406
op_func=inplace_fused_experts,
1407
mutates_args=["hidden_states"],
1408
fake_impl=inplace_fused_experts_fake,
1409
- tags=(
1410
- ()
1411
- if is_torch_equal_or_newer("2.7.0")
1412
- else (torch.Tag.needs_fixed_stride_order,)
1413
- ),
1409
)
1410
1411
@@ -1501,11 +1496,6 @@ direct_register_custom_op(
1496
op_name="outplace_fused_experts",
1497
op_func=outplace_fused_experts,
1498
fake_impl=outplace_fused_experts_fake,
1504
- tags=(
1505
- ()
1506
- if is_torch_equal_or_newer("2.7.0")
1507
- else (torch.Tag.needs_fixed_stride_order,)
1508
- ),
1499
)
1500
1501
vllm/model_executor/layers/quantization/mxfp4.py
-3
@@ -56,7 +56,6 @@ from vllm.scalar_type import scalar_types
56
from vllm.utils.flashinfer import has_flashinfer
57
from vllm.utils.import_utils import has_triton_kernels
58
from vllm.utils.math_utils import round_up
59
-from vllm.utils.torch_utils import is_torch_equal_or_newer
59
60
logger = init_logger(__name__)
61
@@ -89,7 +88,6 @@ def get_mxfp4_backend_with_lora() -> Mxfp4Backend:
88
# If FlashInfer is not available, try either Marlin or Triton
89
triton_kernels_supported = (
90
has_triton_kernels()
92
- and is_torch_equal_or_newer("2.8.0")
91
# NOTE: triton_kernels are only confirmed to work on SM90 and SM100
92
# SM110 fails with this error: https://github.com/vllm-project/vllm/issues/29317
93
# SM120 needs this fix: https://github.com/triton-lang/triton/pull/8498
@@ -151,7 +149,6 @@ def get_mxfp4_backend(with_lora_support: bool) -> Mxfp4Backend:
149
# If FlashInfer is not available, try either Marlin or Triton
150
triton_kernels_supported = (
151
has_triton_kernels()
154
- and is_torch_equal_or_newer("2.8.0")
152
# NOTE: triton_kernels are only confirmed to work on SM90 and SM100
153
# SM110 fails with this error: https://github.com/vllm-project/vllm/issues/29317
154
# SM120 needs this fix: https://github.com/triton-lang/triton/pull/8498
vllm/model_executor/layers/quantization/torchao.py
-14
@@ -108,20 +108,6 @@ class TorchAOConfig(QuantizationConfig):
108
skip_modules: list[str] | None = None,
109
is_checkpoint_torchao_serialized: bool = False,
110
) -> None:
111
- """
112
- # TorchAO quantization relies on tensor subclasses. In order,
113
- # to enable proper caching this needs standalone compile
114
- if is_torch_equal_or_newer("2.8.0.dev"):
115
- os.environ["VLLM_TEST_STANDALONE_COMPILE"] = "1"
116
- logger.info(
117
- "Using TorchAO: Setting VLLM_TEST_STANDALONE_COMPILE=1")
118
-
119
- # TODO: remove after the torch dependency is updated to 2.8
120
- if is_torch_equal_or_newer(
121
- "2.7.0") and not is_torch_equal_or_newer("2.8.0.dev"):
122
- os.environ["VLLM_DISABLE_COMPILE_CACHE"] = "1"
123
- logger.info("Using TorchAO: Setting VLLM_DISABLE_COMPILE_CACHE=1")
124
- """
111
super().__init__()
112
self.torchao_config = torchao_config
113
self.skip_modules = skip_modules or []
vllm/utils/torch_utils.py
+1
-3
@@ -709,9 +709,7 @@ def is_torch_equal(target: str) -> bool:
709
710
# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
711
def supports_xccl() -> bool:
712
- return (
713
- is_torch_equal_or_newer("2.8.0.dev") and torch.distributed.is_xccl_available()
714
- )
712
+ return torch.distributed.is_xccl_available()
713
714
715
# create a library to hold the custom op