[justinchuby]

Refs #837 (item 2 — host→device streaming efficiency). Do not merge yet; do not close #837.

What this changes

CudaWeightResidency::resident_mapped (the hot page-in path) page-locked a fresh
cuMemHostAlloc staging buffer and freed it with cuMemFreeHost on every
page-in — 5,535 alloc/free pairs per 16-token window on qwen14b-zp. Those calls
take the CUDA context lock to pin/unpin OS pages and serialize against the
device, defeating the async, fence-ordered page-in the design already
implements correctly.

This adds a bounded reuse pool (crate::pinned_pool::PinnedStagingPool) and
draws the staging buffer from it instead of allocating one per miss.

Functions/areas I touched (kept deliberately narrow so this merges cleanly
against vmm-churn's work in the same file):

• new file crates/onnx-runtime-ep-cuda/src/pinned_pool.rs
• CudaWeightResidency struct + its two constructors (new, new_leased) — one field
• CudaWeightResidency::resident_mapped — acquire-from-pool instead of alloc_pinned
• counter plumbing in GlobalOffloadStats + profile_native reporting

Not touched: admit_committed_span, admit, drain_copy_stream,
resident_vmm_with's body (beyond passing a pooled buffer into its fill
closure), and all VMM commit/release/mapping code (vmm-churn's domain).
ONNX_GENAI_WEIGHT_OFFLOAD_ASYNC_PAGEIN semantics are unchanged.

Pooling design and bounded memory cost

• Structure: a Mutex<Vec<PinnedStaging>> free-list, keyed by capacity with
  a reuse-if-large-enough rule (smallest fit is handed out; a request larger
  than every retained buffer allocates a fresh one).
• Bound: at most 8 retained buffers and ≤ 512 MiB retained pinned
  host bytes; any buffer that would breach either bound is dropped (freed) on
  release rather than retained.
• Actual cost: page-ins on this path are host-synchronous and serialized
  under the residency lock, so at most one staging buffer is live at a time and
  the pool converges to a single retained buffer sized to the largest page-in
  (tens of MiB for qwen14b). The 8-buffer / 512-MiB ceiling only caps a
  pathological worst case; it never trades unbounded host RAM for speed.

Why return-to-pool cannot happen before the copy fence (safety)

The staging buffer is the source of the H2D copy, so returning it to the
pool before the copy's DMA read completes would silently corrupt weights. This
is safe because every copy on the offload page-in path is host-blocking:

• Both branches of resident_mapped issue the copy through
  CudaRuntime::htod_async_elapsed_ms, which records a CUDA end event after the
  cuMemcpyHtoDAsync and calls event.elapsed_ms.
• cudarc's Event::elapsed_ms calls end.synchronize() (verified in
  cudarc-0.19.8/src/driver/safe/core.rs), which blocks the host thread until
  the end event — and therefore the copy — has completed.
• The buffer is returned to the pool only after that call returns: on the VMM
  branch the PooledStaging guard is moved into the fill closure and dropped
  when the closure ends (strictly after the blocking copy); on the non-VMM
  branch the buffer is released after upload_staged_async (also
  host-blocking) returns.

So on the host timeline the DMA read has provably finished before the buffer is
reusable. PooledStaging::Drop performs only the return-to-pool and never
asserts (an assert in Drop triggers STATUS_STACK_BUFFER_OVERRUN here).

⚠️ Guard rail: if this path is ever switched to non-blocking htod_async +
deferred fence, the return-to-pool must be re-ordered after that fence is awaited
(e.g. via drain_copy_stream) or the safety argument no longer holds. This is
documented at the top of pinned_pool.rs.

Regression guard

pinned_pool::tests::reused_buffers_do_not_reallocate_per_page_in drives 64
same-size acquire→fill→release cycles and asserts exactly one real pinned
allocation with 63 reuses — i.e. a per-page-in pinned allocation cannot silently
return. A deterministic pinned_alloc_calls counter (per-instance + global) is
also surfaced in profile_native so the property is visible at benchmark scale
(pinned_alloc_calls must stay far below page_ins).

Measurements (qwen14b-zp, 16 tokens, release, CUDA_GRAPH=1, VMM=1)

This box's wall-clock is extremely noisy (an identical after-binary produced
both 0.02 tok/s under external contention and 0.55–0.59 tok/s clean). So the
result is led by the deterministic, contention-immune counter, and bandwidth is
reported from clean, uncontended matched pairs run back-to-back on an
otherwise-idle GPU (verified 0 other compute PIDs).

Counter	Before (main)	After (pooled)
pinned_alloc_calls (steady window)	5,535¹	0
pinned_reuses	0¹	5,535
page_ins	5,535	5,535
hit_rate	60.10%	60.10%
evictions	4,799	4,799
bypassed_page_ins	733	733
htod_bytes	65,772,419,072	65,772,419,072
effective H2D GB/s (clean pairs)	5.5 – 7.0	7.3 – 9.5
materialize_ms	7,300 – 9,463	8,244 – 9,354
htod_ms	9,398 – 11,997	6,931 – 9,036
vram_free_ms	16,177 – 39,675	5,624 – 5,710
vram_alloc_ms	0.000	0.000
commits / releases	57,117 / 53,999	57,117 / 53,999
captures / fallbacks	2 / 0	2 / 0
peak_committed_physical_bytes	6,987,710,464	6,987,710,464
ref/byte_underflows, unaccounted	0 / 0 / 0	0 / 0 / 0
decode ms/token (noisy)	3,545 – 6,125	1,700 – 1,831
throughput (noisy)	0.16 – 0.28 tok/s	0.55 – 0.59 tok/s

¹ Baseline main has no pool; the equivalent "before" is one cuMemHostAlloc/cuMemFreeHost
pair per page-in = 5,535 allocations and 5,535 frees per window.

Token IDs byte-identical across every run (baseline and pooled):
[96347, 3375, 724, 11, 358, 2776, 14589, 311, 6723, 429, 498, 3003, 2581, 6617, 315, 752]

Reading the result honestly

• Primary, contention-proof: the steady-state window issues 0 pinned
  allocations for 5,535 page-ins — the pool serves 100% by reuse. This is the
  deterministic proof the defect is fixed.
• Effective H2D bandwidth improved ~7.0 → 9.5 GB/s in the best matched pair
  and ~5.5 → 7.3 GB/s in a second — a consistent uplift, but the magnitude is
  noisy and should not be over-claimed.
• vram_free_ms collapsing from 16–40 s to a stable ~5.6 s is a secondary
  effect I did not target (it's vmm-churn's item 1). The plausible mechanism:
  removing ~5,535 cuMemFreeHost calls — which take the context lock and
  serialize against the device — stops interleaving driver-lock stalls into the
  unmap path. Flagged as a finding for vmm-churn, not a claim I optimized it.

What still dominates (in scope to report, not fixed here)

The offload page-in is fully synchronous: htod_async_elapsed_ms
host-blocks on the copy's end event, so materialize_ms (host fill, ~7–9 s) and
htod_ms (device copy, ~7–9 s) sit on the critical path back-to-back with no
overlap between host fill and device transfer, and none with compute — despite
the async/fence infrastructure existing. That serialized fill+copy, not pinned
alloc/free, is now the largest streaming term. Overlapping them (double-buffered
fill while the previous copy is in flight, and a genuinely non-blocking copy so
transfer overlaps decode compute) is the next lever, and it interacts with
vmm-churn's coarser mappings. Note the current instrumentation
(htod_async_elapsed_ms) is itself what forces the copy synchronous — measuring
honest per-copy ms and preserving overlap are in tension here.

Fence-safety made structural (review follow-up)

The invariant is: a staging buffer is the source of an asynchronous H2D
copy and must not be reused or freed until that copy has actually completed,
not merely been enqueued. Today that holds because htod_async_elapsed_ms
host-blocks: cudarc's Event::elapsed_ms (cudarc-0.19.8/src/driver/safe/core.rs:603-611)
calls end.synchronize(), and its own doc comment says it "does block the host
until the DMA completes." So on this path the copy is provably finished before
the buffer can return to the pool.

Per review, that coupling is now enforced by the type system, not by a
comment, because removing exactly that blocking is the scheduled next piece of
work on this path (see "What still dominates" — the instrumentation is the
serialization):

• htod_async_elapsed_ms now returns a CopyCompleted witness — a zero-sized
  token whose only field is private to the runtime module, so it can be minted
  only by a copy primitive there that host-synchronizes the copy before
  returning. Nothing outside runtime can fabricate one.
• The pool's reuse path — PinnedStagingPool::release and
  PooledStaging::retire — consumes a CopyCompleted. Reuse is therefore
  unreachable without presenting proof the copy finished.
• A future switch to a non-blocking htod_async + deferred fence produces no
  witness at the reuse site, so the code fails to compile until the author
  threads a completion witness through after awaiting that fence. The unsafe
  combination cannot be reached by accident.

Which construction and why. Drop cannot consume an argument, so a
Drop-based return-to-pool could never require the witness. I therefore made
the return an explicit method (retire(self, CopyCompleted)) and demoted
PooledStaging::Drop to a leak-safe fallback that frees the buffer (drops
the inner pinned block) rather than returning it to the pool. Consequence: if a
future refactor forgets to retire, the buffer is freed and re-allocated next
time — a performance regression that the pinned_alloc_calls counter and its
regression test catch immediately — but it can never silently reuse a buffer
whose copy is still in flight. Drop never asserts (STATUS_STACK_BUFFER_OVERRUN
on this codebase). Unit tests that exercise the pool without a real copy use a
#[cfg(test)] pub(crate) CopyCompleted::new_for_test(), which keeps the tests
honest about requiring a witness without granting non-test code the ability to
forge one.

This change is compile-time only — no runtime behavior change — so the
before/after measurements above remain valid and were not re-run for numbers
(the mandatory memory-governance parity gate was re-run and passes).

Validation (all green, local)

• cargo fmt -p onnx-runtime-ep-cuda -- --check, -p onnx-genai-bench -- --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 (all pass)
• cargo test -p onnx-runtime-ep-cuda --features cuda --lib pinned_pool (4/4 pass, incl. regression guard)
• cargo test -p onnx-genai-engine --features cuda,native-backend --test mobius_seqmajor_growth_parity_native_cuda -- --ignored --test-threads=1 (2/2 pass — memory-governance parity)

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.67%. Comparing base (6c714cb) to head (36cb827).
⚠️ Report is 15 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #843      +/-   ##
==========================================
- Coverage   78.26%   77.67%   -0.59%     
==========================================
  Files         350      351       +1     
  Lines      142164   145346    +3182     
  Branches   142164   145346    +3182     
==========================================
+ Hits       111265   112902    +1637     
- Misses      26460    27969    +1509     
- Partials     4439     4475      +36     

Flag	Coverage Δ
cli-ort-linux	83.42% <ø> (ø)
cli-ort-windows	82.92% <ø> (-0.10%)	⬇️
offline	77.46% <ø> (-0.61%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 43 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[justinchuby]

Validated and merging.

The type-level guard is exactly right, and the reasoning behind its shape is the good part. Drop cannot consume an argument, so a Drop-based return could never require the witness — recognizing that and making the return an explicit PooledStaging::retire(self, CopyCompleted) instead, with Drop demoted to a leak-safe path that frees rather than recycles, is the correct resolution. It also fails in the safe direction: forgetting to retire costs a re-allocation that the pinned_alloc_calls guard catches, instead of silently reusing a buffer whose DMA read is still in flight.

I checked the witness for the hole I was most worried about — one witness authorizing the retire of a different buffer whose copy is still running. CopyCompleted is pub struct CopyCompleted(()) with a module-private field, a module-private new(), a cfg(test)-only forge, and no Clone, Copy or Default derive or impl anywhere. So it is a linear move-only token: one completed copy authorizes exactly one retire. That closes it. #[must_use] is a nice touch. Minting one on the src.is_empty() early return is sound — an empty copy has trivially completed.

I also verified the blocking premise in the dependency rather than trusting the comment: cudarc-0.19.8/src/driver/safe/core.rs:603-611, Event::elapsed_ms calls self.synchronize() then end.synchronize().

Gates I ran personally on a worktree with main merged in: cargo fmt -p onnx-runtime-ep-cuda -p onnx-genai-bench -- --check pass; cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings pass (exact CI gate); and the mandatory mobius_seqmajor_growth_parity_native_cuda -- --ignored --test-threads=1 — 2 passed in 57.37s with nvidia-smi verified at 0 compute apps immediately beforehand. That solo qualification matters here specifically: see #851, where I retracted a flaky-gate claim of my own after discovering my "solo" runs were contended by two loops I had started myself.

Your pooling is not implicated by any of that, and the discriminator you supplied is why I can say so cleanly: a genuine return-to-pool-before-fence corruption would fail the bit-identical parity subtest deterministically, not crash the process. Every red anyone saw was a whole-process crash with zero token-stream divergence.

Result stands: pinned_alloc_calls 5,535 → 0 per 16-token window, token IDs byte-identical, fallbacks=0/captures=2 unchanged, hit_rate 60.10% unchanged, commits/releases untouched at 57,117/53,999. The bounded pool (≤8 buffers, ≤512 MiB, converging in practice to a single buffer sized to the largest page-in) is a sensible ceiling rather than an open-ended trade of host RAM for speed.

And thank you for adding the htod_async_elapsed_ms finding to the PR body without implementing it. That restraint is right: it is now filed as sequenced follow-up on #837, and it has to land after this PR precisely because your guard is what makes it safe to attempt at all — a future switch to a deferred fence now fails to compile until someone threads a witness through after awaiting it.

Merging as item 2 of #837. #837 stays open; the de-instrumentation half and item 3 remain.

[github-actions]

🔴 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_preexpanded_dense_oncelock_like_proxy/1x1024x1024	38.39 µs	86.95 µs	+126.5%
🔴	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	38.00 µs	74.93 µs	+97.2%
🔴	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	515.38 µs	945.36 µs	+83.4%
🔴	matmul/medium_generic_bf16_threads=8/32x512x512	407.35 µs	639.11 µs	+56.9%
🔴	matmul/medium_generic_bf16_threads=1/32x512x512	566.30 µs	873.73 µs	+54.3%
🔴	matmul/large_generic_bf16_threads=8/32x1024x1024	1.35 ms	1.99 ms	+47.9%
🔴	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	351.15 µs	503.68 µs	+43.4%
🔴	tokenization/encode_tokens_per_second	353.23 µs	462.59 µs	+31.0%
⚠️	tokenization/decode_tokens_per_second	5.76 ms	7.37 ms	+27.8%
⚠️	matmul/large_generic_f16_threads=8/32x1024x1024	92.41 µs	111.71 µs	+20.9%
⚠️	matmul/large_generic_f32_threads=1/32x1024x1024	9.18 ms	11.05 ms	+20.5%
⚠️	gather/medium_f32_threads=1-internal/32768	3.65 µs	4.40 µs	+20.4%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	525.14 µs	583.41 µs	+11.1%
✅	add/medium_bf16_threads=1-internal/262144	3.11 ms	3.45 ms	+10.9%
✅	grammar_masking/llguidance_compute_mask/32	77.10 µs	83.88 µs	+8.8%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	73.81 µs	78.82 µs	+6.8%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	82.91 µs	88.43 µs	+6.7%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.13 ms	2.27 ms	+6.5%
✅	add/large_f32_threads=1-internal/4194304	47.24 ms	50.26 ms	+6.4%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.63 ms	3.85 ms	+6.2%
✅	qwen3_sampling_processors/top_k_top_p_fast	744.09 µs	785.57 µs	+5.6%
✅	qwen3_sampling_processors/top_k_partial_selection	138.76 µs	144.57 µs	+4.2%
✅	gather/medium_f16_threads=1-internal/32768	2.26 µs	2.34 µs	+3.9%
✅	gather/small_f16_threads=1-internal/4096	503.5 ns	519.1 ns	+3.1%
✅	sampling_latency/min_p_per_token	226.14 µs	233.08 µs	+3.1%
✅	add/large_f16_threads=1-internal/4194304	49.31 ms	50.75 ms	+2.9%
✅	kv_cache/alloc_dealloc_pages	40.12 µs	41.28 µs	+2.9%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	4.51 ms	4.64 ms	+2.9%
✅	gather/small_f32_threads=1-internal/4096	709.2 ns	728.4 ns	+2.7%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.18 ms	2.24 ms	+2.6%
✅	matmul/medium_generic_f16_threads=8/32x512x512	32.32 µs	32.85 µs	+1.6%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.72 ms	5.79 ms	+1.1%
✅	sampling_latency/top_k_per_token	56.90 µs	57.21 µs	+0.5%
✅	gather/small_bf16_threads=1-internal/4096	500.6 ns	497.8 ns	-0.6%
✅	logit_processing/seven_processor_chain_per_step	343.07 µs	337.12 µs	-1.7%
✅	sampling_latency/greedy_per_token	3.49 µs	3.36 µs	-3.9%
✅	add/large_bf16_threads=1-internal/4194304	46.77 ms	44.72 ms	-4.4%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.10 ms	1.04 ms	-5.4%
✅	reduce_mean/medium_f32_threads=1-internal/65536	278.43 µs	253.15 µs	-9.1%
✅	add/medium_f16_threads=1-internal/262144	3.07 ms	2.78 ms	-9.4%
✅	reduce_mean/small_f32_threads=1-internal/4096	18.25 µs	15.91 µs	-12.8%
✅	add/small_f16_threads=1-internal/1024	15.81 µs	13.78 µs	-12.9%
✅	gather/medium_bf16_threads=1-internal/32768	2.89 µs	2.48 µs	-14.3%
🟢	sampling_latency/top_p_per_token	478.33 µs	405.25 µs	-15.3%
🟢	matmul/medium_generic_f16_threads=1/32x512x512	33.68 µs	27.97 µs	-17.0%
🟢	matmul/medium_generic_f32_threads=1/32x512x512	2.70 ms	2.23 ms	-17.3%
🟢	add/medium_f32_threads=1-internal/262144	2.94 ms	2.35 ms	-20.0%
🟢	matmul/small_generic_f32_threads=1/1x256x256	43.36 µs	34.03 µs	-21.5%
🟢	add/small_bf16_threads=1-internal/1024	16.68 µs	13.06 µs	-21.7%
🟢	matmul/small_generic_bf16_threads=8/1x256x256	38.31 µs	29.01 µs	-24.3%
🟢	add/small_f32_threads=1-internal/1024	268.3 ns	202.8 ns	-24.4%
🟢	matmul/small_generic_f16_threads=1/1x256x256	45.93 µs	33.27 µs	-27.6%
🟢	matmul/small_generic_bf16_threads=1/1x256x256	43.63 µs	29.95 µs	-31.3%
🟢	gather/large_bf16_threads=1-internal/131072	19.57 µs	12.73 µs	-34.9%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.46 ms	908.26 µs	-37.7%
🟢	gather/large_f16_threads=1-internal/131072	25.93 µs	14.45 µs	-44.3%
🟢	matmul/small_generic_f16_threads=8/1x256x256	55.72 µs	28.20 µs	-49.4%
🟢	matmul/small_generic_f32_threads=8/1x256x256	77.39 µs	32.43 µs	-58.1%
🟢	gather/large_f32_threads=1-internal/131072	68.28 µs	25.54 µs	-62.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.39 4.10 8.24 }

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)