[ROCm][CI] Relax fused layernorm quant test tolerances for one-ULP outliers (#46658)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Jun 26, 2026 at 22:44 UTC
17a71d87020e163a291bd34f91dad3eb05b448e5
4 files changed
+90
-55
tests/kernels/core/test_fused_quant_layernorm.py
+40
-16
@@ -8,7 +8,7 @@ import pytest
8
import torch
9
10
import vllm._custom_ops as ops
11
-from tests.kernels.utils import opcheck
11
+from tests.kernels.utils import fp8_ulp_distance, opcheck
12
from vllm.model_executor.layers.layernorm import RMSNorm
13
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
14
per_token_group_quant_fp8,
@@ -250,30 +250,54 @@ def test_rms_norm(
250
251
assert ref_out.dtype == quant_dtype
252
assert ops_out.dtype == quant_dtype
253
+
254
+ # Per-block bf16 scales: allow a small relative tolerance for a few groups
255
+ # whose abs-max flips by one ULP between the fused and reference paths. The
256
+ # per-token and fp32 paths stay strict.
257
+ relax_block_rocm = (
258
+ group_size is not None
259
+ and dtype == torch.bfloat16
260
+ and current_platform.is_rocm()
261
+ )
262
+
263
+ def scales_close(rtol: float, atol: float) -> bool:
264
+ if torch.allclose(ref_scales, ops_scales, rtol=rtol, atol=atol):
265
+ return True
266
+ return relax_block_rocm and torch.allclose(
267
+ ref_scales, ops_scales, rtol=1e-2, atol=atol
268
+ )
269
+
270
if quant_dtype == torch.int8:
254
- assert torch.allclose(ref_scales, ops_scales, atol=1e-6)
271
+ assert scales_close(rtol=1e-5, atol=1e-6)
272
# big atol to account for round-off errors.
273
assert torch.allclose(ref_out, ops_out, atol=1)
274
else:
258
- assert torch.allclose(ref_scales, ops_scales)
275
+ assert scales_close(rtol=1e-5, atol=1e-8)
276
a = ref_out.to(dtype=torch.float32)
277
b = ops_out.to(dtype=torch.float32)
278
ok = torch.allclose(a, b, atol=1e-6)
279
if not ok:
263
- # fallback: compare dequantized values with relaxed tolerance
264
- if group_size is None:
265
- a_deq = a * ref_scales.view(-1, 1)
266
- b_deq = b * ops_scales.view(-1, 1)
280
+ if relax_block_rocm:
281
+ # ULP-flipped group scale can cross an E4M3 tie; tolerate a
282
+ # bounded count of isolated fp8 outliers.
283
+ ulp = fp8_ulp_distance(ref_out, ops_out)
284
+ max_outliers = ulp.numel() // 100_000 + 8
285
+ ok = int((ulp > 0).sum().item()) <= max_outliers
286
else:
268
- a_deq = a * ref_scales.repeat_interleave(group_size[1], dim=1)
269
- b_deq = b * ops_scales.repeat_interleave(group_size[1], dim=1)
270
- # NOTE: It is possible that some future test cases trigger this
271
- # max diff due to precision issues. If such an error is
272
- # encountered, it's recommended to inspect the differences between
273
- # all corresponding elements from each tensor (e.g. by looping over
274
- # them) and checking how many the max diff error shows up on (just
275
- # a few bad elements should still be considered acceptable).
276
- ok = torch.allclose(a_deq, b_deq, rtol=5e-2, atol=5e-2)
287
+ # CUDA (& non-bf16): compare dequantized values with relaxed tolerance.
288
+ if group_size is None:
289
+ a_deq = a * ref_scales.view(-1, 1)
290
+ b_deq = b * ops_scales.view(-1, 1)
291
+ else:
292
+ a_deq = a * ref_scales.repeat_interleave(group_size[1], dim=1)
293
+ b_deq = b * ops_scales.repeat_interleave(group_size[1], dim=1)
294
+ # NOTE: It is possible that some future test cases trigger this
295
+ # max diff due to precision issues. If such an error is
296
+ # encountered, it's recommended to inspect the differences between
297
+ # all corresponding elements from each tensor (e.g. by looping over
298
+ # them) and checking how many the max diff error shows up on (just
299
+ # a few bad elements should still be considered acceptable).
300
+ ok = torch.allclose(a_deq, b_deq, rtol=5e-2, atol=5e-2)
301
assert ok
302
if add_residual:
303
assert torch.allclose(ref_residual, ops_residual)
tests/kernels/core/test_layernorm.py
+17
-7
@@ -5,7 +5,7 @@ import pytest
5
import torch
6
7
from tests.kernels.quant_utils import FP8_DTYPE
8
-from tests.kernels.utils import opcheck
8
+from tests.kernels.utils import fp8_ulp_distance, opcheck
9
from vllm import ir
10
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
11
from vllm.platforms import current_platform
@@ -204,12 +204,22 @@ def test_fused_rms_norm_quant(
204
(out_quant_fused, x, weight, quant_scale_t, 1e-6),
205
)
206
207
- torch.testing.assert_close(
208
- out_quant.to(dtype=torch.float32),
209
- out_quant_fused.to(dtype=torch.float32),
210
- atol=1e-3,
211
- rtol=1e-3,
212
- )
207
+ if current_platform.is_rocm():
208
+ # Fused and unfused FP8 paths can land on opposite sides of an E4M3 tie;
209
+ # tolerate a tiny number of isolated fp8 outliers on ROCm.
210
+ ulp = fp8_ulp_distance(out_quant, out_quant_fused)
211
+ max_outliers = ulp.numel() // 100_000 + 8
212
+ num_outliers = int((ulp > 0).sum().item())
213
+ assert num_outliers <= max_outliers, (
214
+ f"FP8 quant mismatch: {num_outliers} fp8 outliers (allowed {max_outliers})"
215
+ )
216
+ else:
217
+ torch.testing.assert_close(
218
+ out_quant.to(dtype=torch.float32),
219
+ out_quant_fused.to(dtype=torch.float32),
220
+ atol=1e-3,
221
+ rtol=1e-3,
222
+ )
223
224
225
@torch.inference_mode()
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+4
-32
@@ -19,6 +19,7 @@ The kernel is imported via
19
import pytest
20
import torch
21
22
+from tests.kernels.utils import bf16_ulp_distance, fp8_ulp_distance
23
from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
get_fp8_min_max,
25
)
@@ -160,35 +161,6 @@ def _call_fused(
161
)
162
163
163
-def _bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
164
- """Representable-step distance between two bf16 tensors.
165
-
166
- Reinterprets the bf16 bit patterns under the IEEE-754 total ordering so
167
- that adjacent representable values differ by exactly 1.
168
- """
169
-
170
- def key(t: torch.Tensor) -> torch.Tensor:
171
- u = t.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF
172
- return torch.where(u >= 0x8000, 0xFFFF - u, u + 0x8000)
173
-
174
- return (key(a) - key(b)).abs()
175
-
176
-
177
-def _fp8_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
178
- """Representable-step distance between two 8-bit fp8 tensors.
179
-
180
- Reinterprets the fp8 bytes under a sign-magnitude total ordering so that
181
- adjacent representable values differ by exactly 1. Inputs must already share
182
- the same fp8 encoding (e.g. both FP8_STORE_DTYPE).
183
- """
184
-
185
- def key(t: torch.Tensor) -> torch.Tensor:
186
- u = t.contiguous().view(torch.uint8).to(torch.int64)
187
- return torch.where(u >= 0x80, 0xFF - u, u + 0x80)
188
-
189
- return (key(a) - key(b)).abs()
190
-
191
-
164
def _as_stored_fp8(t: torch.Tensor) -> torch.Tensor:
165
"""Reinterpret a float8_e4m3fn-typed kernel output under the real (FNUZ on
166
gfx942) encoding the kernel actually wrote, without touching the bytes."""
@@ -235,7 +207,7 @@ def _assert_kv_cache_parity(
207
rec_fused[:, :NOPE_DIM], rec_ref[:, :NOPE_DIM], rtol=0, atol=0
208
)
209
max_ulp = int(
238
- _bf16_ulp_distance(rec_fused[:, NOPE_DIM:], rec_ref[:, NOPE_DIM:]).max().item()
210
+ bf16_ulp_distance(rec_fused[:, NOPE_DIM:], rec_ref[:, NOPE_DIM:]).max().item()
211
)
212
assert max_ulp <= 1, f"RoPE bf16 region differs by {max_ulp} ULP (>1)"
213
@@ -709,7 +681,7 @@ def test_full_cache_per_tensor_fp8_matches_reference(
681
# reduction and RoPE rotation can land the kernel and the torch reference on
682
# opposite sides of an fp8 round-to-nearest tie, so allow <=1 fp8 ULP.
683
q_fused = _as_stored_fp8(q_fp8_fused)
712
- q_max_ulp = int(_fp8_ulp_distance(q_fused, q_fp8_ref).max().item())
684
+ q_max_ulp = int(fp8_ulp_distance(q_fused, q_fp8_ref).max().item())
685
assert q_max_ulp <= 1, f"Q fp8 differs by {q_max_ulp} ULP (>1)"
686
687
# K-cache NoPE region [0, NOPE_DIM) is a deterministic per-tensor fp8 quant
@@ -723,7 +695,7 @@ def test_full_cache_per_tensor_fp8_matches_reference(
695
atol=0,
696
)
697
k_max_ulp = int(
726
- _fp8_ulp_distance(k_fused[..., NOPE_DIM:], k_cache_ref[..., NOPE_DIM:])
698
+ fp8_ulp_distance(k_fused[..., NOPE_DIM:], k_cache_ref[..., NOPE_DIM:])
699
.max()
700
.item()
701
)
tests/kernels/utils.py
+29
@@ -809,6 +809,35 @@ def fp8_allclose(
809
)
810
811
812
+def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
813
+ """Representable-step distance between two bf16 tensors.
814
+
815
+ Reinterprets the bf16 bit patterns under the IEEE-754 total ordering so
816
+ that adjacent representable values differ by exactly 1.
817
+ """
818
+
819
+ def key(t: torch.Tensor) -> torch.Tensor:
820
+ u = t.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF
821
+ return torch.where(u >= 0x8000, 0xFFFF - u, u + 0x8000)
822
+
823
+ return (key(a) - key(b)).abs()
824
+
825
+
826
+def fp8_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
827
+ """Representable-step distance between two 8-bit fp8 tensors.
828
+
829
+ Reinterprets the fp8 bytes under a sign-magnitude total ordering so that
830
+ adjacent representable values differ by exactly 1. Inputs must already share
831
+ the same fp8 encoding (e.g. both FP8_STORE_DTYPE).
832
+ """
833
+
834
+ def key(t: torch.Tensor) -> torch.Tensor:
835
+ u = t.contiguous().view(torch.uint8).to(torch.int64)
836
+ return torch.where(u >= 0x80, 0xFF - u, u + 0x80)
837
+
838
+ return (key(a) - key(b)).abs()
839
+
840
+
841
# Marlin MoE test utils
842
843