[justinchuby]

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 Pool CUDA VMM physical handles across reservations #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 KV layout should be a stride descriptor negotiated per EP, not a growing enum — and the choice is not memory-neutral #783/Probe: token-major-across-all-layers KV — TLB risk cleared, floor measured (#750, #777, #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.

[justinchuby]

Coordinator validation on the merged branch (merged current main, no
conflicts): cargo fmt -p onnx-runtime-cuda-memory -- --check clean, and all
four new GPU test binaries pass on device, single-threaded.

Scope verified: probe only — 4 test files plus 2 docs, zero production code.
Splitting one question per test binary was the right call, since a sticky CUDA
error in one probe would otherwise poison the others.

The result

All five questions cleared, and the one that could have killed the design did not
fire.

Q3 was the kill risk. The shared prefix region is read-only by construction,
so cuMemSetAccess(CU_MEM_ACCESS_FLAGS_PROT_READ) is what stops one request
silently corrupting another's KV. But if that fault were sticky, the
"protection" would destroy the CUDA context and be worse than the corruption it
prevents. Measured: the violation faults loudly, non-stickily and
non-corruptingly, on both the synchronous and asynchronous paths.

Q2 is the one that makes the economics real. Charged once, alive until the
last sharer releases, and 0 incremental bytes per additional sharer —
against the real CudaVmmAllocator and LedgerGovernor, not a mock. This is
the most extreme sharing case the 0->1 / 1->0 granule-ref attribution has ever
been asked to handle, and it holds. The consequence for admission is direct: the
Nth concurrent request costs only its private bytes.

Q1 confirms #727's "one handle at two VAs" generalizes to N=8 sequences, each
reading identically, including under captured-graph replay — which is the
decode hot path, not a synthetic case.

Q4 answers the composition risk I flagged early: the read-only shared prefix,
the read/write live region, and the read-only #759 dummy tail all coexist in one
reservation. Those three want cuMemSetAccess semantics on the same address
space, and finding a conflict after integration would have been expensive.

Q5 puts copy-on-write in proportion: ~5.5 ms pooled / ~9.8 ms cold, once
at divergence — a one-time burst, not a per-token tax.

Why this matters more than the rest of the line

Every other lever here moves bytes by a percentage. This removes a whole
multiple, and it scales with concurrency rather than context length — the
2K-token / 8-way qwen14b cell is 2.6 GiB, and the capacity consequence is
roughly doubled admissible concurrency. On an 8 GiB card that is the
difference between 2 and 8 concurrent requests, which is a capability change, not
a tuning gain.

It also compounds with the layout work rather than competing with it: a shared
prefix is layers x 2 x kv_heads scattered fragments under head-major,
layers x 2 contiguous ranges under seq-major (#782), and one contiguous
range covering all layers under token-major (#787). So the cheapest possible
implementation is also the one the layout direction was already heading toward.

Recommendation accepted. The next increment is an explicit commit_shared_prefix()
entry point on the token-major reservation model, with the explicit pinned-prefix
API rather than hash-based detection — unambiguous, and it covers the
highest-value scenario (a shared system prompt) without guessing.

Merging so the primitive's guarantees become executable guards.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.42%. Comparing base (26f2657) to head (69da9f9).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #793      +/-   ##
==========================================
- Coverage   80.08%   79.42%   -0.66%     
==========================================
  Files         331      331              
  Lines      137782   137782              
  Branches   137782   137782              
==========================================
- Hits       110337   109437     -900     
- Misses      23092    23997     +905     
+ Partials     4353     4348       -5     

Flag	Coverage Δ
cli-ort-linux	83.50% <ø> (ø)
cli-ort-windows	82.99% <ø> (-0.10%)	⬇️
mlas	80.73% <ø> (-0.56%)	⬇️
offline	79.24% <ø> (-0.69%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 8 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

🔴 Benchmark Regression Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
🔴	gather/large_bf16_threads=1-internal/131072	12.25 µs	22.53 µs	+84.0%
🔴	gather/large_f16_threads=1-internal/131072	11.79 µs	20.08 µs	+70.3%
🔴	gather/small_f16_threads=1-internal/4096	528.8 ns	781.1 ns	+47.7%
🔴	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	90.16 µs	128.89 µs	+43.0%
⚠️	add/medium_bf16_threads=1-internal/262144	2.65 ms	3.40 ms	+28.2%
⚠️	gather/medium_f16_threads=1-internal/32768	2.44 µs	3.13 µs	+28.2%
⚠️	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.75 ms	7.01 ms	+21.9%
⚠️	matmul/small_generic_f32_threads=1/1x256x256	39.56 µs	48.14 µs	+21.7%
⚠️	qwen3_sampling_processors/top_k_top_p_fast	669.78 µs	798.77 µs	+19.3%
✅	gather/large_f32_threads=1-internal/131072	30.04 µs	33.86 µs	+12.7%
✅	add/large_f32_threads=1-internal/4194304	45.91 ms	51.33 ms	+11.8%
✅	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	1.14 ms	1.26 ms	+10.6%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	523.25 µs	577.07 µs	+10.3%
✅	tokenization/decode_tokens_per_second	6.40 ms	6.98 ms	+9.1%
✅	gather/medium_f32_threads=1-internal/32768	4.00 µs	4.35 µs	+8.9%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.02 ms	1.11 ms	+8.1%
✅	sampling_latency/top_p_per_token	378.71 µs	408.61 µs	+7.9%
✅	add/large_f16_threads=1-internal/4194304	42.73 ms	45.77 ms	+7.1%
✅	gather/medium_bf16_threads=1-internal/32768	2.42 µs	2.59 µs	+6.8%
✅	gather/small_f32_threads=1-internal/4096	753.2 ns	798.9 ns	+6.1%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	538.21 µs	570.71 µs	+6.0%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	78.04 µs	82.59 µs	+5.8%
✅	kv_cache/alloc_dealloc_pages	40.51 µs	42.58 µs	+5.1%
✅	sampling_latency/top_k_per_token	51.90 µs	54.43 µs	+4.9%
✅	sampling_latency/greedy_per_token	3.24 µs	3.39 µs	+4.5%
✅	tokenization/encode_tokens_per_second	376.26 µs	391.75 µs	+4.1%
✅	gather/small_bf16_threads=1-internal/4096	499.2 ns	519.6 ns	+4.1%
✅	matmul/small_generic_bf16_threads=1/1x256x256	35.94 µs	37.29 µs	+3.7%
✅	logit_processing/seven_processor_chain_per_step	323.91 µs	330.27 µs	+2.0%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.12 ms	2.14 ms	+1.1%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.59 ms	3.61 ms	+0.7%
✅	add/large_bf16_threads=1-internal/4194304	44.69 ms	44.98 ms	+0.6%
✅	add/small_bf16_threads=1-internal/1024	13.09 µs	13.17 µs	+0.6%
✅	reduce_mean/small_f32_threads=1-internal/4096	15.72 µs	15.75 µs	+0.2%
✅	sampling_latency/min_p_per_token	209.96 µs	209.80 µs	-0.1%
✅	grammar_masking/llguidance_compute_mask/32	79.00 µs	78.74 µs	-0.3%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	9.65 ms	9.50 ms	-1.5%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	85.60 µs	83.88 µs	-2.0%
✅	qwen3_sampling_processors/top_k_partial_selection	143.21 µs	135.90 µs	-5.1%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.36 ms	2.23 ms	-5.5%
✅	matmul/small_generic_f16_threads=1/1x256x256	35.55 µs	33.49 µs	-5.8%
✅	matmul/large_generic_bf16_threads=8/32x1024x1024	2.26 ms	2.10 ms	-6.8%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.01 ms	1.84 ms	-8.3%
✅	add/medium_f16_threads=1-internal/262144	3.36 ms	3.03 ms	-10.0%
✅	reduce_mean/medium_f32_threads=1-internal/65536	269.64 µs	240.85 µs	-10.7%
✅	matmul/medium_generic_f16_threads=1/32x512x512	35.93 µs	32.01 µs	-10.9%
✅	matmul/small_generic_f16_threads=8/1x256x256	36.17 µs	31.94 µs	-11.7%
✅	matmul/small_generic_bf16_threads=8/1x256x256	38.53 µs	33.40 µs	-13.3%
✅	matmul/medium_generic_f16_threads=8/32x512x512	41.29 µs	35.76 µs	-13.4%
✅	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	519.22 µs	449.17 µs	-13.5%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	101.89 µs	87.63 µs	-14.0%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	667.82 µs	565.47 µs	-15.3%
🟢	add/small_f16_threads=1-internal/1024	15.38 µs	12.71 µs	-17.3%
🟢	add/medium_f32_threads=1-internal/262144	3.05 ms	2.40 ms	-21.3%
🟢	matmul/small_generic_f32_threads=8/1x256x256	51.13 µs	39.67 µs	-22.4%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.76 ms	1.35 ms	-22.9%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	6.93 ms	5.04 ms	-27.3%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	125.59 µs	76.41 µs	-39.2%
🟢	add/small_f32_threads=1-internal/1024	377.7 ns	195.4 ns	-48.3%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 3.25 3.26 5.42 }

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)