Use device-agnostic helpers for Mamba tests and core ops (#20234)

Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

python/sglang/srt/layers/attention/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.get_device_module(x.device).device(x.device.index):
         _layer_norm_fwd_1pass_kernel[grid](
             x,
             out,

python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py

@@ -427,7 +427,7 @@ def selective_state_update(
         else (0, 0)
     )
 
-    with torch.cuda.device(x.device.index):
+    with torch.get_device_module(x.device).device(x.device.index):
         _selective_scan_update_kernel[grid](
             state,
             x,

python/sglang/srt/layers/attention/mamba/ops/ssd_bmm.py

@@ -179,7 +179,7 @@ def _bmm_chunk_fwd(a, b, chunk_size, seq_idx=None, causal=False, output_dtype=No
         batch,
         nchunks if not has_groups else nchunks * ngroups,
     )
-    with torch.cuda.device(a.device.index):
+    with torch.get_device_module(a.device).device(a.device.index):
         _bmm_chunk_fwd_kernel[grid](
             a,
             b,

python/sglang/srt/layers/attention/mamba/ops/ssd_chunk_state.py

@@ -460,7 +460,7 @@ def _chunk_cumsum_fwd(
         nchunks,
         triton.cdiv(nheads, META["BLOCK_SIZE_H"]),
     )
-    with torch.cuda.device(dt.device.index):
+    with torch.get_device_module(dt.device).device(dt.device.index):
         _chunk_cumsum_fwd_kernel[grid_chunk_cs](
             dt,
             A,
@@ -520,7 +520,7 @@ def _chunk_state_fwd(
         batch * nchunks,
         nheads,
     )
-    with torch.cuda.device(x.device.index):
+    with torch.get_device_module(x.device).device(x.device.index):
         _chunk_state_fwd_kernel[grid](
             x,
             B,
@@ -596,7 +596,7 @@ def chunk_state_varlen(
         batch,
         nheads,
     )
-    with torch.cuda.device(x.device.index):
+    with torch.get_device_module(x.device).device(x.device.index):
         _chunk_state_varlen_kernel[grid](
             x,
             B,

python/sglang/srt/layers/attention/mamba/ops/ssd_state_passing.py

@@ -214,7 +214,7 @@ def _state_passing_fwd(
         (batch, nheads, dim), device=states.device, dtype=torch.float32
     )
     grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE"]), batch, nheads)
-    with torch.cuda.device(states.device.index):
+    with torch.get_device_module(states.device).device(states.device.index):
         _state_passing_fwd_kernel[grid](
             states,
             out,