From 9b17220f4a72d15debe5fcca6894fe34e23da45e Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 17:31:34 -0400 Subject: [PATCH] test(qwen4_exp): compare chunked prefill to one-shot at bf16 tolerance `torch.equal` fails for all three cut points on torch 2.11.0+cu130 / flashinfer 0.6.18 / triton 3.6.0 (sm_89): cuBLAS picks a different kernel for M=5000 and M=5000-cut, so the indexer's q and k round differently, a few rows flip one block at the top-k margin, and the layer output differs by up to 2.6e-3 (a few bf16 ulps). Both paths are deterministic run to run. Use the tolerance of the neighbouring oracle tests. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- tests/models/qwen4_exp/test_qsa_backend.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py index 1d3b944ce..d7643cf0c 100644 --- a/tests/models/qwen4_exp/test_qsa_backend.py +++ b/tests/models/qwen4_exp/test_qsa_backend.py @@ -4,7 +4,8 @@ tokens every complete block is selected, so QSA IS dense attention: the selection must be exactly the causal prefix and the layer output must match ``TorchDenseQSAReference`` (fp32) and a flashinfer dense run over the same pool; -(b) chunked prefill at unaligned cut points equals one-shot prefill (the dual-source compress); +(b) chunked prefill at unaligned cut points matches one-shot prefill to bf16 tolerance (the + dual-source compress); (c) a captured decode replay equals the eager decode step. """ @@ -133,7 +134,11 @@ def test_flashinfer_dense_matches_the_sparse_path(): @requires_cuda @pytest.mark.parametrize("cut", [1001, 4096, 4097], ids=["unaligned", "page-boundary", "boundary+1"]) def test_chunked_prefill_matches_one_shot(cut: int): - """Cut points that are not multiples of index_ratio exercise the dual-source compress.""" + """Cut points that are not multiples of index_ratio exercise the dual-source compress. + + Not bit for bit: the projections see a different M (5000 vs 5000 - cut rows) and cuBLAS + may pick a different kernel per shape, so the indexer's q and k round differently and a + few rows flip one block at the top-k margin. Same tolerance as the oracle tests above.""" config = parsed_config() fixture = Fixture(config, num_pages=512) attn = fixture.layer(QSA_LAYER) @@ -145,7 +150,7 @@ def test_chunked_prefill_matches_one_shot(cut: int): attn.forward(x[:cut], fixture.batch([head], "prefill")) tail = fixture.req(1, cut, length) got = attn.forward(x[cut:], fixture.batch([tail], "prefill")) - assert torch.equal(got, one_shot[cut:]) + torch.testing.assert_close(got.float(), one_shot[cut:].float(), rtol=2e-2, atol=2e-2) @requires_cuda