Repository navigation
feat(native-decode): make KV growth and commit geometry batch-aware, refuse seq-major relocation (2b-impl-2, #750) - #920
Merged
Conversation
…refuse seq-major relocation (2b-impl-2, #750) 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). Batch-general after this increment (no relocation): - 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. - Head-major (BNSH) growing-bucket KV: the flat 0..bucket range and the copy/zero block loop are already batch-general (batch is outermost, block count = batch*kv_heads); made explicit + unit-tested at batch=3. - Seq-major *fixed full-context stride* commit: live_prefix_ranges, seq_major_kv_commit_requests and zero_seq_major_committed_tail now emit one dense run per sequence, each a fixed capacity*bytes_per_token apart. This is relocation-free (the per-sequence stride is the pinned full context), so it is batch-general with CPU unit tests at batch=3. Refused (named, tested): - Seq-major *growing-bucket / realloc* growth at batch>1: kv_growth_byte_layout now bails with a named error, because a growing bucket's per-sequence stride is the mutable bucket capacity, so growth would relocate every sequence b>0's KV. Head-major has no such problem (batch outermost). This refusal is the correct deliverable rather than a silently wrong capacity-dependent stride. Scoping corrections (file/line evidence in the PR): - live_prefix_ranges lives in crates/onnx-genai-engine (native_decode/kv_commit.rs), NOT "another crate" as SCOPING doc section 3 item 6 states. - The doc's "KV growth stride math is already batch-general" is true only for head-major; seq-major growing-bucket at batch>1 relocates and is refused here. Validation (all local, solo where GPU): - cargo fmt -p onnx-genai-engine -- --check - cargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings - 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) - MANDATORY mobius_seqmajor_growth_parity_native_cuda solo: 2/2 passed, seq-major kept graph across 3 growth events, kv_bytes_moved=0, growth_keeps=3, 48-token stream byte-identical to head-major (VMM in-place capture ON/OFF, realloc). - Hardware byte-identity solo (qwen14b-zp, 16 tokens, nvidia-smi verified empty): token IDs EXACT match reference; captures=2 fallbacks=0 invalidations=1; cuda_graph_measured captures=1; vmm_arena ref/byte underflows=0 unaccounted=0; oversubscribed=0. Throughput 7.32 tok/s within the 3.9-28 band (not a finding). Does not close #750 (two increments remain). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
Owner
Author
|
Validated independently before merging, and I caught one thing worth recording about my own process. The mobius gate first came back vacuously green for me. My invocation filtered the test out ( Hardware byte-identity, solo ( Gates: fmt clean, exact CI CUDA clippy gate clean, 96/96 On the substance. The distinction this PR draws is the valuable part, and it was missing from my scoping: a fixed full-context stride commit is batch-general in both layouts, because the batch axis is outermost and each sequence sits a constant The refusal is the right call and is written the way a refusal should be: it names the constraint, explains why head-major escapes it, and points the reader at the path that is relocation-free rather than just failing. A silently capacity-dependent stride would have passed every batch-1 test we have and corrupted the first batch-N run. Also accepted: Merging. 2b-impl-3 (device argmax) next. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #920 +/- ##
==========================================
- Coverage 78.03% 77.79% -0.25%
==========================================
Files 353 354 +1
Lines 143320 148042 +4722
Branches 143320 148042 +4722
==========================================
+ Hits 111842 115162 +3320
- Misses 27036 28387 +1351
- Partials 4442 4493 +51
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
justinchuby
added a commit
that referenced
this pull request
Aug 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-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-fixedbatch=1through the shape/IO layer) by making the KV growth and commit geometry compute overbatchsequences instead of assuming a single sequence at batch index 0.batchstays 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_bytes/full_mask_bytes),grow_vmm_mask_in_placememset,vmm_growth_requestsmask island, and the reallocbuild_grown_maskmemset now all spanbatchrows (they zeroed/committed one row before).0..bucketKV range (checked_shape_bytesalready includes the batch axis) and thecopy_kv_prefix*/zero_kv_suffixblock 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).live_prefix_ranges(now takesbatch),seq_major_kv_commit_requests, andzero_seq_major_committed_tailemit one dense run per sequence, each a fixedcapacity × bytes_per_tokenapart. 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:
batch > 1.kv_growth_byte_layoutnowbail!s with a named error. A growing bucket's per-sequence stride is the mutable bucket capacity, so growing the bucket would relocate every sequenceb > 0's already-written KV — defeating the seq-major fixed-stride "no KV moves / keep the captured graph" contract and themoved_bytes_per_token == 0accounting. 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-smiverified empty)profile_native --model qwen14b-zp --ep cuda --tokens 16 --steady: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-smiverified empty: 2/2 passed. Seq-major kept the captured graph across 3 growth events withkv_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 -- --checkcargo clippy -p onnx-runtime-ep-cuda --features cuda -- -D warnings(exact CI gate)cargo clippy -p onnx-genai-engine --features native-backend,cuda -- -D warningscargo 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)
live_prefix_rangesis NOT "another crate." SCOPING §3 item 6 and §4 say the seq-majorlive_prefix_rangespath "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 isseq_major_kv_commit_requests(cuda.rs). No cross-crate signature change was needed; the driver-level residency measurement inonnx-runtime-cuda-memoryis a mirror, not a caller, so it was unaffected (its clippy gate still passes).batch > 1the block loop would silently relocate sequencesb > 0(capacity-dependent block stride) and break themoved_bytes_per_token == 0/ "no KV moves" invariant. This PR turns that silent misbehavior into an explicit refusal.Judge-it-yourself
Merge
origin/main, then:cargo test -p onnx-genai-engine --features native-backend,cuda --lib native_decodereproduces 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_nreproduces the refusal proof.