perf(gqa): autotune the decomposed prefill GEMMs instead of trusting one heuristic pick - #718
Draft
BoarQing wants to merge 2 commits into
Draft
perf(gqa): autotune the decomposed prefill GEMMs instead of trusting one heuristic pick#718BoarQing wants to merge 2 commits into
BoarQing wants to merge 2 commits into
Conversation
…ristic pick The decomposed GQA pipeline asked hipBLASLt for exactly one heuristic candidate and used it unconditionally. Nothing ever measured it, so a misprediction was invisible -- and it was mispredicting badly: at a 16K prompt the score GEMM's first-ranked algorithm runs 30-32 ms against 8.5-9.9 ms for the fastest candidate of the same shape (3.2-4.5x), which is what made the score GEMM cost 26.25 ms/head against the value GEMM's 9.76 ms for an identical 137.5 GFLOP. The value GEMM was mispredicted too, by 1.4-1.8x. Ports matmul.cpp's approach (MAX_ALGO_CANDIDATES=60, AUTOTUNE_TIMING_ITERS=3) into the GQA GEMM cache. Split in two, because the shape is known where the descriptors are built but the device pointers are not: queryOrCreateGemmState now keeps the whole candidate list, and the first real GEMM of that shape times them with its own operands and caches the winner. Candidate 0 stays active until tuning replaces it, so any failure degrades to the old behaviour. Warm-up doubles as the validity check -- a candidate that cannot run these operands is skipped rather than timed. Two details the port could not inherit: - Workspace. Sizing to the incumbent's requirement would skip the hungrier candidates and narrow the sweep to small-workspace algorithms, so an untuned shape reserves the largest workspace any candidate asks for. - Sweep cost. matmul.cpp's fixed 3 timed iterations assume microsecond launches; these are tens of ms, where 60 candidates x 3 is seconds per shape. Probing the incumbent first and dropping to one timed iteration above 1 ms keeps the sweep affordable without measuring noise. HIPDNN_EP_GQA_AUTOTUNE overrides the runtime-wide HIPDNN_EP_AUTOTUNE for GQA alone. Flipping the shared flag would retune every other GEMM in the model, so a TTFT delta could not be attributed here; the override gives both A/B legs from one binary. Measured, same binary, Gemma-4 26B-A4B 16K VLM prefill, 3 reps + 1 warmup: 37,091 ms -> 25,917 ms warm TTFT, an 11.17 s (30.1%) gain. Run-to-run stddev also drops from 371 ms to 23 ms. Tuning cost is paid once per shape, on the warmup iteration. Numerics: greedy output tracks the untuned run for 25 of 32 tokens and then takes a different, coherent continuation, consistent with fp accumulation order differing between algorithms. This is the same class of change matmul.cpp already makes for every other GEMM in this model. Not verified against a reference at this shape: the decomposed path has no numeric test at head_dim 256 / seqlen 16384 (test/numeric covers attention only to head_dim 128, seqlen 256). Co-authored-by: Cursor <cursoragent@cursor.com>
…e memory Tries the score GEMM with outputFp32=false, which the plan flagged as worth measuring: the fp32 S x S score buffer plus its fp16 probability copy is the decomposed prefill's dominant memory traffic at a long prompt, and a single fp16 buffer should remove roughly half of it. Under HIPDNN_EP_GQA_SCORE_FP16 the GEMM writes fp16 scores and bias, mask and softmax all run in place there, so the fp32 buffer is not allocated at all. The result is negative on time and positive on memory. Measured at the default 16 GiB score budget, same binary, 3 reps + 1 warmup: 26,003 ms without it against 26,178 ms with it -- 175 ms slower, against per-leg stddevs of 29 and 7 ms. The chain is evidently not bandwidth-limited at this shape once the GEMMs are tuned, and the fp32<->fp16 conversion in the bias pass is not free. It is kept, default off, for the memory: halving the per-head score footprint doubles how many heads fit a given score budget, and that budget is the only reason head grouping exists. It is also what a constrained budget wants for speed -- at HIPDNN_EP_GQA_SCORE_BUDGET_MB=1700, where fp32 fits one head per group and fp16 fits two, the same A/B is 25,592 ms against 21,615 ms, a 3.98 s win. So the flag's value is entirely in relieving the budget, not in the narrower dtype, and it should stay off wherever the budget is not binding. Numerics are clean: greedy decode over 32 tokens is token-for-token identical to the fp32-score run at the same 16K prompt. The risk the plan named is real but was not hit -- fp32 scores exist so that softmax and the causal mask cannot produce inf - inf = NaN, and fp16 only bounds that instead of excluding it. The bias add is still evaluated in fp32 before the narrowing store, so a masked entry lands at score + (-65504) as a finite fp16 rather than saturating to -inf; the fp16 causal mask writes -65504 where the fp32 one writes -INFINITY; and all softmax reductions stay fp32 regardless of the buffer dtype, so only the stored logits lose precision. Kernel changes are type parameters on the existing kernels rather than new ones, so every fp32 instantiation is unchanged: softmax_f32_to_out_kernel gains an input type (fp16 in place is safe -- one block owns one column, so a pass-3 store is never read by another block), and add_attention_bias_kernel gains a score type. Correcting the record on the previous commit: its 37,091 -> 25,917 ms was taken with HIPDNN_EP_GQA_SCORE_BUDGET_MB=1700 leaking in from the shell, not at the default budget. Both legs had it, so the A/B stands, but the default-budget figure for that change is 38,012 -> 26,003 ms, a 12.01 s (31.6%) gain. The 38,012 ms untuned baseline also independently reproduces the plan's 38,384 ms starting point. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Thanks for opening a PR! This project follows LLVM's incremental-development and AI-tool-use Before requesting review, please check that:
Reviewers are assigned through |
L2 Accuracy Results (EP vs CPU)
Threshold: 0.01 | Run: 4424 - Commit: |
MorphiZen EP Performance Results
EPContext Export Performance
EPContext Import Performance
OGA Benchmark Results
OGA Wheel Smoke (Python benchmark_e2e.py)
Run: 4424 - Commit: |
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
1 participant
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.
What
GQA's decomposed prefill asked hipBLASLt for exactly one heuristic candidate and never benchmarked it, while
lib/Runtime/real/matmul.cppnext door has had a real autotuner on the same API for a while. This ports that loop intoqueryOrCreateGemmState: request up to 60 candidates, benchmark them on the first real call with the actual operands, cache the winner per shape in the existingGqaGemmCache.The heuristic's first pick was not close. On the score GEMM it ran 30-32 ms against the tuned winner's 8.5-9.9 ms -- a 3.2-4.5x gap on the same 137.5 GFLOP the value GEMM was already doing in 9.76 ms. The score and value GEMMs having identical FLOPs but a 2.7x latency difference was the clue that pointed here.
Tuning happens on real operands rather than scratch buffers, and is cost-aware: a shape whose first candidate exceeds 20 ms gets an abbreviated schedule (2 rounds, 1 iteration, no warmup), because the full 4x30 schedule was calibrated for microsecond launches and cost 13 minutes at a 16K prompt. Workspace sizing reserves the largest candidate's requirement while a shape is still untuned.
Second commit is an honest negative result: asking the score GEMM for fp16 output instead of fp32, so the whole score/bias/softmax chain runs over one fp16 S x S buffer. It halves that buffer and doubles the head-grouping budget, but buys no TTFT (26,003 ms off against 26,178 ms on). Left in, default off, behind
HIPDNN_EP_GQA_SCORE_FP16, for memory-constrained deployments -- fp32 stays the default because it is what keeps softmax clear ofinf - inf = NaN.Measurement
Gemma-4 26B-A4B, 16K VLM prefill, gfx1151, same binary both arms, default score budget, 1 warmup + 3 reps:
HIPDNN_EP_GQA_AUTOTUNE=0-12.01 s (31.6%). Plan estimated 8.7 s.
Verification
HIPDNN_EP_DEBUGconfirms the mechanism, not just the clock:[GQA-AUTOTUNE] score: tested 7/7 algos, best=#5 0.164 ms vs heuristic #0 0.466 ms (2.84x), and 4.41x on another shape. The winner is never candidate #0, which is what the old code took.matmul.cppalready ships this same property.Notes for review
HIPDNN_EP_GQA_AUTOTUNE=0restores the single-heuristic behaviour in the same binary, which is how the number above was taken.Made with Cursor