From 1b6798a6a47d76953125aa4e9f00f0d603a98499 Mon Sep 17 00:00:00 2001 From: EduardDurech <39579228+EduardDurech@users.noreply.github.com> Date: Thu, 29 Jan 2026 20:55:13 +0100 Subject: [PATCH] Fix `torch.__version__` for PEP440 (#15682) --- python/sglang/multimodal_gen/runtime/layers/usp.py | 4 ++-- python/sglang/srt/compilation/compiler_interface.py | 9 +++++---- .../distributed/device_communicators/pynccl_allocator.py | 4 ++-- python/sglang/srt/layers/attention/fla/utils.py | 9 +++------ python/sglang/srt/layers/quantization/fp8_utils.py | 8 ++------ python/sglang/srt/server_args.py | 9 +++------ python/sglang/srt/utils/common.py | 3 ++- python/sglang/srt/utils/patch_torch.py | 5 ++--- 8 files changed, 21 insertions(+), 30 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/usp.py b/python/sglang/multimodal_gen/runtime/layers/usp.py index 04c020320..3bdc6b6cd 100644 --- a/python/sglang/multimodal_gen/runtime/layers/usp.py +++ b/python/sglang/multimodal_gen/runtime/layers/usp.py @@ -5,13 +5,13 @@ from typing import TYPE_CHECKING import torch import torch.distributed._functional_collectives as ft_c -from packaging.version import parse from torch.distributed.tensor.experimental._attention import _cp_options from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_group, get_ulysses_parallel_world_size, ) +from sglang.srt.utils.common import torch_release _cp_options.enable_load_balance = False @@ -226,7 +226,7 @@ def ring_attn( # Starting from torch 2.6.0, _templated_ring_attention expects an integer # segment_id for the attention function. - use_segment_id = parse(torch.__version__).release >= parse("2.6.0").release + use_segment_id = torch_release >= (2, 6) attn_kwargs = dict( op=attn_callable_adapter, diff --git a/python/sglang/srt/compilation/compiler_interface.py b/python/sglang/srt/compilation/compiler_interface.py index 8310f75c9..df7f28b10 100644 --- a/python/sglang/srt/compilation/compiler_interface.py +++ b/python/sglang/srt/compilation/compiler_interface.py @@ -14,6 +14,7 @@ import torch.fx as fx from sglang.srt.compilation.compilation_counter import compilation_counter from sglang.srt.compilation.inductor_pass import pass_context +from sglang.srt.utils.common import torch_release class CompilerInterface: @@ -226,7 +227,7 @@ class InductorAdaptor(CompilerInterface): hash_str, file_path = None, None from torch._inductor.codecache import FxGraphCache, compiled_fx_graph_hash - if torch.__version__.startswith("2.5"): + if torch_release[:2] == (2, 5): original_load = FxGraphCache.load original_load_name = "torch._inductor.codecache.FxGraphCache.load" @@ -252,7 +253,7 @@ class InductorAdaptor(CompilerInterface): hijacked_compile_fx_inner = ( torch._inductor.compile_fx.compile_fx_inner ) # noqa - elif torch.__version__ >= "2.6": + elif torch_release >= (2, 6): # function renamed in 2.6 original_load_name = None @@ -405,7 +406,7 @@ class InductorAdaptor(CompilerInterface): # Dynamo metrics context, see method for more details. exit_stack.enter_context(self.metrics_context()) - if torch.__version__.startswith("2.5"): + if torch_release[:2] == (2, 5): inductor_compiled_graph = FxGraphCache._lookup_graph( hash_str, example_inputs, True, False ) @@ -413,7 +414,7 @@ class InductorAdaptor(CompilerInterface): "Inductor cache lookup failed. Please remove" f"the cache directory and try again." # noqa ) - elif torch.__version__ >= "2.6": + elif torch_release >= (2, 6): from torch._inductor.output_code import CompiledFxGraphConstantsWithGm constants = CompiledFxGraphConstantsWithGm(graph) diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py index c9cfd9ffc..7dfc08c29 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py @@ -3,13 +3,13 @@ import tempfile from contextlib import nullcontext import torch -from packaging import version from torch.cuda.memory import CUDAPluggableAllocator from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils.common import torch_release -after_2_8_0 = version.parse(torch.__version__) >= version.parse("2.8.0") +after_2_8_0 = torch_release >= (2, 8) nccl_allocator_source = """ diff --git a/python/sglang/srt/layers/attention/fla/utils.py b/python/sglang/srt/layers/attention/fla/utils.py index 8613d611d..af6ca3d6e 100644 --- a/python/sglang/srt/layers/attention/fla/utils.py +++ b/python/sglang/srt/layers/attention/fla/utils.py @@ -14,6 +14,8 @@ import torch import triton from packaging import version +from sglang.srt.utils.common import torch_release + logger = logging.getLogger(__name__) COMPILER_MODE = os.getenv("FLA_COMPILER_MODE") == "1" @@ -204,11 +206,6 @@ def checkpoint(fn): return wrapper -@lru_cache(maxsize=None) -def check_pytorch_version(version_s: str = "2.4") -> bool: - return version.parse(torch.__version__) >= version.parse(version_s) - - def _cpu_device_warning(): import warnings @@ -309,7 +306,7 @@ def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: return False -if check_pytorch_version("2.4"): +if torch_release >= (2, 4): device = "cuda" if device == "cpu" else device autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device) autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 6268066fe..dda333b14 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -11,6 +11,7 @@ from sglang.srt.environ import envs from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8 from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil +from sglang.srt.utils.common import torch_release if TYPE_CHECKING: from sglang.srt.server_args import ServerArgs @@ -80,17 +81,12 @@ TORCH_DEVICE_IDENTITY = None def use_rowwise_torch_scaled_mm(): - _TORCH_VERSION = torch.__version__.split("+")[0] - try: - _TORCH_VERSION_TUPLE = tuple(map(int, _TORCH_VERSION.split(".")[:3])) - except ValueError: - _TORCH_VERSION_TUPLE = (0, 0, 0) if _is_hip: # The condition to determine if it is on a platform that supports # torch._scaled_mm rowwise feature. # The condition is determined once as the operations # are time consuming. - return get_device_capability() >= (9, 4) and _TORCH_VERSION_TUPLE >= (2, 7, 0) + return get_device_capability() >= (9, 4) and torch_release >= (2, 7) return False diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index b12400c75..d8d6c6dde 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -62,6 +62,7 @@ from sglang.srt.utils.common import ( json_list_type, nullable_str, parse_connector_type, + torch_release, wait_port_available, xpu_has_xmx_support, ) @@ -4978,10 +4979,7 @@ class ServerArgs: # NOTE: CUDA Green Context may encounter potential issues with CudaGraph on torch 2.7.x – 2.8.x, leading to performance degradation. import torch - parts = torch.__version__.split("+", 1)[0].split(".") - major = int(parts[0]) if len(parts) > 0 and parts[0].isdigit() else 0 - minor = int(parts[1]) if len(parts) > 1 and parts[1].isdigit() else 0 - if (major, minor) > (2, 6): + if torch_release >= (2, 7): logger.warning( "WARNING: PD-Multiplexing may experience performance degradation with torch versions > 2.6.x.\n" f" Current torch version is {torch.__version__}.\n" @@ -5049,8 +5047,7 @@ class ServerArgs: if self.get_model_config().is_multimodal: import torch - torch_version = torch.__version__.split("+", 1)[0] - if torch_version == "2.9.1": + if torch_release[:3] == (2, 9, 1): cudnn_version = None try: cudnn_version = torch.backends.cudnn.version() diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 17a3360d2..219a36b24 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -101,6 +101,7 @@ if TYPE_CHECKING: from sglang.srt.server_args import ServerArgs logger = logging.getLogger(__name__) +torch_release = pkg_version.parse(torch.__version__).release # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip @@ -1927,7 +1928,7 @@ def init_custom_process_group( # https://github.com/pytorch/pytorch/commit/a0c7029a75628cd5fa8df83c0de0ea98ee7fd844 # We need to determine the appropriate parameter name based on PyTorch version pg_options_param_name = ( - "backend_options" if str(torch.__version__) >= "2.6" else "pg_options" + "backend_options" if torch_release >= (2, 6) else "pg_options" ) pg, _ = _new_process_group_helper( world_size, diff --git a/python/sglang/srt/utils/patch_torch.py b/python/sglang/srt/utils/patch_torch.py index 37868c80c..7546502cd 100644 --- a/python/sglang/srt/utils/patch_torch.py +++ b/python/sglang/srt/utils/patch_torch.py @@ -14,10 +14,9 @@ from typing import Callable, Union import torch -from packaging import version from torch.multiprocessing import reductions -from sglang.srt.utils.common import is_npu +from sglang.srt.utils.common import is_npu, torch_release _is_npu = is_npu() @@ -104,7 +103,7 @@ def _modify_tuple(t, index: int, modifier: Callable): def monkey_patch_torch_compile(): - if version.parse(torch.__version__) < version.parse("2.8.0"): + if torch_release < (2, 8): # These things are cacheable by torch.compile. torch.compile just doesn't know it. # This was fixed in PyTorch 2.8, but until then, we monkey patch. import torch._higher_order_ops.auto_functionalize as af