[justinchuby]

Part of #1186

Context

PR #1233 (Phase 2 of #1186, branch justinchuby-1186-memory-api-phase2-round5-fix) went through six review rounds trying to make an "unforgeable mechanism identity + atomic release" apparatus sound enough to safely bind capability allocation/release to one mechanism. Round 6 found it was still unsound (a #[repr(C)] split-inner wrapper could forward a foreign capability from an offset-zero field and defeat any self-reported identity check; the CUDA release path's "atomicity" claim did not match what the code actually does; the design kept adding machinery rather than closing the gap). The coordinator ordered a full stop: PR #1233 is now marked draft and is being left completely untouched by this PR — no commits, pushes, or state changes were made to it or its branch.

This PR is a corrective, design-only response, on a fresh branch off main, that removes the unsound apparatus rather than patching it further, and is scoped to what actually needs fixing on main today.

What this does

• Deletes MechanismId and capability_shares_mechanism entirely. No self-reported identity token remains anywhere in the public API. DeviceAllocator::mechanism_id() and the HostAllocator/CudaVmmAllocator overrides are removed. No new 'static/Any bound is added to the base DeviceAllocator trait — non-'static borrowed implementors remain valid (a dedicated compile/runtime test proves this).
• CudaExecutionProvider::with_memory keeps only its real checks. It still rejects an injected allocator whose device disagrees with the provider's own device, and still rejects injection once a VMM arena is already selected (unchanged, real, narrower guarantee). It no longer claims to reject foreign-capability composition via identity — because no self-attested identity check can actually prove that (a type can honestly return Some(self) from a capability accessor while its methods secretly operate on unrelated state; closing this for real needs a registry that issues capability access rather than trusting self-report, which is [memory] Separate allocator, backing, and lifetime capabilities #1186 Phase 3's "provider/context pinning and binding identity" work, not Phase 2's).
• SharedMapping::release_shared_mapping's doc no longer overclaims atomicity. It now honestly describes the bookkeeping-level exact-once guarantee this crate's own locking already provides (repeated/unmatched calls on the same ptr refund zero rather than double-counting), explicitly distinguished from genuine CUDA-driver-level atomicity, which is deferred to Phase 4's owning-token/generation design.
• Documents (rather than hides) the pre-existing deallocate_span/release_granules_report partial-unmap gap. If a partial CUDA release fails for one granule run, the address range is still handed back to the free list unconditionally, which can let a later allocation overlap granules the bookkeeping still believes are committed. This predates Phase 2 by months and is not something Phase 2 introduced or altered — it's called out explicitly as deferred to Phase 4 (owning allocation token / generation counter to keep a VA range out of the free list until release reaches a consistent terminal state), instead of being patched with more unproven machinery under time pressure.

Why blockers 3 and 5 needed no fix here

Two of the five round-6 blockers (ReleaseReport.complete being discarded by callers, and a generic with_memory<T> regression) only existed on the un-merged PR #1233 branch. main never had ReleaseReport — release methods here are still bare u64 — and with_memory is already the trait-object Arc<dyn DeviceAllocator> form. No code changes were needed on main for those two.

What's preserved (round-1 findings 2-5, from the original Phase 2 work)

• as_shared_mapping() still returns None for detached/pool-less allocators.
• incremental_owned_bytes_for_shared_prefix still checks device/authority before treating a foreign prefix as zero-cost.
• mapped_bytes_for_allocation_ranges's default is still a real, granularity-aware conservative estimate (not a naive sum).
• deallocate_with_unmapped mapped-release reporting still lives on VirtualBacking/SharedMapping rather than silently defaulting to 0 for capability-bearing mechanisms.
• Capability discovery still follows the one mechanism selected at construction (no stale-retarget concern — an OnceLock-style selection remains in place).

Validation

• cargo check --workspace --tests --exclude onnx-genai-bench — clean (only pre-existing unrelated warnings in onnx-genai-ort).
• cargo clippy -p onnx-runtime-memory-api -p onnx-runtime-memory-governor -p onnx-runtime-cuda-memory -p onnx-runtime-ep-cuda --tests -- -D warnings — zero new findings in any file this PR touches; all remaining errors are confirmed pre-existing on main via git stash comparison (all in onnx-runtime-ep-cpu/onnx-runtime-ep-cuda kernel/optimizer files unrelated to memory code).
• cargo test -p onnx-runtime-memory-api -p onnx-runtime-memory-governor -p onnx-runtime-cuda-memory -p onnx-runtime-ep-cuda — all green, 0 failures (381 tests in the ep-cuda unit binary, 13 in memory-api, 61 in memory-governor, remaining GPU-gated binaries correctly ignored in this no-CUDA sandbox).
• cargo doc --no-deps -p onnx-runtime-memory-api -p onnx-runtime-cuda-memory with RUSTDOCFLAGS="-D warnings" — clean.
• cargo doc --no-deps -p onnx-runtime-memory-governor -p onnx-runtime-ep-cuda with RUSTDOCFLAGS="-D warnings" — remaining errors are byte-for-byte identical to main's baseline (diffed the exact error/line sets against a git stashd tree); this PR introduces zero new rustdoc regressions.

Scope discipline

• Base is main (not justinchuby-1186-memory-api-phase-1/PR fix(memory): structurally pair capability allocation/release (#1186 Phase 2 round-5 fix) #1233's branch), since this is a corrective fix to what's already merged, not a continuation of the draft.
• PR fix(memory): structurally pair capability allocation/release (#1186 Phase 2 round-5 fix) #1233 itself is left completely untouched — still draft, no commits/pushes/state changes.
• No eviction/residency policy change, no ProcessMemoryManager, no RAII allocation handle/deferred-free queue, no plugin ABI, no ExecutionProvider::allocator(), no Phase 3/4 work started.

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

[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
🔴	matmul/large_generic_f16_threads=8/32x1024x1024	84.73 µs	153.38 µs	+81.0%
🔴	sampling_latency/top_k_per_token	54.62 µs	74.33 µs	+36.1%
🔴	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	54.93 µs	74.15 µs	+35.0%
🔴	sampling_latency/top_p_per_token	399.64 µs	530.30 µs	+32.7%
🔴	kv_cache/alloc_dealloc_pages	37.87 µs	49.74 µs	+31.3%
🔴	matmul/medium_generic_f16_threads=1/32x512x512	30.11 µs	39.19 µs	+30.2%
⚠️	matmul/small_generic_bf16_threads=1/1x256x256	31.61 µs	40.90 µs	+29.4%
⚠️	matmul/small_generic_bf16_threads=8/1x256x256	32.17 µs	40.92 µs	+27.2%
⚠️	matmul/small_generic_f32_threads=8/1x256x256	53.87 µs	67.81 µs	+25.9%
⚠️	matmul/small_generic_f32_threads=1/1x256x256	42.23 µs	53.14 µs	+25.8%
⚠️	matmul/medium_generic_f16_threads=8/32x512x512	30.52 µs	38.34 µs	+25.6%
⚠️	matmul/large_generic_f32_threads=8/32x1024x1024	4.01 ms	4.83 ms	+20.3%
⚠️	gather/large_bf16_threads=1-internal/131072	15.74 µs	18.58 µs	+18.1%
⚠️	matmul/large_generic_f32_threads=1/32x1024x1024	9.39 ms	11.00 ms	+17.1%
⚠️	grammar_masking/llguidance_compute_mask/32	81.63 µs	95.39 µs	+16.9%
⚠️	add/medium_f32_threads=1-internal/262144	26.79 µs	31.27 µs	+16.7%
⚠️	matmul/small_generic_f16_threads=8/1x256x256	32.54 µs	37.93 µs	+16.6%
⚠️	matmul/small_generic_f16_threads=1/1x256x256	32.24 µs	37.55 µs	+16.5%
⚠️	add/large_f32_threads=1-internal/4194304	786.74 µs	907.85 µs	+15.4%
⚠️	add/large_f16_threads=1-internal/4194304	1.93 ms	2.22 ms	+15.2%
✅	gather/medium_f32_threads=1-internal/32768	4.54 µs	5.21 µs	+14.8%
✅	add/medium_bf16_threads=1-internal/262144	107.03 µs	122.79 µs	+14.7%
✅	gather/medium_f16_threads=1-internal/32768	2.59 µs	2.94 µs	+13.7%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	378.57 µs	424.70 µs	+12.2%
✅	logit_processing/seven_processor_chain_per_step	324.27 µs	358.57 µs	+10.6%
✅	add/small_bf16_threads=1-internal/1024	482.6 ns	532.2 ns	+10.3%
✅	gather/small_f16_threads=1-internal/4096	526.9 ns	574.6 ns	+9.0%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.41 ms	6.90 ms	+7.8%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	549.16 µs	591.83 µs	+7.8%
✅	tokenization/decode_tokens_per_second	7.11 ms	7.63 ms	+7.3%
✅	qwen3_sampling_processors/top_k_top_p_fast	702.25 µs	743.77 µs	+5.9%
✅	matmul/medium_generic_f32_threads=8/32x512x512	917.19 µs	970.56 µs	+5.8%
✅	gather/small_f32_threads=1-internal/4096	766.0 ns	806.8 ns	+5.3%
✅	qwen3_sampling_processors/top_k_partial_selection	158.94 µs	166.32 µs	+4.6%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	79.86 µs	83.19 µs	+4.2%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	586.30 µs	609.76 µs	+4.0%
✅	gather/small_bf16_threads=1-internal/4096	539.3 ns	557.4 ns	+3.4%
✅	sampling_latency/greedy_per_token	3.45 µs	3.55 µs	+2.9%
✅	add/large_bf16_threads=1-internal/4194304	2.02 ms	2.08 ms	+2.9%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.34 ms	2.39 ms	+2.4%
✅	gather/medium_bf16_threads=1-internal/32768	2.76 µs	2.83 µs	+2.3%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.73 ms	2.79 ms	+1.9%
✅	add/medium_f16_threads=1-internal/262144	129.34 µs	131.16 µs	+1.4%
✅	tokenization/encode_tokens_per_second	453.33 µs	447.01 µs	-1.4%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.14 ms	1.12 ms	-1.8%
✅	gather/large_f16_threads=1-internal/131072	13.57 µs	13.07 µs	-3.7%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	4.16 ms	3.96 ms	-4.9%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.45 ms	2.32 ms	-5.3%
✅	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	625.69 µs	576.93 µs	-7.8%
✅	sampling_latency/min_p_per_token	253.43 µs	229.67 µs	-9.4%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	218.72 µs	188.44 µs	-13.8%
✅	reduce_mean/medium_f32_threads=1-internal/65536	329.48 µs	281.36 µs	-14.6%
🟢	reduce_mean/small_f32_threads=1-internal/4096	19.63 µs	16.48 µs	-16.0%
🟢	add/small_f32_threads=1-internal/1024	248.2 ns	205.4 ns	-17.2%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	1.73 ms	1.34 ms	-22.6%
🟢	add/small_f16_threads=1-internal/1024	643.1 ns	492.7 ns	-23.4%
🟢	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	707.62 µs	497.60 µs	-29.7%
🟢	gather/large_f32_threads=1-internal/131072	43.31 µs	30.30 µs	-30.0%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	139.95 µs	67.60 µs	-51.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: { 4.99 4.96 6.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)

[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 corrective branch targeted code no longer present on main.