From 2ac125123a8d312823cf4ef56ca165f04f579dd1 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 23 Jul 2026 22:08:14 -0500 Subject: [PATCH] [CI] Use explicit devices in IR tests (#49513) Signed-off-by: Andreas Karatzas --- tests/kernels/ir/test_layernorm.py | 75 ++++++++++++++++++++---------- vllm/ir/ops/layernorm.py | 22 ++++++--- 2 files changed, 66 insertions(+), 31 deletions(-) diff --git a/tests/kernels/ir/test_layernorm.py b/tests/kernels/ir/test_layernorm.py index 6fc9d7543d..37ec1468a8 100644 --- a/tests/kernels/ir/test_layernorm.py +++ b/tests/kernels/ir/test_layernorm.py @@ -50,13 +50,13 @@ def test_rms_norm_registration(): reason="Currently only kernels on CUDA, ROCm and XPU", ) class TestRMSNorm: - @classmethod - def setup_class(cls, **kwargs): - torch.set_default_device(current_platform.device_type) - def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon): x, weight, epsilon = ir.ops.rms_norm.generate_inputs( - num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon + num_tokens=4, + hidden_size=8, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) out = rms_norm_native(x, weight, epsilon=epsilon) @@ -87,7 +87,11 @@ class TestRMSNorm: def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider): impl = ir.ops.rms_norm.impls[provider] x, weight, eps = ir.ops.rms_norm.generate_inputs( - num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + num_tokens=n_tokens, + hidden_size=hidden_size, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) args = (x, weight, eps) @@ -119,7 +123,11 @@ class TestRMSNorm: pytest.skip(f"{provider} impl not supported on this platform") args = ir.ops.rms_norm.generate_inputs( - num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + num_tokens=n_tokens, + hidden_size=hidden_size, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) # When checking the torch op, we have to set priority and use dispatch @@ -132,11 +140,14 @@ class TestRMSNorm: reason="aiter is only supported on ROCm", ) def test_aiter_rejects_unsupported_dtypes(): - torch.set_default_device(current_platform.device_type) impl = ir.ops.rms_norm.impls["aiter"] for dtype in [torch.float32, torch.float64]: args = ir.ops.rms_norm.generate_inputs( - num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5 + num_tokens=8, + hidden_size=4096, + dtype=dtype, + epsilon=1e-5, + device=current_platform.device_type, ) assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}" @@ -146,15 +157,16 @@ def test_aiter_rejects_unsupported_dtypes(): reason="ROCm vllm_c RMSNorm needs explicit ND input handling", ) def test_vllm_c_rms_norm_accepts_nd_input(): - torch.set_default_device(current_platform.device_type) impl = ir.ops.rms_norm.impls["vllm_c"] if not impl.supported: pytest.skip("vllm_c impl not supported on this platform") - base = torch.randn(3, 8, 192, dtype=torch.float16) + base = torch.randn( + 3, 8, 192, dtype=torch.float16, device=current_platform.device_type + ) x = base.split(64, dim=-1)[0].view(3, 8, 4, 16) assert not x.is_contiguous() - weight = torch.randn(16, dtype=torch.float16) + weight = torch.randn(16, dtype=torch.float16, device=current_platform.device_type) epsilon = 1e-5 output = impl.impl_fn(x, weight, epsilon) @@ -195,18 +207,21 @@ def test_fused_add_rms_norm_registration(): reason="ROCm vllm_c fused_add_rms_norm needs explicit ND input handling", ) def test_vllm_c_fused_add_rms_norm_accepts_nd_input(): - torch.set_default_device(current_platform.device_type) impl = ir.ops.fused_add_rms_norm.impls["vllm_c"] if not impl.supported: pytest.skip("vllm_c impl not supported on this platform") - base = torch.randn(3, 8, 192, dtype=torch.float16) - residual_base = torch.randn(3, 8, 192, dtype=torch.float16) + base = torch.randn( + 3, 8, 192, dtype=torch.float16, device=current_platform.device_type + ) + residual_base = torch.randn( + 3, 8, 192, dtype=torch.float16, device=current_platform.device_type + ) x = base.split(64, dim=-1)[0].view(3, 8, 4, 16) x_residual = residual_base.split(64, dim=-1)[0].view(3, 8, 4, 16) assert not x.is_contiguous() assert not x_residual.is_contiguous() - weight = torch.randn(16, dtype=torch.float16) + weight = torch.randn(16, dtype=torch.float16, device=current_platform.device_type) epsilon = 1e-5 output, residual = impl.impl_fn(x.clone(), x_residual.clone(), weight, epsilon) @@ -227,13 +242,13 @@ def test_vllm_c_fused_add_rms_norm_accepts_nd_input(): reason="Currently only kernels on CUDA, ROCm and XPU", ) class TestFusedAddRMSNorm: - @classmethod - def setup_class(cls, **kwargs): - torch.set_default_device(current_platform.device_type) - def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon): x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs( - num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon + num_tokens=4, + hidden_size=8, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) out, residual_out = fused_add_rms_norm_native(x, x_residual, weight, eps) @@ -278,7 +293,11 @@ class TestFusedAddRMSNorm: def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider): impl = ir.ops.fused_add_rms_norm.impls[provider] x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs( - num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + num_tokens=n_tokens, + hidden_size=hidden_size, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) args = (x, x_residual, weight, eps, None) @@ -324,7 +343,11 @@ class TestFusedAddRMSNorm: pytest.skip(f"{provider} impl not supported on this platform") x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs( - num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + num_tokens=n_tokens, + hidden_size=hidden_size, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) # Test default overload - should NOT modify inputs even with inplace impl @@ -368,7 +391,11 @@ class TestFusedAddRMSNorm: @pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm)) def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider): args = ir.ops.fused_add_rms_norm.generate_inputs( - num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + num_tokens=n_tokens, + hidden_size=hidden_size, + dtype=dtype, + epsilon=epsilon, + device=current_platform.device_type, ) args = args + (None,) # Add variance_size parameter diff --git a/vllm/ir/ops/layernorm.py b/vllm/ir/ops/layernorm.py index 33a71b8f85..a40028a4d0 100644 --- a/vllm/ir/ops/layernorm.py +++ b/vllm/ir/ops/layernorm.py @@ -23,10 +23,14 @@ def rms_norm( @rms_norm.register_input_generator def _rms_norm_input_generator( - num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5 + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + epsilon: float = 1e-5, + device: torch.device | str | None = None, ) -> tuple: - x = torch.randn(num_tokens, hidden_size, dtype=dtype) - weight = torch.randn(hidden_size, dtype=dtype) + x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + weight = torch.randn(hidden_size, dtype=dtype, device=device) return x, weight, epsilon @@ -64,9 +68,13 @@ fused_add_rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3) @fused_add_rms_norm.register_input_generator def _fused_add_rms_norm_input_generator( - num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5 + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + epsilon: float = 1e-5, + device: torch.device | str | None = None, ) -> tuple: - x = torch.randn(num_tokens, hidden_size, dtype=dtype) - x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype) - weight = torch.randn(hidden_size, dtype=dtype) + x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + weight = torch.randn(hidden_size, dtype=dtype, device=device) return x, x_residual, weight, epsilon