Govern Attention Phase-2a scratch via prepare-only workspace (#736) by justinchuby · Pull Request #753 · justinchuby/onnx-genai

justinchuby · 2026-08-10T19:40:45Z

Summary

Next increment on #736: audit the CUDA activation/scratch/workspace allocations that still bypass the device memory authority (beyond the QMoE path already governed by #747/#748), then govern the largest live bypass.

References #736. Does not close the issue.

Step 1 — Audit

A full inventory of production scratch/workspace allocations across onnx-runtime-ep-cuda, onnx-runtime-session and the CUDA kernels is added to docs/MEMORY_ARCHITECTURE.md §2.1, with file/line, byte formula, size class (static / config-derived / prompt-dependent) and lifetime class (model-fixed / session-persistent / request-prefill / step-scoped / already-governed) for every remaining raw path.

Ranking by live bytes, the largest ungoverned bypass is the com.microsoft::Attention Phase-2a scratch:

scores buffer batch · num_heads · sq · sk · elem — quadratic in sequence length (≈268–536 MB at B=1, H=32, S=2048), the single largest prompt-dependent scratch; plus
a static 32 MiB cuBLASLt workspace.

The fused (flash) prefill path allocates no device scratch (shared memory only), so run_attention_phase2a is the only scratch source in AttentionKernel — a clean StepScoped slice.

Step 2 — Govern one vertical slice

attention.rs: phase2a_scratch_layout + Phase2aScratch carve the scores buffer and cuBLASLt workspace from one governed blob. run_attention_phase2a uses an executor-prepared WorkspaceView when present; the compatibility/opt-out path (direct execute, GQA reuse, unit tests) keeps its self-owned alloc_raw scratch. New Kernel::workspace_requirement reports the exact bytes (StepScoped, NONE for the fused path) and execute_with_workspace consumes the reserved slot.
executor/bindings.rs: prepare-only planning (Plan governed QMoE workspace before admission #747) now tracks one workspace peak per lifetime class (SessionPersistent + StepScoped) and prepares both slots, so a graph mixing QMoE (session-persistent) and Attention (step-scoped) governs both. Reserve logic factored into reserve_prepared_workspace; node matching centralised in is_planned_workspace_node.
Same helper for plan and exec: planning and execution both size through use_fused + phase2a_scratch_layout, so the reserved slot always covers the dispatch. A shortfall would error deterministically rather than silently under-allocate.
429, not 500/OOM: reservation happens inside prepare_with_device_bindings (pre-header, prefill admission). reserve_workspace/prepare_mapped_growth emit TierExhausted/CapacityUnavailable, which DriverFailure::from_engine_error already maps to HTTP 429 (Return overload response for KV admission refusal #743) — no server change needed.
Lifetime correctness: step-scoped scratch is released by the executor's existing release_step_workspace on success/error/cancellation; the governed kernel path no longer owns or frees it. No asserts in any Drop.
Regression guard: new unit test pins phase2a_scratch_layout's byte formula; a GPU-gated test asserts workspace_requirement equals the layout for a Phase-2a shape and returns 0 (NONE) for a fused shape — preventing a raw-allocation regression on the governed slice.

Step 3 — Document

docs/MEMORY_ARCHITECTURE.md §2.1 records the governed-today table (QMoE, Attention Phase-2a) and every remaining raw path. It identifies GroupQueryAttention's WS_SCORES slot (batch·heads·sq·sk·4, session-persistent growable) as the recommended next increment, with the three shared 32 MiB cuBLASLt GEMM workspaces as a cheap follow-up.

Validation

cargo fmt --all -- --check ✅
cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings ✅ (exact CI gate, no --all-targets)
cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings ✅
Session executor lib tests (76 incl. prepare_reserves_static_nested_qmoe_workspace_and_child_reuses_it) ✅
CUDA EP attention lib tests incl. new guard tests ✅ (GPU present)
GPU tests: attention_gpu (5), qmoe_gpu (27), allocation_counts_gpu (2), group_query_attention_gpu (25, single-threaded), standard_attention_gpu ✅

Note: gqa_gpu_fixed_decode_capture_replays_bit_identically fails only when many capture tests share one binary in parallel (pre-existing CUDA graph-capture stream interaction, CUDA_ERROR_STREAM_CAPTURE_INVALIDATED); it passes in isolation and single-threaded. Unrelated to this change (the GQA edit only adds a None arg to the prefill helper).

Per instructions: local validation only, not waiting for CI; PR opened but not merged and issue not closed.

Audit and govern the largest ungoverned CUDA scratch bypass beyond the QMoE path (#747/#748): the com.microsoft::Attention Phase-2a quadratic scores buffer (batch*num_heads*sq*sk*elem) plus its 32 MiB cuBLASLt workspace. - attention.rs: add phase2a_scratch_layout + Phase2aScratch; route run_attention_phase2a through an executor-prepared WorkspaceView when present; add Kernel::workspace_requirement (StepScoped, NONE for the shared-memory fused path) and execute_with_workspace. The compat/opt-out path keeps self-owned alloc_raw scratch. - executor/bindings.rs: extend prepare-only planning to track one workspace peak per lifetime class (SessionPersistent + StepScoped) and prepare both slots, so a graph mixing QMoE and Attention governs both. Factor reserve_prepared_workspace; match on is_planned_workspace_node. - Planning and execution size through the identical phase2a_scratch_layout + use_fused helpers, so the reserved slot always covers execution; capacity refusal surfaces as a pre-header HTTP 429 (#743). - Add layout/requirement guard tests preventing a raw-allocation regression on the governed slice. - docs/MEMORY_ARCHITECTURE.md: add the §2.1 scratch/workspace governance audit table listing every remaining raw path with byte formula and lifetime class; identify GroupQueryAttention's WS_SCORES as the next increment. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

codecov · 2026-08-10T19:48:09Z

Codecov Report

❌ Patch coverage is 68.91892% with 23 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.40%. Comparing base (ca8b2bf) to head (92ccd97).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
...ates/onnx-runtime-session/src/executor/bindings.rs	68.91%	17 Missing and 6 partials ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #753   +/-   ##
=======================================
  Coverage   79.40%   79.40%           
=======================================
  Files         331      331           
  Lines      137528   137568   +40     
  Branches   137528   137568   +40     
=======================================
+ Hits       109198   109234   +36     
- Misses      23986    23990    +4     
  Partials     4344     4344

Flag	Coverage Δ
cli-ort-linux	`83.41% <ø> (ø)`
cli-ort-windows	`82.90% <ø> (-0.10%)`	⬇️
mlas	`80.92% <ø> (ø)`
offline	`79.21% <68.91%> (+<0.01%)`	⬆️

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

Files with missing lines	Coverage Δ
...ates/onnx-runtime-session/src/executor/bindings.rs	`51.60% <68.91%> (+2.62%)`	⬆️

... and 1 file 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 · 2026-08-10T20:09:17Z

🔴 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_uncached_dequant_each_call/1x1024x1024`	722.45 µs	1.72 ms	+137.4%
⚠️	`gather/large_f16_threads=1-internal/131072`	20.54 µs	26.18 µs	+27.5%
⚠️	`gather/large_bf16_threads=1-internal/131072`	17.26 µs	21.33 µs	+23.6%
⚠️	`block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024`	89.61 µs	107.46 µs	+19.9%
✅	`sampling_latency/min_p_per_token`	206.56 µs	234.49 µs	+13.5%
✅	`matmul/large_generic_f32_threads=8/32x1024x1024`	6.66 ms	7.35 ms	+10.3%
✅	`matmul/small_generic_f16_threads=8/1x256x256`	35.42 µs	38.96 µs	+10.0%
✅	`matmul/small_generic_f16_threads=1/1x256x256`	38.00 µs	41.18 µs	+8.4%
✅	`tokenization/encode_tokens_per_second`	404.08 µs	435.05 µs	+7.7%
✅	`gather/medium_bf16_threads=1-internal/32768`	2.82 µs	3.01 µs	+6.8%
✅	`block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1`	83.24 µs	87.47 µs	+5.1%
✅	`matmul/medium_generic_f16_threads=8/32x512x512`	47.42 µs	47.80 µs	+0.8%
✅	`matmul/medium_generic_f16_threads=1/32x512x512`	41.71 µs	42.01 µs	+0.7%
✅	`grammar_masking/llguidance_compute_mask/32`	78.78 µs	79.04 µs	+0.3%
✅	`gather/small_bf16_threads=1-internal/4096`	540.3 ns	540.1 ns	-0.0%
✅	`matmul/large_generic_bf16_threads=1/32x1024x1024`	2.55 ms	2.55 ms	-0.3%
✅	`qwen3_sampling_processors/top_k_full_sort_baseline`	2.40 ms	2.39 ms	-0.4%
✅	`qwen3_sampling_processors/top_k_top_p_full_sort_baseline`	6.17 ms	6.09 ms	-1.4%
✅	`matmul/medium_generic_f32_threads=1/32x512x512`	2.88 ms	2.83 ms	-1.6%
✅	`matmul/large_generic_f32_threads=1/32x1024x1024`	11.51 ms	11.24 ms	-2.3%
✅	`matmul/large_generic_bf16_threads=8/32x1024x1024`	2.43 ms	2.37 ms	-2.4%
✅	`gather/large_f32_threads=1-internal/131072`	54.69 µs	53.35 µs	-2.5%
✅	`sampling_latency/greedy_per_token`	3.51 µs	3.30 µs	-5.9%
✅	`sampling_latency/top_k_per_token`	66.67 µs	62.67 µs	-6.0%
✅	`matmul/large_generic_f16_threads=1/32x1024x1024`	107.83 µs	100.41 µs	-6.9%
✅	`qwen3_sampling_processors/top_k_top_p_fast`	716.32 µs	659.49 µs	-7.9%
✅	`gather/small_f16_threads=1-internal/4096`	548.0 ns	496.3 ns	-9.4%
✅	`matmul/small_generic_f32_threads=8/1x256x256`	49.75 µs	45.00 µs	-9.6%
✅	`qwen3_sampling_processors/top_k_partial_selection`	177.50 µs	159.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`	558.16 µs	501.57 µs	-10.1%
✅	`add/medium_f32_threads=1-internal/262144`	2.96 ms	2.66 ms	-10.3%
✅	`matmul/large_generic_f16_threads=8/32x1024x1024`	115.80 µs	103.73 µs	-10.4%
✅	`matmul/small_generic_f32_threads=1/1x256x256`	49.49 µs	44.17 µs	-10.7%
✅	`logit_processing/seven_processor_chain_per_step`	364.63 µs	321.04 µs	-12.0%
✅	`gather/medium_f16_threads=1-internal/32768`	3.06 µs	2.67 µs	-12.9%
✅	`matmul/medium_generic_bf16_threads=1/32x512x512`	706.00 µs	612.06 µs	-13.3%
🟢	`kv_cache/alloc_dealloc_pages`	45.06 µs	38.26 µs	-15.1%
🟢	`qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline`	4.48 ms	3.79 ms	-15.4%
🟢	`sampling_latency/top_p_per_token`	504.94 µs	422.31 µs	-16.4%
🟢	`reduce_mean/large_f32_threads=1-internal/262144`	1.18 ms	984.28 µs	-16.7%
🟢	`qwen3_sampling_processors/top_p_fast_after_top_k`	624.21 µs	519.49 µs	-16.8%
🟢	`matmul/medium_generic_f32_threads=8/32x512x512`	1.81 ms	1.48 ms	-18.0%
🟢	`add/large_bf16_threads=1-internal/4194304`	51.50 ms	42.14 ms	-18.2%
🟢	`gather/medium_f32_threads=1-internal/32768`	5.28 µs	4.30 µs	-18.5%
🟢	`add/medium_f16_threads=1-internal/262144`	3.50 ms	2.85 ms	-18.6%
🟢	`tokenization/decode_tokens_per_second`	8.04 ms	6.54 ms	-18.6%
🟢	`add/small_bf16_threads=1-internal/1024`	17.42 µs	14.02 µs	-19.5%
🟢	`reduce_mean/small_f32_threads=1-internal/4096`	18.59 µs	14.92 µs	-19.8%
🟢	`matmul/small_generic_bf16_threads=1/1x256x256`	43.48 µs	33.92 µs	-22.0%
🟢	`add/large_f32_threads=1-internal/4194304`	56.47 ms	42.86 ms	-24.1%
🟢	`matmul/small_generic_bf16_threads=8/1x256x256`	44.45 µs	33.56 µs	-24.5%
🟢	`add/medium_bf16_threads=1-internal/262144`	4.05 ms	2.96 ms	-26.9%
🟢	`gather/small_f32_threads=1-internal/4096`	912.1 ns	666.9 ns	-26.9%
🟢	`reduce_mean/medium_f32_threads=1-internal/65536`	333.84 µs	242.81 µs	-27.3%
🟢	`add/large_f16_threads=1-internal/4194304`	58.82 ms	41.01 ms	-30.3%
🟢	`add/small_f32_threads=1-internal/1024`	330.0 ns	219.3 ns	-33.6%
🟢	`add/small_f16_threads=1-internal/1024`	22.10 µs	13.94 µs	-36.9%
🟢	`block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024`	138.43 µs	87.24 µs	-37.0%
🟢	`matmul/medium_generic_bf16_threads=8/32x512x512`	1.05 ms	591.93 µs	-43.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: { 4.00 3.65 6.40 }

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)

Resolve conflicts in executor/bindings.rs and kernels/attention.rs. bindings.rs (vs #751): centralise governed-node detection in is_planned_workspace_node, extended to cover BlockQuantizedMoE, IndexShare and com.microsoft::Attention Phase-2a, with per-lifetime-class workspace peaks preserved. attention.rs (vs #757): keep the use_fused helper shared by planning and execution while preserving the capture-safe store and Phase-2a capture refusal so capture_support stays truthful and capture_sync_contract holds. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

justinchuby · 2026-08-11T16:51:09Z

Merge conflict resolution + re-validation (16 commits behind `origin/main`)

Merged origin/main into squad/736-scratchaudit. Two files conflicted; both were semantic (main changed the same mechanism this PR touches). Resolved as follows:

1. `crates/onnx-runtime-session/src/executor/bindings.rs` — vs #751 (`d2db1f5a`, Govern IndexShare CUDA workspace)

#751 widened prepare-only traversal from BlockQuantizedMoE only to matches!(op_type, "BlockQuantizedMoE" | "IndexShare") inline in three places; this PR refactored the same predicate into the centralised is_planned_workspace_node + reserve_prepared_workspace with per-lifetime-class peaks. Resolution: kept this PR's centralised is_planned_workspace_node at all three call sites and folded IndexShare into the predicate itself, so it now governs all three node kinds — BlockQuantizedMoE, IndexShare, and com.microsoft::Attention Phase-2a — through one predicate, with the SessionPersistent/StepScoped per-class peaks preserved. IndexShare governance from #751 is not dropped.

Carried forward and verified from #751:

The mock EP's reserve_workspace in executor/tests.rs still handles both static-byte fixtures (exact assert_eq!) and per-row fixtures (workspace_bytes_per_row, asserting a non-vacuous positive multiple — still fails on zero/mis-sized). Merged cleanly, present.
IndexShare's present-K/V staging still sizes to actual use (3-output decode path routes present K/V to outputs, not workspace). index_share.rs merged cleanly from main; my predicate-only change doesn't touch its sizing.

2. `crates/onnx-runtime-ep-cuda/src/kernels/attention.rs` — vs #757 (`a3856b6a`, close LinearAttention C2 capture seams)

The merge cleanly pulled in main's last_call_capture_safe field, its init, and the dynamic capture_support(). The only conflict was the fused-vs-Phase-2a selection block in run, where this PR extracted the logic into the use_fused helper (shared by planning + execution) while main kept it inline plus the capture-safe store and Phase-2a capture-refusal. Resolution: kept this PR's use_fused helper and re-applied #757's last_call_capture_safe.store(fused, …) + if !fused && is_capturing() { return Err(…) } guard. This keeps capture_support() truthful (the flag is actually set) and keeps run_attention_phase2a's allowlisted unconditional sync valid — the Phase-2a path stays capture-unsupported, so #757's capture-aware sync contract holds. §736 governs the Phase-2a scratch for admission but does not make the sync capture-safe. The capture-aware sync guard from #757 is not silently restored to an unconditional sync.

Validation (all local, per the exact CI gates)

cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --check — clean
cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings — clean
cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings — clean
cargo test -p onnx-runtime-session — full suite green (this is where Govern IndexShare CUDA workspace #751 was previously broken by a mock expectation main had changed; verified intact)
cargo test -p onnx-runtime-ep-cuda --features cuda --test capture_sync_contract — 1 passed (unconditional_syncs_are_limited_to_capture_unsupported_paths; attention.rs::run_attention_phase2a remains correctly allowlisted)
GPU, single-threaded: attention_gpu 5 passed, standard_attention_gpu 24 passed, qmoe_gpu 29 passed, allocation_counts_gpu 2 passed, group_query_attention_gpu 25 passed — including gqa_gpu_fixed_decode_capture_replays_bit_identically, which passes in isolation/single-threaded (the construction_gpu: CUDA graph capture is invalidated under parallel test execution #626 parallel-capture flake was not triggered), confirming that behaviour is unchanged.

Pushed as 92ccd970. Both #751 (IndexShare governance) and #757 (capture-aware sync) are preserved.

justinchuby · 2026-08-11T16:53:41Z

Coordinator pre-merge validation on the merged branch (Windows, RTX box):

cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --check clean
cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings clean (exact CI gate, no --all-targets)
full cargo test -p onnx-runtime-session: 151 passed, 0 failed — mandatory here, because Govern IndexShare CUDA workspace #751 was broken by exactly this class of merge interaction

Both conflicts were semantic rather than cosmetic, and both were resolved in the
direction that preserves what already landed on main:

bindings.rs vs Govern IndexShare CUDA workspace #751 — this PR's centralised is_planned_workspace_node
predicate is kept and IndexShare folded into it, so prepare-only planning now
governs BlockQuantizedMoE, IndexShare and Attention Phase-2a through one
predicate, with per-lifetime-class peaks preserved. Govern IndexShare CUDA workspace #751's dual-mode mock
reserve_workspace assertion and its actual-use present-K/V staging fix both
survive.
attention.rs vs perf(cuda-ep): close LinearAttention C2 capture seams #757 — the shared use_fused helper is kept while perf(cuda-ep): close LinearAttention C2 capture seams #757's
last_call_capture_safe store and Phase-2a capture refusal are re-applied, so
capture_support() stays truthful and the allowlisted sync remains valid.
capture_sync_contract passes.

On the substance: this governs the largest ungoverned bypass the #736 audit
found — the Phase-2a scores buffer at batch x heads x sq x sk x elem, which is
quadratic in sequence length (~268-536 MB at B=1, H=32, S=2048).

Two things make it a better increment than the byte count alone suggests. The
audit established that the fused (flash) prefill path allocates no device
scratch at all (shared memory only), so run_attention_phase2a is the only
scratch source in AttentionKernel — a clean StepScoped slice rather than a
partial one. And extending prepare-only planning to track a peak per lifetime
class fixes a structural gap in #747: a graph mixing QMoE (session-persistent)
with Attention (step-scoped) previously could not govern both.

The audit table in docs/MEMORY_ARCHITECTURE.md §2.1 also names the next
increment from evidence rather than guesswork — GroupQueryAttention's
WS_SCORES slot, with the three shared 32 MiB cuBLASLt GEMM workspaces as a
cheap follow-up.

Merging.

## Govern the `GroupQueryAttention` `WS_SCORES` workspace (#736) Next increment of the CUDA allocation-governance program (#736), chosen from the audit table in `docs/MEMORY_ARCHITECTURE.md` §2.1, which named this slice explicitly. Routes the `GroupQueryAttention` f32 reference-attention score buffer (`WS_SCORES`) — quadratic in `sq * present_capacity` and the largest remaining live GQA scratch — through the merged workspace-governance contracts instead of the raw `self.runtime.alloc_raw` pooled bypass. Closes #736 — no, **does not** close: this is one increment; the issue stays open. ### What changed - **`crates/onnx-runtime-ep-cuda/src/kernels/group_query_attention.rs`** - `Kernel::workspace_requirement` reports the exact governed bytes; `execute_with_workspace` consumes the executor-prepared slot; the compatibility `execute` path keeps its self-owned pooled scratch. - Planning and execution size the reservation through the **same** `gqa_reference_scores_bytes` helper, so the reserved and consumed byte counts cannot drift. A shortfall in the prepared buffer **errors deterministically** rather than silently under-allocating. - New free helpers: `gqa_reference_scores_path(dtype, q_seq, head_dim)` (the static reachability predicate) and `gqa_reference_scores_bytes(batch, heads, q_seq, kv_capacity)` (the shared, overflow-checked byte formula `rows * kv * sizeof(f32)`, `.max(1)`). - CPU-only source-scan **regression guard** (`raw_allocation_guard`) that fails if the governed slot reintroduces an unconditional raw pooled allocation — matches the #751 bar (a unit test that runs on CI without a GPU, not a `Select-String` external scan). Needle assembled at runtime so the literal does not self-match. - **`crates/onnx-runtime-session/src/executor/bindings.rs`** - Extended the centralised `is_planned_workspace_node` predicate to cover `com.microsoft::GroupQueryAttention` alongside `BlockQuantizedMoE`, `IndexShare`, and `Attention`. No parallel predicate added. - **`docs/MEMORY_ARCHITECTURE.md` §2.1** - Moved the GQA `WS_SCORES` slot into the governed-today rows, recorded derived bytes, documented the lifetime justification and the over-reservation finding, and reselected the next slice from evidence. ### Measured / derived bytes `gqa_reference_scores_bytes = (batch · num_heads · q_seq · present_capacity).max(1) · sizeof(f32)` — identical to the old `score_count.max(1) * 4`. - Reference path, `B=1, H=32, Sq=1024, present_capacity=1024, f32` ≈ **128 MiB**. - fp16/bf16 decode & prefill, and **all** f32 single-token decode: **0 bytes** (reported `WorkspaceRequirement::NONE`). ### Lifetime justification — `SessionPersistent` (not `StepScoped`) GQA's slot allocator grows each slot to the largest geometry seen and **retains** it across decode/prefill steps — a standing claim held for the whole session. So `WS_SCORES` is charged as `WorkspaceLifetime::SessionPersistent` + `MemoryRole::Workspace { step_scoped: false }`. This deliberately differs from the Attention Phase-2a scratch landed in #753, which was per-call `StepScoped`. Because `sq * present_capacity` is prompt-dependent, the executor grows the slot transactionally via `MappedGrowthGrant` (#748) against the device authority, and capacity refusal surfaces as the pre-header **429** typed path (`TierExhausted` / `CapacityUnavailable`, #743) rather than a late 500/OOM. A graph mixing QMoE, IndexShare, Attention Phase-2a (step-scoped) and GQA (session-persistent) now exercises **both** per-lifetime-class peaks; the classes reserve into separate executor slots so neither peak stands in for the other. ### Over-reservation finding (the #751 defect class) **There is over-reservation to avoid, and this PR avoids it.** `WS_SCORES` is materialized on **exactly one** GQA route: the f32 *reference* path, reachable only when `query dtype == f32 && !gqa_decode::supported(q_seq, head_dim)` (i.e. f32 prefill `Sq > 1`, or f32 `head_dim > 128`). The fused-flash path, the capture-safe f32/fp16 split-K decode kernels, and the phase-2a path never allocate a device score matrix — they stream softmax through registers/shared memory (phase-2a owns its own scratch). Reserving `WS_SCORES` unconditionally would charge the device authority for bytes the common fp16 decode path never touches — exactly the defect #751 found in IndexShare's present-K/V staging. `workspace_requirement` threads that static signal (query dtype + `gqa_decode::supported`) through the shared helper and reports **NONE** on every fp16/bf16 path and every f32 *decode* path, so they charge zero. ### cuBLASLt GEMM workspaces The audit named three shared 32 MiB cuBLASLt GEMM workspaces (`fused_gemm.rs`, `gemm.rs`, `matmul.rs`) as a cheap follow-up. They do **not** fall out of this change — they are separate step-scoped cuBLASLt workspaces, not GQA slots — and are left for a dedicated slice, as the task permitted. ### Validation (all local, green) - `cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --check` — clean. - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` — 0 warnings. - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` — 0 warnings. - `cargo test -p onnx-runtime-session` — **green** (mandatory 151+ suite). - `cargo test -p onnx-runtime-ep-cuda --features cuda --lib` — 318 passed (incl. the 4 new GQA governance tests). - GPU suites single-threaded (`--features cuda,gpu-tests -- --test-threads=1`): `group_query_attention_gpu` **25 passed** (incl. `gqa_gpu_fixed_decode_capture_replays_bit_identically`, which passed here — the #626 flake only appears under parallel capture contention), `qmoe_gpu` 29, `attention_gpu` 5, `allocation_counts_gpu` 2, `gqa_seqmajor_parity_gpu` 1. ### Next slice (recorded in the doc, chosen from evidence) `standard_attention.rs:602` (default-domain `Attention` pooled scores slot) — same quadratic shape and session-persistent retention — through the same contract with the same over-reservation discipline, then fold in the legacy `attention.rs:711` per-call path. ### Notes - Edits confined to workspace/allocation concerns; KV indexing untouched, to merge cleanly with the concurrent flash/prefill stride-descriptor work in this crate. - No `assert!` in any `Drop` (STATUS_STACK_BUFFER_OVERRUN on this platform). - Do **not** merge and do **not** close #736. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>

…#755) (#798) Closes #755. Makes the authority-scoped **managed no-spill VMM** the **default** on the native CUDA path — no flag required — with **automatic weight streaming** when a model exceeds the resolved device budget. Being larger than the budget becomes a supported configuration rather than an error. Builds directly on the prerequisite #716 (PR #796), which made weight offload capture-compatible by paging weights under a stable virtual address, so flipping the default no longer disables CUDA graph capture for offloaded models. ### What "default" now means 1. **Managed VMM is on without a flag.** Native CUDA loads run memory-strategy inference (#752) unconditionally and it *drives* policy: allocation goes through the physical-handle pool (#740) with committed-granule admission (#745), governed workspace planning (#747/#753/#795), transactional growth grants (#748), against one server-owned authority per device (#744). 2. **Over-budget is not an error.** When the resolved budget cannot hold the weights, the runtime auto-enables weight streaming/offload instead of failing. 3. **`--vram-limit` is an override**, not the trigger. 4. **The legacy allocator stays reachable** for one release via `ONNX_GENAI_LEGACY_ALLOCATOR=1` (back-compat: `ONNX_GENAI_DYNAMIC_KV_WEIGHT_LENDING=0`). ### Observability The resolved budget, chosen strategy and offload state are printed at startup and exposed in `/v1/resources` under a new `memory_strategy` block, so the new default is observable rather than mysterious. ### Invariants preserved - **No silent WDDM spill** — committed physical bytes stay within the managed cap; `cuMemCreate` structurally cannot spill (#704). - **Capacity refusal stays a pre-header HTTP 429** (#743); committed-granule admission (#745) and growth grants (#748) unchanged. - **A model that fits does not start paging** under the new default — tested explicitly (the most likely regression). - **Capture stays ON when offload is active on the stable-VA path** (#796 property), pinned by the existing unit tests. - No prefetch re-introduced (#715); no assert inside any `Drop`. ### Measurement (native CUDA, deterministic counters lead) Ignored live e2e `qwen2_5_0_5b_managed_vmm_default_e2e` (Qwen2.5-0.5B int4 mobius, **same process**, interleaved managed-default vs. legacy opt-out): | scenario | strategy | offload | managed_no_spill | resolved budget | committed device bytes | page-ins / evictions | |---|---|---|---|---|---|---| | **A** fits, managed default (no flag) | `FullResident` | false | true | 7,730,940,928 | 381,681,664 | **0 / 0** (does not page) | | **B** synthetic over-budget (`--vram-limit 268,435,456`) | `DynamicWeightResidency` | true | true | 268,435,456 | 90,177,536 (« cap; no spill) | **434 / 432** | | **C** legacy opt-out (`ONNX_GENAI_LEGACY_ALLOCATOR=1`) | `FullResident` | — | **false** | — | — | — | | **D** explicit `--vram-limit 12 GiB` | override applied | — | — | 12,884,901,888 | — | — | In **B**, at that extreme synthetic cap the residency arithmetic **refused cleanly** ("no page is evictable") rather than spilling — the no-spill contract holds. **Caveats (measurement discipline):** - Wall-clock / tok/s is **not** reported: this box ranged 3.9–28 tok/s across identical runs and other agents build concurrently. Deterministic counters only. - A genuine native large-model run that exceeds VRAM was **not measured**: `qwen14b-zp` lacks `inference_metadata.yaml` (#384) so it is not native-loadable. The over-budget condition was **synthesized** with a small explicit budget, as stated. - **CUDA-graph segment count** is not a public counter; capture-ON-under-offload on the stable-VA path is pinned by the #796 unit tests, not re-measured here. ### Validation - `cargo fmt -p onnx-genai-engine -p onnx-genai-server -- --check` clean (`--all-targets --all` fmt hits Windows os error 206 on this box's long worktree paths — a command-line-length artifact, not a formatting issue). - clippy `-D warnings`: engine (`native-backend`; also `cuda,native-backend`), server, `onnx-runtime-ep-cuda --features cuda` — all clean. - Full `cargo test -p onnx-runtime-session` passes. - Engine `memory_strategy` unit tests (12, incl. 4 new) and server resources tests pass; cuda-gated engine tests incl. #796 capture arms pass. - **Pre-existing unrelated failure:** `failed_runtime_shrink_preserves_policy_and_ledger_limit` (server) fails identically on `origin/main` (confirmed via stash); its error is on the `host` authority, untouched by this native-CUDA-only change. Not fixed here (out of scope). Do **not** merge and do **not** close the issue via this PR beyond the `Closes #755` reference; leaving final disposition to review. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>

## Summary - record the three-slice over-reservation finding and path-first audit guidance - refresh governed and raw workspace status from current CUDA/session code - correct the default-domain Attention score sizing/lifetime and name it as the next slice ## Validation - `git diff --check` - confirmed exactly one Markdown file changed and no non-Markdown files changed References #736 References #751 References #753 References #795 References #799 References #745 References #750 References #777 References #781 References #791 Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>

References #736. Do not merge; do not close the issue. ## Summary Governs the default-domain `Attention` (`StandardAttentionKernel`) f32 score matrix in the CUDA EP, routing its device scratch through the merged `Kernel::workspace_requirement` + `execute_with_workspace` contracts with prepare-only planning (#747), so it is reserved against the memory authority and a capacity shortfall surfaces as the pre-header **429** typed path (`TierExhausted` / `CapacityUnavailable`, #743) via `MappedGrowthGrant` (#748) rather than a late 500/OOM. ## Measured bytes `batch * q_heads * q_seq * total_seq * sizeof(f32)` — quadratic in sequence length. At `B=1, H=32, q_seq=total_seq=2048` = **512 MiB** (536,870,912 bytes), larger than the already-governed Phase-2a range (#753). The buffer is always fp32 (fp32 accumulate) regardless of Q/K/V operand dtype. ## Which routes materialize the matrix (negative result) **Every route.** Unlike GQA's f32-reference-only `WS_SCORES` (#795) or the fused Attention prefill path that streams softmax through shared memory (#753), this kernel's `attention_row` has **no** flash/shared-memory route: it always stages the score matrix in global device memory (writes `scores[srow+j]`, re-reads for softcap/mask/softmax/`probs·V`). So every valid dispatch reserves it and **none reports `WorkspaceRequirement::NONE`** — only unresolvable or non-float input metadata reports `NONE`. This is a genuine-use result, recorded in the audit so the next reader does not re-investigate expecting a zero-charge route. ## Lifetime split Modelled truthfully as **both** classes: - **StepScoped** — per-call prefill/batched route (freed on each exit); this is the 512-MiB-class quadratic worst case. - **SessionPersistent** — the capture-eligible single-token decode route (`batch == 1 && q_seq == 1`), pooled/retained in the kernel workspace (small: linear in `total_seq`, ~256 KiB at the example shape). The classifier keys the split on static input metadata (`batch`, `q_seq`) as a conservative proxy for capture eligibility. Prepare-only planning reserves each class into its own executor slot (#753). The frozen fixed-capacity decode path over-estimates `total_seq` by padding rows, which is safe (reserve >= consume) — execution re-derives the exact size through the **same** `std_attention_scores_bytes` helper and refuses deterministically on any shortfall. ## Over-reservation found None. Planning and execution size through the same helper. The only sibling bypass — the staged-K/V + capture-metadata slots (`standard_attention.rs:702,1396`) — is a genuinely different lifetime/geometry, did **not** fall out of this change, and is left raw (tracked as the next slice). ## Scope / safety - Extends the central `is_planned_workspace_node` predicate (default-domain `"Attention"`); **no** parallel predicate. - Preserves the compatibility/opt-out direct-`execute()` path with self-owned scratch (pooled on the capture-eligible route, per-call otherwise) — unit/GPU allocation-count tests unaffected. - Released on success, error, and cancellation with truthful role/holder attribution. No assert inside any `Drop`. - Adds CPU-only governance unit tests + a source-scan regression guard (runtime-assembled needle) that fails on CI without a GPU if a raw score allocation is reintroduced on the governed site. ## Updated audit table `docs/MEMORY_ARCHITECTURE.md` §2.1: the default-domain `Attention` scores move into the governed rows (both the summary "Governed today" table and the site-by-site audit), with the route/lifetime split and the no-flash-route/no-`NONE` note verified against `attention_row`. The next slice is named from evidence: **`GroupQueryAttention` pooled QKV-projection staging** (`group_query_attention.rs:1033`, `batch * seq * hidden * elem`, session-persistent), now the largest remaining evidenced non-QMoE workspace bypass after the three quadratic score matrices are governed; present K/V there is KV-cache residency (govern via the KV residency contract, #791), not this table. ## Validation (local; did not wait for CI) - `cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --check` — clean - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` — clean (exact CI gate, no `--all-targets`) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` — clean - `cargo test -p onnx-runtime-session` — green - GPU suites single-threaded (`--features cuda,gpu-tests`): `standard_attention_gpu` (25), `attention_gpu`, `group_query_attention_gpu` (25), `qmoe_gpu`, `allocation_counts_gpu` — all green Working as the workspace/allocation slice; edits kept to workspace/allocation concerns to merge cleanly with concurrent KV shape/stride work in `native_decode/`. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>

Governs the `GroupQueryAttention` packed QKV-projection staging buffers — the largest remaining evidenced non-QMoE workspace bypass after the three quadratic score matrices (Phase-2a #753, GQA `WS_SCORES` #795, default-domain `Attention` #736). References #736. **Do not merge; do not close the issue.** ## What was ungoverned `crates/onnx-runtime-ep-cuda/src/kernels/group_query_attention.rs` — the pooled slot allocator (`bytes.max(1)` per slot) raw-allocated the packed Q/K/V projection-staging buffers (`WS_PACKED_Q/K/V`). The unfused prep path splits an interleaved `[B, S, (H + 2·Hkv)·D]` query tensor into three device scratch buffers before the BSH→BNSH transpose and cache append. Those bytes were invisible to the device authority. ## Measured / derived bytes Formula: `align256(B·sq·num_heads·head_dim·elem) + 2·align256(B·sq·kv_num_heads·head_dim·elem)`, `elem = dtype.byte_size()` (f32=4, f16/bf16=2). At `B=1, num_heads=32, head_dim=128` (`q_hidden=4096`), `kv_num_heads=8` (`k_hidden=1024`), `sq=1024`: - **f32 packed prefill:** 16 MiB (Q) + 2·4 MiB (K,V) = **24 MiB** staging, plus the ~128 MiB f32 reference score region (#795) when that route is taken. - **fp16/bf16 packed prefill:** **12 MiB** staging, **0** scores. - **unpacked-input and fused-decode routes:** **0** staging. ## Which routes populate staging vs report `NONE` Staging is materialized on **exactly one** family of routes: a *packed* QKV tensor split on the unfused prep path (`packed_qkv && !fuse_prep`). - **Unpacked inputs** (Q/K/V arrive as three separate tensors) never split → staging regions are zero-length → charge **0**. This is the primary "size from use" finding (#751 class): the pre-#736 pool grew `WS_PACKED_*` on any path reaching `reserve`, but the separate-input route writes Q/K/V straight through. - **fp16/bf16 fused decode** (`q_seq==1` with an aliased device KV cache) fuses prep into the decode kernel and never splits. - The requirement returns `WorkspaceRequirement::NONE` when neither staging nor f32 reference scores are needed. ## Over-reservation found Yes (this is **not** a #736-default-Attention-style genuine-use-everywhere result). Two routes charge zero. One residual conservative reservation: `fuse_prep` aliases the present cache via a **runtime pointer comparison** not visible at planning, so planning reserves staging for every packed dispatch; on fused decode the reserved region is tiny (`q_seq==1`) and left untouched rather than corrupting. Every packed *prefill* genuinely splits, so this is a fallback-of-a-fallback idle. ## Lifetime model + justification **Session-persistent, single class** — verified against the `GqaWorkspace` slot allocator (`group_query_attention.rs:1191`), which grows each slot to the largest geometry seen and **retains** it across decode/prefill steps. Unlike the default-domain `Attention` scores (route-dependent lifetime needing both classes modelled), this lifetime is not route-dependent, so one class models the truth. The staging shares the same session-persistent composite reservation as the #795 `WS_SCORES` region (the executor reserves one slot per lifetime class and hands the kernel one view; the f32 reference-prefill route needs both live at once → disjoint regions of one buffer). Staging regions come first (shape-derived offsets, identical in planning and execution); the score region is last and unpadded so its capacity-dependent size extends into the reserved peak without perturbing staging offsets, and the composite total equals exactly the score bytes when no staging exists (preserves #795 accounting). ## Contract wiring - `Kernel::workspace_requirement` + prepare-only planning (#747); `execute_with_workspace` consumes the reserved slot. - Planning and execution size through the **same** `gqa_workspace_layout` / `gqa_packed_staging_bytes` / `gqa_reference_scores_bytes` helpers; `gqa_carve` refuses a short prepared view **deterministically** rather than under-allocating. - Growth via `MappedGrowthGrant` (#748); capacity refusal surfaces as the pre-header **429** typed path (`TierExhausted` / `CapacityUnavailable`, #743) — never a late 500/OOM. - Central `is_planned_workspace_node` predicate already covers GQA — **no parallel predicate added** (doc comment updated only). - Compatibility/opt-out `execute()` path keeps its self-owned pooled scratch. - No asserts in any `Drop`. - CPU-only source-scan regression test (`gqa_packed_qkv_staging_is_governed_not_raw_allocated`) + layout/requirement unit tests guard the governed site on CI without a GPU. ## Updated audit table `docs/MEMORY_ARCHITECTURE.md` §2.1 and the site-by-site audit: QKV staging moved to governed-today rows with the byte formula + route/lifetime split; the `WS_SCORES`/`:1033` line references corrected to current code (`:1191`/`:1203`, scores carved at `:2662`); the next slice named from evidence — **GQA BNSH transpose scratch** (`WS_Q_BNSH` `:2127` always-reserved, `WS_OUT_BNSH` `:2131` for `q_seq>1`), the always-used sibling (expect a genuine-use prefill result with a possible decode `NONE`). Present K/V remains a KV-residency concern, not workspace. ## Validation (local; CI not waited on) - `cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --check` ✅ - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` ✅ (no `--all-targets`, exact CI gate) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` ✅ - Full `cargo test -p onnx-runtime-session` ✅ - GPU (single-threaded): `group_query_attention_gpu` (25), `gqa_seqmajor_parity_gpu`, `gqa_fp16_gpu`, `standard_attention_gpu` (24), `attention_gpu` (5), `qmoe_gpu` (29), `allocation_counts_gpu` (2), plus new lib unit tests ✅ Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>

justinchuby merged commit d74ca93 into main Aug 11, 2026
13 of 14 checks passed

justinchuby deleted the squad/736-scratchaudit branch August 11, 2026 16:53

This was referenced Aug 11, 2026

Seq-major (BSNH) KV layout: native fp16 decode pair #782

Merged

Infer memory strategy plan at load #752

Merged

Govern GroupQueryAttention WS_SCORES workspace (#736) #795

Merged

justinchuby mentioned this pull request Aug 12, 2026

Make managed no-spill VMM the default with automatic weight streaming (#755) #798

Merged

justinchuby mentioned this pull request Aug 12, 2026

Document repeated CUDA allocation audit finding #800

Merged

justinchuby mentioned this pull request Aug 12, 2026

Govern default-domain Attention f32 score matrix (#736) #802

Merged

justinchuby mentioned this pull request Aug 12, 2026

Govern GroupQueryAttention packed QKV-projection staging (#736) #806

Merged

justinchuby mentioned this pull request Aug 12, 2026

Provision free-budget headroom for the session-persistent workspace zone (fix #810 regression) #816

Closed

Navigation Menu

Repository navigation

Govern Attention Phase-2a scratch via prepare-only workspace (#736) - #753

Govern Attention Phase-2a scratch via prepare-only workspace (#736)#753
justinchuby merged 2 commits into
mainfrom
squad/736-scratchaudit

justinchuby commented Aug 10, 2026

Uh oh!

codecov Bot commented Aug 10, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 10, 2026 •

edited

Loading

Uh oh!

justinchuby commented Aug 11, 2026

Uh oh!

justinchuby commented Aug 11, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

justinchuby commented Aug 10, 2026

Summary

Step 1 — Audit

Step 2 — Govern one vertical slice

Step 3 — Document

Validation

Uh oh!

codecov Bot commented Aug 10, 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 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔴 Benchmark Regression Detected

Uh oh!

justinchuby commented Aug 11, 2026

Merge conflict resolution + re-validation (16 commits behind origin/main)

1. crates/onnx-runtime-session/src/executor/bindings.rs — vs #751 (d2db1f5a, Govern IndexShare CUDA workspace)

2. crates/onnx-runtime-ep-cuda/src/kernels/attention.rs — vs #757 (a3856b6a, close LinearAttention C2 capture seams)

Validation (all local, per the exact CI gates)

Uh oh!

justinchuby commented Aug 11, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Aug 10, 2026 •

edited

Loading

github-actions Bot commented Aug 10, 2026 •

edited

Loading

Merge conflict resolution + re-validation (16 commits behind `origin/main`)

1. `crates/onnx-runtime-session/src/executor/bindings.rs` — vs #751 (`d2db1f5a`, Govern IndexShare CUDA workspace)

2. `crates/onnx-runtime-ep-cuda/src/kernels/attention.rs` — vs #757 (`a3856b6a`, close LinearAttention C2 capture seams)