@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 53ec16a705
Fork 0

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 6, 19:15 UTC

[Hardware] Replace torch.cuda.device_count/current_device/set_device API (#36145)

Signed-off-by: Kunshang Ji <jikunshang95@gmail.com> Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

Kunshang Ji committed Mar 12, 2026 at 22:57 UTC 53ec16a705f27dc72d5b824a5b7ccd490f235383
89 files changed +254 -219
benchmarks/attention_benchmarks/mla_runner.py
+1 -1
@@ -757,7 +757,7 @@ def _run_mla_benchmark_batched(
757
758 backend_cfg = _get_backend_config(backend)
759 device = torch.device(configs_with_params[0][0].device)
760 - torch.cuda.set_device(device)
760 + torch.accelerator.set_device_index(device)
761
762 # Determine block size
763 config_block_size = configs_with_params[0][0].block_size
benchmarks/attention_benchmarks/runner.py
+1 -1
@@ -443,7 +443,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
443 BenchmarkResult with timing and memory statistics
444 """
445 device = torch.device(config.device)
446 - torch.cuda.set_device(device)
446 + torch.accelerator.set_device_index(device)
447
448 backend_cfg = _get_backend_config(config.backend)
449
benchmarks/kernels/benchmark_cutlass_moe_fp8.py
+1 -1
@@ -64,7 +64,7 @@ def bench_run(
64 per_out_ch: bool,
65 mkn: tuple[int, int, int],
66 ):
67 - init_workspace_manager(torch.cuda.current_device())
67 + init_workspace_manager(torch.accelerator.current_device_index())
68 (m, k, n) = mkn
69
70 dtype = torch.half
benchmarks/kernels/benchmark_device_communicators.py
+1 -1
@@ -495,7 +495,7 @@ def main():
495
496 # Set device
497 device = torch.device(f"cuda:{rank}")
498 - torch.cuda.set_device(device)
498 + torch.accelerator.set_device_index(device)
499
500 # Get CPU process group
501 cpu_group = dist.new_group(backend="gloo")
benchmarks/kernels/benchmark_fused_collective.py
+2 -2
@@ -392,7 +392,7 @@ def benchmark_operation(
392 num_op_per_cudagraph = 10
393
394 # Use vLLM's graph_capture to make tensor_model_parallel_all_reduce graph-safe
395 - device = torch.device(f"cuda:{torch.cuda.current_device()}")
395 + device = torch.device(f"cuda:{torch.accelerator.current_device_index()}")
396 with graph_capture(device=device), torch.cuda.graph(graph):
397 for _ in range(num_op_per_cudagraph):
398 operation_func(*args, **kwargs)
@@ -984,7 +984,7 @@ def main():
984 world_size = int(os.environ["WORLD_SIZE"])
985
986 device = torch.device(f"cuda:{rank}")
987 - torch.cuda.set_device(device)
987 + torch.accelerator.set_device_index(device)
988 torch.set_default_device(device)
989
990 init_distributed_environment()
benchmarks/kernels/benchmark_grouped_gemm_cutlass.py
+1 -1
@@ -50,7 +50,7 @@ def bench_run(
50 per_out_ch: bool,
51 mkn: tuple[int, int, int],
52 ):
53 - init_workspace_manager(torch.cuda.current_device())
53 + init_workspace_manager(torch.accelerator.current_device_index())
54 label = "Quant Matmul"
55
56 sub_label = (
benchmarks/kernels/benchmark_w8a8_block_fp8.py
+2 -2
@@ -285,7 +285,7 @@ def tune_on_gpu(args_dict):
285 weight_shapes = args_dict["weight_shapes"]
286 args = args_dict["args"]
287
288 - torch.cuda.set_device(gpu_id)
288 + torch.accelerator.set_device_index(gpu_id)
289 print(f"Starting tuning on GPU {gpu_id} with batch sizes {batch_sizes}")
290
291 block_n = args.block_n
@@ -334,7 +334,7 @@ def distribute_batch_sizes(batch_sizes, num_gpus):
334
335 def main(args):
336 print(args)
337 - num_gpus = torch.cuda.device_count()
337 + num_gpus = torch.accelerator.device_count()
338 if num_gpus == 0:
339 raise RuntimeError("No GPU available for tuning")
340 print(f"Found {num_gpus} GPUs for parallel tuning")
docs/configuration/conserving_memory.md
+1 -1
@@ -15,7 +15,7 @@ llm = LLM(model="ibm-granite/granite-3.1-8b-instruct", tensor_parallel_size=2)
15 ```
16
17 !!! warning
18 - To ensure that vLLM initializes CUDA correctly, you should avoid calling related functions (e.g. [torch.cuda.set_device][])
18 + To ensure that vLLM initializes CUDA correctly, you should avoid calling related functions (e.g. [torch.accelerator.set_device_index][])
19 before initializing vLLM. Otherwise, you may run into an error like `RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
20
21 To control which devices are used, please instead set the `CUDA_VISIBLE_DEVICES` environment variable.
docs/usage/troubleshooting.md
+3 -3
@@ -91,8 +91,8 @@ If GPU/CPU communication cannot be established, you can use the following Python
91 import torch
92 import torch.distributed as dist
93 dist.init_process_group(backend="nccl")
94 - local_rank = dist.get_rank() % torch.cuda.device_count()
95 - torch.cuda.set_device(local_rank)
94 + local_rank = dist.get_rank() % torch.accelerator.device_count()
95 + torch.accelerator.set_device_index(local_rank)
96 data = torch.FloatTensor([1,] * 128).to("cuda")
97 dist.all_reduce(data, op=dist.ReduceOp.SUM)
98 torch.accelerator.synchronize()
@@ -337,7 +337,7 @@ import vllm
337 import torch
338
339 print(f"CUDA available: {torch.cuda.is_available()}")
340 -print(f"CUDA device count: {torch.cuda.device_count()}")
340 +print(f"CUDA device count: {torch.accelerator.device_count()}")
341 EOF
342 ```
343
examples/online_serving/new_weight_syncing/rlhf_http_ipc.py
+1 -1
@@ -106,7 +106,7 @@ def main():
106 # IPC requires the training model to be on the same GPU as the vLLM server
107 # The server should be started on GPU 0 with reduced memory utilization
108 device = "cuda:0"
109 - torch.cuda.set_device(device)
109 + torch.accelerator.set_device_index(device)
110
111 # Load the training model on the same GPU as the server
112 # Use bfloat16 to reduce memory footprint
examples/online_serving/new_weight_syncing/rlhf_http_nccl.py
+1 -1
@@ -131,7 +131,7 @@ def main():
131 inference_world_size = get_world_size(BASE_URL)
132 world_size = inference_world_size + 1 # +1 for the trainer
133 device = f"cuda:{inference_world_size}"
134 - torch.cuda.set_device(device)
134 + torch.accelerator.set_device_index(device)
135
136 # Load the training model
137 print(f"Loading training model: {MODEL_NAME}")
tests/compile/passes/distributed/test_async_tp.py
+1 -1
@@ -300,7 +300,7 @@ def async_tp_pass_on_test_model(
300 set_random_seed(0)
301
302 device = torch.device(f"cuda:{local_rank}")
303 - torch.cuda.set_device(device)
303 + torch.accelerator.set_device_index(device)
304 torch.set_default_device(device)
305 torch.set_default_dtype(dtype)
306
tests/compile/passes/distributed/test_fusion_all_reduce.py
+1 -1
@@ -262,7 +262,7 @@ def all_reduce_fusion_pass_on_test_model(
262 set_random_seed(0)
263
264 device = torch.device(f"cuda:{local_rank}")
265 - torch.cuda.set_device(device)
265 + torch.accelerator.set_device_index(device)
266 torch.set_default_device(device)
267 torch.set_default_dtype(dtype)
268
tests/compile/passes/distributed/test_sequence_parallelism.py
+1 -1
@@ -228,7 +228,7 @@ def sequence_parallelism_pass_on_test_model(
228 set_random_seed(0)
229
230 device = torch.device(f"cuda:{local_rank}")
231 - torch.cuda.set_device(device)
231 + torch.accelerator.set_device_index(device)
232 torch.set_default_device(device)
233 torch.set_default_dtype(dtype)
234
tests/conftest.py
+2 -2
@@ -428,7 +428,7 @@ class HfRunner:
428 )
429
430 # don't put this import at the top level
431 - # it will call torch.cuda.device_count()
431 + # it will call torch.accelerator.device_count()
432 from transformers import AutoProcessor
433
434 self.processor = AutoProcessor.from_pretrained(
@@ -1535,7 +1535,7 @@ def clean_gpu_memory_between_tests():
1535
1536 from tests.utils import wait_for_gpu_memory_to_clear
1537
1538 - num_gpus = torch.cuda.device_count()
1538 + num_gpus = torch.accelerator.device_count()
1539 if num_gpus > 0:
1540 try:
1541 wait_for_gpu_memory_to_clear(
tests/cuda/scripts/check_device_count_respects_env.py
+1 -1
@@ -14,7 +14,7 @@ import torch # noqa: E402
14 from vllm.platforms import current_platform # noqa: F401, E402
15
16 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
17 -count = torch.cuda.device_count()
17 +count = torch.accelerator.device_count()
18
19 if count == 0:
20 sys.exit(0) # Skip: no GPUs available
tests/distributed/eplb_utils.py
+1 -1
@@ -42,7 +42,7 @@ def set_env_vars_and_device(env: dict[str, str]) -> None:
42 update_environment_variables(env)
43 local_rank = os.environ["LOCAL_RANK"]
44 device = torch.device(f"cuda:{local_rank}")
45 - torch.cuda.set_device(device)
45 + torch.accelerator.set_device_index(device)
46
47 # Create a minimal vllm config for init_distributed_environment
48 vllm_config = VllmConfig()
tests/distributed/test_comm_ops.py
+6 -6
@@ -43,7 +43,7 @@ def all_reduce_test_worker(
43 monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
44
45 device = torch.device(f"cuda:{rank}")
46 - torch.cuda.set_device(device)
46 + torch.accelerator.set_device_index(device)
47 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
48 num_elements = 8
49 all_tensors = [
@@ -69,7 +69,7 @@ def reduce_scatter_test_worker(
69 # they will be able to set the device to the correct GPU
70 monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
71 device = torch.device(f"cuda:{rank}")
72 - torch.cuda.set_device(device)
72 + torch.accelerator.set_device_index(device)
73 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
74
75 num_elements = 8
@@ -100,7 +100,7 @@ def all_gather_test_worker(
100 # they will be able to set the device to the correct GPU
101 monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
102 device = torch.device(f"cuda:{rank}")
103 - torch.cuda.set_device(device)
103 + torch.accelerator.set_device_index(device)
104 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
105 num_dimensions = 3
106 tensor_size = list(range(2, num_dimensions + 2))
@@ -134,7 +134,7 @@ def broadcast_tensor_dict_test_worker(
134 # they will be able to set the device to the correct GPU
135 monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
136 device = torch.device(f"cuda:{rank}")
137 - torch.cuda.set_device(device)
137 + torch.accelerator.set_device_index(device)
138 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
139 test_dict = {
140 # device tensor
@@ -171,7 +171,7 @@ def send_recv_tensor_dict_test_worker(
171 ):
172 monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
173 device = torch.device(f"cuda:{rank}")
174 - torch.cuda.set_device(device)
174 + torch.accelerator.set_device_index(device)
175 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
176
177 test_dict = {
@@ -317,7 +317,7 @@ def send_recv_test_worker(
317 ):
318 monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
319 device = torch.device(f"cuda:{rank}")
320 - torch.cuda.set_device(device)
320 + torch.accelerator.set_device_index(device)
321 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
322
323 size = 64
tests/distributed/test_custom_all_reduce.py
+7 -9
@@ -35,7 +35,7 @@ def graph_allreduce(
35 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
36 m.delenv("HIP_VISIBLE_DEVICES", raising=False)
37 device = torch.device(f"cuda:{rank}")
38 - torch.cuda.set_device(device)
38 + torch.accelerator.set_device_index(device)
39 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
40 ensure_model_parallel_initialized(tp_size, pp_size)
41 group = get_tp_group().device_group
@@ -62,12 +62,10 @@ def graph_allreduce(
62 for dtype in [torch.float32, torch.float16, torch.bfloat16]:
63 with graph_capture(device=device) as graph_capture_context:
64 # use integers so result matches NCCL exactly
65 - inp1 = torch.randint(
66 - 1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()
67 - )
68 - inp2 = torch.randint(
69 - 1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()
70 - )
65 + device_idx = torch.accelerator.current_device_index()
66 + inp1 = torch.randint(1, 16, (sz,), dtype=dtype, device=device_idx)
67 + inp2 = torch.randint(1, 16, (sz,), dtype=dtype, device=device_idx)
68 +
69 torch.accelerator.synchronize()
70 graph = torch.cuda.CUDAGraph()
71 with torch.cuda.graph(graph, stream=graph_capture_context.stream):
@@ -95,7 +93,7 @@ def eager_allreduce(
93 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
94 m.delenv("HIP_VISIBLE_DEVICES", raising=False)
95 device = torch.device(f"cuda:{rank}")
98 - torch.cuda.set_device(device)
96 + torch.accelerator.set_device_index(device)
97 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
98
99 # we use the first group to communicate once
@@ -129,6 +127,6 @@ def test_custom_allreduce(
127 test_target,
128 ):
129 world_size = tp_size * pipeline_parallel_size
132 - if world_size > torch.cuda.device_count():
130 + if world_size > torch.accelerator.device_count():
131 pytest.skip("Not enough GPUs to run the test.")
132 multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
tests/distributed/test_eplb_execute.py
+4 -4
@@ -442,7 +442,7 @@ def test_rearrange_expert_weights_with_redundancy(
442 ):
443 """Test the functionality of rearranging expert weights with redundancy."""
444
445 - if torch.cuda.device_count() < world_size:
445 + if torch.accelerator.device_count() < world_size:
446 pytest.skip(f"Need at least {world_size} GPUs to run the test")
447 distributed_run(
448 _test_rearrange_expert_weights_with_redundancy,
@@ -528,7 +528,7 @@ def test_async_transfer_layer_without_mtp(
528 ):
529 """Exercise async EPLB transfer path without MTP/spec decode."""
530
531 - if torch.cuda.device_count() < world_size:
531 + if torch.accelerator.device_count() < world_size:
532 pytest.skip(f"Need at least {world_size} GPUs to run the test")
533
534 distributed_run(
@@ -547,7 +547,7 @@ def test_rearrange_expert_weights_no_change(world_size):
547 unchanged.
548 """
549
550 - if torch.cuda.device_count() < world_size:
550 + if torch.accelerator.device_count() < world_size:
551 pytest.skip(f"Need at least {world_size} GPUs to run the test")
552 distributed_run(_test_rearrange_expert_weights_no_change, world_size)
553
@@ -623,6 +623,6 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:
623 def test_rearrange_expert_weights_profile_mode(world_size):
624 """Test profile mode (should not copy actual weights)"""
625
626 - if torch.cuda.device_count() < world_size:
626 + if torch.accelerator.device_count() < world_size:
627 pytest.skip(f"Need at least {world_size} GPUs to run the test")
628 distributed_run(_test_rearrange_expert_weights_profile_mode, world_size)
tests/distributed/test_eplb_fused_moe_layer.py
+1 -1
@@ -257,7 +257,7 @@ def test_eplb_fml(
257 intermediate_size: int,
258 column_major_scales: bool,
259 ):
260 - if torch.cuda.device_count() < world_size:
260 + if torch.accelerator.device_count() < world_size:
261 pytest.skip(f"Need at least {world_size} GPUs to run the test")
262
263 num_local_experts = num_experts // world_size
tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py
+1 -1
@@ -253,7 +253,7 @@ def test_eplb_fml(
253 monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1")
254 monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", backend)
255
256 - if torch.cuda.device_count() < world_size:
256 + if torch.accelerator.device_count() < world_size:
257 pytest.skip(f"Need at least {world_size} GPUs to run the test")
258
259 num_local_experts = num_experts // world_size
tests/distributed/test_nccl_symm_mem_allreduce.py
+2 -2
@@ -38,7 +38,7 @@ def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
38 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
39 dtype = torch.bfloat16
40 device = torch.device(f"cuda:{local_rank}")
41 - torch.cuda.set_device(device)
41 + torch.accelerator.set_device_index(device)
42 torch.set_default_device(device)
43 torch.set_default_dtype(dtype)
44 update_environment_variables(
@@ -84,7 +84,7 @@ def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
84 @pytest.mark.parametrize("world_size", [2])
85 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
86 def test_nccl_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, world_size):
87 - if world_size > torch.cuda.device_count():
87 + if world_size > torch.accelerator.device_count():
88 pytest.skip("Not enough GPUs to run the test.")
89
90 # Enable SymmMemCommunicator
tests/distributed/test_pynccl.py
+12 -12
@@ -54,7 +54,7 @@ def worker_fn_wrapper(fn):
54 update_environment_variables(env)
55 local_rank = os.environ["LOCAL_RANK"]
56 device = torch.device(f"cuda:{local_rank}")
57 - torch.cuda.set_device(device)
57 + torch.accelerator.set_device_index(device)
58 init_distributed_environment()
59 fn()
60
@@ -73,7 +73,7 @@ def worker_fn():
73
74
75 @pytest.mark.skipif(
76 - torch.cuda.device_count() < 2, reason="Need at least 2 GPUs to run the test."
76 + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
77 )
78 def test_pynccl():
79 distributed_run(worker_fn, 2)
@@ -102,7 +102,7 @@ def multiple_allreduce_worker_fn():
102
103
104 @pytest.mark.skipif(
105 - torch.cuda.device_count() < 4, reason="Need at least 4 GPUs to run the test."
105 + torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs to run the test."
106 )
107 def test_pynccl_multiple_allreduce():
108 # this tests pynccl for multiple tp groups, in a standalone way
@@ -130,7 +130,7 @@ def multiple_allreduce_with_vllm_worker_fn():
130
131
132 @pytest.mark.skipif(
133 - torch.cuda.device_count() < 4, reason="Need at least 4 GPUs to run the test."
133 + torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs to run the test."
134 )
135 def test_pynccl_multiple_allreduce_with_vllm():
136 # this tests pynccl for multiple tp groups, together with vllm
@@ -185,7 +185,7 @@ def all_gather_worker_fn():
185
186
187 @pytest.mark.skipif(
188 - torch.cuda.device_count() < 2, reason="Need at least 2 GPUs to run the test."
188 + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
189 )
190 def test_pynccl_all_gather():
191 distributed_run(all_gather_worker_fn, 2)
@@ -220,7 +220,7 @@ def all_gatherv_worker_fn():
220
221
222 @pytest.mark.skipif(
223 - torch.cuda.device_count() < 2, reason="Need at least 2 GPUs to run the test."
223 + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
224 )
225 def test_pynccl_all_gatherv():
226 distributed_run(all_gatherv_worker_fn, 2)
@@ -260,7 +260,7 @@ def reduce_scatter_worker_fn():
260
261
262 @pytest.mark.skipif(
263 - torch.cuda.device_count() < 2, reason="Need at least 2 GPUs to run the test."
263 + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
264 )
265 def test_pynccl_reduce_scatter():
266 distributed_run(reduce_scatter_worker_fn, 2)
@@ -298,14 +298,14 @@ def reduce_scatterv_worker_fn():
298
299
300 @pytest.mark.skipif(
301 - torch.cuda.device_count() < 2, reason="Need at least 2 GPUs to run the test."
301 + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
302 )
303 def test_pynccl_reduce_scatterv():
304 distributed_run(reduce_scatterv_worker_fn, 2)
305
306
307 @pytest.mark.skipif(
308 - torch.cuda.device_count() < 2, reason="Need at least 2 GPUs to run the test."
308 + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
309 )
310 def test_pynccl_with_cudagraph():
311 distributed_run(worker_fn_with_cudagraph, 2)
@@ -330,7 +330,7 @@ def send_recv_worker_fn():
330
331
332 @pytest.mark.skipif(
333 - torch.cuda.device_count() < 2, reason="Need at least 2 GPUs to run the test."
333 + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
334 )
335 def test_pynccl_send_recv():
336 distributed_run(send_recv_worker_fn, 2)
@@ -363,14 +363,14 @@ def multiple_send_recv_worker_fn():
363
364
365 @pytest.mark.skipif(
366 - torch.cuda.device_count() < 4, reason="Need at least 4 GPUs to run the test."
366 + torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs to run the test."
367 )
368 def test_pynccl_multiple_send_recv():
369 distributed_run(multiple_send_recv_worker_fn, 4)
370
371
372 @pytest.mark.skipif(
373 - torch.cuda.device_count() < 4, reason="Need at least 4 GPUs to run the test."
373 + torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs to run the test."
374 )
375 def test_pynccl_broadcast():
376 distributed_run(broadcast_worker_fn, 4)
tests/distributed/test_quick_all_reduce.py
+12 -15
@@ -39,7 +39,7 @@ def graph_quickreduce(
39 with monkeypatch.context() as m:
40 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
41 device = torch.device(f"cuda:{rank}")
42 - torch.cuda.set_device(device)
42 + torch.accelerator.set_device_index(device)
43 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
44 ensure_model_parallel_initialized(tp_size, pp_size)
45 group = get_tp_group().device_group
@@ -65,12 +65,10 @@ def graph_quickreduce(
65 for sz in test_sizes:
66 for dtype in [torch.float16, torch.bfloat16]:
67 with graph_capture(device=device) as graph_capture_context:
68 - inp1 = torch.randint(
69 - 1, 23, (sz,), dtype=dtype, device=torch.cuda.current_device()
70 - )
71 - inp2 = torch.randint(
72 - -23, 1, (sz,), dtype=dtype, device=torch.cuda.current_device()
73 - )
68 + device_idx = torch.accelerator.current_device_index()
69 + inp1 = torch.randint(1, 23, (sz,), dtype=dtype, device=device_idx)
70 + inp2 = torch.randint(-23, 1, (sz,), dtype=dtype, device=device_idx)
71 +
72 torch.accelerator.synchronize()
73 graph = torch.cuda.CUDAGraph()
74 with torch.cuda.graph(graph, stream=graph_capture_context.stream):
@@ -95,7 +93,7 @@ def eager_quickreduce(
93 with monkeypatch.context() as m:
94 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
95 device = torch.device(f"cuda:{rank}")
98 - torch.cuda.set_device(device)
96 + torch.accelerator.set_device_index(device)
97
98 init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
99
@@ -130,7 +128,7 @@ def test_custom_quick_allreduce(
128 quant_mode,
129 ):
130 world_size = tp_size * pipeline_parallel_size
133 - if world_size > torch.cuda.device_count():
131 + if world_size > torch.accelerator.device_count():
132 pytest.skip("Not enough GPUs to run the test.")
133
134 monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode)
@@ -145,7 +143,7 @@ def qr_variable_input(rank, world_size):
143 has been observed with the gpt_oss model).
144 """
145 device = torch.device(f"cuda:{rank}")
148 - torch.cuda.set_device(device)
146 + torch.accelerator.set_device_index(device)
147 qr_max_size = None # MB
148 _ptr = ops.init_custom_qr(rank, world_size, qr_max_size)
149 ranks = []
@@ -169,14 +167,13 @@ def qr_variable_input(rank, world_size):
167 s1 = 1024
168 while num < 50000: # 50000 is sufficient to identify issues.
169 dtype = torch.float16
170 + device_idx = torch.accelerator.current_device_index()
171 if num % 2 == 0:
172 s2 = 1024
174 - inp1 = torch.zeros(
175 - (s1, s2), dtype=dtype, device=torch.cuda.current_device()
176 - )
173 + inp1 = torch.zeros((s1, s2), dtype=dtype, device=device_idx)
174 else:
175 s2 = 2048
179 - inp1 = torch.ones((s1, s2), dtype=dtype, device=torch.cuda.current_device())
176 + inp1 = torch.ones((s1, s2), dtype=dtype, device=device_idx)
177 result = torch.empty_like(inp1)
178 # FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 NONE = 4
179 ops.qr_all_reduce(_ptr, inp1, result, 3, cast_bf2half=True)
@@ -198,7 +195,7 @@ def qr_variable_input(rank, world_size):
195 @pytest.mark.parametrize("pipeline_parallel_size", [1])
196 def test_custom_quick_allreduce_variable_input(tp_size, pipeline_parallel_size):
197 world_size = tp_size * pipeline_parallel_size
201 - if world_size > torch.cuda.device_count():
198 + if world_size > torch.accelerator.device_count():
199 pytest.skip("Not enough GPUs to run the test.")
200
201 multiprocessing.set_start_method("spawn", force=True)
tests/distributed/test_symm_mem_allreduce.py
+3 -3
@@ -39,7 +39,7 @@ def symm_mem_allreduce_worker(local_rank: int, world_size: int, q: mp.Queue):
39 m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
40 dtype = torch.bfloat16
41 device = torch.device(f"cuda:{local_rank}")
42 - torch.cuda.set_device(device)
42 + torch.accelerator.set_device_index(device)
43 torch.set_default_device(device)
44 torch.set_default_dtype(dtype)
45 update_environment_variables(
@@ -105,7 +105,7 @@ def test_symm_mem_allreduce(
105 monkeypatch: pytest.MonkeyPatch, tp_size, pipeline_parallel_size
106 ):
107 world_size = tp_size * pipeline_parallel_size
108 - if world_size > torch.cuda.device_count():
108 + if world_size > torch.accelerator.device_count():
109 pytest.skip("Not enough GPUs to run the test.")
110 q = mp.get_context("spawn").Queue()
111 mp.spawn(symm_mem_allreduce_worker, args=(world_size, q), nprocs=world_size)
@@ -126,7 +126,7 @@ def test_symm_mem_allreduce(
126 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
127 def test_dp_with_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch):
128 world_size = 4
129 - if world_size > torch.cuda.device_count():
129 + if world_size > torch.accelerator.device_count():
130 pytest.skip("Not enough GPUs to run the test.")
131 # Verify that the DataParallel runs without error
132 engine_args = EngineArgs(
tests/distributed/test_utils.py
+1 -1
@@ -66,7 +66,7 @@ def cpu_worker(rank, WORLD_SIZE, port1, port2):
66
67
68 def gpu_worker(rank, WORLD_SIZE, port1, port2):
69 - torch.cuda.set_device(rank)
69 + torch.accelerator.set_device_index(rank)
70 pg1 = StatelessProcessGroup.create(
71 host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE
72 )
tests/distributed/test_weight_transfer.py
+12 -12
@@ -203,7 +203,7 @@ class TestEngineRegistry:
203
204 def test_nccl_receive_weights_without_init_raises():
205 """Test that receive_weights raises if init_transfer_engine wasn't called."""
206 - if torch.cuda.device_count() < 1:
206 + if torch.accelerator.device_count() < 1:
207 pytest.skip("Need at least 1 GPU for this test")
208
209 config = WeightTransferConfig(backend="nccl")
@@ -336,7 +336,7 @@ def inference_receive_tensor(
336
337
338 @pytest.mark.skipif(
339 - torch.cuda.device_count() < 2,
339 + torch.accelerator.device_count() < 2,
340 reason="Need at least 2 GPUs to run NCCL weight transfer test.",
341 )
342 def test_nccl_weight_transfer_between_processes():
@@ -382,7 +382,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
382
383 def test_valid_update_info(self):
384 """Test creating valid IPCWeightTransferUpdateInfo."""
385 - if torch.cuda.device_count() < 1:
385 + if torch.accelerator.device_count() < 1:
386 pytest.skip("Need at least 1 GPU for this test")
387
388 # Create a dummy tensor and IPC handle
@@ -404,7 +404,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
404
405 def test_mismatched_dtype_names_raises(self):
406 """Test that mismatched dtype_names length raises ValueError."""
407 - if torch.cuda.device_count() < 1:
407 + if torch.accelerator.device_count() < 1:
408 pytest.skip("Need at least 1 GPU for this test")
409
410 dummy_tensor = torch.ones(10, 10, device="cuda:0")
@@ -422,7 +422,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
422
423 def test_mismatched_shapes_raises(self):
424 """Test that mismatched shapes length raises ValueError."""
425 - if torch.cuda.device_count() < 1:
425 + if torch.accelerator.device_count() < 1:
426 pytest.skip("Need at least 1 GPU for this test")
427
428 dummy_tensor = torch.ones(10, 10, device="cuda:0")
@@ -440,7 +440,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
440
441 def test_mismatched_ipc_handles_raises(self):
442 """Test that mismatched ipc_handles length raises ValueError."""
443 - if torch.cuda.device_count() < 1:
443 + if torch.accelerator.device_count() < 1:
444 pytest.skip("Need at least 1 GPU for this test")
445
446 dummy_tensor = torch.ones(10, 10, device="cuda:0")
@@ -458,7 +458,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
458
459 def test_valid_update_info_from_pickled(self, monkeypatch):
460 """Test creating IPCWeightTransferUpdateInfo from pickled handles."""
461 - if torch.cuda.device_count() < 1:
461 + if torch.accelerator.device_count() < 1:
462 pytest.skip("Need at least 1 GPU for this test")
463
464 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
@@ -493,7 +493,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
493
494 def test_both_handles_and_pickled_raises(self):
495 """Test that providing both ipc_handles and ipc_handles_pickled raises."""
496 - if torch.cuda.device_count() < 1:
496 + if torch.accelerator.device_count() < 1:
497 pytest.skip("Need at least 1 GPU for this test")
498
499 dummy_tensor = torch.ones(10, 10, device="cuda:0")
@@ -540,7 +540,7 @@ class TestIPCEngineParsing:
540
541 def test_parse_update_info_valid(self):
542 """Test parsing valid update info dict."""
543 - if torch.cuda.device_count() < 1:
543 + if torch.accelerator.device_count() < 1:
544 pytest.skip("Need at least 1 GPU for this test")
545
546 config = WeightTransferConfig(backend="ipc")
@@ -572,7 +572,7 @@ class TestIPCEngineParsing:
572
573 def test_parse_update_info_pickled(self, monkeypatch):
574 """Test parsing update info with pickled IPC handles (HTTP path)."""
575 - if torch.cuda.device_count() < 1:
575 + if torch.accelerator.device_count() < 1:
576 pytest.skip("Need at least 1 GPU for this test")
577
578 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
@@ -731,7 +731,7 @@ def inference_receive_ipc_tensor(
731
732
733 @pytest.mark.skipif(
734 - torch.cuda.device_count() < 1,
734 + torch.accelerator.device_count() < 1,
735 reason="Need at least 1 GPU to run IPC weight transfer test.",
736 )
737 @pytest.mark.parametrize("mode", ["ray", "http"])
@@ -789,7 +789,7 @@ def test_ipc_weight_transfer_between_processes(mode: str):
789
790 def test_ipc_receive_weights_missing_gpu_uuid_raises():
791 """Test that receive_weights raises if GPU UUID not found in IPC handles."""
792 - if torch.cuda.device_count() < 1:
792 + if torch.accelerator.device_count() < 1:
793 pytest.skip("Need at least 1 GPU for this test")
794
795 config = WeightTransferConfig(backend="ipc")
tests/entrypoints/llm/test_collective_rpc.py
+1 -1
@@ -13,7 +13,7 @@ from ...utils import create_new_process_for_each_test
13 @pytest.mark.parametrize("backend", ["mp", "ray"])
14 @create_new_process_for_each_test()
15 def test_collective_rpc(tp_size, backend, monkeypatch):
16 - if torch.cuda.device_count() < tp_size:
16 + if torch.accelerator.device_count() < tp_size:
17 pytest.skip(f"Not enough GPUs for tensor parallelism {tp_size}")
18 if tp_size == 1 and backend == "ray":
19 pytest.skip("Skip duplicate test case")
tests/entrypoints/weight_transfer/test_weight_transfer_llm.py
+5 -5
@@ -106,7 +106,7 @@ def mock_create_engine(config, parallel_config):
106 @create_new_process_for_each_test()
107 def test_get_world_size_tp1():
108 """Test world_size is correctly configured for TP=1."""
109 - if torch.cuda.device_count() < 1:
109 + if torch.accelerator.device_count() < 1:
110 pytest.skip("Need at least 1 GPU for this test")
111
112 llm = LLM(
@@ -125,7 +125,7 @@ def test_get_world_size_tp1():
125 def test_init_weight_transfer_engine_calls_engine():
126 """Test that init_weight_transfer_engine calls the engine's
127 init_transfer_engine method."""
128 - if torch.cuda.device_count() < 1:
128 + if torch.accelerator.device_count() < 1:
129 pytest.skip("Need at least 1 GPU for this test")
130
131 # Run in-process so mock.patch works (spawn won't inherit the mock)
@@ -174,7 +174,7 @@ def test_init_weight_transfer_engine_calls_engine():
174 @create_new_process_for_each_test()
175 def test_update_weights_calls_engine():
176 """Test that update_weights calls the engine's receive_weights method."""
177 - if torch.cuda.device_count() < 1:
177 + if torch.accelerator.device_count() < 1:
178 pytest.skip("Need at least 1 GPU for this test")
179
180 # Run in-process so mock.patch works (spawn won't inherit the mock)
@@ -233,7 +233,7 @@ def test_update_weights_calls_engine():
233 @create_new_process_for_each_test()
234 def test_full_weight_transfer_flow():
235 """Test the complete weight transfer flow: init -> update."""
236 - if torch.cuda.device_count() < 1:
236 + if torch.accelerator.device_count() < 1:
237 pytest.skip("Need at least 1 GPU for this test")
238
239 # Run in-process so mock.patch works (spawn won't inherit the mock)
@@ -294,7 +294,7 @@ def test_full_weight_transfer_flow():
294 @create_new_process_for_each_test()
295 def test_weight_transfer_config_backend():
296 """Test that WeightTransferConfig backend is properly configured."""
297 - if torch.cuda.device_count() < 1:
297 + if torch.accelerator.device_count() < 1:
298 pytest.skip("Need at least 1 GPU for this test")
299
300 # Test with nccl backend
tests/kernels/attention/test_attention.py
+3 -1
@@ -36,7 +36,9 @@ BLOCK_SIZES = [16, 32]
36 USE_ALIBI = [False, True]
37 KV_CACHE_DTYPE = ["auto", "fp8"]
38 SEEDS = [0]
39 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
39 +CUDA_DEVICES = [
40 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
41 +]
42
43
44 def ref_masked_attention(
tests/kernels/attention/test_cache.py
+8 -6
@@ -35,7 +35,9 @@ NUM_BLOCKS = [1024, 10000]
35
36 NUM_MAPPINGS = [256] # Arbitrary values for testing
37 SEEDS = [0]
38 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
38 +CUDA_DEVICES = [
39 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
40 +]
41
42 # We assume fp8 is always enabled for testing.
43 KV_CACHE_DTYPE = ["auto", "fp8"]
@@ -69,7 +71,7 @@ def test_reshape_and_cache(
71 pytest.skip()
72 set_random_seed(seed)
73 torch.set_default_device(device)
72 - torch.cuda.set_device(device)
74 + torch.accelerator.set_device_index(device)
75 # Create a random slot mapping.
76 num_slots = block_size * num_blocks
77 slot_mapping_lst = random.sample(range(num_slots), num_tokens)
@@ -192,7 +194,7 @@ def test_reshape_and_cache_flash(
194 ) -> None:
195 set_random_seed(seed)
196 torch.set_default_device(device)
195 - torch.cuda.set_device(device)
197 + torch.accelerator.set_device_index(device)
198 assert implementation in ["cuda", "triton"]
199 if implementation == "triton" and kv_cache_layout == "HND":
200 pytest.skip("Triton implementation only supports NHD layout.")
@@ -553,7 +555,7 @@ def test_concat_and_cache_mla(
555 ) -> None:
556 set_random_seed(seed)
557 torch.set_default_device(device)
556 - torch.cuda.set_device(device)
558 + torch.accelerator.set_device_index(device)
559
560 total_slots = num_blocks * block_size
561 slot_mapping_lst = random.sample(range(total_slots), num_tokens)
@@ -632,7 +634,7 @@ def test_concat_and_cache_ds_mla(
634 kv_cache_dtype = "fp8_ds_mla"
635 set_random_seed(seed)
636 torch.set_default_device(device)
635 - torch.cuda.set_device(device)
637 + torch.accelerator.set_device_index(device)
638
639 total_slots = num_blocks * block_size
640 slot_mapping_lst = random.sample(range(total_slots), num_tokens)
@@ -744,7 +746,7 @@ def test_swap_blocks_mla(
746 ) -> None:
747 set_random_seed(seed)
748 torch.set_default_device(device)
747 - torch.cuda.set_device(device)
749 + torch.accelerator.set_device_index(device)
750
751 entry_size = kv_lora_rank + qk_rope_head_dim
752
tests/kernels/attention/test_cutlass_mla_decode.py
+1 -1
@@ -69,7 +69,7 @@ def test_cutlass_mla_decode(
69 init_dtype = torch.bfloat16 if torch_dtype == torch.float8_e4m3fn else torch_dtype
70 torch.set_default_dtype(init_dtype)
71 torch.set_default_device(device)
72 - torch.cuda.set_device(device)
72 + torch.accelerator.set_device_index(device)
73 torch.manual_seed(42)
74 random.seed(42)
75
tests/kernels/attention/test_flashmla.py
+1 -1
@@ -57,7 +57,7 @@ def test_flash_mla(
57 init_dtype = torch.bfloat16 if torch_dtype == torch.float8_e4m3fn else torch_dtype
58 torch.set_default_dtype(init_dtype)
59 torch.set_default_device(device)
60 - torch.cuda.set_device(device)
60 + torch.accelerator.set_device_index(device)
61 torch.manual_seed(0)
62 random.seed(0)
63
tests/kernels/attention/test_prefix_prefill.py
+5 -3
@@ -21,7 +21,9 @@ NUM_HEADS = [64]
21 NUM_QUERIES_PER_KV = [1, 64]
22 HEAD_SIZES = [24, 128]
23 DTYPES = [torch.float16]
24 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
24 +CUDA_DEVICES = [
25 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
26 +]
27 SLIDING_WINDOW = [0, 16, 2048]
28 KV_CACHE_DTYPES = ["auto", "fp8", "fp8_e5m2"]
29
@@ -135,7 +137,7 @@ def test_contexted_kv_attention(
137 # for GPU 1 would run on both GPU0 and GPU1 and things would hang
138 #
139 # see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523
138 - torch.cuda.set_device(device)
140 + torch.accelerator.set_device_index(device)
141
142 MAX_SEQ_LEN = 1024
143 MAX_CTX_LEN = 1024
@@ -356,7 +358,7 @@ def test_contexted_kv_attention_alibi(
358 # for GPU 1 would run on both GPU0 and GPU1 and things would hang
359 #
360 # see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523
359 - torch.cuda.set_device(device)
361 + torch.accelerator.set_device_index(device)
362
363 def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
364 # Fork from: vllm/vllm/model_executor/models/bloom.py#L44
tests/kernels/core/test_activation.py
+3 -1
@@ -26,7 +26,9 @@ DTYPES = [torch.half, torch.bfloat16, torch.float]
26 NUM_TOKENS = [7, 83, 2048] # Arbitrary values for testing
27 D = [512, 13824] # Arbitrary values for testing
28 SEEDS = [0]
29 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
29 +CUDA_DEVICES = [
30 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
31 +]
32
33
34 @pytest.mark.parametrize(
tests/kernels/core/test_fused_quant_layernorm.py
+4 -2
@@ -33,7 +33,9 @@ SCALE_UBS = [True, False]
33 GROUP_SIZES = [None, [1, 64], [1, 128]]
34 TMA_ALIGNMENTS = [0, 4]
35 SEEDS = [0]
36 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
36 +CUDA_DEVICES = [
37 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
38 +]
39
40 EPS = 1e-6
41
@@ -182,7 +184,7 @@ def test_rms_norm(
184 if torch.cuda.is_available():
185 torch.cuda.manual_seed(seed)
186 torch.set_default_device(device)
185 - torch.cuda.set_device(device)
187 + torch.accelerator.set_device_index(device)
188
189 if group_size is not None and hidden_size % group_size[1] != 0:
190 # skip
tests/kernels/core/test_layernorm.py
+3 -1
@@ -14,7 +14,9 @@ NUM_TOKENS = [7, 83, 4096] # Arbitrary values for testing
14 HIDDEN_SIZES = [8, 768, 769, 5120, 5125, 8192] # Arbitrary values for testing
15 ADD_RESIDUAL = [False, True]
16 SEEDS = [0]
17 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
17 +CUDA_DEVICES = [
18 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
19 +]
20
21
22 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
tests/kernels/core/test_pos_encoding.py
+3 -1
@@ -19,7 +19,9 @@ NUM_HEADS = [17] # Arbitrary values for testing
19 BATCH_SIZES = [5] # Arbitrary values for testing
20 SEQ_LENS = [11, 8192] # Arbitrary values for testing
21 SEEDS = [0]
22 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
22 +CUDA_DEVICES = [
23 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
24 +]
25 USE_KEY = [True, False]
26
27
tests/kernels/core/test_rotary_embedding_mla_cache_fused.py
+2 -1
@@ -28,7 +28,8 @@ from vllm.utils.torch_utils import set_random_seed
28 @pytest.mark.parametrize("block_size", [16, 64, 256])
29 @pytest.mark.parametrize("seed", [0])
30 @pytest.mark.parametrize(
31 - "device", [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
31 + "device",
32 + [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)],
33 )
34 @torch.inference_mode()
35 def test_concat_and_cache_mla_rope_fused(
tests/kernels/core/test_uva.py
+3 -1
@@ -6,7 +6,9 @@ import torch
6 from vllm.utils.platform_utils import is_uva_available
7 from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
8
9 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
9 +CUDA_DEVICES = [
10 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
11 +]
12
13
14 @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
tests/kernels/mamba/test_mamba_mixer2.py
+1 -1
@@ -71,7 +71,7 @@ def mixer2_gated_norm_tensor_parallel(
71 set_random_seed(0)
72
73 device = torch.device(f"cuda:{local_rank}")
74 - torch.cuda.set_device(device)
74 + torch.accelerator.set_device_index(device)
75 torch.set_default_device(device)
76 torch.set_default_dtype(dtype)
77
tests/kernels/moe/modular_kernel_tools/common.py
+9 -7
@@ -322,7 +322,7 @@ class WeightTensors:
322 )
323
324 def to_current_device(self):
325 - device = torch.cuda.current_device()
325 + device = torch.accelerator.current_device_index()
326 self.w1 = self.w1.to(device=device)
327 self.w2 = self.w2.to(device=device)
328
@@ -392,7 +392,8 @@ class RankTensors:
392 Return hidden_states
393 """
394 m, k, dtype = (config.M, config.K, config.dtype)
395 - a = torch.randn((m, k), device=torch.cuda.current_device(), dtype=dtype) / 15.0
395 + device = torch.accelerator.current_device_index()
396 + a = torch.randn((m, k), device=device, dtype=dtype) / 15.0
397
398 if config.quant_dtype is None:
399 return a, None
@@ -428,9 +429,10 @@ class RankTensors:
429 topk_weights, topk_ids, _ = fused_topk(hidden_states, score, topk, False)
430
431 # distribute topk_ids evenly
432 + device = torch.accelerator.current_device_index()
433 for mi in range(m):
434 topk_ids[mi] = torch.randperm(config.E)[:topk]
433 - topk_ids = topk_ids.to(device=torch.cuda.current_device())
435 + topk_ids = topk_ids.to(device=device)
436
437 expert_map = None
438 if config.world_size > 1 and config.supports_expert_map():
@@ -440,9 +442,7 @@ class RankTensors:
442 s = pgi.rank * num_local_experts
443 e = s + num_local_experts
444 expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
443 - expert_map = expert_map.to(
444 - device=torch.cuda.current_device(), dtype=torch.int32
445 - )
445 + expert_map = expert_map.to(device=device, dtype=torch.int32)
446
447 return RankTensors(
448 hidden_states=hidden_states,
@@ -558,7 +558,9 @@ def reference_moe_impl(
558
559 def _make_gscale(num_experts: int) -> torch.Tensor:
560 return torch.ones(
561 - (num_experts,), device=torch.cuda.current_device(), dtype=torch.float32
561 + (num_experts,),
562 + device=torch.accelerator.current_device_index(),
563 + dtype=torch.float32,
564 )
565
566
tests/kernels/moe/modular_kernel_tools/parallel_utils.py
+1 -1
@@ -66,7 +66,7 @@ def _worker_parallel_launch(
66 **kwargs: P.kwargs,
67 ) -> None:
68 rank = node_rank * world_local_size + local_rank
69 - torch.cuda.set_device(local_rank)
69 + torch.accelerator.set_device_index(local_rank)
70 device = torch.device("cuda", local_rank)
71 torch.distributed.init_process_group(
72 backend="cpu:gloo,cuda:nccl",
tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py
+2 -1
@@ -34,7 +34,8 @@ def do_profile(
34 record_shapes=True,
35 ) as tprof:
36 fn(**fn_kwargs)
37 - torch.accelerator.synchronize(torch.cuda.current_device())
37 + device = torch.accelerator.current_device_index()
38 + torch.accelerator.synchronize(device=device)
39
40 # TODO (varun): Add a descriptive trace file name
41 tprof.export_chrome_trace(
tests/kernels/moe/parallel_utils.py
+1 -1
@@ -52,7 +52,7 @@ def _worker_parallel_launch(
52 **kwargs: P.kwargs,
53 ) -> None:
54 rank = node_rank * world_local_size + local_rank
55 - torch.cuda.set_device(local_rank)
55 + torch.accelerator.set_device_index(local_rank)
56 device = torch.device("cuda", local_rank)
57 torch.distributed.init_process_group(
58 backend="cpu:gloo,cuda:nccl",
tests/kernels/moe/test_deepep_deepgemm_moe.py
+13 -11
@@ -134,10 +134,8 @@ class TestTensors:
134
135 fp8_info = torch.finfo(torch.float8_e4m3fn)
136 fp8_max, fp8_min = fp8_info.max, fp8_info.min
137 -
138 - rank_tokens = (
139 - torch.randn((m, k), device=torch.cuda.current_device(), dtype=dtype) / 10.0
140 - )
137 + device = torch.accelerator.current_device_index()
138 + rank_tokens = torch.randn((m, k), device=device, dtype=dtype) / 10.0
139 rank_tokens = rank_tokens.clamp(min=fp8_min, max=fp8_max)
140 rank_token_scales = None
141
@@ -145,11 +143,13 @@ class TestTensors:
143 low=0,
144 high=config.num_experts,
145 size=(m, topk),
148 - device=torch.cuda.current_device(),
146 + device=device,
147 ).to(dtype=torch.int64)
148
149 topk_weights = torch.randn(
152 - topk_ids.shape, dtype=torch.float32, device=torch.cuda.current_device()
150 + topk_ids.shape,
151 + dtype=torch.float32,
152 + device=device,
153 )
154
155 return TestTensors(
@@ -296,7 +296,8 @@ def deepep_deepgemm_moe_impl(
296 s = pgi.rank * num_local_experts
297 e = s + num_local_experts
298 expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
299 - return expert_map.to(device=torch.cuda.current_device(), dtype=torch.int32)
299 + device = torch.accelerator.current_device_index()
300 + return expert_map.to(device=device, dtype=torch.int32)
301
302 quant_config = fp8_w8a8_moe_quant_config(
303 w1_scale=w1_scale,
@@ -376,10 +377,11 @@ def _test_deepep_deepgemm_moe(
377
378 set_random_seed(pgi.rank)
379
379 - w1 = w1.to(device=torch.cuda.current_device())
380 - w2 = w2.to(device=torch.cuda.current_device())
381 - w1_scale = w1_scale.to(device=torch.cuda.current_device())
382 - w2_scale = w2_scale.to(device=torch.cuda.current_device())
380 + device = torch.accelerator.current_device_index()
381 + w1 = w1.to(device=device)
382 + w2 = w2.to(device=device)
383 + w1_scale = w1_scale.to(device=device)
384 + w2_scale = w2_scale.to(device=device)
385
386 pg = torch.distributed.new_group(list(range(pgi.world_size)))
387 test_tensors = TestTensors.make(config, pgi.rank)
tests/kernels/moe/test_deepep_moe.py
+8 -9
@@ -210,7 +210,8 @@ def deep_ep_moe_impl(
210 s = pgi.rank * num_local_experts
211 e = s + num_local_experts
212 expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
213 - return expert_map.to(device=torch.cuda.current_device(), dtype=torch.int32)
213 + device = torch.accelerator.current_device_index()
214 + return expert_map.to(device=device, dtype=torch.int32)
215
216 hidden_size = test_tensors.rank_tokens.size(1)
217 is_quantized = w1.dtype == torch.float8_e4m3fn
@@ -365,15 +366,13 @@ def _deep_ep_moe(
366 )
367
368 is_quantized = w1.dtype == torch.float8_e4m3fn
368 - w1 = w1.to(device=torch.cuda.current_device())
369 - w2 = w2.to(device=torch.cuda.current_device())
369 + device_idx = torch.accelerator.current_device_index()
370 + w1 = w1.to(device=device_idx)
371 + w2 = w2.to(device=device_idx)
372 if is_quantized:
371 - w1_scale = w1_scale.to( # type: ignore
372 - device=torch.cuda.current_device()
373 - )
374 - w2_scale = w2_scale.to( # type: ignore
375 - device=torch.cuda.current_device()
376 - )
373 + assert w1_scale is not None and w2_scale is not None
374 + w1_scale = w1_scale.to(device=device_idx)
375 + w2_scale = w2_scale.to(device=device_idx)
376
377 pg = torch.distributed.new_group(list(range(pgi.world_size)))
378 test_tensors = TestTensors.make(config, low_latency_mode)
tests/kernels/moe/test_moe.py
+1 -1
@@ -716,7 +716,7 @@ def test_mixtral_moe(
716 monkeypatch.setenv("MASTER_ADDR", "localhost")
717 monkeypatch.setenv("MASTER_PORT", "12345")
718 init_distributed_environment()
719 - init_workspace_manager(torch.cuda.current_device())
719 + init_workspace_manager(torch.accelerator.current_device_index())
720
721 # Instantiate our and huggingface's MoE blocks
722 vllm_config.compilation_config.static_forward_context = dict()
tests/kernels/moe/test_ocp_mx_moe.py
+2 -2
@@ -71,10 +71,10 @@ def enable_pickle(monkeypatch):
71 )
72 @pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
73 def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
74 - if torch.cuda.device_count() < model_case.tp:
74 + if torch.accelerator.device_count() < model_case.tp:
75 pytest.skip(
76 f"This test requires >={model_case.tp} gpus, got only "
77 - f"{torch.cuda.device_count()}"
77 + f"{torch.accelerator.device_count()}"
78 )
79
80 # `cudagraph_capture_sizes=[16]` to reduce load time.
tests/kernels/quantization/test_cutlass_2of4_sparse.py
+3 -1
@@ -15,7 +15,9 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
15 )
16 from vllm.platforms import current_platform
17
18 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
18 +CUDA_DEVICES = [
19 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
20 +]
21
22 capability = current_platform.get_device_capability()
23 capability = capability[0] * 10 + capability[1]
tests/kernels/quantization/test_cutlass_scaled_mm.py
+3 -1
@@ -40,7 +40,9 @@ MNK_FACTORS = [
40 (512, 24576, 128),
41 ]
42
43 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
43 +CUDA_DEVICES = [
44 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
45 +]
46
47 # -1 means full extent in that dimension
48 TENSORWISE_GROUP_SHAPE = (-1, -1)
tests/kernels/quantization/test_machete_mm.py
+3 -1
@@ -29,7 +29,9 @@ if current_platform.is_rocm():
29 allow_module_level=True,
30 )
31
32 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
32 +CUDA_DEVICES = [
33 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
34 +]
35
36 # TODO: in future PR refactor this and `is_quant_method_supported` in the kernel
37 # unit tests to a common utility function. Currently the use of
tests/kernels/test_cache_kernels.py
+1 -1
@@ -13,7 +13,7 @@ except ImportError:
13 )
14
15
16 -@pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Need CUDA device")
16 +@pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="Need CUDA device")
17 def test_gather_cache_oob():
18 """
19 Tests for OOB read in gather_and_maybe_dequant_cache (Issue #27909).
tests/kernels/test_fused_quant_activation.py
+3 -1
@@ -13,7 +13,9 @@ QUANT_DTYPES = [current_platform.fp8_dtype()]
13 NUM_TOKENS = [1, 17, 86, 1234, 3045] # Arbitrary values for testing
14 HIDDEN_SIZES = [16, 48, 128, 1562, 4096] # Arbitrary values for testing
15 SEEDS = [0]
16 -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
16 +CUDA_DEVICES = [
17 + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
18 +]
19
20
21 def ref_impl(
tests/lora/test_fused_moe_lora_kernel.py
+1 -1
@@ -638,7 +638,7 @@ def use_fused_moe_lora_kernel_tensor_parallel(
638 set_random_seed(seed)
639
640 device = torch.device(f"cuda:{local_rank}")
641 - torch.cuda.set_device(device)
641 + torch.accelerator.set_device_index(device)
642 torch.set_default_device(device)
643 torch.set_default_dtype(dtype)
644
tests/lora/test_layers.py
+8 -8
@@ -61,7 +61,7 @@ pytestmark = pytest.mark.skipif(
61 )
62
63 DEVICES = (
64 - [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
64 + [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
65 if current_platform.is_cuda_alike()
66 else ["cpu"]
67 )
@@ -260,7 +260,7 @@ def test_embeddings(
260 # device, see: https://github.com/triton-lang/triton/issues/2925
261 # Same below.
262 if current_platform.is_cuda_alike():
263 - torch.cuda.set_device(device)
263 + torch.accelerator.set_device_index(device)
264
265 torch.set_default_device(device)
266 max_loras = 8
@@ -359,7 +359,7 @@ def test_lm_head_logits_processor(
359 default_vllm_config, dist_init, num_loras, device, vocab_size, stage
360 ) -> None:
361 if current_platform.is_cuda_alike():
362 - torch.cuda.set_device(device)
362 + torch.accelerator.set_device_index(device)
363
364 torch.set_default_device(device)
365 max_loras = 8
@@ -476,7 +476,7 @@ def test_lm_head_logits_processor_invalid_vocab_size(
476 ) -> None:
477 """Test that LogitsProcessorWithLoRA raises ValueError for invalid vocab sizes."""
478 if current_platform.is_cuda_alike():
479 - torch.cuda.set_device(device)
479 + torch.accelerator.set_device_index(device)
480
481 torch.set_default_device(device)
482 max_loras = 8
@@ -505,7 +505,7 @@ def test_linear_replicated(
505 stage,
506 ) -> None:
507 if current_platform.is_cuda_alike():
508 - torch.cuda.set_device(device)
508 + torch.accelerator.set_device_index(device)
509
510 max_loras = 8
511 torch.set_default_device(device)
@@ -612,7 +612,7 @@ def test_linear_parallel(
612 default_vllm_config, dist_init, num_loras, orientation, fully_shard, device, stage
613 ) -> None:
614 if current_platform.is_cuda_alike():
615 - torch.cuda.set_device(device)
615 + torch.accelerator.set_device_index(device)
616
617 max_loras = 8
618 torch.set_default_device(device)
@@ -737,7 +737,7 @@ def test_column_parallel_packed(
737 default_vllm_config, dist_init, num_loras, repeats, fully_shard, device, stage
738 ) -> None:
739 if current_platform.is_cuda_alike():
740 - torch.cuda.set_device(device)
740 + torch.accelerator.set_device_index(device)
741
742 max_loras = 8
743 torch.set_default_device(device)
@@ -885,7 +885,7 @@ def test_merged_column_parallel_variable_slice(
885 default_vllm_config, dist_init, num_loras, num_slices, device, stage
886 ) -> None:
887 if current_platform.is_cuda_alike():
888 - torch.cuda.set_device(device)
888 + torch.accelerator.set_device_index(device)
889
890 max_loras = 8
891 torch.set_default_device(device)
tests/lora/test_lora_manager.py
+1 -1
@@ -37,7 +37,7 @@ EMBEDDING_MODULES = {
37
38
39 DEVICES = (
40 - [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
40 + [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
41 if current_platform.is_cuda_alike()
42 else ["cpu"]
43 )
tests/lora/test_mixtral.py
+1 -1
@@ -34,7 +34,7 @@ def do_sample(
34 def test_mixtral_lora(mixtral_lora_files, tp_size):
35 """Original test, the LoRA model has the common target modules, not all"""
36 if (
37 - torch.cuda.device_count() < tp_size
37 + torch.accelerator.device_count() < tp_size
38 and tp_size > 1
39 and current_platform.is_cuda_alike()
40 ):
tests/lora/test_punica_ops.py
+2 -2
@@ -395,7 +395,7 @@ def test_kernels(
395 Tests LoRA kernels.
396 """
397 torch.set_default_device(device)
398 - torch.cuda.set_device(device)
398 + torch.accelerator.set_device_index(device)
399 set_random_seed(seed)
400
401 if op_type == "shrink":
@@ -448,7 +448,7 @@ def test_kernels_hidden_size(
448 Tests SGMV and LoRA kernels.
449 """
450 torch.set_default_device(device)
451 - torch.cuda.set_device(device)
451 + torch.accelerator.set_device_index(device)
452 set_random_seed(seed)
453
454 if op_type == "shrink":
tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py
+2 -2
@@ -203,7 +203,7 @@ def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd, model_ref)
203 torch.accelerator.empty_cache()
204
205
206 -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires 2 GPUs")
206 +@pytest.mark.skipif(torch.accelerator.device_count() < 2, reason="Requires 2 GPUs")
207 def test_tensorizer_with_tp_path_without_template(vllm_runner, capfd):
208 try:
209 model_ref = "EleutherAI/pythia-1.4b"
@@ -231,7 +231,7 @@ def test_tensorizer_with_tp_path_without_template(vllm_runner, capfd):
231 ) in combined_output
232
233
234 -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires 2 GPUs")
234 +@pytest.mark.skipif(torch.accelerator.device_count() < 2, reason="Requires 2 GPUs")
235 def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
236 vllm_runner, tmp_path
237 ):
tests/model_executor/test_eagle_quantization.py
+2 -2
@@ -11,7 +11,7 @@ from vllm.model_executor.models.utils import get_draft_quant_config
11 from vllm.platforms import current_platform
12
13 DEVICES = (
14 - [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
14 + [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
15 if current_platform.is_cuda_alike()
16 else ["cpu"]
17 )
@@ -61,7 +61,7 @@ def test_fc_layer_quant_config_usage(default_vllm_config, dist_init, device) ->
61 from vllm.model_executor.layers.linear import ReplicatedLinear
62
63 if current_platform.is_cuda_alike():
64 - torch.cuda.set_device(device)
64 + torch.accelerator.set_device_index(device)
65
66 torch.set_default_device(device)
67
tests/models/test_vision.py
+4 -4
@@ -102,7 +102,7 @@ def run_dp_sharded_vision_model_vs_direct(
102 set_random_seed(0)
103
104 device = f"{current_platform.device_name}:{local_rank}"
105 - current_platform.set_device(device)
105 + torch.accelerator.set_device_index(device)
106 torch.set_default_device(device)
107
108 update_environment_variables(
@@ -288,7 +288,7 @@ def run_dp_sharded_mrope_vision_model_vs_direct(
288 # Set random seed for reproducibility
289 set_random_seed(0)
290 device = f"{current_platform.device_name}:{local_rank}"
291 - current_platform.set_device(device)
291 + torch.accelerator.set_device_index(device)
292 torch.set_default_device(device)
293
294 update_environment_variables(
@@ -365,7 +365,7 @@ def run_dp_sharded_mrope_vision_model_empty_input_worker(
365 """Test run_dp_sharded_mrope_vision_model with empty input."""
366 # Set up distributed environment
367 device = f"{current_platform.device_name}:{local_rank}"
368 - current_platform.set_device(device)
368 + torch.accelerator.set_device_index(device)
369 torch.set_default_device(device)
370
371 update_environment_variables(
@@ -414,7 +414,7 @@ def run_dp_sharded_mrope_vision_model_uneven_load_worker(
414 # Set up distributed environment
415 set_random_seed(123)
416 device = f"{current_platform.device_name}:{local_rank}"
417 - current_platform.set_device(device)
417 + torch.accelerator.set_device_index(device)
418 torch.set_default_device(device)
419
420 update_environment_variables(
tests/quantization/test_quark.py
+6 -8
@@ -210,10 +210,9 @@ WIKITEXT_ACCURACY_CONFIGS = [
210 @pytest.mark.parametrize("config", WIKITEXT_ACCURACY_CONFIGS)
211 @pytest.mark.parametrize("tp_size", [1, 2])
212 def test_ocp_mx_wikitext_correctness(config: AccuracyTestConfig, tp_size: int):
213 - if torch.cuda.device_count() < tp_size:
214 - pytest.skip(
215 - f"This test requires >={tp_size} gpus, got only {torch.cuda.device_count()}"
216 - )
213 + device_count = torch.accelerator.device_count()
214 + if device_count < tp_size:
215 + pytest.skip(f"This test requires >={tp_size} gpus, got only {device_count}")
216
217 task = "wikitext"
218 rtol = 0.1
@@ -246,10 +245,9 @@ def test_ocp_mx_wikitext_correctness(config: AccuracyTestConfig, tp_size: int):
245 reason="Read access to huggingface.co/amd is required for this test.",
246 )
247 def test_mxfp4_gsm8k_correctness(config: AccuracyTestConfig):
249 - if torch.cuda.device_count() < 8:
250 - pytest.skip(
251 - f"This test requires >=8 gpus, got only {torch.cuda.device_count()}"
252 - )
248 + device_count = torch.accelerator.device_count()
249 + if device_count < 8:
250 + pytest.skip(f"This test requires >=8 gpus, got only {device_count}")
251
252 task = "gsm8k"
253 rtol = 0.03
tests/v1/e2e/test_spec_decode.py
+1 -1
@@ -32,7 +32,7 @@ MTP_SIMILARITY_RATE = 0.8
32
33 def _skip_if_insufficient_gpus_for_tp(tp_size: int):
34 """Skip test if available GPUs < tp_size on ROCm."""
35 - available_gpus = torch.cuda.device_count()
35 + available_gpus = torch.accelerator.device_count()
36 if available_gpus < tp_size:
37 pytest.skip(
38 f"Test requires {tp_size} GPUs, but only {available_gpus} available"
tests/v1/kv_connector/unit/test_example_connector.py
+1 -1
@@ -148,7 +148,7 @@ def test_shared_storage_connector_hashes(tmp_path, attn_backend):
148 )
149
150 # don't put this import at the top level
151 - # it will call torch.cuda.device_count()
151 + # it will call torch.accelerator.device_count()
152 from transformers import AutoProcessor
153
154 # Create processor to handle the chat prompt
tests/v1/kv_connector/unit/test_nixl_connector.py
+1 -1
@@ -1570,7 +1570,7 @@ def test_register_kv_caches(
1570 ]
1571 ],
1572 cache_dtype=torch.bfloat16,
1573 - device=torch.cuda.current_device(),
1573 + device=torch.accelerator.current_device_index(),
1574 kernel_block_sizes=[block_size],
1575 )
1576 )
tests/v1/spec_decode/test_acceptance_length.py
+1 -1
@@ -141,7 +141,7 @@ def get_attention_backend_params() -> list[str]:
141
142
143 def get_tp_size_params() -> list[pytest.param]:
144 - num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1
144 + num_gpus = torch.accelerator.device_count() if torch.cuda.is_available() else 1
145 return [pytest.param(tp, id=f"tp{tp}") for tp in TP_SIZES if tp <= num_gpus]
146
147
tests/v1/worker/test_worker_memory_snapshot.py
+2 -1
@@ -117,7 +117,8 @@ def worker_process(
117
118
119 @pytest.mark.skipif(
120 - torch.cuda.device_count() < 2, reason="Need at least 2 GPUs for tensor parallelism"
120 + torch.accelerator.device_count() < 2,
121 + reason="Need at least 2 GPUs for tensor parallelism",
122 )
123 def test_init_distributed_is_called_before_memory_snapshot():
124 """Test that distributed env is setup before memory snapshot.
tools/pre_commit/check_torch_cuda.py
+5 -3
@@ -8,8 +8,8 @@ import regex as re
8 # Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
9 # --------------------------------------------------------------------------- #
10 _TORCH_CUDA_PATTERNS = [
11 - r"\btorch\.cuda\.(empty_cache|synchronize|device\()\b",
12 - r"\bwith\btorch\.cuda\.device\b",
11 + r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|set_device|device\()\b",
12 + r"\bwith\storch\.cuda\.device\b",
13 ]
14
15 ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}
@@ -25,7 +25,9 @@ def scan_file(path: str) -> int:
25 print(
26 f"{path}:{line_num}: "
27 "\033[91merror:\033[0m " # red color
28 - "Found torch.cuda API call"
28 + "Found torch.cuda API call. Please refer RFC "
29 + "https://github.com/vllm-project/vllm/issues/30679, use "
30 + "torch.accelerator API instead."
31 )
32 return 1
33 return 0
vllm/distributed/device_communicators/all2all.py
+1 -1
@@ -491,7 +491,7 @@ class FlashInferAllToAllManager(All2AllManagerBase):
491 self.initialize(
492 world_size=self.world_size,
493 rank=self.rank,
494 - gpus_per_node=torch.cuda.device_count,
494 + gpus_per_node=torch.accelerator.device_count,
495 )
496 return self.initialized
497
vllm/distributed/device_communicators/pynccl_allocator.py
+1 -1
@@ -151,7 +151,7 @@ class nccl_symm_mem_context:
151 self.pynccl_comm = pynccl_comm
152 self._mem_pool_ctx = torch.cuda.use_mem_pool(get_nccl_mem_pool())
153 self.is_graph_capture = torch.cuda.is_current_stream_capturing()
154 - self.device = torch.cuda.current_device()
154 + self.device = torch.accelerator.current_device_index()
155
156 def __enter__(self):
157 if self.disabled:
vllm/distributed/device_communicators/symm_mem.py
+1 -1
@@ -50,7 +50,7 @@ class SymmMemCommunicator:
50 device = torch.device(f"cuda:{device}")
51 elif isinstance(device, str):
52 device = torch.device(device)
53 - torch.cuda.set_device(device)
53 + torch.accelerator.set_device_index(device)
54 self.dtype = torch.bfloat16
55 self.device = device
56 self.group = group
vllm/distributed/eplb/async_worker.py
+1 -1
@@ -33,7 +33,7 @@ def start_async_worker(
33
34 def thread_target() -> None:
35 assert device_index is not None
36 - torch.cuda.set_device(device_index)
36 + torch.accelerator.set_device_index(device_index)
37 cuda_stream = torch.cuda.Stream(device=device_index)
38 loop = asyncio.new_event_loop()
39 asyncio.set_event_loop(loop)
vllm/distributed/eplb/eplb_state.py
+1 -1
@@ -314,7 +314,7 @@ class EplbState:
314 if self.device.type == "cuda":
315 self.cuda_device_index = self.device.index
316 if self.cuda_device_index is None and torch.cuda.is_available():
317 - self.cuda_device_index = torch.cuda.current_device()
317 + self.cuda_device_index = torch.accelerator.current_device_index()
318
319 @staticmethod
320 def build_initial_global_physical_to_logical_map(
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
+2 -2
@@ -483,9 +483,9 @@ def _init_lmcache_engine(
483 )
484
485 # Change current device.
486 - num_gpus = torch.cuda.device_count()
486 + num_gpus = torch.accelerator.device_count()
487 local_rank = parallel_config.rank % num_gpus
488 - torch.cuda.set_device(local_rank)
488 + torch.accelerator.set_device_index(local_rank)
489 device = torch.device(f"cuda:{local_rank}")
490 metadata = LMCacheEngineMetadata(
491 model_config.model,
vllm/distributed/weight_transfer/ipc_engine.py
+2 -2
@@ -169,7 +169,7 @@ class IPCWeightTransferEngine(
169 update_info.shapes,
170 update_info.ipc_handles,
171 ):
172 - device_index = torch.cuda.current_device()
172 + device_index = torch.accelerator.current_device_index()
173 props = torch.cuda.get_device_properties(device_index)
174 physical_gpu_id = str(props.uuid)
175
@@ -242,7 +242,7 @@ class IPCWeightTransferEngine(
242 args = trainer_args
243
244 # Get physical GPU UUID
245 - device_index = torch.cuda.current_device()
245 + device_index = torch.accelerator.current_device_index()
246 props = torch.cuda.get_device_properties(device_index)
247 gpu_uuid = str(props.uuid)
248
vllm/distributed/weight_transfer/nccl_engine.py
+9 -3
@@ -140,13 +140,14 @@ class NCCLWeightTransferEngine(
140 worker_rank = dp_rank * world_size_per_dp + rank_within_dp
141 rank = worker_rank + init_info.rank_offset
142 # Create stateless process group
143 + device = torch.accelerator.current_device_index()
144 self.model_update_group = (
145 NCCLWeightTransferEngine._stateless_init_process_group(
146 init_info.master_address,
147 init_info.master_port,
148 rank,
149 init_info.world_size,
149 - torch.cuda.current_device(),
150 + device=device,
151 )
152 )
153
@@ -275,7 +276,7 @@ class NCCLWeightTransferEngine(
276 Initialize NCCL process group for trainer-side weight transfer.
277
278 The trainer is always rank 0 in the process group. Uses the current
278 - CUDA device (torch.cuda.current_device()).
279 + CUDA device (torch.accelerator.current_device_index()).
280
281 Args:
282 init_info: Either an NCCLWeightTransferInitInfo object or a dict with keys:
@@ -309,8 +310,13 @@ class NCCLWeightTransferEngine(
310 world_size = init_info.world_size
311
312 # Trainer is always rank 0
313 + device = torch.accelerator.current_device_index()
314 return NCCLWeightTransferEngine._stateless_init_process_group(
313 - master_address, master_port, 0, world_size, torch.cuda.current_device()
315 + master_address,
316 + master_port,
317 + 0,
318 + world_size,
319 + device,
320 )
321
322 @staticmethod
vllm/model_executor/layers/attention/static_sink_attention.py
+1 -1
@@ -190,7 +190,7 @@ class StaticSinkAttention(Attention, CustomOp):
190 sink_kv_slot_mapping = torch.arange(
191 self.block_size,
192 self.sink_len + self.block_size,
193 - device=torch.cuda.current_device(),
193 + device=torch.accelerator.current_device_index(),
194 dtype=torch.long,
195 )
196 triton_reshape_and_cache_flash_diffkv(
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+5 -2
@@ -295,14 +295,17 @@ class DefaultMoERunner(MoERunner):
295 states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim)
296 logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts)
297
298 + device = torch.accelerator.current_device_index()
299 self.batched_hidden_states = torch.zeros(
299 - states_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()
300 + states_shape,
301 + dtype=moe.in_dtype,
302 + device=device,
303 )
304
305 self.batched_router_logits = torch.zeros(
306 logits_shape,
307 dtype=moe.router_logits_dtype,
305 - device=torch.cuda.current_device(),
308 + device=device,
309 )
310
311 def must_reduce_shared_expert_outputs(self) -> bool:
vllm/model_executor/layers/fused_moe/trtllm_moe.py
+1 -1
@@ -28,7 +28,7 @@ class TrtLlmGenExperts(mk.FusedMoEExpertsModular):
28 max_capture_size,
29 ):
30 super().__init__(moe_config, quant_config)
31 - self.device = torch.cuda.current_device()
31 + self.device = torch.accelerator.current_device_index()
32 self.num_experts = moe_config.num_local_experts
33 self.gemm1_alpha = torch.tensor(
34 [1.702] * self.num_experts, dtype=torch.float32, device=self.device
vllm/model_executor/layers/layernorm.py
+1 -1
@@ -202,7 +202,7 @@ class RMSNorm(CustomOp):
202 # external Oink initialization work in this case.
203 else:
204 try:
205 - device_index = torch.cuda.current_device()
205 + device_index = torch.accelerator.current_device_index()
206 if _oink_ops.is_oink_available_for_device(device_index):
207 self._use_oink_rmsnorm = True
208 self._use_oink_fused_add_rmsnorm = (
vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py
+2 -1
@@ -36,7 +36,8 @@ class DualChunkRotaryEmbedding(CustomOp):
36 self.chunk_size = chunk_size
37 self.local_size = local_size
38 self.dtype = dtype
39 - self.device = torch.device(f"cuda:{torch.cuda.current_device()}")
39 + device_idx = torch.accelerator.current_device_index()
40 + self.device = torch.device(f"cuda:{device_idx}")
41 (q_cache, qc_cache, k_cache, qc_no_clamp_cache, q_inter_cache) = (
42 self._compute_cos_sin_cache()
43 )
vllm/model_executor/model_loader/tensorizer.py
+3 -3
@@ -539,6 +539,8 @@ def deserialize_tensorizer_model(
539 )
540 before_mem = get_mem_usage()
541 start = time.perf_counter()
542 + device_index = torch.accelerator.current_device_index()
543 + device_type = current_platform.device_type
544 with (
545 open_stream(
546 tensorizer_config.tensorizer_uri, mode="rb", **tensorizer_args.stream_kwargs
@@ -546,9 +548,7 @@ def deserialize_tensorizer_model(
548 TensorDeserializer(
549 stream,
550 dtype=tensorizer_config.dtype,
549 - device=f"xpu:{torch.xpu.current_device()}"
550 - if current_platform.is_xpu()
551 - else f"cuda:{torch.cuda.current_device()}",
551 + device=f"{device_type}:{device_index}",
552 **tensorizer_args.deserialization_kwargs,
553 ) as deserializer,
554 ):
vllm/utils/torch_utils.py
+1 -1
@@ -624,7 +624,7 @@ 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.cuda.device_count()
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
vllm/v1/engine/utils.py
+1 -1
@@ -134,7 +134,7 @@ class CoreEngineProcManager:
134 for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
135 # Adjust device control in DP for non-CUDA platforms
136 # as well as external and ray launchers
137 - # For CUDA platforms, we use torch.cuda.set_device()
137 + # For CUDA platforms, we use torch.accelerator.set_device_index()()
138 if is_dp and (
139 not current_platform.is_cuda_alike()
140 or vllm_config.parallel_config.use_ray
vllm/v1/worker/gpu_ubatch_wrapper.py
+3 -3
@@ -73,8 +73,8 @@ class SMControlContextManager:
73 assert current_platform.is_cuda(), (
74 "SM control is currently only supported on CUDA"
75 )
76 -
77 - total_sms = num_compute_units(torch.cuda.current_device())
76 + device = torch.accelerator.current_device_index()
77 + total_sms = num_compute_units(device)
78
79 assert comm_sms < total_sms
80 self.total_sms = total_sms
@@ -204,7 +204,7 @@ class UBatchWrapper:
204
205 @torch.inference_mode()
206 def _capture_ubatch_thread(results, ubatch_metadata):
207 - torch.cuda.set_device(self.device)
207 + torch.accelerator.set_device_index(self.device)
208 ubatch_context = ubatch_metadata.context
209 with torch.cuda.stream(ubatch_context.compute_stream):
210 _ = torch.cuda.current_blas_handle()
vllm/v1/worker/gpu_worker.py
+3 -3
@@ -239,11 +239,11 @@ class Worker(WorkerBase):
239
240 # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
241 self.local_rank += dp_local_rank * tp_pp_world_size
242 - assert self.local_rank < torch.cuda.device_count(), (
242 + assert self.local_rank < torch.accelerator.device_count(), (
243 f"DP adjusted local rank {self.local_rank} is out of bounds. "
244 )
245 visible_device_count = (
246 - torch.cuda.device_count() if torch.cuda.is_available() else 0
246 + torch.accelerator.device_count() if torch.cuda.is_available() else 0
247 )
248 assert self.parallel_config.local_world_size <= visible_device_count, (
249 f"local_world_size ({self.parallel_config.local_world_size}) must "
@@ -252,7 +252,7 @@ class Worker(WorkerBase):
252 )
253
254 self.device = torch.device(f"cuda:{self.local_rank}")
255 - current_platform.set_device(self.device)
255 + torch.accelerator.set_device_index(self.device)
256
257 current_platform.check_if_supports_dtype(self.model_config.dtype)
258
vllm/v1/worker/xpu_worker.py
+1 -1
@@ -60,7 +60,7 @@ class XPUWorker(Worker):
60 and current_platform.is_xpu()
61 ):
62 self.device = torch.device(f"xpu:{self.local_rank}")
63 - current_platform.set_device(self.device)
63 + torch.accelerator.set_device_index(self.device)
64 current_platform.check_if_supports_dtype(self.model_config.dtype)
65 torch.accelerator.empty_cache()
66 self.init_gpu_memory = torch.xpu.get_device_properties(