From 0e36e3bbd1a55c5ac3602304031e7d7e86bf0ca0 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 23 Jul 2026 11:54:26 -0500 Subject: [PATCH] [CI] Use explicit devices in quantization tests (#49512) Signed-off-by: Andreas Karatzas --- tests/kernels/quantization/test_block_int8.py | 20 ++++++++--------- .../kernels/quantization/test_int8_kernel.py | 22 ++++++++++--------- 2 files changed, 22 insertions(+), 20 deletions(-) diff --git a/tests/kernels/quantization/test_block_int8.py b/tests/kernels/quantization/test_block_int8.py index 310091b6a5..3a13a53118 100644 --- a/tests/kernels/quantization/test_block_int8.py +++ b/tests/kernels/quantization/test_block_int8.py @@ -28,12 +28,6 @@ BLOCK_SIZE = [[128, 128]] SEEDS = [0] -@pytest.fixture(autouse=True, scope="module") -def setup_cuda(): - """Sets the default CUDA device for all tests in this module.""" - torch.set_default_device("cuda") - - @pytest.mark.parametrize( "M,N,K,block_size,out_dtype,seed", itertools.product(M, N, K, BLOCK_SIZE, DTYPES, SEEDS), @@ -41,22 +35,28 @@ def setup_cuda(): @torch.inference_mode() def test_w8a8_block_int8_matmul(M, N, K, block_size, out_dtype, seed): torch.manual_seed(seed) + device = current_platform.device_type factor_for_scale = 1e-2 int8_info = torch.iinfo(torch.int8) int8_max, int8_min = int8_info.max, int8_info.min - A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * int8_max + A_fp32 = torch.rand(M, K, dtype=torch.float32, device=device) + A_fp32 = (A_fp32 - 0.5) * 2 * int8_max A_fp8 = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn) - B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * int8_max + B_fp32 = torch.rand(N, K, dtype=torch.float32, device=device) + B_fp32 = (B_fp32 - 0.5) * 2 * int8_max B_fp8 = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn) block_n, block_k = block_size[0], block_size[1] n_tiles = (N + block_n - 1) // block_n k_tiles = (K + block_k - 1) // block_k - As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale - Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale + As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale + Bs = ( + torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device) + * factor_for_scale + ) ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) out = w8a8_block_int8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) diff --git a/tests/kernels/quantization/test_int8_kernel.py b/tests/kernels/quantization/test_int8_kernel.py index 0daa488922..66385f04eb 100644 --- a/tests/kernels/quantization/test_int8_kernel.py +++ b/tests/kernels/quantization/test_int8_kernel.py @@ -82,12 +82,6 @@ def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids ).sum(dim=1) -@pytest.fixture(autouse=True, scope="module") -def setup_cuda(): - """Sets the default CUDA device for all tests in this module.""" - torch.set_default_device("cuda") - - DTYPES = [torch.half, torch.bfloat16] M = [1, 33] N = [128, 1024] @@ -104,6 +98,7 @@ SEEDS = [0] @torch.inference_mode() def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed): torch.manual_seed(seed) + device = current_platform.device_type # Initialize int8 quantization parameters factor_for_scale = 1e-2 int8_max = 127 @@ -111,19 +106,26 @@ def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed): # Input tensor # M * K - a = torch.randn((M, K), dtype=dtype) / 10 + a = torch.randn((M, K), dtype=dtype, device=device) / 10 # Generate int8 weights - w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2 + w1_fp32 = ( + torch.rand( + (E, 2 * N, K), + dtype=torch.float32, + device=device, + ) + - 0.5 + ) * 2 w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8) - w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2 + w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32, device=device) - 0.5) * 2 w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8) # Generate scale for each column (per-column quantization) w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale - score = torch.randn((M, E), dtype=dtype) + score = torch.randn((M, E), dtype=dtype, device=device) score = torch.softmax(score, dim=-1, dtype=torch.float32) topk_weights, topk_ids = torch.topk(score, topk)