Repository navigation
Govern GQA BNSH transpose workspace - #810
Merged
Merged
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Owner
Author
|
Coordinator pre-merge validation on the merged branch: The lead paid off, and the split is the honest kindI pointed this slice at a specific clue: the source notes that the BSH->BNSH
That is exactly the shape this audit keeps finding — bytes charged against the The tally, six slices in
Five of six. The guidance recorded in #800 — start from use, not from allocation Worth noting the pattern within the pattern: three of the five over-reservations Node matching went into the centralised Merging. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #810 +/- ##
=======================================
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:
|
🔴 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
Follow-up to #810 / issue #736. ## Regression The newly governed GQA prefill transpose workspace exposed a mapped-growth authority gap in the native CUDA VMM path. The executor prepares session-persistent and step-scoped workspace slots independently, but both share the canonical arena mapped zone. That zone starts with a zero-byte allowance. Its first 2 MiB granule could only grow by transferring capacity from a registered holder; when no such holder existed, prefill failed with `TierExhausted` despite unassigned device-authority capacity. ## Fix - Assign mapped capacity that belongs to no zone before attempting holder transfer/reclaim. - Keep committed capacity with the arena zone. - Return newly assigned capacity when a grant is rolled back or the allocator maps fewer bytes than planned (including reuse of an existing granule). - Preserve typed `TierExhausted` refusal once the authority is genuinely exhausted. - Document the mapped-zone growth order. This preserves the truthful session-persistent GQA transpose lifetime and the `Sq == 1` route-sensitive `NONE` behavior from #810; it does not restore self-owned scratch or weaken authority accounting. ## Validation - `cargo test -p onnx-runtime-memory-governor` (39 + 5 passed) - memory-governor Clippy with `-D warnings` - exact CUDA EP Clippy gate - exact engine all-target Clippy gate - full `cargo test -p onnx-runtime-session` - Mobius seq-major growth parity native-CUDA regression (all VMM/realloc, capture on/off configurations passed) - GPU suites single-threaded: GQA 25, seq-major parity 2, Attention 5, allocation counts 2 Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
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>
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.
References #736
Result
WS_Q_BNSHandWS_OUT_BNSHthrough the existing GQAworkspace_requirement/ prepared composite;execute_with_workspacecarves the same helper-computed layout and rejects short views deterministically.Sq == 1decode now uses input Q and output directly, so both transpose regions charge zero and the whole GQA requirement isNONEwhen staging/scores are also absent.Sq == 1overlays Q extraction/BNSH with packed-Q staging rather than summing mutually exclusive uses. EverySq > 1route needs Q and output transpose buffers.executeself-owned pooled compatibility fallback. Prepared grants retain the existing typed pre-header 429 capacity-refusal path.Bytes and lifetime
Derived from
dtype.byte_size()and code geometry (not allocator measurements): at B=1, H=32, head_dim=128, Sq=1024, Q+output transpose scratch is 32 MiB f32 or 16 MiB fp16/bf16. At Sq=1, unpacked non-RoPE is 0; unpacked RoPE is Q-only: 16 KiB f32 or 8 KiB fp16/bf16.The old
GqaWorkspaceslots grew and remained resident across calls, so this is truthfullyWorkspaceLifetime::SessionPersistent. The executor-owned prepared grant supplies success/error/cancellation cleanup and device-authority attribution.Audit/docs
Updates
MEMORY_ARCHITECTURE.md§2.1 and the site audit, and selects default-domainAttentionstaged K/V as the next evidence-backed slice. Adds a CPU-only source-scan regression preventing unconditional raw transpose-slot reservations.Local validation
cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --checkcargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warningscargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warningscargo test -p onnx-runtime-session(153 core tests plus all integration/doc suites green)group_query_attention_gpu(25 passed),gqa_seqmajor_parity_gpu(2),attention_gpu(5),allocation_counts_gpu(2)