Repository navigation
npu: quantise the activation row on an AVX2 path, byte-identical to the shipped loop - #908
Conversation
…he shipped loop (target 2)
Every NPU GEMM quantises one activation row to int8 on the host before its launch, in a scalar loop
built at the baseline x86-64 targets where roundf() is a libm call. On the container's own plan
(bonsai2-27b-t2s2-serve/plan.json: 288 projections, 2,310,144 A elements per decode token) that loop
costs 3.96-4.05 ms/token measured; the AVX2 path costs 0.55 ms/token (0.238 ns/element), 7.2-7.4x,
bytes identical on every projection of the mix.
The loop moves to npu/forward/include/quant_i8.h as a scalar reference plus an AVX2 path dispatched on
__builtin_cpu_supports("avx2"), with the reference used for the tail. Identity is by construction, not
by luck: trunc + an exactly-compared half adjustment is roundf's round-half-away-from-zero, and the
vector path reproduces the shipped loop's three sharp edges -- NaN/+-inf zeroed on the *input*, a finite
input whose product overflows left to saturate through the (int) cast (which x86 makes INT_MIN, so the
clamp yields -127, not 127), and |v*ais| just inside 2^31 taking the adjustment normally.
tests/quant_i8_row_test.cpp proves it with no device and no model: the .5 table, both clamps, NaN, +-inf,
fmax, +-2^31 and just inside it, denormals, -0.0, six real layer scales, lengths that leave a tail, and
6.5 M elements of random *bit patterns* per run (env-tunable rounds), where NaN payloads and denormals
turn up at their real density. It is load-bearing, not decorative: GE -> GT on the half comparison makes
it fail with round-half-to-even answers (67.5 -> 67 vs 68); 97.5 M elements of the same fuzz ran clean
against the shipped loop.
This is the other half of audit row 1: #873 removed 99.2% of the pad memset (+2.60 ms/token) and left
this loop alone.
PR Reviewer Guide 🔍Here are some key observations to aid the review process:
|
bong-water-water-bong
left a comment
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY
At 87dbe494f. The claim is "byte-identical", so I fuzzed the two arms against each other myself, in the engine's configuration (no -mavx2: row_avx2 carries target("avx2") and row() dispatches at run time, so the scalar arm is what a baseline x86-64 build compiles).
82,241 comparisons, 0 mismatches (-O2 and -O3, and again with no -mavx2)
over, in one run: named edges — NaN, ±inf, ±0.0, ±0.5 and nextafter either side of it, denormals, ±2^31 exact and nextafter(2^31, 0), ±3e38, ±2^24 — crossed with nine scales (1, 0.5, 127, 63.5, 1e-30, 1e30, +inf, -1, 2) and every length 1..33 so the tail after the last 8-lane block is covered; then every edge value placed alone in each of the eight lanes; then 40,000 random raw bit patterns (every exponent, every NaN payload) and 40,000 normal activations at a realistic scale. Zero mismatches anywhere.
The three cases the comment says are modelled rather than assumed, checked in the same run:
- NaN in, 0 out —
andwith the ordered mask zeroes the lane, matchingif (!isfinite(v)) v = 0. - A finite input whose product overflows — scalar
roundf(±inf)casts to INT_MIN and clamps to -127, and the vector path reproduces that (theinrangemask suppresses the rounding adjustment so the saturatedcvttpsvalue survives to the clamp). Both signs. - |v| ≥ 2^31 — same suppression, so it saturates rather than being "fixed" to 127. The boundary is exact because at that magnitude the float spacing is 256, so there is no representable value between
2^31 - 256and2^31.
Their own test passes here too: quant_i8_row: PASS (scalar + avx2 arms), 200,000 rounds / 6,535,086 elements / 0 failures. Worth saying that the test re-spells the scalar reference instead of calling the header's copy — that is what stops an edit to the reference from moving both sides together, and it is the right shape.
The speed claim, measured here (rows of 5120 f32, the shipping element count of 2,310,144/token, best of 7, pinned core, under thermal-run):
scalar 1.539 ns/element -> 3.56 ms/token
avx2 0.235 ns/element -> 0.54 ms/token
speedup 6.56x
So the AVX2 side matches the body (0.235 vs 0.238 ns/element, 0.54 vs 0.55 ms/token) but I get 6.6x, not 7.2-7.4x — the scalar figure in the comment (1.72-1.76 ns/element) is about 12 % slower than this box measures today. Not a gate; just make the number in the comment one that a re-run reproduces, so the next person does not go looking for a regression.
Checked by reading, for the wiring: both call sites (quantize_async, quantize_async_rows) now call q4nx_quant_i8::row(...) with the memset(Am, 0, MD*KD) still in front and only the first ak columns written, so the padded columns keep the zero behaviour they had; the CMake target is bound to a target ci.yml already builds (the #635 "Not Run" trap), and a non-AVX2 host still compiles and takes the scalar loop.
No device work: this path is host-side, and the device effect (3.4 ms of a ~415 ms token) is the lane's to measure on the array.
bong-water-water-bong
added
the
tprs-ready
…duces (follow-up to #908) (#919) * npu: make the quantiser header's speed figure one that a re-run reproduces #908's reviewer measured the same mix on this box and got 6.56x (1.539 ns/element scalar, 0.235 AVX2) where the comment claimed 7.2-7.4x (1.72-1.76, 0.238). The AVX2 side agrees; the scalar figure was ~12 % optimistic, so a later re-run would look like a regression. The comment now carries the range both runs measured, labelled as two independent measurements on the same box and mix. Comment-only change: no code, no behaviour, no re-gate. * npu: the quantiser comment's saving range must include both runs' endpoints Copilot's review of #919: the comment said 6.6-7.4x is 'about 3.4-3.5 ms' saved, but the slower scalar run (3.56) minus the faster AVX2 run (0.54) is 3.02 ms, so the range excluded the second measurement. Now stated as the subtraction it is: 3.0-3.5 ms/token (3.56-0.54 to 4.05-0.55), i.e. 0.73-0.85 % of a ~415 ms token. Comment-only; no code, no behaviour, no re-gate. --------- Co-authored-by: agent <agent@local>
What
Every NPU GEMM quantises one activation row to int8 on the host before its launch, in the scalar loop in
I8Ctx::quantize_async/quantize_async_rows. The engine builds at the baseline x86-64 targets, whereroundf()is a libm call, so that loop is a large part of the host time insidenpu_gemm.This moves the row loop into
npu/forward/include/quant_i8.has a scalar reference plus an AVX2 path, dispatched at run time on__builtin_cpu_supports("avx2")with the reference used for the tail. Nothing else changes: the callers keep theiramax/scale handling and pass the same1.0f / scale.This is the second half of the audit row #873 started: #873 removed 99.2 % of the activation pad memset (measured +2.60 ms/token) and left this loop alone.
Measured
Workload is the shipping container's own plan —
bonsai2-27b-t2s2-serve/plan.json: 288 projections, 2,310,144 A elements per decode token (48 gdn layers × K = 5120/5120/6144/5120/17408 + 16 attention layers × 5120/5120/17408). Host replay of exactly that mix, best of 7, five runs:7.2 – 7.4×, bytes identical on every projection of the mix in all five runs.
In situ,
npu-hold, 3 prompts × 24 greedy tokens on two packs, one round each:npu_gemmtotalbonsai2-27b-t2s2-servebonsai2-27b-pb16iBoth in the direction and order of the host measurement. Caveat, stated rather than buried: this is a single round per arm, and the derived
[lh] gemm hostrow is not used here because the two pb16i arms ran at 53 °C and 63 °C and that row moved 31 ms while the top-level total moved 8.7.rounds=2alternates the arms so drift cancels; not run, because the identity gate is the gate and the array is shared.Gate: byte-identical output, at three strengths, on two packs
Tokens, every
logprobfloat, the top-5 floats per token, and the sha256 of the final hidden state — all identical to the same-commit base:t2s2-serve(what :18097 serves)pb16i0
amdxdnaerrors, 0 TDR, no canary short reads. Both arms are the same commit (a2879378), one with the header, one without — symbol check: 0 vs 1q4nx_quant_i8::row_avx2.Why it is byte-identical by construction, not by luck
roundfis round-half-away-from-zero, andso
trunc+ an exact half adjustment is the same integerroundfproduces. The vector path reproduces the shipped loop's three sharp edges, each checked against the shipped code rather than assumed:isfinite), so a large finite input whose product overflows is not zeroed;(int)cast, which x86 makesINT_MIN, so the clamp yields −127, not 127 — the vector path must not "fix" this, and theinrangemask exists precisely becauset + adjwould wrap toINT_MAXfor negative overflow;|v·ais|just inside 2³¹ takes the adjustment normally.On non-x86 hosts, or without AVX2, the scalar loop runs unchanged.
Tests
tests/quant_i8_row_test.cpp— no device, no model, 0.11 s in CI: the .5 table (including126.5 → 127), both clamps, NaN, ±inf,FLT_MAX, ±2³¹ and just inside, denormals,-0.0, six real layer scales, tail lengths, and 6.5 M elements of random bit patterns per run (env-tunable viaQUANT_I8_FUZZ_ROUNDS), where NaN payloads and denormals appear at their real density instead of only where a case list remembers to put them.The test is load-bearing, verified by mutation rather than asserted: flipping
_CMP_GE_OQ→_CMP_GT_OQon the half comparison makes it fail with round-half-to-even answers (in=67.5 … got 67 want 68). A 97.5 M-element run of the same fuzz against the shipped loop was clean.tools/copyright.py --checkclean;1bitbuilds under the thermal guard.