[Test] Sweep the MQA logits indexer's M and N across prefill and decode - #5434
Open
zhiding512 wants to merge 2 commits into
Open
zhiding512 wants to merge 2 commits into
zhiding512 wants to merge 2 commits into
Conversation
The MQA-logits lightning indexer has its head geometry fixed by the model --
single-head KV, head_dim 128, 32 or 64 Q heads -- so M (query rows) and N (KV
length) are the axes that decide which kernel config runs. Nothing in-tree
swept them: the Triton UT tops out at (1024, 1560) and the two benchmarks take
one shape per process.
Add op_tests/test_mqa_logits.py, one table per phase, because prefill and
decode are different kernels with different calling conventions:
prefill fp8_mqa_logits, contiguous KV, causal chunked-prefill window
decode deepgemm_fp8_paged_mqa_logits, paged cache, ATOM's convention
(preshuffled 64-token blocks, ChunkK=256, WavePerEU=2, caller-owned
output buffer, fragmented block tables)
Both references dequantize the same fp8 bytes the kernels read, so err
isolates the kernel from the quantization. Prefill's reference is row-sampled:
the full [H, M, N] fp32 score tensor is ~69 TB at H=32/M=8192/N=65536, so the
sample takes the grid edges and the BLOCK_M=2 block seam first.
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Correctness assertions, decode mask validation, and benchmark methodology/reporting require fixes.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds a configurable M×N correctness and performance sweep for MQA logits prefill and decode kernels.
Changes:
- Adds sampled prefill and paged decode validation.
- Adds randomized cache construction and configurable shape sweeps.
- Reports correctness and benchmark metrics.
File summaries
| File | Description |
|---|---|
op_tests/test_mqa_logits.py |
New prefill/decode sweep, validation harness, and reporting. |
Review details
Suppressed comments (4)
op_tests/test_mqa_logits.py:192
- [verified] The PR description says 64-head prefill at M=1024 is 480–630 TFLOPS and below the 32-head results, but the included sweep table reports 1252–1618 TFLOPS for those same rows; this code computes the column directly as
flops / us / 1e6. The tuning conclusion is therefore not reproducible from the stated sweep. Author must reconcile the table, measurement run, and performance narrative.
ret[f"{name} TFLOPS"] = flops / us / 1e6
ret[f"{name} TB/s"] = nbytes / us / 1e6
op_tests/test_mqa_logits.py:50
- The stated size is incorrect:
[32, 8192, 65536]fp32 elements occupy 68,719,476,736 bytes (about 64 GiB or 69 GB), not 69 TB. The row-sampling rationale remains valid, but this comment overstates the reference memory footprint by three orders of magnitude. Author must correct the unit.
# The full [H, M, N] fp32 score tensor a naive reference builds is ~69 TB at
# H=32/M=8192/N=65536, so prefill accuracy is always checked on a row sample.
op_tests/test_mqa_logits.py:193
checkAllclosereturns a mismatch ratio and only raises for catastrophic errors; with its defaultcatastrophic_check=False, even a 100% mismatch merely populates thiserrfield. Since this is the numerical correctness check for the sampled prefill rows, the sweep can report a nonzero error and still exit successfully. Author must assert the returned ratio is within the intended tolerance before recording the row.
ret[f"{name} err"] = err
op_tests/test_mqa_logits.py:268
checkAllclosereturns a mismatch ratio and only raises for catastrophic errors; with its defaultcatastrophic_check=False, even a 100% mismatch merely populates thiserrfield. The decode sweep therefore can log a failed accuracy result and still exit successfully. Author must assert the returned ratio is within the intended tolerance before recording the row.
ret[f"{name} err"] = err
- Files reviewed: 1/1 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+159
to
+175
| candidates = { | ||
| "triton": lambda: fp8_mqa_logits( | ||
| q_fp8, kv_fp8, kv_scales, weights, ks, ke, clean_logits | ||
| ), | ||
| } | ||
|
|
||
| valid = int((ke.clamp(max=n) - ks).clamp(min=0).sum().item()) | ||
| n_pad = (n + 255) // 256 * 256 | ||
| flops = 2.0 * num_heads * head_dim * valid | ||
| read_bytes = q_fp8.nbytes + kv_fp8.nbytes + kv_scales.nbytes + weights.nbytes | ||
| # Compulsory traffic only. KV is re-read per query row, so once N*128 stops | ||
| # fitting in L2 the achieved bandwidth is well above this number. | ||
| nbytes = read_bytes + (m * n_pad * 4 if clean_logits else valid * 4) | ||
|
|
||
| ret = {"gfx": get_gfx(), "ref_rows": len(rows)} | ||
| for name, fn in candidates.items(): | ||
| out, us = run_perftest(fn, num_rotate_args=_rotate_args(read_bytes)) |
Comment on lines
+258
to
+264
| err = checkAllclose( | ||
| ref[inside], | ||
| out[inside], | ||
| rtol=1e-2, | ||
| atol=1e-2, | ||
| msg=f"{name}: decode b={batch} n={n} h={num_heads} mtp={next_n}", | ||
| ) |
black 26.5.1 and ruff 0.15.18 disagree on the assert wrap, and running ruff format after black undid black's choice. CI checks black and ruff check (pinned 0.16.0), not ruff format, so black goes last. 0.16.0 also adds I001, which wants aiter in its own first-party block.
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Four unresolved findings remain, including three critical issues and one moderate benchmark-validity issue.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (3)
op_tests/test_mqa_logits.py:263
- [verified] The decode result is likewise only reported: for example, one completely wrong row at
batch=64,next_n=1is belowcheckAllclose's 5% non-raising threshold, so the advertised correctness sweep can pass with a bad row. Author must asserterr == 0(or otherwise make any mismatch fail) before returning the row.
err = checkAllclose(
ref[inside],
out[inside],
rtol=1e-2,
atol=1e-2,
op_tests/test_mqa_logits.py:263
- [verified] This decode validation has two holes: it only compares
ref[inside]/out[inside], ignoring any finite logits outside the causal window, and it never asserts the nonzero error fraction returned bycheckAllclose. A future-token write or any numeric mismatch can therefore pass CI. Author must compareout == -float("inf")with~insideand fail whenerr != 0.
err = checkAllclose(
ref[inside],
out[inside],
rtol=1e-2,
atol=1e-2,
op_tests/test_mqa_logits.py:176
- [verified]
_rotate_args()is passed torun_perftestwith zero-argument closures here (and at the decode call on line 258), butrun_perftestonly deep-copies tensors in its explicitargs; it never sees the captured Q/KV/weights/output. The intended cache rotation is therefore never performed, so small shapes can measure warm-cache/L2-hit latency and make the reported bandwidth numbers unreliable. Author must pass the benchmark tensors as explicitrun_perftestarguments or explicitly rotate the captured inputs.
out, us = run_perftest(fn, num_rotate_args=_rotate_args(read_bytes))
- Files reviewed: 1/1 changed files
- Comments generated: 3
- Review effort level: Lite
Comment on lines
+184
to
+190
| err = checkAllclose( | ||
| ref[inside], | ||
| got[inside].float(), | ||
| rtol=1e-2, | ||
| atol=1e-2, | ||
| msg=f"{name}: prefill m={m} n={n} h={num_heads}", | ||
| ) |
|
|
||
|
|
||
| @benchmark() | ||
| def test_mqa_logits_decode(batch, n, num_heads, head_dim, next_n, kv_block): |
Comment on lines
+233
to
+234
| out_logits = torch.full((rows, n), -float("inf"), dtype=dtypes.fp32) | ||
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
op_tests/test_mqa_logits.py, an M x N sweep for the MQA-logits lightningindexer, with one table for prefill and one for decode.
Motivation
The model fixes the indexer's head geometry -- KV is a single head,
head_dimis 128, Q has 32 or 64 heads -- so M (query rows) and N (KV length) are the axes
that decide which kernel config actually runs. Nothing in-tree sweeps them:
op_tests/triton_tests/attention/test_fp8_mqa_logits.pyis a correctness UTthat tops out at
(s_q, s_k) = (1024, 1560), plus two >2 GiB regression shapes.bench_fp8_mqa_logits.pyandbench_deepgemm_attention.pyeach take oneshape per process and report no accuracy.
So there was no single place to answer "what does this op do across the shapes
a model actually serves, and is it still correct there".
Changes
One new file. Prefill and decode are different kernels with different calling
conventions, so they get one
@benchmarkfunction and one table each:fp8_mqa_logits[N, 128], Q[M, H, 128], causal chunked-prefill windowdeepgemm_fp8_paged_mqa_logits[B, next_n, H, 128], M =B * next_nDecode reproduces ATOM's
Indexer._score_topk_decodeconvention: preshuffled64-token KV blocks,
ChunkK=256,WavePerEU=2, caller-owned output buffer, andrandomized (fragmented) block tables rather than a sequential one a real KV
cache never hands the kernel.
Two things worth flagging for review:
errmeasures the kernel, not the quantization. Both referencesdequantize the same fp8 bytes the kernels read, so a non-zero
errcolumn isa real bug rather than fp8 rounding noise.
[H, M, N]fp32 score tensor is ~69 TB at
H=32/M=8192/N=65664. The sample takes thegrid edges and the
BLOCK_M=2block seam first (where a row-indexing buglands), then an even spread; the
ref_rowscolumn records how many rows wereactually compared.
TB/sis compulsory traffic only. Prefill re-reads KV per query row, so thatcolumn is a lower bound there and TFLOPS is the metric to read; decode is
genuinely bandwidth-bound and TB/s is the one that matters.
How to run it
Every flag is a swept list (
-m,-n,-b,-hq,-dh,-mtp,-kb,-c,-p), so a different model's shapes go in without touching the file. Prefillskips
n < m;-nis shared by both phases.Performance
Environment: MI355X (gfx950, 256 CU, 288 GB HBM), run on this branch. All 76 shapes report
err == 0. Both tables below come from one invocation:Prefill: 64 heads is the faster configuration at every shared shape (18/18 of them), at 1215-1990 TFLOPS against 32 heads' 927-1567. Both climb with M --
BLOCK_M=2engages above M=4096, and the slowest column is M=1024 where the grid cannot fill 256 CUs.Decode is bandwidth-bound: 1.9-6.2 TB/s of compulsory traffic once batch >= 16. 64 heads roughly doubles decode TFLOPS there (229-740 against 130-373) at the same TB/s, so head count is nearly free and KV traffic is the whole story.
batch=1never reaches that band -- it is launch-bound, with N from 4K to 128K moving latency only from 3.8 to 4.9 us at 32 heads, and 4.3 to 5.8 us at 64.prefill --
fp8_mqa_logits, 36 shapesdecode --
deepgemm_fp8_paged_mqa_logits, 40 shapesTesting
err == 0, on this branchblack --checkandruff check(0.16.0, the version CI pins) clean on thenew file
SUPPORTED_GFXgates it in and the Triton wrappersdispatch per arch, but op_tests CI only has MI35x runners, so nobody has
run it there yet.
Follow-ups (not in this PR)
length, so the kernel's varctx scheduling path is never exercised.
-mtpdefaults to 1, so thenext_n > 1decode layout is untested.gfx942-only and I have no gfx942 to validate it on, so it is left out rather
than shipped unrun.