perf(native-decode): default-on CUDA-graph capture for multi-component step-inputs

[justinchuby]

What & why

Flips the native decoder multi-component / routed step-inputs CUDA-graph capture from opt-in to default-on, the high-performance win authorized by Justin ("确保高性能").

The captured per-step-input decode path (Inc3c, #384) writes the one-token inputs_embeds/routed tensors into persistent device bindings and reuses the run_one_token captured graph — instead of the eager owned-input uploads that forfeited CUDA-graph capture on every routed decode step. That eager fallback was the root cause of the ~6.5× gemma4-e2b native-vs-ORT decode gap. It was already built and parity-tested but gated OFF; this PR simply makes it the default.

ONNX_GENAI_NATIVE_DECODER_CAPTURE_STEP_INPUTS is inverted opt-in → opt-out:

• unset / truthy / unrecognized → capture-on (default);
• 0 / false / no / off → force the eager owned path (escape hatch).

Scope: this only affects the step-inputs (multi-component / routed) decode path. Monolithic single-component decoders (qwen2.5/qwen3) already captured by default via the token-id run_one_token path — unchanged. No capture-core / plan_capture_region / standard_attention / GAP-3 KV edits.

Regression safety (the point)

Structural eligibility gates are unchanged (graph_enabled, non-empty captured_step_inputs), so ineligible decoders — growing-logical bindings, GLM-style mask-exposed indexers, recurrent Scan/LinearAttention hybrids — auto-decline to eager. No silent-wrong. Default-on output is byte-identical to the eager path.

Tests (GPU, CUDA_VISIBLE_DEVICES=2)

• native_cuda_captured_step_inputs_parity (--test-threads=1): tokens=[0,5,6,7] default_captured_decodes=3 opt_in_captured_decodes=3 opt_out_captured_decodes=0 — the default (no env) engages capture and is byte-identical to the =0 opt-out eager baseline; the counter makes it non-vacuous (a silent no-capture fallback fails).
• New capture_step_inputs_gate_tests unit module: default-on, falsy set opts out, truthy/unknown stay on.
• qwen3_0_6b_capture_step_inputs_decline (--ignored): single-component decoder declines under the default (no env) — default-on never mis-engages an ineligible decoder.
• gemma3n_native_cuda_capture_realmodel (--ignored): default engages + token parity + =0 opt-out declines; graceful-skips on the orthogonal vision-export blocker.
• cargo fmt --all --check clean.

Caveats

• Recurrent Scan / LinearAttention hybrids (27B / 35B-A3B recurrent path) still decline through the structural gate until their separate capture lane lands.
• A real gemma4-e2b end-to-end benchmark remains blocked by the orthogonal stale vision/audio export (vision_encoder OneHot Depth is negative even text-only; the pipeline forces vision load). The real-model harness is ready the moment a clean export / text-only-skip exists.
• gemma4-e2b's absolute headline number is additionally embedding-bound (~5.5 GB every-step embedding upload) — a separate follow-up. This change brings capture-eligible native decode to ORT-parity-or-better broadly; gemma4-e2b's headline is not solely gated by this.

Decision note: .squad/decisions/inbox/cohaagen-capture-default-on.md. Scope note: .squad/decisions/inbox/cohaagen-perfcap-scope.md.

Refs #384 (Inc3c lineage).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.60%. Comparing base (f3c6e79) to head (e1422a2).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #571      +/-   ##
==========================================
- Coverage   81.31%   80.60%   -0.72%     
==========================================
  Files         315      315              
  Lines      123574   123574              
  Branches   123574   123574              
==========================================
- Hits       100489    99606     -883     
- Misses      19026    19916     +890     
+ Partials     4059     4052       -7     

Flag	Coverage Δ
cli-ort-linux	83.27% <ø> (ø)
cli-ort-windows	82.67% <ø> (-0.11%)	⬇️
mlas	77.91% <ø> (ø)
offline	80.52% <ø> (-0.76%)	⬇️

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

🚀 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
🔴	gather/large_f16_threads=1-internal/131072	9.50 µs	14.98 µs	+57.7%
⚠️	gather/large_bf16_threads=1-internal/131072	9.77 µs	11.87 µs	+21.5%
✅	add/small_bf16_threads=1-internal/1024	12.04 µs	13.36 µs	+11.0%
✅	add/medium_f16_threads=1-internal/262144	2.37 ms	2.62 ms	+10.6%
✅	gather/large_f32_threads=1-internal/131072	28.71 µs	30.99 µs	+7.9%
✅	add/small_f16_threads=1-internal/1024	12.06 µs	12.96 µs	+7.5%
✅	gather/small_f32_threads=1-internal/4096	621.4 ns	647.8 ns	+4.2%
✅	reduce_mean/large_f32_threads=1-internal/262144	916.43 µs	936.28 µs	+2.2%
✅	matmul/small_generic_f32_threads=1/1x256x256	33.82 µs	34.48 µs	+2.0%
✅	add/medium_f32_threads=1-internal/262144	2.33 ms	2.35 ms	+0.8%
✅	add/large_f16_threads=1-internal/4194304	38.54 ms	38.46 ms	-0.2%
✅	gather/medium_f32_threads=1-internal/32768	3.79 µs	3.77 µs	-0.5%
✅	gather/medium_bf16_threads=1-internal/32768	2.26 µs	2.25 µs	-0.5%
✅	reduce_mean/medium_f32_threads=1-internal/65536	230.13 µs	227.76 µs	-1.0%
✅	gather/medium_f16_threads=1-internal/32768	2.29 µs	2.27 µs	-1.0%
✅	add/large_bf16_threads=1-internal/4194304	39.62 ms	39.19 ms	-1.1%
✅	reduce_mean/small_f32_threads=1-internal/4096	14.31 µs	14.12 µs	-1.3%
✅	gather/small_bf16_threads=1-internal/4096	443.5 ns	436.4 ns	-1.6%
✅	matmul/small_generic_bf16_threads=1/1x256x256	29.71 µs	29.23 µs	-1.6%
✅	gather/small_f16_threads=1-internal/4096	446.0 ns	434.7 ns	-2.5%
✅	add/medium_bf16_threads=1-internal/262144	2.52 ms	2.43 ms	-3.4%
✅	kv_cache/alloc_dealloc_pages	37.28 µs	35.74 µs	-4.1%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	1.94 ms	1.86 ms	-4.5%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	9.36 ms	8.76 ms	-6.4%
✅	add/large_f32_threads=1-internal/4194304	41.06 ms	37.82 ms	-7.9%
✅	sampling_latency/top_p_per_token	1.04 ms	914.55 µs	-12.2%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	85.05 µs	73.90 µs	-13.1%
✅	sampling_latency/min_p_per_token	360.96 µs	313.07 µs	-13.3%
✅	add/small_f32_threads=1-internal/1024	235.0 ns	202.9 ns	-13.7%
✅	logit_processing/seven_processor_chain_per_step	1.25 ms	1.08 ms	-13.8%
✅	grammar_masking/llguidance_compute_mask/32	81.03 µs	69.62 µs	-14.1%
🟢	matmul/medium_generic_bf16_threads=1/32x512x512	578.27 µs	490.27 µs	-15.2%
🟢	matmul/small_generic_f16_threads=1/1x256x256	32.56 µs	27.53 µs	-15.4%
🟢	sampling_latency/greedy_per_token	3.56 µs	2.96 µs	-16.9%
🟢	matmul/small_generic_f16_threads=8/1x256x256	34.74 µs	28.78 µs	-17.2%
🟢	matmul/large_generic_f16_threads=8/32x1024x1024	96.65 µs	78.32 µs	-19.0%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	476.69 µs	381.91 µs	-19.9%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	1.56 ms	1.24 ms	-20.4%
🟢	matmul/medium_generic_f16_threads=1/32x512x512	37.69 µs	29.66 µs	-21.3%
🟢	sampling_latency/top_k_per_token	566.57 µs	441.33 µs	-22.1%
🟢	matmul/medium_generic_f32_threads=1/32x512x512	2.85 ms	2.15 ms	-24.6%
🟢	matmul/small_generic_bf16_threads=8/1x256x256	40.66 µs	30.30 µs	-25.5%
🟢	matmul/medium_generic_f16_threads=8/32x512x512	37.75 µs	28.00 µs	-25.8%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	5.16 ms	3.56 ms	-31.0%
🟢	tokenization/decode_tokens_per_second	8.31 ms	5.65 ms	-32.0%
🟢	matmul/small_generic_f32_threads=8/1x256x256	48.67 µs	32.46 µs	-33.3%
🟢	tokenization/encode_tokens_per_second	547.51 µs	345.78 µs	-36.8%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	2.02 ms	903.49 µs	-55.3%

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: { 4.24 3.68 6.08 }

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)

[justinchuby]

VERDICT: APPROVE

Independent review by Harry (author Cohaagen locked out). Re-ran everything from PR #571 head e1422a2 on GPU 6; did not trust the author report. This is a regression-sensitive perf change, so the bar was: prove it cannot change decode output and cannot silently mis-capture.

1. Diff scope — CLEAN

git diff origin/main...HEAD touches only: cuda.rs (flag parse inversion + comments + a new pure unit-test module), the 3 named tests, docs/CUDA_GRAPH_CAPTURE.md, and 2 .squad/decisions notes. Grep for plan_capture_region / standard_attention / run_one_token / GAP-3 / KV-mirror shows every hit is in comments or docs — NO capture-core, plan_capture_region, standard_attention, KV path, or GAP-3 code edits. Confirmed the structural eligibility assignment (graph_enabled = ..., mask_exposes_logical, has_dynamic_logical_input_shape, !captured_step_inputs.is_empty()) is byte-for-byte untouched by the diff — only the final AND-term capture_step_inputs_enabled() default flipped.

2. Eligibility-gate trace — nothing can silently mis-capture

capture_step_inputs = graph_enabled && !captured_step_inputs.is_empty() && capture_step_inputs_enabled().
The PR changes ONLY the third term's default (unset: false -> true). The first two terms are unchanged. Therefore the set of decoders that become capture-eligible under default-on is exactly the set already deemed capture-safe by the unchanged structural gates:

• GLM-style mask-exposed indexer -> mask_exposes_logical clears graph_enabled -> term1 false.
• growing-logical / Scan-LinearAttention binding -> dynamic_logical non-empty clears graph_enabled -> term1 false.
• single-component (input_ids, no per-step embeds/routed ports) -> captured_step_inputs empty -> term2 false.
  The default flip mathematically cannot make an ineligible decoder eligible. Conclusion: no silent-wrong path is introduced.

3. Tests re-run (results)

• native_cuda_captured_step_inputs_parity (--test-threads=1): PASS. tokens=[0,5,6,7]; default (no env) captured_decodes=3; explicit =1 captured_decodes=3; opt-out =0 captured_decodes=0. Full generated token STREAM (Vec, whole autoregressive rollout) is asserted equal between default-on and eager, plus asserts default captured_decodes>=1. Runtime also guards each captured step with a device capture-error poll and a non-finite-logits check.
• qwen3_0_6b_capture_step_inputs_decline (--ignored, real int4 model): PASS under DEFAULT-on. Single-component decoder yields captured_step_input_decodes=0 and correct tokens=[279,3146,374,279,1172,1992,1380,279]. Real-model decline-to-eager proven under the new default.
• cargo fmt --all --check: clean.
• gemma3n_native_cuda_capture_realmodel (--ignored, real package): its engagement/parity assertions did NOT execute in my run — the harness gracefully skips because the vision_encoder aborts (OneHot "Depth is negative") before decode. This is a pre-existing fixture/env limitation unrelated to this PR (no decode assertion was reached, nothing failed). Real-model default-on engagement therefore rests on the qwen3 decline run + synthetic parity + the gate trace above.

4. Mutation testing — BOTH mutations FAIL the suite (teeth confirmed)

• Mutation A (force default back to OFF: None => false): parity FAILS — "expected the DEFAULT (no env) to engage the captured per-step-input path at least once, saw 0". So a silent revert-to-eager / no-capture regression is caught.
• Mutation B (make captured != eager): a position offset was a no-op on this fixture (fixture is position-insensitive), so I escalated to corrupting the captured logits (row.reverse()). Parity then FAILS — captured tokens [0,2,0,2] vs eager [0,5,6,7]. So a captured-path output divergence is caught.
  Both mutations reverted; working tree restored to pristine (git diff empty). No commits made.

5. Oracle-strength caveat (non-blocking)

The parity oracle is full-token-stream argmax equality over the whole rollout (errors compound across steps), NOT raw per-logit byte equality — the position-offset no-op shows a sub-argmax-threshold perturbation would not be caught by tokens alone. For a CUDA-graph capture regression the realistic failure modes (stale/wrong persistent bindings, no-capture, bad replay) produce materially wrong logits that DO flip argmax (demonstrated by Mutation B) and are additionally caught by the device capture-error poll + non-finite guard. I judge the oracle non-vacuous and adequate. Recommendation for a follow-up (NOT a merge blocker): add a logit L-inf/allclose assertion between default and eager to harden against sub-threshold drift, and re-enable a real-model multi-component engagement test that doesn't depend on the gemma vision path.

Summary

Scope is exactly the default-flag inversion + eligibility-preserving plumbing + tests + docs. Structural eligibility gates are untouched, so default-on captures only already-capture-safe decoders and ineligible decoders still decline to the byte-identical eager fallback (proven on the real qwen3 model). Parity holds and both negative-control mutations fail the suite. Approving.