Repository navigation
fix python/c++ parity, runner throughput, and int8 kernel overflow - #5
Merged
guirguispierre merged 8 commits intoAug 11, 2026
Merged
Conversation
guirguispierre
marked this pull request as draft
August 11, 2026 04:46
guirguispierre
marked this pull request as ready for review
August 11, 2026 04:46
guirguispierre
marked this pull request as draft
August 11, 2026 04:52
guirguispierre
marked this pull request as ready for review
August 11, 2026 04:52
There was a problem hiding this comment.
Auto-approved: pull request authored by repository owner @guirguispierre.
BitFeedForward used the default exact-erf nn.GELU() while both C++ paths implement the tanh approximation. The ~4e-4 per-layer gap compounds and forks greedy decoding after six tokens; swapping the activation takes a depth-2 model from 6/12 to 12/12 identical tokens.
vmulq_s8 truncates to int8, so an activation of -128 against a weight of -1 yields +128, wraps to -128, and flips the sign of the accumulator: the old kernel returned -8192 where numpy gives +8192. vmull_s8 widens to int16 instead. The model's quantizer clamps to [-127,127], but ternary_matmul is a public entry point that accepts any int8.
bit_linear looped j-outer over row-major (in_dim, out_dim) weights, striding by out_dim and missing cache on every read, with an unpredictable branch per element. Streaming each weight row contiguously with a branchless multiply moves the stories config from 11.9 to ~170-185 tps. Integer addition is associative, so the accumulator totals are unchanged; the stage-by-stage trace against pytorch is identical. Also drops the leftover [BitLinear Debug] prints, which corrupted the generation stream including the new --stream output, and restores --parity, whose dumps were all commented out and called a print_tensor that no longer existed.
BitAttention has returned (output, cache) since kv-caching landed, so x + attn_out raised TypeError. CI ran this step with continue-on-error, which hid it.
Covers what tests/test_kernel_parity.py does not: that an exported model run through the standalone engine agrees with pytorch. The activation contract is asserted directly rather than through a tolerance. Erf-vs-tanh gelu moves the logits by the same ~1e-2 that a single quantizer boundary flip does, and x86 and arm already differ by that much, so no numeric bound separates the two. Greedy trajectories are not asserted at all, since they fork for good once any step ties.
Without embedded/runner the parity module skips silently, so the claim went unverified. The pull_request trigger was also filtered to master, which meant a pr stacked on another pr's branch got no run at all.
Whole-model bit-exactness is not achievable: the per-tensor int8 quantizer rounds, so one ulp of float difference can push a value across a .5 boundary and move an accumulator a whole step. Replaced with the two guarantees that hold, and documented why. Test count was 67, now 149. pocket_4k was listed at ~100K parameters and is 5.28M; stories_base at ~1.33M and is 30.5M, since its yaml ships the 50257 gpt-2 vocabulary while the quick start passes 4096. The 5.3 MB to 2.0 MB model matches no shipped config. The quick start also failed on a clean clone without make -C atomic_1bit/core. Two API_REFERENCE notes about unstored weight scales were made false by the loader fix.
Per AGENTS.md, commits carry no agent attribution.
guirguispierre
force-pushed
the
fix/parity-perf-hardening
branch
from
August 11, 2026 04:53
bfacf57 to
e493656
Compare
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.
Stacked on #4, which correctly fixes the atomic_lib.h loader but leaves three things broken underneath it. BitFeedForward used exact-erf gelu while both C++ paths implement the tanh form, which forks greedy decoding after six tokens; bit_linear strided across row-major weights with a branch per element and ran at 11.9 tps rather than the ~170 the readme claims; and the neon kernel's vmulq_s8 wrapped on an activation of -128, returning -8192 where numpy gives +8192. It also drops the leftover debug prints that were corrupting #4's own new --stream output on first use, restores --parity (every dump was commented out and called a function that no longer existed), and fixes tools/parity_check.py, which had been crashing behind continue-on-error since kv-caching changed BitAttention's return type.
The readme corrections are the part worth checking, because they contradict published numbers: whole-model bit-exactness is not achievable by design and I replaced that claim rather than trying to reach it, pocket_4k is 5.28M parameters rather than the ~100K listed, stories_base is 30.5M rather than ~1.33M, and the quick start's first command failed on a clean clone. Each of those is measured, and the reasoning for the bit-exactness claim is in the readme.
Be suspicious of the parity tests. I originally asserted a tight numeric bound and CI caught it: x86 and arm disagree by the same ~1e-2 that a wrong activation produces, because both are one quantizer boundary flip, so no tolerance separates them. The gelu contract is asserted directly instead, the numeric bound is deliberately loose, and greedy trajectories are not asserted at all since they fork permanently once any step ties. Rebase this onto master if #4 is dropped, since both branches touch atomic_runner.cpp heavily. The runner still has no kv-cache, so throughput roughly halves by 100 tokens.