[justinchuby]

Summary

• Stage 1 for Make CUDA weight offload graph-capture-compatible with stable slots #716: add an opt-in VMM-backed CUDA weight residency path (ONNX_GENAI_WEIGHT_OFFLOAD_VMM=1).
• Reuses the existing onnx-runtime-cuda-memory VMM arena for reserve/create/map/set-access/granularity instead of adding new CUDA driver wrappers.
• Keeps one stable virtual address per weight key and unmaps/remaps physical granules on eviction/page-in; CUDA graph capture remains declined under offload until stage 2 adds the explicit captured-segment residency guarantee.

Why VMM rather than application slots

The original stable-slot design would swap contents through fixed application slots and needs deterministic scheduling plus RAW/WAR ordering. This stage uses CUDA VMM to get the same captured-pointer stability from hardware: kernels see a stable virtual address while residency maps physical pages behind it outside captured regions. That is smaller and composes with #727's verified remap behavior.

WDDM behavior

This is opt-in because CUDA VMM physical memory (cuMemCreate) cannot oversubscribe into WDDM shared memory the way cudaMalloc can. Enabling ONNX_GENAI_WEIGHT_OFFLOAD_VMM=1 can therefore hard-OOM on Windows where the current non-VMM weight path survives by spilling.

Counters / evidence

GPU unit test proves:

• A's device pointer is identical before and after eviction/remap.
• Readback bytes are correct after remap.
• runtime cuMemAlloc/cuMemFree counts are unchanged across page-ins.

14B end-to-end, same rebuilt binary, ONNX_GENAI_WEIGHT_OFFLOAD=1, ONNX_GENAI_WEIGHT_OFFLOAD_DEVICE_BYTES=6000000000, max-new-tokens=8:

arm	vram_alloc	vram_free	page_ins	hits	hit rate	evictions
VMM off	646.59 ms	4335.79 ms	1,791	5,145	74.18%	0
VMM on	0.00 ms	0.00 ms	4,416	2,520	36.33%	2,833

VMM granule floor on 14B: peak committed 6,050,283,520 B; peak content resident 5,132,642,304 B; granule overhead 917,641,216 B (~17.9% of content). The allocator churn result is good, but the residency counters are not neutral: accounting in 2 MiB granules forces evictions/transient bypasses under the 6 GB budget. This is the important follow-up result before merge.

This PR still does not relax resolve_graph_capture_enabled.

Validation

• cargo fmt --all -- --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 clippy -p onnx-genai-cli --features native-cuda -- -D warnings
• GPU test with wheel CUDA/ORT paths set: cargo test -p onnx-runtime-ep-cuda --features cuda stable_vmm_weight_slots_keep_their_virtual_address_across_eviction -- --nocapture (1 passed)

Refs #716

[codecov]

Codecov Report

❌ Patch coverage is 50.00000% with 29 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.13%. Comparing base (ec2c14b) to head (34ecb06).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...ates/onnx-runtime-cuda-memory/src/vmm_allocator.rs	4.34%	22 Missing ⚠️
crates/onnx-genai-cli/src/generate.rs	80.00%	5 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #733      +/-   ##
==========================================
- Coverage   80.85%   80.13%   -0.73%     
==========================================
  Files         331      331              
  Lines      132743   132800      +57     
  Branches   132743   132800      +57     
==========================================
- Hits       107330   106414     -916     
- Misses      21168    22143     +975     
+ Partials     4245     4243       -2     

Flag	Coverage Δ
cli-ort-linux	83.35% <80.00%> (-0.03%)	⬇️
cli-ort-windows	82.84% <80.00%> (-0.12%)	⬇️
mlas	80.92% <ø> (-0.05%)	⬇️
offline	79.98% <4.34%> (-0.77%)	⬇️

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

Files with missing lines	Coverage Δ
crates/onnx-genai-cli/src/generate.rs	90.32% <80.00%> (-0.87%)	⬇️
...ates/onnx-runtime-cuda-memory/src/vmm_allocator.rs	18.45% <4.34%> (-0.54%)	⬇️

... 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]

✅ Benchmarks — No Regression

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
✅	reduce_mean/small_f32_threads=1-internal/4096	15.41 µs	17.36 µs	+12.6%
✅	add/medium_f32_threads=1-internal/262144	2.64 ms	2.88 ms	+8.8%
✅	reduce_mean/medium_f32_threads=1-internal/65536	256.15 µs	273.85 µs	+6.9%
✅	add/medium_f16_threads=1-internal/262144	2.63 ms	2.79 ms	+6.2%
✅	gather/small_f16_threads=1-internal/4096	484.2 ns	504.8 ns	+4.3%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.02 ms	1.06 ms	+4.0%
✅	add/large_bf16_threads=1-internal/4194304	44.09 ms	45.60 ms	+3.4%
✅	add/small_f16_threads=1-internal/1024	13.02 µs	13.37 µs	+2.7%
✅	add/large_f16_threads=1-internal/4194304	41.60 ms	42.68 ms	+2.6%
✅	add/small_bf16_threads=1-internal/1024	13.00 µs	13.33 µs	+2.5%
✅	gather/medium_f16_threads=1-internal/32768	2.34 µs	2.39 µs	+2.1%
✅	gather/small_f32_threads=1-internal/4096	689.0 ns	689.1 ns	+0.0%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	541.92 µs	539.74 µs	-0.4%
✅	add/large_f32_threads=1-internal/4194304	41.87 ms	41.70 ms	-0.4%
✅	gather/medium_f32_threads=1-internal/32768	4.26 µs	4.21 µs	-1.3%
✅	add/medium_bf16_threads=1-internal/262144	2.80 ms	2.77 ms	-1.3%
✅	qwen3_sampling_processors/top_k_top_p_fast	684.16 µs	674.73 µs	-1.4%
✅	sampling_latency/greedy_per_token	3.39 µs	3.32 µs	-2.0%
✅	sampling_latency/min_p_per_token	211.75 µs	206.70 µs	-2.4%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	79.21 µs	77.01 µs	-2.8%
✅	sampling_latency/top_p_per_token	390.24 µs	377.52 µs	-3.3%
✅	gather/medium_bf16_threads=1-internal/32768	2.45 µs	2.35 µs	-4.0%
✅	grammar_masking/llguidance_compute_mask/32	78.89 µs	75.65 µs	-4.1%
✅	tokenization/encode_tokens_per_second	409.57 µs	392.25 µs	-4.2%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.21 ms	2.11 ms	-4.4%
✅	tokenization/decode_tokens_per_second	6.62 ms	6.31 ms	-4.7%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.01 ms	5.70 ms	-5.1%
✅	gather/small_bf16_threads=1-internal/4096	511.2 ns	484.3 ns	-5.3%
✅	kv_cache/alloc_dealloc_pages	41.01 µs	38.81 µs	-5.3%
✅	sampling_latency/top_k_per_token	54.99 µs	51.91 µs	-5.6%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.72 ms	3.51 ms	-5.7%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	588.88 µs	547.66 µs	-7.0%
✅	logit_processing/seven_processor_chain_per_step	335.87 µs	311.42 µs	-7.3%
✅	qwen3_sampling_processors/top_k_partial_selection	148.55 µs	136.85 µs	-7.9%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.24 ms	2.06 ms	-8.2%
✅	matmul/small_generic_f32_threads=1/1x256x256	42.93 µs	38.85 µs	-9.5%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	10.92 ms	9.66 ms	-11.5%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	93.58 µs	82.56 µs	-11.8%
✅	gather/large_f16_threads=1-internal/131072	18.58 µs	16.24 µs	-12.6%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.69 ms	2.35 ms	-12.6%
✅	matmul/small_generic_bf16_threads=1/1x256x256	37.61 µs	32.04 µs	-14.8%
🟢	gather/large_bf16_threads=1-internal/131072	19.83 µs	16.77 µs	-15.4%
🟢	add/small_f32_threads=1-internal/1024	247.4 ns	207.4 ns	-16.2%
🟢	matmul/medium_generic_f16_threads=1/32x512x512	37.67 µs	31.57 µs	-16.2%
🟢	matmul/large_generic_f16_threads=8/32x1024x1024	100.82 µs	83.64 µs	-17.0%
🟢	matmul/small_generic_f16_threads=1/1x256x256	39.10 µs	31.78 µs	-18.7%
🟢	gather/large_f32_threads=1-internal/131072	47.84 µs	38.43 µs	-19.7%
🟢	matmul/small_generic_bf16_threads=8/1x256x256	40.03 µs	31.81 µs	-20.5%
🟢	matmul/small_generic_f16_threads=8/1x256x256	39.90 µs	31.20 µs	-21.8%
🟢	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	538.42 µs	391.16 µs	-27.4%
🟢	matmul/medium_generic_f16_threads=8/32x512x512	44.18 µs	30.83 µs	-30.2%
🟢	matmul/small_generic_f32_threads=8/1x256x256	53.06 µs	35.21 µs	-33.6%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.80 ms	1.08 ms	-40.1%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	2.37 ms	1.40 ms	-41.1%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	696.77 µs	400.83 µs	-42.5%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	7.12 ms	4.07 ms	-42.8%
🟢	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	1.20 ms	569.75 µs	-52.5%
🟢	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	99.91 µs	47.09 µs	-52.9%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	128.79 µs	52.10 µs	-59.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.47 3.33 4.17 }

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]

Follow-up on Option A vs B:

I tried the contained shared-reservation version locally: keep the existing VMM arena as the one reservation, allocate each stable weight slot at packed offsets instead of rounding each slot to 2 MiB, and fix range decommit so a decommit releases any overlapping granule only when all touching slots are gone.

That removes the per-weight padding shape, but the real decode loop still fails under the one-ledger constraint before producing counters. The reason is the #723 policy's zero-eviction behavior: once the stable resident subset fills the 6 GB residency budget, misses outside that subset are transient bypass pages. The non-VMM path can cudaMalloc those transient pages outside the cache budget and then free them; a VMM-backed path must lease/commit the transient granule from the same device ledger. At this model's 6 GB weights + 1.6 GiB KV limit there is only ~0.7 MB headroom, less than one 2 MiB granule, so the first transient bypass commit is correctly refused.

Evicting resident weights to make room makes the run complete, but it is exactly the regression we measured: hit rate drops and evictions become non-zero. Padding is not the only problem; transient bypass headroom also needs a unified reservation/budget owner that can coordinate weights with KV/other zones. So my judgement is Option B: keep #733 as evidence/do-not-merge and let #736 subsume it.

[justinchuby]

Correction to my previous Option-A framing: the VMM arena already has the shared reservation, byte-granularity Spans::carve, and per-granule refcounts. The first #733 implementation bypassed that by rounding every stable weight slot to a separate granule-sized allocation.

I tried the contained fix locally (packed stable weight slots inside the existing arena). With extra headroom (--vram-limit 14GiB, same 6 GB weight budget) it does exactly what we want for residency: page_ins=1,791, hits=5,145, hit rate 74.18%, evictions=0, and vram_alloc/free=0 ms.

But under the original tight limit it still fails before counters: VMM needs to lease a 2 MiB committed granule when only ~742 KB is left in the device ledger. So the remaining blocker is not padding; it is transient/overhead headroom under the one-ledger constraint. The non-VMM path can temporarily cudaMalloc outside the residency budget (and WDDM may spill), while VMM cannot.

So I am keeping #733 draft/do-not-merge and treating this as Option B for the original acceptance bar. The useful conclusion for #736 is sharper now: reuse the existing shared suballocator, but make the unified budget owner reserve/coordinate transient headroom across zones instead of letting weights/KV independently fill their local budgets.

[justinchuby]

Committed-byte admission probe result:

I added the contained arena preflight query (incremental_commit_bytes) and wired stable VMM admission to use committed granules rather than tensor length. It avoids the tight-limit OOM, but the original acceptance counters do not pass.

Packed/shared with extra headroom proves the design is neutral at the policy level: 1,791 / 5,145 / 74.18% / 0, vram_alloc/free=0 ms. But that same resident set is ~6.00 GB content and ~6.51 GB committed. Under the strict 6 GB committed budget, honest admission has to evict/bypass: best tight run I measured was page_ins=4,985, hits=1,951, hit rate 28.13%, evictions=3,217, with peak committed staying under 6 GB.

Conclusion: the preflight query is contained, but the needed budget/headroom policy is not. #733 stays draft/do-not-merge; #736 should own the unified committed-byte budget decision across weights/KV/scratch.

[justinchuby]

Closing this draft as superseded by #736, preserving the branch and measurements.

What it proved:

• per-weight private reservations eliminate cuMemAlloc/cuMemFree but add 17.9% granule overhead, cutting hit rate 74.18% → 36.33% and adding 2,833 evictions;
• using the arena's shared suballocator removes that floor and, at equal physical committed bytes, matches baseline exactly: 1,791 page-ins / 5,145 hits / 74.18% / 0 evictions, with alloc/free at 0 ms;
• the baseline's nominal 6 GB content budget physically consumes ~6.51 GB, exposing the content-vs-committed second ledger;
• under a strict 6 GB committed budget, honest admission necessarily lowers residency, making the allocator win roughly cancel the extra staging/H2D cost.

The conclusion is not that stable VA is wrong. It is that it becomes a speed win only with the shared physical authority/pool and dynamic zone lending in #736. Landing a separate opt-in path now would duplicate ownership and preserve the wrong budget unit.

The stable-pointer mechanism is retained in this branch and #727 proves graph replay survives remap. #716 remains open for capture+offload after #736's pool/committed-byte admission lands.