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