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