[bong-water-water-bong]

The cached graph returned logits [1,seq,151936] for every prompt position while the runner
reads a single row (the final one, where the argmax comes from).

What changed

• export_qwen3.py --cached returns out.logits[:, -1:, :], so a 2048-token prompt stops
  materialising 1.2 GB of logits it never reads.
• ort_run.cpp reads the last row of whatever the graph returned (p + count - vocab) instead of
  indexing by seq - 1. The old expression was correct for a full [1,seq,V] output and read past
  the end of a one-row one — that was a real bug and it showed up as garbage token ids the moment
  the output got sliced.

What it bought: peak RSS 7.98 → 6.60 GB. Nothing else.

What it did not buy, and the measurement that says so: a second variant that runs the backbone
and then applies the head to the last hidden row (1 row instead of seq through the 151936-wide
matmul) is not faster — 2048 prefill 567 tok/s sliced vs 640 head-last vs 621 before, all within the
30–460 ms per-length stddev. The reason is arithmetic: a 2048-token prefill is
2 × 0.596e9 × 2048 = 2.4 TFLOP in ~3.2 s = ~760 GFLOP/s on eight cores, i.e. compute-bound at the
hardware, and the head is 155M of the model's 596M parameters. Both variants are in BENCHMARK-PARITY.md
(Round 3) with the policy decision recorded: the one-line slice is kept over the backbone-plus-head
version because it is equal in memory and does not depend on self.inner.model / self.inner.lm_head
attribute names.

Numerics unchanged: ids 13 16 13 16, mean KL 1.238e-07 vs transformers' own cached greedy loop,
max|dlogit| 2.12e-05, ctest 4/4.

Also in here: the NPU row's open question is answered — a 2048-token prompt does complete on our
dx route (2,678 tokens after the chat template, 54.76 tok/s prefill = 48.9 s, then 41.11 tok/s decode,
exit 0, device released under window GO-7). The earlier "hang" was the harness asking for six
consecutive 52-second repetitions.

[bong-water-water-bong]

VERDICT: READY (tprs-ready applied). I reproduced the bug and the fix on this box, not just the claim.

The off-by-end read is real, and the new expression fixes it. The comment says the old p + vocab*(seq-1) "read past the end of a one-row output" — so I built a model with exactly the shape your export now produces (the existing ~/wt/onnx-artifacts/qwen3-0.6b-cached.onnx with a Slice(logits[:, -1:, :]) inserted before the graph output, output name kept as logits), and ran main's ort_run and this branch's against it, same ids, 8 threads:

main  (p + vocab*(seq-1))      ids: 140518 39423 13 358          <- garbage, as you said
#686  (p + count - vocab)      ids: 264 11416 3110 315           <- matches the full-logits run

and on the un-sliced artifact both binaries agree exactly (264 11416 3110 315, prefill 54.8 vs 41.7 ms within noise), i.e. the change is a no-op for a graph that still returns [1,seq,V] and only matters for a sliced one. That is the whole risk of this PR and it is closed.

The arithmetic checks. 2 × 0.596e9 × 2048 = 2.44 TFLOP over ~3.2 s = 763 GFLOP/s, matching your ~760; the head is 151936 × 1024 = 155.6M of 596M parameters, matching "155M of 596M". The sliced logits a 2048-token prefill would stop materialising are 2048 × 151936 × 4 B = 1.245 GB, against a measured peak drop of 7.98 → 6.60 GB — same direction, same order, with the rest being allocator/arena effects.

What I could not reproduce: the 7.98 → 6.60 GB peak itself, because re-exporting needs torch's ONNX exporter and onnxscript is not installed on this box (ModuleNotFoundError: No module named 'onnxscript'). So that number is yours; the injection above is what I could build without it, and it is the part that matters for correctness.

Numerics unchanged is consistent with everything I ran (identical ids on the untouched artifact, and the KL figure you quote matches what docs/onnx.md already records). The negative result on the head-last variant is properly recorded in Round 3 with its per-length stddev, and the policy reason for keeping the one-liner (not depending on self.inner.model / self.inner.lm_head attribute names) is the right kind of reason to write down.

Two notes, neither blocking:

1. The doc's "Round 3" table is good, but the RSS column mixes two things — the slice's 1.4 GB and the head-last variant's equal memory. One clause saying which variants those peaks belong to would stop a reader attributing the whole drop to the slice.
2. The NPU datapoint you added (2,678 tokens after the template, 54.76 tok/s prefill = 48.9 s, 41.11 tok/s decode, exit 0, window GO-7) answers the open question in BENCHMARK-PARITY.md, but the earlier "hang" framing ("the harness asking for six consecutive 52-second repetitions") is only in the PR body. Worth one line in the doc, where the next person will look.

[bong-water-water-bong]

Re-verdict at the new head d9dfd40 (merge of main into the branch): READY, tprs-ready stays. This head is a merge commit, not a rebase, so the tree is different from the one I verified — what matters is that the PR's own content against main is unchanged, and it is:

$ git diff --stat origin/main d9dfd40
 onnx/BENCHMARK-PARITY.md | 102 ++++++++++++++++++++++++++++++++++++++++++++---
 onnx/export_qwen3.py     |   7 +++-
 onnx/ort_run.cpp         |   5 ++-
 3 files changed, 106 insertions(+), 8 deletions(-)
$ git log --oneline --no-merges origin/main..d9dfd40
7060b2e onnx: return only the final logits row, and measure what that buys

That is byte-for-byte the change I reviewed at 7060b2e1 — one non-merge commit, the same three files, the same +106/−8 — with main's later work merged in underneath (origin/main is an ancestor, so GitHub's MERGEABLE is the real thing). build pass (4m3s) at this head, Analyze x4 + CodeQL green, pr_agent red as always.

The two doc notes from my first review stay notes, not gates. Nothing else needed from the owner: this is mergeable as it stands.

One heads-up for the merge, not about this PR: #695 also edits onnx/ort_run.cpp (it adds --basic and rewires make_session). Whichever lands second will need the usual look at that file — #686's delta there is 5 lines, #695's is ~20, so I expect a trivial merge, but it is the one file they share.

[bong-water-water-bong]

VERDICT: READY

Gate line for the protocol, at head d9dfd400 (current). build pass, MERGEABLE, tprs-ready applied.

What that rests on, unchanged from my re-verdict above: this head is main merged in, not a rebase, so the tree moved — but the PR's own content against main is byte-identical to the head I reviewed (git diff --stat origin/main d9dfd400 -> onnx/BENCHMARK-PARITY.md, onnx/export_qwen3.py, onnx/ort_run.cpp, +106 −8; one non-merge commit, 7060b2e), and origin/main is an ancestor. The change itself — return only the final logits row and measure what that buys — was verified at 7060b2e1. The two doc notes I raised are notes, not gates.

[bong-water-water-bong]

VERDICT: READY

At head 3ac18e3 (post-update, main merged in), build pass, MERGEABLE, tprs-ready applied.

The merge added no author change, as expected, and I checked that rather than assuming it:

$ git log --oneline --no-merges origin/main..3ac18e3
7060b2e onnx: return only the final logits row, and measure what that buys
$ git diff --stat origin/main 3ac18e3
 3 files changed, 106 insertions(+), 8 deletions(-)      # BENCHMARK-PARITY.md, export_qwen3.py, ort_run.cpp
$ git merge-base --is-ancestor origin/main 3ac18e3      # yes

Same one non-merge commit, same three files, same +106/−8 as the head I reviewed — so the content verdict carries unchanged, and this line is the one that moves. Awaiting merge.

[bong-water-water-bong]

VERDICT: READY

At head 88055c4 (main merged in again after #699 landed), build pass, MERGEABLE, tprs-ready applied.

Third time this file has moved and nothing about it changed:

$ git log --oneline --no-merges origin/main..88055c4
7060b2e onnx: return only the final logits row, and measure what that buys
$ git diff --stat origin/main 88055c4
 3 files changed, 106 insertions(+), 8 deletions(-)
$ git merge-base --is-ancestor origin/main 88055c4      # yes

Same single non-merge commit, same three files, same +106/−8 as every head since 7060b2e1. Content verdict stands; this is the current gate line.

[bong-water-water-bong]

VERDICT: READY

At head 4cca180 (main merged in again after #695 landed; it is the last of the three still open), build pass, tprs-ready applied.

Same check, fourth head, nothing new to review:

$ git log --oneline --no-merges origin/main..4cca180
7060b2e onnx: return only the final logits row, and measure what that buys
$ git diff --stat origin/main 4cca180
 3 files changed, 106 insertions(+), 8 deletions(-)      # BENCHMARK-PARITY.md, export_qwen3.py, ort_run.cpp
$ git merge-base --is-ancestor origin/main 4cca180      # yes
$ parent of 4cca180: 88055c4                            # the head I gated at 14:58Z

Content verdict unchanged since 7060b2e1: one non-merge commit, the same three files, the same +106/−8. #699 and #695 have both landed around it, so this is now the only labelled PR that is waiting on a merge — nothing ahead of it, nothing shared left in onnx/ort_run.cpp except #688 (still CHANGES).