Skip to content

onnx EP: per-axis QDQ scales (s_x[m], s_w[n]) in the fused int8 GEMM (TQDQ task 2) - #690

Merged
bong-water-water-bong merged 1 commit into
mainfrom
1bit/onnx-ep-qdq-per-axis
Oct 31, 2026
Merged

bong-water-water-bong merged 1 commit into
mainfrom
1bit/onnx-ep-qdq-per-axis

Conversation

@bong-water-water-bong

@bong-water-water-bong bong-water-water-bong commented Oct 31, 2026 •

Copy link
Copy Markdown
Collaborator

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 (DequantizeLinear axis 1 on a [K, N] weight). This is
    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, 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 the
kernel's arithmetic is unchanged: gemm_q8 gained pointer arguments beside the scalars, and the
column half of the scale is pre-folded with 1/s_y once per output element (the per-tensor path keeps
its 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 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 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_ReleaseAllocator now returns the slot and new_allocator
reuses 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)

case result
per-tensor scales worst |d| = 0 LSB vs the CPU EP
per-axis (sx=8 sw=6 in the EP trace) worst |d| = 0 LSB
per-axis along K refused: claimed=1, qdq_groups=0, per-node path, worst |d| = 0
peak RSS, 256×1024×1024 CPU EP ~7 MB vs fused 264 kB; delta 6356/6140/6680/6688/6140 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 still exact

The 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

  • Asymmetric zero points (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_w belong in the epilogue with column sums precomputed at weight
    load.
  • The device leg — the fused kernel's GEMM is still the host int8 gemm_q8; feeding int8 straight
    into the device BO needs an NPU-runtime entry (Am from int8, float-free packB) plus a window.
    Still blocked on pi-ortep's answer in ~/briefs/TORTEP-STATE.md.
  • A per-axis scale whose axis length is symbolic (an activation row count that is not static) is
    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 garbage zp (all 48 outputs off by up to 68 LSB, deterministically; optimized builds hid it). Fixed with CreateTensorAsOrtValue + memcpy in commit f3b17f7. Verified in a fresh no-optimization build: 5/5 tests pass, worst |d| = 0.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_y is folded once per output element and applied as (acc * s_x[m]) * wmul[n], which is the same three factors as the documented C_s32 · s_x[m] · s_w[n] / s_y with 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 bong-water-water-bong added the tprs-ready TPRS: review verdict READY posted by the standing reviewer; merge gate satisfied label Oct 31, 2026

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 (tensor 0/48, axis 0/48 fused, xaxis-K refused on the per-node path, rss 0/262144 with a 6220 kB peak delta).
  • But its base moved: 1bit/onnx-ep-qdq-fused is now 1f7d64605 (a main merge plus 4fa2fb3, "onnx EP proof: put initializer data in ORT-owned memory"), and GitHub now reports this PR CONFLICTING/DIRTY against 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's add_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).

@bong-water-water-bong bong-water-water-bong removed the tprs-ready TPRS: review verdict READY posted by the standing reviewer; merge gate satisfied label Oct 31, 2026
Base automatically changed from 1bit/onnx-ep-qdq-fused to main October 31, 2026 12:16
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.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 bong-water-water-bong added the tprs-ready TPRS: review verdict READY posted by the standing reviewer; merge gate satisfied label Oct 31, 2026
@bong-water-water-bong
bong-water-water-bong merged commit 3b950ac into main Oct 31, 2026
9 of 10 checks passed
@bong-water-water-bong
bong-water-water-bong deleted the 1bit/onnx-ep-qdq-per-axis branch October 31, 2026 12:24

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

tprs-ready TPRS: review verdict READY posted by the standing reviewer; merge gate satisfied

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant