[CI] Use explicit devices in quantization tests (#49512)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Andreas Karatzas committed
Jul 23, 2026 at 11:54 UTC
0e36e3bbd1a55c5ac3602304031e7d7e86bf0ca0
2 files changed
+22
-20
tests/kernels/quantization/test_block_int8.py
+10
-10
@@ -28,12 +28,6 @@ BLOCK_SIZE = [[128, 128]]
28
SEEDS = [0]
29
30
31
-@pytest.fixture(autouse=True, scope="module")
32
-def setup_cuda():
33
- """Sets the default CUDA device for all tests in this module."""
34
- torch.set_default_device("cuda")
35
-
36
-
31
@pytest.mark.parametrize(
32
"M,N,K,block_size,out_dtype,seed",
33
itertools.product(M, N, K, BLOCK_SIZE, DTYPES, SEEDS),
@@ -41,22 +35,28 @@ def setup_cuda():
35
@torch.inference_mode()
36
def test_w8a8_block_int8_matmul(M, N, K, block_size, out_dtype, seed):
37
torch.manual_seed(seed)
38
+ device = current_platform.device_type
39
factor_for_scale = 1e-2
40
int8_info = torch.iinfo(torch.int8)
41
int8_max, int8_min = int8_info.max, int8_info.min
42
48
- A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * int8_max
43
+ A_fp32 = torch.rand(M, K, dtype=torch.float32, device=device)
44
+ A_fp32 = (A_fp32 - 0.5) * 2 * int8_max
45
A_fp8 = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn)
46
51
- B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * int8_max
47
+ B_fp32 = torch.rand(N, K, dtype=torch.float32, device=device)
48
+ B_fp32 = (B_fp32 - 0.5) * 2 * int8_max
49
B_fp8 = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn)
50
51
block_n, block_k = block_size[0], block_size[1]
52
n_tiles = (N + block_n - 1) // block_n
53
k_tiles = (K + block_k - 1) // block_k
54
58
- As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale
59
- Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
55
+ As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale
56
+ Bs = (
57
+ torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device)
58
+ * factor_for_scale
59
+ )
60
61
ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
62
out = w8a8_block_int8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
tests/kernels/quantization/test_int8_kernel.py
+12
-10
@@ -82,12 +82,6 @@ def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids
82
).sum(dim=1)
83
84
85
-@pytest.fixture(autouse=True, scope="module")
86
-def setup_cuda():
87
- """Sets the default CUDA device for all tests in this module."""
88
- torch.set_default_device("cuda")
89
-
90
-
85
DTYPES = [torch.half, torch.bfloat16]
86
M = [1, 33]
87
N = [128, 1024]
@@ -104,6 +98,7 @@ SEEDS = [0]
98
@torch.inference_mode()
99
def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed):
100
torch.manual_seed(seed)
101
+ device = current_platform.device_type
102
# Initialize int8 quantization parameters
103
factor_for_scale = 1e-2
104
int8_max = 127
@@ -111,19 +106,26 @@ def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed):
106
107
# Input tensor
108
# M * K
114
- a = torch.randn((M, K), dtype=dtype) / 10
109
+ a = torch.randn((M, K), dtype=dtype, device=device) / 10
110
111
# Generate int8 weights
117
- w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2
112
+ w1_fp32 = (
113
+ torch.rand(
114
+ (E, 2 * N, K),
115
+ dtype=torch.float32,
116
+ device=device,
117
+ )
118
+ - 0.5
119
+ ) * 2
120
w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
121
120
- w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2
122
+ w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32, device=device) - 0.5) * 2
123
w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
124
125
# Generate scale for each column (per-column quantization)
126
w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale
127
w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale
126
- score = torch.randn((M, E), dtype=dtype)
128
+ score = torch.randn((M, E), dtype=dtype, device=device)
129
score = torch.softmax(score, dim=-1, dtype=torch.float32)
130
topk_weights, topk_ids = torch.topk(score, topk)
131