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:
co-authored by
Ma Mingfei
parent
c5f1339773
commit
8a9e424faa
@@ -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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user