diff --git a/tests/kernels/core/test_fused_rms_norm_gated.py b/tests/kernels/core/test_fused_rms_norm_gated.py index f411f563ac..86b8e79656 100644 --- a/tests/kernels/core/test_fused_rms_norm_gated.py +++ b/tests/kernels/core/test_fused_rms_norm_gated.py @@ -7,11 +7,14 @@ matching the eager triton kernel output.""" import pytest import torch +from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.fused_norm_gate import ( FusedRMSNormGated, ) from vllm.utils.torch_utils import set_random_seed +DEVICE = "xpu:0" if current_platform.is_xpu() else "cuda:0" + DTYPES = [torch.bfloat16] HIDDEN_SIZES = [128, 512] NUM_TOKENS = [64, 128] @@ -39,7 +42,7 @@ def test_compiled_vs_eager( """forward_native decomposition matches forward_cuda triton kernel.""" torch._dynamo.reset() set_random_seed(seed) - device = torch.device("cuda:0") + device = torch.device(DEVICE) module = FusedRMSNormGated( hidden_size, @@ -88,7 +91,7 @@ def test_compiled_vs_eager_multidim( """forward_native decomposition handles multi-dimensional inputs.""" torch._dynamo.reset() set_random_seed(seed) - device = torch.device("cuda:0") + device = torch.device(DEVICE) head_dim = shape[-1] module = FusedRMSNormGated( diff --git a/tests/kernels/quantization/test_block_int8.py b/tests/kernels/quantization/test_block_int8.py index 3a13a53118..4a12ecd50d 100644 --- a/tests/kernels/quantization/test_block_int8.py +++ b/tests/kernels/quantization/test_block_int8.py @@ -14,7 +14,15 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import ( ) from vllm.platforms import current_platform -if current_platform.get_device_capability() < (7, 0): +if not (current_platform.is_cuda_alike() or current_platform.is_xpu()): + pytest.skip( + "INT8 Triton kernels require a CUDA-alike or XPU device", + allow_module_level=True, + ) + +if current_platform.is_cuda_alike() and not current_platform.has_device_capability( + (7, 0) +): pytest.skip("INT8 Triton requires CUDA 7.0 or higher", allow_module_level=True) vllm_config = VllmConfig() diff --git a/tests/kernels/quantization/test_int8_kernel.py b/tests/kernels/quantization/test_int8_kernel.py index 66385f04eb..704bb90f84 100644 --- a/tests/kernels/quantization/test_int8_kernel.py +++ b/tests/kernels/quantization/test_int8_kernel.py @@ -7,6 +7,7 @@ import itertools import pytest import torch +from tests.kernels.quant_utils import native_per_token_group_quant_int8 from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.fused_moe import fused_experts from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig @@ -15,7 +16,15 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import ( ) from vllm.platforms import current_platform -if current_platform.get_device_capability() < (7, 0): +if not (current_platform.is_cuda_alike() or current_platform.is_xpu()): + pytest.skip( + "INT8 Triton kernels require a CUDA-alike or XPU device", + allow_module_level=True, + ) + +if current_platform.is_cuda_alike() and not current_platform.has_device_capability( + (7, 0) +): pytest.skip("INT8 Triton requires CUDA 7.0 or higher", allow_module_level=True) @@ -155,3 +164,25 @@ def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed): torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)) ) / torch.mean(torch.abs(ref_out.to(torch.float32))) assert rel_diff < 0.05 + + +@pytest.mark.parametrize("num_tokens", [1, 33, 64]) +@pytest.mark.parametrize("hidden_dim", [128, 1024]) +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode() +def test_per_token_quant_int8(num_tokens, hidden_dim, dtype): + """Covers the round_int8 Triton kernel via per_token_quant_int8.""" + torch.manual_seed(0) + device = current_platform.device_type + x = torch.randn((num_tokens, hidden_dim), dtype=dtype, device=device) * 8 + + q, s = per_token_quant_int8(x) + # Per-token int8 quant is per-token-group quant with group_size = hidden_dim. + ref_q, ref_s = native_per_token_group_quant_int8(x, hidden_dim) + + assert q.dtype == torch.int8 + assert q.shape == (num_tokens, hidden_dim) + assert s.shape == (num_tokens, 1) + assert torch.allclose(s.float(), ref_s.float(), atol=1e-3, rtol=1e-3) + # Quantized values may differ by at most 1 ULP due to rounding. + assert torch.all((q.float() - ref_q.float()).abs() <= 1) diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py index 1cef5eb93a..e0f585c76a 100644 --- a/tests/kernels/quantization/test_triton_scaled_mm.py +++ b/tests/kernels/quantization/test_triton_scaled_mm.py @@ -13,7 +13,7 @@ import torch from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed -device = "cuda" +device = current_platform.device_type triton_scaled_mm_module = importlib.import_module( "vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm" diff --git a/tests/kernels/test_fla_layernorm_guard.py b/tests/kernels/test_fla_layernorm_guard.py index bb4de93201..f669bd1814 100644 --- a/tests/kernels/test_fla_layernorm_guard.py +++ b/tests/kernels/test_fla_layernorm_guard.py @@ -5,6 +5,7 @@ import pytest import torch import torch.nn.functional as F +from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( layer_norm_fwd, layernorm_fn, @@ -12,6 +13,8 @@ from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( ) from vllm.utils.torch_utils import set_random_seed +DEVICE = "xpu" if current_platform.is_xpu() else "cuda" + def layer_norm_ref( x, @@ -115,7 +118,7 @@ def test_layer_norm_fwd_basic( ) -> None: """Test basic layer norm forward pass without z (gate) tensor.""" set_random_seed(seed) - device = torch.device("cuda:0") + device = torch.device(DEVICE) # Create inputs x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) @@ -157,7 +160,7 @@ def test_layer_norm_fwd_with_gate( ) -> None: """Test layer norm forward pass with z (gate) tensor.""" set_random_seed(42) - device = torch.device("cuda:0") + device = torch.device(DEVICE) # Create inputs x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) @@ -214,7 +217,7 @@ def test_layer_norm_fwd_with_groups( ) set_random_seed(42) - device = torch.device("cuda:0") + device = torch.device(DEVICE) # Create inputs x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) @@ -254,7 +257,7 @@ def test_layer_norm_rows_per_block( ) -> None: """Test that rows_per_block logic works correctly for various M sizes.""" set_random_seed(42) - device = torch.device("cuda:0") + device = torch.device(DEVICE) hidden_size = 1024 # Create inputs @@ -279,7 +282,7 @@ def test_strided_input(dtype: torch.dtype) -> None: """Test that the kernel handles non-contiguous (strided) inputs correctly.""" set_random_seed(42) - device = torch.device("cuda:0") + device = torch.device(DEVICE) num_tokens = 128 hidden_size = 1024 @@ -319,7 +322,7 @@ def test_output_buffer_provided( ) -> None: """Test that the kernel works when an output buffer is provided.""" set_random_seed(42) - device = torch.device("cuda:0") + device = torch.device(DEVICE) # Create inputs x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) @@ -360,7 +363,7 @@ def test_multidimensional_input( ) -> None: """Test that the autograd function handles multidimensional inputs.""" set_random_seed(42) - device = torch.device("cuda:0") + device = torch.device(DEVICE) hidden_size = shape[-1] # Create inputs @@ -404,7 +407,7 @@ def test_rmsnorm_gated_forward_native_dtype( from vllm.model_executor.layers.layernorm import RMSNormGated - device = torch.device("cuda:0") + device = torch.device(DEVICE) set_random_seed(42) layer = RMSNormGated(