Plan governed QMoE workspace before admission by justinchuby · Pull Request #747 · justinchuby/onnx-genai

justinchuby · 2026-08-09T15:55:33Z

Implements workspace-planning Stage 0 for block-quantized MoE prefill.

adds a default-zero typed Kernel workspace requirement API
uses one checked QMoE layout helper for planning and execution
prepares/reserves governed workspace before cold and ordinary persistent-session admission callbacks
cleans scheduler state on preparation failure and maps device-limit refusals through the existing pre-header 429 path
removes QMoE persistent scratch calls to CudaRuntime::alloc_raw
preserves routed pipelines when QMoE input shapes do not depend on runtime-produced values
conservatively binds inputs_embeds as [1, request_len, static_hidden]; genuinely unresolved routed QMoE dependencies fail precisely before admission
recursively plans static QMoE in If/Loop/Scan children, takes the peak across reachable children, and passes the parent prepared workspace through nested executors for sequential reuse
grows scratch workspace without transiently double-charging old and replacement allocations

Validation:

targeted session/engine checks and library clippy with -D warnings passed
nested control-flow QMoE prepare/reuse regression passed
block_quantized_moe_gpu: 5 passed on CUDA
glm_tiny_qmoe_native_cuda_e2e: 3 passed serially on CUDA, including ordinary non-cold session admission ordering
native_cuda_routed_pipeline_decoder_parity: passed on CUDA with CUDA_VISIBLE_DEVICES=0 and ONNX_GENAI_CUDA_KV_MAX_LEN=16
server workspace overload classification test passed

The PR remains draft pending production-server ordinary streaming coverage and the Qwen14B 6 GiB long-prompt reproducer.

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

codecov · 2026-08-09T16:02:54Z

Codecov Report

❌ Patch coverage is 48.42932% with 197 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.57%. Comparing base (a9f4b2d) to head (e44d190).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...ates/onnx-runtime-session/src/executor/bindings.rs	35.77%	121 Missing and 19 partials ⚠️
...ates/onnx-runtime-session/src/executor/dispatch.rs	50.00%	36 Missing and 7 partials ⚠️
crates/onnx-runtime-session/src/lib.rs	0.00%	10 Missing ⚠️
crates/onnx-runtime-session/src/executor/mod.rs	50.00%	1 Missing and 2 partials ⚠️
.../onnx-runtime-session/src/executor/control_flow.rs	96.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #747      +/-   ##
==========================================
- Coverage   80.35%   79.57%   -0.79%     
==========================================
  Files         331      331              
  Lines      135311   135684     +373     
  Branches   135311   135684     +373     
==========================================
- Hits       108735   107971     -764     
- Misses      22310    23423    +1113     
- Partials     4266     4290      +24

Flag	Coverage Δ
cli-ort-linux	`83.41% <ø> (ø)`
cli-ort-windows	`82.90% <ø> (-0.10%)`	⬇️
mlas	`80.96% <ø> (+0.04%)`	⬆️
offline	`79.39% <48.42%> (-0.83%)`	⬇️

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

Files with missing lines	Coverage Δ
crates/onnx-runtime-ep-api/src/kernel.rs	`83.72% <100.00%> (+1.52%)`	⬆️
crates/onnx-runtime-ep-api/src/provider.rs	`70.39% <100.00%> (+0.46%)`	⬆️
crates/onnx-runtime-session/src/executor/run.rs	`65.28% <100.00%> (ø)`
crates/onnx-runtime-session/src/executor/state.rs	`83.33% <ø> (ø)`
.../onnx-runtime-session/src/executor/control_flow.rs	`65.50% <96.00%> (+0.56%)`	⬆️
crates/onnx-runtime-session/src/executor/mod.rs	`51.83% <50.00%> (-0.04%)`	⬇️
crates/onnx-runtime-session/src/lib.rs	`63.63% <0.00%> (-0.75%)`	⬇️
...ates/onnx-runtime-session/src/executor/dispatch.rs	`69.60% <50.00%> (-1.71%)`	⬇️
...ates/onnx-runtime-session/src/executor/bindings.rs	`43.34% <35.77%> (-4.61%)`	⬇️

... and 7 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 · 2026-08-09T16:21:37Z

🔴 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`	39.59 µs	94.02 µs	+137.5%
🔴	`matmul/medium_generic_f32_threads=8/32x512x512`	889.56 µs	1.75 ms	+96.9%
🔴	`block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024`	39.65 µs	70.01 µs	+76.6%
🔴	`matmul/large_generic_f32_threads=8/32x1024x1024`	3.59 ms	6.34 ms	+76.3%
🔴	`matmul/medium_generic_f16_threads=1/32x512x512`	28.42 µs	43.91 µs	+54.5%
🔴	`gather/large_bf16_threads=1-internal/131072`	10.63 µs	15.99 µs	+50.4%
🔴	`matmul/medium_generic_bf16_threads=8/32x512x512`	363.25 µs	531.87 µs	+46.4%
🔴	`matmul/medium_generic_f16_threads=8/32x512x512`	27.79 µs	40.30 µs	+45.0%
🔴	`block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024`	666.34 µs	953.55 µs	+43.1%
🔴	`matmul/large_generic_f16_threads=8/32x1024x1024`	78.30 µs	109.49 µs	+39.8%
🔴	`matmul/medium_generic_bf16_threads=1/32x512x512`	499.29 µs	679.26 µs	+36.0%
🔴	`matmul/small_generic_bf16_threads=1/1x256x256`	29.57 µs	39.93 µs	+35.0%
🔴	`matmul/small_generic_bf16_threads=8/1x256x256`	30.61 µs	40.73 µs	+33.0%
🔴	`gather/large_f16_threads=1-internal/131072`	10.90 µs	14.48 µs	+32.8%
🔴	`gather/large_f32_threads=1-internal/131072`	29.57 µs	38.82 µs	+31.3%
⚠️	`matmul/medium_generic_f32_threads=1/32x512x512`	2.19 ms	2.84 ms	+29.8%
⚠️	`block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1`	361.92 µs	468.75 µs	+29.5%
⚠️	`matmul/small_generic_f32_threads=8/1x256x256`	35.98 µs	46.13 µs	+28.2%
⚠️	`matmul/large_generic_f16_threads=1/32x1024x1024`	74.19 µs	94.39 µs	+27.2%
⚠️	`matmul/large_generic_bf16_threads=1/32x1024x1024`	1.93 ms	2.37 ms	+22.8%
⚠️	`reduce_mean/medium_f32_threads=1-internal/65536`	233.02 µs	284.93 µs	+22.3%
⚠️	`matmul/large_generic_f32_threads=1/32x1024x1024`	8.82 ms	10.75 ms	+21.9%
⚠️	`gather/small_bf16_threads=1-internal/4096`	436.8 ns	529.1 ns	+21.1%
⚠️	`matmul/small_generic_f16_threads=1/1x256x256`	28.04 µs	33.62 µs	+19.9%
⚠️	`reduce_mean/large_f32_threads=1-internal/262144`	941.98 µs	1.13 ms	+19.8%
⚠️	`matmul/small_generic_f32_threads=1/1x256x256`	34.00 µs	40.70 µs	+19.7%
⚠️	`gather/small_f32_threads=1-internal/4096`	603.8 ns	721.4 ns	+19.5%
⚠️	`matmul/small_generic_f16_threads=8/1x256x256`	29.61 µs	35.34 µs	+19.3%
⚠️	`add/large_f16_threads=1-internal/4194304`	39.16 ms	46.37 ms	+18.4%
⚠️	`matmul/large_generic_bf16_threads=8/32x1024x1024`	1.50 ms	1.78 ms	+18.4%
⚠️	`gather/medium_f16_threads=1-internal/32768`	2.19 µs	2.53 µs	+15.2%
✅	`reduce_mean/small_f32_threads=1-internal/4096`	14.50 µs	16.62 µs	+14.7%
✅	`gather/medium_f32_threads=1-internal/32768`	3.45 µs	3.91 µs	+13.4%
✅	`block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1`	72.44 µs	81.75 µs	+12.9%
✅	`gather/small_f16_threads=1-internal/4096`	459.7 ns	515.3 ns	+12.1%
✅	`add/medium_bf16_threads=1-internal/262144`	2.58 ms	2.87 ms	+11.3%
✅	`add/medium_f16_threads=1-internal/262144`	2.50 ms	2.74 ms	+9.9%
✅	`sampling_latency/min_p_per_token`	214.82 µs	233.57 µs	+8.7%
✅	`add/large_f32_threads=1-internal/4194304`	40.40 ms	43.17 ms	+6.9%
✅	`qwen3_sampling_processors/top_k_partial_selection`	139.42 µs	147.06 µs	+5.5%
✅	`qwen3_sampling_processors/top_k_top_p_full_sort_baseline`	5.86 ms	6.16 ms	+5.1%
✅	`add/small_bf16_threads=1-internal/1024`	12.41 µs	12.97 µs	+4.5%
✅	`qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline`	3.65 ms	3.80 ms	+4.1%
✅	`sampling_latency/top_p_per_token`	392.38 µs	406.78 µs	+3.7%
✅	`add/small_f16_threads=1-internal/1024`	12.79 µs	13.19 µs	+3.1%
✅	`qwen3_sampling_processors/top_k_full_sort_baseline`	2.15 ms	2.21 ms	+2.9%
✅	`qwen3_sampling_processors/top_k_top_p_fast`	668.85 µs	686.05 µs	+2.6%
✅	`grammar_masking/llguidance_compute_mask/32`	79.40 µs	81.18 µs	+2.2%
✅	`qwen3_sampling_processors/top_p_fast_after_top_k`	540.03 µs	548.02 µs	+1.5%
✅	`logit_processing/seven_processor_chain_per_step`	330.63 µs	334.75 µs	+1.2%
✅	`kv_cache/alloc_dealloc_pages`	40.43 µs	40.50 µs	+0.2%
✅	`add/medium_f32_threads=1-internal/262144`	2.58 ms	2.55 ms	-1.1%
✅	`sampling_latency/top_k_per_token`	55.82 µs	54.78 µs	-1.9%
✅	`gather/medium_bf16_threads=1-internal/32768`	2.70 µs	2.59 µs	-3.9%
✅	`sampling_latency/greedy_per_token`	3.54 µs	3.36 µs	-5.3%
✅	`add/small_f32_threads=1-internal/1024`	213.1 ns	201.7 ns	-5.4%
✅	`add/large_bf16_threads=1-internal/4194304`	50.26 ms	46.60 ms	-7.3%
✅	`tokenization/encode_tokens_per_second`	460.61 µs	403.37 µs	-12.4%
🟢	`tokenization/decode_tokens_per_second`	8.52 ms	6.68 ms	-21.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.82 3.53 5.11 }

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)

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

justinchuby · 2026-08-09T19:13:50Z

Second re-review blockers addressed in 919640f:\n\n- Pipeline decoder now prepares exact workspace from actual routed tensors before the admission callback; the server waits for that callback before headers.\n- Eager CUDA allocators retain a role-aware workspace lease from the shared device authority (VMM allocators continue charging committed mappings).\n- Fallback session dispatch synchronizes, releases disposable old scratch, and safely grows/re-aligns; failed replacement leaves a retryable empty slot.\n- Routed fixture now contains a live runtime BlockQuantizedMoE, with forced exhaustion proving TierExhausted before admission/token and successful recovery with ledger baseline restored.\n\nValidation: routed CUDA 3/3; GLM QMoE CUDA 3/3; block_quantized_moe_gpu 5/5; fallback growth regression passed; EP API 39/39; server streaming admission 5/5; changed-library clippy -D warnings passed. Full server lib run had 3 unrelated existing/environment failures (two registry authority assertions and a missing vision.onnx fixture).\n\nPR remains draft; production-server streaming and Qwen14B 6 GiB reproducer remain outstanding.

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

justinchuby · 2026-08-09T20:10:52Z

Final blocker fixed in 6e12226 without broadening engine features: onnx-runtime-session now re-exports WorkspaceRequirement, and native decode uses that already-required abstraction boundary.\n\nValidation:\n- exact gate: cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings passed\n- cargo check --locked --all-targets -p onnx-genai-engine --no-default-features --features native-backend passed\n- CUDA+native engine check passed\n- real qwen14b-zp server E2E at an exact 6GB authority limit: bounded 8K streaming prompt returned pre-header HTTP 429 JSON (no SSE/token), rejection metric incremented once, active sessions returned to 0, a subsequent small request succeeded, and VRAM stayed 5,167,382,528 <= 6,000,000,000 bytes.\n\nPR remains draft.

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

justinchuby · 2026-08-09T20:47:52Z

Speculative-width blocker fixed in c1f1491. The native speculation plan is now resolved before admission and its exact plan.width raises prepare-only QMoE query rows to max(prompt_rows, plan.width). Prompt-lookup and shared-KV widths, including per-request overrides, are covered; draft/MTP/EAGLE remain explicitly rejected on native before preparation.\n\nValidation:\n- exact native-backend clippy gate passed\n- runtime plan tests cover non-speculative, prompt-lookup width changes, prompt > width, and shared-KV's differing width semantics\n- native speculative CPU suite 3/3 passed\n- GLM QMoE CUDA suite 4/4 passed, including prompt length 1, configured width 4, an actual multi-row verify, parity with non-speculative output, and no prepared-workspace growth failure\n- CUDA+native library clippy passed.

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

justinchuby · 2026-08-09T21:18:54Z

Precision clamp fixed in e44d190. Planner and runtime now share
ative_speculative::verification_width; preparation runs after scheduler caps are applied and uses min(plan.width, scheduled.max_tokens, remaining_context), then maxes that with prompt rows.\n\nTests cover width 1024 capped to one output token, context remaining below width, normal width, non-speculative behavior, shared-KV width, and direct planner/runtime helper equality. Native speculative tests 3/3, GLM QMoE CUDA 4/4, and exact native-backend clippy all pass.

justinchuby · 2026-08-09T21:28:45Z

Final merge-safety review passed at e44d1904 after closing the full preparation/admission/lifecycle surface.

What lands

A typed kernel workspace planning contract and live governed QMoE implementation:

Kernel::workspace_requirement defaults to zero for existing kernels;
block-quantized MoE planning and execution share one checked sizing helper;
native ordinary, routed/pipeline and control-flow paths prepare workspace without launching kernels;
workspace is reserved through the same device authority before Return overload response for KV admission refusal #743 admission acknowledgment/SSE headers;
persistent workspace is retained/reused and released at session teardown; fallback sessions safely grow scratch without old+new double charging;
QMoE persistent scratch no longer uses raw CudaRuntime::alloc_raw;
resource refusal maps to pre-header 429, not late backend 500.

Speculative verification uses the exact shared runtime clamp:

min(plan.width, post_scheduler_max_new_tokens, remaining_context)

then takes max with prompt rows, preventing both under-allocation and false width-based rejection.

Review bugs closed

Normal native session bypassed preparation; only cold path used it.
Routed/inputs_embeds pipelines were unconditionally disabled.
QMoE in If/Loop/Scan child graphs was unplanned.
Persistent scratch growth double-charged old+new allocation.
Pipeline decoder still admitted before actual routed workspace preparation.
Fallback InferenceSession::run could not grow an undersized slot.
Native-backend-only feature build lacked EP API type visibility.
Speculative verify could require more rows than prompt preparation.
Raw plan width over-reserved beyond token/context caps.

Validation

QMoE CUDA suite: 4/4
native speculative suite: 3/3
routed parity and nested workspace regressions pass
exact native-backend and CUDA-native Clippy pass
qwen2.5-14b, 6 GB, 8K prompt: pre-header 429, no late 500; following smaller request succeeds; VRAM stays below limit
code review: no significant issues

What remains

This solves governed workspace planning/allocation. Dynamic weight↔KV/workspace lending still needs the transactional MappedGrowthGrant vertical slice; draft #746 remains hold until rebuilt on this contract. #736 remains open.

Merging; latest CI is pending, not waiting on matrix/coverage.

## Summary - authority-owned weak `ReclaimableMappedHolder` registration and transactional RAII `MappedGrowthGrant` - unused-allowance-first transfer, external reclaim callbacks, grant-bound capacity tokens, rollback, FIFO serialization, and atomic reserved-to-mapped attribution - batched native CUDA KV/workspace growth with exact global granule 0→1 charging and 1→0 refund - one canonical mapped allowance per suballocating VMM arena; content lifetime remains separate from physical zone ownership - explicit byte `serve --vram-limit` selects managed no-spill VMM/pool mode; `ONNX_GENAI_DYNAMIC_KV_WEIGHT_LENDING=0` restores compatibility fallback - typed capacity refusals map to pre-header HTTP 429 with `Retry-After` Refs #736 Refs #717 Refs #745 Refs #747 ## Review fixes - managed VMM initialization failure is fatal before model allocation; compatibility mode alone may fall back - managed physical-pool ownership avoids double-charging resident weights when offload is unnecessary - allocator transactions charge only exact global granule transitions, including packed/shared workspace allocations - CUDA provider centrally refunds the canonical arena allowance on full deallocation and partial range decommit, including when an ordinary buffer drops the final shared reference - specialized workspace/session cleanup cannot omit or double-refund attribution - state-only native models load without KV bytes/token when paged geometry is not required ## Validation - `cargo fmt --all -- --check` - memory governor: 36 passed - session: 135 passed - CUDA EP library: 295 passed, 4 ignored - engine CUDA/native library: 434 passed, 1 ignored - hardware VMM GPU: 16 passed on RTX 4060 Laptop GPU - provider canonical full/partial decommit GPU regression: passed - changed-library clippy with `-D warnings`: passed - changed-package all-target clippy with `-D warnings`: passed ## 14B / 6 GiB E2E Model: `C:\Users\justinchu\dev\models\qwen14b-zp`; release native-CUDA binary. - managed mode logged VMM/pool strategy and registered the weight holder - short `max_tokens=24`: HTTP 200, correct `Paris`; weight allowance 4,831,838,208 bytes; owned 5,534,384,128 <= 6 GiB - first physical KV growth transferred/attributed 201,326,592 bytes - 8K refused pre-header with HTTP 429, never late 500/OOM - four queued short requests completed FIFO; server did not demonstrate simultaneous live KV accumulation - opt-out did not auto-enable VMM/pool and retained compatibility behavior Latest commit: `440a1df4`. PR remains open and unmerged. --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>

The prepare-only workspace planning path added by #747 calls get_or_create without the growing-symbol capture classification introduced here. Compute it with the same helper used by build and dispatch so prepare-only planning keys the same kernel that execution later uses. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624

…753) ## 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 (#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 (#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. 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>

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>

Plan governed QMoE workspace before admission

78bd334

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

Address QMoE workspace review blockers

7d4ec38

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

justinchuby marked this pull request as draft August 9, 2026 16:45

Copilot AI added 2 commits August 9, 2026 10:28

Plan routed and nested QMoE workspace

39bfa3c

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

Fix routed QMoE workspace admission

919640f

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

Fix native backend workspace type layering

6e12226

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

Size QMoE workspace for speculative verify

c1f1491

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

Clamp speculative QMoE workspace rows

e44d190

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

justinchuby marked this pull request as ready for review August 9, 2026 21:28

justinchuby merged commit f861905 into main Aug 9, 2026
16 checks passed

justinchuby deleted the squad/736-governed-qmoe-workspace branch August 9, 2026 21:28

This was referenced Aug 9, 2026

Dynamic lending between CUDA weights and KV growth #746

Closed

Add transactional mapped growth grants #748

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

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

Plan governed QMoE workspace before admission - #747

Plan governed QMoE workspace before admission#747
justinchuby merged 7 commits into
mainfrom
squad/736-governed-qmoe-workspace

justinchuby commented Aug 9, 2026 •

edited

Loading

Uh oh!

codecov Bot commented Aug 9, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 9, 2026 •

edited

Loading

Uh oh!

justinchuby commented Aug 9, 2026

Uh oh!

justinchuby commented Aug 9, 2026

Uh oh!

justinchuby commented Aug 9, 2026

Uh oh!

justinchuby commented Aug 9, 2026

Uh oh!

justinchuby commented Aug 9, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

justinchuby commented Aug 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Aug 9, 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 9, 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 9, 2026

Uh oh!

justinchuby commented Aug 9, 2026

Uh oh!

justinchuby commented Aug 9, 2026

Uh oh!

justinchuby commented Aug 9, 2026

Uh oh!

justinchuby commented Aug 9, 2026

What lands

Review bugs closed

Validation

What remains

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

justinchuby commented Aug 9, 2026 •

edited

Loading

codecov Bot commented Aug 9, 2026 •

edited

Loading

github-actions Bot commented Aug 9, 2026 •

edited

Loading