perf(cuda): split-KV FlashDecoding for attention_row decode (+10% V2-Lite wide-ctx)

[justinchuby]

Summary

attention_row launches one block per (batch, q_head, query) row, so at decode it uses only ~16 of 132 SMs and is memory-latency bound. Profiling V2-Lite eager decode at wide context (GPU3, ~390-word prompt, 256 tok) showed attention_row is the #1 decode kernel at 36% of GPU decode — ncu: grid=16, waves/SM=0.02 (~116 idle SMs), DRAM 0.79%, achieved occupancy 9.8% → a definitive under-utilization signal.

This PR adds a two-kernel FlashDecoding split-KV path that spreads each row's key reduction across many blocks to fill the idle SMs.

Design

Behind ONNX_GENAI_ATTN_SPLITKV (default ON; =0 reproduces the monolithic baseline for A/B):

• attention_split — grid (total_rows, num_splits). Each block reduces one KV slice with a chunk-local online softmax (running max + rescale), writing an unnormalized partial P·V plus per-split max/sum meta. A single-split fast path (total_seq <= chunk) reproduces attention_row exactly and writes the final normalized output with sentinel meta (max=0, sum=1).
• attention_combine — grid (total_rows). Uniform log-sum-exp merge of the per-split partials. The single-split sentinel makes it a bit-exact pass-through, so any context <= chunk (default 256) stays byte-identical to the monolithic kernel.

Capture-safety: num_splits and the engage decision derive only from the fixed KV capacity (cap), never live seqlen, so eager and capture make identical launch decisions → the eager==capture long-context lock holds. Split scratch is drawn from the internal capture-warmed ws pool (WS_SPLIT slot), leaving the governed workspace layout (and its exact-byte unit tests) untouched.

Multi-split reorders fp32 partials, so wide context is validated to the #1150 f64 tolerance (2e-4.max(·*2e-5)), not byte-identity. Greedy token output is unchanged in practice.

Numerics gates (H200, GPU3, pinned, single-threaded)

Gate	Result
V2-Lite golden 24-tok lock (matches_golden)	✅ byte-identical (single-split path)
V2-Lite 340-tok long-context lock (eager==capture + golden prefix)	✅ pass (multi-split engaged, capture-safe)
standard_attention_capture_gpu / _fp16_gpu / _bf16_gpu	✅ pass
standard_attention_gpu	✅ 23/24 (the 1 failure ..._requires_homogeneous_floating_input_dtypes is pre-existing on clean origin/main, unrelated)
layout unit tests (--lib standard_attention)	✅ 14/14
Dense qwen2.5-0.5b token md5	✅ byte-identical (no regression)

Perf A/B (H200 GPU3, pinned, medians-of-5)

Wide-context kernel (nsys, same binary, split ON vs =0):

	decode attention kernel (median)	DRAM throughput
baseline (attention_row)	147.6 µs	0.79%
split (attention_split+combine)	~62.5 + 2.5 µs (~2.3× faster)	9.26%

E2E throughput (tok/s):

context	split OFF (=0)	split ON (default)	Δ
wide (~390-word prompt, 128 tok)	51.64	56.90	+10.2%
short "Hello" eager (128 tok)	62.25	63.04	+1.3% (neutral)
short "Hello" capture (128 tok)	128.48	128.17	−0.2% (neutral)
dense qwen2.5-0.5b capture	610.36	610.25	neutral

num_splits = cap.div_ceil(chunk).clamp(1, 64) with chunk default 256 (tunable via ONNX_GENAI_ATTN_SPLIT_CHUNK). At the true long-context regime where attention_row reaches ~40% of decode (cap≈2048), this yields ~8 splits → grid ~128, filling the machine; smaller chunks extract a bit more at moderate context at the cost of engaging split earlier.

Verdict: GO, default-ON

Short context stays byte-identical (comfortable margin below chunk); wide context gets +10.2% E2E / ~2.3× kernel; capture-safe; no dense regression. All hard gates green.

Do not self-merge — reporting to coordinator for admin-merge after review.

Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com

[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
🔴	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	43.18 µs	81.64 µs	+89.1%
🔴	gather/large_f32_threads=1-internal/131072	24.81 µs	43.58 µs	+75.7%
🔴	gather/large_f16_threads=1-internal/131072	12.10 µs	20.04 µs	+65.6%
🔴	matmul/medium_generic_f16_threads=8/32x512x512	29.78 µs	48.83 µs	+64.0%
🔴	matmul/large_generic_bf16_threads=8/32x1024x1024	1.29 ms	2.00 ms	+55.3%
🔴	matmul/medium_generic_bf16_threads=1/32x512x512	502.91 µs	763.77 µs	+51.9%
🔴	gather/large_bf16_threads=1-internal/131072	10.55 µs	16.02 µs	+51.8%
🔴	matmul/medium_generic_f16_threads=1/32x512x512	27.45 µs	40.99 µs	+49.4%
🔴	matmul/small_generic_bf16_threads=8/1x256x256	29.27 µs	42.19 µs	+44.2%
🔴	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	359.10 µs	516.23 µs	+43.8%
🔴	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	539.46 µs	757.26 µs	+40.4%
🔴	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	43.11 µs	58.82 µs	+36.4%
🔴	matmul/small_generic_f16_threads=1/1x256x256	30.99 µs	42.17 µs	+36.1%
🔴	gather/medium_f32_threads=1-internal/32768	3.54 µs	4.78 µs	+35.1%
🔴	matmul/small_generic_bf16_threads=1/1x256x256	30.93 µs	41.02 µs	+32.6%
⚠️	gather/small_bf16_threads=1-internal/4096	457.3 ns	585.7 ns	+28.1%
⚠️	matmul/medium_generic_bf16_threads=8/32x512x512	355.71 µs	448.84 µs	+26.2%
⚠️	add/large_bf16_threads=1-internal/4194304	1.98 ms	2.45 ms	+24.0%
⚠️	gather/medium_f16_threads=1-internal/32768	2.34 µs	2.89 µs	+23.6%
⚠️	matmul/large_generic_f16_threads=8/32x1024x1024	78.95 µs	96.75 µs	+22.5%
⚠️	matmul/large_generic_f16_threads=1/32x1024x1024	73.40 µs	86.80 µs	+18.3%
⚠️	matmul/small_generic_f32_threads=1/1x256x256	34.35 µs	40.55 µs	+18.0%
⚠️	matmul/medium_generic_f32_threads=1/32x512x512	2.19 ms	2.58 ms	+17.7%
⚠️	reduce_mean/large_f32_threads=1-internal/262144	958.06 µs	1.12 ms	+16.8%
⚠️	matmul/medium_generic_f32_threads=8/32x512x512	896.67 µs	1.05 ms	+16.7%
✅	gather/small_f32_threads=1-internal/4096	697.5 ns	791.7 ns	+13.5%
✅	gather/small_f16_threads=1-internal/4096	565.0 ns	630.9 ns	+11.7%
✅	gather/medium_bf16_threads=1-internal/32768	2.34 µs	2.59 µs	+10.6%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	3.71 ms	4.06 ms	+9.5%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	9.10 ms	9.87 ms	+8.5%
✅	tokenization/encode_tokens_per_second	347.77 µs	372.20 µs	+7.0%
✅	matmul/small_generic_f32_threads=8/1x256x256	36.05 µs	37.53 µs	+4.1%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	482.94 µs	499.28 µs	+3.4%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	1.87 ms	1.92 ms	+2.4%
✅	sampling_latency/top_k_per_token	48.65 µs	49.75 µs	+2.3%
✅	reduce_mean/small_f32_threads=1-internal/4096	16.50 µs	16.69 µs	+1.1%
✅	sampling_latency/top_p_per_token	352.84 µs	356.35 µs	+1.0%
✅	sampling_latency/min_p_per_token	193.25 µs	194.84 µs	+0.8%
✅	sampling_latency/greedy_per_token	3.05 µs	3.07 µs	+0.8%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.29 ms	3.24 ms	-1.5%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	1.99 ms	1.96 ms	-1.6%
✅	add/medium_f16_threads=1-internal/262144	127.69 µs	125.49 µs	-1.7%
✅	add/medium_f32_threads=1-internal/262144	26.15 µs	25.70 µs	-1.7%
✅	qwen3_sampling_processors/top_k_partial_selection	133.16 µs	130.34 µs	-2.1%
✅	tokenization/decode_tokens_per_second	5.94 ms	5.70 ms	-4.1%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.81 ms	5.48 ms	-5.7%
✅	kv_cache/alloc_dealloc_pages	40.75 µs	38.32 µs	-6.0%
✅	reduce_mean/medium_f32_threads=1-internal/65536	281.44 µs	258.14 µs	-8.3%
✅	logit_processing/seven_processor_chain_per_step	333.32 µs	296.82 µs	-11.0%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	55.54 µs	49.44 µs	-11.0%
✅	add/small_bf16_threads=1-internal/1024	490.9 ns	420.6 ns	-14.3%
✅	matmul/small_generic_f16_threads=8/1x256x256	42.68 µs	36.55 µs	-14.4%
✅	add/large_f16_threads=1-internal/4194304	2.50 ms	2.14 ms	-14.4%
🟢	qwen3_sampling_processors/top_k_top_p_fast	727.25 µs	608.95 µs	-16.3%
🟢	grammar_masking/llguidance_compute_mask/32	82.67 µs	68.79 µs	-16.8%
🟢	add/large_f32_threads=1-internal/4194304	926.13 µs	738.11 µs	-20.3%
🟢	add/medium_bf16_threads=1-internal/262144	137.57 µs	104.68 µs	-23.9%
🟢	add/small_f32_threads=1-internal/1024	292.6 ns	196.8 ns	-32.7%
🟢	add/small_f16_threads=1-internal/1024	666.8 ns	437.1 ns	-34.5%

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.09 3.45 7.46 }

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)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.80%. Comparing base (06c62e0) to head (975123c).
⚠️ Report is 76 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1340      +/-   ##
==========================================
- Coverage   80.88%   80.80%   -0.09%     
==========================================
  Files         364      362       -2     
  Lines      160729   157677    -3052     
  Branches   160729   157677    -3052     
==========================================
- Hits       130005   127405    -2600     
+ Misses      26069    25632     -437     
+ Partials     4655     4640      -15     

Flag	Coverage Δ
mlas	?
offline	80.80% <ø> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 3 files with indirect coverage changes

🚀 New features to boost your workflow:

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