Skip to content

feat(memory): split optional allocator capabilities - #1263

Closed
justinchuby wants to merge 2 commits into
justinchuby-1186-memory-api-phase1-relandfrom
justinchuby-1186-memory-capabilities-phase2
Closed

justinchuby wants to merge 2 commits into
justinchuby-1186-memory-api-phase1-relandfrom
justinchuby-1186-memory-capabilities-phase2

Conversation

@justinchuby

@justinchuby justinchuby commented Aug 18, 2026 •

Copy link
Copy Markdown
Owner

Summary

  • reduce DeviceAllocator to device identity, ordinary allocation, canonical whole-allocation release, and an explicit default-false governed commits_on_demand() signal
  • add independently discoverable VirtualBacking and SharedMapping capabilities in onnx-runtime-memory-api
  • keep CUDA on one construction-selected memory mechanism; reject with_memory injection when VMM is already selected instead of silently ignoring it
  • preserve governor-specific mapped-capacity operations as inherent in-tree VMM methods without adding identity or Phase 3/4 ownership machinery

Part of #1186

Public API

  • DeviceAllocator, HostAllocator, MemoryRole, and MemoryError now live in onnx-runtime-memory-api; governor paths re-export them for compatibility
  • DeviceAllocator::{as_virtual_backing, as_shared_mapping} discover optional borrowed capabilities from the selected allocator reference
  • DeviceAllocator::commits_on_demand() remains an explicit accounting opt-in: true requires both lazy physical mapping and governor charging; capability presence alone is not enough
  • new VirtualBacking covers reserve/commit/decommit/committed-byte operations; new SharedMapping covers shared-prefix creation, validated incremental cost, and mapping
  • SharedMapping::incremental_owned_bytes_for_shared_prefix returns Result<u64, MemoryError> so foreign device/type/authority inputs are rejected before any cost is reported
  • with_memory(Arc<dyn DeviceAllocator>) signatures are unchanged and exercised through an already-erased trait object
  • optional capabilities have no whole-allocation release operation

Intentional EP decommit migration

ExecutionProvider::decommit_allocation_range no longer inherits a successful Ok(0) no-op. Partial decommit has no eager equivalent, so an EP without VirtualBacking now 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 on ExecutionProvider::deallocate / DeviceAllocator.

This migration is also recorded in docs/ep-plugin/EP_PLUGIN_EXPORT_INVENTORY.md.

Accounting and release evidence

  • eager fallback allocates and accounts for the complete requested footprint; partial decommit without VirtualBacking is explicitly unsupported
  • CUDA provider-managed mapped growth uses one eligibility predicate everywhere: construction-selected in-tree CudaVmmAllocator + explicit governed signal + discovered VirtualBacking; arbitrary injected capabilities never enter the concrete capacity-token path
  • VMM mapped-byte estimation and commit report physical granularity; the disjoint-tiny-range test now obtains expected granularity independently from cuMemGetAllocationGranularity
  • shared-prefix physical bytes remain charged once; valid added mappings report zero owned bytes, while wrong-device, wrong-authority, foreign-type, and foreign zero-byte prefixes are rejected before cost reporting
  • terminal release remains DeviceAllocator::deallocate[_with_unmapped] through the owning EP, including allocations reserved through VirtualBacking
  • eager batched commit validates every buffer device before taking the no-capability fallback
  • out-of-crate tests implement and independently discover VirtualBacking and SharedMapping through Arc<dyn DeviceAllocator> without allocator Any, 'static, or identity-token requirements

Validation

Head: d120140a6de480e62c01163f8b2fc6821dceead5

  • affected all-target check: passed
  • targeted tests: 550 passed, 0 failed
  • CUDA hardware cases: 56 compiled and ignored on this non-CUDA host, including VMM, prefix-sharing, construction selection/injection, device/authority rejection, and driver-granularity cases
  • changed-file rustfmt: passed
  • scoped Clippy -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
  • strict rustdoc -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 failures
  • git diff --check: passed
  • independent follow-up review: no significant findings
  • final added-line scan found no Phase 3/4 or prohibited identity/release apparatus; virtual_memory.rs is unchanged

Existing base failures

  • workspace cargo fmt --all -- --check still flags only unchanged onnx-genai-server/src/routes/completions.rs
  • unallowlisted broader Clippy scopes remain blocked by unchanged CPU/CUDA kernel warnings (accelerate_gemm.rs, matmul.rs, device_argmax.rs, matmul_nbits.rs, standard_attention.rs, optimizer.rs)
  • broader strict rustdoc remains blocked by pre-existing private/unresolved links in governor, EP API, CUDA EP, ORT, and session documentation; the corrected capability crates are clean under -D warnings

justinchuby and others added 2 commits August 18, 2026 08:36
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

codecov Bot commented Aug 18, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 45.24422% with 213 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.99%. Comparing base (9e0e20a) to head (d120140).

Files with missing lines Patch % Lines
...ates/onnx-runtime-cuda-memory/src/vmm_allocator.rs 0.00% 124 Missing ⚠️
crates/onnx-runtime-memory-api/src/capability.rs 54.97% 77 Missing ⚠️
crates/onnx-runtime-memory-api/src/allocator.rs 85.71% 12 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@                              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     
Flag Coverage Δ
mlas 85.19% <ø> (?)
offline 79.89% <45.24%> (-0.74%) ⬇️

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

Files with missing lines Coverage Δ
crates/onnx-runtime-ep-api/src/provider.rs 63.99% <100.00%> (+1.89%) ⬆️
crates/onnx-runtime-memory-api/src/lib.rs 100.00% <100.00%> (ø)
crates/onnx-runtime-memory-governor/src/lib.rs 83.26% <ø> (-0.06%) ⬇️
...s/onnx-runtime-memory-governor/src/shareability.rs 96.44% <ø> (ø)
...ates/onnx-runtime-session/src/executor/bindings.rs 59.50% <ø> (ø)
crates/onnx-runtime-memory-api/src/allocator.rs 86.95% <85.71%> (-13.05%) ⬇️
crates/onnx-runtime-memory-api/src/capability.rs 54.97% <54.97%> (ø)
...ates/onnx-runtime-cuda-memory/src/vmm_allocator.rs 10.72% <0.00%> (-0.21%) ⬇️

... and 12 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

Copy link
Copy Markdown

🔴 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/medium_generic_bf16_threads=8/32x512x512 438.67 µs 1.54 ms +250.1%
🔴 matmul/small_generic_bf16_threads=8/1x256x256 32.48 µs 111.71 µs +244.0%
🔴 matmul/medium_generic_f16_threads=8/32x512x512 40.26 µs 103.74 µs +157.6%
🔴 matmul/medium_generic_f32_threads=8/32x512x512 922.34 µs 2.20 ms +138.7%
🔴 block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1 64.61 µs 139.70 µs +116.2%
🔴 matmul/medium_generic_f32_threads=1/32x512x512 2.31 ms 4.52 ms +95.7%
🔴 matmul/small_generic_f32_threads=8/1x256x256 55.92 µs 101.10 µs +80.8%
🔴 matmul/small_generic_f16_threads=8/1x256x256 31.11 µs 49.34 µs +58.6%
🔴 block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1 431.01 µs 638.34 µs +48.1%
🔴 matmul/small_generic_f16_threads=1/1x256x256 30.69 µs 43.48 µs +41.7%
🔴 block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024 88.89 µs 117.73 µs +32.5%
⚠️ matmul/medium_generic_bf16_threads=1/32x512x512 600.86 µs 775.82 µs +29.1%
⚠️ matmul/medium_generic_f16_threads=1/32x512x512 38.26 µs 48.54 µs +26.9%
⚠️ gather/large_bf16_threads=1-internal/131072 11.64 µs 13.86 µs +19.1%
⚠️ matmul/large_generic_bf16_threads=8/32x1024x1024 1.73 ms 2.05 ms +18.6%
✅ matmul/small_generic_bf16_threads=1/1x256x256 32.00 µs 36.26 µs +13.3%
✅ add/large_f16_threads=1-internal/4194304 1.75 ms 1.98 ms +13.0%
✅ matmul/large_generic_f32_threads=1/32x1024x1024 11.71 ms 13.09 ms +11.8%
✅ sampling_latency/greedy_per_token 3.34 µs 3.62 µs +8.3%
✅ gather/large_f16_threads=1-internal/131072 11.70 µs 12.38 µs +5.8%
✅ qwen3_sampling_processors/top_p_fast_after_top_k 518.54 µs 538.01 µs +3.8%
✅ qwen3_sampling_processors/top_k_top_p_full_sort_baseline 5.61 ms 5.80 ms +3.3%
✅ sampling_latency/top_p_per_token 399.87 µs 409.38 µs +2.4%
✅ sampling_latency/top_k_per_token 59.58 µs 60.35 µs +1.3%
✅ grammar_masking/llguidance_compute_mask/32 74.14 µs 74.87 µs +1.0%
✅ qwen3_sampling_processors/top_k_top_p_fast 656.00 µs 659.72 µs +0.6%
✅ gather/small_bf16_threads=1-internal/4096 475.3 ns 477.3 ns +0.4%
✅ logit_processing/seven_processor_chain_per_step 316.59 µs 317.55 µs +0.3%
✅ reduce_mean/large_f32_threads=1-internal/262144 984.42 µs 986.94 µs +0.3%
✅ gather/medium_f32_threads=1-internal/32768 3.90 µs 3.89 µs -0.2%
✅ add/large_bf16_threads=1-internal/4194304 1.67 ms 1.67 ms -0.2%
✅ reduce_mean/small_f32_threads=1-internal/4096 15.17 µs 15.10 µs -0.5%
✅ qwen3_sampling_processors/top_k_partial_selection 159.24 µs 158.15 µs -0.7%
✅ qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline 3.94 ms 3.91 ms -0.7%
✅ gather/small_f32_threads=1-internal/4096 674.8 ns 668.0 ns -1.0%
✅ matmul/large_generic_bf16_threads=1/32x1024x1024 2.28 ms 2.25 ms -1.5%
✅ gather/medium_bf16_threads=1-internal/32768 2.42 µs 2.39 µs -1.5%
✅ gather/small_f16_threads=1-internal/4096 491.5 ns 479.9 ns -2.4%
✅ reduce_mean/medium_f32_threads=1-internal/65536 251.92 µs 245.69 µs -2.5%
✅ gather/medium_f16_threads=1-internal/32768 2.44 µs 2.37 µs -2.7%
✅ matmul/large_generic_f16_threads=1/32x1024x1024 104.13 µs 100.96 µs -3.0%
✅ add/large_f32_threads=1-internal/4194304 742.78 µs 708.33 µs -4.6%
✅ gather/large_f32_threads=1-internal/131072 31.40 µs 29.88 µs -4.9%
✅ add/medium_bf16_threads=1-internal/262144 108.85 µs 103.32 µs -5.1%
✅ sampling_latency/min_p_per_token 240.77 µs 226.16 µs -6.1%
✅ matmul/small_generic_f32_threads=1/1x256x256 38.83 µs 36.44 µs -6.2%
✅ qwen3_sampling_processors/top_k_full_sort_baseline 2.59 ms 2.41 ms -6.6%
✅ kv_cache/alloc_dealloc_pages 41.13 µs 38.37 µs -6.7%
✅ add/medium_f32_threads=1-internal/262144 26.85 µs 24.63 µs -8.3%
✅ add/small_bf16_threads=1-internal/1024 476.8 ns 436.0 ns -8.6%
✅ add/small_f16_threads=1-internal/1024 500.1 ns 456.6 ns -8.7%
✅ tokenization/encode_tokens_per_second 439.41 µs 398.40 µs -9.3%
✅ add/small_f32_threads=1-internal/1024 218.6 ns 195.4 ns -10.6%
✅ add/medium_f16_threads=1-internal/262144 118.26 µs 103.12 µs -12.8%
✅ tokenization/decode_tokens_per_second 7.59 ms 6.45 ms -15.0%
🟢 matmul/large_generic_f16_threads=8/32x1024x1024 148.87 µs 125.27 µs -15.8%
🟢 matmul/large_generic_f32_threads=8/32x1024x1024 7.00 ms 5.58 ms -20.2%
🟢 block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024 951.17 µs 749.46 µs -21.2%
🟢 block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024 92.31 µs 59.14 µs -35.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: { 2.85 4.20 5.64 }
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 added a commit that referenced this pull request Aug 22, 2026
> **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
@justinchuby

Copy link
Copy Markdown
Owner Author

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 d120140a6 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant