Repository navigation
feat(native-decode): thread batch through CUDA decode shape/IO layer, fixed to 1 (2b-impl-1, #750) - #913
Merged
Conversation
… fixed to 1 (#750) Stage 2b-impl-1 of the native batch-N decode split. Introduce a �atch field on DecodeCudaState (constructor-fixed to 1) and thread it through the shape and IO-binding layer so the batch extent is no longer a hard-coded 1 literal at the sites a batch-N step will need: - persistent_state_shapes: KV/recurrent axis-0 -> batch (both arms), BNSH axis order preserved (GQA reads present_capacity from axis 2). - persistent_output_shape: batch-axis (0) symbolic collapse -> batch; other symbolic axes stay 1. Logits binding passes batch; auxiliary outputs keep the historical collapse-to-1 (their batch-symbol pinning is 2b-impl-4). - mask/input_ids/embeds/position_ids allocations, extend_mask, rewind and both mask-growth paths (VMM in-place + realloc) use batch. Behaviour-preserving by construction: every current entry point constructs at batch 1, so shapes and decode are byte-identical. KV growth/commit geometry, the token writer, device argmax and batch-symbol capture pinning stay batch-1 and are owned by 2b-impl-2..4. Validated: fmt, clippy (engine native-backend all-targets; ep-cuda cuda), engine native-backend tests (+2 new shape-threading unit tests), onnx-runtime-session, the six fused_batch_prefill guards, the mandatory mobius_seqmajor_growth_parity_native_cuda GPU gate solo (2/2, seq-major byte-identical to head-major across growth), and a solo on-hardware batch-1 byte-identity profile_native run on qwen14b-zp: tokens exact-match the #750 reference and htod_bytes_per_token=2,349,010,944 / page_ins_per_token=188 / captures=1 replays=14 fallbacks=0 invalidations=1 / peak_committed 7,725,907,968 < managed_limit / oversubscribed=0 all match the #905 baseline. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #913 +/- ##
==========================================
+ Coverage 77.78% 78.43% +0.64%
==========================================
Files 354 354
Lines 148042 148042
Branches 148042 148042
==========================================
+ Hits 115161 116114 +953
+ Misses 28390 27433 -957
- Partials 4491 4495 +4
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
Owner
Author
|
Validated independently before merging. Byte-identity on hardware, solo ( Gates: fmt clean, 93/93 I read the diff rather than trusting the batch-1 result, because a change that is byte-identical by construction can hide a wrong axis that only appears at batch N. \persistent_state_shapes\ keeps batch on axis 0 and the grown seq extent on axis 2 for both KV layouts, which is the binding contract the CUDA GQA node depends on; the new \cuda_persistent_state_shapes_thread_batch_axis_0\ test pins that at \�atch=3\ where a transposed axis would actually show up, rather than at 1 where it could not. That is the right test for this increment. The scope discipline is what makes this mergeable: \�atch\ is constructor-fixed to 1 and the doc comment names exactly what is not generalized (token writer, KV growth/commit geometry, device argmax, capture symbol pinning). #905's inventory held up under implementation — worth recording, since the last three increments on this line each returned a correction instead. Merging. 2b-impl-2 (KV growth/commit) is next. |
justinchuby
added a commit
that referenced
this pull request
Aug 14, 2026
…l it the route (#919) `PROGRESS.md` and `WEIGHT_OFFLOAD.md` both still described the #864 hybrid as the open route to beating the OS, and the north-star row went as far as *"its two blocking risks are now measured away"*. True when written; now stale **in the worst direction**, because it points the next person at a lever that has already been built and measured to fail (#912). Corrected with what was actually found. The two blocking risks really were measured away first (#880: ~5.6 GB/s real strided GEMV, bit-identical capture with host-mapped pointers) — so what killed it was **neither of them**: - aggregate *distinct* host-mapped reads above ~0.44–0.65 GB/step **silently return stale data** (48 cold weights collapsed generation 16 → 3 tokens; a single read is bit-identical at 1/8/16/32, and a copy-instead A/B pins the fault on the read, so it is an aperture ceiling rather than a per-read bug); - capped at a provably safe 256 MiB, the arm ran **0.73 tok/s against WDDM's 7.84**. The safe aperture (0.26 GB/step) is smaller than the traffic it must displace (~0.6 GB/step) — short **by construction**, not by tuning. Row 3's honest status is now stated as such: **on WDDM consumer hardware we do not beat the OS, we correctly decline to compete with it, and no route currently on the table changes that.** It then names the two levers that do remain — batching (which creates the intra-step reuse that makes residency pay at all) and the admission side of the residency gap. Also in this pass: - Removes a justification the residency-policy item was **borrowing** from the hybrid, and adds what #892/#901 since established: eviction order is not the cause and is worth at most ~10% of the gap, so ~90% is admission. Anyone picking that up should not spend it on eviction policy. - Refreshes the #750 batching item, which still listed stage 2b as remaining after #891 and #913 landed. - Adds seven link-reference definitions (#783, #863, #864, #874, #877, #880, #912) that were used in bracket-link form but never defined, so they rendered as literal text. Docs only. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
justinchuby
added a commit
that referenced
this pull request
Aug 14, 2026
…refuse seq-major relocation (2b-impl-2, #750) (#920) ## Stage 2b-impl-2 (#750): make KV growth and commit geometry batch-aware Increment 2 of 4 from the split in `docs/NATIVE_BATCH_DECODE_2B_IMPL_SCOPING.md`. Extends #913 (which threaded a constructor-fixed `batch=1` through the shape/IO layer) by making the **KV growth and commit geometry** compute over `batch` sequences instead of assuming a single sequence at batch index 0. `batch` stays **constructor-fixed to 1**, so every change is **byte-identical at batch 1** (the primary gate). Batch-N is **not** turned on (2b-impl-4). Does **not** close #750. ### What is batch-general after this increment (per layout) **Relocation-free, made batch-general + tested at batch=3:** - **Mask commit byte extents** — construction (`mask_bytes` / `full_mask_bytes`), `grow_vmm_mask_in_place` memset, `vmm_growth_requests` mask island, and the realloc `build_grown_mask` memset now all span `batch` rows (they zeroed/committed one row before). - **Head-major (BNSH) growing bucket** — the flat `0..bucket` KV range (`checked_shape_bytes` already includes the batch axis) and the `copy_kv_prefix*` / `zero_kv_suffix` block loop are already batch-general because batch is the **outermost** axis (block count = `batch × kv_heads`, each `(batch, head)` stripe re-strides independently). Made explicit + unit-tested (`kv_growth_byte_layout_refuses_seq_major_batch_n_and_allows_head_major_batch_n`). - **Seq-major fixed full-context stride commit** — `live_prefix_ranges` (now takes `batch`), `seq_major_kv_commit_requests`, and `zero_seq_major_committed_tail` emit **one dense run per sequence**, each a fixed `capacity × bytes_per_token` apart. The per-sequence stride is the *pinned* full context, so this is **relocation-free** and batch-general. New CPU tests: `seq_major_batch_n_is_one_dense_run_per_sequence`, `head_major_batch_n_is_one_fragment_per_sequence_head`. **Refused (named, tested) — NOT batch-general:** - **Seq-major growing-bucket / realloc growth at `batch > 1`.** `kv_growth_byte_layout` now `bail!`s with a named error. A growing bucket's per-sequence stride **is** the mutable bucket capacity, so growing the bucket would **relocate every sequence `b > 0`'s** already-written KV — defeating the seq-major fixed-stride "no KV moves / keep the captured graph" contract and the `moved_bytes_per_token == 0` accounting. Head-major has no such problem (batch outermost). A clear refusal is the correct deliverable here rather than a silently wrong capacity-dependent stride. The relocation-free batch-N seq-major path is the fixed-stride commit above, which never calls this helper. ### Byte-identity + counters (hardware, solo, `nvidia-smi` verified empty) `profile_native --model qwen14b-zp --ep cuda --tokens 16 --steady`: ``` generated_token_ids [96347, 3375, 724, 11, 358, 2776, 14589, 311, 6723, 429, 498, 3003, 2581, 6617, 315, 752] # EXACT #750 reference cuda_graph captures=2 replays=28 fallbacks=0 invalidations=1 cuda_graph_measured captures=1 replays=14 fallbacks=0 invalidations=1 vmm_arena ref_underflows=0 byte_underflows=0 unaccounted_committed_bytes=0 resource_vram oversubscribed_bytes=0 throughput 7.32 tok/s # within the 3.9-28 band vs 7.84 baseline -> NOT a finding (#863) ``` Token identity is exact and the deterministic counters (`captures=2 fallbacks=0`, underflows 0) match the baseline conditions of the run command. ### Mobius parity gate (MANDATORY, run SOLO) `mobius_seqmajor_growth_parity_native_cuda -- --ignored --test-threads=1`, `nvidia-smi` verified empty: **2/2 passed**. Seq-major **kept** the captured graph across **3 growth events** with `kv_bytes_moved=0`, `growth_keeps=3`, and a **byte-identical 48-token stream** to head-major across VMM in-place (capture ON/OFF), realloc, and the no-growth control. This is the most relevant gate since this increment touches growth/commit, and it confirms the seq-major fixed-commit + tail-zero changes are byte-identical at batch 1. ### Other validation (all green, local) - `cargo fmt -p onnx-genai-engine -- --check` - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` (exact CI gate) - `cargo clippy -p onnx-genai-engine --features native-backend,cuda -- -D warnings` - `cargo test -p onnx-genai-engine --features native-backend,cuda --lib native_decode` — **96 passed**, incl. the 3 new batch>1 geometry tests. ### What did not survive contact with the code (scoping corrections) 1. **`live_prefix_ranges` is NOT "another crate."** SCOPING §3 item 6 and §4 say the seq-major `live_prefix_ranges` path "must be proven batch-correct (another crate)." It actually lives **in this crate**: `crates/onnx-genai-engine/src/native_decode/kv_commit.rs:106`. Its only non-test caller is `seq_major_kv_commit_requests` (`cuda.rs`). No cross-crate signature change was needed; the driver-level residency measurement in `onnx-runtime-cuda-memory` is a **mirror**, not a caller, so it was unaffected (its clippy gate still passes). 2. **"KV growth stride math is already batch-general" is only true for head-major.** SCOPING §3 ("Already batch-general") claims the block-count stride math re-strides N batch blocks correctly for both layouts and "this is the part people assume is the blocker; it is not." That holds for **head-major** (batch outermost). For **seq-major growing bucket at `batch > 1`** the block loop would *silently relocate* sequences `b > 0` (capacity-dependent block stride) and break the `moved_bytes_per_token == 0` / "no KV moves" invariant. This PR turns that silent misbehavior into an explicit refusal. 3. **The finer split is fixed-stride vs growing-bucket, not head- vs seq-major.** The relocation trap is specific to a *growing bucket*; the seq-major *fixed full-context stride* commit is relocation-free and was made fully batch-general here. ### Judge-it-yourself Merge `origin/main`, then: `cargo test -p onnx-genai-engine --features native-backend,cuda --lib native_decode` reproduces the 3 batch>1 geometry tests; the solo mobius run reproduces the batch-1 byte-identity; `kv_growth_byte_layout_refuses_seq_major_batch_n_and_allows_head_major_batch_n` reproduces the refusal proof. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
justinchuby
added a commit
that referenced
this pull request
Aug 14, 2026
…batch-aware (2b-impl-3, #750) (#924) ## What this is **Increment 3 of 4** of native batch-N decode (#750), stage 2b-impl-3: make the **device argmax / greedy token selection** path batch-aware. `batch` stays **constructor-fixed to 1** — this delivers the *capability*, it does not turn batch-N on (that is 2b-impl-4). Byte-identity at batch 1 is the primary gate. Threads a `batch` count through the whole device-argmax stack: - the CUDA kernel (`onnx-runtime-ep-cuda` `kernels::device_argmax`): partials grid gains a sequence row (`blockIdx.y`), finalize gains one block per sequence (`blockIdx.x`), result header is `2 × batch` words (a `[token_id, capture_error]` pair per sequence), scratch is `batch × 2 × partial_count`; - `device_argmax_scratch_words(elements, batch)` and `scratch_words(elements, batch)`; - the `ExecutionProvider::device_argmax` trait + the CUDA impl; - `DeviceIoBinding::device_argmax(elements, batch, result)`; - the `__native_greedy_argmax` buffer allocation (now `2·batch + scratch`); - `read_greedy_result` now delegates to a new `read_greedy_result_batch` returning `Vec<(TokenId, u32)>` of length `batch`. ## What is now batch-general, and what is not **Batch-general:** the argmax kernel selects `batch` tokens from `[batch, 1, vocab]`, one grid row per sequence, writing `batch` `[token_id, capture_error]` pairs. `read_greedy_result_batch` reads all N rows. **Not batch-general (correctly out of scope):** the **token writer** `write_decode_inputs` (`cuda.rs`) still writes one i64 at offset 0. Per the scoping doc that writer is item 3 = **2b-impl-1**, and the batch-N entry point that would present N tokens and consume N rows is **2b-impl-4** (doc §2.1, §3). So nothing yet writes N selected tokens into N input slots — there is no batch-N caller until impl-4. I did **not** touch `write_decode_inputs`. ## Is batch-1 tie-breaking provably unchanged? Yes — two independent reasons 1. **Launch geometry is bit-identical at batch 1.** `partial_count` is per-sequence and unchanged; at `batch == 1` the partials grid is `(partial_count, 1, 1)` and finalize is `(1, 1, 1)` — exactly the prior launch. The sequence offset (`blockIdx.y`, `blockIdx.x`) is 0, the result header is 2 words, and the scratch offset/layout are unchanged, so the kernel writes the same words and the 8-byte read-back is unchanged. 2. **The reduction operator is order-invariant by construction.** `argmax_update` keeps *max value, min index on ties* — commutative and associative — so the selected `(value, index)` is independent of how the vocab is partitioned into blocks/partials. Even if the grid shape changed, ties could not drift. This is why greedy selection here is not a numerics risk the way a floating-point sum would be. ## Byte-identity result (gate 5, solo) `profile_native --model qwen14b-zp --ep cuda --tokens 16 --steady`, `nvidia-smi` verified empty before the run (`0 MiB`), greedy (`sampling: OFF`): ``` generated_token_ids: [96347, 3375, 724, 11, 358, 2776, 14589, 311, 6723, 429, 498, 3003, 2581, 6617, 315, 752] ← exact match to the #750 reference cuda_graph: enabled=true captures=2 replays=28 fallbacks=0 invalidations=1 cuda_graph_measured: captures=1 replays=14 fallbacks=0 invalidations=1 vmm_arena: ref_underflows=0 byte_underflows=0 unaccounted_committed_bytes=0 resource_vram: oversubscribed_bytes=0 ``` Deterministic counters all green (`captures=2`, `fallbacks=0`, underflows `0`, `oversubscribed=0`), matching the doc's batch-1 baseline. No throughput delta reported (this box's wall-clock ranges widely; #863/#884). Note: this run had `MANAGED_WEIGHT_STREAMING`/`CUDA_GRAPH` env unset, so the `weight_offload_*` byte counters read 0 — that is an env difference, not a code effect; token identity and the graph/underflow counters are the gates and they hold. ## Mobius parity (gate 4, solo): **2 passed** ``` cargo test -p onnx-genai-engine --features cuda,native-backend --test mobius_seqmajor_growth_parity_native_cuda -- --ignored --nocapture --test-threads=1 => test result: ok. 2 passed; 0 failed ``` (Requires the anaconda site-packages CUDA DLLs on PATH, else CublasLt is not found and the child emits no measurement — a failing gate, not a pass.) ## Other gates - `cargo fmt -p onnx-genai-engine -- --check` ✓ (also ep-cuda/ep-api/session) - `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` ✓ (exact CI gate, no `--all-targets`) - `cargo test -p onnx-genai-engine --features native-backend,cuda --lib native_decode` ✓ 96 passed - device test `device_argmax_batch_selects_per_row_max_with_deliberate_ties` ✓ (GPU, solo) — plus the two pre-existing 151,936-element argmax tests still pass. ## The falsifying control `device_argmax_batch_selects_per_row_max_with_deliberate_ties` puts the max at a **different position in every row**, makes each row larger than one block's grid-stride reach (so multiple partials are reduced), and plants a tie at an **earlier index than the peak** in two rows. It asserts (a) the expected per-row token, (b) `capture_error == 0` per row, and (c) each batched row equals a standalone single-sequence argmax of the same logits. A shared reduction buffer, a wrong per-sequence stride, or global (rather than per-sequence) tie-breaking would each fail this. ## Scoping doc: what survived contact with the code The doc's item 5 (device argmax → 2b-impl-3) matched the code exactly: `read_greedy_result` calls `binding.device_argmax(vocab, …)` returning one token, and the kernel + `device_argmax_scratch_words` live in `onnx-runtime-ep-cuda`. **No correction to the split.** One observation worth recording (not a re-split): the doc lists `write_decode_inputs` under 2b-impl-1 (item 3, "Need N tokens × N rows"), but merged #913 left it single-token and it is still single-token here — because the N-token write is inseparable from a batch-N *caller*, which the doc itself places in 2b-impl-4 (§2.1). So "where the N selected tokens go" genuinely lands with the entry points in impl-4, consistent with §2.1/§3; impl-1's item 3 could only ever write the one token a batch-1 constructor presents. Closes nothing — leaves #750 open. Do not merge on my behalf. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
justinchuby
added a commit
that referenced
this pull request
Aug 14, 2026
…n batch-N on (2b-impl-4, #750) (#930) ## What this is Increment **4 of 4** of native batch-N decode (issue #750), stage 2b-impl-4: pin the batch symbol, wire the token writer, turn batch-N on. Increments 1-3 (#913, #920, #924) were byte-identical by construction with `batch` constructor-fixed to `1`. **This is the increment that changes behaviour**, so it is measured rather than asserted. **#750 stays open** for the 2c caller. > **Correction (2nd commit).** The first version of this PR reported that > "capture and weight-streaming are mutually exclusive on this build." **That > was a harness artifact, not a runtime property**, and it has been corrected > in place. The sweep loaded the native session via `load_with_resolved_io`, > which passes `cuda_offload_policy: None`; `load_with_cuda_options_and_io` > then leaves `weight_offload_stable_va` at its conservative `false` default > (`native_decode/load.rs:793`), so the capture predicate > `weight_offload_enabled && !weight_offload_stable_va` declined capture *by > construction* whenever offload was on. The knob I set (the entry point) set > the flag that produced the decline that I read as a finding. The harness now > drives the **fully governed** `Engine` session, whose effective offload > policy carries the managed no-spill stable-VA authority — and **capture and > streaming coexist**, exactly as #796/#836 built and verified. All the > capture/streaming numbers below are re-measured on that governed path. ## What is on and under what flag - **`ONNX_GENAI_NATIVE_DECODE_BATCH`** (default `1`). At the default the decode path is **byte-identical to `main`**. Setting `N>1` binds a stable batch extent at session construction and drives N rows through the persistent captured decode path. - Batch symbol pinned for aux outputs via `persistent_output_shape` (axis-0 → `batch`); KV seq-axis pinning unchanged. - Token writer wired: `write_decode_inputs_batch` (N token IDs → N slots), `write_position_binding_batch` (rank-R mrope per sequence), `extend_mask` (N mask rows). Single-token entry points delegate to the batch writers (byte-identical at batch 1). - `decode_cuda_greedy_batch` / `NativeDecodeSession::decode_greedy_batch` is the batch-N greedy entry point; `native_decode_batch()` exposes the extent. - `Engine::native_decode_session_mut()` (new, native-backend) lets the measurement harness drive batch-N decode against the governed session. - `profile_native --native-decode-batch-sweep <list> [--native-decode-batch-context L]` builds the governed `Engine` per batch shape and exercises the captured decode path. **The batch-N generation driver (real per-sequence prompt seeding) is stage 2c and out of scope here** — the harness fans one identical token to all N rows, which is a strict correctness probe and, because the weight/page/capture counters are content-invariant (#884), the trustworthy measurement vehicle. ## Mandatory gates — all re-run locally after the correction | Gate | Result | |------|--------| | `cargo fmt -p onnx-genai-engine -- --check` (+ `-p onnx-genai-bench`) | **pass** | | `cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings` | **exit 0** | | `cargo test -p onnx-genai-engine --features native-backend,cuda --lib native_decode` | **96 passed, 0 failed** | | **mobius parity** `--test mobius_seqmajor_growth_parity_native_cuda -- --ignored --nocapture --test-threads=1` | **2 passed** (SOLO) | | **batch-1 byte-identity on hardware** | **exact match** (below) | ### Batch-1 byte-identity with counters Reference token IDs matched **exactly** with the feature compiled in and batch configured to 1 (`qwen14b-zp`, `--steady --tokens 16`): ``` [96347, 3375, 724, 11, 358, 2776, 14589, 311, 6723, 429, 498, 3003, 2581, 6617, 315, 752] ``` Counters matched the documented baseline: `captures=2`, `replays=28`, `fallbacks=0`, `htod_bytes_per_token=2349010944`, `page_ins_per_token=188`, `hit_rate=78.32%`, `byte_hit_rate=70.16%`. ## Correctness of batch-N (row identity) `qwen14b-zp`, all N, every step, on the governed captured+streaming arm: **all N rows == row 0** (the #892 cross-row-corruption detector) and **row 0 == the batch-1 stream** at N=2,4,8 (the scoping-doc §3.8 guard). Token *values* at N>1 are content-arbitrary (real per-sequence seeding is 2c); the counters are not. ## Recapture per shape (#891) — **measured, capture survives batch-N WHILE streaming** `qwen14b-zp`, governed engine, managed streaming forced (`ONNX_GENAI_MANAGED_WEIGHT_STREAMING=1`, `ONNX_GENAI_CUDA_GRAPH=1`), empty past, measured window, per batch shape. **Non-zero htod confirms weights are actually streaming during these captured steps:** | N | graph enabled | captures | replays | fallbacks | invalidations | htod_bytes_per_token | |---|---------------|----------|---------|-----------|---------------|----------------------| | 1 | true | 1 | 6 | 0 | 0 | 2,994,222,848 | | 2 | true | 1 | 6 | 0 | 0 | 1,502,098,688 | | 4 | true | 1 | 6 | 0 | 0 | 736,831,040 | | 8 | true | 1 | 6 | 0 | 0 | 424,642,624 | **Capture survives batch-N as implemented, on the real streaming decode path** — one capture per shape, zero invalidations, zero fallbacks at every N, with ~3 GB/token of weights streaming underneath. This is the positive resolution of #891's open risk the line has been waiting for, and it is the part the original PR got wrong by measuring the wrong arm. ## 1/N weight-stream amortization on the captured decode path — **measured** Same governed captured+streaming arm as above (managed path, #874 note: the WDDM default is the OS shared-memory path with htod=0, so managed streaming is forced). Model `qwen14b-zp` (8.3 GB weights), empty past, `--tokens 8`. Emitted tokens = steps × N. | N | htod_bytes_per_token | page_ins_per_token | abs page_ins/forward | evictions | byte_hit_rate | |---|----------------------|--------------------|----------------------|-----------|---------------| | 1 | 2,994,222,848 | 276.8 | 2214 | 1273 | 61.96% | | 2 | 1,502,098,688 | 139.9 | 2239 | 1293 | 61.83% | | 4 | 736,831,040 | 66.6 | 2132 | 1185 | 62.55% | | 8 | 424,642,624 | 41.6 | 2662 | 1678 | 56.84% | `htod_bytes_per_token` falls **~1/N** (2.99 GB → 1.50 → 0.74 → 0.42; a 7.05× reduction N=1→8 against the ideal 8×). Amortization vs the N=1 ideal extrapolation: N=2 **+0.3%**, N=4 **−1.6%**, N=8 **+13.5%**. The N=8 shortfall is not noise — it is the #866 offset (next section). All four shapes ran SOLO with `nvidia-smi` verified empty; no contention flags. **The amortization is approximate here, not byte-exact.** On this governed managed arm the evictor's working set and page order vary per session load, so the per-forward `abs page_ins` drifts (2214 / 2239 / 2132 / 2662) rather than staying identical. That is the honest behaviour of the *captured streaming* path. For contrast, the earlier *eager* offload arm (offload on, capture declined, fixed `DEVICE_BYTES=6e9`) gave **byte-exact** 1/N with an identical `abs page_ins=1791` at every N — because a fixed budget with no capture is fully deterministic. Both arms agree that the weight stream is paged once per forward and the per-produced-token cost falls ~1/N; the captured arm is the one that matters for #750 and it now measures cleanly. Per the task instruction, **no throughput delta is reported as a finding** (this box ranges 3.9–28 tok/s across identical configs). The deterministic counters above are the evidence. ## The #866 offset — **measured within-sweep on the governed arm** #866 makes the weight budget elastic against occupancy, so batch-N (N× the working set) should shrink the weight residency and push `byte_hit_rate` down. The governed captured arm shows exactly this **within a single sweep**: - N = 1, 2, 4: `byte_hit_rate` ≈ **62%** (61.96 / 61.83 / 62.55, flat within run-to-run noise), evictions ≈ 1200–1300. - N = 8: `byte_hit_rate` drops to **56.84% — a −5.12-point offset**, evictions rise to 1678, and `abs page_ins/forward` rise 2214 → 2662. That offset is precisely why the N=8 per-token htod runs **+13.5% above** the ideal 1/N line: the N× working set ate into the elastic weight budget, so each forward had to re-stream more. This is the #866 offset the line predicted, measured on the real decode path rather than inferred. (The earlier PR could only show it as an *arm-to-arm* budget difference because the un-governed sweep never let the elastic budget respond to batch; the governed engine does.) ## Honest summary of what is and isn't proven - **Proven:** batch symbol pinned; token/mask/position writer wired; batch-N runs end-to-end N=1,2,4,8; batch-1 byte-identity; row identity and cross-batch identity; **capture survives batch-N *while weights stream*** (captures=1, invalidations=0, htod≈3 GB/token) on the governed engine path; ~1/N weight-stream amortization on the captured decode path; the within- sweep #866 byte-hit-rate offset (−5.12 pts at N=8). - **Corrected:** the original "capture and streaming are mutually exclusive" claim was a harness artifact of `load_with_resolved_io`'s conservative `weight_offload_stable_va=false` default; the governed path proves they coexist. - **Named limitation (in the spirit of #920's seq-major `bail!`):** on the governed managed-streaming arm the 1/N amortization is *approximate*, not byte-exact, because the evictor's per-load working set varies and the #866 offset grows the stream at N=8. This is a property of managed streaming, not a batch-N defect. - **Out of scope (2c):** real per-sequence prompt seeding / the batch-N generation driver. The harness fans one identical token to all rows, so token *values* at N>1 are content-arbitrary; the counters are content- invariant (#884) and are the measurement. **Do not merge** — leaving #750 open for the 2c caller. --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
justinchuby
added a commit
that referenced
this pull request
Aug 14, 2026
…mpl-4 (#939) docs(progress): record the #750 batch-N line as complete through 2b-impl-4 The four-way split (#905) is done: #913, #920, #924, #930 all landed. The entry still described increments 2-4 as remaining and carried none of what they found. The headline is the result the line was started for, reproduced independently rather than taken from the PR: capture survives batch-N while weights stream. Solo, nvidia-smi verified empty, managed streaming forced, qwen14b-zp at context 64 -- N=2 and N=4 both show enabled=true captures=1 invalidations=0 fallbacks=0 with htod_bytes_per_token 1,234,271,232 and 547,528,704, byte_hit_rate 68.64% -> 72.17%. captures=0 replays=8 inside the measured window is capture-once/replay-many working as intended. Also records what each increment actually contributed, because three of them returned a correction to my scoping rather than the scoping: - #920 drew the distinction the scoping had missed -- a fixed full-context stride commit is batch-general in both layouts, a growing bucket only in head-major -- and refuses seq-major growing-bucket growth at batch>1 with a named bail! rather than computing a capacity-dependent stride that would have passed every batch-1 test and corrupted the first batch-N run. - #924 proved batch-1 tie-breaking unchanged two ways, including that argmax_update is a total order so the reduction cannot depend on tree shape. And it records the near-miss, because that is the part worth remembering. #930's first version reported capture and weight streaming as mutually exclusive, from a real decline message naming a real predicate. The sweep harness called load_with_resolved_io, which passes cuda_offload_policy: None, so stable_va fell to its conservative unwrap_or(false) and capture declined by construction. The harness asserted pointer-instability; the runtime was never asked. That would have closed the most valuable open question on this line with a convincing wrong answer. Finally, the #866 offset is now measured within-sweep rather than inferred: byte_hit_rate flat ~62% at N=1,2,4 then -5.12 points at N=8 as evictions go 1273 -> 1678, which is why N=8 sits +13.5% above the ideal 1/N line. Batch-N multiplies KV occupancy by N and eats the elastic weight budget -- the mechanism #866 predicted, now observed rather than argued. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
justinchuby
added a commit
that referenced
this pull request
Aug 14, 2026
…increment corrected (#941) docs(#750): mark the 2b-impl scoping complete and annotate what each increment corrected All four increments landed (#913, #920, #924, #930), but the document still read as an open plan, and 2c is being scoped against it right now. The corrections are the more useful part of it now, because three of the four returned one, so they are recorded both as a table up front and inline at the exact line that was wrong -- a reader trips over the claim, not over the summary. - §3's "the pleasant surprise: KV growth stride math is already batch-general" holds only for head-major. A fixed full-context stride commit is batch-general in both layouts because the batch axis is outermost; a growing bucket is not, because seq-major's per-sequence stride IS the mutable capacity, so growth relocates every sequence b>0. #920 refuses that case with a named bail! rather than computing a capacity-dependent stride, which would have passed every batch-1 test and corrupted the first batch-N run. The section conflated the two paths. - §4 item 6 places live_prefix_ranges in onnx-genai-kv; it is in this crate at native_decode/kv_commit.rs:122, so no cross-crate coordination was needed. Line number checked against current main rather than copied from the PR, since it moved when the batch parameter was added. Records the outcome the line was started for, reproduced by me rather than quoted from the PRs: capture survives batch-N while weights stream -- enabled=true captures=1 invalidations=0 fallbacks=0 with htod_bytes_per_token 1,234,271,232 at N=2 and 547,528,704 at N=4, byte_hit_rate 68.64% -> 72.17%, solo with managed streaming forced. That closes #891's flagged risk positively. And records the near-miss, because it is the part most worth carrying into 2c: #930's first version reported capture and streaming as mutually exclusive, from a real decline naming a real predicate, because its harness called an entry point that passed cuda_offload_policy: None and stable_va fell to a conservative default. The harness asserted pointer-instability; the runtime was never asked. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
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.
Stage 2b-impl-1 (#750): thread
batchthrough the CUDA decode shape / IO-binding layer, fixed to 1Increment 1 of 4 from the split proposed in #905 (
docs/NATIVE_BATCH_DECODE_2B_IMPL_SCOPING.md). Behaviour-preserving by construction — abatchvalue is introduced and threaded so the batch extent is no longer a hard-coded1literal at the shape/IO sites a batch-N step will later need, but every current entry point still constructs at batch 1, so decode is byte-identical. Batch-N is not turned on.Does not close #750 (four-way split, three increments remain).
What changed
DecodeCudaStategains abatch: usizefield, set to1innew().persistent_state_shapes: axis-01→batchin both the growing-KV and fixed-recurrent arms. BNSH axis order preserved — the GQA node readspresent_capacityfrom axis 2 regardless ofkv_layout, so axis 0 stays batch / axis 2 stays seq.persistent_output_shape: batch axis (0) symbolic collapse →batch; every other symbolic axis stays1. The logits binding passesbatch; auxiliary outputs keep the historical collapse-to-1 (pass1explicitly) — their batch-symbol pinning iscollect_unit_symbolswork owned by 2b-impl-4.input_ids/inputs_embeds/position_idsallocations,extend_mask,rewind's mask reshape, and both mask-growth paths (VMM in-placegrow_vmm_mask_in_place+ legacy reallocbuild_grown_mask) usebatch.Validation (all green)
cargo fmt -p onnx-genai-engine -- --checkcargo clippy --locked --all-targets -p onnx-genai-engine --features native-backend -- -D warningscargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings(untouched crate; run anyway as the CI gate)cargo test -p onnx-genai-engine --features native-backend(incl. 2 new tests)cargo test -p onnx-runtime-sessionfused_batch_prefillguardsmobius_seqmajor_growth_parity_native_cuda -- --ignored --test-threads=1, run solo withnvidia-smiverified empty: 2/2 passed, seq-major byte-identical to head-major across forced growth, capture ON/OFF, VMM/realloc (growth_keeps=3,kv_bytes_moved=0for seq-major; identical 48-token stream). (I hit self-inflicted GPU contention from a sibling agent running the same test in another worktree, stopped my run, coordinated, and re-ran solo per Latent gap: seq-major capture-retention proof covers KV binding pointers only, not weight buffers (original flaky-gate claim RETRACTED — contended measurement) #851.)Hardware batch-1 byte-identity (solo,
nvidia-smiverified empty before the run)profile_native --model qwen14b-zp --ep cuda --backend native --tokens 16 --steady:Every deterministic counter matches the #905 baseline exactly. Wall-clock is not headlined (WDDM-paging noise, #863).
Was #905's file/line inventory accurate?
Yes, substantially — it held up under implementation. Every §3 site I touched exists and pins batch at axis 0 as described, and the shape layer threaded cleanly without dragging in a later increment. Four refinements worth recording:
write_decode_inputsandread_logits(named in the 2b-impl-1 list) needed no edit at fixed batch=1.write_decode_inputswrites a singlei64at offset 0 — the per-row loop is a behaviour change that belongs to a batch-N writer/entry point (2b-impl-4), not a shape literal.read_logitsalready reads throughself.logits_shapeand returnsVec<Vec<f32>>, so it is already batch-general once the logits shape carriesbatch. The 2b-impl-1 intent for these two is satisfied by threading the shapes they consume, not by editing the functions. Minor scoping refinement, not a defect.position_idsbatch axis is index 1 for rank-N mrope ([N, B, S]), not index 0. The doc's[1,1] / [rank,1,1]is correct but doesn't call this out; the fix places batch at[rank, batch, 1](and[batch, 1]for rank 1).rewind's mask logical reshape andbuild_grown_mask(legacy realloc). Forgrow_vmm_mask_in_placeI threaded only the shape (batch axis) and left the commit byte-range (0..bytes) to 2b-impl-2 — the shape/commit boundary the doc drew is real and is respected.[1, heads, seq_len, head_dim](head-major host-mirror path, not the batch-N decode target — it alreadybail!s under seq-major), and the QMoE workspace-planning zeros[1, input_len, hidden](prepare-only planner). Both are batch-1 by nature of the non-batch-N paths they serve; 2b-impl-4 can revisit if a batch-N caller ever reaches them.The mobius gate was runnable this time — both
qwen2.5-0.5b-q4_0-mobius(head) and...-mobius-seqmajor(seq-major-patched) model dirs are present on this box, so unlike #905 this increment cleared the mandatory KV-governance gate solo.