diff --git a/python/sglang/test/attention/test_prefix_chunk_info.py b/python/sglang/test/attention/test_prefix_chunk_info.py index 2b85b695b..5002a0b09 100644 --- a/python/sglang/test/attention/test_prefix_chunk_info.py +++ b/python/sglang/test/attention/test_prefix_chunk_info.py @@ -4,6 +4,7 @@ import torch from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.utils.common import get_device from sglang.test.test_utils import CustomTestCase TEST_CASES = [ @@ -131,14 +132,17 @@ def check_kv_indices(forward_batch): assert torch.allclose(computed_kv_indices, ref_kv_indices) -@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") +@unittest.skipIf( + not (torch.cuda.is_available() or torch.xpu.is_available()), + "Test requires CUDA or XPU", +) class TestPrefixChunkInfo(CustomTestCase): def setUp(self): # Common test parameters self.num_local_heads = 128 self.kv_lora_rank = 512 self.qk_rope_head_dim = 64 - self.device = torch.device("cuda") + self.device = get_device() self.dtype = torch.bfloat16 self.extend_len = 64 self.max_bs = 4 diff --git a/test/registered/attention/test_kda_kernels.py b/test/registered/attention/test_kda_kernels.py index 42bdea4f0..0c11e236e 100644 --- a/test/registered/attention/test_kda_kernels.py +++ b/test/registered/attention/test_kda_kernels.py @@ -11,33 +11,40 @@ from sglang.srt.layers.attention.fla.kda import ( fused_recurrent_kda, kda_gate_chunk_cumsum, ) +from sglang.srt.utils.common import get_device from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=12, suite="stage-b-test-1-gpu-large") -@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") +@unittest.skipIf( + not (torch.cuda.is_available() or torch.xpu.is_available()), + "Test requires CUDA or XPU", +) class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase): def setUp(self): + self.device = get_device() self.token_num = 4 - self.query_start_loc = torch.tensor([0, 1, 2, 3, 4], device="cuda") - self.cache_indices = torch.tensor([0, 2, 5, 8], device="cuda") + self.query_start_loc = torch.tensor([0, 1, 2, 3, 4], device=self.device) + self.cache_indices = torch.tensor([0, 2, 5, 8], device=self.device) self.local_num_heads = 8 self.head_dim = 128 self.cache_len = 64 self.A_log = torch.randn( - 1, 1, self.local_num_heads, 1, dtype=torch.float32, device="cuda" + 1, 1, self.local_num_heads, 1, dtype=torch.float32, device=self.device ) self.a = torch.randn( 1, self.token_num, self.local_num_heads * self.head_dim, dtype=torch.bfloat16, - device="cuda", + device=self.device, ) self.dt_bias = torch.randn( - self.local_num_heads * self.head_dim, dtype=torch.bfloat16, device="cuda" + self.local_num_heads * self.head_dim, + dtype=torch.bfloat16, + device=self.device, ) self.softplus_beta = 1.0 self.softplus_threshold = 20.0 @@ -47,7 +54,7 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase): self.local_num_heads, self.head_dim, dtype=torch.bfloat16, - device="cuda", + device=self.device, ) self.k = torch.randn( 1, @@ -55,7 +62,7 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase): self.local_num_heads, self.head_dim, dtype=torch.bfloat16, - device="cuda", + device=self.device, ) self.v = torch.randn( 1, @@ -63,10 +70,14 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase): self.local_num_heads, self.head_dim, dtype=torch.bfloat16, - device="cuda", + device=self.device, ) self.beta = torch.randn( - 1, self.token_num, self.local_num_heads, dtype=torch.bfloat16, device="cuda" + 1, + self.token_num, + self.local_num_heads, + dtype=torch.bfloat16, + device=self.device, ) self.ssm_states = torch.zeros( @@ -75,7 +86,7 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase): self.head_dim, self.head_dim, dtype=torch.float32, - device="cuda", + device=self.device, ) def run_fused(self): diff --git a/test/registered/quant/test_triton_scaled_mm.py b/test/registered/quant/test_triton_scaled_mm.py index 8057fc593..d56a0f131 100644 --- a/test/registered/quant/test_triton_scaled_mm.py +++ b/test/registered/quant/test_triton_scaled_mm.py @@ -5,6 +5,7 @@ import torch import torch.testing from sglang.srt.layers.quantization.fp8_kernel import triton_scaled_mm +from sglang.srt.utils.common import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -31,20 +32,25 @@ def torch_scaled_mm( class TestScaledMM(CustomTestCase): @classmethod def setUpClass(cls): - if not torch.cuda.is_available(): - raise unittest.SkipTest("This test requires a CUDA device.") - torch.set_default_device("cuda") + if not (torch.cuda.is_available() or torch.xpu.is_available()): + raise unittest.SkipTest("No CUDA or XPU device available") + cls._device = get_device() + torch.set_default_device(cls._device) def _make_inputs(self, M, K, N, in_dtype): if in_dtype == torch.int8: - a = torch.randint(-8, 8, (M, K), dtype=in_dtype, device="cuda") - b = torch.randint(-8, 8, (K, N), dtype=in_dtype, device="cuda") + a = torch.randint(-8, 8, (M, K), dtype=in_dtype, device=self._device) + b = torch.randint(-8, 8, (K, N), dtype=in_dtype, device=self._device) else: # fp8 a = torch.clamp( - 0.1 * torch.randn((M, K), dtype=torch.float16, device="cuda"), -0.3, 0.3 + 0.1 * torch.randn((M, K), dtype=torch.float16, device=self._device), + -0.3, + 0.3, ).to(in_dtype) b = torch.clamp( - 0.1 * torch.randn((K, N), dtype=torch.float16, device="cuda"), -0.3, 0.3 + 0.1 * torch.randn((K, N), dtype=torch.float16, device=self._device), + -0.3, + 0.3, ).to(in_dtype) return a, b @@ -56,7 +62,7 @@ class TestScaledMM(CustomTestCase): ] try: - torch.tensor([1.0], dtype=torch.float8_e4m3fn, device="cuda") + torch.tensor([1.0], dtype=torch.float8_e4m3fn, device=self._device) test_configs.append((32, 32, 32, torch.float8_e4m3fn, torch.float16, False)) except: print("FP8 not supported, skipping") @@ -68,13 +74,13 @@ class TestScaledMM(CustomTestCase): input, weight = self._make_inputs(M, K, N, in_dtype) scale_a = 0.1 + 0.05 * torch.rand( - (M, 1), dtype=torch.float32, device="cuda" + (M, 1), dtype=torch.float32, device=self._device ) scale_b = 0.1 + 0.05 * torch.rand( - (N, 1), dtype=torch.float32, device="cuda" + (N, 1), dtype=torch.float32, device=self._device ) bias = ( - 0.01 * torch.randn((M, N), dtype=out_dtype, device="cuda") + 0.01 * torch.randn((M, N), dtype=out_dtype, device=self._device) if with_bias else None )