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, 20:30 UTC
replace `with torch.cuda.device` with `with torch.accelerator.device_index` (#36144)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Yan Ma committed
Mar 12, 2026 at 14:12 UTC
894843eb25ddbdedec93b68140f2eb14fceea7ce
10 files changed
+17
-15
benchmarks/kernels/benchmark_moe.py
+5
-1
@@ -626,7 +626,11 @@ class BenchmarkWorker:
626
if visible_device != f"{self.device_id}":
627
need_device_guard = True
628
629
- with torch.cuda.device(self.device_id) if need_device_guard else nullcontext():
629
+ with (
630
+ torch.accelerator.device_index(self.device_id)
631
+ if need_device_guard
632
+ else nullcontext()
633
+ ):
634
for idx, config in enumerate(tqdm(search_space)):
635
try:
636
kernel_time = benchmark_config(
tools/pre_commit/check_torch_cuda.py
+2
-2
@@ -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\b",
12
- r"\btorch\.cuda\.synchronize\b",
11
+ r"\btorch\.cuda\.(empty_cache|synchronize|device\()\b",
12
+ r"\bwith\btorch\.cuda\.device\b",
13
]
14
15
ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}
vllm/distributed/device_communicators/pynccl.py
+1
-3
@@ -133,9 +133,7 @@ class PyNcclCommunicator:
133
assert isinstance(device, torch.device)
134
self.device = device
135
# nccl communicator and stream will use this device
136
- # `torch.cuda.device` is a context manager that changes the
137
- # current cuda device to the specified one
138
- with torch.cuda.device(device):
136
+ with torch.accelerator.device_index(device.index):
137
self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
138
self.world_size, self.unique_id, self.rank
139
)
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py
+2
-2
@@ -218,7 +218,7 @@ class P2pNcclEngine:
218
data = {"cmd": "NEW", "unique_id": bytes(unique_id.internal)}
219
sock.send(msgpack.dumps(data))
220
221
- with torch.cuda.device(self.device):
221
+ with torch.accelerator.device_index(self.device.index):
222
rank = 0
223
with set_p2p_nccl_context(self.nccl_num_channels):
224
comm: ncclComm_t = self.nccl.ncclCommInitRank(2, unique_id, rank)
@@ -377,7 +377,7 @@ class P2pNcclEngine:
377
data = msgpack.loads(message)
378
if data["cmd"] == "NEW":
379
unique_id = self.nccl.unique_id_from_bytes(bytes(data["unique_id"]))
380
- with torch.cuda.device(self.device):
380
+ with torch.accelerator.device_index(self.device.index):
381
rank = 1
382
with set_p2p_nccl_context(self.nccl_num_channels):
383
comm: ncclComm_t = self.nccl.ncclCommInitRank(
vllm/model_executor/layers/fla/ops/utils.py
+1
-1
@@ -105,7 +105,7 @@ def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:
105
break
106
107
if tensor is not None:
108
- ctx = torch.cuda.device(tensor.device.index)
108
+ ctx = torch.accelerator.device_index(tensor.device.index)
109
else:
110
ctx = contextlib.nullcontext()
111
vllm/model_executor/layers/mamba/ops/layernorm_gated.py
+1
-1
@@ -119,7 +119,7 @@ def _layer_norm_fwd(
119
# heuristics for number of warps
120
num_warps = min(max(BLOCK_N // 256, 1), 8)
121
grid = (M, ngroups)
122
- with torch.cuda.device(x.device.index):
122
+ with torch.accelerator.device_index(x.device.index):
123
_layer_norm_fwd_1pass_kernel[grid](
124
x,
125
out,
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+1
-1
@@ -419,7 +419,7 @@ def selective_state_update(
419
and dt.stride(-1) == 0
420
and dt_bias.stride(-1) == 0
421
)
422
- with torch.cuda.device(x.device.index):
422
+ with torch.accelerator.device_index(x.device.index):
423
_selective_scan_update_kernel[grid](
424
state,
425
x,
vllm/model_executor/layers/mamba/ops/ssd_bmm.py
+1
-1
@@ -185,7 +185,7 @@ def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtyp
185
* triton.cdiv(chunk_size, META["BLOCK_SIZE_N"]),
186
nchunks * ngroups,
187
)
188
- with torch.cuda.device(a.device.index):
188
+ with torch.accelerator.device_index(a.device.index):
189
_bmm_chunk_fwd_kernel[grid](
190
a_ptr=a,
191
b_ptr=b,
vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py
+2
-2
@@ -323,7 +323,7 @@ def _chunk_cumsum_fwd(
323
nheads, nchunks, chunk_size, device=dt.device, dtype=torch.float32
324
)
325
grid_chunk_cs = lambda META: (nchunks, triton.cdiv(nheads, META["BLOCK_SIZE_H"]))
326
- with torch.cuda.device(dt.device.index):
326
+ with torch.accelerator.device_index(dt.device.index):
327
_chunk_cumsum_fwd_kernel[grid_chunk_cs](
328
dt_ptr=dt,
329
A_ptr=A,
@@ -378,7 +378,7 @@ def _chunk_state_fwd(
378
nchunks,
379
nheads,
380
)
381
- with torch.cuda.device(x.device.index):
381
+ with torch.accelerator.device_index(x.device.index):
382
_chunk_state_fwd_kernel[grid](
383
x_ptr=x,
384
b_ptr=B,
vllm/model_executor/layers/mamba/ops/ssd_state_passing.py
+1
-1
@@ -120,7 +120,7 @@ def _state_passing_fwd(
120
)
121
122
grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE"]), batch, nheads)
123
- with torch.cuda.device(states.device.index):
123
+ with torch.accelerator.device_index(states.device.index):
124
_state_passing_fwd_kernel[grid](
125
states_ptr=states,
126
out_ptr=out,