Make UTs compatible for XPU (#27106)

test/registered/debug_utils/test_tensor_dump_forward_hook.py

@@ -7,6 +7,7 @@ from sglang.srt.debug_utils.tensor_dump_forward_hook import (
     register_forward_hook_for_model,
 )
 from sglang.srt.distributed.parallel_state import (
+    get_default_distributed_backend,
     init_distributed_environment,
     initialize_model_parallel,
 )
@@ -14,7 +15,7 @@ from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import LinearBase
 from sglang.srt.models.qwen2 import Qwen2MLP
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
-from sglang.srt.utils import add_prefix
+from sglang.srt.utils import add_prefix, get_device
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(
@@ -78,8 +79,10 @@ def init_weights(module):
 
 def test_model_forward_dump(tmp_path):
     set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+    device = get_device()
+    backend = get_default_distributed_backend(device)
     init_distributed_environment(
-        backend="nccl",
+        backend=backend,
         world_size=1,
         rank=0,
         local_rank=0,
@@ -88,14 +91,14 @@ def test_model_forward_dump(tmp_path):
     initialize_model_parallel()
     model = MockCausalLM()
     model.apply(init_weights)
-    model = model.cuda().bfloat16()
+    model = model.to(device=device, dtype=torch.bfloat16)
     dumper = register_forward_hook_for_model(
         model, tmp_path / "sglang_dump", [0], 0, 0, 0
     )
 
     dir_path = dumper.get_dump_dir()
     inp = torch.randn(4, TEST_HIDDEN_SIZE, dtype=torch.bfloat16) * 0.01
-    result = model(inp.cuda())
+    result = model(inp.to(device))
     data = torch.load(f"{dir_path}/Pass00000.pt")
     assert "model.layernorm" in data
     assert "model.mlp.down_proj" in data

test/registered/kernels/test_fused_topk_deepseek.py

@@ -4,6 +4,7 @@ import pytest
 import torch
 
 from sglang.srt.layers.moe.topk import biased_grouped_topk_gpu, biased_grouped_topk_impl
+from sglang.srt.utils import get_device
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=2, suite="nightly-1-gpu", nightly=True)
@@ -29,11 +30,12 @@ def test_fused_topk_deepseek(seq_length, params, apply_routed_scaling_factor_on_
     """
     num_experts, num_expert_group, topk_group, topk = params
     dtype = torch.float32
+    device = get_device()
 
     torch.manual_seed(seq_length)
-    hidden_states = torch.randn(seq_length, 128, dtype=dtype, device="cuda")
-    gating_output = torch.randn(seq_length, num_experts, dtype=dtype, device="cuda")
-    correction_bias = torch.randn(num_experts, dtype=dtype, device="cuda")
+    hidden_states = torch.randn(seq_length, 128, dtype=dtype, device=device)
+    gating_output = torch.randn(seq_length, num_experts, dtype=dtype, device=device)
+    correction_bias = torch.randn(num_experts, dtype=dtype, device=device)
 
     routed_scaling_factor = 2.5 if apply_routed_scaling_factor_on_output else None
 
@@ -71,8 +73,8 @@ def test_fused_topk_deepseek(seq_length, params, apply_routed_scaling_factor_on_
     sum_check = torch.allclose(output_sum, ref_output_sum, rtol=1e-03, atol=1e-04)
 
     # Check 2: Scatter-based comparison with allowance for tie-breaking
-    res = torch.zeros(seq_length, num_experts, dtype=torch.float32, device="cuda")
-    ref = torch.zeros(seq_length, num_experts, dtype=torch.float32, device="cuda")
+    res = torch.zeros(seq_length, num_experts, dtype=torch.float32, device=device)
+    ref = torch.zeros(seq_length, num_experts, dtype=torch.float32, device=device)
 
     res.scatter_(1, indices.long(), output)
     ref.scatter_(1, ref_indices.long(), ref_output)

test/registered/unit/mem_cache/test_radix_cache_unit.py

@@ -39,6 +39,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
 )
 from sglang.srt.mem_cache.mamba_radix_cache import TreeNode as MambaTreeNode
 from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
+from sglang.srt.utils import get_device
 
 # Test constants
 DEFAULT_PAGE_SIZE = 4
@@ -774,7 +775,7 @@ class TestRadixCache(unittest.TestCase):
         base_prefix_len = 10000
         suffix_len = 100
 
-        torch_allocated_before = torch.cuda.memory_allocated()
+        torch_allocated_before = torch.get_device_module().memory_allocated()
 
         # build dataset with common prefix
         common_prefix = [
@@ -784,7 +785,7 @@ class TestRadixCache(unittest.TestCase):
             suffix = [random.randint(1, vocab_size - 1) for _ in range(suffix_len)]
             seq = common_prefix + suffix
             keys.append(seq)
-            values.append(torch.zeros(len(seq), device="cuda", dtype=torch.int32))
+            values.append(torch.zeros(len(seq), device=get_device(), dtype=torch.int32))
 
         cache: RadixCache = RadixCache.create_simulated()
 
@@ -793,7 +794,9 @@ class TestRadixCache(unittest.TestCase):
 
         del values
 
-        torch_allocated = torch.cuda.memory_allocated() - torch_allocated_before
+        torch_allocated = (
+            torch.get_device_module().memory_allocated() - torch_allocated_before
+        )
         cache_size_bytes = cache.total_size() * 4
         print(f"\nCache size (MB): {cache_size_bytes / (1024 * 1024)}")
         print(f"Torch allocated (MB): {torch_allocated / (1024 * 1024)}")