Repository navigation
Pool CUDA VMM physical handles across reservations - #740
Merged
Merged
Conversation
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #740 +/- ##
==========================================
- Coverage 80.77% 79.86% -0.91%
==========================================
Files 331 331
Lines 134104 134786 +682
Branches 134104 134786 +682
==========================================
- Hits 108318 107643 -675
- Misses 21523 22883 +1360
+ Partials 4263 4260 -3
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
|
justinchuby
marked this pull request as draft
August 7, 2026 18:50
Owner
Author
|
Follow-up on the production-wiring review: I attempted the shared registry and production I stopped and discarded that change because correct governor accounting is not representable with the current abstractions:
Therefore production wiring would either create a second set of books or falsely report unaccounted/free memory. Per #736's instruction, I am not adding a private ledger or shipping unreachable production wiring. This PR is marked draft pending a larger accounting change: a process/device governor (or an explicit stable shared accounting authority/token) that all engines on that device use. Once that exists, the tested registry/wiring is straightforward. |
This was referenced Aug 7, 2026
Closed
Owner
Author
|
Stage 0 authority contract is now in #741. It adds stable device-keyed Production status is explicit: clones/holders inside one |
justinchuby
added a commit
that referenced
this pull request
Aug 7, 2026
## Summary - add stable `MemoryAuthorityId` identity to `MemoryGovernor`, keyed by `DeviceKey` - make `LeaseLedger` own that identity so every governor wrapper over one ledger agrees - add explicit buffer-owned vs backing-owned physical accounting to `VirtualBacking` - reject backing/governor authority mismatches before address reservation or physical commit - document that pooled-unmapped bytes remain physically charged while mapped zone bytes are attribution only ## Production identity wiring Production constructors that now carry/reuse identity: - `LeaseLedger::new` / `LeaseLedger::new_for_device` create the authority identity owned by that ledger. - `LedgerGovernor::new` reads the ledger-owned identity, so every wrapper/clone over one ledger agrees. - `EngineResourceGovernor::new_with_capacities` creates one ledger; its `ModelMemoryPlan`, admission ceiling, and engine lease holders reuse wrappers/clones over that ledger and therefore share identity **within that engine**. - `HostGovernor::new` stores one machine-host identity; every `HostLeaseGovernor::new` adapter over that same `HostGovernor` reuses it. - the ORT test/support `RecordingGovernor` delegates identity to its wrapped `LedgerGovernor`. Production constructors still creating unique authority instances: - every separate `EngineResourceGovernor::new` / `new_with_capacities` call constructs a new `LeaseLedger`, so multiple engines targeting the same CUDA device still have distinct authority IDs and distinct accounting books. - independent `LeaseLedger::new` calls are intentionally distinct, even for the same `DeviceKey`. **Scope boundary:** this PR is necessary identity/compatibility infrastructure, but it does **not** unblock #740's cross-engine process/device physical-handle pool wiring and does not make the global pool production-ready. The explicit remaining step is to construct/inject multiple engines for one physical device compatibility domain from one shared device-scoped `DeviceGovernor`/ledger/authority. Only after those production engine constructors deliberately reuse that authority can #740 attach one process/device pool and charge retained-unmapped handles correctly. ## Accounting contract - backing-owned physical bytes are charged once to their named authority - pooled-unmapped bytes remain charged while no holder maps them - mapped holder/zone accounting is attribution, not another physical lease - a `VirtualBuffer` cannot silently combine backing authority A with governor B ## Validation - workspace search found all four `MemoryGovernor` implementations: `LedgerGovernor`, `HostLeaseGovernor`, third-party contract test `MyOwnGovernor`, and ORT test/support `RecordingGovernor` - `cargo test --locked -p onnx-genai-ort --no-run` - `cargo test --locked -p onnx-genai-cli --no-run` (Windows CLI ORT-equivalent test-target compilation) - `cargo test --locked -p onnx-genai-ort governed_allocator::tests -- --test-threads=1` (20 passed) - `cargo fmt --all -- --check` - `cargo clippy --locked --all-targets -p onnx-runtime-memory-governor -- -D warnings` - `cargo clippy --locked --all-targets -p onnx-runtime-virtual-memory -- -D warnings` - `cargo test --locked -p onnx-runtime-memory-governor` - `cargo test --locked -p onnx-runtime-virtual-memory` - targeted scheduler authority-sharing test - sabotage: removed the mismatch return, observed `backing_accounting_rejects_a_different_authority_before_reservation` fail, restored the guard, and observed it pass Stage 0 for #736. Follow-up prerequisite infrastructure for draft #740; not production pool wiring. --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 8, 2026
## Summary - add a server-owned, concurrency-safe device-authority map keyed by backend/device compatibility domain - inject the shared authority into production engine and pipeline construction while keeping host/disk ledgers engine-local - report VRAM used/limit/headroom directly from the shared device ledger - reject per-model device limits that conflict with the server authority before model allocation ## Production wiring The same server-owned authority provider is used by: - `--model` single-model startup - `--models-dir` eager loads - `--models-config` eager loads - first-request lazy loads - admin model loads/reloads - pipeline package construction through those same paths All of those paths converge on `ModelRegistry` and `build_handle_with_authorities`. Concurrent loads serialize authority creation with the server-owned map, not a process-global singleton. Standalone callers remain backward compatible: `Engine::from_dir`, `Engine::from_dir_with_session_options`, `Engine::from_pipeline_dir`, and `PipelineEngine::from_dir*` without an injected provider create a unique device authority as before. New provider-taking constructors opt into sharing. ## Lifecycle and metrics - each engine owns only its leases; unloading it drops those leases - the registry-owned authority survives individual model unloads - the final registry/authority drop releases the ledger - `/metrics` and `/v1/resources` overlay VRAM used/limit/headroom from the one live shared ledger rather than summing model snapshots - host and disk ledgers remain private per engine ## Tests - shared headroom, combined-limit refusal, and per-lease release - production two-engine server path shares authority and aggregate metrics ledger - unloading one engine releases only its leases - conflicting limits name both values - concurrent creation returns one authority - different device keys return different authorities ## Validation - `cargo fmt --all -- --check` - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` - `cargo clippy --locked --all-targets -p onnx-genai-server -- -D warnings` - targeted engine/server authority and production-load tests - `cargo test --locked -p onnx-genai-cli --no-run` Prerequisite for #736 and draft #740. Builds on #741. This PR does not implement `PhysicalHandlePool`. --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Share physical handle pools by CUDA compatibility and memory authority, propagate production authorities through CUDA EP/native pipeline construction, and prove cross-reservation reuse through the production allocator path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
force-pushed
the
squad/736-global-physical-handle-pool
branch
from
August 8, 2026 20:35
c742600 to
87823c9
Compare
Unify standalone pipeline authority, synchronize pooled unmaps, trim retained handles during limit shrink, and route public standalone CUDA providers through pooled VMM construction. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Owner
Author
|
Pushed
Validation completed:
The new public-provider/in-flight GPU proof compiled but could not execute on this host because |
Synchronize reservation teardown before pooled reuse, block claims during limit reconfiguration, and apply server multi-device limit changes atomically. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Owner
Author
|
Pushed Fixed the final three blockers:
Passed:
Still environment-blocked: |
Validate every authority's pooled-unmapped trim capacity before releasing any handles, and keep limits and policy unchanged when preflight or commit cannot complete. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Owner
Author
|
Pushed Multi-device shrink is now true two-phase under deterministic authority gates:
Targeted validation:
|
Serialize pool checkout, return, teardown, snapshot, and trim with authority-scoped read/write gates held across multi-device reconfiguration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Owner
Author
|
Pushed The remaining preflight race is closed with authority-scoped pool read/write gates:
Targeted barrier GPU test passed: checkout and return both block behind the transaction writer; failed/no-trim transaction preserves the pooled handle, while successful trim makes the later checkout create a new handle rather than observe stale preflight state. Also passed:
|
justinchuby
marked this pull request as ready for review
August 9, 2026 00:28
Owner
Author
|
Final merge-safety review passed at What landsA process/device/authority-scoped pool of granule-sized CUDA VMM physical handles:
GPU proofAcross separate reservations A → B:
Public provider construction and teardown tests execute on the real CUDA path. Provider teardown synchronizes before reuse; sync failure quarantines rather than reuses possibly-live handles. Review bugs closed
Final design uses authority-scoped read/write gates: normal checkout/return/teardown use shared access; multi-device reconfiguration acquires exclusive gates in deterministic order and holds them through snapshot, trim and commit. Atomic counters remain observability only, never synchronization. What remains for #736
Merging. |
justinchuby
added a commit
that referenced
this pull request
Aug 9, 2026
## Summary - add transactional VMM span admission with separate mapped-zone and global-owned constraints - coordinate weight-zone allowances through the shared memory authority, so multiple caches cannot each claim the device - distinguish pooled handle reuse (0 new owned bytes, one mapped granule) from handle creation (new owned bytes) - force-release transaction-created handles on rollback while returning borrowed handles to the pool - evict by actual unmapped bytes, bound zero-progress loops, and expose content/mapped/owned metrics - avoid double-charging the production VMM pool with the temporary content-sized startup reservation ## Validation - `cargo fmt --all -- --check` - exact CUDA EP CI clippy: `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` - CUDA-memory, memory-governor, engine, CLI, and server clippy - memory-governor, CUDA-memory, and weight-paging tests - real CUDA production-pool rollback/admission test - real CUDA VMM weight residency test - release CUDA CLI build ## 14B production-pool E2E Same branch binary, `C:\Users\justinchu\dev\models\qwen14b-zp`, native CUDA, greedy, explicit `--vram-limit 6000000000`, VMM and production pooling enabled: - output: `The capital of France is Paris.` (EOS, 8 generated tokens) - derived weight mapped allowance: **4,389,387,264 B** after preserving **1,610,612,736 B** for KV/recurrent state - content resident: **3,875,440,640 B** - mapped weight physical: **3,938,451,456 B** - total authority-owned physical: **5,414,846,464 B**, below the 6 GB hard limit - page-ins/hits/evictions: **6098 / 838 / 4533** (12.08% hit rate; reported honestly at the stricter physical budget) - weight runtime allocation time: **0 ms**; VMM peak committed: **5.0 GiB** - shutdown reported the arena at **0 mapped granules**; the pre-shutdown owned metric remained above mapped bytes, truthfully charging retained pool ownership until authority teardown This is an incremental admission/attribution change; full zones and broader lifecycle work remain. Refs #736. Extends #740 and preserves deterministic counters from #723. --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
This was referenced Aug 10, 2026
justinchuby
added a commit
that referenced
this pull request
Aug 11, 2026
## Summary A **truth pass** over the memory / offload / KV / capture documentation. Over the past days a series of hardware measurements (RTX 4060 Laptop GPU, Windows/WDDM) overturned several claims still asserted as fact in `docs/`. This PR finds those claims, corrects them against what was measured, and — where a document described an intended design in the present tense as though it were implemented — marks it as a design proposal with its current status. No wrong claim is silently deleted: each is replaced with what was measured and the PR/issue that superseded it, so a reader who remembers the old number can see it was tested and retired. **Docs-only. No code or tests changed.** ## Corrected claims | File | What it said | What is true | Superseded by | |---|---|---|---| | `ORT2.md` §30.4–30.8 | "Offloading doesn't mean stalling — prefetch overlaps with compute"; "If compute time > transfer time: zero visible overhead"; "Overlap compute + transfer: **Always**"; `// Async prefetch next layer's weights (overlaps with current compute)` | Sequential weight prefetch produced no usable overlap and was **deleted**; the bottleneck is disk/transfer bandwidth, not overlap. On a model that doesn't fit there isn't enough independent compute to hide the transfer behind. | PR #715; AirLLM analysis #718 (rates prefetch ~10%) | | `ORT2.md` §30.6 / §30.7 | "30% weights on CPU → prefetch mostly hidden by compute"; `memory.prefetch_layers` presented as an active knob | Estimate, not measured; prefetch removed. Managed offload was ~27× slower than doing nothing (WDDM) on a non-fitting model. | PR #715 | | `DESIGN.md` §39.5 | KV table "**tokens per granule**, 8B GQA (2048 B/token): 2 MiB → 1024"; floor `num_layers × 2 × 2 MiB` | KV is per-binding/head-major: a fixed full-context stride commits **objects × granule**, not tokens/granule. qwen2.5-0.5b commits **96 head-stripes × 2 MiB = 192 MiB for ~12 KiB** live. Crossover is `granule/(head_dim × sizeof(dtype))`, **model-size independent** (~8,192 tok @ head_dim 128/fp16; ~16,384 @ 64/fp16); quantized KV **doubles** it. 2 MiB is min **and** recommended — no finer granule on this device. | PR #776 | | `MEMORY_ARCHITECTURE.md` | h2d attribution 1.7% (already partly recorded); KV geometry, crossover, read-bounded commit, offload/capture exclusion not consolidated | Re-measured h2d transfer **18.8%**. Added the central KV reservation-geometry section: per-object granule floor, the crossover formula, 2 MiB min==recommended, the read-bounded-commit result (read—not reservation—forces the commit; #721 stage-4 32× regression = 1.5 GB vs 48 MB), offload/capture mutual exclusion, and prefetch deletion. | PR #715, #716, #772, #776 | | `MEMORY_ARCHITECTURE.md` (KV page storage row) | "still host-only" without VMM ownership | Device KV paging is owned by the CUDA VMM layer (`CudaVmmAllocator`, #740 pool), with committed-granule admission (#745) and growth grants (#748); `native_decode/cuda.rs` has no `PageTable` consumer and `paged_gqa.rs` is a batch-1 CPU primitive. | #740, #745, #748 | | `WEIGHT_OFFLOAD.md` §7.3, Phase 4 | "Prefetch the next exact/predicted expert wave … while the current wave computes"; Phase 4 predictive prefetch as the plan | Marked as a design proposal whose central premise (prefetch overlap) is measured false so far; plumbing (stream-ordered H2D, fences, double-buffer scheduler) is distinct from prefetch producing a throughput gain. | PR #715, #718 | | `WEIGHT_OFFLOAD.md` §0.4 | Did not mention offload/capture exclusion, the residency-cache defect, or committed-physical budgeting | Added: offload disables capture today (#716 is the fix; #708/#728 = 154→34 segments); residency cache had a **0% hit rate for its whole life**, fixed to 74.18% with evictions 6,286→0 (#723); a nominal 6 GB budget physically consumes **~6.51 GB**. | #716, #723 | | `CUDA_GRAPH_CAPTURE.md` | No mention of offload interaction or VMM remap survival | Added: offload/capture mutual exclusion (#716); a captured graph **survives a VMM remap at the same VA**, writing into the new physical pages — sentinel-proven; unmapping during in-flight replay and `cuMemMap`-during-capture are **not** proven. | #716, #727 | ## Deliberately left unchanged `GQA_KV_MATERIALIZATION_DESIGN.md`, `KV_INSERTION_DESIGN.md`, `DECODE_POOL_RESIDENCY_DESIGN.md`, `NATIVE_CUDA_DECODE.md`, `CUDA_STRATEGY.md`, `native-ort-kv-capacity.md` — checked against the relevant facts and found accurate/self-consistent for their scope (CPU-kernel data movement, CPU decode-pool threading, the native-decode milestone plan, the KV-capacity policy). Correct docs were not restructured, and subject areas outside this pass (CPU EP, quantization formats, CLI, diffusion) were not touched. ## Preventing recurrence Added a documentation-hygiene note to `MEMORY_ARCHITECTURE.md`: a performance claim in docs must cite the measurement that produced it; when measurement overturns a claim, correct it in place and name what superseded it rather than deleting it; and keep measured fact, design intent, and not-yet-built visibly distinct. Refs #715, #716, #718, #721, #723, #727, #740, #745, #748, #772, #776. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 11, 2026
… forces commit, head-major floor stands (#772) ## What & why The owner's directive for #750: give each sequence one **stable, flat, contiguous virtual address** for its whole context so the existing flat GQA attention kernel keeps reading a plain `[heads, seq, head_dim]` buffer while KV is physically scattered and grown on demand — paging invisibly underneath, the kernel never learning page tables. If that holds, batching is just N sequences each with its own stable contiguous VA and the flat kernel works unchanged. This PR **confronts the known blocker first** (#721 stage 4: a stable full-context VA committed **1.5 GB** vs bucket growth's **48 MB**, a 32x regression, because the decode kernel read the full padded shape) and answers the crux with **evidence on this hardware**. It is primarily a **truthful negative for the fixed-stride variant** plus a confirmation that the already-landed bucket-stride path is the correct realization. **No engine behavior change — isolating GPU test + docs only. Do not merge.** ## The crux, answered by an isolating GPU test `crates/onnx-runtime-cuda-memory/tests/vmm_kv_contiguous_tail_gpu.rs` (3 tests, pass on device; measured granule **2 MiB**): 1. **`bounded_read_of_contiguous_va_leaves_tail_granules_uncommitted`** — a read bounded to the committed live length, *including a captured graph replay* (the decode hot path), succeeds while the reservation tail stays physically **uncommitted**. The flat contiguous VA is real and its tail is physically absent. 2. **`padded_read_into_uncommitted_tail_faults_forcing_commit`** — the *same* buffer, read **one byte past** the live length into the uncommitted tail, faults `CUDA_ERROR_INVALID_VALUE` (synchronous, non-sticky). This is **#721 stage 4 reproduced in isolation**: the **read pattern, not the reservation, forces the commit**. A kernel that reads the full padded shape dereferences the tail, so the tail can never be decommitted. 3. **`head_major_full_context_stride_pays_one_granule_per_object`** — even a perfectly bounded read cannot make a **fixed full-context stride** cheap: KV is per-binding/head-major, so each head's valid prefix lands in its own granule and the commit floor is `objects x granule`. Reproduces the documented qwen2.5-0.5b floor on hardware: **96 head-stripes × 2 MiB = 192 MiB committed for ~12 KiB of live content.** ## Committed-bytes verdict (the number that decides adoption) | Scheme | Committed bytes | Re-capture on growth? | |---|---|---| | Bucket growth (today) | ~48 MB (baseline) | yes | | **Fixed full-context stride** (the literal "one flat VA, never re-strided") | **192 MiB near-empty / 1.5 GB @32k — 32x worse** | no | | **Bucket stride under stable VA** (landed `kv_commits_on_demand`, #682/#740/#748) | **parity with bucket growth**, + stable VA | **still yes** | Verified on real qwen2.5-0.5b (`native_cuda_vmm_kv_grows_in_place_and_commits_more_granules`): full-context VA reserved, only the initial **bucket** committed at load (`< full/2`), growth keeps `device_ptrs` **stable** (contiguous VA) and commits exactly the **bucket delta**. **Conclusion:** a length-bounded kernel is **necessary but not sufficient**. Low committed bytes *and* no re-capture cannot both hold under the head-major layout: the fixed stride that removes re-capture pays the granule floor (loses 32x), and the bucket stride that keeps committed bytes at parity still re-strides/re-captures on growth. The GQA op already takes `seqlens_k`/`total_sequence_length` so its *compute* is length-bounded; the floor is set by the layout stride, not the compute bound. Full write-up: **`docs/VMM_KV_CONTIGUOUS_VA.md`**. ## Does `onnx-genai-kv`'s page machinery still matter, and who owns paging? Direct answer to the owner's question: - The **native CUDA decode path has no `PageTable`/`PagedKvCache` consumer** (`native_decode/cuda.rs` references neither; `paged_gqa.rs` is a batch-1 **CPU** primitive). The KV page store is **host-only** (`docs/MEMORY_ARCHITECTURE.md`: "never touched device memory"). - Under the VMM contiguous-VA design, **device KV paging is owned by the CUDA VMM layer** — `CudaVmmAllocator` (#740 pool) mapping granules under a fixed per-binding reservation, with governor grants (#748) and committed-granule admission (#745). The kernel sees a flat contiguous buffer and never learns page tables. - **Implication:** for the CUDA batch path, `onnx-genai-kv`'s page machinery is **not needed** as the device pager — a batch is N sequences each with its own stable VMM reservation, flat kernel unchanged. `onnx-genai-kv` stays relevant for the **CPU** paged-GQA path and host-side concerns (page store, prefix sharing/CoW, quant layout). ## Next experiments (not in this PR) 1. **Commit below the bucket** — since GQA is length-bounded, only `ceil(logical_len)` granules per head need backing. Open question the isolating test frames: does the ORT GQA **CUDA** kernel / graph capture ever touch the bucket tail `[logical_len, bucket)`? If not, sub-bucket commit beats bucket growth. Needs an engine change + byte-identical decode run. 2. **Seq-major KV layout** — the only layout giving both low commit and no re-capture, but requires re-exported models / a custom kernel. ## Validation - `cargo fmt --all -- --check` ✓ - `cargo clippy -p onnx-runtime-cuda-memory --features gpu-tests --all-targets -- -D warnings` ✓ - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` ✓ (CI gate, no `--all-targets`) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` ✓ - New GPU tests (3) + existing `vmm_graph_remap_gpu` (4) pass on device. - Real-model measurement: `native_cuda_vmm_kv_grows_in_place_and_commits_more_granules` passes on qwen2.5-0.5b. Refs #750, #721. Does not close any issue. **Do not merge.** Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 11, 2026
…sured (#750, #777, #783) (#787) ## Token-major-across-all-layers KV: viability probe A **measurement probe with one open risk settled** — no production kernel converted this round (microbenchmark + one isolating GPU floor test only, per scope). References #750 (multi-request batching), #777 (page-level prefix sharing), #783 (KV layout as a stride descriptor). Builds on the merged #772 head-major floor, #776 granularity probe, and the landed #782 seq-major decode pair / #778 investigation. Device: **RTX 4060 Laptop (sm_89, Ada), 8 GiB**, driver 591.55, CUDA 13. ### The question Under token-major, reading one layer's K strides by the **full per-token KV size** (~192 KB for qwen14b) instead of ~2 KB. Each head's run stays a 256-byte contiguous segment (#778 measured that as fine), so the risk is **not coalescing — it is TLB pressure**: each run lands on a different page. ### Result 1 — TLB risk is cleared (`bench-tokenmajor/`) Stride-only bandwidth sweep (identical bytes/occupancy/run; only the token stride changes via a plane-group parameter). Median GB/s, `ratio` = stride ÷ the 2 KB stride at the same L: | qwen14b L | working set | 2 KB | **192 KB (token-major)** | ratio | |---|---|---|---|---| | 8192 | 1.5 GiB | 206.3 | 206.1 | 0.999 | | 32768 | **6.0 GiB** | 207.1 | **207.1** | **1.000** | At the realistic working sets the 192 KB stride is **dead-flat**. Reads are DRAM-bandwidth-bound (~207 GB/s ≈ 80% of peak) *independent of stride* → **2 MiB device pages cover the 192 KB stride; the token-major read is free.** (`ncu` is not installed on this box, so the stride-only BW curve is the deterministic TLB proxy — every other variable held constant.) ### Result 2 — the floor payoff, measured (`vmm_kv_token_major_floor_gpu.rs`) Through the **real #740 authority-scoped pool via `carve()`** (no second allocator, no per-sequence physical reservation), reporting committed **physical** bytes: > head-major: 768 objects × 2 MiB granule = **1536 MiB** committed for 196608 B of live content. Token-major commits **2 MiB** for the identical content — a **768× reduction**. One token-major reservation commits **one granule per sequence** (stays one granule up to 10 tokens = `floor(2 MiB / 192 KiB)`). ### Implementation cost & #777 (in the doc) - Recommend **`KvLayout` → stride descriptor (#783)**, not a third enum variant — token-major collapses the per-(layer,side) buffer boundary the two existing variants keep; a descriptor shares the kernel edit with seq-major instead of duplicating it. - Per-layer bindings (`engine/model_io.rs` positional vectors) must survive as **views into one reservation** — the real structural work. - #777: a shared prefix is **one contiguous multi-map across all layers** (vs `layers×2` under seq-major). ~**2.56 GiB** saved for a 2,000-token prompt at 8 concurrent on qwen14b, robust to short prefixes (sharing granule floor is per-sequence, not per-stripe). ### Recommendation **GO**, staged: descriptor + binding-views first (validate on the existing output-level parity oracle), decode-pair kernel flip after, then bank #777. The read path is free; the cost is on the binding layer. ### Validation - `cargo fmt -p onnx-runtime-cuda-memory -- --check` ✅ · bench workspace `cargo fmt -- --check` ✅ - `cargo clippy -p onnx-runtime-cuda-memory --features gpu-tests --all-targets -- -D warnings` ✅ - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` ✅ (baseline; unchanged) - Both GPU tests pass; bench runs the full sweep. > Do **not** merge; do **not** close any issue. This closes a real question either way — the answer happens to be positive. --- ### Update: vLLM-paging granularity comparison (section 6, verified) Corrected framing folded in — vLLM's real *per-sequence* allocation quantum is `block_size x all layers`, not the ~32 KB per-layer block: | | Per-sequence quantum | Kernel walks a page table? | Physical mem returnable? | |---|---|---|---| | vLLM paging | **~3 MiB** (16 tok x 48L, qwen14b) | **yes** | no — pool pre-committed | | **token-major + VMM** | **2 MiB** (~10.7 tok) | **no** | **yes** | Verified arithmetic (`granule / kv_bytes_per_token`): qwen14b **10.67** tok/granule, qwen0.5b **170.67**; vLLM = `block_size` (16). Internal fragmentation per sequence: token-major <= one 2 MiB granule (model-independent) vs paged <= one block (scales with model). Honest both ways — finer than vLLM on qwen14b (2 vs 3 MiB), coarser on qwen0.5b (2 vs 0.19 MiB). **Quantized KV doubles token-major tokens/granule (denser) while doubling head-major's crossover (worse)** — the asymmetry reinforces the quantized-KV direction. So on granularity token-major is **equivalent-or-finer than paging** while keeping flat kernels and on-demand release. --- ### Correction: quantization is a split result (verified fp16/fp8/int4, both models) Earlier "quantization is a uniform win for token-major" was half right. Two separate, verified claims: **1 (true, good) — tokens/granule auto-scale** (byte-denominated quantum, the 2 MiB hardware granule), no tuning: | qwen14b | fp16 | fp8 | int4 | |---|---|---|---| | token-major tokens/granule | 10.7 | 21.3 | 42.7 | **2 (true, against us) — the per-sequence BYTE floor stays fixed at 2 MiB**, while a paged (token-denominated, block=16) design's waste shrinks: | worst-case waste/seq | fp16 | fp8 | int4 | |---|---|---|---| | Paged (16 tok) | ~3 MiB | ~1.5 MiB | ~0.75 MiB | | **token-major + VMM** | **2 MiB** | **2 MiB** | **2 MiB** | So **under aggressive KV quantization a token-denominated paged design wastes fewer bytes per sequence than we do** — stated plainly (absolute waste ≤ one partial granule, ~16 MiB at 8 concurrent: a nuance, not a threat). **The claim that survives intact:** head-major makes quantization and residency **antagonistic** (crossover `granule/(head_dim·dtype)` doubles as bytes halve: 8192/16384/32768 tok on qwen14b), while **token-major makes them independent** (no crossover — floor is already one granule/sequence at any dtype). Also now marked **hardware-imposed** (2 MiB granule, #776 `MINIMUM==RECOMMENDED`) vs **design choice** (vLLM `block_size`; token-major's flat-VA decision), and internal fragmentation stated in the same units for both. TLB remains the primary risk — measured, ratio 1.000 at 6 GiB. --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 11, 2026
…robenchmark (#778) ## Seq-major (BSNH) KV layout — decision-critical investigation **Draft. Do not merge; closes no issue.** Investigation + measurement + two narrow microbenchmarks. No production kernel conversion in this round. Deliverable: `docs/SEQ_MAJOR_KV_INVESTIGATION.md` + standalone GPU benches under `bench-seqmajor/` (workspace-excluded; do not perturb the main build): `seqmajor_bench` (strided read) and `remap_burst` (real-VMM boundary spike). References #750, #759, **#776**, **#777**; builds on #522/#726 and #727/#740/#745. ### Owner corrections folded in (this revision) 1. **Seq-major shrinks the granule floor by `kv_heads`× — it does NOT remove it.** K and V stay separate per-layer bindings (verified: `model_io.rs` per-layer paired `key/value_cache_{inputs,outputs}`), so the floor unit goes `layers·2·kv_heads` → `layers·2` (qwen14b **768 → 96** granules), and the crossover becomes `granule/(kv_heads·head_dim·dtype)`: | model | head-major crossover | **seq-major crossover** | factor | |---|---|---|---| | qwen14b (kv8, hd128) | 8,192 tok | **1,024 tok** | **8×** = kv_heads | | qwen0.5b (kv2, hd64) | 16,384 tok | **8,192 tok** | **2×** = kv_heads | The seq-major crossover is therefore **model-dependent** (scales with kv_heads), not model-independent — big win on wide-GQA, modest on narrow. Token-major (1-granule floor) is *evaluated, not prototyped*: worth a future stride-sweep experiment given §2's cheap strided reads, gated on TLB/locality. 2. **Remap cost is a lockstep BURST, not a per-token tax — measured on HW** (`remap_burst.rs`, real `cuMemCreate`/`cuMemMap`/`cuMemSetAccess`): | scenario | N granules | measured burst | |---|---|---| | qwen14b head-major | 768 | **49.5 ms** | | qwen14b seq-major | 96 | **4.5 ms** | | qwen0.5b head-major | 96 | 4.9 ms | | qwen0.5b seq-major | 48 | 2.8 ms | Per-granule ~48–64 µs (below #776's 150 µs — that folded in first-touch); granule re-confirmed 2 MiB min==recommended. Seq-major cuts the single-token growth spike by `kv_heads`. Fully removable via **commit-ahead** (graph-safe because the VA is fixed — unlike bucket-growth realloc which forces re-capture) and/or **boundary staggering** (achievable; lockstep is not inherent). ### The three questions a go/no-go needs — answered 1. **ORT BSNH crux — definitively NO.** ORT GQA past/present KV is **BNSH-only on every dispatch path** (cuDNN-SDPA, FlashAttention, XQA, memory-efficient). `group_query_attention_helper.h` hardcodes `past_kv_format = Q_K_V_BNSH`; `CheckInputs` rejects `dim[1] != kv_num_heads` with a literal "BNSH Input" error. ⇒ Seq-major must be a **per-backend capability**: native seq-major, ORT BNSH, neither reads the other's KV bytes. Only a layout tag on the already-per-backend KV ownership contract; no unacceptable break. Per-step transpose explicitly rejected (full extra KV pass on the hot path). 2. **Strided-read falsifier — cleared (≤2%).** Fused online-softmax decode microbench, interleaved, sm_89, L ∈ {512…32K}. Seq/head ratio **0.80–1.02** everywhere. The `head_dim×dtype < 128 B` danger case does not materialize — including a **32 B contiguous run** (int4 KV), a quarter cache line: 0.83–1.00. Quantized KV is safe under seq-major. 3. **Prefix-sharing prize (#777, co-primary) — quantified & mechanism-confirmed.** qwen14b 2048-tok × 8 concurrent removes **2.62 GiB**; capacity **3 → 8 at 2 GiB, 10 → 30 at 6 GiB (≈2.7–3.0×)**. Mechanism verified against the **real** allocator/ledger: multi-map proven (#727), charge-once refcount-gated `granule_refs` (#740/#745) is the right shape; gap is 1:N handle bookkeeping (out of scope). Read-only `PROT_READ` shared prefix composes with the #759 dummy-page `cuMemSetAccess` (independent per-subrange descriptors). Head-major impossibility confirmed & quantified (needs P ≥ 8192/16384 tok/head-stripe; seq-major divides by kv_heads). New finding: page sharing has its own granule floor (14B P≥1024, 0.5b P≥8192). ### Context (#776): the granule lever is dead `GRANULARITY_MINIMUM == RECOMMENDED == 2 MiB` (re-confirmed by `remap_burst`) — no 64 KiB granule, so the crossover cannot be shrunk by a finer granule, and quantized KV *doubles* the head-major crossover. Seq-major is thus the only *layout* lever, cutting floor + crossover + growth-spike by `kv_heads`. #759 fixed-stride+dummy-tail is reframed as a long-context hardening layer that coexists with seq-major. ### Change inventory & recommendation: **GO**, staged 7+ BNSH kernels; no byte-exact KV oracle exists (parity already compares outputs); native KV is device-only so physical layout may diverge from the declared BNSH shape. Smallest next increment: add a `KvLayout` tag; convert only the decode pair (GQA append + `gqa_decode_fp16` read); reuse the output-level parity oracle on qwen0.5b; then add 1:N handle bookkeeping to bank #777. Not a big-bang rewrite. A negative result was in scope; this is a positive, staged one. ### Validation `bench-seqmajor` passes `cargo fmt -- --check` and `cargo clippy --release --all-targets -- -D warnings` clean. No existing package source changed, so the repo CI gates (`onnx-runtime-ep-cuda`, `onnx-genai-engine`) are unaffected. --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 11, 2026
…ng (#777) (#793) ## VMM page-level prefix sharing — isolating GPU probe + design (#777) **Scope:** an isolating GPU probe in `crates/onnx-runtime-cuda-memory` plus a written design. **No KV integration, no engine or production-kernel change.** References #777, and builds on #727 (1-handle-at-2-VAs), #740 (physical-handle pool + 0→1/1→0 granule-ref attribution), #776 (2 MiB granule), #787 (token-major floor), #759 (dummy page). Read `docs/PREFIX_SHARE_INVESTIGATION.md` for the full measured answers, tables, and design; `docs/MEMORY_ARCHITECTURE.md` points to it. ### Result: the primitive is sound enough to build KV integration on — no kill finding All five questions cleared, each with its own GPU test binary (one question per binary so a poisoned context or a process-global counter cannot contaminate another): | Q | Test binary | Finding | |---|---|---| | **Q1** N-way sharing | `vmm_prefix_share_gpu` | One physical granule mapped into **N=8** sequences' reservations reads identically at every VA, including under **captured-graph replay** (the decode hot path). A write through any alias is visible at all — one physical page, not N copies. | | **Q2** ledger charge-once | `vmm_prefix_share_ledger_gpu` | Against the **real** `CudaVmmAllocator` + `LedgerGovernor` (no mock): a shared granule is charged **once**, stays alive until the **last** sharer leaves, and each additional sharer needs **0** incremental owned bytes. Admission arithmetic: **the Nth request costs only its private bytes.** | | **Q3** write protection | `vmm_prefix_share_write_protect_gpu` | A synchronous and an asynchronous write into a `PROT_READ` shared prefix both fault; the fault is **non-sticky** (context survives, later work succeeds) and **non-corrupting** (the other sharer's copy is intact). **The potential kill finding did not fire.** | | **Q4** compose with #759 | `vmm_prefix_share_compose_gpu` | A read-only shared prefix head, a read/write private live granule, and a read-only shared dummy tail **coexist in one reservation** with three distinct `cuMemSetAccess` postures. | | **Q5** copy-on-write | `vmm_prefix_share_gpu` | Boundary COW (copy 2 MiB + remap) is a **one-time burst**: ~**5.5 ms pooled** (#740 retained handle) / ~9.8 ms cold on WDDM — ~22% / ~39% of an assumed 25 ms decode step — never a per-token tax. | ### Concurrency-scaled saving (physical, granule-rounded) Under token-major the whole-model prefix is one contiguous range. Sharing an `S`-token prefix across `C` concurrent requests removes `(C−1)` duplicate copies. Example — qwen14b, **2,048-token** prefix, **8** concurrent: **2,688 MiB (≈2.6 GiB)** removed, matching the ~2.56 GiB figure in `MEMORY_ARCHITECTURE.md`. Full 512/2K/8K × 2/4/8/16 tables for qwen14b and qwen2.5-0.5b are in the doc. **Capacity consequence** (the number that matters): admissible concurrency rises from `floor(B / (prefix + private))` to `1 + floor((B − prefix) / private)`. For a 50/50 prefix/private split this roughly **doubles** concurrency on a fixed budget (e.g. 3.0 GiB KV budget: **4 → 8** requests), and the multiple grows as the shared prefix grows — reproducing the issue's "2 vs 8 concurrent" illustration. ### Design (documented, not implemented) - **Detection:** start with the **explicit pinned-prefix API** (unambiguous, covers the highest-value multi-tenant system-prompt scenario); hash-based detection is a later layer. - **Lifetime:** the **union** of all sharers' lifetimes — exactly the #740 granule refcount, plus an owner reference for a pinned prefix. - **Eviction:** a shared prefix is one unit affecting every sharer at once; the first increment makes a pinned prefix non-evictable for its lifetime. - **Layout:** impractical head-major, workable seq-major, **cleanest token-major** (one multi-map) — a further argument for #783/#787. ### Smallest next increment An explicit `commit_shared_prefix(handle, offset, len)` entry point on the token-major arena that maps a pooled physical handle into a new sequence's reservation at offset 0 and takes a granule-ref on it (owned axis → Nth sharer bills 0). Pinned, read-only, no detection/COW/eviction yet. Validate on the existing output-level parity oracle. The attention kernel and the flat per-sequence VA are unchanged. ### Constraints honoured Never unmap under a replay in flight (divergence syncs first); captures only read already-mapped stable VAs; no assertion inside any `Drop`; physical bytes reported from handle counts / the governor's owned axis, never nominal content bytes. ### Validation (local, this box — RTX 4060 Laptop, 8 GiB, WDDM) - `cargo fmt -p onnx-runtime-cuda-memory -- --check` ✅ - `cargo clippy -p onnx-runtime-cuda-memory --features gpu-tests --all-targets -- -D warnings` ✅ - New GPU tests + existing `vmm_graph_remap_gpu`, `vmm_dummy_page_gpu`, `vmm_kv_token_major_floor_gpu` — **13 tests, all pass**, single-threaded. ✅ **Do not merge, do not close #777.** Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 11, 2026
…capture (#716) (#796) ## Summary Weight offload previously forced whole-step CUDA graph capture **OFF**: the pager returned a different device pointer on every page-in (`alloc_raw`/`free_raw`), and a captured graph bakes pointers into its recorded nodes, so the two were mutually exclusive (`resolve_graph_capture_enabled` returned `false` whenever offload was on). The models that most need offload — those that do not fit in VRAM — were exactly the ones that lost graph capture and all of the #708/#728 capture-fragmentation wins (35B-A3B decode 154 -> 34 segments). This PR gives each retained weight `key` a **reserved-once device virtual address** and remaps physical granules underneath on page-in/eviction, so a captured graph replays correctly across evict->repage. It then removes the offload/capture mutual exclusion, **gated on the paging path actually being stable-VA** (the managed no-spill authority). Hard prerequisite for #755 (managed no-spill VMM as the default with automatic weight streaming). ## What changed, by step **Step 1 - prove the mechanism in isolation first** (`crates/onnx-runtime-cuda-memory/tests/vmm_stable_va_weight_slot_gpu.rs`, **passes on this HW**): a captured graph that reads from a stable VA still produces correct results after the pages under that VA are swapped for different physical granules holding **different weight values**. Extends the #727 remap-at-same-VA proof to the weight-value-swap case. **Step 2 - stable-VA slot conversion** (`crates/onnx-runtime-ep-cuda/src/weight_paging.rs`): - Per-key `StableWeightSlot { va, len }` reserved once through the **existing #740 authority-scoped pool** via `allocate_committed`/`carve()` suballocation. No second allocator, no per-weight reservation (the mistake that made #733 net-negative). Charging stays on the global granule-ref **0->1 / 1->0** attribution. - `CudaWeightPage::Drop` decommits granules but **keeps** the VA for stable slots (`deallocate_span` only for transient bypass pages). Never asserts in `Drop`. - Respects committed-granule admission (#745): reserving VA is free; only committed physical bytes are charged. 2 MiB granule sizing per #776. - **Enforceable ordering rule (not a comment):** page-ins - and thus eviction-driven `decommit` of any stable slot - run entirely under the residency mutex; capture is capture-once/replay-many with no page-ins during replay; the engine additionally declines capture when a step reports a bypass/eviction (`GLOBAL_BYPASSED_PAGE_INS`). So no `decommit` of a baked VA can occur while a replay is in flight - the case #727 explicitly did **not** prove safe. No `cuMemMap`/grow inside a captured region. - Residency GPU test `vmm_retained_weight_key_keeps_a_stable_virtual_address_across_repage` proves the VA is identical across evict->repage with one reused granule and **no accounting regression**. The existing `vmm_weight_admission_reuses_owned_granules_without_runtime_alloc_free` and `weight_offload_gpu` suite still pass (no regression). **Step 3 - remove the exclusion, gated on stable-VA** (`native_decode/{cuda,mod,load}.rs`): `resolve_graph_capture_enabled` now takes `weight_offload_stable_va` and keeps capture OFF only when `weight_offload_enabled && !weight_offload_stable_va`. The signal is `policy.enabled && policy.managed_no_spill` (the explicit byte `--vram-limit` authority that installs the VMM arena + physical pool). The pointer-unstable `alloc_raw`/`free_raw` compatibility path keeps the old capture-off exclusion. **#723's `StableResident` eviction policy is unchanged** - this changes only the addressing mechanism underneath it. **Step 4 - measure** (`crates/onnx-genai-engine/tests/stable_va_capture_offload_e2e.rs`): drives the managed no-spill VMM stable-VA authority and verifies **token-identical** greedy decode with capture OFF and ON in one process, dumping the deterministic offload counters (`page_ins`, `hits`, `evictions`, `bypassed_page_ins`, `vram_alloc_ns`/`vram_free_ns`). ## Honest measurement caveat The full sustained-paging throughput/counter comparison the issue asks for could **not** be completed in this environment: - The genuinely-offloaded model (`qwen14b-zp`, 15.5 GB) is **not native-loadable** - it has no `inference_metadata.yaml` (`model.io` ports), which is the separate #384 load-chain gap. The native decode path (where this change lives) cannot load it. - On the native-loadable 0.5B model, the managed no-spill authority is **pin-or-fail**: either the ~381 MB weight set fits fully resident (0 page-ins) or a step cannot commit its working set (admission/headroom failure after heavy eviction, ~504 evictions observed). There is no clean budget window that both pages sustainedly *and* completes a full decode - that is exactly what #755's automatic weight streaming is for. What **is** verified end-to-end here: under the managed no-spill VMM stable-VA authority, decode with capture OFF and capture ON is **byte-identical** (32/32 greedy tokens), i.e. capture is transparent and no longer force-disabled. The addressing mechanism itself is proven **deterministically** by the two GPU tests in steps 1-2. A negative/partial result was an acceptable deliverable per the task; this is a positive mechanism result with a partial end-to-end result. ## Validation (local; CI not waited on) - `cargo fmt -p onnx-runtime-ep-cuda -p onnx-genai-engine -p onnx-runtime-cuda-memory -- --check` PASS - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` PASS (no `--all-targets`, the CI gate) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` PASS (also linted the new cuda-gated test under `--features cuda,native-backend`) - `cargo test -p onnx-runtime-session` PASS (full suite) - GPU tests single-threaded PASS: `vmm_stable_va_weight_slot_gpu`, `vmm_retained_weight_key_keeps_a_stable_virtual_address_across_repage`, `vmm_weight_admission_reuses_owned_granules_without_runtime_alloc_free`, `weight_offload_gpu` (5/5) - engine `native_decode` unit tests PASS incl. new `weight_offload_on_stable_va_path_keeps_graph_capture` and updated `weight_offload_forces_graph_capture_off` Refs #716, #755, #727, #740, #723. **Do not merge; do not close #716.** Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
## Headline: layout now moves committed physical bytes PR #794 converted flash prefill and decode to the shared `KvCacheStrides` descriptor and then measured that head-major and seq-major commit **identical** physical bytes — `100,663,296 B` on qwen2.5-0.5b and `402,653,184 B` on qwen14b. The cause was on the binding layer: the native CUDA KV binding always committed a single flat bucket range `0..(capacity × kv_heads × head_dim × elem)` regardless of layout, so seq-major changed **kernel indexing only**, not reservation or commit geometry. #787 reached the same conclusion from the allocator side: the read path is free; the cost is on the binding layer. This PR builds the layout-aware commit **geometry**, makes the binding layer **consume** the KV layout descriptor, and **measures** the residency delta directly at the CUDA driver level. ### Measured committed physical bytes (deterministic, `vmm_kv_layout_residency_gpu`) On a fixed full-context stride (the stable-VA regime), at the near-empty floor: | Model | Per-binding head-major | Per-binding seq-major | Ratio | Model-wide head-major | Model-wide seq-major | |---|---:|---:|---:|---:|---:| | qwen2.5-0.5b | 4 MiB (2 granules) | 2 MiB (1 granule) | **2× = kv_heads** | 192 MiB (96 granules) | 96 MiB (48 granules) | | qwen14b | 16 MiB (8 granules) | 2 MiB (1 granule) | **8× = kv_heads** | 1536 MiB (768 granules) | 192 MiB (96 granules) | This is the number #794 measured as *equal*. It moves by the `kv_heads` factor once committed geometry follows the descriptor — the `layers × 2` floor materialising, on real 2 MiB CUDA granules. ### Growth / re-capture The same GPU test grows a seq-major fixed-stride binding **under a captured replay** — committing tail granules at a stable VA and stride (granule `0 → 3`) — and the originally captured graph replays unchanged: **0 re-captures**. This is the stable-stride win #782 could not demonstrate end to end, shown here on hardware. Bucket growth, by contrast, re-strides and forces a re-capture. ### Boundary-crossing step cost Not measured end to end in this PR, and reported as such rather than extrapolated. #778's 49.5 ms vs 4.5 ms microbenchmark needs a seq-major **generation** through the engine to reproduce, which requires the BSNH fixed-stride shape build described below (there is no seq-major model on record to drive it). The driver-level residency measurement is the correctness-independent result this PR stands on. ## What changed - **`crates/onnx-genai-engine/src/native_decode/kv_commit.rs`** (new) — pure, unit-tested `KvCommitLayout` geometry. `live_prefix_ranges` returns the commit ranges a layout implies on a fixed full-context stride: seq-major is one dense run `0..ceil(live_bytes/granule)`; head-major is one live fragment per head stripe. `KvCommitLayout::default() == HeadMajor`, so **absent metadata means head-major, byte-identical to today**. 9 unit tests. - **Engine consumes the descriptor** — `resolve_cuda_kv_layout` maps `ModelIoSpec.kv_layout` (`#792` view-carrying descriptor) to `KvCommitLayout` (env override `ONNX_GENAI_CUDA_KV_LAYOUT` for diagnostics only), threaded from `load.rs` into `DecodeCudaState`, surfaced as `CudaKvDebugStats.kv_layout_seq_major`. Head-major keeps its exact behavior; the growth path (`apply_vmm_growth` / `vmm_growth_requests`, which #716 edits) is **untouched** for a clean merge. - **`crates/onnx-runtime-cuda-memory/tests/vmm_kv_layout_residency_gpu.rs`** (new) — driver-level GPU measurement (2 tests), the headline above. Uses the #740 pool's granule model; commits/reads back real committed bytes. - **`docs/MEMORY_ARCHITECTURE.md`** — floor table + note updated from "not realized today" to the measured driver-level result, plus the precise structural boundary. ## Honest structural boundary (per the task's "stop and report" clause) The layout-aware commit *geometry* is built and measured, and the engine now consumes the layout metadata. What is **not** wired is the seq-major dense-prefix commit on the end-to-end **generation** path, because: - `persistent_state_shapes` hard-codes a **BNSH** physical shape whose axis-2 (seq) grows; there is no **BSNH fixed-stride** physical-shape build to hang the dense-prefix commit on, and - no model on this machine declares seq-major, so a byte-identical seq-major *generation* cannot be validated here. The per-layer `key_cache_*`/`value_cache_*` bindings must survive as views (they do — this PR does not collapse them). The BSNH fixed-stride shape build is the shared prerequisite for token-major and #777 page-level prefix sharing, and is the documented next step. This is reported as information, not a failure. ## Correctness `gqa_seqmajor_parity_gpu::seq_major_full_generation_prefill_and_decode_is_bit_identical` (independent CPU prefill + decode oracles) passes — decode output is byte-identical between head-major and seq-major for a full generation, capture on and off. ## Validation (all green locally; not waiting for CI) - `cargo fmt -p onnx-genai-engine -p onnx-runtime-cuda-memory -- --check` ✓ - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` ✓ - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` ✓ - full `cargo test -p onnx-runtime-session` ✓ - GPU single-threaded: `vmm_kv_layout_residency_gpu`, `vmm_kv_contiguous_tail_gpu`, `vmm_kv_token_major_floor_gpu`, `gqa_seqmajor_parity_gpu` (2), and `group_query_attention_gpu` (25 passed, 1 ignored H200 bench) ✓ - `kv_commit` unit tests (9) + `cuda_kv_layout_resolves_from_metadata_with_env_override` ✓ ## Notes / caveats - The box is heavily contended; wall-clock is unreliable (identical runs have ranged 3.9–28 tok/s). Results above lead with **byte and granule counters**. - **qwen14b** needs weight offload on this 8 GiB GPU, which still disables capture (#716, separate) — so capture counters are unavailable at that scale. The re-capture demonstration is at the driver level and on the small model. Refs #794, #787, #782, #750, #777. **Do not merge; does not close any issue.** --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
…#755) (#798) Closes #755. Makes the authority-scoped **managed no-spill VMM** the **default** on the native CUDA path — no flag required — with **automatic weight streaming** when a model exceeds the resolved device budget. Being larger than the budget becomes a supported configuration rather than an error. Builds directly on the prerequisite #716 (PR #796), which made weight offload capture-compatible by paging weights under a stable virtual address, so flipping the default no longer disables CUDA graph capture for offloaded models. ### What "default" now means 1. **Managed VMM is on without a flag.** Native CUDA loads run memory-strategy inference (#752) unconditionally and it *drives* policy: allocation goes through the physical-handle pool (#740) with committed-granule admission (#745), governed workspace planning (#747/#753/#795), transactional growth grants (#748), against one server-owned authority per device (#744). 2. **Over-budget is not an error.** When the resolved budget cannot hold the weights, the runtime auto-enables weight streaming/offload instead of failing. 3. **`--vram-limit` is an override**, not the trigger. 4. **The legacy allocator stays reachable** for one release via `ONNX_GENAI_LEGACY_ALLOCATOR=1` (back-compat: `ONNX_GENAI_DYNAMIC_KV_WEIGHT_LENDING=0`). ### Observability The resolved budget, chosen strategy and offload state are printed at startup and exposed in `/v1/resources` under a new `memory_strategy` block, so the new default is observable rather than mysterious. ### Invariants preserved - **No silent WDDM spill** — committed physical bytes stay within the managed cap; `cuMemCreate` structurally cannot spill (#704). - **Capacity refusal stays a pre-header HTTP 429** (#743); committed-granule admission (#745) and growth grants (#748) unchanged. - **A model that fits does not start paging** under the new default — tested explicitly (the most likely regression). - **Capture stays ON when offload is active on the stable-VA path** (#796 property), pinned by the existing unit tests. - No prefetch re-introduced (#715); no assert inside any `Drop`. ### Measurement (native CUDA, deterministic counters lead) Ignored live e2e `qwen2_5_0_5b_managed_vmm_default_e2e` (Qwen2.5-0.5B int4 mobius, **same process**, interleaved managed-default vs. legacy opt-out): | scenario | strategy | offload | managed_no_spill | resolved budget | committed device bytes | page-ins / evictions | |---|---|---|---|---|---|---| | **A** fits, managed default (no flag) | `FullResident` | false | true | 7,730,940,928 | 381,681,664 | **0 / 0** (does not page) | | **B** synthetic over-budget (`--vram-limit 268,435,456`) | `DynamicWeightResidency` | true | true | 268,435,456 | 90,177,536 (« cap; no spill) | **434 / 432** | | **C** legacy opt-out (`ONNX_GENAI_LEGACY_ALLOCATOR=1`) | `FullResident` | — | **false** | — | — | — | | **D** explicit `--vram-limit 12 GiB` | override applied | — | — | 12,884,901,888 | — | — | In **B**, at that extreme synthetic cap the residency arithmetic **refused cleanly** ("no page is evictable") rather than spilling — the no-spill contract holds. **Caveats (measurement discipline):** - Wall-clock / tok/s is **not** reported: this box ranged 3.9–28 tok/s across identical runs and other agents build concurrently. Deterministic counters only. - A genuine native large-model run that exceeds VRAM was **not measured**: `qwen14b-zp` lacks `inference_metadata.yaml` (#384) so it is not native-loadable. The over-budget condition was **synthesized** with a small explicit budget, as stated. - **CUDA-graph segment count** is not a public counter; capture-ON-under-offload on the stable-VA path is pinned by the #796 unit tests, not re-measured here. ### Validation - `cargo fmt -p onnx-genai-engine -p onnx-genai-server -- --check` clean (`--all-targets --all` fmt hits Windows os error 206 on this box's long worktree paths — a command-line-length artifact, not a formatting issue). - clippy `-D warnings`: engine (`native-backend`; also `cuda,native-backend`), server, `onnx-runtime-ep-cuda --features cuda` — all clean. - Full `cargo test -p onnx-runtime-session` passes. - Engine `memory_strategy` unit tests (12, incl. 4 new) and server resources tests pass; cuda-gated engine tests incl. #796 capture arms pass. - **Pre-existing unrelated failure:** `failed_runtime_shrink_preserves_policy_and_ledger_limit` (server) fails identically on `origin/main` (confirmed via stash); its error is on the `host` authority, untouched by this native-CUDA-only change. Not fixed here (out of scope). Do **not** merge and do **not** close the issue via this PR beyond the `Closes #755` reference; leaving final disposition to review. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
## Land the `commit_shared_prefix` allocator entry point (#777) Implements the smallest next increment named by the #777 probe (PR #793): the explicit pinned-prefix API on the CUDA VMM arena. This is the **allocator primitive only** — not a kernel change, not copy-on-write, and not wired into the KV path or engine. Refs #777, #793, #727, #740, #745. ### What this adds An explicit pinned-prefix API on `CudaVmmAllocator` in `crates/onnx-runtime-cuda-memory`: - **`create_shared_prefix(bytes)`** — reserves a private window and acquires N physical granules from the existing #740 authority-scoped pool, maps them read/write for the owner to fill, and charges them **once** on the owned axis (`governor.used(Tier::Device)`). Returns a `SharedPrefix` holding one owner-reference (the registry-entry lifetime model, so the prefix survives momentarily-zero sharers). - **`commit_shared_prefix(&prefix, ptr, allocation_bytes, byte_offset)`** — maps those already-owned handles into a *different* sequence's reservation at a granule-aligned `byte_offset`, **`PROT_READ`** (the Q3 non-sticky write-protection posture). Goes through the existing #740 pool and `carve()` suballocation on the existing global granule-ref **0→1 / 1→0** attribution. **No second allocator. No per-sequence physical reservation** (what made #733 net-negative). - **`incremental_owned_bytes_for_shared_prefix(&prefix)` → 0** — the Nth sharer bills only its private bytes, Q2's admission arithmetic made explicit on the entry point (#745). ### How the invariants hold in code (not comments) - **Union lifetime, structurally.** A pool-level `shared[handle]` refcount composes with the arena's per-granule allocation refcount; the physical handle returns to the pool only when the **last** mapping (owner or any sharer) unmaps. Freeing N−1 sharers does not release it; freeing the last one does. - **Charged once.** Shared granules are charged on `create_shared_prefix`; each `commit_shared_prefix` adds **0** owned bytes. - **Never map inside a captured region.** `enter_graph_capture()` raises a depth counter that `commit_shared_prefix` refuses under — the ordering rule is enforced, not documented. - **Loud write-protection.** Shared granules are `CU_MEM_ACCESS_FLAGS_PROT_READ`, so a stray write faults rather than silently corrupting another request's KV; Q3 proved this is non-sticky and this preserves that property. - **Physical bytes only** (`committed_physical_bytes`, `mapped_bytes`), never nominal content bytes. **No `assert!` in any `Drop`** (STATUS_STACK_BUFFER_OVERRUN on this platform). - **Errors, never mis-maps.** Misaligned target VA, granularity mismatch, over-long prefix, out-of-arena pointer, non-live allocation, a target granule that is already committed, and any attempt under an open capture all return `Err` and roll back cleanly (unmap + pool return), leaving the region uncommitted. ### Tests (`vmm_commit_shared_prefix_gpu.rs`) Against the **real** `CudaVmmAllocator` + `LedgerGovernor` (no mock), one concern per test: - N sequences share one pinned prefix; the ledger charges it **once**; freeing N−1 sharers does not release it, freeing the last one does; teardown returns `total_owned_bytes`, `governor.used`, and `creates − releases` all to zero. - Admitting an additional sharer costs only its private bytes (asserted against the real ledger). - A write into the shared region faults and the context survives (reuses the Q3 pattern; sync and async). - An unsupported request errors rather than mis-mapping. Per the #797 harness warning, all device ops route through synchronous copies or a single created stream (avoiding the legacy-default-stream vs non-blocking-stream silent partial-fill race), and each test builds its own allocator so no alphabetically-earlier sibling warms the context. ### Validation (local; not waiting for CI) - `cargo fmt -p onnx-runtime-cuda-memory -- --check` — clean. - `cargo clippy -p onnx-runtime-cuda-memory --features gpu-tests --all-targets -- -D warnings` — clean. - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` — clean. - GPU suites, single-threaded, all pass: the new `vmm_commit_shared_prefix_gpu` (4/4) plus the required regression set — the four `vmm_prefix_share_*`, `vmm_graph_remap_gpu`, `vmm_kv_token_major_floor_gpu`, `vmm_kv_layout_residency_gpu`, `vmm_stable_va_weight_slot_gpu`. ### What is now possible vs. what still needs KV integration The allocator can now physically share a pinned prefix across independent sequence reservations — charged once, read-only, correct union lifetime and admission arithmetic. The capacity mechanism (roughly doubled admissible concurrency for a large shared system prompt) now exists at the allocator layer. It does **not** yet: construct a sequence's KV as `shared prefix + private tail` in the engine, detect that two requests share a prefix, or handle divergence past the shared region. The immediate next increment named from evidence is **KV-path integration** (validated on the output-level parity oracle), which depends on the token-major KV shape/stride work (#783/#787) since only token-major reduces the prefix to a single multi-map per sequence. Then hash-based detection (layer 2), then boundary copy-on-write (layer 3, cost already characterized in Q5). `docs/PREFIX_SHARE_INVESTIGATION.md` is updated to move `commit_shared_prefix` from "next increment" to a "Landed this round" section and to name the following increment. > Not merging and not closing #777 per the task. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
…809) ## What this does (#777) PR #803 landed the prefix-sharing primitive (`create_shared_prefix` / `commit_shared_prefix` / `incremental_owned_bytes_for_shared_prefix`) on the production `CudaVmmAllocator`, but **nothing in production called it** — only its definition and GPU tests. An unconsumed contract reads as working infrastructure while proving nothing; that is exactly the shape that left #721 stage 3 stranded. This PR gives the primitive a **first real, deliberately narrow consumer**. ### What is now shared in production - **An allocator-agnostic seam on the `DeviceAllocator` trait** (`onnx-runtime-memory-governor`): `create_shared_prefix`, `incremental_owned_bytes_for_shared_prefix`, `commit_shared_prefix`, plus an opaque `dyn SharedDevicePrefix` handle and `SharedPrefixCommitInfo`. A caller holding only `dyn DeviceAllocator` can now pin a token prefix once and have subsequent sequences map it. Default impls **refuse** (`InvalidRequest`) rather than mis-map, so every existing non-VMM allocator keeps compiling and refuses safely (unsupported requests error rather than mismap, per #745/#750). - **`CudaVmmAllocator` implements the seam** by delegating to the merged inherent methods (#803) and downcasting the opaque handle via `as_any()`. Granules come from the existing #740 authority-scoped pool through `carve()`; no second allocator, no per-sequence physical reservation. - **First consumer:** the real **seq-major** fused fp16 GQA decode kernel (#782/#792/#794) reads shared-prefix VMM KV through the `dyn DeviceAllocator` seam. The prefix is physically contiguous as `layers × 2` ranges (seq-major, the gate #793 requires); token-major (one multi-map/sequence) is measured (#787) but not built here. The API is **explicit** (a caller declares the shared prefix), as #793 recommended for the highest-value case (a shared system prompt), and is **restricted to prefixes that are provably read-only for the sharers' lifetime** — copy-on-write at divergence is out of scope. ### Proof + measured numbers New GPU parity test `crates/onnx-runtime-ep-cuda/tests/gqa_shared_prefix_parity_gpu.rs` drives the real seq-major GQA decode kernel over shared-prefix VMM KV and proves two sequences sharing one pinned prefix produce **byte-identical** output to two independent sequences. Measured (KV_HEADS=8, HEAD_DIM=128, f16, 1024-token prefix + 1024-token private tail per sequence; 2 MiB granule → 1024 tokens/granule): | quantity | value | | --- | --- | | independent (2 sequences) committed **physical** bytes | **16,777,216 B** (8 granules) | | shared (2 sequences) committed **physical** bytes | **12,582,912 B** (6 granules) | | removed by sharing | **4,194,304 B** (2 granules = `(C−1)×(K_prefix+V_prefix)`) | | second-sharer admission cost | **4,194,304 B** — private tails only | | prefix `incremental_owned_bytes_for_shared_prefix` | **0** (charged once) | | output | **byte-identical** for both sequences | Deterministic counters lead; wall-clock is omitted (this box's identical runs range 3.9–28 tok/s and other agents build concurrently — no number is extrapolated). ### Validation (all local; not waiting for CI, fix-forward) - `cargo fmt -p onnx-runtime-memory-governor -p onnx-runtime-cuda-memory -p onnx-runtime-ep-cuda -- --check` — clean (`--all` fails on this box with os error 206, a Windows command-line-length artifact, not formatting). - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` — clean. - `cargo clippy -p onnx-runtime-cuda-memory --features gpu-tests --all-targets -- -D warnings` — clean. - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` — clean. - Full `cargo test -p onnx-runtime-session` — pass. - Single-threaded GPU suites — pass: `vmm_commit_shared_prefix_gpu` (4), `vmm_prefix_share_gpu` (2), `vmm_prefix_share_ledger_gpu` (1), `vmm_prefix_share_write_protect_gpu` (1), `vmm_prefix_share_compose_gpu` (1), `gqa_seqmajor_parity_gpu` (2), and the new `gqa_shared_prefix_parity_gpu` (1). Harness respects the known traps: VMM allocator built on the EP's own CUDA context; all device ops route through the single EP stream (#797); pool env set before any allocator is constructed and the process is isolated (#804); no assertions inside any `Drop`; live reservations explicitly deallocated before allocators drop. ### What remains (structural / out of scope) - **Auto engine use is blocked structurally.** The engine generation loop cannot call this seam automatically until `native_decode/cuda.rs` can build a seq-major fixed-stride physical KV shape — `persistent_state_shapes` is hard-coded BNSH and no model declares a seq-major end-to-end physical shape (#794 showed seq-major changed only kernel indexing, not commit geometry). A BNSH/seq-major fixed-stride physical-shape build is the prerequisite. This PR delivers the seam + the seq-major GQA-kernel consumer + proof; the automatic decode-path wiring is the next structural increment. - Hash-based automatic prefix detection (later; explicit path proven first). - Token-major (one multi-map per sequence; measured #787, not built). - Copy-on-write at divergence (#793 measured ~5.5 ms pooled / ~9.8 ms cold, once). Docs updated: `docs/MEMORY_ARCHITECTURE.md` (concurrency-scaled prize section) and `docs/PREFIX_SHARE_INVESTIGATION.md` (smallest next increment) now record the first production consumer, the measured numbers, and the remaining structural gap. Refs #777 #803 #793 #782 #745 #750 > Please do **not** merge and do **not** close #777 from this PR. Edits are > confined to the KV/prefix integration path to merge cleanly with concurrent work. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
…ssage (#821) ## Root Cause **Case (b)**: the test was over-fitted to an exact error string that legitimately changed. Commit `c7633eec4` ("Pool CUDA VMM physical handles across reservations", PR #740) refactored the server memory-shrink rejection path in `crates/onnx-genai-server/src/state.rs`. The old message was: ``` cannot satisfy lowered resource limit of {resolved} bytes: {domain} currently has {used} committed bytes ``` The new message is: ``` cannot satisfy lowered resource limit of {resolved} bytes: {domain} has {used} leased bytes but only {releasable} pooled-unmapped bytes are releasable ``` The test `failed_runtime_shrink_preserves_policy_and_ledger_limit` (added one commit earlier in `7fa131833`, PR #744) asserted `error.contains("committed bytes")`, which broke when #740 landed. **Introducing commit**: `c7633eec4` (PR #740) ## Fix Update the assertion to check for `"cannot satisfy lowered resource limit"` — the stable prefix shared by all rejection paths in both `state.rs` and `memory_authority.rs`. **Why this is correct, not merely green**: The test verifies two properties: 1. Shrinking below committed memory *fails* with a limit-rejection error (not silently succeeds or panics). 2. The policy and ledger limits are *rolled back* to their original values. Property (1) is now checked by asserting the error contains the stable semantic prefix. Property (2) is checked by the subsequent `assert_eq!` lines on the snapshot and policy, which remain unchanged. ## Validation ``` $ cargo test -p onnx-genai-server --lib --no-fail-fast test result: ok. 183 passed; 0 failed; 2 ignored; 0 measured; 0 filtered out; finished in 10.42s $ cargo clippy -p onnx-genai-server -- -D warnings Finished (clean) ``` Integration tests (`http`, `vlm_image_bundle`) have 10 pre-existing failures unrelated to this change (they fail identically on origin/main). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This was referenced Aug 12, 2026
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
…verdict for native batching (#750) (#844) Stage 1 of native batched decode (#750). Per the reprioritization comment, stage 1 is **(1) validate the premise, (2) scope the smallest honest increment, (3) land it or state the staging plainly.** This PR does not deliver batching — it lands the measurement instrument the effort needs and reports the verdict below. **Do not merge / do not close #750.** ## 1. Verdict on the premise: does native batching amortize the weight stream? **CONFIRMED — conditional on a single-fused-forward implementation.** Established by code reading plus the existing #837 counters (GPU is contended by three sibling agents; per the protocol I did the CPU-side work first and did not run a contended headline). **Evidence the weight-residency cache is *not* per-sequence:** - `CudaWeightResidency` (`crates/onnx-runtime-ep-cuda/src/weight_paging.rs`) is owned per-EP/per-device and keyed purely by weight identity: `pages: HashMap<u64, Arc<CudaWeightPage>>`. There is no sequence or batch dimension anywhere in the cache, its policy, or its stable-VA slot map. - The page-in entry point `resident_mapped(key, ...)` / `page_lazy_weight(key, ...)` (`provider.rs`) returns a cache **hit** for a resident key regardless of which row/sequence triggered it. A miss streams once and admits under the LRU/stable-resident budget. - The ORT batched path already exploits exactly this: `Engine::continuous_batch_manager` (`batched.rs:649`) binds `batch_size = N` into **one** `BatchedStaticCacheDecodeSession` run per step (GEMM `M = N` rows). Each weight op executes once per step and streams once for N tokens. The native path simply has no equivalent — it structurally pins batch to 1. **Corroborating counters (#837 baseline, qwen14b-zp, 16 tokens):** `htod_bytes = 65,772,419,072` ⇒ 4.11 GB/token; `page_ins = 5,535` ⇒ ~346 page-ins/step. ~346 distinct weight tensors streamed once per forward is consistent with "one forward per decode step, each offloaded weight paged exactly once." A single-fused batch-N forward keeps ~346 page-ins/step while emitting N tokens ⇒ `htod_bytes/token` and `page_ins/token` fall ~1/B. **The premise FAILS only if** batch-N is implemented as N independent forwards (or N per-sequence graph replays). On this over-budget model (16.65 GB weights vs 7.73 GB budget) the working set is evicted every step, so a per-sequence loop would miss each weight N times per step and amortize nothing. So the premise is a statement about **implementation discipline**: batch must be one fused forward with `M = N`, never a per-sequence loop. **The knee is unlocated.** As B grows, per-step compute grows and KV grows `B × context`; at some B the step becomes compute-bound (or KV+activations shrink the resident weight budget enough to lower `hit_rate`) and scaling flattens. Honest claim: "scales with B up to a knee we have not located," never "batch 8 = 8x." ## 2. Chosen scope, and why the batch-N increment itself is multi-PR The issue's suggested first increment (make the native path report `batch_supported = false`) is **already landed and tested** (#758 — `driver.rs:697-713`, `routes/admin.rs`, server test at `server/tests.rs:1876`). It is not available to re-land. The smallest *structural* increment — uniform-length batch-N (all sequences stepping together, no ragged admission, sidestepping the paged path) — is **genuinely multi-PR** on the native path. `native_decode/cuda.rs` (3,963 lines) assumes batch=1 structurally in at least four coupled places, each of which must change together to keep a valid decode: 1. **Input binding** collapses batch/query-seq to 1 (single-token `[1,1,·]` device bindings, `cuda.rs:574`). 2. **KV cache** is a single-sequence "batch-0 addressing" contiguous/VMM buffer with no row dimension (`cuda.rs:1782`, `ensure_capacity_*`); N sequences need N reservations or a batched KV layout. 3. **CUDA-graph capture** bakes the batch=1 shape; batch-N forces a recapture per batch shape (must be quantified against the `captures > 0 / fallbacks == 0` gate, #796). 4. **Symbolic-axis resolution** pins axis 0 to 1 (`cuda.rs:2371-2432`), rejecting any non-collapsible batch axis. Given the hard constraints (batch-1 byte-identical; capture preserved; `peak_committed < managed_limit` with KV now `B × context`; ledger counters at 0), landing a *correct* batch-N in one PR would be large and risky. Per the explicit guidance ("prefer a smaller increment that lands over a larger one that does not," and "a well-evidenced staging plan is a good stage-1 outcome"), I chose to land the **measurement instrument** now and stage the implementation: - **Stage 2a** (next): batch-N *input binding + fused forward* behind the existing admission contracts, uniform length only, with a batch-1 byte-identical guard test. No KV layout change yet — validate the fused-forward amortization on a 2-row synthetic before touching KV growth. - **Stage 2b**: batched KV (N reservations under the #740 pool, or batched layout) + recapture accounting. - **Stage 2c**: wire `continuous_batch_manager` to a native uniform-batch manager; ragged/continuous batching follows. ## 3. What this PR lands `weight_offload_amortization: emitted_tokens=… htod_bytes_per_token=… page_ins_per_token=…` on both native report paths in `profile_native`. These are the batch-invariant ratios the protocol mandates leading every batch comparison; `profile_native` previously printed only the raw accumulated counters, leaving the ratio to hand-computation and to be masked by wall-clock variance. Pure bench-only reporting — **runtime decode is byte-identical**. Unit test pins the ratio against the #837 baseline (65,772,419,072 / 16 = 4.11 GB/token). **Baseline (from #837, to be re-emitted by the new line once the GPU frees):** | metric | value | |---|---| | htod_bytes_per_token | 4.11 GB (65,772,419,072 / 16) | | page_ins_per_token | ~346 (5,535 / 16) | An "after (batch-N)" row is deliberately absent — this PR lands no batching. Stage 2a will populate it, and the ~1/B drop in these two rows is the pass/fail signal. ## Validation - `cargo fmt -p onnx-genai-bench -- --check` ✓ - `cargo clippy -p onnx-genai-bench --features bench-native,cuda --bin profile_native -- -D warnings` ✓ - `cargo test -p onnx-genai-bench --features bench-native,cuda --bin profile_native` (new `per_emitted_token_divides_counter_by_tokens`) ✓ No memory/workspace-governance code touched, so the mobius parity gate is N/A for this PR. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
This was referenced Aug 12, 2026
justinchuby
added a commit
that referenced
this pull request
Aug 14, 2026
) (#969) ## #956: route plugin-path device allocations through the pooled VMM arena Closes #956. ### Corrected diagnosis (verify, do not trust) The task's hypothesis was that `adopt_memory_governor` is the only installer of the VMM arena and is never called on the plugin path, so ORT scratch falls through to `CudaDeviceAllocator` (`cuMemAlloc`). **That is outdated.** The provider constructor (`new_with_policy_and_governor`) installs the arena **at construction** when `ONNX_GENAI_CUDA_VMM` is enabled — this has been true since PR #740 (constructor-time install), not only at governor adoption. `factory_create_allocator` in `onnx-runtime-ep-plugin` wraps the shared EP in a `DeviceAllocator` that projects the EP's `allocate`/`deallocate` through ORT's `OrtAllocator` vtable, so **ORT scratch already routes through `CudaExecutionProvider::allocate` → `self.memory()` → the arena** when the arena is on. The plumbing exists. The real gap: the **standalone** (plugin, no-governor) constructor branch built the arena via `standalone_with_teardown_synchronizer`, which took its physical-handle-pool bound from `physical_handle_pool_bytes()` (env only). With no env override and no governor to supply a default, the standalone arena retained **zero** physical granules: every scratch free unmapped *and released* its physical granules, and the next same-size request re-created them. That trades `cuMemAlloc` per dispatch for `cuMemCreate`/`cuMemRelease` per dispatch — still a per-dispatch driver allocation, still capture-unsafe. That is the exact residual #956 tracks. ### The fix Give the standalone arena a **default 256 MiB retained physical-handle pool** (matching the governor path's existing default), so repeated same-size scratch requests reuse committed memory instead of churning `cuMemCreate`/`cuMemRelease`. - `vmm_allocator.rs`: `standalone_with_teardown_synchronizer` gains a `default_pool_bytes: Option<usize>` param; pool bound is now `physical_handle_pool_bytes().or(default_pool_bytes)` (explicit env still wins). - `provider.rs`: new `DEFAULT_STANDALONE_PHYSICAL_POOL_BYTES = 256 << 20`; the standalone (None-governor) branch passes `Some(DEFAULT_STANDALONE_PHYSICAL_POOL_BYTES)`. - `device_allocator.rs`: added a `cumemalloc_calls` counter to `CudaDeviceAllocator` (the direct `cuMemAlloc` instrument used by the proof test). **No default flipped.** The capability lives behind the existing `ONNX_GENAI_CUDA_VMM` gate. The default pool only affects the standalone/plugin branch (reached by `new_default`/`new(0)` with `governor = None`); native/CLI/governed paths take a different branch and are byte-for-byte unchanged. ### Evidence — the arena *served* the requests (measured, not inferred) GPU proof tests (`--features cuda,gpu-tests`, RTX 4060, run solo per #851): - **Default `cuMemAlloc` path scales one-for-one:** `#956 default path cuMemAlloc calls: 16 requests -> 16, 64 requests -> 64` - **Standalone VMM arena does NOT scale, and the pool actually served each request:** `#956 standalone scratch reuse: warm(creates=1 hits=0 owned=2097152B) ` `+16cyc(creates=1 hits=16 owned=2097152B) +64cyc(creates=1 hits=80 owned=2097152B)` → `cuMemCreate` stays at **1** regardless of step count; `pool_hits` grow one-for-one (16 → 80), proving the retained pool *served* the requests; committed bytes bounded at **2 MiB** (1 granule) at both 16 and 64 cycles — **no leak**; `releases` flat (handles reused, not released per cycle). - **Plugin construction path routes through the arena:** `plugin_construction_path_routes_device_memory_through_pooled_vmm_arena` builds the provider exactly as the CUDA plugin does (`new(0)` + `ONNX_GENAI_CUDA_VMM=1`), asserts `commits_on_demand()` (arena, not `cuMemAlloc`, serves), and shows a real EP `allocate`/`deallocate` moves the pool's `creates + pool_hits` — the arena served an actual EP allocation, not merely installed. **This is a local counter-based proxy for the H200 `nsys` gate named in the issue's acceptance criteria, NOT the gate itself.** `scripts/validate_ep_workspace_h200.py --nsys` and `compute-sanitizer` are not runnable on this box. The counters measure the same underlying property the nsys gate measures (driver allocation calls vs. decode steps). ### Numerics (criterion 4) - The diff is a strict **no-op on all non-plugin paths**: the changed standalone branch is only reached with `governor = None` (the plugin path). Native/CLI/governed generation executes identical code before/after, so generation output on those paths is provably byte-identical. - The only new numerics risk is whether pooled reuse returns correct data. The reuse proof test writes a per-cycle-tagged 256-byte header and reads it back **every cycle** (all 80 cycles), asserting the reused committed memory returns exactly the bytes written that cycle — a stale read from a prior cycle would mismatch. All cycles passed. - **Honest limitation:** plugin-path token-id generation is **not runnable on this box** — the CUDA plugin crate is documented "UNVALIDATED ON HARDWARE" and there is no wired plugin-generation harness here. I did not fabricate token ids. The no-op-on-governed- paths argument plus the per-cycle read-back integrity check is the numerics evidence I can stand behind on this hardware. ### Gates (verbatim; check passed count, not exit code) - `cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-cuda-memory -p onnx-runtime-ep-plugin` — clean (exit 0, `--check` clean). - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` — clean (exit 0, no warnings). - `cargo test -p onnx-runtime-ep-cuda --features cuda --lib` — `test result: ok. 374 passed; 0 failed; 10 ignored; 0 measured; 0 filtered out` (baseline restored; run solo on idle GPU). - `cargo test -p onnx-runtime-cuda-memory` — lib `test result: ok. 8 passed; 0 failed`. - `cargo test -p onnx-runtime-ep-plugin` — lib `test result: ok. 207 passed; 0 failed`. ### Scope Out of scope per the issue: flipping any default. Capability is gated behind `ONNX_GENAI_CUDA_VMM`; measurements reported so the owner can decide the default after measuring the arena against `cuMemAlloc` on real models. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
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
CudaVmmAllocatorPolicy and contract
ONNX_GENAI_CUDA_PHYSICAL_HANDLE_POOL_BYTESenables production pooling and sets its retained-byte boundGPU proof
Production-path A→B transfer, constructed through two real
CudaVmmAllocator::newcalls with separate context wrappers and one authority:Direct three-granule proof:
The production test was sabotaged by disabling allocator pool selection and failed at
production allocator must install the pool, confirming it exercises production wiring. A separate GPU test verifies different authorities receive different pools and cannot exchange handles.Validation
cargo fmt --all -- --checkcargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warningscargo clippy --locked --all-targets -p onnx-runtime-cuda-memory -- -D warningscargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warningscargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend,cuda -- -D warningscargo test --locked -p onnx-runtime-virtual-memory(21 passed)cargo test --locked -p onnx-runtime-cuda-memory --features gpu-tests --test physical_pool_production_gpu -- --nocapture --test-threads=1(1 passed on GPU)cargo test --locked -p onnx-runtime-cuda-memory --features gpu-tests --test virtual_memory_gpu -- --nocapture --test-threads=1(9 passed on GPU)Closes #736