Repository navigation
perf(cuda-ep): capture-eligible pointwise ops via build-time growing-symbol classification (C1) + re-anchor 35B QMoE oracle lock on fp32 teacher-forcing (#722) - #728
Conversation
…symbol classification (C1) Classify the graph's growing (KV-length) symbols once at build/graph-load time from attention past/present KV axes (reusing the #700 recurrent-state machinery), then mark a pointwise/elementwise/bitwise/prelu op capture-eligible iff NONE of its symbolic output dims is a growing symbol. Pinned symbolic dims (batch, query seq_len=1, heads) are admitted; anything carrying a KV-length symbol stays eager. This collapses the 35B-A3B QMoE hybrid GDN decode from 154 to 34 CUDA-graph segments. Removes the earlier strict all-Dim::Static rule, which was too strict: it treated a PINNED symbolic dim the same as a GROWING one and kept every GDN pointwise op eager. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The C1 capture classifier legitimately flips the 35B-A3B greedy token@119 from 33803 to 46283 — an inherent fp16 inter-layer near-tie (top-2 gap < 0.1 logit), both within fp16 noise of the fp32 oracle (see #722, DeepSeek-R1 #158 precedent). The old lock asserted the autoregressive captured token@119 == 33803, which is capture-dependent (a coin-flip), so it broke under C1. Re-anchor the correctness lock on the capture-INDEPENDENT invariant that BOTH captured and eager paths satisfy: the fp32-oracle argmax via fresh-engine teacher-forcing of a FIXED canonical context (prompt + recorded 33803-path prefix). Teacher-forcing is prefill-dominated and capture-independent, so it re-derives 33803 with capture on or off. - Primary lock: QMoE fresh-engine teacher-forced argmax == 33803 + margin band. - Dense int4 teacher-forced argmax == 33803 + margin (replaces the broken autoregressive dense assertions). - fp32-oracle CPU cross-check at the same canonical context (env-gated). - The autoregressive token@119 signal is preserved as a DOCUMENTED, NON-FATAL #722 observation (logs 33803-or-not) so the captured-vs-eager fragility stays visible as a tripwire. - #695 continuation test: demote the 33803 pin to a #722 note; assert the capture-independent reused==fresh invariant instead. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #728 +/- ##
==========================================
- Coverage 79.85% 79.36% -0.50%
==========================================
Files 329 331 +2
Lines 135046 137462 +2416
Branches 135046 137462 +2416
==========================================
+ Hits 107840 109094 +1254
- Misses 22880 24028 +1148
- Partials 4326 4340 +14
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
|
Revise the build-time growing-symbol CUDA-graph-capture classifier to
address all three of Harry's blocking findings while preserving the
35B-A3B QMoE decode collapse (154->34 captured segments).
Finding 1 (representative aliasing): apply the growing-symbol denylist to
both an op's INPUTS and OUTPUTS. Broadcast can substitute a pinned
representative on an aliasing op's output, but the first op in the chain
still carries the raw growing KV symbol on an input, so the input-side
check keeps it eager without a union-find pass.
Finding 2 (incomplete coverage): derive growing symbols from two
complementary sources instead of a 3-op allowlist -- recognized attention
ops now include CompressedSparseAttention, plus a generic scan of the
model's declared past*/present* rank-4 KV I/O. The rank-4 +
symbolic-penultimate guard excludes fixed-capacity conv/recurrent state so
GDN pointwise ops stay capturable.
Finding 3 (permissive oracle tripwire): restore a FATAL autoregressive
tripwire in qwen36_35b_a3b_qmoe_divergence -- assert token@119 is one of
the two known-benign fp16-tie outcomes {33803, 46283}; any other token
(real capture corruption) fails CI. fp32-teacher-forcing primary lock and
#722 documentation retained.
The task-preferred pinned allowlist was implemented and measured but
collapses 34->94 segments (it excludes benign fresh warm-decode symbols),
so the accepted denylist fallback is used. Adds 5 unit tests; validated
34 segments / 83.11 tok/s on 35B, 0.5b byte-exact capture ON==OFF, and the
oracle test passing capture ON and OFF.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…#728) Close Harry's two remaining findings on the capture-eligibility classifier: Finding 1 (downstream aliasing escape): shape inference unifies broadcast symbols into a lower-id/named representative, so a growing KV symbol can be aliased into a pinned-looking representative that a downstream pointwise consumer carries on both edges -- wrongly admitting it as capture-seq- independent. Fix: make the growing property equivalence-class-closed. Rebuild the broadcast union-find (mirroring shape-inference broadcast_dim's lower-id representative rule) over default-domain elementwise-broadcasting nodes, mark a class growing if ANY member is a collected growing symbol, and add all members of growing classes to the growing set. Over-inclusion is fail-safe (stays eager, never admits a growing op). This closes the downstream-consumer hole transitively. Finding 2 (CSA output-5 omitted): dynamic ratio-4 CompressedSparseAttention mints a fresh GROWING `selections` symbol on output 5's LAST axis, which the generic past/present rank-4 penultimate scan misses. Added last_axis_outputs/last_axis_growing_symbol to collect it. Shipped the canonicalized, equivalence-class-closed growing DENYLIST (with CSA output-5 coverage) rather than the pinned allowlist: Deckard measured the naive allowlist at 154->94 (dissolves the collapse by failing to recognize legit pinned warm-decode symbols), while the closed denylist recovers 154->34 and is provably complete for every attention/KV op here (GQA, Attention, IndexShare, CompressedSparseAttention incl. output-5 selections, and the generic declared past/present scan) while staying fail-safe under the closure. Tests: added growing_symbol_alias_keeps_downstream_consumer_eager (Finding 1 downstream regression -- fails on un-closed logic) and csa_output5_selections_symbol_keeps_consuming_ops_eager (Finding 2); kept/ updated the pinned->capturable, growing->eager, and recurrent-state classifier tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…drift-free (PR #728) Round-4 revision (Batty) fixing Harry's round-3 Finding 1 escape. The executor's growing-symbol closure previously re-implemented a partial copy of shape inference's broadcast unification covering only elementwise ops, so MatMul batch dims / Einsum ellipsis / Concat non-concat axes / Expand could erase a growing KV-length symbol into a pinned-looking representative and let a downstream consumer stay wrongly capturable (silent CUDA-graph decode corruption). Path B (authoritative, complete-by-construction): shape inference now records every (loser, winner) symbol unification at its single broadcast_dim chokepoint and persists it onto Graph::symbol_unifications. The executor closes its growing set over that authoritative record, covering every current and future broadcasting handler with zero per-op enumeration. Recording is additive: the returned representative and every inferred dim are unchanged (inference output byte-identical; full shape-inference suite passes). - onnx-runtime-ir: add Graph::symbol_unifications field. - onnx-runtime-shape-inference: SymbolInterner records unifications in broadcast_dim's two-distinct-symbolic branch; infer_graph_scoped persists them. - onnx-runtime-session: close_growing_under_symbol_unification reads the authoritative record; removed op_broadcasts_elementwise/broadcast_dim_from_right. - Tests: new matmul_batch_alias_keeps_downstream_consumer_eager (Harry's escape, fails on HEAD 817eee5), inference completeness tests for MatMul + Concat. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…rd; default classifier to fail-safe (PR #728) Round-4 fix: SymbolInterner::lower interned a derived DimExpr (e.g. seq_kv*8 from Reshape([-1])/Flatten) to a fresh SymbolId while recording NOTHING, so the capture-eligibility classifier never learned the fresh symbol depended on a growing KV length. A downstream pointwise consumer carrying only that fresh symbol was wrongly classified capture-safe -> silent decode corruption. Step 1 (byte-identical, additive): record derivation provenance at the `lower` chokepoint. SymbolInterner now records (derived -> source) edges for every constituent symbol of an interned derived expression, opaque markers for overflow/negative degrades, and the inference symbol floor. These persist onto Graph::{symbol_derivations, symbol_opaque, inference_symbol_floor}. lower still returns the identical Dim, so inference stays byte-identical. Executor: one worklist closure over the directed adjacency of broadcast unifications (undirected) + derivation edges (directed source->derived), so a growing source poisons its derived products without over-poisoning pinned roots. Step 2: implement BOTH classifiers behind ONNX_GENAI_CAPTURE_CLASSIFIER, defaulting to FAIL-SAFE (an op is capturable only if every edge symbol provably traces via the lineage record to pinned, non-growing roots; any inference-minted symbol without recorded provenance -> eager). The denylist-over-provenance fallback is switchable for the coordinator's GPU A/B measurement. Tests: reshape/flatten derived-growing tests (confirmed FAIL on 571ea0d), fail-safe pinned-derived-capturable and untraceable-minted-eager tests; all prior broadcast/MatMul/Concat/CSA/alias tests kept green. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Harry's round-5 REJECT: the classifier's verdict was an additive HINT, not the authoritative veto the PR premise requires. `capture_shape_eligible` OR-ed the classifier bit (`seq_independent`) with runtime-extent fallbacks (`product == 1`, `is_fixed_decode_shape`), so a node the classifier CORRECTLY disqualified (e.g. `Reshape([seq_kv*8]) -> Sigmoid`, growing symbol) was still admitted to CUDA-graph capture via the rank-1 `[N]` fallback, baking a stale growing extent and silently corrupting decode on replay. `seq_independent == false` is a definitive "not provably safe to capture" in BOTH classifiers (fail-safe: not-all-pinned; denylist: some-symbol-growing), so the OR terms can only ever WRONGLY override a real disqualification. Make the classifier the sole authority: `capture_shape_eligible(seq_independent, _shape) = seq_independent`. Signature unchanged, so all callers (bitwise, pointwise, prelu, elementwise binary/silu-mul) are covered centrally. Remove the now-unused `is_fixed_decode_shape` helper. Add capture-LEVEL regression tests (not just the classifier predicate): a disqualified growing rank-1 consumer (`seq_independent=false`, `[320]`) is never eligible and yields no capture-safe signature (confirmed FAIL on HEAD 0fd87df, PASS now); a pinned single-token decode shape (`seq_independent=true`, `[1,1,heads,dim]`) stays eligible. cargo test -p onnx-runtime-ep-cuda --lib: 284 pass. -p onnx-runtime-session --lib: 145 pass. --features cuda --lib: 148 pass. clippy/fmt clean. 35B oracle + segment-count A/B left to the coordinator's warm GPU. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…round-7) Round-7 REJECT (Harry): classifier-disqualified nodes could still enter CUDA-graph capture because many CUDA kernels return CaptureSupport::Supported unconditionally and never override set_capture_seq_independent (UnaryMathKernel, NotKernel, BitwiseNotKernel, and dozens more). A disqualified GROWING rank-1 value consumed by such a kernel bakes a stale launch grid into the replayed device graph -> silent decode corruption. Fix the whole class centrally instead of per-kernel whack-a-mole: make the build-time capture classifier an authoritative HARD VETO at the single capture-admission chokepoint (Executor::node_capture_reason). Before consulting the kernel's own capture_support(), decline capture with a new SeamReason::ClassifierDisqualified whenever node_capture_seq_independent(graph, node, capture_growing_symbols) is false. The verdict comes from executor/classifier state already computed independently of any kernel, so no kernel opinion can re-admit a disqualified node. The veto is strictly additive: capture-eligible iff (classifier seq-independent) AND (kernel Supported); over-declining is correctness-safe, under-declining is corruption. Tests (crates/onnx-runtime-session/src/executor/tests.rs): - classifier_disqualified_node_is_vetoed_despite_supported_kernel: real capture-admission path with an unconditional-Supported mock kernel; fails on pre-veto code (node admitted), passes post-veto. - classifier_qualified_node_with_supported_kernel_is_admitted: additive check. - warm_decode_seeding_admits_previously_unresolved_capture_safe_node: clears the disqualifying set to isolate the seeding transition from the classifier veto. onnx-runtime-session --lib: 147 passed. onnx-runtime-ep-cuda --lib --features cuda: 289 passed, 2 ignored. cargo fmt --all --check clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…(PR #728 round-8) The round-7 central capture veto ran as the FIRST check in Executor::node_capture_reason, shadowing the EP structural policy. A control-flow/sequence node (If/Loop/Scan) that also references a disqualifying growing symbol was reported as SeamReason::ClassifierDisqualified (EagerDeviceSeam) instead of HostControlFlowOrSequence (HostSeam) — numerically safe (both stay eager) but corrupting the public capture-segmentation contract that distinguishes host round-trips from eager device seams. Reorder so the veto runs AFTER ep.plan_capture_region structural classification (host control-flow/sequence and unresolved-shape declines keep precedence) but BEFORE the kernel admission check, so an unconditional-Supported kernel still can never re-admit a disqualified device node. Add disqualified_control_flow_node_reports_host_seam_not_device_seam pinning the precedence. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ROVE for PR #728 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The prepare-only workspace planning path added by #747 calls get_or_create without the growing-symbol capture classification introduced here. Compute it with the same helper used by build and dispatch so prepare-only planning keys the same kernel that execution later uses. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
|
Rebase/merge validation from the coordinator: this branch was 21 commits behind One real semantic conflict, now fixed in 090c63c: the prepare-only workspace planning path added by #747 ( Local validation on the merged tree (Windows, RTX GPU box):
Merging. |
…ecisions (#756) Records two owner decisions that redirect in-flight work, each with the measured counter-evidence it has to beat. **Contiguous-VA KV for batching (#750).** Each sequence gets its own contiguous device virtual address with physical granules mapped underneath on demand; the attention kernel keeps seeing a flat buffer and never learns page tables. This supersedes #721 stage 3 unless the route fails. The record carries the result it must beat: #721 stage 4 measured a stable full-context VA committing 1.5 GB where bucket growth commits 48 MB, because the decode kernel reads the full padded shape and masks, forbidding tail decommit. **Managed no-spill VMM as default (#755)**, with automatic weight streaming instead of hard OOM. Sequenced explicitly behind #716, because weight offload and CUDA graph capture are mutually exclusive today — flipping the default first would silently disable capture for exactly the models that need offload, giving back the 154 -> 34 segment collapse from #708/#728. Docs only, no code paths touched. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
## Summary C2 follow-up to #728. This closes the capture-contract bug behind its 33 residual LinearAttention seams and audits the rest of the CUDA kernel directory for the same class of violation. ### Established mechanism A direct capture of `LinearAttentionKernel::execute` reproduced the failure before this change: ``` CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED: operation not permitted when stream is capturing ``` The unconditional trailing `synchronize()` aborted recording. Session quarantine then classified every sibling LinearAttention node as `CaptureRecordingFailed`. | 35B-A3B decode | Segments | LinearAttention declines | Decode | |---|---:|---:|---:| | Before (#728 measurement) | 34 | 33 × `CaptureRecordingFailed` | 11.959 ms/tok (83.62 tok/s) | | After | Not measurable on this host | Direct kernel capture: 0; capture/replay succeeds | Not measurable on this host | Because #728's complete seam census was exactly those 33 LinearAttention nodes, the deterministic expected result on the same artifact is 34 → 1 segment. I am not presenting that as measured: the required 35B artifacts are absent here and the available GPU is an 8 GB RTX 4060. ### Why the synchronization is not load-bearing LinearAttention launches one kernel on the executor stream and performs no host readback or second-stream work. Native decode deliberately aliases past/present recurrent state, but each thread loads its full state column before overwriting it. The new regression captures and replays four changing decode steps with that aliasing and verifies output plus recurrent state byte-for-byte against eager CUDA. The sync is therefore only eager asynchronous-error surfacing; same-stream ordering supplies visibility during capture. ### CUDA kernel audit Capture-aware trailing synchronization: - Activations - ArgReduce - ConvTranspose - FlashAttention - FusedGelu - GridSample - LinearAttention - SpaceToDepth and EyeLike (`structural.rs`) Truthful/dynamic admission for allocation or host-staged paths: - Attention: warmed fused path supported; Phase-2a allocation path unsupported - CompressedSparseAttention: host-staged fallback explicitly rejected during capture - NonZero: unsupported because it allocates/uploads coordinate metadata per call - LpPool: unsupported because it allocates/uploads/frees metadata per call `capture_sync_contract.rs` source-audits every kernel synchronization and fails on any new unreviewed unconditional sync reachable by a capture-supported path. Intentional unsupported/fallback sites are explicitly allowlisted. No assertions were added to `Drop`. ## Validation Passed locally with the requested CUDA wheel paths, including cuBLASLt: - `cargo fmt --all -- --check` - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` - CUDA EP library suite: 308 passed - `linear_attention_gpu`: 7 passed - `attention_gpu`: 5 passed, 1 benchmark ignored - capture synchronization contract test: passed - session quarantine classification test: passed Pure-attention `qwen2.5-0.5b-int4-onnx-native` produced byte-identical 64-token output with capture off/on: - capture off: 31.932 ms/tok (31.32 tok/s) - capture on: 7.057 ms/tok (141.70 tok/s) The re-anchored 35B oracle binaries compiled and were invoked with `ONNX_GENAI_CUDA_GRAPH=0` and `=1`; both real-model cases skipped because `/home/justinchu/qwen36-35b-a3b-qmoe-artifacts` is unavailable. The same-artifact 34 → 1 segment measurement and 35B oracle run are required before merge. --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
## 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>
…capture (#716) (#796) ## Summary Weight offload previously forced whole-step CUDA graph capture **OFF**: the pager returned a different device pointer on every page-in (`alloc_raw`/`free_raw`), and a captured graph bakes pointers into its recorded nodes, so the two were mutually exclusive (`resolve_graph_capture_enabled` returned `false` whenever offload was on). The models that most need offload — those that do not fit in VRAM — were exactly the ones that lost graph capture and all of the #708/#728 capture-fragmentation wins (35B-A3B decode 154 -> 34 segments). This PR gives each retained weight `key` a **reserved-once device virtual address** and remaps physical granules underneath on page-in/eviction, so a captured graph replays correctly across evict->repage. It then removes the offload/capture mutual exclusion, **gated on the paging path actually being stable-VA** (the managed no-spill authority). Hard prerequisite for #755 (managed no-spill VMM as the default with automatic weight streaming). ## What changed, by step **Step 1 - prove the mechanism in isolation first** (`crates/onnx-runtime-cuda-memory/tests/vmm_stable_va_weight_slot_gpu.rs`, **passes on this HW**): a captured graph that reads from a stable VA still produces correct results after the pages under that VA are swapped for different physical granules holding **different weight values**. Extends the #727 remap-at-same-VA proof to the weight-value-swap case. **Step 2 - stable-VA slot conversion** (`crates/onnx-runtime-ep-cuda/src/weight_paging.rs`): - Per-key `StableWeightSlot { va, len }` reserved once through the **existing #740 authority-scoped pool** via `allocate_committed`/`carve()` suballocation. No second allocator, no per-weight reservation (the mistake that made #733 net-negative). Charging stays on the global granule-ref **0->1 / 1->0** attribution. - `CudaWeightPage::Drop` decommits granules but **keeps** the VA for stable slots (`deallocate_span` only for transient bypass pages). Never asserts in `Drop`. - Respects committed-granule admission (#745): reserving VA is free; only committed physical bytes are charged. 2 MiB granule sizing per #776. - **Enforceable ordering rule (not a comment):** page-ins - and thus eviction-driven `decommit` of any stable slot - run entirely under the residency mutex; capture is capture-once/replay-many with no page-ins during replay; the engine additionally declines capture when a step reports a bypass/eviction (`GLOBAL_BYPASSED_PAGE_INS`). So no `decommit` of a baked VA can occur while a replay is in flight - the case #727 explicitly did **not** prove safe. No `cuMemMap`/grow inside a captured region. - Residency GPU test `vmm_retained_weight_key_keeps_a_stable_virtual_address_across_repage` proves the VA is identical across evict->repage with one reused granule and **no accounting regression**. The existing `vmm_weight_admission_reuses_owned_granules_without_runtime_alloc_free` and `weight_offload_gpu` suite still pass (no regression). **Step 3 - remove the exclusion, gated on stable-VA** (`native_decode/{cuda,mod,load}.rs`): `resolve_graph_capture_enabled` now takes `weight_offload_stable_va` and keeps capture OFF only when `weight_offload_enabled && !weight_offload_stable_va`. The signal is `policy.enabled && policy.managed_no_spill` (the explicit byte `--vram-limit` authority that installs the VMM arena + physical pool). The pointer-unstable `alloc_raw`/`free_raw` compatibility path keeps the old capture-off exclusion. **#723's `StableResident` eviction policy is unchanged** - this changes only the addressing mechanism underneath it. **Step 4 - measure** (`crates/onnx-genai-engine/tests/stable_va_capture_offload_e2e.rs`): drives the managed no-spill VMM stable-VA authority and verifies **token-identical** greedy decode with capture OFF and ON in one process, dumping the deterministic offload counters (`page_ins`, `hits`, `evictions`, `bypassed_page_ins`, `vram_alloc_ns`/`vram_free_ns`). ## Honest measurement caveat The full sustained-paging throughput/counter comparison the issue asks for could **not** be completed in this environment: - The genuinely-offloaded model (`qwen14b-zp`, 15.5 GB) is **not native-loadable** - it has no `inference_metadata.yaml` (`model.io` ports), which is the separate #384 load-chain gap. The native decode path (where this change lives) cannot load it. - On the native-loadable 0.5B model, the managed no-spill authority is **pin-or-fail**: either the ~381 MB weight set fits fully resident (0 page-ins) or a step cannot commit its working set (admission/headroom failure after heavy eviction, ~504 evictions observed). There is no clean budget window that both pages sustainedly *and* completes a full decode - that is exactly what #755's automatic weight streaming is for. What **is** verified end-to-end here: under the managed no-spill VMM stable-VA authority, decode with capture OFF and capture ON is **byte-identical** (32/32 greedy tokens), i.e. capture is transparent and no longer force-disabled. The addressing mechanism itself is proven **deterministically** by the two GPU tests in steps 1-2. A negative/partial result was an acceptable deliverable per the task; this is a positive mechanism result with a partial end-to-end result. ## Validation (local; CI not waited on) - `cargo fmt -p onnx-runtime-ep-cuda -p onnx-genai-engine -p onnx-runtime-cuda-memory -- --check` PASS - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` PASS (no `--all-targets`, the CI gate) - `cargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warnings` PASS (also linted the new cuda-gated test under `--features cuda,native-backend`) - `cargo test -p onnx-runtime-session` PASS (full suite) - GPU tests single-threaded PASS: `vmm_stable_va_weight_slot_gpu`, `vmm_retained_weight_key_keeps_a_stable_virtual_address_across_repage`, `vmm_weight_admission_reuses_owned_granules_without_runtime_alloc_free`, `weight_offload_gpu` (5/5) - engine `native_decode` unit tests PASS incl. new `weight_offload_on_stable_va_path_keeps_graph_capture` and updated `weight_offload_forces_graph_capture_off` Refs #716, #755, #727, #740, #723. **Do not merge; do not close #716.** Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>
…detail (#825) Justin asked to update `docs/PROGRESS.md` with the latest progress and drop stale info. It had grown to **1166 lines** of day-by-day changelog spanning 2026-07-15 → 2026-07-27, with the entire recent arc missing. ## Changes (1166 → 303 lines) - **Added recent milestones (2026-07-28 → 2026-08-12), newest-first:** - EP plugins run inside ORT via the plugin-EP C ABI (#762) and ship on PyPI — **`nxrt-ep-cpu` + `nxrt-ep-cuda` 0.1.0.dev5 both LIVE** (#819/#824/#820). - Absent-optional-output memory-safety wave (#762 review; Miri-clean). - VMM / KV-layout / offload / batching residency wave — managed no-spill VMM default, layout-governed KV floor (768× reduction), offload+capture coexist, #736 over-reservation audit (#755–#814). - Issue triage + autonomous fixes (#784/#785/#786/#788/#789); CPU-EP TopK correctness+perf (#774/#775). - Qwen3.5/3.6-27B hybrid Gated-DeltaNet native CUDA (#779); GLM-4-9B + DeepSeek-V2-Lite native + ORT 1.28 3-config fairness (#766/#770/#771). - Fused QMoE decode kernel + QMoE arc concluded; CUDA-graph capture trilogy, version-selectable CUDA (#708/#728/#757/#760), megakernel study (#769). - **Compressed stale detail:** the ~40 dated July day-by-day sections + verbose diffusion/LoRA/ORT2 narrative collapsed into a ~35-line "Foundational milestones" spine; superseded perf-wave numbers dropped; §1–§40 table condensed to a compact grid. - **Header refreshed:** last-updated → 2026-08-12; Published line now lists both EP plugins on PyPI + Miri CI. - **Preserved & updated** the DeepSeek/GLM known-gaps checklists and added new open items. Docs-only; full narrative remains in `.squad/decisions-archive/` and `docs/benchmarks/`. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Summary
Ships C1 — the build-time growing-symbol pointwise-capture classifier — collapsing the 35B-A3B QMoE hybrid GDN decode from 154 → 34 CUDA-graph segments, together with a re-anchor of the 35B QMoE oracle correctness lock onto fp32 teacher-forcing (#722) so the lock is capture-independent.
This is the culmination of the 35B-A3B CUDA-graph capture-fragmentation wave (roofline investigation → GDN-Split C3 #708 → this C1). C3 removed the host-sync seams; C1 collapses the dominant pointwise segment count.
The classifier (pinned-vs-growing rationale)
Classify the graph's growing (KV-length) symbols once at build/graph-load time — structurally, from attention
past/presentKV sequence axes only (reusing #700'srecurrent_past_names()/is_recurrent_state_shape()/ KV-length-symbol machinery, DRY). Collect them into aHashSet<SymbolId>.A pointwise/elementwise/bitwise/prelu op is
seq_independent(capture-eligible) iff NONE of its symbolic output dims is in the growing set. Pinned symbolic dims — batch=1, query seq_len=1, heads — are admitted; anything carrying a KV-length symbol stays eager.Dim::Static? It was proven a no-op: GDN pointwise ops carry[Symbolic(batch=1), Symbolic(seq=1), Static(16), Static(128)]— both symbolic dims are pinned, never all-Static — so strict-C1 kept every GDN op eager and collapsed nothing.[Symbolic(heads),Static(128)], heads ∉ growing → capturable;[Symbolic(seq_kv),Static(128)], seq_kv ∈ growing → eager. Runtime shape values cannot distinguish these; the KV-axis set can.sequence_lensymbol is NOT marked growing — it is symbolic but PINNED (=1, held constant across every decode replay). Only the KV-cache length axis actually grows. This distinction is the entire win.EP-API / build-time plumbing: the flag is set at kernel CREATION (
KernelFactory::create), before the single warm step — no runtime diff, no 2nd warm step, never gates the main-executor arm. (Earlier runtime-diff designs were structurally wrong: main arms at warm#1, so a warm#2 diff corrupts decode.) One shared helper feeds all 5 pointwise kernels; no per-op copies, no per-model hardcoding.seq_len>1 guard: the capture region is keyed to the concrete pinned-symbol binding (batch=1, query seq_len=1). A seq_len>1 entry forms a different (or uncaptured) region and cannot replay a seq_len==1 capture. The flag only asserts "no growing-symbol dependence," so the pinned query symbol is never misclassified.
The #722 re-anchor (why it's needed and why it's a stronger invariant)
C1 legitimately flips the 35B-A3B greedy token@119 from 33803 → 46283 — an inherent fp16 inter-layer near-tie (top-2 gap < 0.1 logit), both within fp16 noise of the fp32 oracle. Heavy accumulators are already fp32; this is an inter-layer fp16 reassociation/rounding coin-flip, not a reduction bug (diagnosed in #722, same class as DeepSeek-R1 #158).
The OLD lock asserted the autoregressive captured token@119 == 33803 — capture-dependent, so it broke under C1 (which moves the captured stream onto the eager stream). Per #158 precedent, an fp16 near-tie satisfied by BOTH paths is not a bug; the lock belongs on the fp32 oracle.
Re-anchored
qwen36_35b_a3b_qmoe_divergence.rs:prompt + SHARED_GENERATED_PREFIX, the recorded 119-token 33803-path prefix). Teacher-forcing is prefill-dominated → capture-independent → both capture-on and capture-off re-derive 33803. The context is fixed/recorded, NOT reconstructed from the live stream (under C1 that stream is the 46283-path).reused_argmax == fresh_argmaxinvariant.The fp32-teacher-forcing anchor is stronger: it's the ground-truth correctness invariant that both capture arms satisfy, instead of pinning on which side of a coin-flip one autoregressive stream lands.
Results (post-rebase onto origin/main, KV_MAX_LEN=4096, GPU0)
(Isolated earlier run measured 12.263→11.763, +4.3%; this re-measure ran on a busier box. The 154→34 collapse is the robust invariant; the 33 residual seams are LinearAttention's load-bearing C2 sync, out of scope.)
Correctness — byte-exact proof on BOTH models AND capture on/off
ONNX_GENAI_CUDA_GRAPH=0/1) — the fp32-teacher-forcing anchor holds regardless of capture.qwen2.5-0.5b-int4-onnx-native: decode byte-identical, native CUDA == ORT CUDA == expected, capture-ON and capture-OFF (KV-length-dependent ops stay eager).cargo fmt --all --checkclean;cargo clippyclean on changed crates.Blast radius
This touches the SHARED elementwise/bitwise/pointwise/prelu capture gate used by ALL models — a wrong broadening would silently corrupt any model. Mitigations: the classifier keys off graph metadata (the growing KV-axis symbol set), never runtime shape values; the negative-assertion unit test guards the growing→eager direction; byte-exact validation on both a hybrid (35B) and a pure-attention (qwen2.5-0.5b) model, capture on and off.
Notes
Do not self-merge.
Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com