Repository navigation
feat(native-decode): make the device argmax / greedy token selection batch-aware (2b-impl-3, #750) - #924
Merged
Conversation
…batch-aware (2b-impl-3, #750) Select `batch` greedy tokens from a `[batch, 1, vocab]` logits tensor instead of one token from `[1, 1, vocab]`, threading a `batch` count through the whole device-argmax stack: the CUDA kernel launch, the `ExecutionProvider` trait, the `onnx-runtime-ep-cuda` provider, `DeviceIoBinding::device_argmax`, the result buffer allocation, and `read_greedy_result`. `batch` stays constructor-fixed to 1; this delivers the capability, not batch-N on (that is 2b-impl-4). Batch-1 is byte-identical by construction: the grid stays 1-deep (gridDim.y == 1 / gridDim.x == 1), the result header stays 2 words, and the scratch layout is unchanged, so the launch geometry and the 8-byte read-back are bit-for-bit the prior single-sequence path. The reduction operator is additionally order-invariant (max value / min index on ties), so tie-breaks cannot drift with the grid shape. New device test `device_argmax_batch_selects_per_row_max_with_deliberate_ties` puts the max at a different position per row, forces multi-partial reductions, and plants ties at an earlier index than the peak, asserting each batched row equals a standalone single-sequence argmax — the control that a shared reduction buffer or a wrong per-sequence stride would fail. 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
… capability negotiation (#926) Two corrections, the first prompted by a simple observation: **Linux has no WDDM.** ## The #912 closure was written as though it settled the hybrid everywhere It settles it on **one consumer GPU under WDDM**. On Linux neither half of the argument survives: - **The comparison arm does not exist.** There is no OS shared-memory fallback, so an over-budget model that does not stream fails outright. The hybrid competes with managed streaming (0.11–0.86 tok/s on this box) or with *not running* — not with WDDM's 7.84 tok/s. A result that loses badly on Windows can be the best available option on Linux. - **The ceiling that closed it is plausibly VidMm's, not the hardware's.** #863 already measured WDDM demoting our own `cuMemCreate`/`cuMemMap` granules invisibly; a host registration being silently remapped is the same family of behaviour. On Linux `cuMemHostRegister` pins pages in the driver with nothing above it. That second point is a **hypothesis**, so it is filed as #925 to be measured rather than asserted in either direction. This is #783's lesson turned on our own result. Worth noting: **the code was already right.** `shared_memory_weight_fallback` is `cfg!(windows)`, Linux still auto-enables managed streaming, unit-locked by `managed_default_no_flag_over_budget_model_auto_streams`. This is documentation catching up to implementation, which is the safer direction to be wrong in. The decision table now splits the single "defer to the platform" row **into two**, conditioned on whether a demand-paging path exists at all — because those two rows are the same workload on two operating systems with opposite correct answers. ## That split motivates the second change The runtime chooses between those rows using `cfg!(windows)`: a compile-time proxy for "is this WDDM" that is wrong for TCC-mode Windows and silent about any third EP. Serviceable as an interim, indefensible as a design. New section: **Platform and EP capability negotiation.** Every row traced to the measurement that established the need: | Reported fact | Evidence | Status today | |---|---|---| | Oversubscription behaviour (demand-page / fail / unified) | #864/#874 vs #783 | `cfg!(windows)` | | Ledger truthfulness | #863 | absent; assumed true | | Aggregate host-mapped read capacity | #912; unmeasured on Linux (#925) | absent; assumed unbounded, which corrupted output | | VMM capability and granule | 2 MiB CUDA granule | known to the allocator, not reported | | Resident vs host bandwidth for the access pattern | #877 → #880, a 2× error in the sizing input | absent | | KV layout as a stride descriptor | #783 | a runtime-side enum | | Reclaim capability | C5 | absent | The point that organises it: **these are quantities, not flags.** A boolean "supports host mapping" would not have prevented #912, whose entire finding was about *how much*. Today's channel is `ExecutionProviderCapabilities` — an opaque string set carrying exactly one entry, `"nxrt"`. ## Do we need new EP API? Three different answers - **Five reporting facts** want **one** versioned, `struct_size`-prefixed record rather than seven entry points, negotiated the way the plugin ABI already negotiates major/minor. A later capability then costs a field, not an ABI break. - **KV layout is a negotiation, not a report** — propose/accept/counter over a stride descriptor (#783) — because the runtime must still bind exactly one physical form per binding, so an EP that merely announces a preference leaves disagreements unresolvable. - **Reclaim is behavioural** — a `reclaim(target_bytes) -> released_bytes` callback, so C5 pressure can reach an EP-side holder, with the holder choosing victims (I6) and free to release zero. Plus four rules that keep this from fragmenting, of which the load-bearing one is: **an unreported capability degrades to its most conservative reading, never its most convenient one.** Unknown host-mapped capacity is zero, not unbounded. #912 is the cautionary case — the optimistic default produced silently wrong tokens rather than an error. ## A citation of my own, corrected Three documents referenced `#924` as the Linux issue **before I had created it**, and the number was taken by an unrelated PR in the meantime. Filed for real as #925 and the references repointed. Citing a number before verifying it exists is exactly what the measurement-discipline skill tells everyone else not to do. Docs only. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
Owner
Author
|
Validated independently before merging. The tie-breaking claim is the one that mattered here, and I verified it rather than accepting it. The new device test is the right control: per-row maxima at different positions plus deliberate earlier-index ties, asserting each batched row equals a standalone single-sequence argmax. A wrong sequence stride or a shared scratch region shows up there and cannot show up at batch 1. Gates, all re-run by me:
On the scoping observation. Reporting that Merging. 2b-impl-4 is the last one, and the only one that changes behaviour: pin the batch symbol, flip batch-N on, and report measured recapture-per-shape plus the #866 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #924 +/- ##
==========================================
+ Coverage 78.12% 78.43% +0.30%
==========================================
Files 353 354 +1
Lines 143317 148049 +4732
Branches 143317 148049 +4732
==========================================
+ Hits 111966 116116 +4150
- Misses 26908 27436 +528
- Partials 4443 4497 +54
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.
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.
batchstays 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
batchcount through the whole device-argmax stack:onnx-runtime-ep-cudakernels::device_argmax): partials grid gains a sequence row (blockIdx.y), finalize gains one block per sequence (blockIdx.x), result header is2 × batchwords (a[token_id, capture_error]pair per sequence), scratch isbatch × 2 × partial_count;device_argmax_scratch_words(elements, batch)andscratch_words(elements, batch);ExecutionProvider::device_argmaxtrait + the CUDA impl;DeviceIoBinding::device_argmax(elements, batch, result);__native_greedy_argmaxbuffer allocation (now2·batch + scratch);read_greedy_resultnow delegates to a newread_greedy_result_batchreturningVec<(TokenId, u32)>of lengthbatch.What is now batch-general, and what is not
Batch-general: the argmax kernel selects
batchtokens from[batch, 1, vocab], one grid row per sequence, writingbatch[token_id, capture_error]pairs.read_greedy_result_batchreads 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 touchwrite_decode_inputs.Is batch-1 tie-breaking provably unchanged? Yes — two independent reasons
partial_countis per-sequence and unchanged; atbatch == 1the 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.argmax_updatekeeps 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-smiverified empty before the run (0 MiB), greedy (sampling: OFF):Deterministic counters all green (
captures=2,fallbacks=0, underflows0,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 hadMANAGED_WEIGHT_STREAMING/CUDA_GRAPHenv unset, so theweight_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
(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 passeddevice_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_tiesputs 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 == 0per 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_resultcallsbinding.device_argmax(vocab, …)returning one token, and the kernel +device_argmax_scratch_wordslive inonnx-runtime-ep-cuda. No correction to the split. One observation worth recording (not a re-split): the doc listswrite_decode_inputsunder 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.