fix(memory): keep Phase 3 selection live and allocator pins alive

[justinchuby]

Summary

• stop selection withdrawal from restoring a dead or unregistered predecessor, and stop two failing selects from leaving a dead mechanism selected
• make retirement and device loss terminal before they drop the selection, closing the mirror-image race that wedged a device just as permanently
• run the allocator destructor while its provider-context and authority pins are still alive

Part of #1186; third revision of rejected Phase-3 work #1279 and #1283, stacked on justinchuby-fix-phase3-binding-races.

Selection can no longer point at a dead mechanism

select publishes a candidate and then re-checks it, because a mechanism can be retired or lost in between. The old withdrawal restored the recorded predecessor unconditionally, which had three failure modes — each terminal, because a stale identity in the slot makes bind fail forever and makes registration's selected.entry(device).or_insert(..) self-heal a no-op:

• a predecessor retired or device-lost after it was recorded was resurrected;
• a predecessor removed after it was recorded was resurrected as an unregistered identity;
• with two concurrently failing selects, the later withdrawal restored the earlier, already-dead candidate.

withdraw_failed_selection now (1) touches the slot only while the candidate still owns it, so a losing candidate never overwrites a newer selection; (2) restores the predecessor only while it is still registered and Active, re-confirming registration under the same lock that publishes it; (3) otherwise clears the slot, because an absent selection is the only state a later registration heals. A restored predecessor is re-checked exactly like the candidate was, and that retry carries no fallback, so the loop runs at most twice.

Fixing only the withdrawal left a mirror-image window, found in review of this revision. retire and invalidate_device dropped the selection under the registry lock and then flipped the lifecycle under the mechanism lock. Since the two lock classes may never be held together, a select that had already validated the mechanism could publish it after the clear and still observe Active at its re-check — so both calls returned Ok leaving a retired or lost mechanism permanently selected. Both now make the lifecycle terminal first, so any selection published afterwards must have validated earlier and withdraws itself. Device loss additionally drops only a selection naming a mechanism it actually invalidated.

Allocator teardown outlives its pins

MechanismEntry declared its provider-context and authority pins before the allocator, and Rust drops fields in declaration order, so a third-party allocator that releases device state from Drop ran after both pins were gone. The allocator and its pins now live in one MechanismResources owner whose declaration order is documented as load-bearing: allocator, then authority, then provider context.

Preserved

BoundSharedPrefix field order and its drop-order test from #1283, the Phase-4 disclaimers, and the EP mapped-attribution refund adoption warning are unchanged. This PR still adds no deferred-free queue, fence/event scheduling, owning allocation RAII, physical-release completeness state, partial-unmap recovery, quarantine, pointer-only retry API, Phase-5 policy, identity redesign, or virtual_memory change. Registry and mechanism locks remain unnested, and no allocator or capability callback runs under a lock.

Validation

Nine barrier-sequenced tests drive the real BindingRegistry — no stub — through healthy predecessor restore, a retired predecessor, a removed predecessor, a newer selection, two failing selects, later-registration healing, the retire/select race, the device-loss/select race, and the plain retire and device-loss selection drops. Each was confirmed to fail under a controlled revert of the exact behaviour it covers:

• reverting the withdrawal to the unvalidated restore fails the retired-predecessor, removed-predecessor, two-failing-selects, and healing tests (the healthy-restore and newer-selection tests still pass, as expected for preserved properties);
• reverting retire/invalidate_device to the old phase order fails both race tests, with select wrongly returning Ok(());
• deleting either selection clear fails the selection-drop test;
• inverting MechanismResources field order fails the allocator drop test with context alive: false, authority alive: false.

Suites: onnx-runtime-memory-api 34 passed (18 lib, 15 integration, 1 doc; 25 at base), onnx-runtime-memory-governor 63, onnx-runtime-ep-cpu --test shared_allocator 8, onnx-runtime-ep-cuda --lib 382 passed / 22 ignored, onnx-runtime-cuda-memory --lib 8, onnx-genai-ort governed_allocator 20. Registry concurrency stress 20/20 runs (4,800 iterations); the barrier race tests 20/20 runs. Scoped cargo clippy -p onnx-runtime-memory-api -p onnx-runtime-memory-governor --all-targets -- -D warnings clean, strict memory-API rustdoc clean, git diff --check clean, and an independent high-confidence review found no remaining issues.

Compared against the exact #1283 base at 5f1b249e, inherited failures are unchanged:

• workspace cargo fmt --all -- --check reports only the same two pre-existing onnx-genai-server/src/routes/completions.rs diffs at lines 2891 and 2914;
• workspace all-target check reports only the same three pre-existing onnx-genai-bench/tests/fused_batch_prefill.rs import errors;
• strict governor rustdoc reports only the same pre-existing private BUDGET_SHARE_DENOMINATOR and unresolved KvLayout links.

[justinchuby]

Final independent review: the #1279 → #1283 rejection chain is resolved. Verified candidate CAS guards, registered+active predecessor restoration with bounded retry, B/C failure orders, C preservation, retire/device-loss lifecycle ordering, registration self-healing, unnested lock classes, allocator-before-authority/context destruction, and deterministic production-registry regressions. No Phase-4/5 scope creep found. Combined Phase 3 is review-ready and remains open/unmerged for human review.

[codecov]

Codecov Report

❌ Patch coverage is 93.18182% with 27 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.92%. Comparing base (5f1b249) to head (640a9d9).

Files with missing lines	Patch %	Lines
crates/onnx-runtime-memory-api/src/binding.rs	93.18%	13 Missing and 14 partials ⚠️

Additional details and impacted files

@@                           Coverage Diff                            @@
##           justinchuby-fix-phase3-binding-races    #1301      +/-   ##
========================================================================
- Coverage                                 80.36%   79.92%   -0.45%     
========================================================================
  Files                                       363      363              
  Lines                                    159665   159976     +311     
  Branches                                 159665   159976     +311     
========================================================================
- Hits                                     128315   127853     -462     
- Misses                                    26649    27422     +773     
  Partials                                   4701     4701              

Flag	Coverage Δ
mlas	85.09% <ø> (-0.23%)	⬇️
offline	79.81% <93.18%> (-0.45%)	⬇️

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

Files with missing lines	Coverage Δ
crates/onnx-runtime-memory-api/src/binding.rs	71.71% <93.18%> (+7.66%)	⬆️

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

⚠️ Benchmark Change 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
⚠️	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.59 ms	4.61 ms	+28.2%
⚠️	add/small_f16_threads=1-internal/1024	502.8 ns	643.2 ns	+27.9%
⚠️	qwen3_sampling_processors/top_k_partial_selection	146.73 µs	185.86 µs	+26.7%
⚠️	sampling_latency/greedy_per_token	3.25 µs	4.10 µs	+26.3%
⚠️	qwen3_sampling_processors/top_k_full_sort_baseline	2.18 ms	2.74 ms	+25.6%
⚠️	sampling_latency/min_p_per_token	213.85 µs	267.67 µs	+25.2%
✅	add/medium_bf16_threads=1-internal/262144	113.45 µs	129.38 µs	+14.0%
✅	tokenization/encode_tokens_per_second	391.40 µs	443.06 µs	+13.2%
✅	add/large_f32_threads=1-internal/4194304	767.76 µs	854.16 µs	+11.3%
✅	sampling_latency/top_k_per_token	57.75 µs	64.19 µs	+11.2%
✅	sampling_latency/top_p_per_token	389.06 µs	428.27 µs	+10.1%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	529.47 µs	572.16 µs	+8.1%
✅	qwen3_sampling_processors/top_k_top_p_fast	666.18 µs	716.97 µs	+7.6%
✅	add/small_bf16_threads=1-internal/1024	543.5 ns	582.7 ns	+7.2%
✅	add/medium_f16_threads=1-internal/262144	124.21 µs	131.65 µs	+6.0%
✅	kv_cache/alloc_dealloc_pages	38.35 µs	40.04 µs	+4.4%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.80 ms	5.89 ms	+1.6%
✅	grammar_masking/llguidance_compute_mask/32	78.87 µs	79.75 µs	+1.1%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	115.37 µs	116.44 µs	+0.9%
✅	logit_processing/seven_processor_chain_per_step	326.84 µs	329.73 µs	+0.9%
✅	gather/small_f32_threads=1-internal/4096	738.9 ns	741.0 ns	+0.3%
✅	tokenization/decode_tokens_per_second	6.82 ms	6.78 ms	-0.6%
✅	add/medium_f32_threads=1-internal/262144	26.91 µs	26.68 µs	-0.8%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.09 ms	1.08 ms	-1.5%
✅	reduce_mean/medium_f32_threads=1-internal/65536	276.29 µs	271.22 µs	-1.8%
✅	reduce_mean/small_f32_threads=1-internal/4096	16.57 µs	16.11 µs	-2.8%
✅	matmul/medium_generic_f16_threads=8/32x512x512	47.20 µs	45.28 µs	-4.1%
✅	gather/large_f16_threads=1-internal/131072	11.98 µs	11.18 µs	-6.7%
✅	gather/small_f16_threads=1-internal/4096	528.6 ns	493.3 ns	-6.7%
✅	gather/large_bf16_threads=1-internal/131072	12.61 µs	11.69 µs	-7.3%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	620.53 µs	567.15 µs	-8.6%
✅	gather/small_bf16_threads=1-internal/4096	530.1 ns	479.2 ns	-9.6%
✅	gather/medium_bf16_threads=1-internal/32768	2.69 µs	2.43 µs	-9.9%
✅	gather/medium_f16_threads=1-internal/32768	2.71 µs	2.42 µs	-10.7%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.59 ms	2.31 ms	-11.1%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	106.60 µs	94.52 µs	-11.3%
✅	gather/large_f32_threads=1-internal/131072	29.27 µs	25.88 µs	-11.6%
✅	add/small_f32_threads=1-internal/1024	290.1 ns	256.2 ns	-11.7%
✅	gather/medium_f32_threads=1-internal/32768	4.40 µs	3.80 µs	-13.6%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	10.87 ms	9.24 ms	-15.0%
🟢	add/large_bf16_threads=1-internal/4194304	2.17 ms	1.83 ms	-15.6%
🟢	add/large_f16_threads=1-internal/4194304	2.12 ms	1.71 ms	-19.3%
🟢	matmul/small_generic_f32_threads=1/1x256x256	45.98 µs	37.08 µs	-19.4%
🟢	matmul/small_generic_bf16_threads=1/1x256x256	40.05 µs	31.88 µs	-20.4%
🟢	matmul/small_generic_f16_threads=1/1x256x256	39.38 µs	31.02 µs	-21.2%
🟢	matmul/medium_generic_f16_threads=1/32x512x512	41.57 µs	32.42 µs	-22.0%
🟢	matmul/small_generic_bf16_threads=8/1x256x256	40.51 µs	31.57 µs	-22.1%
🟢	matmul/large_generic_f16_threads=1/32x1024x1024	105.86 µs	80.55 µs	-23.9%
🟢	matmul/large_generic_bf16_threads=1/32x1024x1024	2.66 ms	2.02 ms	-23.9%
🟢	matmul/small_generic_f16_threads=8/1x256x256	41.40 µs	30.69 µs	-25.9%
🟢	matmul/small_generic_f32_threads=8/1x256x256	55.36 µs	35.39 µs	-36.1%
🟢	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	602.58 µs	382.80 µs	-36.5%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	6.01 ms	3.79 ms	-37.0%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	2.31 ms	1.35 ms	-41.5%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.58 ms	911.16 µs	-42.5%
🟢	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	109.65 µs	60.77 µs	-44.6%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	656.07 µs	361.97 µs	-44.8%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	133.37 µs	56.32 µs	-57.8%
🟢	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	1.45 ms	596.90 µs	-58.9%

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.85 4.84 5.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]

Closing: superseded by #1579, which collapsed phases 1–7 into a single merge — landed as a36964280 (squash, 126 files, +43956 −3286).

This PR was approved by a session that was not its author, and its work is in main.

Containment verified, not assumed:

$ git merge-base --is-ancestor 640a9d9a7 ebceab071   # ebceab071 = #1579 head
→ ancestor

Every one of the 15 stack heads (#1252 → #1533) is a literal ancestor of the merged head, and I separately confirmed the merged head's content reached main: of the 126 files #1579 touched, exactly one differs from main — crates/onnx-runtime-session/src/executor/tests.rs, where main carries 120 extra lines from #1703 (Expand-broadcast mask tests). That is the conflict resolution correctly preserving the other side, not content loss.

Why this didn't close itself: #1579's body says Closes the stack: #1252 #1263 …, which GitHub does not parse — the closing keyword must be immediately followed by the reference, and in any case closing keywords only ever close issues, never pull requests. So all 15 stayed open by mechanism, not by intent.

#1186 stays open: #1579 explicitly gates it on memory-plugin-provider-wiring (the back half of Phase 6 criterion 8, provider/context pinning), which has not landed.