[justinchuby]

Part of #1186

Important context: PRs #1192 and #1206 were merged before round-5 review fixes landed

While the round-5 review of PR #1206 (Phase 2) was in progress, both
#1206 (base justinchuby-1186-memory-api-phase-1, containing Phase 2
through round-4 fixes) and #1192 (base main, "phases 1-2") were
merged, and both source branches were deleted. This happened outside
this session's control — the round-5 review findings below were
never applied to those merged PRs, so main currently contains the
round-4 state, which round-5 review explicitly found structurally
unsafe.

Since the original branch stack (justinchuby-1186-memory-api-phase-1
-> justinchuby-1186-memory-api-phase-2) no longer exists, this PR
carries the completed round-5 redesign as a direct follow-up against
main. It contains only the round-5 fix commit — verified to
cherry-pick cleanly with zero conflicts, because the file contents at
the round-4 head (6845826e) are byte-for-byte identical to what
squash-merging left in main (confirmed via direct diff of every
touched file). No Phase 1/Phase 2 content is duplicated; this is a
pure incremental fix on top of what's already in main.

What round-5 review found (structurally unsafe, not just buggy)

1. release_shared_mapping-style code removed live tracking (VA
   spans, granule refs) and gave the address/space back to the free
   list before, or regardless of whether, the underlying CUDA
   unmap/handle-return actually succeeded — a UAF/double-free-adjacent
   race: another thread could reuse the same address while the first
   release was still in flight or had failed.
2. A bare 0 return conflated "nothing was mapped" with "release
   failed", so a caller (including the base-trait default path) had no
   way to distinguish a successful no-op from a failure, and would
   proceed as if it had succeeded either way.
3. CudaVmmAllocator's granule-release paths removed live tracking and
   returned free-list space unconditionally, so a partial CUDA
   failure (some granules unmapped, others not) could not be
   represented — tracking/free-list state would diverge from the
   actual mapping state.
4. The round-4 MechanismId (a self-reported {ptr, TypeId} pair) is
   not an unforgeable identity proof: an implementation can report a
   foreign object's identity, ZST/address-reuse can produce spurious
   collisions, and it can also wrongly reject legitimate transparent
   wrappers.

The redesign (this commit)

Rather than continuing to patch the identity check, the release seam
was redesigned so capability release is structurally paired with
capability allocation instead of being validated by self-reported
identity:

• ReleaseReport { unmapped_bytes: u64, complete: bool } replaces
  the bare u64 return of every release-shaped capability method:
  DeviceAllocator::deallocate_with_unmapped,
  VirtualBacking::deallocate_committed,
  VirtualBacking::decommit_allocation_range, and
  SharedMapping::release_shared_mapping. Failure is now
  Err(MemoryError); a genuine "0 bytes were mapped, and release
  completed" is Ok(ReleaseReport::complete(0)); a partial CUDA
  failure is Ok(ReleaseReport::partial(n)). 0 and "failed" can no
  longer be confused at any point in the chain (closes finding 2).

• DeviceMemoryMechanism, one owned capability bundle, replaces
  the forgeable identity check entirely. Its constructors (eager,
  with_virtual_backing, with_shared_mapping,
  with_virtual_backing_and_shared_mapping) each take exactly one
  Arc<T> and populate the allocator/virtual-backing/shared-mapping
  fields via Arc::clone of that same object, requiring the same
  concrete T to implement every requested trait simultaneously. A
  composite/wrapper object whose as_virtual_backing() override
  claims a foreign backing can still compile against the minimal
  with_memory (which only needs DeviceAllocator), but
  DeviceMemoryMechanism::eager never calls that override, so the
  foreign capability can never surface. This closes finding 4
  structurally — there is no self-reported identity left to forge,
  so no check can be fooled by it. Proven by two rewritten tests:
  with_memory_never_surfaces_a_composites_foreign_virtual_backing
  and with_memory_never_surfaces_a_foreign_vmm_embedded_at_offset_zero
  (the latter specifically reproduces the offset-zero
  #[repr(C)]-adversarial construction from the prior review round).
  Any is no longer a supertrait bound on the base DeviceAllocator
  trait anywhere, so non-'static borrowed adapters remain legal
  (this also closes round-4's blocker 2, which is still in scope here
  since it was never separately fixed/merged).

• vmm_allocator.rs's release paths (deallocate_span,
  decommit_allocation_range, and the allocate_committed rollback
  path) now follow a reconcile-before-mutating-tracking discipline:
  granule/span state is removed from live tracking and returned to
  the free list only after the underlying CUDA unmap/handle-return
  call has succeeded for that granule/span. release_granules_reconciled
  returns both the bytes actually released and the set of granules
  that remain owned (i.e., were not successfully released), so a
  partial failure leaves the allocator's internal state consistent
  with what's actually mapped -- nothing is double-counted, nothing is
  silently given back that wasn't actually freed. This closes findings
  1 and 3. Uncommitted-span rollback that cannot fully unwind is now
  counted through a bounded GLOBAL_ROLLBACK_LEAKS counter instead of
  forcing success by corrupting live/free-list state.

• CudaExecutionProvider now stores one DeviceMemoryMechanism
  instead of a separate memory + OnceLock<vmm> pair, so exactly one
  mechanism is selected per provider, chosen once at
  construction/binding and never retargeted after allocations exist
  (this preserves the construction-time/OnceLock analysis from the
  earlier architectural clarification in this epic). The
  with_memory/with_memory_and_virtual_backing/
  with_memory_and_shared_mapping/with_memory_and_capabilities
  family rejects only on a genuine after-VMM-selected conflict or a
  device mismatch -- never by re-validating a capability's
  self-reported identity, since there is no longer any such identity
  to validate.

Compatibility

• DeviceAllocator's minimal contract (device identity +
  allocate/deallocate) is unchanged; eager CPU/CUDA allocators, ORT
  allocator integration, and third-party governor/allocator
  implementations are unaffected except for the ReleaseReport return
  type change on the one release-shaped method they may override
  (default implementation preserved for allocators that don't need
  mapped-zone accounting).
• No eviction/residency policy change, no ProcessMemoryManager, no
  RAII allocation handle/deferred-free queue, no plugin ABI change, no
  ExecutionProvider::allocator(). EP stream/context responsibilities
  are untouched.
• Built-in CUDA eager allocation is not removed in this PR (per the
  epic-level decision that only a future, separately-labeled Phase 7
  will do that once VMM fully subsumes it).

Validation

cargo test -p onnx-runtime-memory-api -p onnx-runtime-memory-governor \
           -p onnx-runtime-cuda-memory -p onnx-runtime-ep-cuda
# 17 + 61 + 8 + (381 lib, 24 correctly ignored/GPU-only) passed, 0 failed

cargo test -p onnx-genai-ort --test governed_allocator_session
# 6 passed, 0 failed (ORT allocator integration)

cargo clippy -p onnx-runtime-memory-api -p onnx-runtime-memory-governor \
             -p onnx-runtime-cuda-memory --all-targets -- -D warnings
# clean

cargo clippy -p onnx-runtime-ep-cuda --no-deps --all-targets -- -D warnings
# clean except confirmed pre-existing baseline errors in unrelated files
# (matmul_nbits.rs, device_argmax.rs, standard_attention.rs, optimizer.rs),
# verified identical before/after this diff via git stash

RUSTDOCFLAGS="-D warnings" cargo doc -p onnx-runtime-memory-api \
             -p onnx-runtime-cuda-memory --no-deps
# clean

RUSTDOCFLAGS="-D warnings" cargo doc -p onnx-runtime-ep-cuda --no-deps
# 30 pre-existing errors (kernels/*, runtime.rs, weight_paging.rs doc
# links unrelated to this change), confirmed identical count before/
# after this diff via git stash

cargo check --workspace --tests --exclude onnx-genai-bench
# clean (onnx-genai-bench has pre-existing unrelated missing-feature
# errors, confirmed via git stash)

No Phase 3/4 work is included. This closes out the round-5 review
findings for #1186 Phase 2.

[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	53.86 µs	89.29 µs	+65.8%
🔴	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	51.83 µs	82.14 µs	+58.5%
🔴	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	856.02 µs	1.34 ms	+56.8%
🔴	gather/small_f32_threads=1-internal/4096	633.2 ns	918.0 ns	+45.0%
🔴	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	454.25 µs	650.21 µs	+43.1%
🔴	add/large_bf16_threads=1-internal/4194304	1.50 ms	2.13 ms	+42.1%
🔴	add/large_f32_threads=1-internal/4194304	602.94 µs	798.63 µs	+32.5%
🔴	add/large_f16_threads=1-internal/4194304	1.56 ms	2.05 ms	+31.5%
⚠️	matmul/medium_generic_f16_threads=8/32x512x512	34.51 µs	44.33 µs	+28.5%
⚠️	gather/small_f16_threads=1-internal/4096	460.5 ns	585.8 ns	+27.2%
⚠️	grammar_masking/llguidance_compute_mask/32	76.84 µs	97.65 µs	+27.1%
⚠️	add/medium_bf16_threads=1-internal/262144	94.82 µs	116.73 µs	+23.1%
⚠️	matmul/medium_generic_bf16_threads=1/32x512x512	580.83 µs	697.38 µs	+20.1%
⚠️	matmul/large_generic_f32_threads=8/32x1024x1024	3.83 ms	4.58 ms	+19.8%
⚠️	add/medium_f16_threads=1-internal/262144	102.64 µs	122.38 µs	+19.2%
⚠️	reduce_mean/large_f32_threads=1-internal/262144	933.56 µs	1.11 ms	+19.0%
⚠️	gather/medium_f32_threads=1-internal/32768	3.52 µs	4.19 µs	+18.9%
⚠️	gather/large_f16_threads=1-internal/131072	17.92 µs	21.18 µs	+18.2%
⚠️	matmul/small_generic_bf16_threads=8/1x256x256	36.05 µs	42.37 µs	+17.5%
⚠️	matmul/small_generic_f16_threads=8/1x256x256	32.05 µs	37.54 µs	+17.1%
⚠️	add/small_f16_threads=1-internal/1024	427.7 ns	493.2 ns	+15.3%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	87.60 µs	99.90 µs	+14.0%
✅	reduce_mean/medium_f32_threads=1-internal/65536	226.72 µs	256.34 µs	+13.1%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	163.02 µs	183.03 µs	+12.3%
✅	reduce_mean/small_f32_threads=1-internal/4096	14.22 µs	15.81 µs	+11.2%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	481.71 µs	534.30 µs	+10.9%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.30 ms	5.83 ms	+10.0%
✅	matmul/small_generic_f16_threads=1/1x256x256	33.73 µs	36.93 µs	+9.5%
✅	sampling_latency/top_k_per_token	52.37 µs	56.93 µs	+8.7%
✅	gather/small_bf16_threads=1-internal/4096	454.3 ns	490.1 ns	+7.9%
✅	add/medium_f32_threads=1-internal/262144	22.76 µs	24.55 µs	+7.9%
✅	logit_processing/seven_processor_chain_per_step	301.58 µs	323.91 µs	+7.4%
✅	add/small_bf16_threads=1-internal/1024	429.5 ns	456.2 ns	+6.2%
✅	matmul/large_generic_bf16_threads=8/32x1024x1024	1.59 ms	1.69 ms	+6.2%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.40 ms	3.60 ms	+6.1%
✅	kv_cache/alloc_dealloc_pages	37.33 µs	39.54 µs	+5.9%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.03 ms	2.15 ms	+5.9%
✅	qwen3_sampling_processors/top_k_top_p_fast	629.88 µs	659.88 µs	+4.8%
✅	sampling_latency/min_p_per_token	212.72 µs	221.82 µs	+4.3%
✅	tokenization/encode_tokens_per_second	383.38 µs	398.01 µs	+3.8%
✅	matmul/small_generic_bf16_threads=1/1x256x256	36.39 µs	37.49 µs	+3.0%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.13 ms	2.19 ms	+2.9%
✅	tokenization/decode_tokens_per_second	6.52 ms	6.69 ms	+2.5%
✅	gather/large_bf16_threads=1-internal/131072	17.73 µs	18.04 µs	+1.7%
✅	sampling_latency/top_p_per_token	392.86 µs	389.35 µs	-0.9%
✅	sampling_latency/greedy_per_token	3.46 µs	3.39 µs	-2.0%
✅	qwen3_sampling_processors/top_k_partial_selection	144.23 µs	140.91 µs	-2.3%
✅	gather/medium_bf16_threads=1-internal/32768	2.48 µs	2.42 µs	-2.5%
✅	add/small_f32_threads=1-internal/1024	210.7 ns	203.9 ns	-3.2%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	10.20 ms	9.68 ms	-5.1%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	96.33 µs	88.09 µs	-8.5%
✅	matmul/medium_generic_f16_threads=1/32x512x512	43.22 µs	39.41 µs	-8.8%
✅	gather/medium_f16_threads=1-internal/32768	2.77 µs	2.44 µs	-12.0%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	473.76 µs	412.55 µs	-12.9%
✅	gather/large_f32_threads=1-internal/131072	34.92 µs	30.07 µs	-13.9%
🟢	matmul/medium_generic_f32_threads=1/32x512x512	3.00 ms	2.41 ms	-19.6%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	2.11 ms	1.68 ms	-20.5%
🟢	matmul/small_generic_f32_threads=8/1x256x256	57.96 µs	42.96 µs	-25.9%
🟢	matmul/small_generic_f32_threads=1/1x256x256	53.40 µs	39.13 µs	-26.7%

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: { 2.83 3.31 4.99 }

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]

Superseded by the safety revert in #1247. Phase 1 and a narrower Phase 2 will be re-landed as a fresh reviewed stack; this fix-forward design must not be merged.