test(st): right-size too-tight golden tolerances across s-tests (bf16/fp16/fp32-K-split)#1945
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the test configuration in test_batch_paged_attention.py by adjusting the absolute and relative tolerance (atol and rtol) to 1e-3 for BF16-precision. This change addresses intermittent test flakiness caused by hardware approximation differences in BF16 calculations. There are no review comments, and I have no additional feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7cd7f45a59
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self.config.atol = 1e-3 | ||
| self.config.rtol = 1e-3 |
There was a problem hiding this comment.
These tolerances still do not cover the BF16 one-ULP mismatch the comment is trying to admit. If a pij_batch value, or a lij_batch summand, around 0.5 lands on opposite sides of a BF16 rounding boundary, the actual/expected difference is 2^-8 ≈ 3.9e-3, while torch.allclose with rtol=atol=1e-3 only allows about 1.5e-3 for a 0.5 value. That leaves the same hardware-exp-vs-torch-exp flake in place for different random inputs; use a BF16-scale tolerance or per-output tolerance large enough for at least one BF16 ULP.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch — addressed in 0d711ec9.
You're right that rtol=atol=1e-3 doesn't cover a full BF16 ULP for values near 0.5: a summand/pij_batch element in [0.5, 1) has ULP 2^-8 ≈ 3.9e-3, while 1e-3 only admits 1e-3 + 1e-3·0.5 = 1.5e-3 there — so the hardware-exp-vs-torch-exp flake would survive for different random inputs.
Fix: rtol now covers ≥ one BF16 ULP relative everywhere. BF16 eps = 2^-7 ≈ 7.8e-3 (max relative ULP, at the low end of each binade), so I set rtol=1.6e-2 (PyTorch's bf16 assert_close default, ~2-ULP margin) with atol=1e-3 for near-zero summands. At the 0.5 worst case the threshold is now 1e-3 + 1.6e-2·0.5 = 9e-3 > 3.9e-3 (covers even 2 ULP). The comment and commit message were updated to reflect the ULP analysis.
📝 WalkthroughWalkthroughThis change updates a single test case's numerical comparison tolerances. ChangesTest tolerance adjustment
Estimated code review effort: 1 (Trivial) | ~3 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…sion pij_batch is a BF16 output and lij is a row_sum of BF16-quantized exp values (exp -> cast BF16 -> FP32 -> row_sum), but the test used the default rtol=atol=1e-5. The device's hardware exp is a ~1-2 ULP approximation, so an exp value near a BF16 rounding boundary rounds to a different BF16 than the torch.exp golden -- a full one-ULP mismatch on that element (observed on a2a3: 1/32 elements, diff = exactly 2^-10). BF16 has 8 significand bits, so eps = 2^-7 ~= 7.8e-3 and one ULP is up to 2^-8 ~= 3.9e-3 absolute near 0.5. A tolerance that admits this flake must allow >= one BF16 ULP relative, so set rtol=1.6e-2 (PyTorch's bf16 assert_close default, ~2 ULP margin) + atol=1e-3 -- e.g. at 0.5 the threshold is 9e-3 > 3.9e-3. The previous 1e-5 (and an intermediate 1e-3) were too tight to cover a single BF16 ULP.
tonibohnlein
force-pushed
the
fix/batch-softmax-prepare-bf16-tolerance
branch
from
7cd7f45 to
0d711ec
Compare
…-K-split) Companion to the batch_softmax_prepare fix: an audit of tests/st/runtime/ found the same too-tight-golden-tolerance flake class in other kernels that validate reduced-precision compute against the default rtol=atol=1e-5 (or a bf16-insufficient 1e-3). Right-size each to admit >= one ULP of the precision it actually computes in: - fp16 transcendentals (test_unary_math TileSqrtTestCase, fp16 sqrt): the HW Newton approx differs from torch.sqrt by ~1 fp16 ULP -> rtol=atol=1e-3. - fp32 K-split matmuls (test_mat_slice_to_left, test_tensor_batch_matmul, all 5 batch cases): device K-reduction order differs from the single-pass torch golden, drifting ~K*eps_fp32 (~1.5e-5 at K=128) -> rtol=atol=1e-4 (AutoL0). - bf16 softmax/attention (test_paged_attention_multi_config SoftmaxPrepare, test_batch_paged_attention full attention, test_qwen3_decode_scope3_mixed): bf16 exp/matmul output at rtol=1e-3 is below one bf16 ULP (eps_bf16=2^-7) -> rtol 1.6e-2 (softmax-prepare) / 2e-2 (attention), matching existing bf16 s-tests. All one-line tolerance edits; no kernel/codegen change. Surfaced by CI flakes on tensor_batch_matmul (fp32 K-split) and tile_sqrt_fp16.
tonibohnlein
changed the title
- fp16 transcendental (test_unary_math): bump rtol=atol 1e-3 -> 2e-3. 1e-3 gave
only ~1 fp16 ULP of margin, so a full-ULP HW sqrt error plus golden rounding at
a binade bottom could still flake. Fix the ULP comment (max ULP near 1 is
2^-10 ~ 9.8e-4, not 2^-11) and pin the observed CI flake (2^-11 ~ 4.9e-4).
- BatchPagedAttention (full attention): the softmax pij intermediate is FP16
(the assembled program's pij_b = pl.FP16; golden casts through torch.float16),
not BF16 -- fix the comment. rtol=2e-2 stays, justified by the BF16 QK/PV
matmul reductions (bf16 ULP > fp16 pij's).
- Override discipline: guard the BatchSoftmaxPrepare / BatchPagedAttention /
multi_config SoftmaxPrepare tolerance sets with `if kwargs.get("config") is
None` so a caller-supplied (e.g. stricter) config is respected, matching the
matmul/sqrt cases.
- Pin the observed max diff on the fp32 K-split matmul comment (1.685e-5).
Hzfengsy
left a comment
Hzfengsy
left a comment
There was a problem hiding this comment.
Test tolerance right-sizing looks sound; the BF16-ULP tolerance concern from the automated reviewer was addressed. LGTM.
2 participants
Summary
Several device s-tests validate reduced-precision compute against a golden with a tolerance tighter than one ULP of the precision the kernel actually computes in — most at the default
rtol=atol=1e-5, some at a bf16-insufficient1e-3. That guarantees intermittent single-element flakes whenever a value lands on a rounding boundary (the device's hardwareexp/sqrt/matmul-reduction rounds differently than torch). This PR right-sizes each to admit ≥ 1 ULP of its precision.This started from
batch_softmax_prepare(bf16 softmax at1e-5) and, after an audit oftests/st/runtime/, was extended to the rest of the same flake class — two of which have since flaked live in CI (tensor_batch_matmulfp32 K-split;tile_sqrt_fp16).Tolerance policy (admit ≥ 1 ULP of the compute precision)
1.6e-2 / 1e-3(softmax) ·2e-2 / 1e-3(attention)ε = 2⁻⁷ ≈ 7.8e-3; one ULP up to2⁻⁸ ≈ 3.9e-3.1.6e-2= PyTorch bf16 default;2e-2matches existing bf16 attention s-tests2e-3 / 2e-3sqrtHW Newton approx differs fromtorch.sqrtby ~1–2 fp16 ULP (max ULP near 1 is2⁻¹⁰ ≈ 9.8e-4).2e-3clears the observed flake (2⁻¹¹ ≈ 4.9e-4) with ~8× margin —1e-3gave only ~1 ULP, too thin1e-4 / 1e-4~K·eps_fp32 ≈ 1.5e-5at K=128 (larger for cancellation-heavy elements). Same class the AutoL0 matmul s-tests already loosenTests fixed
batch_softmax_prepare(1.6e-2);test_paged_attention_multi_config::SoftmaxPrepare(+ its PTOAS subclass, 1.6e-2);test_batch_paged_attention::BatchPagedAttentionfull attention (2e-2 — BF16 QK/PV matmuls dominate; its softmaxpijintermediate is FP16);test_qwen3_decode_scope3_mixed(2e-2). Tolerance sets guarded byif kwargs.get("config") is Noneso a caller-supplied config is respected.test_unary_math::TileSqrtTestCasefp16 sqrt (set in_UnaryMathBase.__init__fordtype==FP16, so it also covers any future fp16 sin/cos).test_mat_slice_to_left; all 5test_tensor_batch_matmulfp32 batch-matmul cases.Each is a one-line tolerance edit (
self.config.rtol/atolor theRunConfig(...)literal). No kernel or codegen change.Deliberately not changed (audit found these safe)
Pure data-movement bf16/fp16 outputs (scatter/gather/copy/pad/reshape/transpose) are device-exact; device-order-matched goldens (e.g.
test_col_reductionreplicates the fp16 reduction tree); constant/integer-input matmuls (bf16-exact products); and ops already appropriately loosened. Two fp16 elementwise cases (test_vector_miscmuls / col_expand_add) are borderline — a single fp16 mul/add is normally bit-exact — so they're left as-is pending an actual flake, to avoid loosening genuinely-exact ops.Why the earlier reviewed part changed
The
batch_softmax_preparevalue was corrected from1e-3up to1.6e-2after review (a1e-3threshold at 0.5 is1.5e-3, below one bf16 ULP of3.9e-3). The same reasoning is applied consistently across the bf16 tests here.