Repository navigation
Make KV tier migration a real storage operation - #742
Merged
Merged
Conversation
Replace mutable page residency labels with store-owned residency and explicit byte-copy migration. Keep GPU residency documented as host-backed emulation until CudaPageStore lands. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #742 +/- ##
==========================================
+ Coverage 79.94% 80.94% +1.00%
==========================================
Files 330 331 +1
Lines 127929 134104 +6175
Branches 127929 134104 +6175
==========================================
+ Hits 102272 108557 +6285
+ Misses 21478 21285 -193
- Partials 4179 4262 +83
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
|
Roll back failed CoW allocations and any incidental eviction, report migration replacement bytes as an unleased transient, and prevent best-effort prefetch from exceeding hot capacity after demotion failure. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Reserve RAII sibling leases from the pool's existing governor authority before replacement stores or rollback snapshots are allocated. Refuse migration pressure without state changes and release transient grants on every path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Reserve the maximum source, free-page, eviction-victim, and replacement transient before cloning. Keep the guard alive until every snapshot drops, and make direct Page migration private so accounting cannot be bypassed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Split persistent page growth from transient snapshots and absorb successful growth into the pool lease. Rollback releases growth only after the allocated page is removed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Owner
Author
|
Five review rounds closed the accounting and transactional boundaries before merge. What changesKV tier migration is now a real storage operation:
Bugs found and fixed during review
Final tests cover free-page and no-free-page branches, repeated growth, copy failure, directional migration failure, exact ledger peak/release, mixed non-zero storage components, and sabotage of reserve enforcement. Final review found no significant issue. Remaining limitationBoth CPU and emulated GPU locations still use Merging. |
justinchuby
added a commit
that referenced
this pull request
Aug 24, 2026
Four doc sites still told a reader that a `CudaPageStore` was on its way and that the host-backed tiers were *"emulation until Stage 3"*. #721's own body has said otherwise since 2026-08-18: | stage | status | |---|---| | 1. Extract `KvPageStore` + `HostPageStore` | **merged** (#726) | | 2. Make tier migration real, charged to the governor | **merged** (#742) | | 3. `CudaPageStore` backed by the VMM arena | **superseded** | | 4. Native CUDA decode consumes `PageTable` | **superseded** | On native CUDA, device KV paging is owned by the **VMM layer** — `CudaVmmAllocator` with its physical-handle pool (#740), committed-granule admission (#745) and growth grants (#748). It already reserves a stable virtual range and maps committed granules behind it. A second page allocator inside `onnx-genai-kv` would **duplicate that ownership rather than complete it**. ## Why this is worth its own commit Comments are where a cancelled plan survives longest. *"until Stage 3"* reads as a scheduled prerequisite, so the next person to open the file either builds a dead stage or spends an afternoon working out that they shouldn't. The issue body was updated; the code that people actually read was not. ## What changed | site | before → after | |---|---| | `page_table.rs:19` `DevicePageSpan` | *"a future `CudaPageStore` … would reserve a stable virtual range"* → the VMM layer already does exactly that; the span stays because the trait boundary is what keeps `head_token_row()` off the store contract and what lets an out-of-tree store exist at all | | `page_table.rs:94` `KvPageStoreFactory` | *"Stage 3 can supply a factory"* → an out-of-tree extension point, not an in-crate hook | | `tiered.rs:1` module doc | *"until Stage 3 supplies a CUDA store"* → host-addressable tiers are this crate's **end state**, not a waypoint | | `paged_cache.rs:658` `preempt_sequence` | *"host-backed emulation until Stage 3"* → declared residencies that stay declared | | `MEMORY_ARCHITECTURE.md:73` | *"Stages 2–5 in #721"* — names a stage 5 that does not exist and implies three open ones → the real per-stage status, pointing at the issue body as authoritative | That last row already said device KV paging belongs to the VMM layer, so this only stops the two halves of one table cell contradicting each other. ## Deliberately left alone `copy_from_host`'s *"a future device store can implement this with an explicit host-to-device transfer"* is **not** the same claim and stays. An out-of-tree device store remains possible and supported; what is not coming is one **in this crate**, on this repository's roadmap. ## Validation Doc comments and one doc table cell. No behaviour change — `cargo check -p onnx-genai-kv` and `cargo doc -p onnx-genai-kv --no-deps` both clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
KvPageStore, removing mutablePage::deviceCudaPageStorewithout cache-caller changesPage::migrateprivate so all production migration is table-mediated and cannot bypass leasingBoth hot and cold stores remain host-addressable
HostPageStoreinstances.Device::Gpu(0)is still an emulated residency label until Stage 3 addsCudaPageStore; this PR makes migration semantics and ownership real and replaceable, not device memory.Validation
cargo fmt --all -- --checkcargo test -p onnx-runtime-memory-governor -p onnx-genai-kv(140 KV unit tests plus integrations; 28 governor tests)cargo clippy --locked --all-targets -p onnx-runtime-memory-governor -p onnx-genai-kv -- -D warningscargo test -p onnx-genai-engine --features native-backendcargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings.migrate(search finds only the private helper invocation insidePageTable::migrate_pagepage.device =writes; field removedStage 2 of #721.
Final CoW persistent-accounting fix