biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

replace with torch.cuda.device with with torch.accelerator.device_index (#36144)

Browse Source 
Signed-off-by: Yan Ma <yan.ma@intel.com>
This commit is contained in:
Yan Ma
2026-03-12 14:12:57 +08:00
committed by GitHub
parent 584a3f56de
commit 894843eb25
10 changed files with 17 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/distributed/device_communicators/pynccl.py
View File

@@ -133,9 +133,7 @@ class PyNcclCommunicator:
assert isinstance(device, torch.device)
self.device = device
# nccl communicator and stream will use this device
# `torch.cuda.device` is a context manager that changes the
# current cuda device to the specified one
with torch.cuda.device(device):
with torch.accelerator.device_index(device.index):
self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
self.world_size, self.unique_id, self.rank
)

4
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py
View File

@@ -218,7 +218,7 @@ class P2pNcclEngine:
data = {"cmd": "NEW", "unique_id": bytes(unique_id.internal)}
sock.send(msgpack.dumps(data))
with torch.cuda.device(self.device):
with torch.accelerator.device_index(self.device.index):
rank = 0
with set_p2p_nccl_context(self.nccl_num_channels):
comm: ncclComm_t = self.nccl.ncclCommInitRank(2, unique_id, rank)
@@ -377,7 +377,7 @@ class P2pNcclEngine:
data = msgpack.loads(message)
if data["cmd"] == "NEW":
unique_id = self.nccl.unique_id_from_bytes(bytes(data["unique_id"]))
with torch.cuda.device(self.device):
with torch.accelerator.device_index(self.device.index):
rank = 1
with set_p2p_nccl_context(self.nccl_num_channels):
comm: ncclComm_t = self.nccl.ncclCommInitRank(