Repository navigation
Add transactional mapped growth grants - #748
Merged
Merged
Conversation
Coordinate weight reclaim with native KV and workspace mapping through authority-owned RAII reservations and batched VMM commits. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #748 +/- ##
==========================================
- Coverage 80.10% 79.34% -0.76%
==========================================
Files 331 331
Lines 135684 137216 +1532
Branches 135684 137216 +1532
==========================================
+ Hits 108686 108876 +190
- Misses 22704 24007 +1303
- Partials 4294 4333 +39
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
|
| Status | Scenario | Base | PR | Change |
|---|---|---|---|---|
matmul/small_generic_bf16_threads=1/1x256x256 |
42.45 µs | 53.60 µs | +26.3% | |
matmul/large_generic_bf16_threads=8/32x1024x1024 |
1.57 ms | 1.82 ms | +15.8% | |
| ✅ | qwen3_sampling_processors/top_k_top_p_fast |
651.48 µs | 744.41 µs | +14.3% |
| ✅ | gather/small_f16_threads=1-internal/4096 |
523.7 ns | 589.5 ns | +12.6% |
| ✅ | gather/medium_f32_threads=1-internal/32768 |
4.81 µs | 5.35 µs | +11.1% |
| ✅ | matmul/large_generic_bf16_threads=1/32x1024x1024 |
2.01 ms | 2.23 ms | +11.0% |
| ✅ | matmul/medium_generic_f32_threads=8/32x512x512 |
1.21 ms | 1.34 ms | +10.7% |
| ✅ | qwen3_sampling_processors/top_k_top_p_full_sort_baseline |
5.79 ms | 6.34 ms | +9.5% |
| ✅ | gather/small_f32_threads=1-internal/4096 |
739.7 ns | 790.4 ns | +6.9% |
| ✅ | matmul/small_generic_bf16_threads=8/1x256x256 |
45.53 µs | 48.53 µs | +6.6% |
| ✅ | sampling_latency/min_p_per_token |
220.92 µs | 233.78 µs | +5.8% |
| ✅ | matmul/small_generic_f16_threads=1/1x256x256 |
41.09 µs | 42.98 µs | +4.6% |
| ✅ | gather/small_bf16_threads=1-internal/4096 |
517.7 ns | 541.1 ns | +4.5% |
| ✅ | reduce_mean/large_f32_threads=1-internal/262144 |
1.10 ms | 1.14 ms | +3.2% |
| ✅ | matmul/small_generic_f16_threads=8/1x256x256 |
43.56 µs | 44.05 µs | +1.1% |
| ✅ | matmul/large_generic_f32_threads=8/32x1024x1024 |
4.65 ms | 4.69 ms | +0.9% |
| ✅ | kv_cache/alloc_dealloc_pages |
43.77 µs | 44.08 µs | +0.7% |
| ✅ | gather/large_f32_threads=1-internal/131072 |
42.00 µs | 42.16 µs | +0.4% |
| ✅ | gather/medium_bf16_threads=1-internal/32768 |
2.86 µs | 2.83 µs | -0.9% |
| ✅ | qwen3_sampling_processors/top_p_fast_after_top_k |
544.95 µs | 528.71 µs | -3.0% |
| ✅ | add/large_f16_threads=1-internal/4194304 |
47.58 ms | 46.15 ms | -3.0% |
| ✅ | matmul/large_generic_f16_threads=1/32x1024x1024 |
96.47 µs | 93.05 µs | -3.5% |
| ✅ | add/large_bf16_threads=1-internal/4194304 |
47.60 ms | 45.89 ms | -3.6% |
| ✅ | gather/large_f16_threads=1-internal/131072 |
15.92 µs | 15.25 µs | -4.2% |
| ✅ | add/large_f32_threads=1-internal/4194304 |
48.38 ms | 45.79 ms | -5.4% |
| ✅ | sampling_latency/top_p_per_token |
400.32 µs | 370.78 µs | -7.4% |
| ✅ | matmul/medium_generic_f32_threads=1/32x512x512 |
2.83 ms | 2.61 ms | -7.7% |
| ✅ | matmul/large_generic_f32_threads=1/32x1024x1024 |
10.64 ms | 9.77 ms | -8.2% |
| ✅ | sampling_latency/top_k_per_token |
56.50 µs | 51.80 µs | -8.3% |
| ✅ | sampling_latency/greedy_per_token |
3.56 µs | 3.25 µs | -8.8% |
| ✅ | reduce_mean/medium_f32_threads=1-internal/65536 |
311.51 µs | 277.88 µs | -10.8% |
| ✅ | block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1 |
86.91 µs | 76.44 µs | -12.1% |
| ✅ | qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline |
3.92 ms | 3.44 ms | -12.3% |
| ✅ | gather/medium_f16_threads=1-internal/32768 |
3.20 µs | 2.78 µs | -13.1% |
| ✅ | tokenization/encode_tokens_per_second |
425.44 µs | 369.62 µs | -13.1% |
| ✅ | add/medium_f16_threads=1-internal/262144 |
3.03 ms | 2.60 ms | -14.1% |
| ✅ | tokenization/decode_tokens_per_second |
7.13 ms | 6.08 ms | -14.7% |
| 🟢 | gather/large_bf16_threads=1-internal/131072 |
17.36 µs | 14.71 µs | -15.3% |
| 🟢 | qwen3_sampling_processors/top_k_full_sort_baseline |
2.54 ms | 2.14 ms | -15.7% |
| 🟢 | add/medium_f32_threads=1-internal/262144 |
3.17 ms | 2.67 ms | -16.0% |
| 🟢 | logit_processing/seven_processor_chain_per_step |
405.75 µs | 336.69 µs | -17.0% |
| 🟢 | qwen3_sampling_processors/top_k_partial_selection |
167.47 µs | 137.87 µs | -17.7% |
| 🟢 | add/medium_bf16_threads=1-internal/262144 |
3.50 ms | 2.88 ms | -17.7% |
| 🟢 | matmul/large_generic_f16_threads=8/32x1024x1024 |
129.03 µs | 103.12 µs | -20.1% |
| 🟢 | matmul/medium_generic_f16_threads=1/32x512x512 |
48.20 µs | 38.33 µs | -20.5% |
| 🟢 | block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1 |
487.93 µs | 386.60 µs | -20.8% |
| 🟢 | matmul/medium_generic_f16_threads=8/32x512x512 |
52.89 µs | 41.25 µs | -22.0% |
| 🟢 | reduce_mean/small_f32_threads=1-internal/4096 |
22.20 µs | 17.18 µs | -22.6% |
| 🟢 | matmul/small_generic_f32_threads=1/1x256x256 |
58.25 µs | 44.06 µs | -24.4% |
| 🟢 | matmul/small_generic_f32_threads=8/1x256x256 |
66.74 µs | 50.19 µs | -24.8% |
| 🟢 | block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024 |
59.66 µs | 42.71 µs | -28.4% |
| 🟢 | grammar_masking/llguidance_compute_mask/32 |
118.15 µs | 82.59 µs | -30.1% |
| 🟢 | matmul/medium_generic_bf16_threads=1/32x512x512 |
796.84 µs | 551.57 µs | -30.8% |
| 🟢 | add/small_bf16_threads=1-internal/1024 |
22.46 µs | 15.39 µs | -31.5% |
| 🟢 | matmul/medium_generic_bf16_threads=8/32x512x512 |
836.24 µs | 568.20 µs | -32.1% |
| 🟢 | add/small_f16_threads=1-internal/1024 |
22.43 µs | 13.96 µs | -37.8% |
| 🟢 | block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024 |
67.19 µs | 41.44 µs | -38.3% |
| 🟢 | add/small_f32_threads=1-internal/1024 |
372.7 ns | 224.8 ns | -39.7% |
| 🟢 | block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024 |
907.28 µs | 528.55 µs | -41.7% |
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: { 9.13 5.17 6.50 }
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)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Owner
Author
|
Local final review passed at What landsA complete transactional mapped-growth protocol shared by native CUDA KV and governed workspace:
Explicit byte Live 14B / 6 GiB evidence
Current server scheduling serializes those requests, so this does not yet prove simultaneous live-KV accumulation; that limitation is documented. Review defects closedThe review series closed grant self-blocking, allowance arithmetic, capacity theft, rollback, shared-granule union, workspace late allocation, WDDM fallback, typed 429, resident-weight double charge, state-only models, arena attribution ownership, ordinary/range deallocation refunds, and shared allocation attribution symmetry. #736 remains open: activation/scratch expansion beyond QMoE, targeted VA layout, full zone policy/lifecycle, and preemption/swap remain. Merging. |
This was referenced Aug 10, 2026
Merged
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
…753) ## Summary Next increment on #736: audit the CUDA activation/scratch/workspace allocations that still bypass the device memory authority (beyond the QMoE path already governed by #747/#748), then govern the largest live bypass. References #736. Does **not** close the issue. ## Step 1 — Audit A full inventory of production scratch/workspace allocations across `onnx-runtime-ep-cuda`, `onnx-runtime-session` and the CUDA kernels is added to `docs/MEMORY_ARCHITECTURE.md` §2.1, with file/line, byte formula, size class (static / config-derived / prompt-dependent) and lifetime class (model-fixed / session-persistent / request-prefill / step-scoped / already-governed) for every remaining raw path. Ranking by live bytes, the largest ungoverned bypass is the `com.microsoft::Attention` **Phase-2a** scratch: - **scores buffer** `batch · num_heads · sq · sk · elem` — quadratic in sequence length (≈268–536 MB at B=1, H=32, S=2048), the single largest prompt-dependent scratch; plus - a static **32 MiB cuBLASLt workspace**. The fused (flash) prefill path allocates **no** device scratch (shared memory only), so `run_attention_phase2a` is the *only* scratch source in `AttentionKernel` — a clean StepScoped slice. ## Step 2 — Govern one vertical slice - **`attention.rs`**: `phase2a_scratch_layout` + `Phase2aScratch` carve the scores buffer and cuBLASLt workspace from one governed blob. `run_attention_phase2a` uses an executor-prepared `WorkspaceView` when present; the compatibility/opt-out path (direct `execute`, GQA reuse, unit tests) keeps its self-owned `alloc_raw` scratch. New `Kernel::workspace_requirement` reports the exact bytes (StepScoped, `NONE` for the fused path) and `execute_with_workspace` consumes the reserved slot. - **`executor/bindings.rs`**: prepare-only planning (#747) now tracks one workspace peak **per lifetime class** (SessionPersistent + StepScoped) and prepares both slots, so a graph mixing QMoE (session-persistent) and Attention (step-scoped) governs both. Reserve logic factored into `reserve_prepared_workspace`; node matching centralised in `is_planned_workspace_node`. - **Same helper for plan and exec**: planning and execution both size through `use_fused` + `phase2a_scratch_layout`, so the reserved slot always covers the dispatch. A shortfall would error deterministically rather than silently under-allocate. - **429, not 500/OOM**: reservation happens inside `prepare_with_device_bindings` (pre-header, prefill admission). `reserve_workspace`/`prepare_mapped_growth` emit `TierExhausted`/`CapacityUnavailable`, which `DriverFailure::from_engine_error` already maps to HTTP 429 (#743) — no server change needed. - **Lifetime correctness**: step-scoped scratch is released by the executor's existing `release_step_workspace` on success/error/cancellation; the governed kernel path no longer owns or frees it. No asserts in any `Drop`. - **Regression guard**: new unit test pins `phase2a_scratch_layout`'s byte formula; a GPU-gated test asserts `workspace_requirement` equals the layout for a Phase-2a shape and returns `0` (NONE) for a fused shape — preventing a raw-allocation regression on the governed slice. ## Step 3 — Document `docs/MEMORY_ARCHITECTURE.md` §2.1 records the governed-today table (QMoE, Attention Phase-2a) and every remaining raw path. It identifies **`GroupQueryAttention`'s `WS_SCORES` slot** (`batch·heads·sq·sk·4`, session-persistent growable) as the recommended next increment, with the three shared 32 MiB cuBLASLt GEMM workspaces as a cheap follow-up. ## Validation - `cargo fmt --all -- --check` ✅ - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` ✅ (exact CI gate, no `--all-targets`) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` ✅ - Session executor lib tests (76 incl. `prepare_reserves_static_nested_qmoe_workspace_and_child_reuses_it`) ✅ - CUDA EP attention lib tests incl. new guard tests ✅ (GPU present) - GPU tests: `attention_gpu` (5), `qmoe_gpu` (27), `allocation_counts_gpu` (2), `group_query_attention_gpu` (25, single-threaded), `standard_attention_gpu` ✅ Note: `gqa_gpu_fixed_decode_capture_replays_bit_identically` fails only when many capture tests share one binary in parallel (pre-existing CUDA graph-capture stream interaction, `CUDA_ERROR_STREAM_CAPTURE_INVALIDATED`); it passes in isolation and single-threaded. Unrelated to this change (the GQA edit only adds a `None` arg to the prefill helper). Per instructions: local validation only, not waiting for CI; PR opened but **not** merged and issue **not** closed. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 11, 2026
… forces commit, head-major floor stands (#772) ## What & why The owner's directive for #750: give each sequence one **stable, flat, contiguous virtual address** for its whole context so the existing flat GQA attention kernel keeps reading a plain `[heads, seq, head_dim]` buffer while KV is physically scattered and grown on demand — paging invisibly underneath, the kernel never learning page tables. If that holds, batching is just N sequences each with its own stable contiguous VA and the flat kernel works unchanged. This PR **confronts the known blocker first** (#721 stage 4: a stable full-context VA committed **1.5 GB** vs bucket growth's **48 MB**, a 32x regression, because the decode kernel read the full padded shape) and answers the crux with **evidence on this hardware**. It is primarily a **truthful negative for the fixed-stride variant** plus a confirmation that the already-landed bucket-stride path is the correct realization. **No engine behavior change — isolating GPU test + docs only. Do not merge.** ## The crux, answered by an isolating GPU test `crates/onnx-runtime-cuda-memory/tests/vmm_kv_contiguous_tail_gpu.rs` (3 tests, pass on device; measured granule **2 MiB**): 1. **`bounded_read_of_contiguous_va_leaves_tail_granules_uncommitted`** — a read bounded to the committed live length, *including a captured graph replay* (the decode hot path), succeeds while the reservation tail stays physically **uncommitted**. The flat contiguous VA is real and its tail is physically absent. 2. **`padded_read_into_uncommitted_tail_faults_forcing_commit`** — the *same* buffer, read **one byte past** the live length into the uncommitted tail, faults `CUDA_ERROR_INVALID_VALUE` (synchronous, non-sticky). This is **#721 stage 4 reproduced in isolation**: the **read pattern, not the reservation, forces the commit**. A kernel that reads the full padded shape dereferences the tail, so the tail can never be decommitted. 3. **`head_major_full_context_stride_pays_one_granule_per_object`** — even a perfectly bounded read cannot make a **fixed full-context stride** cheap: KV is per-binding/head-major, so each head's valid prefix lands in its own granule and the commit floor is `objects x granule`. Reproduces the documented qwen2.5-0.5b floor on hardware: **96 head-stripes × 2 MiB = 192 MiB committed for ~12 KiB of live content.** ## Committed-bytes verdict (the number that decides adoption) | Scheme | Committed bytes | Re-capture on growth? | |---|---|---| | Bucket growth (today) | ~48 MB (baseline) | yes | | **Fixed full-context stride** (the literal "one flat VA, never re-strided") | **192 MiB near-empty / 1.5 GB @32k — 32x worse** | no | | **Bucket stride under stable VA** (landed `kv_commits_on_demand`, #682/#740/#748) | **parity with bucket growth**, + stable VA | **still yes** | Verified on real qwen2.5-0.5b (`native_cuda_vmm_kv_grows_in_place_and_commits_more_granules`): full-context VA reserved, only the initial **bucket** committed at load (`< full/2`), growth keeps `device_ptrs` **stable** (contiguous VA) and commits exactly the **bucket delta**. **Conclusion:** a length-bounded kernel is **necessary but not sufficient**. Low committed bytes *and* no re-capture cannot both hold under the head-major layout: the fixed stride that removes re-capture pays the granule floor (loses 32x), and the bucket stride that keeps committed bytes at parity still re-strides/re-captures on growth. The GQA op already takes `seqlens_k`/`total_sequence_length` so its *compute* is length-bounded; the floor is set by the layout stride, not the compute bound. Full write-up: **`docs/VMM_KV_CONTIGUOUS_VA.md`**. ## Does `onnx-genai-kv`'s page machinery still matter, and who owns paging? Direct answer to the owner's question: - The **native CUDA decode path has no `PageTable`/`PagedKvCache` consumer** (`native_decode/cuda.rs` references neither; `paged_gqa.rs` is a batch-1 **CPU** primitive). The KV page store is **host-only** (`docs/MEMORY_ARCHITECTURE.md`: "never touched device memory"). - Under the VMM contiguous-VA design, **device KV paging is owned by the CUDA VMM layer** — `CudaVmmAllocator` (#740 pool) mapping granules under a fixed per-binding reservation, with governor grants (#748) and committed-granule admission (#745). The kernel sees a flat contiguous buffer and never learns page tables. - **Implication:** for the CUDA batch path, `onnx-genai-kv`'s page machinery is **not needed** as the device pager — a batch is N sequences each with its own stable VMM reservation, flat kernel unchanged. `onnx-genai-kv` stays relevant for the **CPU** paged-GQA path and host-side concerns (page store, prefix sharing/CoW, quant layout). ## Next experiments (not in this PR) 1. **Commit below the bucket** — since GQA is length-bounded, only `ceil(logical_len)` granules per head need backing. Open question the isolating test frames: does the ORT GQA **CUDA** kernel / graph capture ever touch the bucket tail `[logical_len, bucket)`? If not, sub-bucket commit beats bucket growth. Needs an engine change + byte-identical decode run. 2. **Seq-major KV layout** — the only layout giving both low commit and no re-capture, but requires re-exported models / a custom kernel. ## Validation - `cargo fmt --all -- --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` ✓ (CI gate, no `--all-targets`) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` ✓ - New GPU tests (3) + existing `vmm_graph_remap_gpu` (4) pass on device. - Real-model measurement: `native_cuda_vmm_kv_grows_in_place_and_commits_more_granules` passes on qwen2.5-0.5b. Refs #750, #721. Does not close any issue. **Do not merge.** Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 11, 2026
…le and honest (#750) (#758) ## Summary First increment for #750: make the **native backend's inability to batch observable and honest**, without implementing batched native decode. The native decode path pins both the batch and query-seq dimensions to `1` as a structural invariant (`native_decode/cuda.rs`). Despite that, the server inferred "batching enabled" solely from whether an ORT decoder session existed: `Engine::continuous_batch_manager` returns `Err` on native (`"ORT decoder session is unavailable"`), and `driver.rs` swallowed that into a debug/info "using per-request engine path" log line. So `--max-batch` was accepted and reported while having **no effect** on the native path. ## What changed ### Engine (`onnx-genai-engine`) - New `Engine::batching_capability() -> BatchingCapability`, **sourced from the decode path** (resolved `EngineDecodeBackend` + `ModelDecodePath`), *not* from `session.is_some()`. This is the core honesty fix. - Native → `Some(1)` with a reason naming the backend and citing `native_decode/cuda.rs`. - ORT static-cache and shared-buffer past/present → unlimited (`None`, bounded only by memory / `--max-batch`). - Legacy and non-shared-buffer past/present → `Some(1)` with the reason. - `BatchingCapability` re-exported from the crate root. ### Server (`onnx-genai-server`) - A decode-path-sourced `BatchingReport` is resolved at engine startup and logged at **info** level (`resolved batching capability` with `batch_supported` / `requested_max_batch` / `effective_max_batch` / `reason`). - Surfaced so an operator can *see* the truth instead of inferring it: - `GET /v1/resources` → `batching { supported, requested_max_batch, effective_max_batch, reason }`. - `/v1/debug/kv` → `batch_supported` / `effective_max_batch`. - admin vram-override response includes the same report for parity. - New `--max-batch` CLI flag (env `ONNX_GENAI_MAX_BATCH`). ### Metrics honesty note Documented at the metric site that `onnx_genai_batch_size_current` counts **admitted generations**, not sequences co-decoded in one batched pass (`crates/onnx-genai-server/src/metrics.rs`). It climbs with concurrent requests even on a backend that decodes them one at a time — which is exactly why metrics alone never revealed the missing batching. ## Decision: explicit refusal vs. silent clamp The task asked to decide and justify. I chose a **hybrid**, biased toward explicit refusal: - **Operator explicitly asks for `--max-batch > 1` on a backend that cannot honor it → hard startup refusal** naming the flag and the reason (e.g. `--max-batch 4 cannot be honored: the native decode backend advances exactly one sequence per step…`). No silent acceptance. - **Flag omitted → the server default width is clamped to the decode path's limit** (silently), and the clamp is made **visible** in `/v1/resources` (`effective_max_batch`, `supported`). Rationale for not refusing in the omitted case: the server default (`DEFAULT_MAX_BATCH = 4`) applies to *every* deployment. Refusing on the default would break every existing native / legacy server on startup, which is out of scope for a "report the truth" increment. Refusal is reserved for the case where an operator made an explicit, unhonorable request — there the silent no-op is genuinely misleading. Either way the effective width is never left invisible. ## Tests (CPU-only, no GPU required) - **engine integration** (`tests/batching_capability.rs`): static-cache capability matches a *successful* `continuous_batch_manager`; plain past/present capability matches single-sequence decode (manager errors). Ties the reported capability to what the decode path can actually do, for both paths. - **engine native unit test**: a native engine reports `Some(1)` even with a StaticCache I/O shape (proving capability is a property of the decode path, not the model), and cannot build a `>1` continuous batch manager. - **server**: explicit `--max-batch 4` is **refused** on a non-batching backend (not silently accepted); the default width is clamped and `/v1/resources` reports the honest capability; static-cache reports `supported=true` (positive control). Deleting the capability plumbing breaks these — no vacuous assertions. ## Validation (local; CI not awaited) - `cargo fmt --all -- --check` — clean - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` — clean - `cargo clippy --locked --all-targets -p onnx-genai-server -- -D warnings` — clean - Targeted engine + server tests pass. (One pre-existing failure, `registry::tests::failed_runtime_shrink_preserves_policy_and_ledger_limit`, also fails on the base commit `675b697b` in this CPU-only environment and is unrelated to this change; the other registry tests pass single-threaded — the multi-threaded flakiness is pre-existing global-VRAM-ledger state, not touched here.) ## Scope guardrails No existing admission/refusal behavior is weakened or changed (#743 pre-header 429, #745 committed-granule admission, #748 growth grants). This increment is purely about reporting the truth. Batched native decode is **not** implemented here. Refs #750. Does not close the issue. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 11, 2026
## Govern the `GroupQueryAttention` `WS_SCORES` workspace (#736) Next increment of the CUDA allocation-governance program (#736), chosen from the audit table in `docs/MEMORY_ARCHITECTURE.md` §2.1, which named this slice explicitly. Routes the `GroupQueryAttention` f32 reference-attention score buffer (`WS_SCORES`) — quadratic in `sq * present_capacity` and the largest remaining live GQA scratch — through the merged workspace-governance contracts instead of the raw `self.runtime.alloc_raw` pooled bypass. Closes #736 — no, **does not** close: this is one increment; the issue stays open. ### What changed - **`crates/onnx-runtime-ep-cuda/src/kernels/group_query_attention.rs`** - `Kernel::workspace_requirement` reports the exact governed bytes; `execute_with_workspace` consumes the executor-prepared slot; the compatibility `execute` path keeps its self-owned pooled scratch. - Planning and execution size the reservation through the **same** `gqa_reference_scores_bytes` helper, so the reserved and consumed byte counts cannot drift. A shortfall in the prepared buffer **errors deterministically** rather than silently under-allocating. - New free helpers: `gqa_reference_scores_path(dtype, q_seq, head_dim)` (the static reachability predicate) and `gqa_reference_scores_bytes(batch, heads, q_seq, kv_capacity)` (the shared, overflow-checked byte formula `rows * kv * sizeof(f32)`, `.max(1)`). - CPU-only source-scan **regression guard** (`raw_allocation_guard`) that fails if the governed slot reintroduces an unconditional raw pooled allocation — matches the #751 bar (a unit test that runs on CI without a GPU, not a `Select-String` external scan). Needle assembled at runtime so the literal does not self-match. - **`crates/onnx-runtime-session/src/executor/bindings.rs`** - Extended the centralised `is_planned_workspace_node` predicate to cover `com.microsoft::GroupQueryAttention` alongside `BlockQuantizedMoE`, `IndexShare`, and `Attention`. No parallel predicate added. - **`docs/MEMORY_ARCHITECTURE.md` §2.1** - Moved the GQA `WS_SCORES` slot into the governed-today rows, recorded derived bytes, documented the lifetime justification and the over-reservation finding, and reselected the next slice from evidence. ### Measured / derived bytes `gqa_reference_scores_bytes = (batch · num_heads · q_seq · present_capacity).max(1) · sizeof(f32)` — identical to the old `score_count.max(1) * 4`. - Reference path, `B=1, H=32, Sq=1024, present_capacity=1024, f32` ≈ **128 MiB**. - fp16/bf16 decode & prefill, and **all** f32 single-token decode: **0 bytes** (reported `WorkspaceRequirement::NONE`). ### Lifetime justification — `SessionPersistent` (not `StepScoped`) GQA's slot allocator grows each slot to the largest geometry seen and **retains** it across decode/prefill steps — a standing claim held for the whole session. So `WS_SCORES` is charged as `WorkspaceLifetime::SessionPersistent` + `MemoryRole::Workspace { step_scoped: false }`. This deliberately differs from the Attention Phase-2a scratch landed in #753, which was per-call `StepScoped`. Because `sq * present_capacity` is prompt-dependent, the executor grows the slot transactionally via `MappedGrowthGrant` (#748) against the device authority, and capacity refusal surfaces as the pre-header **429** typed path (`TierExhausted` / `CapacityUnavailable`, #743) rather than a late 500/OOM. A graph mixing QMoE, IndexShare, Attention Phase-2a (step-scoped) and GQA (session-persistent) now exercises **both** per-lifetime-class peaks; the classes reserve into separate executor slots so neither peak stands in for the other. ### Over-reservation finding (the #751 defect class) **There is over-reservation to avoid, and this PR avoids it.** `WS_SCORES` is materialized on **exactly one** GQA route: the f32 *reference* path, reachable only when `query dtype == f32 && !gqa_decode::supported(q_seq, head_dim)` (i.e. f32 prefill `Sq > 1`, or f32 `head_dim > 128`). The fused-flash path, the capture-safe f32/fp16 split-K decode kernels, and the phase-2a path never allocate a device score matrix — they stream softmax through registers/shared memory (phase-2a owns its own scratch). Reserving `WS_SCORES` unconditionally would charge the device authority for bytes the common fp16 decode path never touches — exactly the defect #751 found in IndexShare's present-K/V staging. `workspace_requirement` threads that static signal (query dtype + `gqa_decode::supported`) through the shared helper and reports **NONE** on every fp16/bf16 path and every f32 *decode* path, so they charge zero. ### cuBLASLt GEMM workspaces The audit named three shared 32 MiB cuBLASLt GEMM workspaces (`fused_gemm.rs`, `gemm.rs`, `matmul.rs`) as a cheap follow-up. They do **not** fall out of this change — they are separate step-scoped cuBLASLt workspaces, not GQA slots — and are left for a dedicated slice, as the task permitted. ### Validation (all local, green) - `cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --check` — clean. - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` — 0 warnings. - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` — 0 warnings. - `cargo test -p onnx-runtime-session` — **green** (mandatory 151+ suite). - `cargo test -p onnx-runtime-ep-cuda --features cuda --lib` — 318 passed (incl. the 4 new GQA governance tests). - GPU suites single-threaded (`--features cuda,gpu-tests -- --test-threads=1`): `group_query_attention_gpu` **25 passed** (incl. `gqa_gpu_fixed_decode_capture_replays_bit_identically`, which passed here — the #626 flake only appears under parallel capture contention), `qmoe_gpu` 29, `attention_gpu` 5, `allocation_counts_gpu` 2, `gqa_seqmajor_parity_gpu` 1. ### Next slice (recorded in the doc, chosen from evidence) `standard_attention.rs:602` (default-domain `Attention` pooled scores slot) — same quadratic shape and session-persistent retention — through the same contract with the same over-reservation discipline, then fold in the legacy `attention.rs:711` per-call path. ### Notes - Edits confined to workspace/allocation concerns; KV indexing untouched, to merge cleanly with the concurrent flash/prefill stride-descriptor work in this crate. - No `assert!` in any `Drop` (STATUS_STACK_BUFFER_OVERRUN on this platform). - Do **not** merge and do **not** close #736. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
…#755) (#798) Closes #755. Makes the authority-scoped **managed no-spill VMM** the **default** on the native CUDA path — no flag required — with **automatic weight streaming** when a model exceeds the resolved device budget. Being larger than the budget becomes a supported configuration rather than an error. Builds directly on the prerequisite #716 (PR #796), which made weight offload capture-compatible by paging weights under a stable virtual address, so flipping the default no longer disables CUDA graph capture for offloaded models. ### What "default" now means 1. **Managed VMM is on without a flag.** Native CUDA loads run memory-strategy inference (#752) unconditionally and it *drives* policy: allocation goes through the physical-handle pool (#740) with committed-granule admission (#745), governed workspace planning (#747/#753/#795), transactional growth grants (#748), against one server-owned authority per device (#744). 2. **Over-budget is not an error.** When the resolved budget cannot hold the weights, the runtime auto-enables weight streaming/offload instead of failing. 3. **`--vram-limit` is an override**, not the trigger. 4. **The legacy allocator stays reachable** for one release via `ONNX_GENAI_LEGACY_ALLOCATOR=1` (back-compat: `ONNX_GENAI_DYNAMIC_KV_WEIGHT_LENDING=0`). ### Observability The resolved budget, chosen strategy and offload state are printed at startup and exposed in `/v1/resources` under a new `memory_strategy` block, so the new default is observable rather than mysterious. ### Invariants preserved - **No silent WDDM spill** — committed physical bytes stay within the managed cap; `cuMemCreate` structurally cannot spill (#704). - **Capacity refusal stays a pre-header HTTP 429** (#743); committed-granule admission (#745) and growth grants (#748) unchanged. - **A model that fits does not start paging** under the new default — tested explicitly (the most likely regression). - **Capture stays ON when offload is active on the stable-VA path** (#796 property), pinned by the existing unit tests. - No prefetch re-introduced (#715); no assert inside any `Drop`. ### Measurement (native CUDA, deterministic counters lead) Ignored live e2e `qwen2_5_0_5b_managed_vmm_default_e2e` (Qwen2.5-0.5B int4 mobius, **same process**, interleaved managed-default vs. legacy opt-out): | scenario | strategy | offload | managed_no_spill | resolved budget | committed device bytes | page-ins / evictions | |---|---|---|---|---|---|---| | **A** fits, managed default (no flag) | `FullResident` | false | true | 7,730,940,928 | 381,681,664 | **0 / 0** (does not page) | | **B** synthetic over-budget (`--vram-limit 268,435,456`) | `DynamicWeightResidency` | true | true | 268,435,456 | 90,177,536 (« cap; no spill) | **434 / 432** | | **C** legacy opt-out (`ONNX_GENAI_LEGACY_ALLOCATOR=1`) | `FullResident` | — | **false** | — | — | — | | **D** explicit `--vram-limit 12 GiB` | override applied | — | — | 12,884,901,888 | — | — | In **B**, at that extreme synthetic cap the residency arithmetic **refused cleanly** ("no page is evictable") rather than spilling — the no-spill contract holds. **Caveats (measurement discipline):** - Wall-clock / tok/s is **not** reported: this box ranged 3.9–28 tok/s across identical runs and other agents build concurrently. Deterministic counters only. - A genuine native large-model run that exceeds VRAM was **not measured**: `qwen14b-zp` lacks `inference_metadata.yaml` (#384) so it is not native-loadable. The over-budget condition was **synthesized** with a small explicit budget, as stated. - **CUDA-graph segment count** is not a public counter; capture-ON-under-offload on the stable-VA path is pinned by the #796 unit tests, not re-measured here. ### Validation - `cargo fmt -p onnx-genai-engine -p onnx-genai-server -- --check` clean (`--all-targets --all` fmt hits Windows os error 206 on this box's long worktree paths — a command-line-length artifact, not a formatting issue). - clippy `-D warnings`: engine (`native-backend`; also `cuda,native-backend`), server, `onnx-runtime-ep-cuda --features cuda` — all clean. - Full `cargo test -p onnx-runtime-session` passes. - Engine `memory_strategy` unit tests (12, incl. 4 new) and server resources tests pass; cuda-gated engine tests incl. #796 capture arms pass. - **Pre-existing unrelated failure:** `failed_runtime_shrink_preserves_policy_and_ledger_limit` (server) fails identically on `origin/main` (confirmed via stash); its error is on the `host` authority, untouched by this native-CUDA-only change. Not fixed here (out of scope). Do **not** merge and do **not** close the issue via this PR beyond the `Closes #755` reference; leaving final disposition to review. 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
Governs the `GroupQueryAttention` packed QKV-projection staging buffers — the largest remaining evidenced non-QMoE workspace bypass after the three quadratic score matrices (Phase-2a #753, GQA `WS_SCORES` #795, default-domain `Attention` #736). References #736. **Do not merge; do not close the issue.** ## What was ungoverned `crates/onnx-runtime-ep-cuda/src/kernels/group_query_attention.rs` — the pooled slot allocator (`bytes.max(1)` per slot) raw-allocated the packed Q/K/V projection-staging buffers (`WS_PACKED_Q/K/V`). The unfused prep path splits an interleaved `[B, S, (H + 2·Hkv)·D]` query tensor into three device scratch buffers before the BSH→BNSH transpose and cache append. Those bytes were invisible to the device authority. ## Measured / derived bytes Formula: `align256(B·sq·num_heads·head_dim·elem) + 2·align256(B·sq·kv_num_heads·head_dim·elem)`, `elem = dtype.byte_size()` (f32=4, f16/bf16=2). At `B=1, num_heads=32, head_dim=128` (`q_hidden=4096`), `kv_num_heads=8` (`k_hidden=1024`), `sq=1024`: - **f32 packed prefill:** 16 MiB (Q) + 2·4 MiB (K,V) = **24 MiB** staging, plus the ~128 MiB f32 reference score region (#795) when that route is taken. - **fp16/bf16 packed prefill:** **12 MiB** staging, **0** scores. - **unpacked-input and fused-decode routes:** **0** staging. ## Which routes populate staging vs report `NONE` Staging is materialized on **exactly one** family of routes: a *packed* QKV tensor split on the unfused prep path (`packed_qkv && !fuse_prep`). - **Unpacked inputs** (Q/K/V arrive as three separate tensors) never split → staging regions are zero-length → charge **0**. This is the primary "size from use" finding (#751 class): the pre-#736 pool grew `WS_PACKED_*` on any path reaching `reserve`, but the separate-input route writes Q/K/V straight through. - **fp16/bf16 fused decode** (`q_seq==1` with an aliased device KV cache) fuses prep into the decode kernel and never splits. - The requirement returns `WorkspaceRequirement::NONE` when neither staging nor f32 reference scores are needed. ## Over-reservation found Yes (this is **not** a #736-default-Attention-style genuine-use-everywhere result). Two routes charge zero. One residual conservative reservation: `fuse_prep` aliases the present cache via a **runtime pointer comparison** not visible at planning, so planning reserves staging for every packed dispatch; on fused decode the reserved region is tiny (`q_seq==1`) and left untouched rather than corrupting. Every packed *prefill* genuinely splits, so this is a fallback-of-a-fallback idle. ## Lifetime model + justification **Session-persistent, single class** — verified against the `GqaWorkspace` slot allocator (`group_query_attention.rs:1191`), which grows each slot to the largest geometry seen and **retains** it across decode/prefill steps. Unlike the default-domain `Attention` scores (route-dependent lifetime needing both classes modelled), this lifetime is not route-dependent, so one class models the truth. The staging shares the same session-persistent composite reservation as the #795 `WS_SCORES` region (the executor reserves one slot per lifetime class and hands the kernel one view; the f32 reference-prefill route needs both live at once → disjoint regions of one buffer). Staging regions come first (shape-derived offsets, identical in planning and execution); the score region is last and unpadded so its capacity-dependent size extends into the reserved peak without perturbing staging offsets, and the composite total equals exactly the score bytes when no staging exists (preserves #795 accounting). ## Contract wiring - `Kernel::workspace_requirement` + prepare-only planning (#747); `execute_with_workspace` consumes the reserved slot. - Planning and execution size through the **same** `gqa_workspace_layout` / `gqa_packed_staging_bytes` / `gqa_reference_scores_bytes` helpers; `gqa_carve` refuses a short prepared view **deterministically** rather than under-allocating. - Growth via `MappedGrowthGrant` (#748); capacity refusal surfaces as the pre-header **429** typed path (`TierExhausted` / `CapacityUnavailable`, #743) — never a late 500/OOM. - Central `is_planned_workspace_node` predicate already covers GQA — **no parallel predicate added** (doc comment updated only). - Compatibility/opt-out `execute()` path keeps its self-owned pooled scratch. - No asserts in any `Drop`. - CPU-only source-scan regression test (`gqa_packed_qkv_staging_is_governed_not_raw_allocated`) + layout/requirement unit tests guard the governed site on CI without a GPU. ## Updated audit table `docs/MEMORY_ARCHITECTURE.md` §2.1 and the site-by-site audit: QKV staging moved to governed-today rows with the byte formula + route/lifetime split; the `WS_SCORES`/`:1033` line references corrected to current code (`:1191`/`:1203`, scores carved at `:2662`); the next slice named from evidence — **GQA BNSH transpose scratch** (`WS_Q_BNSH` `:2127` always-reserved, `WS_OUT_BNSH` `:2131` for `q_seq>1`), the always-used sibling (expect a genuine-use prefill result with a possible decode `NONE`). Present K/V remains a KV-residency concern, not workspace. ## Validation (local; CI not waited on) - `cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-session -- --check` ✅ - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` ✅ (no `--all-targets`, exact CI gate) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` ✅ - Full `cargo test -p onnx-runtime-session` ✅ - GPU (single-threaded): `group_query_attention_gpu` (25), `gqa_seqmajor_parity_gpu`, `gqa_fp16_gpu`, `standard_attention_gpu` (24), `attention_gpu` (5), `qmoe_gpu` (29), `allocation_counts_gpu` (2), plus new lib unit tests ✅ Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 12, 2026
…r); floor is granule-bound on qwen0.5b (#797 #801 #812 #811 #787 #777) (#827) ## Summary Realises the `layers × 2` KV floor **through the engine** by wiring the seq‑major fixed‑stride KV commit path through the shared `kv_commit::live_prefix_ranges` geometry, the last increment in the chain named by #797, #801 and #812. `DecodeCudaState::seq_major_kv_commit_requests` now *consumes* the same `kv_commit.rs` geometry the driver‑level residency test (`vmm_kv_layout_residency_gpu`) and the `kv_commit.rs` unit tests exercise, instead of duplicating the byte arithmetic inline — so the live commit geometry and the measured floor are single‑sourced and cannot drift (previously `kv_commit.rs` carried a `#![cfg_attr(not(test), allow(dead_code))]` because nothing in a non‑test build called it). Head‑major stays **byte‑identical**: on a growing packed bucket its per‑head stride *is* the bucket, so the `kv_heads` live‑prefix fragments tile the same contiguous `[0, bucket_bytes)` run — head‑major's dense ranges equal its bucket ranges. `vmm_growth_requests` is left as the flat bucket commit, now with a doc comment saying exactly that. Refs #797 #801 #812 #811 #787 #777. ## Headline: committed‑bytes comparison `physical_committed_bytes`, process‑isolated, `mobius_seqmajor_growth_parity_native_cuda`, qwen2.5‑0.5b: | Config (vmm‑inplace) | head‑major | seq‑major | |---|---:|---:| | capture ON, forced growth | 100,663,296 | 100,663,296 | | capture OFF, forced growth | 100,663,296 | 100,663,296 | | no‑growth control (`min_bucket=64`) | 100,663,296 | 100,663,296 | | realloc, forced growth | 786,432 | 786,432 | **The floor does not move on this model, and that is the correct result, not an un‑wired commit.** At `head_dim=64`/fp16 with a 512‑token hard max, a single binding's *entire reservation* is `512 × 2 × 64 × 2 = 131,072 B = 128 KiB`, far below the 2 MiB CUDA granule. So **both** layouts pin at exactly one granule per binding: `48 bindings × 2 MiB = 100,663,296 B`, which *is* the `layers × 2 = 48`‑granule floor. Seq‑major already sits on that floor; head‑major coincides because its 64‑token bucket is also sub‑granule. The `kv_heads×` separation (`vmm_kv_layout_residency_gpu`, qwen14b, 32,768‑token stride) only materialises once a **single head's live stripe crosses a granule** — `capacity × head_dim × elem ≥ 2 MiB`, i.e. ≈8,192 tokens at head_dim 128/fp16 (the measured head‑major crossover, #776). A 512‑max, 48‑token qwen0.5b run never approaches it, so #797's driver‑level `kv_heads×` win is **not reproducible end‑to‑end on this harness** — its geometry never leaves a single granule. This is the same class of precise negative result as #794 (identical bytes) and #812, stated with the same precision. ## Full per‑config counters (measured) | Config | growth_events | kv_bytes_moved | captures | replays | invalidations | growth_keeps | |---|---:|---:|---:|---:|---:|---:| | head vmm ON, growth | 3 | 688,576 | 4 | 39 | 4 | 0 | | seq vmm ON, growth | 3 | **0** | 1 | 45 | 1 | **3** | | head vmm ON, no‑growth control | 0 | 0 | 1 | 45 | 1 | 0 | | seq vmm ON, no‑growth control | 0 | 0 | 1 | 45 | 1 | 0 | | head vmm OFF, growth | 3 | 688,576 | 0 | 0 | 4 | 0 | | seq vmm OFF, growth | 3 | 0 | 0 | 0 | 1 | 3 | | head realloc OFF, growth | 3 | 688,576 | 0 | 0 | 4 | 0 | | seq realloc OFF, growth | 3 | 688,576 | 0 | 0 | 4 | 0 | - **Token IDs byte‑identical** across every configuration (`" Paris. It is the largest city…"`, 48 tokens). - Seq‑major keeps its captured graph across all 3 growths (`growth_keeps=3`, forced‑growth graph accounting `1/45/1` collapses onto the no‑growth control) — **the #811 keep decision is preserved**; the dense‑range commit changes neither device pointers nor physical shapes, so the four keep dependencies still hold. - Head‑major still moves 688,576 B and re‑captures per growth (`4/39/4`, `growth_keeps=0`) — unchanged. Wall‑clock intentionally omitted (deterministic counters answer the question; this box's identical runs vary 3.9–28 tok/s). ## Invariants held - Idempotent re‑commit; governor bills only newly mapped granules (via `seq_major_kv_commit_requests` → `commit_binding_ranges_with_mapped_growth`). - #748 `MappedGrowthGrant` growth path, #745 committed‑granule admission, #814 governor claim ordering, pre‑header 429 refusal (#743) — all unchanged. - No map/grow inside a captured region; no unmap under in‑flight replay; **no assert inside any `Drop`**. ## What this unblocks / does not - The dense‑prefix commit is now floor‑faithful and single‑sourced, so the `layers × 2` geometry is realised end‑to‑end **wherever a per‑head stripe exceeds a granule** (large `head_dim`/context). Demonstrating the `kv_heads×` win end‑to‑end needs such a model, **not** a change to this commit path. - It does **not** address the seq‑major seed/host‑mirror stride arithmetic (`read_present_kv` / `seed_prefix`), which still refuse under seq‑major (#812). Prefix sharing under seq‑major therefore still needs the BSNH seed as a separate follow‑on; the doc bullet is corrected to say so. ## Validation (all green, on an RTX 4060 Laptop, ORT 1.27, CUDA 13 wheels) - `cargo fmt -p onnx-genai-engine -- --check` ✔ (`--all-targets` full‑workspace fmt hits Windows os error 206, a command‑length artifact, not a formatting problem) - `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` ✔ - `cargo test -p onnx-runtime-session` ✔ (277 passed) - `cargo test -p onnx-genai-engine --features native-backend --lib` ✔ (454 passed; `kv_commit` 9/9) - **Mandatory** `mobius_seqmajor_growth_parity_native_cuda --ignored --test-threads=1` ✔ (2 passed, incl. the head‑major negative oracle) - `gqa_seqmajor_parity_gpu` ✔ (2), `gqa_shared_prefix_parity_gpu` ✔ (1), `group_query_attention_gpu` ✔ (26), `vmm_kv_contiguous_tail_gpu` ✔ (3), `vmm_kv_layout_residency_gpu` ✔ (2), `vmm_kv_token_major_floor_gpu` ✔ (2) — all `--ignored --test-threads=1`. Do not merge / do not close any issue. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 24, 2026
Four doc sites still told a reader that a `CudaPageStore` was on its way and that the host-backed tiers were *"emulation until Stage 3"*. #721's own body has said otherwise since 2026-08-18: | stage | status | |---|---| | 1. Extract `KvPageStore` + `HostPageStore` | **merged** (#726) | | 2. Make tier migration real, charged to the governor | **merged** (#742) | | 3. `CudaPageStore` backed by the VMM arena | **superseded** | | 4. Native CUDA decode consumes `PageTable` | **superseded** | On native CUDA, device KV paging is owned by the **VMM layer** — `CudaVmmAllocator` with its physical-handle pool (#740), committed-granule admission (#745) and growth grants (#748). It already reserves a stable virtual range and maps committed granules behind it. A second page allocator inside `onnx-genai-kv` would **duplicate that ownership rather than complete it**. ## Why this is worth its own commit Comments are where a cancelled plan survives longest. *"until Stage 3"* reads as a scheduled prerequisite, so the next person to open the file either builds a dead stage or spends an afternoon working out that they shouldn't. The issue body was updated; the code that people actually read was not. ## What changed | site | before → after | |---|---| | `page_table.rs:19` `DevicePageSpan` | *"a future `CudaPageStore` … would reserve a stable virtual range"* → the VMM layer already does exactly that; the span stays because the trait boundary is what keeps `head_token_row()` off the store contract and what lets an out-of-tree store exist at all | | `page_table.rs:94` `KvPageStoreFactory` | *"Stage 3 can supply a factory"* → an out-of-tree extension point, not an in-crate hook | | `tiered.rs:1` module doc | *"until Stage 3 supplies a CUDA store"* → host-addressable tiers are this crate's **end state**, not a waypoint | | `paged_cache.rs:658` `preempt_sequence` | *"host-backed emulation until Stage 3"* → declared residencies that stay declared | | `MEMORY_ARCHITECTURE.md:73` | *"Stages 2–5 in #721"* — names a stage 5 that does not exist and implies three open ones → the real per-stage status, pointing at the issue body as authoritative | That last row already said device KV paging belongs to the VMM layer, so this only stops the two halves of one table cell contradicting each other. ## Deliberately left alone `copy_from_host`'s *"a future device store can implement this with an explicit host-to-device transfer"* is **not** the same claim and stays. An out-of-tree device store remains possible and supported; what is not coming is one **in this crate**, on this repository's roadmap. ## Validation Doc comments and one doc table cell. No behaviour change — `cargo check -p onnx-genai-kv` and `cargo doc -p onnx-genai-kv --no-deps` both clean. 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
ReclaimableMappedHolderregistration and transactional RAIIMappedGrowthGrantserve --vram-limitselects managed no-spill VMM/pool mode;ONNX_GENAI_DYNAMIC_KV_WEIGHT_LENDING=0restores compatibility fallbackRetry-AfterRefs #736
Refs #717
Refs #745
Refs #747
Review fixes
Validation
cargo fmt --all -- --check-D warnings: passed-D warnings: passed14B / 6 GiB E2E
Model:
C:\Users\justinchu\dev\models\qwen14b-zp; release native-CUDA binary.max_tokens=24: HTTP 200, correctParis; weight allowance 4,831,838,208 bytes; owned 5,534,384,128 <= 6 GiBLatest commit:
440a1df4. PR remains open and unmerged.