minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Test] Add XPU device support to unit tests (#22236)

Browse Source 
Co-authored-by: vshekhawat-hlab <vshekhawat@habana.ai>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
shubham singhal
2026-05-01 07:18:51 +08:00
committed by GitHub
co-authored by vshekhawat-hlab Ma Mingfei
parent 8b23d32ec1
commit e35ac95cdc
3 changed files with 45 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/quant/test_triton_scaled_mm.py
+17 -11
View File
@@ -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
)