test(qwen4_exp): compare chunked prefill to one-shot at bf16 tolerance

[gdevenyi]

What

test_chunked_prefill_matches_one_shot asserts torch.equal between the chunked and the one-shot prefill output of the QSA layer. On torch 2.11.0+cu130 / flashinfer 0.6.18 / triton 3.6.0 (RTX 6000 Ada, sm_89) all three cut points fail on main (af71ba4), by a few bf16 ulps. This compares at the tolerance the neighbouring oracle tests in the same file already use (rtol=2e-2, atol=2e-2) and says why in the docstring.

Why this is the test's contract and not a chunked-prefill bug

Measured on plain main with the test's own fixture, seed and cut points:

• Both paths are deterministic: one-shot vs one-shot and chunked vs chunked are bit-equal run to run.
• The rounding comes from GEMM shape. qkv_proj over the same rows gives different bf16 results for M=5000 and M=904 (max diff 3.9e-3, one ulp at that magnitude); the indexer's q and k differ for every cut. cuBLAS picks its kernel per shape, which a chunked prefill cannot avoid.
• Those indexer differences flip one block at the top-k margin in a handful of rows; no row moves by more than 1e-2.

cut	rows compared	rows with a different block set	max abs diff (output)	qkv_proj max diff, M=5000 vs M=rows
1001 (unaligned)	3999	3	2.58e-3	0 (bit-equal)
4096 (page-boundary)	904	1	2.56e-3	3.9e-3
4097 (boundary+1)	903	1	2.56e-3	3.9e-3

The bf16 defaults of assert_close (atol 1e-5) also fail, on elements near 1e-6, so the explicit tolerance is the one that matches the rest of the file.

Testing

tests/models/qwen4_exp/test_qsa_backend.py on the stack above: 9 passed (was 3 failed, 6 passed).

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

[gdevenyi]

Re-tested against current main (fb7f732, i.e. after the #418 / #427 / #426 quantization refactor) on 2 x RTX 6000 Ada, TP=2 box.

Method. This PR's head merged onto main, then the full pytest tests suite. The run is CUDA-hidden (CUDA_VISIBLE_DEVICES="") on purpose: this box is serving a model on both GPUs, and with them visible 65-95 GPU tests fail on main itself with AcceleratorError: out of memory, with the count swinging ~10 between identical runs. Hiding CUDA makes the result deterministic, so a failure-set difference against main means something. Baseline: main = 1205 passed, 350 skipped, 0 failed.

Result: 1205 passed, 350 skipped, no new failures.

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

[gdevenyi]

Closing this. It was written with heavy AI assistance, and the maintainers have indicated that do not want such contributions.

The description and the diff stay here for anyone who wants to pick the idea up.