perf(cuda): retune GQA decode flash-decoding by justinchuby · Pull Request #999 · justinchuby/onnx-genai

justinchuby · 2026-08-15T07:50:41Z

Summary

retunes fp16/bf16 GQA decode split-K from ~2 target waves to ~4 target waves
adds ONNX_GENAI_CUDA_GQA_SPLITS rollback/A-B knob
uses exp2 softmax math and parallelizes split-merge output lanes

Measurements (glm-4-9b-int4, GPU1, graph ON, ONNX_GENAI_LMHEAD_CUBLASLT=1)

KV512 attention core: ~663 us -> ~442 us; attention total ~1058 us -> ~888 us
KV2048 attention core: ~2195 us -> ~1315 us; attention total ~2587 us -> ~1758 us
KV512 e2e: 5.175 ms/token, 193.23 tok/s
KV2048 e2e: 6.035 ms/token, 165.70 tok/s
short base decode: 207.46 tok/s, captures=3 replays=474 fallbacks=0
ORT fair graph-off same GPU: 197.07 tok/s; ORT graph-on still fails in harness with constructed OrtValue error

Validation

cargo build --release -p onnx-genai-bench --features bench-native,bench-ort,cuda,cuda-ort --bin profile_native --quiet
cargo test -p onnx-runtime-ep-cuda --features cuda,gpu-tests --test gqa_fp16_gpu --quiet
cargo test -p onnx-runtime-ep-cuda --features cuda,gpu-tests --test matmul_nbits_marlin_numerics --quiet
cargo fmt --all -- --check
cargo clippy -p onnx-runtime-ep-cuda --features cuda --lib -- -D warnings

Accuracy notes

qwen KV512/KV2048 generated streams matched old split-count control in measured runs
glm generated stream changes vs old split count at long context, expected from split-K softmax reduction order; flagged for Chew numerics gate

Increase the split-K decode occupancy target for fp16/bf16 GQA, add a rollback split-count knob, use exp2 softmax math, and parallelize split-merge output lanes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

codecov · 2026-08-15T07:58:52Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.56%. Comparing base (50668cb) to head (f682494).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #999   +/-   ##
=======================================
  Coverage   83.56%   83.56%           
=======================================
  Files          12       12           
  Lines        5270     5270           
  Branches     5270     5270           
=======================================
  Hits         4404     4404           
  Misses        682      682           
  Partials      184      184

Flag	Coverage Δ
cli-ort-linux	`83.52% <ø> (ø)`
cli-ort-windows	`83.02% <ø> (ø)`

Flags with carried forward coverage won't be shown. Click here to find out more.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

github-actions · 2026-08-15T08:16:23Z

🔴 Benchmark Regression Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
🔴	`qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline`	3.54 ms	4.70 ms	+32.9%
⚠️	`qwen3_sampling_processors/top_k_top_p_full_sort_baseline`	5.30 ms	6.48 ms	+22.4%
⚠️	`qwen3_sampling_processors/top_k_partial_selection`	139.01 µs	164.97 µs	+18.7%
⚠️	`grammar_masking/llguidance_compute_mask/32`	69.79 µs	81.36 µs	+16.6%
⚠️	`qwen3_sampling_processors/top_p_fast_after_top_k`	484.07 µs	558.98 µs	+15.5%
✅	`sampling_latency/greedy_per_token`	3.35 µs	3.81 µs	+13.7%
✅	`qwen3_sampling_processors/top_k_full_sort_baseline`	2.15 ms	2.39 ms	+11.1%
✅	`add/medium_f32_threads=1-internal/262144`	2.33 ms	2.57 ms	+10.4%
✅	`sampling_latency/top_k_per_token`	52.54 µs	57.43 µs	+9.3%
✅	`kv_cache/alloc_dealloc_pages`	38.42 µs	41.68 µs	+8.5%
✅	`tokenization/decode_tokens_per_second`	6.43 ms	6.96 ms	+8.4%
✅	`sampling_latency/top_p_per_token`	375.51 µs	404.14 µs	+7.6%
✅	`add/medium_bf16_threads=1-internal/262144`	2.42 ms	2.56 ms	+5.8%
✅	`qwen3_sampling_processors/top_k_top_p_fast`	634.36 µs	668.43 µs	+5.4%
✅	`logit_processing/seven_processor_chain_per_step`	304.01 µs	319.12 µs	+5.0%
✅	`sampling_latency/min_p_per_token`	211.32 µs	218.47 µs	+3.4%
✅	`reduce_mean/large_f32_threads=1-internal/262144`	909.80 µs	934.40 µs	+2.7%
✅	`gather/small_f16_threads=1-internal/4096`	446.5 ns	454.8 ns	+1.9%
✅	`add/medium_f16_threads=1-internal/262144`	2.42 ms	2.44 ms	+0.9%
✅	`add/large_f32_threads=1-internal/4194304`	37.64 ms	37.91 ms	+0.7%
✅	`add/small_f32_threads=1-internal/1024`	181.4 ns	182.1 ns	+0.4%
✅	`reduce_mean/medium_f32_threads=1-internal/65536`	228.15 µs	227.59 µs	-0.2%
✅	`add/small_f16_threads=1-internal/1024`	12.13 µs	12.08 µs	-0.4%
✅	`add/large_f16_threads=1-internal/4194304`	38.36 ms	38.18 ms	-0.5%
✅	`reduce_mean/small_f32_threads=1-internal/4096`	14.19 µs	14.11 µs	-0.6%
✅	`add/large_bf16_threads=1-internal/4194304`	39.59 ms	39.37 ms	-0.6%
✅	`add/small_bf16_threads=1-internal/1024`	11.99 µs	11.79 µs	-1.7%
✅	`gather/small_bf16_threads=1-internal/4096`	458.2 ns	440.2 ns	-3.9%
✅	`gather/small_f32_threads=1-internal/4096`	654.4 ns	626.9 ns	-4.2%
✅	`block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1`	80.27 µs	75.19 µs	-6.3%
✅	`matmul/medium_generic_bf16_threads=8/32x512x512`	403.18 µs	371.76 µs	-7.8%
✅	`tokenization/encode_tokens_per_second`	410.47 µs	376.13 µs	-8.4%
✅	`block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024`	43.87 µs	39.74 µs	-9.4%
✅	`gather/medium_bf16_threads=1-internal/32768`	2.53 µs	2.27 µs	-10.0%
✅	`matmul/large_generic_bf16_threads=1/32x1024x1024`	2.18 ms	1.89 ms	-13.3%
✅	`matmul/small_generic_f32_threads=1/1x256x256`	40.76 µs	34.93 µs	-14.3%
🟢	`matmul/medium_generic_f32_threads=1/32x512x512`	2.61 ms	2.17 ms	-16.9%
🟢	`gather/medium_f16_threads=1-internal/32768`	2.70 µs	2.22 µs	-17.7%
🟢	`matmul/large_generic_f32_threads=1/32x1024x1024`	10.90 ms	8.92 ms	-18.1%
🟢	`matmul/medium_generic_bf16_threads=1/32x512x512`	625.24 µs	509.13 µs	-18.6%
🟢	`matmul/small_generic_bf16_threads=8/1x256x256`	41.75 µs	33.97 µs	-18.6%
🟢	`block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1`	420.89 µs	339.88 µs	-19.2%
🟢	`matmul/large_generic_bf16_threads=8/32x1024x1024`	1.68 ms	1.35 ms	-19.4%
🟢	`matmul/small_generic_bf16_threads=1/1x256x256`	40.69 µs	32.19 µs	-20.9%
🟢	`gather/large_f32_threads=1-internal/131072`	31.78 µs	24.91 µs	-21.6%
🟢	`matmul/small_generic_f32_threads=8/1x256x256`	43.34 µs	33.64 µs	-22.4%
🟢	`block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024`	670.32 µs	513.62 µs	-23.4%
🟢	`matmul/small_generic_f16_threads=1/1x256x256`	37.93 µs	28.47 µs	-24.9%
🟢	`matmul/small_generic_f16_threads=8/1x256x256`	42.67 µs	31.78 µs	-25.5%
🟢	`matmul/large_generic_f16_threads=8/32x1024x1024`	105.65 µs	77.57 µs	-26.6%
🟢	`gather/medium_f32_threads=1-internal/32768`	5.00 µs	3.53 µs	-29.4%
🟢	`matmul/medium_generic_f16_threads=1/32x512x512`	41.00 µs	27.41 µs	-33.2%
🟢	`matmul/large_generic_f16_threads=1/32x1024x1024`	113.22 µs	72.92 µs	-35.6%
🟢	`matmul/medium_generic_f16_threads=8/32x512x512`	44.83 µs	28.32 µs	-36.8%
🟢	`block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024`	69.21 µs	42.76 µs	-38.2%
🟢	`matmul/large_generic_f32_threads=8/32x1024x1024`	6.00 ms	3.70 ms	-38.4%
🟢	`gather/large_bf16_threads=1-internal/131072`	17.40 µs	10.29 µs	-40.8%
🟢	`gather/large_f16_threads=1-internal/131072`	18.70 µs	10.20 µs	-45.5%
🟢	`matmul/medium_generic_f32_threads=8/32x512x512`	1.78 ms	900.37 µs	-49.4%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 3.13 3.16 3.83 }

What this cannot catch

Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
Sub-threshold regressions that compound over multiple PRs
Performance changes that only manifest under GPU execution
Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)

## Summary - stacked on #999 - attributes the fixed GQA decode aux bucket: standalone RoPE/append/transposes are already fused away; split-K merge is the dominant fixed aux cost - tiles fp16/bf16 split-merge across 4 head-dim blocks per query row - precomputes split weights once per merge CTA in shared memory ## Measurements (glm-4-9b-int4, GPU2, graph ON, ONNX_GENAI_LMHEAD_CUBLASLT=1) Before = #999 stacked baseline in this worktree. | KV | core before | merge before | prep before | core after | merge after | prep after | e2e after | |---:|---:|---:|---:|---:|---:|---:|---:| | 512 | 447.6 us | 280.8 us | 104.0 us | 448.0 us | 223.4 us | 101.8 us | 5.088 ms/tok, 196.56 tok/s | | 2048 | 1335.0 us | 280.9 us | 109.1 us | 1335.1 us | 223.2 us | 107.6 us | 5.940 ms/tok, 168.35 tok/s | Aux reduction: ~57 us/step, context-independent. Short base decode: 207.57 tok/s, captures=3 replays=474 fallbacks=0. ## Correctness / validation - glm and qwen KV512/KV2048 generated streams matched #999 exactly in measured before/after runs - capture remained intact (fallbacks=0) - cargo test -p onnx-runtime-ep-cuda --features cuda,gpu-tests --test gqa_fp16_gpu --quiet - cargo test -p onnx-runtime-ep-cuda --features cuda,gpu-tests --test group_query_attention_gpu gqa_gpu_fp16_decode_split_k_long_context_is_deterministic_and_matches_baseline --quiet - cargo fmt --all -- --check - cargo clippy -p onnx-runtime-ep-cuda --features cuda --lib -- -D warnings Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Promote the #991 cuBLASLt lm_head decode MatMul to default-on and harden its CUDA-graph capture guard from a single cached shape to a shape-keyed plan cache, so the out-of-box glm base decode banks the cuBLASLt logits projection without an env flag. Capture guard (Gaff's #991 blocker): `launch_dense_capturable` now keeps a shape-keyed `Vec<DenseGemmPlan>` (LRU, MRU-promote on hit, bounded by DENSE_PLAN_CACHE_CAP=8) instead of a single `Option<DenseGemmPlan>`. Every distinct (dtype, M, K, N) the node runs keeps its own preselected cuBLASLt algorithm + persistent workspace, so alternating shapes (prefill M>1, decode M==1, a speculative M=K verify width) all replay with no per-call heuristic query inside the captured region. The reject-cold-miss-during-capture contract is unchanged (now enforced per shape); eviction only ever runs off the non-capturing creation path, and the hot decode plan is kept at the MRU front, so a plan baked into a live captured graph is never freed. Strictly safer than the single-slot path (evicts less; same recapture-on-change discipline). Default-on: `lmhead_cublaslt_enabled()` inverts to default-TRUE-unless-disabled; set ONNX_GENAI_LMHEAD_CUBLASLT=0/false/off for the hand-GEMV escape hatch. #999 GQA flash-decoding is already default-on (its env vars are rollback-only knobs), so no flip is needed there. GEMV fp16 (#996) stays untouched (not in this base) and is a one-liner default flip once it merges. Measured (glm-4-9b-int4, H200, graph-on, fallbacks=0): - default-on short decode 207.72 tok/s (vs escape-hatch-off hand-GEMV 202.24) - default-on KV2048 171.69 tok/s, fallbacks=0 across 8 bucket-growth recaptures - lm_head cuBLASLt(default) == hand-GEMV(off) top-40 logprobs BYTE-IDENTICAL - qwen no regression (152.16 vs 153.60, within noise), fallbacks=0 - f64 oracle 7/7, GQA fp16 2/2, fmt + clippy clean Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…1004) ## Default-on: lm_head cuBLASLt + multi-shape capture guard Promotes #991's cuBLASLt `lm_head` decode MatMul to **default-on** (out-of-box glm base), after clearing Gaff's #991 blocker: the single-shape cuBLASLt plan cache is hardened into a **shape-keyed multi-shape plan cache** so plan creation never happens inside a captured CUDA-graph region across prefill/decode/verify shape switches. Rebased onto main `ad141370` (post-#996); disjoint from #996's opt-in fp16 (which stays OFF by default). ### Capture guard — path (a): shape-keyed multi-shape plan cache #991 already had a warm-once / reject-cold-miss-during-capture contract, but for a SINGLE cached shape (a shape switch evicted + re-ran the heuristic search). This PR generalizes it to a `Vec<DenseGemmPlan>` keyed by `(dtype,M,K,N)`, MRU-ordered, bounded by `DENSE_PLAN_CACHE_CAP = 8`: - **hit** → MRU-promote, launch preselected algo (no heuristic query); - **cold miss while capturing** → `Err` (reject-before-warm, now per shape); - **cold miss not capturing** → create + insert-front + LRU-truncate. - Eviction runs **only** off the non-capturing path; the hot decode plan (M==1) stays MRU-front and is never the LRU victim. Strictly safer than the single slot (evicts far less; same recapture-on-change discipline). **Shape-transition proof** (temporary `DENSE_PLAN_SEARCH` counter, glm KV2048 run = prompt → decode → **8 KV-bucket-growth recaptures**, 8623 replays): - **exactly 2 plan searches for the whole run** — `M=3` prefill (once) + `M=1` decode (once), each on the non-capturing path; - **0 searches** during any of the 8623 replays or 8 recaptures; - **fallbacks=0** (a cold-miss-inside-capture would surface as a fallback; none did). This is the "no per-transition plan search, not just steady-state" guarantee. Counter reverted (not committed). ### Default-on flip (byte-identical wins only) - `lmhead_cublaslt_enabled()` → default-TRUE-unless-disabled; escape hatch `ONNX_GENAI_LMHEAD_CUBLASLT=0|false|off` → hand fp16 GEMV. - **#999 GQA flash-decoding already default-on** (audited — its env vars are rollback-only knobs); no flip needed. - **GEMV fp16 (#996) stays PERMANENTLY opt-in — NOT flipped.** Chew's #996 verdict is land-as-opt-in: the fp16 K-accumulate is not byte-identity-safe vs fp32 (glm knife-edge argmax flips), so it is not a default-on candidate. Out-of- box GEMV stays fp32 (byte-identical); fp16 via `ONNX_GENAI_GEMV_FP16=1`. ### Measured (glm-4-9b-int4, H200, graph-on, median, no env flags) | config | short tok/s | KV2048 tok/s | fallbacks | |---|---|---|---| | **default-on (no flags)** | **~206–208** | **~166–172** | **0** | | escape-hatch `=0` (hand GEMV) | 202.24 | — | 0 | (short 207.72 GPU3 / 206.44 GPU4; KV2048 171.69 / 168.06 / 165.59 across GPU3/6/4 — cross-GPU + shared-host variance; all fallbacks=0.) - Byte-identity: cuBLASLt lm_head (default) vs hand-GEMV (off) top-40 token-0 logprobs **byte-identical** (both select token 315). - qwen no regression (152.16 vs 153.60, noise), fallbacks=0. - f64 oracle 7/7; GQA fp16 2/2; fmt + clippy clean. **Out-of-box number is ~206–208 short / ~166–172 KV2048 and that is final for this PR.** The ~232 combined figure required fp16 GEMV opt-in ON; since fp16 ships permanently opt-in, 232 is only reachable via `ONNX_GENAI_GEMV_FP16=1`. ⚠️ Please gate with fresh Chew (byte-identity / f64 oracle) + Gaff (capture-safe, fallbacks=0, win real). --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

justinchuby mentioned this pull request Aug 15, 2026

perf(cuda): tile GQA decode split-merge epilogue #1003

Merged

justinchuby marked this pull request as ready for review August 15, 2026 08:30

justinchuby merged commit eeeb18e into main Aug 15, 2026
7 of 15 checks passed

justinchuby deleted the squad/attn-decode-flashdecoding branch August 15, 2026 08:30

justinchuby mentioned this pull request Aug 15, 2026

perf(cuda): default-on lm_head cuBLASLt + multi-shape capture guard #1004

Merged

Navigation Menu

Repository navigation

perf(cuda): retune GQA decode flash-decoding - #999

perf(cuda): retune GQA decode flash-decoding#999
justinchuby merged 1 commit into
mainfrom
squad/attn-decode-flashdecoding

justinchuby commented Aug 15, 2026

Uh oh!

codecov Bot commented Aug 15, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 15, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

justinchuby commented Aug 15, 2026

Summary

Measurements (glm-4-9b-int4, GPU1, graph ON, ONNX_GENAI_LMHEAD_CUBLASLT=1)

Validation

Accuracy notes

Uh oh!

codecov Bot commented Aug 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions Bot commented Aug 15, 2026

🔴 Benchmark Regression Detected

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

codecov Bot commented Aug 15, 2026 •

edited

Loading