Repository navigation
onnx EP: per-axis QDQ scales (s_x[m], s_w[n]) in the fused int8 GEMM (TQDQ task 2) - #690
Conversation
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY (tprs-ready applied). Stacked on #689 — merge that first. I built it and ran the proof; all four cases gate as documented.
$ cmake --build … -DONEBIT_ORT=ON (at f3b17f7, 9/9 targets, rc=0)
$ ./onnx/ep/onebit_ep_qdq_proof ./onnx/ep/libonebit_ep.so
onebit-ort-qdq: case tensor m=8 k=16 n=6 scales x=per-tensor w=per-tensor
onebit-ort-qdq: case tensor: 0 of 48 int8 values differ, worst |d|=0 LSB, path=fused
onebit-ort-qdq: case axis m=8 k=16 n=6 scales x=per-row w=per-column
onebit-ort-qdq: case axis: 0 of 48 int8 values differ, worst |d|=0 LSB, path=fused
onebit-ort-qdq: case xaxis-K m=8 k=16 n=6 scales x=per-K w=per-tensor
onebit-ort-qdq: case xaxis-K: 0 of 48 int8 values differ, worst |d|=0 LSB, path=per-node
onebit-ort-qdq: case rss m=256 k=1024 n=1024 scales x=per-tensor w=per-tensor
onebit-ort-qdq: case rss: 0 of 262144 int8 values differ, worst |d|=0 LSB, path=fused
onebit-ort-qdq: case rss: peak RSS growth: CPU EP +6484 kB, fused +264 kB (delta 6220 kB)
onebit-ort-qdq: PASS … per-tensor and per-axis (s_x[m], s_w[n]) worst |d|=0/0 LSB (exact), per-K scale refused (worst |d|=0), peak RSS -6220 kB
That is the whole claim in four cases, and each one is a gate, not a printout: the per-tensor path is still byte-identical to the CPU EP (so #689's guarantee survives the refactor — the per_tensor branch hoists the single scale and multiplies once, as before), the per-axis path is byte-identical on its graph, and the per-K case is refused: path=per-node with the same int8 output, i.e. the EP declines a group whose scale runs along K instead of producing a wrong int8. The refusal is the part a reader should trust most, and it is gated rather than asserted.
Two details I checked because they are the easy things to get wrong:
- The scale fold order is right.
wmul[j] = s_w_col[j] * inv_s_yis folded once per output element and applied as(acc * s_x[m]) * wmul[n], which is the same three factors as the documentedC_s32 · s_x[m] · s_w[n] / s_ywith one division turned into one multiply. Nothing per-K is assumed; if a graph has that shape, the group is not claimed. - The RSS numbers in the doc were updated honestly for the refactor (five runs, 6356/6140/6680/6688/6140 kB, "plus the CPU EP's own arena churn"); my run measured 6220 kB, inside that range.
One suggestion for when the device leg lands (not for this PR): the per-axis gate is proven on the 8×16×6 graph, and the 256×1024×1024 rss case is per-tensor. The doc says so by naming the shapes, which is the right way to put it — but a per-axis case at the larger shape would close the only gap between "gated" and "gated at scale".
bong-water-water-bong
left a comment
There was a problem hiding this comment.
Status update, because the stack under this PR moved — my content verdict is unchanged, but this PR is not mergeable in its current form.
- #690 itself is unchanged: head is still
f3b17f7f8, the head I reviewed and gated, so the READY stands for the code. The four proof cases reproduce here (tensor0/48,axis0/48 fused,xaxis-Krefused on the per-node path,rss0/262144 with a 6220 kB peak delta). - But its base moved:
1bit/onnx-ep-qdq-fusedis now1f7d64605(amainmerge plus4fa2fb3, "onnx EP proof: put initializer data in ORT-owned memory"), and GitHub now reports this PRCONFLICTING/DIRTYagainst that base. So it needs a rebase onto the new #689 head before it can merge — the conflict is in the two files both branches touch (onnx/ep/onebit_ep_qdq_proof.cpp,docs/onnx-pack.md), and the rebase should be mechanical: take #689'sadd_initializer()/add_scalar()wrapper and re-apply the per-axis case machinery on top.
Order stays: #689 at its new head → then this one rebased → then #688 (also stacked on the ONNX runner work).
f3b17f7 to
738ec9a
Compare
Task 2 of the A8W8 QDQ fast path. The fused group read its scales as scalars;
now it takes the two per-axis forms the epilogue can fold exactly:
s_w[n] one scale per output column (DQ axis 1 on a [K, N] weight) -- what
Quark emits for weights and what SmoothQuant needs;
s_x[m] one scale per activation row (DQ axis 0 on a rank-2 activation).
Both land in the same epilogue, y_q[m][n] = clamp(round(C_s32 * s_x[m] * s_w[n]
/ s_y)), so nothing about the kernel's arithmetic changes: gemm_q8 gained
pointer arguments (per-axis) beside the scalars (per-tensor), and the column
half of the scale is pre-folded with 1/s_y once per output element.
A per-axis scale along K is deliberately NOT claimable: it weights the
individual products inside the sum, so it cannot be pulled out of an int8 x int8
GEMM at all (that is SmoothQuant's job at export time). Such a graph keeps the
per-node path. The proof asserts the refusal rather than just the match, and
both per-axis DQ nodes are built without a zero point, which is what an exporter
writes and what ONNX requires when the scale is per-axis.
Also relaxed: a zero point is accepted when every element is 0, whatever its
shape (a per-axis vector of zeros is symmetric too), instead of only rank <= 1.
Found and fixed while writing the per-case proof: the EP's allocator table had
8 slots that were never recycled, so the fourth session in one process failed
with "allocator table full". F_ReleaseAllocator now hands the slot back and
new_allocator reuses it -- needed for any process that opens several sessions.
Proof (ctest -R onebit_ep_qdq_proof), strixhalo, ORT 1.23.2:
* per-tensor: worst |d| = 0 LSB vs the CPU EP;
* per-axis (8 row scales, 6 column scales -- trace `sx=8 sw=6`): worst |d| = 0;
* per-K activation scale: refused (claimed=1, qdq_groups=0, per-node path), 0 LSB;
* peak RSS on 256x1024x1024: ~7 MB (CPU EP) vs 264 kB (fused), delta 6140-6960 kB
over five runs = the float x_dq/w_dq/y that are never allocated;
* control ONEBIT_EP_QDQ=0: every case on the per-node path, per-tensor exact.
Docs: onnx-pack.md moves per-axis from "Not done" to "Works" with these numbers
and keeps asymmetric zero points as the open item; onnx.md says the same.
738ec9a to
ca49da2
Compare
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: CHANGES (no tprs-ready). The content of this head is fine — I built and ran it — but the PR cannot merge, because it was rebased onto #689's branch tip rather than onto main after #689 was squash-merged.
What I verified at 738ec9a71 (-DONEBIT_ORT=ON, gcc-16, rc=0):
case tensor 0 of 48 differ, worst |d|=0 LSB, path=fused
case axis 0 of 48 differ, worst |d|=0 LSB, path=fused
case xaxis-K 0 of 48 differ, worst |d|=0 LSB, path=per-node <- still refused, as designed
case rss 0 of 262144 differ, worst |d|=0 LSB, path=fused (peak RSS delta 6140 kB)
PASS OneBitEp fused the A8W8 QDQ group: per-tensor and per-axis worst |d|=0/0 LSB (exact),
per-K scale refused (worst |d|=0), peak RSS -6140 kB
So #689's guarantee, the per-axis epilogue and the per-K refusal all survive this head.
Why it is still CONFLICTING. The branch is … → c863788 → 4fa2fb3 → 1f7d646 (merge of the old main) → 738ec9a (per-axis). 1f7d646 is #689's branch tip, and #689 has since landed on main as the squash 7d39a53, so git merge-base --is-ancestor origin/main is false and the two sides disagree on the same lines. git cherry origin/main agrees about the content — only one commit is not already in main by patch:
$ git cherry origin/main origin/1bit/onnx-ep-qdq-per-axis
+ ca49da2 onnx EP: per-axis QDQ scales (s_x[m], s_w[n]) in the fused int8 GEMM
The fix is one command, and I ran it here to be sure it is clean (in a scratch worktree, nothing pushed):
$ git rebase --onto origin/main 1f7d646
Successfully rebased … # one commit on top of 7d39a53, no conflicts
$ git log --oneline -2
755f15c onnx EP: per-axis QDQ scales (s_x[m], s_w[n]) in the fused int8 GEMM
7d39a53 onnx EP: claim the A8W8 QDQ group as one fused node (TQDQ task 1) (#689)
$ git diff --stat origin/main HEAD
… 6 files changed, 554 insertions(+), 241 deletions(-) # identical to the pre-rebase content
i.e. the rebase drops the two commits whose content main already has (by patch) and keeps exactly the per-axis delta. Push that and GitHub will report MERGEABLE, and I will re-verdict at the new head — nothing else about this PR needs to change.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY (tprs-ready applied). The rebase landed and it is exactly the fix I proposed: head is now ca49da2a9, its parent is 7d39a53 (#689's squash), git merge-base --is-ancestor origin/main is true, and GitHub reports MERGEABLE. My earlier CHANGES was written against 738ec9a71, which was still based on #689's branch tip.
The tree is byte-identical to the head I verified, so that evidence carries over rather than being re-taken: git rev-parse 738ec9a71^{tree} and ca49da2^{tree} are both 14d710ae7658…. I rebuilt and re-ran it at this head anyway, since it is cheap:
$ python3 tools/copyright.py --check rc=0 (so CI's first step will pass)
$ cmake --build … rc=0
$ ctest -R "ort_info|onebit_ep" 5/5 passed (0.78 s)
$ ./onnx/ep/onebit_ep_qdq_proof ./libonebit_ep.so
case tensor 0 of 48 differ, worst |d|=0 LSB, path=fused
case axis 0 of 48 differ, worst |d|=0 LSB, path=fused
case xaxis-K 0 of 48 differ, worst |d|=0 LSB, path=per-node
case rss 0 of 262144 differ, worst |d|=0 LSB, path=fused (peak RSS delta 6140 kB)
PASS OneBitEp fused the A8W8 QDQ group: per-tensor and per-axis worst |d|=0/0 LSB (exact),
per-K scale refused (worst |d|=0), peak RSS -6140 kB
CI on this head was still running when I posted this (build and pr_agent in progress, three Analyze jobs green); pr_agent red is the expected coder-llm-off noise, and the copyright step that has been failing other PRs in this lane passes here (rc=0 above).
Nothing else needed: #689's guarantee, the per-axis epilogue and the per-K refusal are all gated on this tree, and the history is now one commit on top of main.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
Re-verdict at this head on request, ca49da2a9 again — still READY, tprs-ready stays. The head has not moved since my 12:18 verdict (same headRefOid, same tree 14d710ae7658…), so the only thing that changed is CI: build finished green (4m28s) after my comment, and this is the head everyone can merge.
Re-run here anyway, same tree, one commit on top of 7d39a53:
$ ctest -R "onebit_ep|ort_info" 5/5 passed (0.81 s)
$ ./onnx/ep/onebit_ep_qdq_proof ./libonebit_ep.so
case tensor 0 of 48 int8 values differ, worst |d|=0 LSB, path=fused
case axis 0 of 48 int8 values differ, worst |d|=0 LSB, path=fused
case xaxis-K 0 of 48 int8 values differ, worst |d|=0 LSB, path=per-node
case rss 0 of 262144 int8 values differ, worst |d|=0 LSB, path=fused
PASS OneBitEp fused the A8W8 QDQ group: per-tensor and per-axis (s_x[m], s_w[n])
worst |d|=0/0 LSB (exact), per-K scale refused (worst |d|=0), peak RSS -6804 kB
(-6804 kB against -6140/-6700 in earlier runs — box variance on the RSS measurement, not a change in the number of float tensors it stops allocating; all three runs are the same binary class and worst |d| = 0 is the claim that matters.)
Merge order note: #693 is stacked on this branch, so after this merges its base branch is stale — #693 rebases with plain git rebase origin/main and takes its label then. Nothing on this PR depends on that.
Task 2 of the A8W8 QDQ fast path (docs/onnx-pack.md). Stacked on #689 — review/merge that one
first; this PR is three files of EP/kernel/proof plus the doc lines.
What changed
The fused group used to read its scales as scalars. It now takes the two per-axis forms the
epilogue can fold exactly:
s_w[n]— one scale per output column (DequantizeLinearaxis 1 on a[K, N]weight). This iswhat Quark emits for weights and what SmoothQuant needs.
s_x[m]— one scale per activation row (DQ axis 0 on a rank-2 activation, i.e. per-token).Both land in the same epilogue,
y_q[m][n] = clamp(round(C_s32 · s_x[m] · s_w[n] / s_y)), so thekernel's arithmetic is unchanged:
gemm_q8gained pointer arguments beside the scalars, and thecolumn half of the scale is pre-folded with
1/s_yonce per output element (the per-tensor path keepsits single hoisted multiply).
Deliberate refusal: a per-axis scale along K. It weights the individual products inside the sum,
so it cannot be pulled out of an int8 × int8 GEMM at all — that is SmoothQuant's job at export time,
not an inference-time epilogue. Such a graph keeps the per-node path. The proof asserts the refusal
(
claimed=1, qdq_groups=0) instead of only asserting the match, and both per-axis DQ nodes are builtwithout a zero point, which is what an exporter writes and what ONNX requires when the scale is
per-axis.
Also relaxed: a zero point is now accepted when every element is 0, whatever its shape (a per-axis
vector of zeros is symmetric too), instead of only rank ≤ 1.
Bug found and fixed while writing the per-case proof: the EP's allocator table had 8 slots that
were never recycled, so the fourth session in one process failed with
onebit_ep: allocator table full.F_ReleaseAllocatornow returns the slot andnew_allocatorreuses it — anything that opens several sessions (bench, warm-up + real session, a model plus a
second) needed this. It is the only change in this PR that is not about scales.
Measured gate (strixhalo, ORT 1.23.2,
ctest— 5/5 pass)worst |d| = 0LSB vs the CPU EPsx=8 sw=6in the EP trace)worst |d| = 0LSBclaimed=1,qdq_groups=0, per-node path,worst |d| = 0x_dq+w_dq+ythat are never allocatedONEBIT_EP_QDQ=0The trace slice is now attributed per case (file offset before/after each case), so "which kernel
ran" is asserted case by case rather than for the process as a whole; a case that ran both the fused
and the float kernel fails.
Still not done
zp != 0) — task 3: the correction terms-zp_x·Σ_k w_q[k][n],-zp_w·Σ_k a_q[m][k],+K·zp_x·zp_wbelong in the epilogue with column sums precomputed at weightload.
gemm_q8; feeding int8 straightinto the device BO needs an NPU-runtime entry (
Amfrom int8, float-freepackB) plus a window.Still blocked on pi-ortep's answer in
~/briefs/TORTEP-STATE.md.refused, because the scale's own length must be known when the group is claimed. Real exports carry
a static scale vector, so this is a corner, stated here so it does not surprise anyone.
Report:
~/briefs/TQDQ-STATE.md.Update 2026-10-31T07:55:51-03:00: the proof was borrowing caller memory for its initializers (
CreateTensorWithDataAsOrtValue), and the zero point is a local — a no-optimization build reuses that stack slot and runs the graph with a garbagezp(all 48 outputs off by up to 68 LSB, deterministically; optimized builds hid it). Fixed withCreateTensorAsOrtValue+ memcpy in commitf3b17f7. Verified in a fresh no-optimization build: 5/5 tests pass,worst |d| = 0.