From 1a16c2ad2cb43b68b99a25965ebdb799c690968d Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Fri, 28 Aug 2026 18:08:20 -0600 Subject: [PATCH] [CI][Test] Deflake the rms_norm scaling-property assertions (#54271) Signed-off-by: Oxana Korzh Co-authored-by: Claude Opus 5 (1M context) --- tests/kernels/ir/test_layernorm.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/kernels/ir/test_layernorm.py b/tests/kernels/ir/test_layernorm.py index c87fb72845..fc79677165 100644 --- a/tests/kernels/ir/test_layernorm.py +++ b/tests/kernels/ir/test_layernorm.py @@ -12,7 +12,7 @@ from tests.ir.ir_test_utils import ( clone_args, supported_providers, ) -from tests.kernels.allclose_default import get_default_rtol +from tests.utils import set_random_seed from vllm import ir from vllm.platforms import current_platform @@ -52,6 +52,7 @@ def test_rms_norm_registration(): ) class TestRMSNorm: def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon): + set_random_seed(0) x, weight, epsilon = ir.ops.rms_norm.generate_inputs( num_tokens=4, hidden_size=8, @@ -66,9 +67,13 @@ class TestRMSNorm: assert out.dtype == x.dtype assert out.device == x.device - # Check the scaling property of rms norm + # Check the scaling property of rms norm. This holds only + # approximately: epsilon does not scale with x, so + # rms_norm(2x) = x / sqrt(mean(x^2) + epsilon/4), which differs from + # rms_norm(x) by the epsilon term. Use the op's declared tolerance + # rather than a tighter hard-coded one. out2 = rms_norm_native(x * 2.0, weight, epsilon=epsilon) - torch.testing.assert_close(out2, out, rtol=get_default_rtol(out), atol=1e-3) + assert_close(ir.ops.rms_norm, out2, out) # Mean square should be approximately 1 (ignoring epsilon and weight scaling) combined_norm = out.float() / weight.float() @@ -266,6 +271,7 @@ def test_vllm_c_fused_add_rms_norm_accepts_nd_input(): ) class TestFusedAddRMSNorm: def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon): + set_random_seed(0) x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs( num_tokens=4, hidden_size=8, @@ -304,7 +310,7 @@ class TestFusedAddRMSNorm: out2, _ = fused_add_rms_norm_native( x * 2.0, torch.zeros_like(x), weight, epsilon=epsilon ) - torch.testing.assert_close(out2, out1, rtol=get_default_rtol(out), atol=1e-3) + assert_close(ir.ops.fused_add_rms_norm, out2, out1) # Check behavior with and without weight weight1 = torch.ones_like(weight)