Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
gdevenyi wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

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

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions tests/models/qwen4_exp/test_qsa_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""

Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down