Enable CUDA-graph capture for causal fixed-capacity decode via on-device KV length

[justinchuby]

Summary

Whole-step CUDA-graph capture was declined for causal decoders such as granite-1b-a400m. The predicate persistent_inputs_have_fixed_logical_shapes declines when a persistent binding exposes a growing logical prefix, and granite's 48 past_key_values.N.{key,value} bindings do exactly that (physical [1,8,256,64] vs logical [1,8,0,64]). This is the root cause of granite's large deficit to ORT (capture is the dominant perf lever on this engine).The engine already has an on-device valid-length ABI (derive_len → dev_len) that lets the mask-driven decode bind KV at fixed physical capacity so capture stays shape-static. It was gated to is_causal == 0. For causal decode the three host-side KV consumers — the append slot (key_write_start), the causal frontier (offset), and the score-loop bound (total_seq) — read the logical length from host shape metadata, which is frozen to physical capacity at capture time. Replaying such a graph at a different true length is a length-baked read: silent numerical corruption, not a crash (measured in a prior falsification run as an all-8279 byte-divergent stream).## Length source (decision + rationale)

Two options were considered: (a) reuse derive_len (scan the additive mask frontier), or (b) maintain an engine-side device scalar KV length. Chosen: reuse derive_len, because it was measured to work for the causal case:- granite's mask input (input 3) is the standard additive causal-mask builder Where(And(Cast(attention_mask), GreaterOrEqual(query_pos, CumSum(attention_mask))), 0, -65504). It carries both causal and padding information.- Measured in the frozen configuration: the mask is materialized at physical width 256, host key_past_seq reports the frozen 256 (wrong), but derive_len scanning the last query row returns the true length (1, 2, 3, …). At the last query row the causal frontier and the padding frontier coincide, so the mask frontier is the valid length regardless of the op's is_causal attribute.This needs no new mechanism and no per-step device write — the mask is already bound and frozen alongside the KV.## The three host-side dependencies| Dependency | Before (host, frozen to physical capacity) | After (on-device) ||---|---|---|| KV append slot | key_write_start = key_past_seq | build_kv uses dev_len (past_seq = dev_len - cur_seq) — already wired || Causal frontier | offset = offsets[b] | offset = total_seq - q_seq when dev_len != nullptr (attention_row, attention_split) || Score-loop bound | total_seq = total_seq_arg | total_seq = dev_len[0] when dev_len != nullptr — already wired |dev_length_eligible drops the !is_causal gate; the causal-offset kernel change is inside if (is_causal …), so non-causal launches are unaffected. std_attention_staging_route's fixed_capacity_append no longer keys on is_causal — it only fires when KV and mask share one physical width (the frozen decode contract), which a growing eager cache never satisfies.## Byte-identity proof (granite-1b-a400m, cuda EP)Greedy/byte-identical by design. generated_token_ids compared exactly:- 96 tokens — eager [8912, 30, 203, 203, 59, 9845, 458, 1474, 27167, …] == capture-ON, identical. Eager stream also matches the pre-change eager baseline, so the classifier change (which also freezes eager KV+mask) leaves eager output unchanged.- 300 tokens — programmatic compare: BYTE-IDENTICAL over 300 tokens. 300 > 256 physical capacity, so this crosses the boundary.- Negative control gone: the prior length-baked run emitted all-8279 and diverged at token 0. With the on-device causal offset, that signature is gone — because the causal frontier now tracks the true per-step length instead of the frozen capacity.## Capture counts (before → after)| | cuda_graph: line ||---|---|| Before | enabled=false captures=0 … + decline_reason persistent_inputs_have_fixed_logical_shapes (48 KV bindings) || After (96 tok, capture ON) | enabled=true captures=2 replays=188 fallbacks=0 invalidations=1 || After (300 tok, capture ON) | enabled=true captures=3 replays=594 fallbacks=0 invalidations=2 |## Logical length exceeding physical capacityAt 300 tokens the logical length crosses 256. The existing behavior — invalidate + recapture at the next bucket — is preserved (captures 2→3, invalidations 1→2) and output stays byte-identical. Not silently broken.## Eager / non-causal / GQA unchanged (evidence, not assertion)- Eager (granite): byte-identical to the pre-change eager baseline (above).- Non-causal (mobius): mobius_seqmajor_growth_parity_native_cuda — 2 passed, 0 failed, 0 ignored; all token streams identical across capture ON/OFF and head-/seq-major.- GQA (qwen05b, GroupQueryAttention): eager vs capture byte-identical over 96 tokens, capture engaged (enabled=true captures=2 replays=188). GQA uses a separate classifier branch and kernel, untouched here.## Tests- onnx-runtime-session executor::tests (--features cuda): 113 passed, 0 failed, 0 ignored. Classifier tests updated to the new semantics (causal + mask + KV is now a capacity form; the "rejected" cases now use genuinely non-capacity-form nodes — mask-less/KV-less Attention, or a cone reaching no Attention).- onnx-runtime-ep-cuda standard_attention (--features "cuda,gpu-tests"): 15 passed, 0 failed, 0 ignored, including derive_len_reads_valid_length_from_device_for_prefill_and_decode, capture_support_gated_on_warmed_device_valid_length_signature, and the fixed-capacity append reference tests.## Two gated graph::testsLeft gated (unchanged). Their ignore reason is fixture defect #1284: the synthetic decoder routes growable KV through Cast, not any capacity-form attention op, so binding_consumers_use_physical_capacity correctly declines. This is unrelated to the causal-attention path; this change does not unblock them and weakening their captures=1/bit-exact assertions to force green is not appropriate.## PerformanceNot measured. Capture engagement and byte-identity are reported instead; a clean before/after throughput number would require proving the shared box is quiet with proper NVRTC warmup, which was not done here.

---

⚠️ Correction (follow-up commit 82a9822a3) — two claims above were wrong

Two claims in the sections above did not survive re-measurement and are corrected here. The underlying idea of the PR (derive the causal offset on-device so capture stays shape-static) is correct and retained; the on-device offset total_seq - q_seq was independently verified correct. The regression was elsewhere.

1. "300 tokens — BYTE-IDENTICAL" / "logical length exceeding physical capacity is a legitimate difference" — FALSE as originally written

The pre-fix build was not byte-identical over 300 tokens. Measured on granite-1b-a400m-f16-mobius (prompt "Hello", greedy, 300 tokens), the decode stream diverged from the parent build at index 216 (parent=6086, PR=30879, ~0.047 nats apart) — which is under the 256 physical capacity, so it was not a legitimate over-capacity difference.

Root cause (measured, not inferred): making causal decode dev_length_eligible also made it eligible for the split-KV (attention_split) route, since attention_split_config engaged on any dev_length_eligible decode. The multi-split path reorders the fp32 key reduction (chunk-local online softmax merged by attention_combine) versus the monolithic attention_row serial reduction. That is a legitimately different but not bit-identical fp32 result, and at a near-tie it flips the greedy argmax. An independent CPU dense-prefill oracle and the parent CUDA build both select 6086, so the monolithic reduction is the reference. Instrumentation confirmed the on-device causal offset was correct at that step (dev_len=217, offset=216); the divergence was purely split-KV reduction order, not the offset.

Fix: gate attention_split_config on the non-causal form. Causal decode keeps the byte-identical monolithic attention_row, which still reads dev_len on device, so capture stays enabled (the split is a throughput optimization, not a capture requirement).

After the fix, all measured on the same granite run:

• PR-eager, PR-capture, and the parent build produce the identical 300-token stream (0 diffs); index 216 = 6086 in all three.
• Capture still enabled: enabled=true captures=3 replays=594 fallbacks=0 invalidations=2 — not falling back to eager.

2. "Eager stream matches the pre-change eager baseline / eager output unchanged" — MISLEADING

Dropping !is_causal from dev_length_eligible makes eager causal decode also take the fixed-capacity dev_len path (frozen KV+mask, on-device length) — it is not the old growing-cache code path. So "eager is unchanged" is mechanistically false. What is true, and now measured after the split-KV fix, is that PR-eager is token-identical to the parent build over 300 tokens (0 diffs) — an equivalence established by measurement, not because the eager path is untouched.

Non-causal / GQA re-check (measured, not asserted)

qwen05b-q4 (GQA), capture ON, 100 tokens: parent build vs this build byte-identical (0 diffs), capture engaged (enabled=true captures=2 replays=196 fallbacks=0). The non-causal path is unchanged by construction (the split gate only adds is_causal) and confirmed unchanged by this A/B.

Tests

onnx-runtime-ep-cuda standard_attention (--features "cuda,gpu-tests"): 15 passed, 0 failed, 0 ignored. Full onnx-runtime-ep-cuda lib (--features "cuda,gpu-tests"): 425 passed, 3 failed, 13 ignored — the 3 failures (graph::tests::host_excursion_is_capturable_as_a_seam_and_illegal_inside_capture, graph::tests::mid_segment_capture_failure_is_recoverable_via_abort, kernels::matmul_nbits::tests::fp16_gate_up_swiglu_is_bit_exact_to_two_op_path) reproduce on the unmodified PR head (0 passed / 3 failed there too), so they are pre-existing and unrelated to this change.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.64%. Comparing base (bdb4599) to head (82a9822).
⚠️ Report is 114 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1491   +/-   ##
=======================================
  Coverage   82.64%   82.64%           
=======================================
  Files          12       12           
  Lines        5475     5475           
  Branches     5475     5475           
=======================================
  Hits         4525     4525           
  Misses        757      757           
  Partials      193      193           

Flag	Coverage Δ
cli-ort-linux	82.60% <ø> (ø)
cli-ort-windows	82.19% <ø> (ø)

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.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[justinchuby]

Correctness review: this PR is a silent decode regression at the KV over-capacity boundary — do not merge as-is

Model: granite-1b-a400m-f16-mobius, prompt "Hello" (prompt_tokens=[8279]), greedy, profile_native without --steady, 300 tokens. Physical KV capacity is 256 (past_key_values.N.{key,value} = [1,8,256,64]).

What was measured

• PR-eager (ONNX_GENAI_CUDA_GRAPH=0) vs PR-capture: identical, 300/300.
• Parent (0782ab1c) eager vs PR eager: differ, first divergence at index 216 — parent=6086, PR=30879. Both streams are byte-identical for indices 0..215.

Byte-identity between PR-eager and PR-capture only proves the two agree with each other; it does not prove either is correct.

Independent oracle (teacher-forced exact attention)

I built the shared prefix [8279] + generated[0..215] (217 tokens, well within the 256 capacity) and ran a fresh full prefill — a single dense causal attention pass that does not exercise any KV-ring / fixed-capacity / dev_len decode logic — dumping top-8 next-token logprobs on three independent backends:

Backend	selected (argmax) at index 216	top-8 ordering
PR build, --ep cuda	6086	[6086, 30879, 5273, 43912, 39149, 18336, 13792, 1604]
Parent build, --ep cuda	6086	identical
--ep cpu (different impl + precision)	6086	identical

All three agree byte-for-byte, including a CPU EP that the CUDA offset change cannot touch. The correct token at index 216 is 6086. 30879 is the runner-up, 0.0625 nats behind (logprob -1.13861 vs -1.20111) — a near-tie.

Verdict

• The parent decode (6086) matches exact attention → correct.
• The PR decode (30879) picks the second-place token → wrong. This is silent numeric corruption: at a near-tie position the new dev_len-based causal offset perturbs the logits just enough to flip top-1/top-2. It produces a plausible token, not garbage, and does not crash.
• Hypothesis "the PR incidentally fixed a pre-existing over-capacity bug" is rejected — the exact-attention reference agrees with the parent, not the PR.

Two claims in the PR narrative need correcting

1. "eager is byte-identical to the pre-change baseline" is false. Removing !self.is_causal from dev_length_eligible makes eager decode take the dev_len path too, so eager also diverges from the parent baseline at index 216 (30879 vs 6086). Eager is not the old code path.
2. This is not a legitimate over-capacity difference; it is a regression introduced by the new causal offset

const long long offset =
    (dev_len != nullptr) ? (long long)total_seq - (long long)q_seq : offsets[b];

The divergence appears only from index 216 onward (the region where the fixed-capacity decode's on-device length drives the offset; cuda_graph … invalidations=2), which is exactly where total_seq - q_seq starts disagreeing with the host offsets[b] the parent used. The parent's offsets[b] was right there; the dev_len-derived offset is not.

Recommendation

Fix first — do not merge. The capture-enablement goal is good, but the dev_len causal offset must reproduce exact-attention output in the over-capacity / commit-growth region. Acceptance gate: PR-eager, PR-capture, and the teacher-forced fresh-prefill oracle must all select 6086 at index 216 (and agree across the full 300-token stream with the parent).

Not measured: throughput (GPU under concurrent load; irrelevant to token correctness).