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>

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)