4dfad17ed1 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/compile/fullgraph/test_basic_correctness.py

+2 -3

@@ -6,7 +6,6 @@ import pytest

6

7 from vllm.config import CompilationMode

8 from vllm.platforms import current_platform

9 -from vllm.utils.torch_utils import cuda_device_count_stateless

9

10 from ...utils import compare_all_settings

11

@@ -109,10 +108,10 @@ def test_compile_correctness(

108 tp_size = test_setting.tp_size

109 attn_backend = test_setting.attn_backend

110 method = test_setting.method

112 - if cuda_device_count_stateless() < pp_size * tp_size:

111 + if current_platform.device_count() < pp_size * tp_size:

112 pytest.skip(

113 f"Need at least {pp_size}*{tp_size} CUDA gpus but got "

115 - f"{cuda_device_count_stateless()}"

114 + f"{current_platform.device_count()}"

115 )

116

117 final_args = [

tests/compile/test_config.py

+1 -1

@@ -412,7 +412,7 @@ def test_cudagraph_sizes_post_init(

412

413 with (

414 ctx,

415 - patch("vllm.config.parallel.cuda_device_count_stateless", return_value=tp_size),

415 + patch.object(current_platform, "device_count", return_value=tp_size),

416 ):

417 kwargs = {}

418 if cudagraph_capture_sizes is not None:

tests/distributed/test_utils.py

+1 -2

@@ -13,7 +13,6 @@ from vllm.distributed.utils import StatelessProcessGroup

13 from vllm.platforms import current_platform

14 from vllm.utils.network_utils import get_open_port

15 from vllm.utils.system_utils import update_environment_variables

16 -from vllm.utils.torch_utils import cuda_device_count_stateless

16

17 from ..utils import multi_gpu_test

18

@@ -21,7 +20,7 @@ from ..utils import multi_gpu_test

20 @ray.remote

21 class _CUDADeviceCountStatelessTestActor:

22 def get_count(self):

24 - return cuda_device_count_stateless()

23 + return current_platform.device_count()

24

25 def set_cuda_visible_devices(self, cuda_visible_devices: str):

26 update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices})

tests/kernels/moe/test_modular_kernel_combinations.py

+3 -3

@@ -15,7 +15,7 @@ from vllm.config import VllmConfig, set_current_vllm_config

15 from vllm.platforms import current_platform

16 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

17 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm

18 -from vllm.utils.torch_utils import cuda_device_count_stateless, set_random_seed

18 +from vllm.utils.torch_utils import set_random_seed

19 from vllm.v1.worker.workspace import init_workspace_manager

20

21 from .modular_kernel_tools.common import (

@@ -310,10 +310,10 @@ def test_modular_kernel_combinations_multigpu(

310 world_size: int,

311 pytestconfig,

312 ):

313 - if cuda_device_count_stateless() < world_size:

313 + if current_platform.device_count() < world_size:

314 pytest.skip(

315 f"Not enough GPUs available to run, got "

316 - f"{cuda_device_count_stateless()} expected "

316 + f"{current_platform.device_count()} expected "

317 f"{world_size}."

318 )

319

tests/model_executor/model_loader/test_reload.py

+3 -4

@@ -19,7 +19,6 @@ from vllm.model_executor.model_loader.reload.meta import (

19 from vllm.model_executor.model_loader.reload.types import LayerReloadingInfo

20 from vllm.model_executor.model_loader.reload.utils import get_layer_tensors

21 from vllm.platforms import current_platform

22 -from vllm.utils.torch_utils import cuda_device_count_stateless

22

23

24 def test_move_metatensors():

@@ -140,7 +139,7 @@ def test_get_numel_loaded():

139 ],

140 )

141 def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):

143 - if cuda_device_count_stateless() < tp_size:

142 + if current_platform.device_count() < tp_size:

143 pytest.skip(reason="Not enough CUDA devices")

144

145 if "FP8" in base_model and not current_platform.supports_fp8():

@@ -206,8 +205,8 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):

205 def test_online_quantize_reload(

206 base_model, mul_model, add_model, quantization, tp_size, vllm_runner

207 ):

209 - if cuda_device_count_stateless() < tp_size:

210 - pytest.skip(reason="Not enough CUDA devices")

208 + if current_platform.device_count() < tp_size:

209 + pytest.skip(reason="Not enough GPU devices")

210

211 if quantization == "fp8" and not current_platform.supports_fp8():

212 pytest.skip(reason="Requires FP8 support")

tests/models/quantization/test_gpt_oss.py

+2 -2

@@ -21,8 +21,8 @@ import lm_eval

21 import pytest

22 from packaging import version

23

24 +from vllm.platforms import current_platform

25 from vllm.platforms.rocm import on_gfx950

25 -from vllm.utils.torch_utils import cuda_device_count_stateless

26

27 MODEL_ACCURACIES = {

28 # Full quantization: attention linears and MoE linears

@@ -89,7 +89,7 @@ def test_gpt_oss_attention_quantization(

89 expected_accuracy: float,

90 monkeypatch: pytest.MonkeyPatch,

91 ):

92 - if tp_size > cuda_device_count_stateless():

92 + if tp_size > current_platform.device_count():

93 pytest.skip("Not enough GPUs to run this test case")

94

95 if "amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8" in model_name and on_gfx950():

tests/utils.py

+3 -4

@@ -58,7 +58,6 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser

58 from vllm.utils.mem_constants import GB_bytes

59 from vllm.utils.network_utils import get_open_port

60 from vllm.utils.torch_utils import (

61 - cuda_device_count_stateless,

61 set_random_seed, # noqa: F401 - re-exported for use in test files

62 )

63

@@ -384,7 +383,7 @@ class RemoteVLLMServer:

383 elif current_platform.is_cuda():

384 with _nvml():

385 total_used = 0

387 - device_count = cuda_device_count_stateless()

386 + device_count = current_platform.device_count()

387 for i in range(device_count):

388 handle = nvmlDeviceGetHandleByIndex(i)

389 mem_info = nvmlDeviceGetMemoryInfo(handle)

@@ -1497,7 +1496,7 @@ def multi_gpu_marks(*, num_gpus: int):

1496 """Get a collection of pytest marks to apply for `@multi_gpu_test`."""

1497 test_selector = pytest.mark.distributed(num_gpus=num_gpus)

1498 test_skipif = pytest.mark.skipif(

1500 - cuda_device_count_stateless() < num_gpus,

1499 + current_platform.device_count() < num_gpus,

1500 reason=f"Need at least {num_gpus} GPUs to run the test.",

1501 )

1502

@@ -1529,7 +1528,7 @@ def gpu_tier_mark(*, min_gpus: int = 1, max_gpus: int | None = None):

1528 @gpu_tier_mark(max_gpus=1) # only on single-GPU

1529 @gpu_tier_mark(min_gpus=2, max_gpus=4) # 2-4 GPUs only

1530 """

1532 - gpu_count = cuda_device_count_stateless()

1531 + gpu_count = current_platform.device_count()

1532 marks = []

1533

1534 if min_gpus > 1:

tests/v1/shutdown/test_delete.py

+3 -3

@@ -11,8 +11,8 @@ from tests.v1.shutdown.utils import (

11 )

12 from vllm import LLM, SamplingParams

13 from vllm.engine.arg_utils import AsyncEngineArgs

14 +from vllm.platforms import current_platform

15 from vllm.sampling_params import RequestOutputKind

15 -from vllm.utils.torch_utils import cuda_device_count_stateless

16 from vllm.v1.engine.async_llm import AsyncLLM

17

18 MODELS = ["hmellor/tiny-random-LlamaForCausalLM"]

@@ -34,7 +34,7 @@ async def test_async_llm_delete(

34 tensor_parallel_size: degree of tensor parallelism

35 send_one_request: send one request to engine before deleting

36 """

37 - if cuda_device_count_stateless() < tensor_parallel_size:

37 + if current_platform.device_count() < tensor_parallel_size:

38 pytest.skip(reason="Not enough CUDA devices")

39

40 engine_args = AsyncEngineArgs(

@@ -83,7 +83,7 @@ def test_llm_delete(

83 enable_multiprocessing: enable workers in separate process(es)

84 send_one_request: send one request to engine before deleting

85 """

86 - if cuda_device_count_stateless() < tensor_parallel_size:

86 + if current_platform.device_count() < tensor_parallel_size:

87 pytest.skip(reason="Not enough CUDA devices")

88

89 with monkeypatch.context() as m:

tests/v1/shutdown/test_forward_error.py

+3 -3

@@ -15,7 +15,7 @@ from tests.v1.shutdown.utils import (

15 from vllm import LLM, AsyncEngineArgs, SamplingParams

16 from vllm.distributed import get_tensor_model_parallel_rank

17 from vllm.model_executor.models.llama import LlamaForCausalLM

18 -from vllm.utils.torch_utils import cuda_device_count_stateless

18 +from vllm.platforms import current_platform

19 from vllm.v1.engine.async_llm import AsyncLLM

20 from vllm.v1.engine.exceptions import EngineDeadError

21

@@ -60,7 +60,7 @@ async def test_async_llm_model_error(

60

61 AsyncLLM always uses an MP client.

62 """

63 - if cuda_device_count_stateless() < tensor_parallel_size:

63 + if current_platform.device_count() < tensor_parallel_size:

64 pytest.skip(reason="Not enough CUDA devices")

65

66 # Monkeypatch an error in the model.

@@ -126,7 +126,7 @@ def test_llm_model_error(

126 TODO(andy) - LLM without multiprocessing; LLM with multiprocessing

127 and >1 rank

128 """

129 - if cuda_device_count_stateless() < tensor_parallel_size:

129 + if current_platform.device_count() < tensor_parallel_size:

130 pytest.skip(reason="Not enough CUDA devices")

131

132 with monkeypatch.context() as m:

tests/v1/shutdown/test_startup_error.py

+3 -3

@@ -15,7 +15,7 @@ from vllm import LLM

15 from vllm.distributed import get_tensor_model_parallel_rank

16 from vllm.engine.arg_utils import AsyncEngineArgs

17 from vllm.model_executor.models.llama import LlamaForCausalLM

18 -from vllm.utils.torch_utils import cuda_device_count_stateless

18 +from vllm.platforms import current_platform

19 from vllm.v1.engine.async_llm import AsyncLLM

20

21 MODELS = ["hmellor/tiny-random-LlamaForCausalLM"]

@@ -57,7 +57,7 @@ def test_async_llm_startup_error(

57 Test profiling (forward()) and load weights failures.

58 AsyncLLM always uses an MP client.

59 """

60 - if cuda_device_count_stateless() < tensor_parallel_size:

60 + if current_platform.device_count() < tensor_parallel_size:

61 pytest.skip(reason="Not enough CUDA devices")

62

63 # Monkeypatch an error in the model.

@@ -99,7 +99,7 @@ def test_llm_startup_error(

99 # If MODELS list grows, each architecture needs its own test variant.

100 if model != "JackFram/llama-68m":

101 pytest.skip(reason="Only test JackFram/llama-68m")

102 - if cuda_device_count_stateless() < tensor_parallel_size:

102 + if current_platform.device_count() < tensor_parallel_size:

103 pytest.skip(reason="Not enough CUDA devices")

104

105 with monkeypatch.context() as m:

tools/pre_commit/check_torch_cuda.py

+2

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

10 _TORCH_CUDA_PATTERNS = [

12 r"\bwith\storch\.cuda\.device\b",

13 + # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally

14 + r"\bcuda_device_count_stateless\(\)\b",

15 ]

16

17 ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}

vllm/config/parallel.py

+2 -3

@@ -16,7 +16,6 @@ from vllm.config.utils import config

16 from vllm.logger import init_logger

17 from vllm.platforms import current_platform

18 from vllm.utils.network_utils import get_open_ports_list

19 -from vllm.utils.torch_utils import cuda_device_count_stateless

19

20 if TYPE_CHECKING:

21 from ray.runtime_env import RuntimeEnv

@@ -726,9 +725,9 @@ class ParallelConfig:

725 backend = "mp"

726 elif (

727 current_platform.is_cuda()

729 - and cuda_device_count_stateless() < self.world_size

728 + and current_platform.device_count() < self.world_size

729 ):

731 - gpu_count = cuda_device_count_stateless()

730 + gpu_count = current_platform.device_count()

731 raise ValueError(

732 f"World size ({self.world_size}) is larger than the number of "

733 f"available GPUs ({gpu_count}) in this node. If this is "

vllm/distributed/device_communicators/all_reduce_utils.py

+2 -2

@@ -19,8 +19,8 @@ import torch.multiprocessing as mp

19 import vllm.envs as envs

20 from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary

21 from vllm.logger import init_logger

22 +from vllm.platforms import current_platform

23 from vllm.utils.system_utils import update_environment_variables

23 -from vllm.utils.torch_utils import cuda_device_count_stateless

24

25 logger = init_logger(__name__)

26

@@ -320,7 +320,7 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:

320

321 is_distributed = dist.is_initialized()

322

323 - num_dev = cuda_device_count_stateless()

323 + num_dev = current_platform.device_count()

324 cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES

325 if cuda_visible_devices is None:

326 cuda_visible_devices = ",".join(str(i) for i in range(num_dev))

vllm/distributed/device_communicators/custom_all_reduce.py

+1 -2

@@ -17,7 +17,6 @@ from vllm.distributed.device_communicators.all_reduce_utils import (

17 from vllm.distributed.parallel_state import in_the_same_node_as

18 from vllm.logger import init_logger

19 from vllm.platforms import current_platform

20 -from vllm.utils.torch_utils import cuda_device_count_stateless

20

21 try:

22 ops.meta_size()

@@ -135,7 +134,7 @@ class CustomAllreduce:

134 if cuda_visible_devices:

135 device_ids = list(map(int, cuda_visible_devices.split(",")))

136 else:

138 - device_ids = list(range(cuda_device_count_stateless()))

137 + device_ids = list(range(current_platform.device_count()))

138

139 physical_device_id = device_ids[device.index]

140 tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")

vllm/distributed/device_communicators/quick_all_reduce.py

+1 -2

@@ -13,7 +13,6 @@ from vllm.config import get_current_vllm_config_or_none

13 from vllm.distributed.parallel_state import in_the_same_node_as

14 from vllm.logger import init_logger

15 from vllm.platforms import current_platform

16 -from vllm.utils.torch_utils import cuda_device_count_stateless

16

17 logger = init_logger(__name__)

18

@@ -137,7 +136,7 @@ class QuickAllReduce:

136 if cuda_visible_devices:

137 device_ids = list(map(int, cuda_visible_devices.split(",")))

138 else:

140 - device_ids = list(range(cuda_device_count_stateless()))

139 + device_ids = list(range(current_platform.device_count()))

140 physical_device_id = device_ids[device.index]

141 tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")

142 gather_list = [

vllm/platforms/cuda.py

+29 -3

@@ -9,7 +9,7 @@ from __future__ import annotations

9 import os

10 from collections.abc import Callable

11 from datetime import timedelta

12 -from functools import cache, wraps

12 +from functools import cache, lru_cache, wraps

13 from typing import TYPE_CHECKING, TypeVar

14

15 import torch

@@ -20,9 +20,9 @@ from typing_extensions import ParamSpec

20 # import custom ops, trigger op registration

21 import vllm._C # noqa

22 import vllm._C_stable_libtorch # noqa

23 +import vllm.envs as envs

24 from vllm.logger import init_logger

25 from vllm.utils.import_utils import import_pynvml

25 -from vllm.utils.torch_utils import cuda_device_count_stateless

26 from vllm.v1.attention.backends.registry import AttentionBackendEnum

27

28 from .interface import DeviceCapability, Platform, PlatformEnum

@@ -47,6 +47,32 @@ pynvml = import_pynvml()

47 torch.backends.cuda.enable_cudnn_sdp(False)

48

49

50 +@lru_cache(maxsize=8)

51 +def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int:

52 + """Get number of CUDA devices, caching based on the value of CUDA_VISIBLE_DEVICES

53 + at the time of call.

54 +

55 + This should be used instead of torch.accelerator.device_count() unless

56 + CUDA_VISIBLE_DEVICES has already been set to the desired value.

57 +

58 + # This can be removed and simply replaced with torch.cuda.get_device_count

59 + # after https://github.com/pytorch/pytorch/pull/122815 is released."""

60 + # Note: cuda_visible_devices is not used, but we keep it as an argument for

61 + # LRU Cache purposes.

62 +

63 + # Code below is based on

64 + # https://github.com/pytorch/pytorch/blob/

65 + # c1cd946818442aca8c7f812b16d187ce1586c3bc/

66 + # torch/cuda/__init__.py#L831C1-L831C17

67 + import torch.cuda

68 +

69 + if not torch.cuda._is_compiled():

70 + return 0

71 + raw_count = torch.cuda._device_count_nvml()

72 + r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count

73 + return r

74 +

75 +

76 @cache

77 def _get_backend_priorities(

78 use_mla: bool,

@@ -456,7 +482,7 @@ class CudaPlatformBase(Platform):

482

483 @classmethod

484 def device_count(cls) -> int:

459 - return cuda_device_count_stateless()

485 + return _cuda_device_count_stateless(envs.CUDA_VISIBLE_DEVICES)

486

487 @classmethod

488 def check_if_supports_dtype(cls, dtype: torch.dtype):

vllm/platforms/rocm.py

+33 -2

@@ -13,7 +13,6 @@ from torch.distributed.distributed_c10d import is_nccl_available

13

14 import vllm.envs as envs

15 from vllm.logger import init_logger

16 -from vllm.utils.torch_utils import cuda_device_count_stateless

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

17

18 from .interface import DeviceCapability, Platform, PlatformEnum

@@ -67,6 +66,38 @@ _ROCM_DEVICE_ID_NAME_MAP: dict[str, str] = {

66 }

67

68

69 +@lru_cache(maxsize=8)

70 +def _rocm_device_count_stateless(cuda_visible_devices: str | None = None) -> int:

71 + """Get number of ROCm devices, caching based on the value of CUDA_VISIBLE_DEVICES

72 + at the time of call.

73 +

74 + This should be used instead of torch.accelerator.device_count() unless

75 + CUDA_VISIBLE_DEVICES has already been set to the desired value.

76 +

77 + # This can be removed and simply replaced with torch.cuda.get_device_count

78 + # after https://github.com/pytorch/pytorch/pull/122815 is released."""

79 + # Note: cuda_visible_devices is not used, but we keep it as an argument for

80 + # LRU Cache purposes.

81 +

82 + # Code below is based on

83 + # https://github.com/pytorch/pytorch/blob/

84 + # c1cd946818442aca8c7f812b16d187ce1586c3bc/

85 + # torch/cuda/__init__.py#L831C1-L831C17

86 + import torch.cuda

87 +

88 + if not torch.cuda._is_compiled():

89 + return 0

90 + # ROCm uses amdsmi instead of nvml for stateless device count

91 + # This requires a sufficiently modern version of Torch 2.4.0

92 + raw_count = (

93 + torch.cuda._device_count_amdsmi()

94 + if (hasattr(torch.cuda, "_device_count_amdsmi"))

95 + else -1

96 + )

97 + r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count

98 + return r

99 +

100 +

101 def _sync_hip_cuda_env_vars():

102 """Ensure HIP_VISIBLE_DEVICES and CUDA_VISIBLE_DEVICES are consistent.

103 Treats empty string as unset. Raises on genuine conflicts."""

@@ -810,7 +841,7 @@ class RocmPlatform(Platform):

841

842 @classmethod

843 def device_count(cls) -> int:

813 - return cuda_device_count_stateless()

844 + return _rocm_device_count_stateless(getattr(envs, cls.device_control_env_var))

845

846 @classmethod

847 def check_if_supports_dtype(cls, dtype: torch.dtype):

vllm/usage/usage_lib.py

+1 -2

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

22 from vllm.connections import global_http_connection

23 from vllm.logger import init_logger

24 from vllm.utils.platform_utils import cuda_get_device_properties

25 -from vllm.utils.torch_utils import cuda_device_count_stateless

25 from vllm.version import __version__ as VLLM_VERSION

26

27 logger = init_logger(__name__)

@@ -196,7 +195,7 @@ class UsageMessage:

195 from vllm.platforms import current_platform

196

197 if current_platform.is_cuda_alike():

199 - self.gpu_count = cuda_device_count_stateless()

198 + self.gpu_count = current_platform.device_count()

199 self.gpu_type, self.gpu_memory_per_device = cuda_get_device_properties(

200 0, ("name", "total_memory")

201 )

vllm/utils/torch_utils.py

-45

@@ -6,7 +6,6 @@ import os

6 import random

7 import threading

8 from collections.abc import Callable, Collection

9 -from functools import lru_cache

9 from typing import TYPE_CHECKING, Any, TypeVar

10

11 import numpy as np

@@ -16,7 +15,6 @@ from packaging import version

15 from packaging.version import Version

16 from torch.library import Library, infer_schema

17

19 -import vllm.envs as envs

18 from vllm.logger import init_logger

19

20 if TYPE_CHECKING:

@@ -590,49 +588,6 @@ def aux_stream() -> torch.cuda.Stream | None:

588 return _aux_stream

589

590

593 -@lru_cache(maxsize=8)

594 -def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int:

595 - # Note: cuda_visible_devices is not used, but we keep it as an argument for

596 - # LRU Cache purposes.

597 -

598 - # Code below is based on

599 - # https://github.com/pytorch/pytorch/blob/

600 - # c1cd946818442aca8c7f812b16d187ce1586c3bc/

601 - # torch/cuda/__init__.py#L831C1-L831C17

602 - import torch.cuda

603 - import torch.version

604 -

605 - from vllm.platforms import current_platform

606 -

607 - if not torch.cuda._is_compiled():

608 - return 0

609 - if current_platform.is_rocm():

610 - # ROCm uses amdsmi instead of nvml for stateless device count

611 - # This requires a sufficiently modern version of Torch 2.4.0

612 - raw_count = (

613 - torch.cuda._device_count_amdsmi()

614 - if (hasattr(torch.cuda, "_device_count_amdsmi"))

615 - else -1

616 - )

617 - else:

618 - raw_count = torch.cuda._device_count_nvml()

619 - r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count

620 - return r

621 -

622 -

623 -def cuda_device_count_stateless() -> int:

624 - """Get number of CUDA devices, caching based on the value of

625 - CUDA_VISIBLE_DEVICES at the time of call.

626 -

627 - This should be used instead of torch.accelerator.device_count()

628 - unless CUDA_VISIBLE_DEVICES has already been set to the desired

629 - value."""

630 -

631 - # This can be removed and simply replaced with torch.cuda.get_device_count

632 - # after https://github.com/pytorch/pytorch/pull/122815 is released.

633 - return _cuda_device_count_stateless(envs.CUDA_VISIBLE_DEVICES)

634 -

635 -

591 def weak_ref_tensor(tensor: Any) -> Any:

592 """

593 Create a weak reference to a tensor.

vllm/v1/executor/ray_utils.py

+1 -3

@@ -369,9 +369,7 @@ def initialize_ray_cluster(

369

370 # Prevalidate GPU requirements before Ray processing

371 if current_platform.is_cuda() and parallel_config.world_size > 1:

372 - from vllm.utils.torch_utils import cuda_device_count_stateless

373 -

374 - available_gpus = cuda_device_count_stateless()

372 + available_gpus = current_platform.device_count()

373 if parallel_config.world_size > available_gpus:

374 logger.warning(

375 "Tensor parallel size (%d) exceeds available GPUs (%d). "

replace cuda_device_count_stateless() to current_platform.device_count() (#37841)