[justinchuby]

Refs #736

Summary

• fold default-domain Attention WS_STAGE_KEY / WS_STAGE_VALUE into the existing prepared score-workspace composite
• size planning and execution through std_attention_workspace_layout, carve staged K/V from execute_with_workspace, and fail deterministically on shortfall
• keep direct execute as the self-owned compatibility/opt-out path
• update the centralized is_planned_workspace_node contract and the authoritative memory audit

Routes and bytes

Dense aliased KV growth genuinely needs disjoint rebuild buffers before copy-back. Each staged region is derived as B * Hkv * present_seq * head_dim * element_bytes; K and V are simultaneously live. At B=1, Hkv=8, present_seq=2048, head_dim=128, total K+V staging is derived (not measured) as 16 MiB f32 or 8 MiB fp16/bf16, plus region alignment.

Fixed-capacity mask-driven append (mask width equals physical past capacity), no-past, and missing present-output routes add zero staged-K/V bytes. The overall requirement is not NONE because default-domain Attention still always materializes its independently governed fp32 score matrix. Output arity is threaded into planning so absent present outputs are not charged.

Lifetimes

• SessionPersistent: dense B=1, Sq=1 growth is capture-eligible and previously retained the stage slots in StdAttnWorkspace.
• StepScoped: multi-token or batched dense growth previously allocated/freed owned stage buffers per call.

Prepared-workspace ownership preserves success/error/cancellation release and routes authority refusal through the existing pre-header 429 TierExhausted / CapacityUnavailable path.

Over-reservation result

The prior allocator was already conditional at execution; staging is not globally unnecessary. The defect risk was charging it unconditionally after governance. This PR preserves the route split, so fixed-capacity and absent-output routes add zero staging rather than introducing a new reservation regression.

Audit

docs/MEMORY_ARCHITECTURE.md moves staged K/V into governed-today with verified line references and derived formulas. The next evidenced slice is the route-dependent MatMulNBits f32 dequantized-weight fallback (K * N * sizeof(f32), derived 64 MiB at 4096x4096 and 172 MiB at 4096x11008).

Validation

• cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --check
• cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings
• cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings
• cargo test -p onnx-runtime-session (153 unit tests plus integration suites)
• CUDA, single-threaded with cuda,gpu-tests: standard_attention_gpu (24/24), attention_gpu (5 passed, 1 H200 benchmark ignored), group_query_attention_gpu (25 passed, 1 H200 benchmark ignored), allocation_counts_gpu (2/2)
• governed dense-alias GPU regression and CPU-only source-scan/route/lifetime tests

[justinchuby]

Coordinator: validated and approved in substance, but holding the merge until the main regression is fixed. Not a comment on this PR''s quality.

What I verified

• merged current origin/main into the branch: no conflicts
• cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --check clean
• full cargo test -p onnx-runtime-session: 153 passed, 0 failed

I also ran the end-to-end reproducer against this branch, because this PR adds session-persistent workspace reservations — the same class that is currently under-provisioned on main. The result is the important part:

cannot reserve 2097152 bytes for Workspace { step_scoped: false }:
0 of 2097152 bytes are already leased, leaving 0

Byte-for-byte identical to the failure on pure origin/main, and still the 2 MiB GQA reservation — so this PR inherits the regression rather than worsening it. That was the specific risk I wanted to rule out before deciding, and it is ruled out.

Why I am still holding

main currently fails mobius_seqmajor_growth_parity_native_cuda because #806 and #810 together added session-persistent workspace reservations that the governor''s budget does not provision headroom for (0 of 2097152 already leased, leaving 0, on a completely idle GPU). A fix is in flight.

Stacking a third governed session-persistent reservation onto a budget already proven insufficient is not something I want to do blind, even though this branch does not make the current failure worse. Once the headroom fix lands I will re-run the reproducer on this branch and merge if it stays clean.

On the substance

The route split is the right shape and matches what the audit table predicted for this slice: dense alias routes use governed staging with dual lifetimes, while fixed-capacity, no-past and absent-output routes charge zero. Modelling both lifetime classes rather than forcing one was the specific difficulty flagged in the table, and #802 had to do the same for its slice.

This is the seventh slice, and the tally now stands at six of seven finding over-reservation — with the sharper observation that most were route-conditional rather than wholly unnecessary (#795 f32-only, #806 packed-input-only, #810 prefill-only, and now this one). The recurring defect is not "this buffer is useless" but "this buffer is charged on routes that do not use it", and the fix is consistently a static signal threaded through the shared sizing helper.

Next slice recorded as MatMulNBits dequantized weights.

One process note this regression exposes

This class of change keeps passing the mandatory pre-merge set and then breaking an end-to-end path, because mobius_seqmajor_growth_parity_native_cuda is #[ignore]d and needs --features cuda,native-backend --ignored. A green cargo test -p onnx-runtime-session says nothing about it. I have asked the headroom fix to recommend whether that test belongs in the standard pre-merge set for anything touching workspace governance — on the evidence of seven slices, I think it does.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.55%. Comparing base (1ea53e4) to head (9e95cb4).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #813   +/-   ##
=======================================
  Coverage   83.55%   83.55%           
=======================================
  Files          12       12           
  Lines        5192     5192           
  Branches     5192     5192           
=======================================
  Hits         4338     4338           
  Misses        677      677           
  Partials      177      177           

Flag	Coverage Δ
cli-ort-linux	83.50% <ø> (ø)
cli-ort-windows	82.99% <ø> (ø)

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 Change 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/medium_bf16_threads=1-internal/32768	2.90 µs	3.74 µs	+28.8%
⚠️	gather/large_f16_threads=1-internal/131072	17.57 µs	21.06 µs	+19.9%
⚠️	gather/large_f32_threads=1-internal/131072	53.78 µs	62.78 µs	+16.7%
✅	matmul/small_generic_f16_threads=1/1x256x256	37.37 µs	42.65 µs	+14.1%
✅	grammar_masking/llguidance_compute_mask/32	82.68 µs	93.80 µs	+13.5%
✅	logit_processing/seven_processor_chain_per_step	329.40 µs	370.43 µs	+12.5%
✅	matmul/small_generic_bf16_threads=1/1x256x256	39.27 µs	43.64 µs	+11.1%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.24 ms	6.72 ms	+7.5%
✅	gather/small_f32_threads=1-internal/4096	736.2 ns	774.9 ns	+5.3%
✅	reduce_mean/medium_f32_threads=1-internal/65536	288.24 µs	297.77 µs	+3.3%
✅	matmul/small_generic_f16_threads=8/1x256x256	43.69 µs	44.97 µs	+2.9%
✅	qwen3_sampling_processors/top_k_top_p_fast	734.76 µs	751.39 µs	+2.3%
✅	kv_cache/alloc_dealloc_pages	42.81 µs	43.75 µs	+2.2%
✅	tokenization/encode_tokens_per_second	459.54 µs	465.03 µs	+1.2%
✅	sampling_latency/greedy_per_token	3.64 µs	3.64 µs	+0.1%
✅	gather/medium_f16_threads=1-internal/32768	2.89 µs	2.88 µs	-0.2%
✅	matmul/small_generic_f32_threads=1/1x256x256	51.10 µs	50.35 µs	-1.5%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.66 ms	2.62 ms	-1.5%
✅	tokenization/decode_tokens_per_second	7.39 ms	7.14 ms	-3.4%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.30 ms	1.25 ms	-3.5%
✅	matmul/small_generic_bf16_threads=8/1x256x256	41.26 µs	39.71 µs	-3.8%
✅	gather/large_bf16_threads=1-internal/131072	21.32 µs	20.42 µs	-4.2%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.34 ms	2.23 ms	-4.7%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	580.75 µs	546.36 µs	-5.9%
✅	sampling_latency/min_p_per_token	240.05 µs	225.30 µs	-6.1%
✅	sampling_latency/top_k_per_token	65.75 µs	61.65 µs	-6.2%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	620.40 µs	579.14 µs	-6.7%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.96 ms	3.69 ms	-6.8%
✅	qwen3_sampling_processors/top_k_partial_selection	160.64 µs	147.75 µs	-8.0%
✅	gather/small_bf16_threads=1-internal/4096	561.3 ns	507.1 ns	-9.7%
✅	add/small_bf16_threads=1-internal/1024	17.36 µs	15.25 µs	-12.2%
✅	add/small_f16_threads=1-internal/1024	15.03 µs	13.18 µs	-12.3%
✅	matmul/medium_generic_f16_threads=1/32x512x512	39.13 µs	34.14 µs	-12.8%
✅	reduce_mean/small_f32_threads=1-internal/4096	18.42 µs	16.01 µs	-13.1%
✅	matmul/medium_generic_f16_threads=8/32x512x512	40.87 µs	35.43 µs	-13.3%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	121.26 µs	104.93 µs	-13.5%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	91.86 µs	79.37 µs	-13.6%
✅	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	479.22 µs	413.97 µs	-13.6%
✅	sampling_latency/top_p_per_token	468.87 µs	404.39 µs	-13.8%
✅	add/medium_f32_threads=1-internal/262144	3.78 ms	3.23 ms	-14.8%
🟢	gather/medium_f32_threads=1-internal/32768	5.57 µs	4.72 µs	-15.2%
🟢	matmul/large_generic_f32_threads=1/32x1024x1024	11.54 ms	9.74 ms	-15.6%
🟢	gather/small_f16_threads=1-internal/4096	633.4 ns	527.6 ns	-16.7%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.31 ms	1.07 ms	-18.1%
🟢	add/medium_f16_threads=1-internal/262144	4.30 ms	3.46 ms	-19.5%
🟢	matmul/large_generic_f16_threads=1/32x1024x1024	109.04 µs	87.78 µs	-19.5%
🟢	add/small_f32_threads=1-internal/1024	260.8 ns	209.7 ns	-19.6%
🟢	add/large_bf16_threads=1-internal/4194304	60.73 ms	48.78 ms	-19.7%
🟢	matmul/large_generic_bf16_threads=1/32x1024x1024	2.66 ms	2.09 ms	-21.6%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	584.69 µs	456.31 µs	-22.0%
🟢	add/medium_bf16_threads=1-internal/262144	4.61 ms	3.43 ms	-25.6%
🟢	add/large_f16_threads=1-internal/4194304	60.27 ms	44.72 ms	-25.8%
🟢	add/large_f32_threads=1-internal/4194304	60.39 ms	44.45 ms	-26.4%
🟢	matmul/small_generic_f32_threads=8/1x256x256	83.12 µs	56.33 µs	-32.2%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	7.12 ms	4.81 ms	-32.5%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	95.74 µs	63.16 µs	-34.0%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	2.31 ms	1.49 ms	-35.4%
🟢	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	110.29 µs	62.49 µs	-43.3%
🟢	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	1.25 ms	697.12 µs	-44.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: { 3.41 3.36 6.32 }

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]

Unblocked and merging. #814 landed the governor headroom fix, and I re-ran the end-to-end reproducer on this branch merged with current \main:

\
test mobius_seq_major_growth_is_bit_identical_to_head_major ... ok
test result: ok. 1 passed; 0 failed
\\

That is the check I said I would run before merging — this PR adds a third governed session-persistent reservation, and I did not want to stack it onto a budget that was proven insufficient until the arithmetic was fixed. It is now clean, so the hold is released.

Worth recording what the fix investigation established about this class: #810 was the tipping byte on the fp16-unpacked-flash-prefill route rather than the sole cause, the governor fix covers the whole \SessionPersistent Workspace\ composite class (so it generalises to this slice), and no other end-to-end path is over-subscribed. That last point is why I am comfortable merging a further reservation now rather than waiting.

Applying the validation-gap policy from #814 to this merge: for anything touching workspace governance, \mobius_seqmajor_growth_parity_native_cuda\ runs before merge, because a green \cargo test -p onnx-runtime-session\ demonstrably says nothing about it — that is exactly how the regression got in.