[CI] Use explicit devices in IR tests (#49513)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Andreas Karatzas committed
Jul 23, 2026 at 22:08 UTC
2ac125123a8d312823cf4ef56ca165f04f579dd1
2 files changed
+66
-31
tests/kernels/ir/test_layernorm.py
+51
-24
@@ -50,13 +50,13 @@ def test_rms_norm_registration():
50
reason="Currently only kernels on CUDA, ROCm and XPU",
51
)
52
class TestRMSNorm:
53
- @classmethod
54
- def setup_class(cls, **kwargs):
55
- torch.set_default_device(current_platform.device_type)
56
-
53
def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
54
x, weight, epsilon = ir.ops.rms_norm.generate_inputs(
59
- num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
55
+ num_tokens=4,
56
+ hidden_size=8,
57
+ dtype=dtype,
58
+ epsilon=epsilon,
59
+ device=current_platform.device_type,
60
)
61
out = rms_norm_native(x, weight, epsilon=epsilon)
62
@@ -87,7 +87,11 @@ class TestRMSNorm:
87
def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
88
impl = ir.ops.rms_norm.impls[provider]
89
x, weight, eps = ir.ops.rms_norm.generate_inputs(
90
- num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
90
+ num_tokens=n_tokens,
91
+ hidden_size=hidden_size,
92
+ dtype=dtype,
93
+ epsilon=epsilon,
94
+ device=current_platform.device_type,
95
)
96
args = (x, weight, eps)
97
@@ -119,7 +123,11 @@ class TestRMSNorm:
123
pytest.skip(f"{provider} impl not supported on this platform")
124
125
args = ir.ops.rms_norm.generate_inputs(
122
- num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
126
+ num_tokens=n_tokens,
127
+ hidden_size=hidden_size,
128
+ dtype=dtype,
129
+ epsilon=epsilon,
130
+ device=current_platform.device_type,
131
)
132
133
# When checking the torch op, we have to set priority and use dispatch
@@ -132,11 +140,14 @@ class TestRMSNorm:
140
reason="aiter is only supported on ROCm",
141
)
142
def test_aiter_rejects_unsupported_dtypes():
135
- torch.set_default_device(current_platform.device_type)
143
impl = ir.ops.rms_norm.impls["aiter"]
144
for dtype in [torch.float32, torch.float64]:
145
args = ir.ops.rms_norm.generate_inputs(
139
- num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5
146
+ num_tokens=8,
147
+ hidden_size=4096,
148
+ dtype=dtype,
149
+ epsilon=1e-5,
150
+ device=current_platform.device_type,
151
)
152
assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}"
153
@@ -146,15 +157,16 @@ def test_aiter_rejects_unsupported_dtypes():
157
reason="ROCm vllm_c RMSNorm needs explicit ND input handling",
158
)
159
def test_vllm_c_rms_norm_accepts_nd_input():
149
- torch.set_default_device(current_platform.device_type)
160
impl = ir.ops.rms_norm.impls["vllm_c"]
161
if not impl.supported:
162
pytest.skip("vllm_c impl not supported on this platform")
163
154
- base = torch.randn(3, 8, 192, dtype=torch.float16)
164
+ base = torch.randn(
165
+ 3, 8, 192, dtype=torch.float16, device=current_platform.device_type
166
+ )
167
x = base.split(64, dim=-1)[0].view(3, 8, 4, 16)
168
assert not x.is_contiguous()
157
- weight = torch.randn(16, dtype=torch.float16)
169
+ weight = torch.randn(16, dtype=torch.float16, device=current_platform.device_type)
170
epsilon = 1e-5
171
172
output = impl.impl_fn(x, weight, epsilon)
@@ -195,18 +207,21 @@ def test_fused_add_rms_norm_registration():
207
reason="ROCm vllm_c fused_add_rms_norm needs explicit ND input handling",
208
)
209
def test_vllm_c_fused_add_rms_norm_accepts_nd_input():
198
- torch.set_default_device(current_platform.device_type)
210
impl = ir.ops.fused_add_rms_norm.impls["vllm_c"]
211
if not impl.supported:
212
pytest.skip("vllm_c impl not supported on this platform")
213
203
- base = torch.randn(3, 8, 192, dtype=torch.float16)
204
- residual_base = torch.randn(3, 8, 192, dtype=torch.float16)
214
+ base = torch.randn(
215
+ 3, 8, 192, dtype=torch.float16, device=current_platform.device_type
216
+ )
217
+ residual_base = torch.randn(
218
+ 3, 8, 192, dtype=torch.float16, device=current_platform.device_type
219
+ )
220
x = base.split(64, dim=-1)[0].view(3, 8, 4, 16)
221
x_residual = residual_base.split(64, dim=-1)[0].view(3, 8, 4, 16)
222
assert not x.is_contiguous()
223
assert not x_residual.is_contiguous()
209
- weight = torch.randn(16, dtype=torch.float16)
224
+ weight = torch.randn(16, dtype=torch.float16, device=current_platform.device_type)
225
epsilon = 1e-5
226
227
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():
242
reason="Currently only kernels on CUDA, ROCm and XPU",
243
)
244
class TestFusedAddRMSNorm:
230
- @classmethod
231
- def setup_class(cls, **kwargs):
232
- torch.set_default_device(current_platform.device_type)
233
-
245
def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
246
x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
236
- num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
247
+ num_tokens=4,
248
+ hidden_size=8,
249
+ dtype=dtype,
250
+ epsilon=epsilon,
251
+ device=current_platform.device_type,
252
)
253
out, residual_out = fused_add_rms_norm_native(x, x_residual, weight, eps)
254
@@ -278,7 +293,11 @@ class TestFusedAddRMSNorm:
293
def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
294
impl = ir.ops.fused_add_rms_norm.impls[provider]
295
x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
281
- num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
296
+ num_tokens=n_tokens,
297
+ hidden_size=hidden_size,
298
+ dtype=dtype,
299
+ epsilon=epsilon,
300
+ device=current_platform.device_type,
301
)
302
args = (x, x_residual, weight, eps, None)
303
@@ -324,7 +343,11 @@ class TestFusedAddRMSNorm:
343
pytest.skip(f"{provider} impl not supported on this platform")
344
345
x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
327
- num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
346
+ num_tokens=n_tokens,
347
+ hidden_size=hidden_size,
348
+ dtype=dtype,
349
+ epsilon=epsilon,
350
+ device=current_platform.device_type,
351
)
352
353
# Test default overload - should NOT modify inputs even with inplace impl
@@ -368,7 +391,11 @@ class TestFusedAddRMSNorm:
391
@pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm))
392
def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
393
args = ir.ops.fused_add_rms_norm.generate_inputs(
371
- num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
394
+ num_tokens=n_tokens,
395
+ hidden_size=hidden_size,
396
+ dtype=dtype,
397
+ epsilon=epsilon,
398
+ device=current_platform.device_type,
399
)
400
args = args + (None,) # Add variance_size parameter
401
vllm/ir/ops/layernorm.py
+15
-7
@@ -23,10 +23,14 @@ def rms_norm(
23
24
@rms_norm.register_input_generator
25
def _rms_norm_input_generator(
26
- num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5
26
+ num_tokens: int,
27
+ hidden_size: int,
28
+ dtype: torch.dtype,
29
+ epsilon: float = 1e-5,
30
+ device: torch.device | str | None = None,
31
) -> tuple:
28
- x = torch.randn(num_tokens, hidden_size, dtype=dtype)
29
- weight = torch.randn(hidden_size, dtype=dtype)
32
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
33
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
34
return x, weight, epsilon
35
36
@@ -64,9 +68,13 @@ fused_add_rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3)
68
69
@fused_add_rms_norm.register_input_generator
70
def _fused_add_rms_norm_input_generator(
67
- num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5
71
+ num_tokens: int,
72
+ hidden_size: int,
73
+ dtype: torch.dtype,
74
+ epsilon: float = 1e-5,
75
+ device: torch.device | str | None = None,
76
) -> tuple:
69
- x = torch.randn(num_tokens, hidden_size, dtype=dtype)
70
- x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype)
71
- weight = torch.randn(hidden_size, dtype=dtype)
77
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
78
+ x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
79
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
80
return x, x_residual, weight, epsilon