fix: use a bf16-relative tolerance in the DSA indexer K kernel test

[alphabetc1]

Motivation

test/registered/kernels/ops/attention/test_dsv32_indexer_fusion.py::test_k_norm_rope_matches_reference is flaky in base-b / 1-gpu-large at roughly a 1% rate. A recent instance:

AssertionError: Tensor-likes are not close!
Mismatched elements: 1 / 4736 (0.0%)
Greatest absolute difference: 0.061458587646484375 at index (24, 34) (up to 0.06 allowed)
Greatest relative difference: 0.002833961509168148 at index (24, 34) (up to 0.0 allowed)

The kernel is correct; the test's tolerance model is not.

The test compares a bf16 kernel output against an fp32 reference with a purely absolute tolerance:

torch.testing.assert_close(out.float(), ref, atol=0.06, rtol=0.0)

bf16 keeps 8 significand bits, so its quantization step is relative: for outputs in [16, 32) one ulp is 0.125 and half an ulp is 0.0625, which already exceeds atol=0.06. The single fp32 → bf16 rounding at the end of fused_k_indexer_norm_rope can therefore break the assertion on its own, with no numerical error in the kernel at all.

The test reaches those magnitudes routinely: cos/sin here are randn (not real cosines bounded by 1) and the LayerNorm gamma/beta are randn as well, so outputs reach |v| ~ 30.

Two further points made this show up as a random CI failure rather than a deterministic one:

• _make_inputs() uses a seeded generator, but key / weight / bias were drawn from the global RNG, so every CI run used different inputs.
• The same atol=0.06 is used for the second bf16 comparison in test_indexer_uses_replaced_rope_cache_for_fused_kernels, which passes today only because its inputs happen to be seeded.

The tolerance has been this way since the test was introduced in #27705; this is not a recent regression.

Modifications

• Replace the absolute tolerance at both bf16-vs-fp32 comparison sites with a bf16-aware one: rtol = 2**-7 (2x margin over the 2**-8 pure-rounding bound) plus atol = 1e-3 for values driven near zero by cancellation.
• Seed key / weight / bias in test_k_norm_rope_matches_reference with the same generator convention the rest of the module already uses, so a failure is reproducible.

No kernel or runtime code is touched.

Accuracy Tests

Measured on an H200, replaying the exact test body.

Reproduction of the flake with the current tolerance (400 draws, unseeded inputs, as CI runs it today):

failures at atol=0.06, rtol=0   : 4 (1.0%)
worst max_abs_err               : 0.062300 at ref=-22.3127
  bf16 half-ulp at that ref     : 0.062500
worst relative error            : 9.615e-03
max |ref| seen                  : 30.426
trials whose err exceeds the bf16 half-ulp bound: 0

The last line is the point: across 400 draws, no element ever deviated from the fp32 reference by more than half a bf16 ulp, so the whole discrepancy is representation, not kernel error.

The new tolerance (2000 draws, unseeded inputs):

trials 2000: failures at (atol=1e-3, rtol=2^-7) = 0
worst err/allowed ratio = 0.490  (>=1.0 would fail)
pinned-seed(1) err/allowed ratio = 0.459

The check never came closer than 2x to failing, and the seed the test now pins sits at the same margin.

The test file itself, run 10 consecutive times on an H200:

7 passed  x 10

Speed Tests and Profiling

Not applicable — test-only change.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #32810657006
Latest PR Test (Extra): ❌ Run #32810656786
Latest PR Test (AMD ROCm 7.2): ❌ Run #32810656953

[alphabetc1]

/rerun-test test_dsv32_indexer_fusion.py

[github-actions]

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit 1ec20fd.

Re-run was not dispatched. What to do:

• Rebase your branch onto the latest main and push again
• Follow issue #21065 for context
• CI-fix PRs may request the bypass-maintenance label to skip this check

[alphabetc1]

/rerun-test test_dsv32_indexer_fusion.py

[github-actions]

Results for /rerun-test test_dsv32_indexer_fusion.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_dsv32_indexer_fusion.py

[alphabetc1]

/rerun-test test_dsv32_indexer_fusion.py

[github-actions]

Results for /rerun-test test_dsv32_indexer_fusion.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_dsv32_indexer_fusion.py