Repository navigation
feat(memory): split optional allocator capabilities - #1263
justinchuby wants to merge 2 commits into
Conversation
Keep ordinary allocation and terminal release on DeviceAllocator while exposing VirtualBacking and SharedMapping as independent capabilities discovered from the selected allocator. Preserve erased allocator injection, VMM granularity accounting, shared-prefix charging, and canonical EP release. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Restore the explicit governed commits-on-demand signal, restrict provider-managed mapped growth to the selected CUDA VMM, reject foreign shared-prefix cost queries, preserve eager device validation, and document the decommit migration. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 46c5d75b-8146-489c-b82f-08ee29c27ce4
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## justinchuby-1186-memory-api-phase1-reland #1263 +/- ##
=============================================================================
- Coverage 80.62% 79.99% -0.64%
=============================================================================
Files 360 362 +2
Lines 155718 158583 +2865
Branches 155718 158583 +2865
=============================================================================
+ Hits 125551 126856 +1305
- Misses 25560 27111 +1551
- Partials 4607 4616 +9
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
> **A100 safety update (head `13c37a7f`):** expected shared-prefix admission failures are now all-or-private: K/V commit is transactional, a failed V share rolls K back before enqueue, and only a successful rollback permits private-KV fallback. Fatal kernel faults remain process-fatal; the new cross-process VMM probe proves a worker fault does not affect an actively executing peer, while the current server is still in-process and does not yet consume physical shared-prefix metadata. See the [final verified conclusion](#1579 (comment)). Collapses phases 1–7 of the #1186 memory architecture rework onto current `main` as a single merge. The stack forked 270 commits ago (176 of them touching `crates/`), so rebasing it layer by layer would mean solving the same 17 conflicts fourteen times, with no reviewer ever looking at the intermediate states. Closes the stack: #1252 #1263 #1279 #1283 #1301 #1341 #1349 #1426 #1440 #1448 #1454 #1462 #1465 #1468 #1533. **A review guide is in the first comment.** It is the part worth reading — this diff is 103 files, but only three decisions in it are ones a compiler cannot check. ## What was already reviewed, and what wasn't Every phase was reviewed and approved by a session that was not its author, under the rejection-lockout rule (a rejected author never writes the next revision). | Phase | PR | Verdict | |---|---|---| | 1–5 | #1252 → #1426 | approved | | 6 | #1440, #1448, #1454 | **rejected** ×3 | | 6 | #1462 | approved | | 7 | #1465 | **rejected** | | 7 | #1468 | approved | | 7 test fixes | #1533 | approved, A100-verified | The four rejected rounds were not replaced — the later rounds are stacked **on top of** them. So the tree here is the approved state, but the history contains the rejected commits. **#1440 / #1448 / #1454 / #1465 should not be reviewed individually**; they close automatically. **Not reviewed anywhere:** the conflict resolutions themselves. That is what this PR is for. ## Verification Run here, on this merged tree: - `cargo check --workspace --all-targets` — clean - `cargo check -p onnx-genai-engine --features cuda,native-backend --all-targets` — clean. Worth calling out: the default feature set does **not** compile the `cfg(cuda)` code, which is where the riskiest edits are. Checking only the default set would have missed a real break (see the guide). - `cargo test` over 7 crates — **1105 passed, 2 failed, 88 ignored** - `cargo clippy --workspace --all-targets` - `cargo fmt` **The 2 failures and the 1 clippy error are pre-existing and proven so, not merge damage:** - `platform_capacity::{disk_capacity_is_measured_for_the_working_directory, an_explicit_byte_limit_is_honored_without_a_device_query}` — `platform_capacity.rs` is byte-identical to `main` (`git diff origin/main HEAD -- ` that file is empty). The cause is a macOS-only FFI layout bug: `fsblkcnt_t` is 4 bytes on macOS (confirmed: `sizeof(struct statvfs)=64`, `sizeof(fsblkcnt_t)=4`) while the Rust struct declares `f_blocks`/`f_bfree`/`f_bavail` as `u64`. CI is Linux, where it is 8 bytes. Left alone: it is a real bug but not this PR's. - `optimizer.rs` clippy `approx_constant` — present on both parents; the merged file is byte-identical to `main`. - Three files had `rustfmt` drift already present on `main` (`matmul_nbits.rs`, `normalization.rs`, `optimizer.rs`). Reverted rather than swept in, to keep the diff readable. **Final A100 revalidation (head `13c37a7f`):** full CUDA-memory GPU suite passed; CUDA EP default-parallel lib suite passed (**488 passed / 17 ignored**); real fp16 GQA covered three requests × two interleaved decode steps with two shared peers and one transactional private fallback, byte-identical to independent GPU and CPU references. A device-started/event-gated worker remained healthy through a peer process `CUDA_ERROR_ILLEGAL_ADDRESS`, owner exit, replacement worker, and process restart. Governor tests, targeted Miri, CI Clippy, fmt, and CUDA honesty also passed. ## Still open after this merges - `memory-plugin-provider-wiring` — the back half of Phase 6 criterion 8 (provider/context pinning), which fell in the gap between phases. **#1186 must not be closed until it lands.** This merge makes it more tractable, not less: see decision 3 in the guide. - `memory-deferred-invariant-asserts` — two surviving mutants found by #1462's reviewer, adjudicated NON-BLOCKING and not defects. - #1533's CUDA honesty guard warns on a legitimately portable anchor. Should be moved out of the CUDA test binary or allow-listed rather than left warning. --------- Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <copilot@github.com> Copilot-Session: 46c5d75b-8146-489c-b82f-08ee29c27ce4 Copilot-Session: 39ff6824-d35f-4d3b-8f5f-043a7119a100 Copilot-Session: c80f8522-983c-47f7-8241-2155a823aabe
|
Closing: superseded by #1579, which collapsed phases 1–7 into a single merge — landed as This PR was approved by a session that was not its author, and its work is in Containment verified, not assumed: 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 Why this didn't close itself: #1579's body says #1186 stays open: #1579 explicitly gates it on |
Summary
DeviceAllocatorto device identity, ordinary allocation, canonical whole-allocation release, and an explicit default-false governedcommits_on_demand()signalVirtualBackingandSharedMappingcapabilities inonnx-runtime-memory-apiwith_memoryinjection when VMM is already selected instead of silently ignoring itPart of #1186
Public API
DeviceAllocator,HostAllocator,MemoryRole, andMemoryErrornow live inonnx-runtime-memory-api; governor paths re-export them for compatibilityDeviceAllocator::{as_virtual_backing, as_shared_mapping}discover optional borrowed capabilities from the selected allocator referenceDeviceAllocator::commits_on_demand()remains an explicit accounting opt-in:truerequires both lazy physical mapping and governor charging; capability presence alone is not enoughVirtualBackingcovers reserve/commit/decommit/committed-byte operations; newSharedMappingcovers shared-prefix creation, validated incremental cost, and mappingSharedMapping::incremental_owned_bytes_for_shared_prefixreturnsResult<u64, MemoryError>so foreign device/type/authority inputs are rejected before any cost is reportedwith_memory(Arc<dyn DeviceAllocator>)signatures are unchanged and exercised through an already-erased trait objectIntentional EP decommit migration
ExecutionProvider::decommit_allocation_rangeno longer inherits a successfulOk(0)no-op. Partial decommit has no eager equivalent, so an EP withoutVirtualBackingnow returns actionable unsupported. Out-of-tree EPs that relied on the old default must implement real decommit with actual unmapped-byte reporting or handle the unsupported result. Ordinary whole-allocation release remains unchanged onExecutionProvider::deallocate/DeviceAllocator.This migration is also recorded in
docs/ep-plugin/EP_PLUGIN_EXPORT_INVENTORY.md.Accounting and release evidence
VirtualBackingis explicitly unsupportedCudaVmmAllocator+ explicit governed signal + discoveredVirtualBacking; arbitrary injected capabilities never enter the concrete capacity-token pathcuMemGetAllocationGranularityDeviceAllocator::deallocate[_with_unmapped]through the owning EP, including allocations reserved throughVirtualBackingVirtualBackingandSharedMappingthroughArc<dyn DeviceAllocator>without allocatorAny,'static, or identity-token requirementsValidation
Head:
d120140a6de480e62c01163f8b2fc6821dceead5-D warnings: passed for memory API, governor, CUDA memory, EP API, ORT, and session; CPU/CUDA EP scopes pass when allowing only unchanged baseline lints listed below-D warnings: passed for memory API and CUDA memory; the restored session binding link resolves and is absent from the session crate's unchanged rustdoc failuresgit diff --check: passedvirtual_memory.rsis unchangedExisting base failures
cargo fmt --all -- --checkstill flags only unchangedonnx-genai-server/src/routes/completions.rsaccelerate_gemm.rs,matmul.rs,device_argmax.rs,matmul_nbits.rs,standard_attention.rs,optimizer.rs)-D warnings