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

Signed-off-by: Yan Ma <yan.ma@intel.com>

vllm/model_executor/layers/fla/ops/utils.py

@@ -105,7 +105,7 @@ def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:
                     break
 
         if tensor is not None:
-            ctx = torch.cuda.device(tensor.device.index)
+            ctx = torch.accelerator.device_index(tensor.device.index)
         else:
             ctx = contextlib.nullcontext()
 

vllm/model_executor/layers/mamba/ops/layernorm_gated.py

@@ -119,7 +119,7 @@ def _layer_norm_fwd(
     # heuristics for number of warps
     num_warps = min(max(BLOCK_N // 256, 1), 8)
     grid = (M, ngroups)
-    with torch.cuda.device(x.device.index):
+    with torch.accelerator.device_index(x.device.index):
         _layer_norm_fwd_1pass_kernel[grid](
             x,
             out,

vllm/model_executor/layers/mamba/ops/mamba_ssm.py

@@ -419,7 +419,7 @@ def selective_state_update(
         and dt.stride(-1) == 0
         and dt_bias.stride(-1) == 0
     )
-    with torch.cuda.device(x.device.index):
+    with torch.accelerator.device_index(x.device.index):
         _selective_scan_update_kernel[grid](
             state,
             x,

vllm/model_executor/layers/mamba/ops/ssd_bmm.py

@@ -185,7 +185,7 @@ def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtyp
         * triton.cdiv(chunk_size, META["BLOCK_SIZE_N"]),
         nchunks * ngroups,
     )
-    with torch.cuda.device(a.device.index):
+    with torch.accelerator.device_index(a.device.index):
         _bmm_chunk_fwd_kernel[grid](
             a_ptr=a,
             b_ptr=b,

vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py

@@ -323,7 +323,7 @@ def _chunk_cumsum_fwd(
         nheads, nchunks, chunk_size, device=dt.device, dtype=torch.float32
     )
     grid_chunk_cs = lambda META: (nchunks, triton.cdiv(nheads, META["BLOCK_SIZE_H"]))
-    with torch.cuda.device(dt.device.index):
+    with torch.accelerator.device_index(dt.device.index):
         _chunk_cumsum_fwd_kernel[grid_chunk_cs](
             dt_ptr=dt,
             A_ptr=A,
@@ -378,7 +378,7 @@ def _chunk_state_fwd(
         nchunks,
         nheads,
     )
-    with torch.cuda.device(x.device.index):
+    with torch.accelerator.device_index(x.device.index):
         _chunk_state_fwd_kernel[grid](
             x_ptr=x,
             b_ptr=B,

vllm/model_executor/layers/mamba/ops/ssd_state_passing.py

@@ -120,7 +120,7 @@ def _state_passing_fwd(
     )
 
     grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE"]), batch, nheads)
-    with torch.cuda.device(states.device.index):
+    with torch.accelerator.device_index(states.device.index):
         _state_passing_fwd_kernel[grid](
             states_ptr=states,
             out_ptr=out,