[justinchuby]

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.
• perf(cuda): retune GQA decode flash-decoding #999 GQA flash-decoding already default-on (audited — its env vars are
  rollback-only knobs); no flip needed.
• GEMV fp16 (perf(cuda): fp16 int4 wide GEMV matching ORT's M=1 arithmetic — glm 199→213 tok/s, accuracy-gated (opt-in) #996) stays PERMANENTLY opt-in — NOT flipped. Chew's perf(cuda): fp16 int4 wide GEMV matching ORT's M=1 arithmetic — glm 199→213 tok/s, accuracy-gated (opt-in) #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).

[codecov]

Codecov Report

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

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1004   +/-   ##
=======================================
  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.11% <ø> (+0.09%)	⬆️

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]

🔴 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
🔴	matmul/medium_generic_f32_threads=8/32x512x512	935.44 µs	1.77 ms	+89.2%
🔴	matmul/small_generic_bf16_threads=8/1x256x256	31.63 µs	43.05 µs	+36.1%
🔴	matmul/medium_generic_f32_threads=1/32x512x512	2.30 ms	3.03 ms	+32.2%
🔴	add/medium_f16_threads=1-internal/262144	2.56 ms	3.38 ms	+32.0%
🔴	grammar_masking/llguidance_compute_mask/32	76.66 µs	100.40 µs	+31.0%
⚠️	matmul/medium_generic_f16_threads=1/32x512x512	30.66 µs	39.45 µs	+28.7%
⚠️	add/large_f32_threads=1-internal/4194304	40.36 ms	51.72 ms	+28.2%
⚠️	add/medium_bf16_threads=1-internal/262144	2.77 ms	3.42 ms	+23.2%
⚠️	matmul/medium_generic_f16_threads=8/32x512x512	30.10 µs	37.07 µs	+23.1%
⚠️	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	1.01 ms	1.24 ms	+22.9%
⚠️	sampling_latency/greedy_per_token	3.51 µs	4.20 µs	+19.5%
⚠️	matmul/large_generic_f16_threads=1/32x1024x1024	91.47 µs	108.21 µs	+18.3%
⚠️	gather/small_f16_threads=1-internal/4096	489.7 ns	576.8 ns	+17.8%
⚠️	add/large_f16_threads=1-internal/4194304	40.57 ms	46.96 ms	+15.8%
⚠️	reduce_mean/small_f32_threads=1-internal/4096	15.14 µs	17.44 µs	+15.1%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.89 ms	6.76 ms	+14.8%
✅	add/small_f32_threads=1-internal/1024	209.4 ns	239.6 ns	+14.5%
✅	gather/small_bf16_threads=1-internal/4096	481.7 ns	548.2 ns	+13.8%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	539.36 µs	612.93 µs	+13.6%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	98.08 µs	110.91 µs	+13.1%
✅	gather/medium_f32_threads=1-internal/32768	3.85 µs	4.29 µs	+11.6%
✅	reduce_mean/medium_f32_threads=1-internal/65536	244.61 µs	271.20 µs	+10.9%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	562.58 µs	622.37 µs	+10.6%
✅	add/large_bf16_threads=1-internal/4194304	42.79 ms	47.11 ms	+10.1%
✅	gather/large_f32_threads=1-internal/131072	29.58 µs	32.56 µs	+10.1%
✅	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	529.36 µs	582.41 µs	+10.0%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.25 ms	2.45 ms	+9.0%
✅	kv_cache/alloc_dealloc_pages	42.93 µs	46.80 µs	+9.0%
✅	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	135.29 µs	145.91 µs	+7.9%
✅	gather/medium_bf16_threads=1-internal/32768	2.42 µs	2.61 µs	+7.8%
✅	tokenization/encode_tokens_per_second	448.81 µs	481.12 µs	+7.2%
✅	sampling_latency/min_p_per_token	236.89 µs	253.90 µs	+7.2%
✅	logit_processing/seven_processor_chain_per_step	362.18 µs	386.94 µs	+6.8%
✅	sampling_latency/top_p_per_token	441.33 µs	471.13 µs	+6.8%
✅	gather/large_bf16_threads=1-internal/131072	11.71 µs	12.41 µs	+6.0%
✅	sampling_latency/top_k_per_token	56.86 µs	59.96 µs	+5.4%
✅	tokenization/decode_tokens_per_second	7.15 ms	7.52 ms	+5.2%
✅	qwen3_sampling_processors/top_k_top_p_fast	737.42 µs	775.17 µs	+5.1%
✅	gather/medium_f16_threads=1-internal/32768	2.43 µs	2.54 µs	+4.6%
✅	reduce_mean/large_f32_threads=1-internal/262144	983.98 µs	1.03 ms	+4.4%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	83.82 µs	87.40 µs	+4.3%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	4.02 ms	4.17 ms	+3.9%
✅	add/small_f16_threads=1-internal/1024	13.95 µs	14.47 µs	+3.7%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.39 ms	2.46 ms	+3.3%
✅	gather/small_f32_threads=1-internal/4096	667.8 ns	680.2 ns	+1.9%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	6.71 ms	6.74 ms	+0.3%
✅	matmul/small_generic_f16_threads=1/1x256x256	31.17 µs	31.05 µs	-0.4%
✅	gather/large_f16_threads=1-internal/131072	12.50 µs	12.38 µs	-0.9%
✅	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	112.53 µs	110.89 µs	-1.5%
✅	matmul/small_generic_f32_threads=1/1x256x256	36.91 µs	36.33 µs	-1.6%
✅	qwen3_sampling_processors/top_k_partial_selection	160.06 µs	157.02 µs	-1.9%
✅	matmul/small_generic_bf16_threads=1/1x256x256	34.54 µs	33.53 µs	-2.9%
✅	add/medium_f32_threads=1-internal/262144	3.03 ms	2.94 ms	-3.0%
✅	add/small_bf16_threads=1-internal/1024	13.73 µs	13.28 µs	-3.3%
✅	matmul/small_generic_f32_threads=8/1x256x256	36.40 µs	33.58 µs	-7.8%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	547.50 µs	498.19 µs	-9.0%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	11.91 ms	10.74 ms	-9.8%
✅	matmul/small_generic_f16_threads=8/1x256x256	35.82 µs	30.63 µs	-14.5%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	2.15 ms	1.71 ms	-20.6%

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: { 5.06 3.53 4.13 }

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)