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,

test/registered/layers/mamba/test_mamba2_mixer.py

@@ -13,6 +13,7 @@ from sglang.srt.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
+from sglang.srt.utils import get_device, get_device_count
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=32, suite="stage-b-test-2-gpu-large")
@@ -35,14 +36,14 @@ def test_mixer2_gated_norm_multi_gpu(
     seq_len: int,
     hidden_size_n_groups: tuple[int, int],
     dtype: torch.dtype,
-    device: str = "cuda",
+    device: str = get_device(),
 ):
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA device not available")
+    if device not in ["cuda", "xpu"]:
+        pytest.skip("Test only supports CUDA and XPU devices")
 
     assert (
-        torch.cuda.device_count() >= NUM_GPUS
-    ), f"This test requires at least {NUM_GPUS} GPUs, but only {torch.cuda.device_count()} available"
+        get_device_count() >= NUM_GPUS
+    ), f"This test requires at least {NUM_GPUS} GPUs, but only {get_device_count()} available"
 
     hidden_size, n_groups = hidden_size_n_groups
     num_processes = NUM_GPUS
@@ -79,8 +80,8 @@ def mixer2_gated_norm_tensor_parallel(
 ):
     torch.manual_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
-    torch.cuda.set_device(device)
+    device = torch.device(get_device(local_rank))
+    torch.get_device_module(device).set_device(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
 

test/registered/layers/mamba/test_mamba_ssm.py

@@ -13,6 +13,7 @@ from einops import rearrange, repeat
 
 from sglang.srt.layers.attention.mamba.causal_conv1d_triton import PAD_SLOT_ID
 from sglang.srt.layers.attention.mamba.ops import selective_state_update
+from sglang.srt.utils import get_device
 
 
 def selective_state_update_ref(
@@ -92,10 +93,9 @@ def selective_state_update_ref(
 @pytest.mark.parametrize("dstate", [16, 32, 64])
 @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
 def test_selective_state_update(dim, dstate, has_z, itype):
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA device not available")
-
-    device = "cuda"
+    device = get_device()
+    if device not in ["cuda", "xpu"]:
+        pytest.skip("Test only supports CUDA and XPU devices")
 
     rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
     if itype == torch.bfloat16:
@@ -136,10 +136,9 @@ def test_selective_state_update(dim, dstate, has_z, itype):
 def test_selective_state_update_with_batch_indices(
     with_padding, dim, dstate, has_z, itype
 ):
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA device not available")
-
-    device = "cuda"
+    device = get_device()
+    if device not in ["cuda", "xpu"]:
+        pytest.skip("Test only supports CUDA and XPU devices")
     rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
     if itype == torch.bfloat16:
         rtol, atol = 1e-1, 1e-1
@@ -229,10 +228,9 @@ def test_selective_state_update_with_batch_indices(
 def test_selective_state_update_with_heads_with_batch_indices(
     dim, dstate, ngroups, has_z, tie_hdim, itype
 ):
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA device not available")
-
-    device = "cuda"
+    device = get_device()
+    if device not in ["cuda", "xpu"]:
+        pytest.skip("Test only supports CUDA and XPU devices")
     rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 3e-2)
     if itype == torch.bfloat16:
         rtol, atol = 1e-1, 1e-1

test/registered/layers/mamba/test_mamba_ssm_ssd.py

@@ -14,6 +14,7 @@ from einops import rearrange, repeat
 
 from sglang.srt.layers.attention.mamba.mamba2_metadata import Mamba2Metadata
 from sglang.srt.layers.attention.mamba.ops import mamba_chunk_scan_combined
+from sglang.srt.utils import get_device
 from sglang.srt.utils.common import is_hip
 from sglang.utils import is_in_ci
 
@@ -99,10 +100,12 @@ def ssd_minimal_discrete(
     return Y, final_state
 
 
-def generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype, device="cuda"):
+def generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype, device=None):
 
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA device not available")
+    if device is None:
+        device = get_device()
+    if device not in ["cuda", "xpu"]:
+        pytest.skip("Test only supports CUDA and XPU devices")
 
     torch.manual_seed(0)
     A = -torch.exp(torch.rand(n_heads, dtype=itype, device=device))
@@ -125,7 +128,7 @@ def generate_continuous_batched_examples(
     n_heads,
     d_head,
     itype,
-    device="cuda",
+    device=None,
     return_naive_ref=True,
 ):
 
@@ -138,8 +141,10 @@ def generate_continuous_batched_examples(
 
     # generate the full-length example
     A, dt, X, B, C = generate_random_inputs(
-        num_examples, full_length, n_heads, d_head, itype
+        num_examples, full_length, n_heads, d_head, itype, device
     )
+    # Capture the resolved device from the tensors
+    device = X.device
 
     if return_naive_ref:
         Y_min, final_state_min = ssd_minimal_discrete(
@@ -227,8 +232,9 @@ if is_in_ci():
 @pytest.mark.parametrize("d_head", SINGLE_DHEAD)
 @pytest.mark.parametrize("seq_len_chunk_size", SINGLE_SEQ_LEN_CHUNK_SIZE)
 def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size, itype):
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA device not available")
+    device = get_device()
+    if device not in ["cuda", "xpu"]:
+        pytest.skip("Test only supports CUDA and XPU devices")
 
     # this tests the kernels on a single example (no batching)
 
@@ -319,8 +325,9 @@ if is_in_ci():
     ],
 )
 def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases, itype):
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA device not available")
+    device = get_device()
+    if device not in ["cuda", "xpu"]:
+        pytest.skip("Test only supports CUDA and XPU devices")
 
     # this test with multiple examples in a continuous batch
     # (i.e. chunked prefill)
@@ -398,8 +405,9 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases,
     ],
 )
 def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA device not available")
+    device = get_device()
+    if device not in ["cuda", "xpu"]:
+        pytest.skip("Test only supports CUDA and XPU devices")
 
     # This test verifies the correctness of the chunked prefill implementation
     # in the mamba2 ssd kernels, by comparing concatenation (in the sequence
@@ -632,8 +640,9 @@ def test_mamba_chunk_scan_intermediate_states(
     seq_len_chunk_size,
     itype,
 ):
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA device not available")
+    device = get_device()
+    if device not in ["cuda", "xpu"]:
+        pytest.skip("Test only supports CUDA and XPU devices")
 
     if itype == torch.bfloat16:
         atol, rtol = 5e-2, 5e-2