Replace hardcoded CUDA device with get_device() for XPU support (#13599)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
Kalyan Kumar
2026-05-01 07:13:46 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent c5f1339773
commit 8a9e424faa
4 changed files with 54 additions and 14 deletions
+20 -6
View File
@@ -11,6 +11,13 @@ from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=10, suite="stage-b-test-1-gpu-large")
from sglang.srt.utils import get_device, is_cuda, is_xpu
_is_cuda = is_cuda()
_is_xpu = is_xpu()
device = get_device()
class TestFP8Base(CustomTestCase):
@classmethod
@@ -26,7 +33,7 @@ class TestFP8Base(CustomTestCase):
@staticmethod
def _make_A(M, K, group_size, out_dtype):
quant_A = torch.rand(
M, K // group_size, group_size, dtype=torch.float32, device="cuda"
M, K // group_size, group_size, dtype=torch.float32, device=device
)
# -1 ~ 1
quant_A = quant_A * 2 - 1
@@ -38,7 +45,7 @@ class TestFP8Base(CustomTestCase):
quant_A = quant_A.to(out_dtype).to(torch.float32)
# create scale and A
scale = torch.rand(M, K // group_size, dtype=torch.float32, device="cuda")
scale = torch.rand(M, K // group_size, dtype=torch.float32, device=device)
scale /= fmax
A = quant_A * scale[..., None]
@@ -60,7 +67,7 @@ class TestFP8Base(CustomTestCase):
N_aligned // group_size,
group_size,
dtype=torch.float32,
device="cuda",
device=device,
)
quant_B = quant_B * 2 - 1
@@ -77,7 +84,7 @@ class TestFP8Base(CustomTestCase):
N_aligned // group_size,
1,
dtype=torch.float32,
device="cuda",
device=device,
)
scale /= fmax
@@ -91,8 +98,9 @@ class TestFP8Base(CustomTestCase):
class TestPerTokenGroupQuantFP8(TestFP8Base):
def test_per_token_group_quant_fp8(self):
if torch.cuda.get_device_capability()[0] < 9:
if _is_cuda and torch.cuda.get_device_capability()[0] < 9:
return
A, A_quant_gt, scale_gt = self._make_A(
M=self.M, K=self.K, group_size=self.group_size, out_dtype=self.quant_type
)
@@ -107,8 +115,14 @@ class TestPerTokenGroupQuantFP8(TestFP8Base):
class TestW8A8BlockFP8Matmul(TestFP8Base):
def test_w8a8_block_fp8_matmul(self):
if torch.cuda.get_device_capability()[0] < 9:
if _is_cuda and torch.cuda.get_device_capability()[0] < 9:
return
elif _is_xpu:
# XPU doesn't provide traditional capability info like CUDA
pass
else:
return
A, A_quant_gt, A_scale_gt = self._make_A(
M=self.M, K=self.K, group_size=self.group_size, out_dtype=self.quant_type
)