Repository navigation
Give the prefix-sharing primitive a first production consumer (#777) - #809
Merged
Merged
Conversation
The merged prefix-sharing primitive (create_shared_prefix / commit_shared_prefix / incremental_owned_bytes_for_shared_prefix on CudaVmmAllocator, #803) had no live caller in production — only its definition and GPU tests. An unconsumed contract reads as working infrastructure while proving nothing, the same shape that stranded #721 stage 3. This wires a first real, deliberately narrow consumer. - Add 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. Default impls refuse (InvalidRequest) rather than mis-map, so every existing non-VMM allocator keeps compiling and refuses safely. - Implement the seam on CudaVmmAllocator by delegating to the merged inherent methods and downcasting the opaque handle via as_any(). Explicit, read-only-prefix API — not hash-based detection; divergence /COW out of scope. - First consumer: the real seq-major fused fp16 GQA decode kernel reads shared-prefix VMM KV through the `dyn DeviceAllocator` seam. New GPU parity test gqa_shared_prefix_parity_gpu proves two sequences sharing one pinned seq-major prefix (layers x 2 contiguous ranges) 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): - independent = 8 granules (16,777,216 B); shared = 6 granules (12,582,912 B); sharing removes 2 granules = (C-1) x (K_prefix + V_prefix). - prefix charged once: incremental_owned_bytes_for_shared_prefix = 0. - second sharer admission = private bytes only (4,194,304 B). - output byte-identical for both sequences. Remaining (structural / out of scope): the engine generation loop cannot call the seam automatically until native_decode/cuda.rs can build a seq-major fixed-stride physical KV shape (persistent_state_shapes is hard-coded BNSH; #794 showed seq-major changed only kernel indexing). Also later: hash-based detection, token-major (one multi-map/sequence), copy-on-write at divergence (~5.5 ms pooled / ~9.8 ms cold, once). Refs #777 #803 #793 #782 #745 #750 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #809 +/- ##
=======================================
Coverage 83.55% 83.55%
=======================================
Files 12 12
Lines 5192 5192
Branches 5192 5192
=======================================
Hits 4338 4338
Misses 677 677
Partials 177 177
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
Owner
Author
|
Coordinator pre-merge validation on the merged branch: The primitive now has a production consumerI flagged this two hours ago as the thing that mattered: #803 landed The numbers match the geometry exactly, which is what I wanted to see:
The seam is the right shape for an interchangeable-backend architecturePutting
So an unsupported request faults loudly at the seam instead of silently producing The consumer itself drives the real seq-major fused fp16 GQA decode kernel over What remains, correctly reportedEngine auto-use is blocked until Merging. |
🔴 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
|
This was referenced Aug 12, 2026
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
…ajor-only KV consumers (#777 #801 #797 #794 #809) (#812) # Remove the last blocker to seq-major KV being usable in production: gate the fourth place the layout lives **Outcome reached: #3 — "it works only under a condition."** Seq-major decode works end-to-end and is byte-identical to head-major, but only on the **pure decode path where our own GQA kernel is the sole consumer of the device KV buffer**. This PR makes the runtime **refuse rather than mis-map** when that condition does not hold, and pins down — with citations — exactly what remains between "growth byte geometry is seq-major" (#801) and "the `layers × 2` floor realised + prefix sharing available." Refs #777 #797 #794 #801 #809 (and #787 for token-major). Does **not** merge, does **not** close any issue. ## The question this PR answers > What remains between "growth byte geometry is seq-major" and "a real model runs seq-major end to end with the floor realised and prefix sharing available"? A fourth place the layout lives, or can the declared BNSH shape and a seq-major physical layout coexist provided only our own kernels read the buffer? **Both.** The declared BNSH shape and a seq-major physical layout *do* coexist on the native decode path (KV never leaves the device; only our GQA kernel indexes it). But there **is a fourth place the layout lives** — two engine-side *non-kernel* consumers of the device KV buffer that still stride it with hard-coded head-major (BNSH) arithmetic: 1. **`DecodeCudaState::read_present_kv`** (`native_decode/cuda.rs`) — the device present-KV **host mirror** used by prefix/paged read-out and diagnostics. It compacts per head with `capacity_head_stride = physical_shape[2] * head_dim * elem` and a per-head fragment read. 2. **`DecodeCudaState::seed_prefix`** (`native_decode/cuda.rs`) — the **paged prefix-reuse device seed** (#777/#809's device counterpart). It writes each head at `head * capacity_head_stride * elem_size`, the same head-major capacity offset. Under a seq-major (BSNH) physical buffer, heads are **interleaved per token** (`slot * kv_heads * head_dim + h * head_dim + d`), not laid out as `max_len × head_dim` capacity stripes. So both offsets index the wrong bytes: the host mirror would **silently return mis-indexed KV**, and the prefix seed would **corrupt the cache**. Neither had a seq-major guard — a real latent-corruption path. This is the fourth place the layout lives, beyond the three that are already converted: 1. kernel indexing — done (#792/#794); 2. commit geometry — done (#797); 3. growth byte geometry — done (#801); 4. **engine-side non-kernel KV consumers — gated here.** ## What this PR changes `crates/onnx-genai-engine/src/native_decode/cuda.rs`: - `read_present_kv` and `seed_prefix` now **error under `kv_layout.is_seq_major()`** with a precise message stating the condition and why (head-major byte arithmetic vs interleaved seq-major bytes). This matches the kernel-level "unsupported readers/writers fail rather than silently mis-index" gate (`kv_stride.rs::require_converted_path_support`) at the engine level. - **Head-major is untouched and byte-identical** — the gate is a pure `if is_seq_major() { bail! }` before the existing code. Absent metadata ⇒ head-major ⇒ no gate. `docs/MEMORY_ARCHITECTURE.md` (§ "KV layout and residency"): records the fourth layout location and the refuse-not-mismap gate as a bullet in the existing residency-caveat block. ## The structural constraint (#801), re-verified The BNSH binding **metadata** shape must stay for both layouts. The CUDA GQA node validates `past_key`/`past_value` as `[batch, kv_heads, seq, head_dim]` and reads `present_capacity` from **axis 2 regardless of `kv_layout`** (`group_query_attention.rs`), and `DeviceIoBinding` requires matching axis order. `persistent_state_shapes` is therefore *correctly* BNSH for seq-major too; seq-major lives in the byte geometry, not a permuted shape. Confirmed still true after merging the concurrent GQA-workspace change. ## Why the `layers × 2` floor is still not realised through the engine (and what it needs) Measured directly, **before** the merge broke the harness, with the validated #801/#805 setup — stock `qwen2.5-0.5b-q4_0-mobius` (BNSH) vs a copy whose 24 GQA nodes carry `kv_layout=1`, forcing growth with `ONNX_GENAI_KV_MIN_BUCKET=8`, 48 greedy tokens, 3 growth events (`mobius_seqmajor_growth_parity_native_cuda`, this branch @ pre-merge): | Growth path | Capture | Output | Physical committed bytes (head = seq) | d2d KV copy bytes head | d2d KV copy bytes seq | |---|---|---|---:|---:|---:| | in-place VMM (default) | OFF | **byte-identical** | 100,663,296 | 688,576 | **0** | | in-place VMM (default) | ON | **byte-identical** | 100,663,296 | 688,576 | **0** | | legacy realloc | OFF | **byte-identical** | 786,432 | 688,576 | 688,576 | Capture-ON, per-process: head-major and seq-major both `(growth 3, captures 4, replays 39, invalidations 4)`. So on this branch the seq-major decode path is **byte-identical head vs seq, capture ON and OFF, and moves 0 KV bytes on in-place growth** — but **`physical_committed_bytes` is identical between layouts**, i.e. the floor is *not* realised through the engine. Two cited reasons: 1. **`vmm_growth_requests` commits the full bucket** (`(index, 0, full_bucket_bytes)`) for every KV binding, for both layouts, instead of the dense live-prefix ranges `native_decode/kv_commit.rs::live_prefix_ranges` computes (those helpers are unit-tested and mirror the driver GPU test but are `allow(dead_code)` — not wired into the engine). 2. **The `kv_heads×` difference only becomes *measurable* under a large fixed stride.** With bucket growth the bucket ≈ 2× valid, so on 0.5b both layouts land in **1 granule/binding** (48 granules = 96 MiB). The `kv_heads=2` gap appears only when a head stripe crosses a granule boundary, i.e. `capacity × head_dim × elem ≥ granule` (`ONNX_GENAI_CUDA_KV_MAX_LEN=32768` in the driver test `vmm_kv_layout_residency_gpu`): head-major = `layers × 2 × kv_heads` granules, seq-major = `layers × 2` (the measured 192 MiB → 96 MiB, 2× = `kv_heads`; 1536 MiB → 192 MiB, 8× on qwen14b). Realising the measurable floor therefore requires **a fixed full-context stride + dense-prefix commit + per-step commit-ahead**, which eliminates growth-driven bucket recapture. That change lives in the **growth/capture-invalidation decision** — the explicit lane of the concurrent `conditional-invalidation` / `bsnh-fixed-stride` work — and outside this PR's shape/binding-construction lane. It is left as the documented next step (shared with token-major #787 and prefix sharing #777), consistent with the existing doc note: *"wiring the dense-prefix commit into the live path is the documented next step."* ## Prefix sharing (#777/#809) under seq-major It does **not** fall out naturally. #809's device consumer is `seed_prefix`, a head-major byte writer; a shared prefix is only physically contiguous under seq-major, but the seed that would install it cannot yet address the seq-major buffer. This PR makes that explicit (the seed errors under seq-major) rather than letting prefix reuse silently corrupt seq-major KV. Realising it means teaching `seed_prefix` (and `read_present_kv`) the BSNH byte geometry — same dense-prefix work as the floor. ## Validation - `cargo fmt -p onnx-genai-engine -- --check` ✓ (`--all` fails with os error 206 on this box — a command-line-length artifact, not formatting) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` ✓ - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` ✓ (post-merge) - full `cargo test -p onnx-runtime-session` ✓ (post-merge) - `mobius_seqmajor_growth_parity_native_cuda` (byte-identical head vs seq, capture ON/OFF, 0 bytes moved) ✓ **both pre-merge and post-merge** (post-merge re-run includes the #814 governor fix, see below); the gate does not regress it. All four configs identical to the pre-merge numbers above (`physical_committed_bytes` 100,663,296 / 786,432; seq `kv_bytes_moved=0`). ### Cross-agent regression surfaced by `merge origin/main` — now fixed by #814 Merging `origin/main` (which brought the `gqa-bnsh-transpose` GQA transpose-scratch change #810) initially broke the mobius end-to-end test **during prefill workspace prep**, before any decode: ``` Error: prepare native CUDA prefill workspace Caused by: cannot reserve 2097152 bytes of device memory for Workspace { step_scoped: false }: 0 of 2097152 bytes are already leased, leaving 0 ``` #810 reserves a new **session-persistent** BNSH transpose-scratch `Workspace { step_scoped: false }` (`WorkspaceLifetime::SessionPersistent`) during prefill; the `vram-limit dynamic KV/weight lending` strategy did not provision headroom for it (GPU otherwise 100% free — 0 MiB used). This was a governor role-budgeting interaction in the **GQA/memory lane**, independent of this PR's gate (which only adds error paths in `read_present_kv`/`seed_prefix`, never on the prefill workspace path). I attributed it precisely and reported it to the `gqa-bnsh-transpose` sibling rather than fixing out of lane. **It has since been fixed forward in #814** (commit `54f89a03`: "mapped growth now claims unassigned authority capacity before holder reclaim, with rollback/overestimate accounting"). Merging #814 locally on top of this branch, the mobius test **passes all four VMM/realloc × capture configurations, byte-identical head vs seq**, confirming this PR's gate and #814's governor fix coexist cleanly. #814 is not yet on `main`, so this branch is left at its clean two-file state (the fix will arrive via `main`); the post-merge validation above was run against `this branch + #814` locally. ## Files - `crates/onnx-genai-engine/src/native_decode/cuda.rs` — seq-major refuse-not-mismap gate on the two head-major-only KV consumers. - `docs/MEMORY_ARCHITECTURE.md` — fourth-layout-location note. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
…T answer (#777) (#822) Makes prefix sharing **layout-general** and answers, definitively and with code citations, whether the ORT backend can use it. Refs #777, #803, #809, #793, #776, #783. ## 1. The shareability decision is now arithmetic, not a layout (core work) Prior framing (mine) said prefix sharing "requires seq-major". That is a **granule-relative** claim stated absolutely, and it was wrong. The real gate: ``` fragment_bytes = prefix_len x (contiguous bytes per fragment in that layout) shareable = fragment_bytes >= granule shareable_granules_per_fragment = floor(fragment_bytes / granule) multi_map_ops = fragments x shareable_granules_per_fragment wasted_boundary_bytes/sequence = fragments x (fragment_bytes mod granule) ``` Layout sets `fragment_bytes` and the **cost**, not the **possibility**. The genuine requirements are (a) a VMM-backed KV buffer and (b) `fragment_bytes >= granule`. New tested predicate `onnx_runtime_memory_governor::shareability::evaluate_prefix_shareability` over `ModelKvGeometry` + `KvFragmentation` (`head_major_bnsh` / `seq_major_bsnh` / `token_major`, or an arbitrary #783 stride descriptor). Its `PrefixShareability` result carries the full accounting and a `refusal_reason()` so a KV path **refuses with a reason** instead of silently making N private copies. `gqa_shared_prefix_parity_gpu` now consults it to admit the share rather than assuming seq-major. Granularity is **queried** (2 MiB CUDA #776, ~64 KiB Level Zero/Vulkan, 4 KiB CPU mmap), never hardcoded. **Head-major becomes shareable when the arithmetic says so** — exactly 8,192 tokens at 2 MiB (`granule / (head_dim x dtype)`), realistic for RAG and long system prompts. ### Correction to the issue's worked table (a term to double-check) The 2 MiB seq-major cell reads "2 granules each". That is `ceil` (granules the prefix *touches* / residency). The **shareable** count is `floor(4,096,000 / 2,097,152) = 1`, giving `96 x 1 = 96` multi-map ops — which matches the design section's own "96 multi-maps per sequence". So the correct cell is **1 granule/fragment, 96 ops**. It is a ceil-vs-floor slip in one cell, not a missing term; every other verified cell matches. The invariant `fragments x contiguous_bytes_per_token = whole-model per-token bytes` holds for all layouts (encoded as a test). Full layout x granule x prefix-length grid for **qwen14b** and **qwen2.5-0.5b** is in `docs/PREFIX_SHARE_INVESTIGATION.md`. ## 2. ORT answer — definitive, with citations The question is not "can ORT share" (the primitive is layout-agnostic) but **"does ORT KV route through an allocator we can back with the VMM arena?"** - **Dynamic decode** (`DecodeKvMode::ZeroCopyRebind`, `decode/mod.rs:66`): ORT allocates `present.*` itself; no allocator of ours is involved. **Cannot route.** - **Shared-buffer decode** (`decode/shared_batch.rs:360`): uses `Value::empty_in` -> `CreateTensorAsOrtValue` with an allocator from `for_session_device` -> `CreateAllocator` (`allocator.rs:239`), which the doc says **"wraps the session's internal EP allocator"** (`allocator.rs:232`) — selects the *device*, not the *implementation*. - **`RegisterAllocator`** (`governed_allocator.rs:738`) installs our vtable but ORT reserves the arena kind for its own arenas (`allocator.rs:65`, `governed_allocator.rs:746`) and fronts a registered device allocator with its **own BFC arena** — each KV tensor is a byte-range inside an arena chunk, not a granule-aligned VMM reservation. **So `commit_shared_prefix` cannot attach to today's automatic ORT paths — a valuable negative.** The one viable route: allocate the shared-buffer KV on **our VMM arena** and bind it as external device memory via `CreateTensorWithDataAsOrtValue` (`value.rs:1103`, "so a caller can hand ORT memory it allocated itself") + `IoBinding`. Bounded by head-major arithmetic (>= 8,192-tok prefixes at 2 MiB). **Scoped, not implemented in this PR.** ## 3. Seam honesty preserved The #809 refusing defaults stay; the predicate supplies the reason so a non-shareable configuration errors rather than mis-maps. Shared regions stay `PROT_READ`, granules charged once, `incremental_owned = 0`, no second allocator. ## Validation (all green locally; CI not awaited, fix-forward) - `cargo fmt -p onnx-runtime-memory-governor -p onnx-runtime-ep-cuda -- --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` - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` - `cargo test -p onnx-runtime-session` (all pass), predicate unit tests (8 pass) - GPU, single-threaded: `vmm_prefix_share_*` (5), `vmm_commit_shared_prefix_gpu` (4), `gqa_shared_prefix_parity_gpu` (1) — all pass - #814 policy: `mobius_seqmajor_growth_parity_native_cuda --ignored` — pass Does **not** merge or close #777. 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.
What this does (#777)
PR #803 landed the prefix-sharing primitive (
create_shared_prefix/commit_shared_prefix/incremental_owned_bytes_for_shared_prefix) on theproduction
CudaVmmAllocator, but nothing in production called it — only itsdefinition 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
DeviceAllocatortrait(
onnx-runtime-memory-governor):create_shared_prefix,incremental_owned_bytes_for_shared_prefix,commit_shared_prefix, plus anopaque
dyn SharedDevicePrefixhandle andSharedPrefixCommitInfo. A callerholding only
dyn DeviceAllocatorcan now pin a token prefix once and havesubsequent sequences map it. Default impls refuse (
InvalidRequest) ratherthan mis-map, so every existing non-VMM allocator keeps compiling and refuses
safely (unsupported requests error rather than mismap, per Use committed granules for VMM weight admission #745/Native backend cannot batch: continuous batching is ORT-only, so --max-batch is silently ignored and every request decodes alone #750).
CudaVmmAllocatorimplements the seam by delegating to the merged inherentmethods (Land the commit_shared_prefix allocator entry point (#777) #803) and downcasting the opaque handle via
as_any(). Granules comefrom the existing Pool CUDA VMM physical handles across reservations #740 authority-scoped pool through
carve(); no secondallocator, no per-sequence physical reservation.
(Seq-major (BSNH) KV layout: native fp16 decode pair #782/Replace KvLayout enum with a KV-cache stride descriptor (#783) #792/Enable stride-described seq-major flash prefill #794) reads shared-prefix VMM KV through the
dyn DeviceAllocatorseam. The prefix is physically contiguous as
layers × 2ranges (seq-major,the gate Probe: VMM page-level prefix sharing — primitive sound, no kill finding (#777) #793 requires); token-major (one multi-map/sequence) is measured
(Probe: token-major-across-all-layers KV — TLB risk cleared, floor measured (#750, #777, #783) #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.rsdrives thereal 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):
(C−1)×(K_prefix+V_prefix))incremental_owned_bytes_for_shared_prefixDeterministic 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 (--allfails 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.cargo test -p onnx-runtime-session— 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 newgqa_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 beforeallocators drop.
What remains (structural / out of scope)
call this seam automatically until
native_decode/cuda.rscan build aseq-major fixed-stride physical KV shape —
persistent_state_shapesishard-coded BNSH and no model declares a seq-major end-to-end physical shape
(Enable stride-described seq-major flash prefill #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.
Docs updated:
docs/MEMORY_ARCHITECTURE.md(concurrency-scaled prize section) anddocs/PREFIX_SHARE_INVESTIGATION.md(smallest next increment) now record thefirst production consumer, the measured numbers, and the remaining structural
gap.
Refs #777 #803 #793 #782 #745 #750