Repository navigation
Probe CUDA graph replay across VMM remap - #727
Merged
Merged
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #727 +/- ##
==========================================
- Coverage 80.69% 80.15% -0.54%
==========================================
Files 329 331 +2
Lines 130463 132633 +2170
Branches 130463 132633 +2170
==========================================
+ Hits 105273 106309 +1036
- Misses 20959 22086 +1127
- Partials 4231 4238 +7
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
Add a CUDA-gated integration probe for issue #721 that captures a graph over a reserved VMM address, replays after full remap and growth-shaped mapping, probes thread-local capture behavior for cuMemMap, and verifies alias mappings across virtual addresses. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
force-pushed
the
squad/721-vmm-graph-probe
branch
from
August 7, 2026 01:59
6fdb57f to
237e986
Compare
|
| Status | Scenario | Base | PR | Change |
|---|---|---|---|---|
block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1 |
556.32 µs | 647.54 µs | +16.4% | |
| ✅ | qwen3_sampling_processors/top_k_partial_selection |
130.57 µs | 144.27 µs | +10.5% |
| ✅ | qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline |
3.27 ms | 3.59 ms | +9.8% |
| ✅ | block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024 |
97.09 µs | 105.47 µs | +8.6% |
| ✅ | tokenization/encode_tokens_per_second |
392.22 µs | 423.12 µs | +7.9% |
| ✅ | qwen3_sampling_processors/top_k_top_p_full_sort_baseline |
5.24 ms | 5.65 ms | +7.8% |
| ✅ | qwen3_sampling_processors/top_k_full_sort_baseline |
1.99 ms | 2.13 ms | +6.7% |
| ✅ | block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1 |
83.83 µs | 89.23 µs | +6.4% |
| ✅ | tokenization/decode_tokens_per_second |
6.31 ms | 6.69 ms | +6.0% |
| ✅ | sampling_latency/top_k_per_token |
54.88 µs | 57.61 µs | +5.0% |
| ✅ | block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024 |
131.62 µs | 137.45 µs | +4.4% |
| ✅ | sampling_latency/min_p_per_token |
212.63 µs | 219.32 µs | +3.1% |
| ✅ | qwen3_sampling_processors/top_p_fast_after_top_k |
510.84 µs | 522.84 µs | +2.4% |
| ✅ | sampling_latency/top_p_per_token |
377.51 µs | 384.50 µs | +1.9% |
| ✅ | kv_cache/alloc_dealloc_pages |
35.88 µs | 36.44 µs | +1.5% |
| ✅ | matmul/small_generic_f16_threads=8/1x256x256 |
41.06 µs | 41.63 µs | +1.4% |
| ✅ | sampling_latency/greedy_per_token |
3.26 µs | 3.29 µs | +1.2% |
| ✅ | gather/large_f16_threads=1-internal/131072 |
16.60 µs | 16.78 µs | +1.1% |
| ✅ | matmul/small_generic_f16_threads=1/1x256x256 |
36.63 µs | 36.58 µs | -0.1% |
| ✅ | qwen3_sampling_processors/top_k_top_p_fast |
645.25 µs | 636.08 µs | -1.4% |
| ✅ | logit_processing/seven_processor_chain_per_step |
303.53 µs | 298.62 µs | -1.6% |
| ✅ | add/medium_f32_threads=1-internal/262144 |
2.93 ms | 2.85 ms | -2.6% |
| ✅ | grammar_masking/llguidance_compute_mask/32 |
71.49 µs | 68.83 µs | -3.7% |
| ✅ | gather/large_bf16_threads=1-internal/131072 |
16.41 µs | 15.53 µs | -5.4% |
| ✅ | matmul/large_generic_f32_threads=1/32x1024x1024 |
9.83 ms | 9.22 ms | -6.2% |
| ✅ | gather/large_f32_threads=1-internal/131072 |
46.30 µs | 42.29 µs | -8.7% |
| ✅ | matmul/large_generic_bf16_threads=1/32x1024x1024 |
2.16 ms | 1.96 ms | -9.4% |
| ✅ | matmul/medium_generic_f32_threads=1/32x512x512 |
2.61 ms | 2.33 ms | -10.6% |
| ✅ | matmul/small_generic_f32_threads=1/1x256x256 |
50.38 µs | 43.84 µs | -13.0% |
| ✅ | matmul/large_generic_f16_threads=1/32x1024x1024 |
93.16 µs | 80.29 µs | -13.8% |
| 🟢 | reduce_mean/small_f32_threads=1-internal/4096 |
18.07 µs | 15.04 µs | -16.8% |
| 🟢 | matmul/small_generic_bf16_threads=1/1x256x256 |
44.70 µs | 36.49 µs | -18.4% |
| 🟢 | gather/medium_f32_threads=1-internal/32768 |
6.08 µs | 4.96 µs | -18.4% |
| 🟢 | matmul/medium_generic_bf16_threads=1/32x512x512 |
624.37 µs | 503.49 µs | -19.4% |
| 🟢 | matmul/large_generic_f16_threads=8/32x1024x1024 |
103.97 µs | 83.09 µs | -20.1% |
| 🟢 | add/small_bf16_threads=1-internal/1024 |
15.20 µs | 11.95 µs | -21.3% |
| 🟢 | matmul/small_generic_bf16_threads=8/1x256x256 |
41.88 µs | 32.73 µs | -21.9% |
| 🟢 | add/medium_bf16_threads=1-internal/262144 |
3.22 ms | 2.50 ms | -22.3% |
| 🟢 | matmul/medium_generic_f16_threads=1/32x512x512 |
40.57 µs | 31.42 µs | -22.5% |
| 🟢 | add/large_bf16_threads=1-internal/4194304 |
50.19 ms | 38.86 ms | -22.6% |
| 🟢 | add/small_f16_threads=1-internal/1024 |
16.25 µs | 12.51 µs | -23.0% |
| 🟢 | gather/small_bf16_threads=1-internal/4096 |
655.8 ns | 496.4 ns | -24.3% |
| 🟢 | add/large_f32_threads=1-internal/4194304 |
51.05 ms | 38.27 ms | -25.0% |
| 🟢 | add/large_f16_threads=1-internal/4194304 |
51.69 ms | 38.36 ms | -25.8% |
| 🟢 | add/medium_f16_threads=1-internal/262144 |
3.40 ms | 2.51 ms | -26.2% |
| 🟢 | gather/medium_f16_threads=1-internal/32768 |
3.27 µs | 2.42 µs | -26.2% |
| 🟢 | reduce_mean/medium_f32_threads=1-internal/65536 |
342.71 µs | 252.73 µs | -26.3% |
| 🟢 | reduce_mean/large_f32_threads=1-internal/262144 |
1.47 ms | 1.05 ms | -28.1% |
| 🟢 | gather/medium_bf16_threads=1-internal/32768 |
3.43 µs | 2.37 µs | -31.0% |
| 🟢 | gather/small_f16_threads=1-internal/4096 |
689.7 ns | 473.7 ns | -31.3% |
| 🟢 | matmul/medium_generic_f16_threads=8/32x512x512 |
43.85 µs | 29.17 µs | -33.5% |
| 🟢 | matmul/large_generic_f32_threads=8/32x1024x1024 |
6.60 ms | 4.23 ms | -35.9% |
| 🟢 | gather/small_f32_threads=1-internal/4096 |
1.11 µs | 698.3 ns | -36.8% |
| 🟢 | matmul/large_generic_bf16_threads=8/32x1024x1024 |
2.39 ms | 1.47 ms | -38.2% |
| 🟢 | matmul/medium_generic_bf16_threads=8/32x512x512 |
688.09 µs | 416.91 µs | -39.4% |
| 🟢 | matmul/small_generic_f32_threads=8/1x256x256 |
79.60 µs | 47.29 µs | -40.6% |
| 🟢 | block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024 |
1.26 ms | 643.42 µs | -48.9% |
| 🟢 | add/small_f32_threads=1-internal/1024 |
355.3 ns | 178.8 ns | -49.7% |
| 🟢 | matmul/medium_generic_f32_threads=8/32x512x512 |
1.82 ms | 913.06 µs | -49.8% |
Visual flags:
Host info
CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 2.50 3.53 7.51 }
What this cannot catch
- Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
- Sub-threshold regressions that compound over multiple PRs
- Performance changes that only manifest under GPU execution
- Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)
Owner
Author
|
Verified the hardening myself before merging. The decisive test now reads: write_host(base, 0x77, PROBE_LEN);
assert_device_bytes(base, 0x77, "fresh mapping before replay");
graph.replay(stream);
assert_device_bytes(base, GRAPH_VALUE, "replay after full remap");That closes the page-recycling confound: a pass now means the replay overwrote memory that demonstrably belonged to the new mapping, whichever physical pages back it. Previously the assertion was only consistent with the conclusion — the driver had just released those pages and The added diagnostic that reports whether the fresh mapping came back recycled, all- What this establishes
The two "untested" rows are labelled as such rather than guessed at. That is what makes the tested rows worth relying on. What it unblocks#721 stage 4 — native CUDA KV grown by committing granules under a reserved, stable virtual address, so kernels see one flat contiguous pointer and paging happens underneath. Its stated falsifier was "if a captured graph does not survive a KV growth event under VMM, the central premise here is wrong and stage 4 should stop". The falsifier did not fire. #721 stage 5 — cross-sequence prefix sharing by mapping shared physical handles into multiple virtual ranges, with CoW diverging at granule granularity. #716 — the same mechanism is the weight-side counterpart. Capture is refused today because "the pager's alloc/copy/free ops are capture-illegal": allocation and copy are illegal inside a captured region, and each page-in hands back a new pointer that goes stale on eviction. A stable VA removes both, and moves map/unmap outside the captured segment. See #716 (comment) — that issue should be re-scoped from application-layer "stable slots" to this. Graph capture measures 2.6x in isolation with non-overlapping ranges, and is currently disabled on every offload run. That is the size of what this unblocks. Merging. |
justinchuby
added a commit
that referenced
this pull request
Aug 7, 2026
Consolidating today's findings into the document rather than leaving them in issues. ## The new section: a third failure mode `MEMORY_ARCHITECTURE.md` already recorded two ways this area fails — code that is right, tested and never reached; and measurements that are reproducible but support a conclusion they do not license. Today produced a third, and it cost more than either. **An instrument that runs, reports a plausible number, and measures the wrong thing.** `h2d_enqueue_copy_ms` put host-to-device transfer at **1.7%** of decode step time. That figure retired an entire line of work — prefetch depth, pinned versus pageable staging, transfer overlap — and was published on #705 and #718. The arithmetic gives it away: the same step stages **7.33 GiB**, and 7.33 GiB in 30 ms would be ~250 GB/s, roughly ten times this machine's PCIe link. The counter bracketed an asynchronous enqueue *returning*. The first repair renamed it `copy_wait_ms`, pointing it at `compute_wait_fence` — whose own doc comment reads *"a stream-ordered, **non host-blocking** cross-stream wait."* The same non-measurement survived under a name that invited the wrong reading **more** strongly than before. Measured properly with CUDA events, transfer is **18.8%**. Two habits, both cheap: - **Divide the bytes by the time and ask whether the answer is a real bandwidth.** A number near link speed, near memcpy speed, or ten times either tells you what the counter actually bracketed. This is what caught it, and it takes seconds. - **State on every timing counter exactly what lies between start and stop, and whether anything there blocks the host.** Enforced in `weight_paging.rs` since #715. ## The pattern underneath all three Telemetry that never reaches the operator, three times in one subsystem: documented env vars nothing reads (#688), counters computed and never printed (#719), counters printed with nothing writing them (#715). The #719 case is the sharpest. A residency cache paging at a **0% hit rate** printed a weight-offload section in which **every visible row read `0.00`** — and my first reading was that the cache was inert. The opposite was true and provably so: the section printing at all required one of the three *hidden* counters to be non-zero. The instrument reported the exact inverse of what it measured. After a third instance the response should be a guard, not more vigilance. ## Status table rows where today changed the facts - **`--vram-limit`** — was "does not limit VRAM"; now enforced at load, auto-enabling offload with the weight budget derived as `limit − device state`. Caveat carried rather than closed: activations and runtime overhead are still `unknown` and unsubtracted (#514). PR #717 - **Weight residency cache (CUDA)** — new row. It had a **0% hit rate for its entire life**: 6,936 page-ins, 0 hits, at both 3 GB and 6 GB, because `evict_to_fit` is LRU against a cyclic sequential scan. Fixed to **74.18%** against a `B/W ≈ 76%` ceiling, evictions 6,286 → **0**. #720, PR #723 - **Weight offload vs the OS** — the 27x gap now has a named mechanism instead of a mystery, with ~39% of step time placed downstream of that one defect. - **KV page storage** — new row. `onnx-genai-kv` never touched device memory; the contract is now open (`Box<dyn KvPageStore>`), host-only pending stages 2–5. PR #726 - **Captured graphs across a VMM remap** — new row, verified on hardware and sentinel-proven. PR #727 Docs only. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
This was referenced Aug 10, 2026
justinchuby
pushed a commit
that referenced
this pull request
Aug 11, 2026
… mechanism rigor Reorders per owner priority: page-level prefix sharing (#777) is now co-primary with the ORT BSNH crux. - Byte saving computed across P in {512,2K,8K} x R in {2,4,8,16} for qwen14b and qwen0.5b, ideal and granule-rounded. Confirms #777 headline (2048 tok x 8 on qwen14b = 2.62 GiB removed). - New rigor finding: page sharing has its own granule floor; it saves nothing when the per-(layer,side) shared stripe is sub-granule (P<~1K tok on 14B, <8K on 0.5b). Large-model/long-prefix win, not universal. - Capacity consequence tabled: qwen14b 2->6 GiB KV budget, 2560-tok ctx w/ 2048 shared prefix goes 3->8 .. 10->30 concurrent requests (2.7-3.0x). Confirms the '2 vs 8' claim, honestly. - Mechanism confirmed against the REAL allocator/ledger (not a mock): multi-map proven (#727), granule_refs charge-once + refcount-gated release (#740/#745); the only gap is 1:N handle bookkeeping (blocks Vec / single arena reservation). cuMemSetAccess is per-subrange so a PROT_READ shared prefix composes with the #759 PROT_NONE dummy page without conflict. - Head-major impossibility confirmed AND quantified: needs P >= granule/ (head_dim*dtype) = 8192/16384 tok/head-stripe to share one granule; seq-major divides that by kv_heads. Structural, no head-major grouping recovers it. - Recommendation and next-increment updated to bank the #777 prize. References #777 alongside #750 and #759. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This was referenced Aug 11, 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
## Probe: the #759 "dummy page mapping" primitive for VMM KV cache Refs #759. **Self-contained GPU probe in `crates/onnx-runtime-cuda-memory` — no engine and no KV-path change.** Establishes or kills the primitive at the driver level *before* KV integration, extending `vmm_graph_remap_gpu` (#727) and `vmm_kv_contiguous_tail_gpu` (#772: a padded read one byte into the uncommitted tail faults `CUDA_ERROR_INVALID_VALUE`). Do **not** merge, do **not** close #759. ### Fill value — masking decides it, and it forbids NaN (corrected) An earlier revision filled the dummy with a NaN "sentinel" for detectability. **That is wrong.** #721 stage 4's decode kernel reads the *full padded shape* and relies on **additive masking** (`score -> -inf` past the live length). NaN defeats additive masking — `q·NaN = NaN`, `NaN + (-inf) = NaN`, `exp(NaN) = NaN` — so a NaN dummy poisons the softmax output masking would otherwise have made correct, **including the live positions**. `dummy_fill_and_crossover::masking_determines_the_safe_dummy_fill` measures the rule (host-side IEEE-754, deterministic): | Case | Measured output | Verdict | |---|---|---| | masking effective + **zero** tail | `1.0000` (tail weight 0) | zeros safe | | masking effective + **NaN** tail | `NaN` | **NaN forbidden** | | masking **not** effective + zero tail | `0.8808` (tail weighted `exp(0)=1`) | even zeros wrong | **Rule: fill with ZERO once the decode kernel's tail masking is verified effective in the EP; never NaN. If masking is unverified, the dummy page is crash-safe only and the kernel must be bounded.** Whether *our* production decode kernel masks those tail positions is an EP fact this memory-crate probe cannot exercise (no attention kernel here) — it must be verified there before zeros are trusted for correctness. The driver-level tests therefore write only an arbitrary `0x5a` readback marker (a test probe, not the production fill). ### The five questions (measured on this hardware, granule = 2 MiB) | # | Question | Test | Result | |---|----------|------|--------| | 1 | Removes the fault? | `vmm_dummy_page_gpu::dummy_backed_tail_removes_the_padded_read_fault` | **Yes.** One dummy handle across the tail turns #772's faulting padded read *and a captured-graph copy that reads the whole padded range* into a success. | | 2 | Fill choice / detectability | `dummy_fill_and_crossover::masking_determines_the_safe_dummy_fill` | **Zeros, never NaN** (see table above). Detectability via NaN is rejected because it breaks the masked case. | | 3 | Write-protection non-sticky? | `vmm_dummy_write_protect_gpu::read_only_dummy_write_fault_is_non_sticky` (own binary) | **Non-sticky — viable.** A `PROT_READ` dummy rejects a synchronous `cuMemcpyHtoD_v2` **and** an async `cuMemsetD8Async` with `CUDA_ERROR_INVALID_VALUE` (validated at enqueue); context survives both, confirmed by a fresh allocate/write/read. | | 4 | Ledger charges once? | `vmm_dummy_page_ledger_gpu` | **Yes.** *Part A:* one handle at **96 VAs = one 2 MiB page** (per-VA would be 192 MiB), all aliasing. *Part B:* the **real** allocator + `LedgerGovernor` charge a shared granule once (2 MiB for 128 allocations), drop the mapping refcount to zero on free, release the handle on teardown (governor → 0). | | 5 | Growth cost? | `vmm_dummy_page_gpu::dummy_to_real_growth_remap_cost` | **Expensive.** Single dummy→real remap **≈ 450–500 µs**; a **96-object fixed-stride token growth step ≈ 30–35 ms = ≈ 300–360 µs/object.** This is the design's main performance risk. | ### Crossover is model-size independent — and the granule is its only lever The crossover between fixed-stride+dummy and bucket growth reduces to a closed form that **does not depend on model size**: ``` objects = layers × 2 × kv_heads kv_bytes/token = objects × head_dim × elem_bytes crossover_tokens = objects × granule / kv_bytes_per_token = granule / (head_dim × elem_bytes) # layers and kv_heads cancel ``` `crossover_is_model_size_independent` verifies the cancellation against a deliberately 480× size range (4 layers/1 head vs 120 layers/16 heads — identical crossover) and both real models. So the crossover is a property of the **layout unit** (`head_dim`, KV dtype, granule), not of the model. It scales **linearly with the granule**, making the granule the single highest-value lever. **Granularity measured on-device (`vmm_granularity_gpu::report_both_allocation_granularities`):** > **`CU_MEM_ALLOC_GRANULARITY_MINIMUM` == `CU_MEM_ALLOC_GRANULARITY_RECOMMENDED` == 2 MiB on this hardware.** The hoped-for finer granule (e.g. 64 KiB, which would collapse the crossover ~32× into the low hundreds of tokens) is **not exposed on this box**. The lever cannot be pulled here; the design is stuck at the 2 MiB crossover. `granularity_cost_minimum_vs_recommended` consequently commits and reads identically at both flags (~2.4 ms / 32 MiB region either way, ~105 GiB/s read both) — because both flags return the same granule. On a device that *does* expose a finer minimum, the same binary would price the trade-off (32× more `cuMemMap`/`cuMemUnmap` calls, TLB pressure). **Crossover table (tokens) = `granule / (head_dim × elem_bytes)`, model-size independent** (`granularity_is_the_crossover_lever`): | granule | hd64 fp16 | hd64 fp8/int8 | hd128 fp16 | hd128 fp8/int8 | |---|---|---|---|---| | **2 MiB (MIN==REC here)** | 16 384 | 32 768 | 8 192 | 16 384 | | 64 KiB (*counterfactual — not on this device*) | 512 | 1 024 | 256 | 512 | fp8/int8 KV halves bytes/token, so it **doubles** the crossover in tokens. At the realizable 2 MiB granule the crossover is **8K–16K tokens** — a high bar that loses fixed-stride+dummy against bucket growth in most serving windows. **Per-model crossover at 2 MiB** (`fixed_stride_plus_dummy_crossover_vs_bucket_growth`, geometry read from each `genai_config.json`): | Model | objects = layers×2×kv_heads | per-head stride | floor | crossover | context | |---|---|---|---|---|---| | qwen2.5-0.5b | 24×2×2 = **96** | 4 MiB (2 granules) | 192 MiB | **16 384 tok** (½ ctx) | 32 768 | | qwen14b | 48×2×8 = **768** | **2 MiB (1 granule)** | 1536 MiB | **8 192 tok** (full ctx) | 8 192 | - **qwen2.5-0.5b:** crossover at **half** its 32 768 context — fixed-stride+dummy wins only in the upper half of long-context serving. - **qwen14b (the #721 stage-4 1.5 GB model):** per-head stride is **exactly one granule**, so there is **no per-head uncommitted tail** — the dummy page saves **no memory** (floor == honest full commit == 1.5 GB); its only benefit for 14b is **stride stability (no re-capture)**. The reviewer's optimistic "1.5 GB → 192 MiB" conflated two models. ### Minimum-granule dummy pattern works (when a finer granule exists) `minimum_granularity_dummy_pattern_works_including_captured_replay` builds the #759 primitive at whatever the driver reports as MINIMUM (2 MiB here): one dummy handle aliased across 64 tail granules, a full-padded read, and a captured-graph device-to-device read into the dummy tail — all succeed and read back correctly. So *if* a device exposed a finer granule, the pattern would carry it; the blocker on this box is that no finer granule is offered. ### What this does **not** change The dummy page does **not** reduce the `objects × granule` floor. #772 measured 96×2 MiB = **192 MiB for ~12 KiB live content** (0.5b, head-major fixed stride); those granules hold **real, distinct live data**. Anyone reading #759 as a fix for the 32× regression is mistaken. ### Accounting caveat for integration (Q4) The production commit path maps a **distinct** handle per granule; a genuine dummy integration maps **one** handle at N VAs, so `mapped_bytes` (VA coverage) would count each dummy VA while `total_owned_bytes` (physical) counts the one handle once. Charge on the **owned/physical** axis — which the governor already bills against `Tier::Device`. ### Ordering rules relied on (per #727) - The dummy→real remap in Q5 runs with **no graph replay in flight** (#727 did not prove unmap-under-replay safe); in production it must precede the step's captured launch, never run concurrently. - No `cuMemMap`/growth **inside** a captured region. ### Recommendation The primitive is **sound as a fault-safety net**: it removes the tail fault (Q1), its write-protection is non-sticky (Q3), and it charges once and frees cleanly against the real ledger (Q4). But three constraints bound where it is worth building: 1. **Fill correctness (Q2):** only zeros, and only after tail-masking is verified in the decode kernel; NaN is forbidden. Until masking is verified, the dummy page is crash-safe, not correctness-safe. 2. **Growth cost (Q5):** remapping ~96–768 granules per token under a fixed full-context stride costs ~30–35 ms/step. 3. **Granule (crossover):** the memory win only appears above the crossover, which at the realizable **2 MiB** granule is 8K–16K tokens; **the finer granule that would lower this is not exposed on this hardware.** For 14b it saves no memory at all. **Do not build it under a fixed full-context head-major stride** (no memory win for 14b; per-token remap cost prohibitive; crossover too high at the only available granule). **Build it, if at all, under a *bucketed* growable path** — one contiguous dummy-backed VA per sequence, real granules committed **per bucket** (remaps rare, only at bucket boundaries), fill = zeros gated on verified masking, dummy `PROT_READ`. **Smallest next increment:** (a) verify in the CUDA EP that the decode attention kernel actually masks positions `>= live_len` (this decides zeros-vs-bounded-kernel); (b) prototype a single bucketed dummy-backed KV buffer and measure end-to-end committed bytes and per-token step time against bucket growth, gating on remaps-per-token staying O(few). Re-run `vmm_granularity_gpu` on any target GPU whose MINIMUM granule is finer than 2 MiB before betting on the fixed-stride variant there. ### Validation run locally - `cargo fmt -p onnx-runtime-cuda-memory -- --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` ✅ (exact CI gate) - GPU: `vmm_dummy_page_gpu` (Q1, Q5), `vmm_dummy_page_ledger_gpu` (Q4 ×2), `vmm_dummy_write_protect_gpu` (Q3), `vmm_granularity_gpu` (granularity report + min-granule dummy pattern + cost), plus existing `vmm_graph_remap_gpu` ✅. CPU: `dummy_fill_and_crossover` (masking rule + crossover + model-size-independence + granularity table) ✅. (`vmm_kv_contiguous_tail_gpu` lives on #772's branch, not on `main`.) --------- 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
## 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>
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
onnx-runtime-cuda-memoryfor Split KV page mapping from page storage so device-resident paged KV becomes possible #721.0x77and asserting that readback before replay, so the decisive assertion cannot be explained by physical-page recycling.cuMemMapbehavior during thread-local capture and validates one physical handle mapped at two virtual addresses.Findings
Verdict: a graph instantiated before a VMM remap replayed successfully afterwards and overwrote a fresh post-remap sentinel in the new mapping on this RTX 4060 Laptop/WDDM setup.
Observed CUDA results in the executed test:
0x77; after replay,cuGraphLaunchand post-launchcuStreamSynchronizereturnedCUDA_SUCCESSand readback was0x5a. The pre-sentinel read before writing was all zero (all 0x5a = false; all zero = true) on this run.CUDA_SUCCESS; original granule was rewritten to0x5a; added granule remained mapped and unchanged (0x77).cuMemMapduringCU_STREAM_CAPTURE_MODE_THREAD_LOCALcapture returnedCUDA_SUCCESS;cuStreamEndCapturealso returnedCUDA_SUCCESS. This does not prove the map effect is a replayable graph node.Synchronization note: the probe synchronizes the graph stream before unmap/remap. I did not run the negative unmapped-replay direction because it can plausibly raise an illegal-address fault and poison the CUDA context/process; that is not worth putting in this integration test.
Validation
cargo fmt --all -- --checkcargo clippy --locked --all-targets -p onnx-runtime-cuda-memory -- -D warnings$env:NXRT_CUDA_WHEEL_ROOTS="C:\Users\justinchu\AppData\Local\anaconda3\Lib\site-packages"; PATH prepended with ORT/CUDA/CUDNN wheel dirs;cargo test --locked -p onnx-runtime-cuda-memory --features gpu-tests --test vmm_graph_remap_gpu -- --nocapture(4 passed, actually executed on GPU; rebuilttarget\debug\deps\vmm_graph_remap_gpu-457b4a0d34ec31c4.exeimmediately before running)Closes #721? No — this is only the decisive experiment for stage 4.