Repository navigation
docs: consolidate KV layout and residency design - #791
Merged
Merged
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Owner
Author
|
Coordinator review — approving and merging. Verified docs-only: all 3 changed I checked the large deletion (-724 lines) rather than assuming it was safe. It is Spot-checked that the load-bearing measurements actually made it in, not just the That last one is worth keeping deliberately: the "tokens per granule" model was Why this needed to land now: the reasoning behind this design line existed almost The section also does the thing I care about most for future readers — it Merging. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #791 +/- ##
==========================================
+ Coverage 79.41% 80.26% +0.85%
==========================================
Files 331 331
Lines 137690 137690
Branches 137690 137690
==========================================
+ Hits 109342 110523 +1181
+ Misses 24000 22816 -1184
- Partials 4348 4351 +3
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
## Summary - record the three-slice over-reservation finding and path-first audit guidance - refresh governed and raw workspace status from current CUDA/session code - correct the default-domain Attention score sizing/lifetime and name it as the next slice ## Validation - `git diff --check` - confirmed exactly one Markdown file changed and no non-Markdown files changed References #736 References #751 References #753 References #795 References #799 References #745 References #750 References #777 References #781 References #791 Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
References #736. Do not merge; do not close the issue. ## Summary Governs the default-domain `Attention` (`StandardAttentionKernel`) f32 score matrix in the CUDA EP, routing its device scratch through the merged `Kernel::workspace_requirement` + `execute_with_workspace` contracts with prepare-only planning (#747), so it is reserved against the memory authority and a capacity shortfall surfaces as the pre-header **429** typed path (`TierExhausted` / `CapacityUnavailable`, #743) via `MappedGrowthGrant` (#748) rather than a late 500/OOM. ## Measured bytes `batch * q_heads * q_seq * total_seq * sizeof(f32)` — quadratic in sequence length. At `B=1, H=32, q_seq=total_seq=2048` = **512 MiB** (536,870,912 bytes), larger than the already-governed Phase-2a range (#753). The buffer is always fp32 (fp32 accumulate) regardless of Q/K/V operand dtype. ## Which routes materialize the matrix (negative result) **Every route.** Unlike GQA's f32-reference-only `WS_SCORES` (#795) or the fused Attention prefill path that streams softmax through shared memory (#753), this kernel's `attention_row` has **no** flash/shared-memory route: it always stages the score matrix in global device memory (writes `scores[srow+j]`, re-reads for softcap/mask/softmax/`probs·V`). So every valid dispatch reserves it and **none reports `WorkspaceRequirement::NONE`** — only unresolvable or non-float input metadata reports `NONE`. This is a genuine-use result, recorded in the audit so the next reader does not re-investigate expecting a zero-charge route. ## Lifetime split Modelled truthfully as **both** classes: - **StepScoped** — per-call prefill/batched route (freed on each exit); this is the 512-MiB-class quadratic worst case. - **SessionPersistent** — the capture-eligible single-token decode route (`batch == 1 && q_seq == 1`), pooled/retained in the kernel workspace (small: linear in `total_seq`, ~256 KiB at the example shape). The classifier keys the split on static input metadata (`batch`, `q_seq`) as a conservative proxy for capture eligibility. Prepare-only planning reserves each class into its own executor slot (#753). The frozen fixed-capacity decode path over-estimates `total_seq` by padding rows, which is safe (reserve >= consume) — execution re-derives the exact size through the **same** `std_attention_scores_bytes` helper and refuses deterministically on any shortfall. ## Over-reservation found None. Planning and execution size through the same helper. The only sibling bypass — the staged-K/V + capture-metadata slots (`standard_attention.rs:702,1396`) — is a genuinely different lifetime/geometry, did **not** fall out of this change, and is left raw (tracked as the next slice). ## Scope / safety - Extends the central `is_planned_workspace_node` predicate (default-domain `"Attention"`); **no** parallel predicate. - Preserves the compatibility/opt-out direct-`execute()` path with self-owned scratch (pooled on the capture-eligible route, per-call otherwise) — unit/GPU allocation-count tests unaffected. - Released on success, error, and cancellation with truthful role/holder attribution. No assert inside any `Drop`. - Adds CPU-only governance unit tests + a source-scan regression guard (runtime-assembled needle) that fails on CI without a GPU if a raw score allocation is reintroduced on the governed site. ## Updated audit table `docs/MEMORY_ARCHITECTURE.md` §2.1: the default-domain `Attention` scores move into the governed rows (both the summary "Governed today" table and the site-by-site audit), with the route/lifetime split and the no-flash-route/no-`NONE` note verified against `attention_row`. The next slice is named from evidence: **`GroupQueryAttention` pooled QKV-projection staging** (`group_query_attention.rs:1033`, `batch * seq * hidden * elem`, session-persistent), now the largest remaining evidenced non-QMoE workspace bypass after the three quadratic score matrices are governed; present K/V there is KV-cache residency (govern via the KV residency contract, #791), not this table. ## Validation (local; did not wait for CI) - `cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --check` — clean - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` — clean (exact CI gate, no `--all-targets`) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` — clean - `cargo test -p onnx-runtime-session` — green - GPU suites single-threaded (`--features cuda,gpu-tests`): `standard_attention_gpu` (25), `attention_gpu`, `group_query_attention_gpu` (25), `qmoe_gpu`, `allocation_counts_gpu` — all green Working as the workspace/allocation slice; edits kept to workspace/allocation concerns to merge cleanly with concurrent KV shape/stride work in `native_decode/`. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
…804) (#807) ## Why This class of defect produced **two wrong conclusions this week**, both of which cost real work before being caught. Both are the same family: **state that is frozen or warmed once per process, making results depend on test order.** On a measurement-driven subsystem a wrong number does not just fail — it redirects the design. - **#797** — a GPU residency test failed ~50% of the time on a cold context and looked exactly like a correctness bug in growth-under-capture. Root cause: a readback on the **legacy default stream** raced a memset on a created **non-blocking stream** (the two are mutually exempt from implicit synchronization), so the readback returned partial fills **with no CUDA error**. `cuCtxSynchronize` did *not* fix it, which refuted the obvious race theory. It only "passed" because an alphabetically-earlier sibling subtest warmed the context. - **#801 / #804** — `captures=0` was reported as the decoder structurally declining CUDA graph capture, and that attribution reached a merged PR body. Root cause (#804): the capture-OFF phase set `ONNX_GENAI_CUDA_GRAPH=0` before the process-wide `RuntimeConfig` `OnceLock` was first read, and that config is frozen on first read — so every later phase in the same process inherited it. #794 and #801 failed to reproduce an end-to-end measurement and misattributed the reason. This is the same telemetry-that-misleads lineage as #719 / #715 (documented in `MEMORY_ARCHITECTURE.md`'s "third failure mode"). After the second instance in a week, the response is a **guard**, not more vigilance. This PR is **test infrastructure + documentation only**. No production behavior changes. ## What ### 1. A guard for process-frozen configuration (`onnx-genai-runtime-config`) `runtime_config()` records the environment it froze the `RuntimeConfig` snapshot from (captured through the same lookup `from_env` uses, so the set can never drift from the parsed fields). On every later call it re-reads those variables and **panics loudly** — with the frozen and current values, and a pointer to #804 — if any now differs. A test can no longer silently believe it set a knob that was already frozen. - **Gated on `debug_assertions`, not `cfg(test)`.** The #804 mutation lives in an *integration* test of a downstream crate, where `cfg(test)` for this crate is never set. `debug_assertions` is enabled for dependencies under the `dev`/`test` profile, so the guard is live for every `cargo test` in the workspace, and it is compiled out of `--release` — **the production path is byte-for-byte unchanged and pays nothing** (a single `OnceLock` read). - The comparison is a pure `find_changed()` helper with unit tests for the unchanged, changed, newly-set (the #804 shape), and cleared cases. ### 2. A helper for stream-correct device test code (`onnx-runtime-cuda-memory::test_support::TestStream`) A shared, test-only (`#[cfg(any(test, feature = "gpu-tests"))]`) RAII wrapper around **one** non-blocking stream. Every device op a test needs — `fill`, `write`, `read`, plus `raw()` for graph launches/memsets — flows through that stream, so a single `cuStreamSynchronize` is a total order. The module doc explains the legacy-default-vs-non-blocking exemption **at the call site** so the next author does not rediscover #797 the hard way. `vmm_kv_layout_residency_gpu.rs` (the exact #797 site) is converted to it and its four local stream helpers removed; **8/8 reliable in isolation** and green in the full single-threaded suite. ### 3. The pattern, documented where it will be read (`docs/MEMORY_ARCHITECTURE.md`) A tight "**a fourth failure mode: order-dependent test state**" subsection next to the existing third-failure-mode material: both case studies, the two mechanisms (frozen process config; context warmed by an earlier sibling), the two remedies (process isolation for anything that reads process-frozen config; single-stream discipline for device tests), and the one-line discipline that names the defect — *run any test you touch in isolation as well as in the full suite*. Cites #797/#801/#804/#794/#719/#715 and links rather than duplicating (#791 consolidated this doc to stop drift). ## Inventory of other order-dependent tests found (not fixed here) Per scope discipline, these are **listed, not converted**. The `TestStream` helper and the config guard make converting them mechanical follow-ups. ### Device tests that mix a created stream with default-stream copies (#797 hazard) Each creates a non-blocking stream **and** issues default-stream `cuMemcpyHtoD_v2`/`cuMemcpyDtoH_v2` copies. Where those touch the same buffers as stream-ordered work on a cold context, they are exposed to the exact #797 race. They should migrate to `test_support::TestStream`: - `crates/onnx-runtime-cuda-memory/tests/vmm_commit_shared_prefix_gpu.rs` - `crates/onnx-runtime-cuda-memory/tests/vmm_dummy_page_gpu.rs` - `crates/onnx-runtime-cuda-memory/tests/vmm_dummy_write_protect_gpu.rs` - `crates/onnx-runtime-cuda-memory/tests/vmm_granularity_gpu.rs` - `crates/onnx-runtime-cuda-memory/tests/vmm_graph_remap_gpu.rs` - `crates/onnx-runtime-cuda-memory/tests/vmm_kv_contiguous_tail_gpu.rs` - `crates/onnx-runtime-cuda-memory/tests/vmm_prefix_share_gpu.rs` - `crates/onnx-runtime-cuda-memory/tests/vmm_prefix_share_write_protect_gpu.rs` - `crates/onnx-runtime-cuda-memory/tests/vmm_stable_va_weight_slot_gpu.rs` (Non-vmm siblings `virtual_memory_gpu.rs` and `physical_pool_production_gpu.rs` also use default-stream copies; they pass but are the same latent shape.) ### Tests that mutate a `RuntimeConfig`-feeding env var within a process These change a variable that feeds the frozen snapshot (`ONNX_GENAI_CUDA_GRAPH`, `ONNX_GENAI_EP`) across phases. The correct pattern is **child-process-per-phase**; the new guard will now fire loudly on any that mutate in-process and then read `runtime_config()`: - `crates/onnx-genai-engine/tests/mobius_seqmajor_growth_parity_native_cuda.rs` — **already correct** (spawns a child per mode; its own doc comment states `ONNX_GENAI_CUDA_GRAPH` is parsed once). This is the reference pattern. - `crates/onnx-genai-engine/tests/stable_va_capture_offload_e2e.rs` — sets `ONNX_GENAI_CUDA_GRAPH` per capture phase. - `crates/onnx-genai-engine/tests/gemma4_e2b_native_cuda_pipeline_bench.rs` — sets `ONNX_GENAI_CUDA_GRAPH` to `1` then `0`. - `crates/onnx-genai-engine/tests/qwen2_5_1_5b_divergence.rs` — switches `ONNX_GENAI_EP` between `cpu` and `cuda`. - `crates/onnx-genai-engine/tests/qwen3_0_6b_capture_step_inputs_decline.rs` and `gemma3n_native_cuda_capture_realmodel.rs` — toggle the native-decoder capture env per phase. Session-crate tests that toggle env (`ONNX_GENAI_DECODE_MEMO`, `ONNX_GENAI_SCAN_INLINE_SINGLE_TRIP`, `ONNX_RT_SIBLING_MERGE`, the EP-CPU weight-offload knob) read those variables **directly at their call sites**, not through `RuntimeConfig`, so they are *not* affected by the freeze guard and the full `onnx-runtime-session` suite stays green. ## Validation (all local; not waiting for CI) - `cargo fmt -p onnx-genai-runtime-config -p onnx-runtime-cuda-memory -- --check` ✅ (`--all` hits the Windows `os error 206` long-command artifact, not a formatting issue) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` ✅ - `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 test -p onnx-runtime-session` ✅ (full suite) - `cargo test -p onnx-runtime-cuda-memory --features gpu-tests -- --test-threads=1` ✅ (all vmm_*/gpu suites, RTX 4060) - The converted `seq_major_fixed_stride_grows_under_captured_replay_without_recapture` run **8× in isolation**: 8/8 pass — the property this PR is about. ## Notes - Kept strictly to test infra + the docs subsection so the two other active workstreams (GQA pooled QKV staging in `group_query_attention.rs`; head-major vs seq-major re-capture measurement) merge cleanly — no overlap. - Does **not** merge and does **not** close any issue. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
## What Documentation-only sync of five memory/KV design docs to the merged, executable-test-backed facts about VMM commit geometry, KV layout, weight offload, CUDA-graph capture, and plugin export. No code or tests changed. `MEMORY_ARCHITECTURE.md` is treated as the single source for KV geometry and the audit table; the other four link into it rather than restating (per #791). ## Most misleading claim found `MEMORY_ARCHITECTURE.md` — the document explicitly designated the single source of truth — still asserted the engine *"explicitly invalidates the graph on every KV bucket growth for both layouts"* and that *"the engine does not yet use that fact to retain the captured graph,"* with a table showing seq-major at `captures=4, invalidations=4`. That was falsified by **#811**: conditional invalidation has landed, seq-major fixed-stride growth now **keeps** the captured graph (growth-attributable invalidations **3 → 0**, `captures=1/invalidations=1`, matching the no-growth control), verified by `mobius_seqmajor_growth_parity_native_cuda.rs` including a negative oracle. Because this doc is the single source others link to, the stale claim had the widest blast radius. ## Changes - **MEMORY_ARCHITECTURE.md** — correct the #811 conditional-invalidation prose + table (verified against `cuda.rs` `binding_growth_signature`/`record_growth_decision` and the parity test); add the ORT-backend KV-routing + #762 plugin-export note (verified `allocator.rs:232` `for_session_device` selects by device, not implementation); thread #822 (granule queried), #812 (engine-side consumer gate), and a #800/#802 record correcting this doc's earlier fp16 sizing of Phase-2a scores (verified `attention.rs:368` uses `dtype.element_size()`). - **CUDA_GRAPH_CAPTURE.md** — replace the mutual-exclusion result with #796 (stable VA removes it), #798 (managed no-spill default), #811 (conditional invalidation), and the #804/#807 `captures=0` process-artifact correction. - **WEIGHT_OFFLOAD.md** — correct the "mutually exclusive today" claim (#796) and the north-star status for the managed no-spill default (#798). - **KV_INSERTION_DESIGN.md** (most stale, 2026-07-15) — add a synchronization banner linking to `MEMORY_ARCHITECTURE.md`; note the layout line post-dates the memo, prefix sharing is a landed primitive (#777), and block-table paged attention remains a proposal (verified: no `block_table` in any CUDA attention kernel). - **DESIGN.md** — forward-pointer in the already-superseded §39.5 noting layout (seq-major/token-major, #787/#792) is now the primary residency lever. ## Method Every retained performance claim cites the PR that measured it; measured fact / design intent / not-yet-built are distinguished; superseded designs are labelled, not deleted. The "still open" item (engine commits flat buckets, not the dense live-prefix ranges `kv_commit.rs` computes, so the `layers×2` floor is not yet realised through the engine) is described as in-flight, not delivered. ## Note on #762 The task framing had #762 (EP plugin export) as open; `gh pr view 762` shows it **merged** today. It is described accurately per that ground truth: the plugin-EP mechanism landed for the CPU EP (validated end-to-end), while the CUDA plugin shim compiles but is unvalidated on hardware — so plugin-EP-hosted VMM KV remains in-progress design intent, not a shipped capability. ## Validation Documentation only — `git diff --stat` touches exactly the five `.md` files; working tree otherwise clean. **Not merged; no issue closed.** Refs #736, #762, #776, #777, #783, #787, #791, #792, #794, #796, #797, #798, #801, #802, #804, #807, #811, #812, #822 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
docs/MEMORY_ARCHITECTURE.mdthe single authoritative explanation of KV layout and residencyValidation
git diff --checkRefs #750, #759, #772, #776, #777, #778, #782, #783, #787.