[ROCm][CI] Use explicit wvSplitKrc skinny-GEMM test tolerance for bf16 (gfx950) (#49309)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Jul 31, 2026 at 05:25 UTC
2773ec31b3c75b994351643e25e019e1168e6aed
1 file changed
+8
-3
tests/kernels/quantization/test_rocm_skinny_gemms.py
+8
-3
@@ -150,8 +150,8 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode
150
xavier = (
151
math.sqrt(2 / k) if xnorm else 1
152
) # normalize to avoid large output-bias deltas
153
- A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
154
- B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
153
+ A = torch.randn(n, k, dtype=dtype, device="cuda") * xavier
154
+ B = torch.randn(m, k, dtype=dtype, device="cuda") * xavier
155
if padded_a:
156
A = pad_fp8(A)
157
@@ -167,7 +167,12 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode
167
out = ops.wvSplitKrc(A, B, cu_count, BIAS)
168
169
if xnorm:
170
- torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)
170
+ # The O(1) bias lifts outputs to ~O(1), where one bf16 ULP (~3.9e-3, the
171
+ # worst measured divergence) exceeds 1e-3. Bump atol to 5e-3 only for
172
+ # biased bf16 (above that ULP, still under finfo(bf16).eps); keep 1e-3
173
+ # otherwise.
174
+ atol = 5e-3 if (dtype == torch.bfloat16 and BIAS is not None) else 1e-3
175
+ torch.testing.assert_close(out, ref_out, atol=atol, rtol=1e-8)
176
else:
177
torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2)
178