From c6a1e6a9a78c66785e9ba111f59396095824793c Mon Sep 17 00:00:00 2001 From: Pierre Guirguis <22091706+guirguispierre@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:52:40 -0700 Subject: [PATCH 1/8] use tanh gelu to match the c++ runtime 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. --- atomic_1bit/model/transformer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/atomic_1bit/model/transformer.py b/atomic_1bit/model/transformer.py index b321b50..7c48972 100644 --- a/atomic_1bit/model/transformer.py +++ b/atomic_1bit/model/transformer.py @@ -90,7 +90,9 @@ def __init__(self, config: AtomicConfig): hidden_dim = 4 * config.dim self.fc1 = BitLinear(config.dim, hidden_dim, bias=False) self.fc2 = BitLinear(hidden_dim, config.dim, bias=False) - self.act = nn.GELU() + # guirguispierre 2026-08-10 - must match the tanh gelu in + # embedded/atomic_runner.cpp or python and c++ output diverge + self.act = nn.GELU(approximate="tanh") def forward(self, x): return self.fc2(self.act(self.fc1(x))) From 12278f1749dcf5db2c8c455bca93300e8586c66c Mon Sep 17 00:00:00 2001 From: Pierre Guirguis <22091706+guirguispierre@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:52:40 -0700 Subject: [PATCH 2/8] fix int8 overflow in neon ternary kernel 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. --- atomic_1bit/core/backends/cpu_kernel.cpp | 27 +++++++++++++----------- tests/test_kernel_parity.py | 23 ++++++++++++++++++++ 2 files changed, 38 insertions(+), 12 deletions(-) diff --git a/atomic_1bit/core/backends/cpu_kernel.cpp b/atomic_1bit/core/backends/cpu_kernel.cpp index 2e4d37b..a54a894 100644 --- a/atomic_1bit/core/backends/cpu_kernel.cpp +++ b/atomic_1bit/core/backends/cpu_kernel.cpp @@ -36,33 +36,36 @@ void ternary_matmul(const int8_t *A, const int8_t *B_transposed, int32_t *C, int32x4_t v_acc0 = vdupq_n_s32(0); int32x4_t v_acc1 = vdupq_n_s32(0); + // guirguispierre 2026-08-10 - vmull_s8 widens to int16 because vmulq_s8 + // truncates: activation -128 against weight -1 wraps and flips the sum + // Process 32 items at a time (2x unrolled from 16) for (; k <= K - 32; k += 32) { // First 16 elements int8x16_t va0 = vld1q_s8(val_a + k); int8x16_t vb0 = vld1q_s8(val_b + k); - int8x16_t vprod0 = vmulq_s8(va0, vb0); - int16x8_t v_pair0 = vpaddlq_s8(vprod0); - int32x4_t v_quad0 = vpaddlq_s16(v_pair0); - v_acc0 = vaddq_s32(v_acc0, v_quad0); + v_acc0 = vaddq_s32( + v_acc0, vpaddlq_s16(vmull_s8(vget_low_s8(va0), vget_low_s8(vb0)))); + v_acc0 = vaddq_s32( + v_acc0, vpaddlq_s16(vmull_s8(vget_high_s8(va0), vget_high_s8(vb0)))); // Second 16 elements int8x16_t va1 = vld1q_s8(val_a + k + 16); int8x16_t vb1 = vld1q_s8(val_b + k + 16); - int8x16_t vprod1 = vmulq_s8(va1, vb1); - int16x8_t v_pair1 = vpaddlq_s8(vprod1); - int32x4_t v_quad1 = vpaddlq_s16(v_pair1); - v_acc1 = vaddq_s32(v_acc1, v_quad1); + v_acc1 = vaddq_s32( + v_acc1, vpaddlq_s16(vmull_s8(vget_low_s8(va1), vget_low_s8(vb1)))); + v_acc1 = vaddq_s32( + v_acc1, vpaddlq_s16(vmull_s8(vget_high_s8(va1), vget_high_s8(vb1)))); } // Handle remaining 16-element block for (; k <= K - 16; k += 16) { int8x16_t va = vld1q_s8(val_a + k); int8x16_t vb = vld1q_s8(val_b + k); - int8x16_t vprod = vmulq_s8(va, vb); - int16x8_t v_pair = vpaddlq_s8(vprod); - int32x4_t v_quad = vpaddlq_s16(v_pair); - v_acc0 = vaddq_s32(v_acc0, v_quad); + v_acc0 = vaddq_s32( + v_acc0, vpaddlq_s16(vmull_s8(vget_low_s8(va), vget_low_s8(vb)))); + v_acc0 = vaddq_s32( + v_acc0, vpaddlq_s16(vmull_s8(vget_high_s8(va), vget_high_s8(vb)))); } // Combine accumulators diff --git a/tests/test_kernel_parity.py b/tests/test_kernel_parity.py index 98910f1..f2e36ef 100644 --- a/tests/test_kernel_parity.py +++ b/tests/test_kernel_parity.py @@ -86,3 +86,26 @@ def test_all_neg_ones_weights(self): C = self.ternary_matmul(A, B) expected = -A.astype(np.int32).sum(axis=1, keepdims=True) * np.ones((1, 4), dtype=np.int32) np.testing.assert_array_equal(C, expected) + + def test_int8_min_activation(self): + # guirguispierre 2026-08-10 - -128 against -1 is +128, which int8 cannot + # hold; K=64 keeps this on the vectorized path + A = np.full((1, 64), -128, dtype=np.int8) + B = np.full((64, 8), -1, dtype=np.int8) + C = self.ternary_matmul(A, B) + expected = A.astype(np.int32) @ B.astype(np.int32) + np.testing.assert_array_equal(C, expected) + assert C[0, 0] == 8192, "int8 overflow in the widening multiply" + + def test_full_int8_activation_range(self): + rng = np.random.default_rng(0) + for _ in range(40): + K = int(rng.integers(1, 300)) + N = int(rng.integers(1, 64)) + A = rng.integers(-128, 128, size=(2, K)).astype(np.int8) + B = rng.integers(-1, 2, size=(K, N)).astype(np.int8) + np.testing.assert_array_equal( + self.ternary_matmul(A, B), + A.astype(np.int32) @ B.astype(np.int32), + err_msg=f"mismatch at K={K}, N={N}", + ) From f4389b53cd69335a105711ff3db5b0a108d6b23c Mon Sep 17 00:00:00 2001 From: Pierre Guirguis <22091706+guirguispierre@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:52:40 -0700 Subject: [PATCH 3/8] fix runner numerics, speed, and debug output 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. --- embedded/atomic_lib.h | 49 ++++++----- embedded/atomic_runner.cpp | 173 +++++++++++++++++-------------------- 2 files changed, 109 insertions(+), 113 deletions(-) diff --git a/embedded/atomic_lib.h b/embedded/atomic_lib.h index e30c2a8..d3e9d31 100644 --- a/embedded/atomic_lib.h +++ b/embedded/atomic_lib.h @@ -48,14 +48,15 @@ inline void softmax(vector &x) { // Affine RMSNorm (used for the layer's ln1 / ln2 / ln_f). inline void rms_norm(const vector &x, const vector &w, vector &out, int dim) { - float ss = 0.0f; + // guirguispierre 2026-08-10 - double sum and a divide rather than a + // reciprocal multiply, to match atomic_runner.cpp and the python reference + double ss = 0.0; for (float f : x) - ss += f * f; - ss /= dim; - float rms = 1.0f / sqrt(ss + 1e-5f); + ss += (double)f * (double)f; + float rms = sqrt((float)(ss / dim) + 1e-5f); for (int i = 0; i < dim; ++i) { - out[i] = x[i] * rms * w[i]; + out[i] = (x[i] / rms) * w[i]; } } @@ -79,10 +80,11 @@ inline void bit_linear(const vector &x, const vector &w, int out_dim = (int)out.size(); // SubLN - float ss = 0.0f; + // guirguispierre 2026-08-10 - double for the same reason as rms_norm + double ss = 0.0; for (float v : x) - ss += v * v; - float rms = sqrt(ss / in_dim + 1e-5f); + ss += (double)v * (double)v; + float rms = sqrt((float)(ss / in_dim) + 1e-5f); // Activation absmax over x/rms = max(|x|)/rms float max_abs = 0.0f; @@ -93,11 +95,12 @@ inline void bit_linear(const vector &x, const vector &w, max_val = 1e-5f; float scale_x = 127.0f / max_val; - float effective_scale = scale_x / rms; + // guirguispierre 2026-08-10 - divide before scaling to match + // activation_quant(); folding it into one multiply rounds differently vector x_q(in_dim); for (int i = 0; i < in_dim; ++i) { - float q = roundf(x[i] * effective_scale); + float q = roundf((x[i] / rms) * scale_x); if (q > 127.0f) q = 127.0f; if (q < -127.0f) @@ -105,18 +108,22 @@ inline void bit_linear(const vector &x, const vector &w, x_q[i] = (int8_t)q; } - float dequant = scale_w / scale_x; - for (int o = 0; o < out_dim; ++o) { - int32_t acc = 0; - for (int i = 0; i < in_dim; ++i) { - int8_t wv = w[i * out_dim + o]; - if (wv == 1) - acc += x_q[i]; - else if (wv == -1) - acc -= x_q[i]; - } - out[o] = (float)acc * dequant; + // guirguispierre 2026-08-10 - i-outer streams each weight row contiguously + // instead of striding by out_dim, worth ~14x; integer sums are associative so + // the totals are unchanged + vector acc(out_dim, 0); + for (int i = 0; i < in_dim; ++i) { + int32_t xv = x_q[i]; + if (xv == 0) + continue; // free sparsity: skip the whole row + const int8_t *wrow = &w[(size_t)i * out_dim]; + for (int o = 0; o < out_dim; ++o) + acc[o] += wrow[o] * xv; } + + float dequant = scale_w / scale_x; + for (int o = 0; o < out_dim; ++o) + out[o] = (float)acc[o] * dequant; } struct AtomicLayer { diff --git a/embedded/atomic_runner.cpp b/embedded/atomic_runner.cpp index 15c0dc2..c4a68cf 100644 --- a/embedded/atomic_runner.cpp +++ b/embedded/atomic_runner.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -63,10 +64,12 @@ struct AtomicModel { // RMSNorm void rms_norm(const vector &x, const vector &weight, vector &out, int dim) { - float sum_sq = 0.0f; + // guirguispierre 2026-08-10 - float32 summation drifts from torch's pairwise + // sum enough to shift the int8 quant_scale downstream + double sum_sq = 0.0; for (float v : x) - sum_sq += v * v; - float rms = sqrt(sum_sq / dim + 1e-5f); + sum_sq += (double)v * (double)v; + float rms = sqrt((float)(sum_sq / dim) + 1e-5f); for (int i = 0; i < dim; ++i) { out[i] = (x[i] / rms) * weight[i]; } @@ -78,16 +81,17 @@ void rms_norm(const vector &x, const vector &weight, // Out: Float output [Out] // Scratch: [In] (Pre-allocated buffer to avoid malloc) void bit_linear(const vector &x, const TensorI8 &w, vector &out, - vector &scratch) { + vector &scratch, vector &acc32) { int in_dim = w.rows; int out_dim = w.cols; // 0. Internal RMSNorm (Non-affine) // Per BitNet paper/impl: BitLinear(x) = Layernorm(x) ... - float sum_sq = 0.0f; + // guirguispierre 2026-08-10 - double for the same reason as rms_norm + double sum_sq_d = 0.0; for (float v : x) - sum_sq += v * v; - float rms = sqrt(sum_sq / in_dim + 1e-5f); + sum_sq_d += (double)v * (double)v; + float rms = sqrt((float)(sum_sq_d / in_dim) + 1e-5f); // 1. Quantize Input // We effectively quantize (x / rms). @@ -100,14 +104,16 @@ void bit_linear(const vector &x, const TensorI8 &w, vector &out, max_val = 1e-5f; float quant_scale = 127.0f / max_val; - // effective_scale = quant_scale / rms - float effective_scale = quant_scale / rms; - vector x_q(in_dim); + // guirguispierre 2026-08-10 - divide before scaling to match + // activation_quant(); folding it into one multiply rounds differently + // guirguispierre 2026-08-10 - caller's buffer avoids a malloc per call in the + // hot loop on microcontrollers + if ((int)scratch.size() < in_dim) + scratch.resize(in_dim); + int8_t *x_q = scratch.data(); for (int i = 0; i < in_dim; ++i) { - // x_norm = x[i] / rms - // val = x_norm * quant_scale - float val = round(x[i] * effective_scale); + float val = round((x[i] / rms) * quant_scale); if (val > 127) val = 127; if (val < -127) @@ -116,42 +122,40 @@ void bit_linear(const vector &x, const TensorI8 &w, vector &out, } // 2. Ternary Kernel - fill(out.begin(), out.end(), 0.0f); - - for (int j = 0; j < out_dim; ++j) { - int32_t acc = 0; - for (int i = 0; i < in_dim; ++i) { - int8_t xv = x_q[i]; - int8_t wv = w.data[i * out_dim + j]; - - if (wv == 1) - acc += xv; - else if (wv == -1) - acc -= xv; - } - // 3. Dequantize - // out = (acc / x_scale) * w_scale - float val = ((float)acc / quant_scale) * w.scale; - out[j] = val; - - // Debug first few outputs of first layer - static int debug_count = 0; - if (debug_count < 5) { - // print diagnostics for first few calls - if (j == 0) { - cout << "[BitLinear Debug] Dim In: " << in_dim << " Out: " << out_dim - << endl; - cout << " Input Max: " << max_val << " Quant Scale: " << quant_scale - << " RMS: " << rms << endl; - cout << " Weight Scale: " << w.scale << endl; // Check if this is 0 - cout << " Acc: " << acc << " Dequant: " << val << endl; - cout << " Sample W: " << (int)w.data[0] << " " << (int)w.data[1] - << endl; - cout << " Sample X_Q: " << (int)x_q[0] << " " << (int)x_q[1] << endl; - debug_count++; - } - } + // guirguispierre 2026-08-10 - i-outer streams each weight row contiguously + // instead of striding by out_dim, worth ~14x; integer sums are associative so + // the totals are unchanged + if ((int)acc32.size() < out_dim) + acc32.resize(out_dim); + int32_t *acc = acc32.data(); + memset(acc, 0, out_dim * sizeof(int32_t)); + + for (int i = 0; i < in_dim; ++i) { + int32_t xv = x_q[i]; + if (xv == 0) + continue; // free sparsity: skip the whole row + const int8_t *wrow = &w.data[(size_t)i * out_dim]; + // guirguispierre 2026-08-10 - w is {-1,0,1} so the multiply is the add/sub + // and vectorizes; the branchy form stalls on mispredicts + for (int j = 0; j < out_dim; ++j) + acc[j] += wrow[j] * xv; } + + // 3. Dequantize + // out = (acc / x_scale) * w_scale + for (int j = 0; j < out_dim; ++j) + out[j] = ((float)acc[j] / quant_scale) * w.scale; +} + +// guirguispierre 2026-08-10 - full float32 precision so tools/parity_check.py +// can diff values exactly rather than by eye +void print_tensor(const char *name, const vector &v, int count = 10) { + int n = min((int)v.size(), count); + cout << "[" << name << "] n=" << v.size() << " "; + cout << std::setprecision(9); + for (int i = 0; i < n; ++i) + cout << v[i] << (i + 1 < n ? " " : ""); + cout << endl; } // Softmax @@ -274,6 +278,7 @@ struct Workspace { vector fc_act, mlp_out; vector xt, xn; vector scratch_q; + vector acc32; void ensure_size(int max_seq_len, int dim) { int max_el = max_seq_len * dim; @@ -339,13 +344,10 @@ void forward(AtomicModel &model, const vector &tokens, } // Print Embedding Sum (last token) - // if (parity_mode) { - // int last_idx = (seq_len - 1) * dim; - // // Define a vector to copy so we can use print_tensor - // vector tmp(dim); - // memcpy(tmp.data(), &ws.X[last_idx], dim * sizeof(float)); - // print_tensor("Embedding_Sum", tmp); - // } + if (parity_mode) { + vector tmp(&ws.X[(seq_len - 1) * dim], &ws.X[seq_len * dim]); + print_tensor("Embedding_Sum", tmp, dim); + } int l_idx = 0; for (auto &layer : model.layers) { @@ -356,22 +358,19 @@ void forward(AtomicModel &model, const vector &tokens, memcpy(&ws.X_norm[t * dim], ws.xn.data(), dim * sizeof(float)); } - // if (parity_mode) { - // int last_idx = (seq_len - 1) * dim; - // vector tmp(dim); - // memcpy(tmp.data(), &ws.X_norm[last_idx], dim * sizeof(float)); - // // Construct string layer0_LN1 - // string name = "Layer" + to_string(l_idx) + "_LN1"; - // print_tensor(name.c_str(), tmp); - // } + if (parity_mode) { + vector tmp(&ws.X_norm[(seq_len - 1) * dim], + &ws.X_norm[seq_len * dim]); + print_tensor(("Layer" + to_string(l_idx) + "_LN1").c_str(), tmp, dim); + } // Attention for (int t = 0; t < seq_len; ++t) { memcpy(ws.xt.data(), &ws.X_norm[t * dim], dim * sizeof(float)); - bit_linear(ws.xt, layer.q_w, ws.q, ws.scratch_q); - bit_linear(ws.xt, layer.k_w, ws.k, ws.scratch_q); - bit_linear(ws.xt, layer.v_w, ws.v, ws.scratch_q); + bit_linear(ws.xt, layer.q_w, ws.q, ws.scratch_q, ws.acc32); + bit_linear(ws.xt, layer.k_w, ws.k, ws.scratch_q, ws.acc32); + bit_linear(ws.xt, layer.v_w, ws.v, ws.scratch_q, ws.acc32); memcpy(&ws.Q[t * dim], ws.q.data(), dim * sizeof(float)); memcpy(&ws.K[t * dim], ws.k.data(), dim * sizeof(float)); @@ -410,24 +409,18 @@ void forward(AtomicModel &model, const vector &tokens, // O Proj for (int t = 0; t < seq_len; ++t) { memcpy(ws.xt.data(), &ws.AttnOut[t * dim], dim * sizeof(float)); - bit_linear(ws.xt, layer.o_w, ws.layer_out, ws.scratch_q); + bit_linear(ws.xt, layer.o_w, ws.layer_out, ws.scratch_q, ws.acc32); // Residual for (int i = 0; i < dim; ++i) ws.X[t * dim + i] += ws.layer_out[i]; } - // if (parity_mode) { - // // We want to print Attn Output (Residual Added? Python script prints - // // `attn_out` BEFORE residual add) Wait, python script printed - // `attn_out = - // // layer.attn(x_norm)`. Here `ws.layer_out` is `attn_out`. But we just - // // added it to X. We have it in `ws.layer_out`. Assuming seq_len=1 for - // // parity check usually. If seq_len > 1, this only captures the LAST - // // token's execution in that loop. - // string name = "Layer" + to_string(l_idx) + "_Attn_Out"; - // print_tensor(name.c_str(), ws.layer_out); - // } + // guirguispierre 2026-08-10 - layer_out is pre-residual here, matching what + // tools/parity_check.py prints on the python side + if (parity_mode) + print_tensor(("Layer" + to_string(l_idx) + "_Attn_Out").c_str(), + ws.layer_out, dim); // MLP for (int t = 0; t < seq_len; ++t) { @@ -442,21 +435,19 @@ void forward(AtomicModel &model, const vector &tokens, // } int hidden = 4 * dim; - bit_linear(ws.xn, layer.fc1_w, ws.fc_act, ws.scratch_q); + bit_linear(ws.xn, layer.fc1_w, ws.fc_act, ws.scratch_q, ws.acc32); for (int i = 0; i < hidden; ++i) ws.fc_act[i] = gelu(ws.fc_act[i]); - bit_linear(ws.fc_act, layer.fc2_w, ws.mlp_out, ws.scratch_q); + bit_linear(ws.fc_act, layer.fc2_w, ws.mlp_out, ws.scratch_q, ws.acc32); for (int i = 0; i < dim; ++i) ws.X[t * dim + i] += ws.mlp_out[i]; } - // if (parity_mode) { - // string name = "Layer" + to_string(l_idx) + "_MLP_Out"; - // print_tensor(name.c_str(), ws.mlp_out); - // cout << endl; // Separator - // } + if (parity_mode) + print_tensor(("Layer" + to_string(l_idx) + "_MLP_Out").c_str(), + ws.mlp_out, dim); l_idx++; } @@ -466,15 +457,13 @@ void forward(AtomicModel &model, const vector &tokens, memcpy(ws.xt.data(), &ws.X[last_t * dim], dim * sizeof(float)); rms_norm(ws.xt, model.ln_f.data, ws.xn, dim); - // if (parity_mode) { - // print_tensor("Final_LN", ws.xn); - // } + if (parity_mode) + print_tensor("Final_LN", ws.xn, dim); - bit_linear(ws.xn, model.head_w, logits, ws.scratch_q); + bit_linear(ws.xn, model.head_w, logits, ws.scratch_q, ws.acc32); - // if (parity_mode) { - // print_tensor("Final_Logits", logits, 20); - // } + if (parity_mode) + print_tensor("Final_Logits", logits, (int)logits.size()); } // Sampling: temperature + optional top-k + optional top-p (nucleus). From 2a54243a68d7613979111b43f5b76cb99cf2c014 Mon Sep 17 00:00:00 2001 From: Pierre Guirguis <22091706+guirguispierre@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:52:40 -0700 Subject: [PATCH 4/8] fix parity check crash on kv-cache tuple 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. --- tools/parity_check.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/parity_check.py b/tools/parity_check.py index 9b816eb..6d9a628 100644 --- a/tools/parity_check.py +++ b/tools/parity_check.py @@ -98,7 +98,7 @@ def run_parity_check(): print(f" Manual Dequant: {dequant}") q_out = layer.attn.q_proj(x_norm) - attn_out = layer.attn(x_norm) + attn_out, _ = layer.attn(x_norm) print_tensor_stats(f"Layer{i}_Attn_Out", attn_out) x = x + attn_out From 7578fb68fbd90ce8cb13ec6aed3c2512043f787d Mon Sep 17 00:00:00 2001 From: Pierre Guirguis <22091706+guirguispierre@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:52:40 -0700 Subject: [PATCH 5/8] add python/c++ runtime parity tests 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. --- tests/test_cpp_runtime_parity.py | 143 +++++++++++++++++++++++++++++++ 1 file changed, 143 insertions(+) create mode 100644 tests/test_cpp_runtime_parity.py diff --git a/tests/test_cpp_runtime_parity.py b/tests/test_cpp_runtime_parity.py new file mode 100644 index 0000000..38be745 --- /dev/null +++ b/tests/test_cpp_runtime_parity.py @@ -0,0 +1,143 @@ +import os +import subprocess +import sys + +import pytest +import torch +import torch.nn as nn + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from atomic_1bit.model.transformer import AtomicTransformer, AtomicConfig +from atomic_1bit.utils.export_to_cpp import export_model + +RUNNER = os.path.join(os.path.dirname(__file__), "..", "embedded", "runner") + +pytestmark = pytest.mark.skipif( + not os.path.exists(RUNNER), + reason="C++ runner not built (cd embedded && g++ -O3 -std=c++17 " + "atomic_runner.cpp -o runner)", +) + +CFG = AtomicConfig(vocab_size=256, dim=64, depth=4, heads=4, context_length=32) + +# guirguispierre 2026-08-10 - loose because the int8 quantizer rounds, so one +# ULP of float difference can flip a value across a .5 boundary and move an +# accumulator a whole step; this catches gross mismatches, not single flips +AGREEMENT_TOL = 5e-2 + + +def _run_runner(model_bin, *args): + return subprocess.run( + [RUNNER, "--model", str(model_bin), *args], + capture_output=True, text=True, timeout=120, + ) + + +def _generated_tokens(res, n_prompt=1): + body = res.stdout.split("Generating:", 1)[1].split("Done.")[0] + fields = body.split() + for tok in fields: + assert tok.lstrip("-").isdigit(), ( + f"non-numeric token {tok!r} in generation stream: {body!r}" + ) + # guirguispierre 2026-08-10 - runner echoes the prompt before generating + return [int(t) for t in fields[n_prompt:]] + + +def _cpp_logits(model_bin, start_token): + res = _run_runner(model_bin, "--parity", "--temp", "0", + "--start_token", str(start_token)) + for line in res.stdout.splitlines(): + if line.startswith("[Final_Logits]"): + vals = line.split("] n=", 1)[1].split()[1:] + return torch.tensor([float(v) for v in vals]) + raise AssertionError( + f"runner emitted no Final_Logits.\nstdout:\n{res.stdout[-2000:]}\n" + f"stderr:\n{res.stderr[-2000:]}" + ) + + +def _py_logits(model, start_token): + with torch.no_grad(): + return model(torch.tensor([[start_token]], dtype=torch.long))[0, -1, :] + + +@pytest.fixture(scope="module") +def exported(tmp_path_factory): + torch.manual_seed(1234) + model = AtomicTransformer(CFG) + model.eval() + path = tmp_path_factory.mktemp("parity") / "model.bin" + export_model(model, str(path)) + return model, path + + +class TestActivationContract: + # guirguispierre 2026-08-10 - the numeric tests below cannot catch a wrong + # activation: erf-vs-tanh gelu perturbs logits by the same ~1e-2 that one + # quantizer boundary flip does, and x86 and arm already differ by that much + def test_mlp_uses_tanh_gelu(self): + act = AtomicTransformer(CFG).layers[0].mlp.act + assert isinstance(act, nn.GELU) + assert act.approximate == "tanh", ( + "embedded/atomic_runner.cpp implements the tanh gelu; exact-erf " + "nn.GELU() here diverges by ~4e-4 per layer and compounds" + ) + + def test_tanh_and_erf_gelu_actually_differ(self): + x = torch.linspace(-3, 3, 64) + delta = (nn.GELU()(x) - nn.GELU(approximate="tanh")(x)).abs().max() + assert delta > 1e-5 + + +class TestCppRuntimeParity: + def test_output_stream_is_not_polluted(self, exported): + _, path = exported + res = _run_runner(path, "--steps", "8", "--temp", "0", + "--start_token", "7") + tokens = _generated_tokens(res) + assert len(tokens) == 8, f"expected 8 generated tokens, got {tokens!r}" + + @pytest.mark.parametrize("start_token", [0, 7, 42, 128, 255]) + def test_logits_agree(self, exported, start_token): + model, path = exported + cpp = _cpp_logits(path, start_token) + py = _py_logits(model, start_token) + + assert cpp.shape == py.shape + rel = ((cpp - py).abs().max() / py.abs().max()).item() + assert rel < AGREEMENT_TOL, ( + f"relative logit error {rel:.3e} exceeds {AGREEMENT_TOL:.0e} -- an " + f"op has diverged between atomic_1bit/ and embedded/" + ) + + @pytest.mark.parametrize("start_token", [0, 7, 42, 128, 255]) + def test_disagreement_is_only_ever_a_tie_break(self, exported, start_token): + # guirguispierre 2026-08-10 - asserting equal argmax would be flaky, so + # require instead that pytorch scores the token c++ picked as good as + # its own; a genuinely wrong engine fails this, a near-tie does not + model, path = exported + cpp = _cpp_logits(path, start_token) + py = _py_logits(model, start_token) + + delta = (cpp - py).abs().max().item() + shortfall = (py.max() - py[int(cpp.argmax())]).item() + assert shortfall <= 2 * delta + 1e-6, ( + f"C++ chose token {int(cpp.argmax())}, which PyTorch scores " + f"{shortfall:.4f} below its own pick {int(py.argmax())} -- more " + f"than the {delta:.3e} numerical difference explains" + ) + + +class TestRunnerCLI: + def test_help_exits_cleanly(self): + res = subprocess.run([RUNNER, "--help"], capture_output=True, + text=True, timeout=60) + assert res.returncode == 0 + assert "--model" in res.stdout and "--steps" in res.stdout + + def test_unknown_flag_is_rejected(self): + res = subprocess.run([RUNNER, "--not-a-flag"], capture_output=True, + text=True, timeout=60) + assert res.returncode != 0 From 075400ce9a7525aed8945aa1bacafe2f9b911421 Mon Sep 17 00:00:00 2001 From: Pierre Guirguis <22091706+guirguispierre@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:52:41 -0700 Subject: [PATCH 6/8] build the runner in ci and run on all prs 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. --- .github/workflows/ci.yml | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5ce1c41..37acd5c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,8 +3,9 @@ name: CI on: push: branches: [ master, feature/** ] + # guirguispierre 2026-08-10 - no branch filter, restricting to master meant a + # pr stacked on another pr's branch got no ci run pull_request: - branches: [ master ] jobs: test: @@ -44,6 +45,19 @@ jobs: test -f atomic_1bit/core/libatomic.so || (echo "libatomic.so not found!" && exit 1) ls -lh atomic_1bit/core/libatomic.so + - name: Build embedded C++ runner + run: | + cd embedded + g++ -O3 -std=c++17 atomic_runner.cpp -o runner + cd .. + + - name: Verify runner was built + # guirguispierre 2026-08-10 - without this binary the parity tests skip + # silently and the claim goes unverified + run: | + test -f embedded/runner || (echo "embedded/runner not found!" && exit 1) + ./embedded/runner --help + - name: Run linting - Black run: | black --check atomic_1bit/ tests/ benchmarks/ tools/ scripts/ --exclude='/(\.git|\.venv|build|dist)/' From 284559eff72e463dcea96f197048c283665efc73 Mon Sep 17 00:00:00 2001 From: Pierre Guirguis <22091706+guirguispierre@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:53:32 -0700 Subject: [PATCH 7/8] correct readme claims against measured numbers 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. --- README.md | 91 +++++++++++++++++++++++++++++++++++-------- docs/API_REFERENCE.md | 4 +- docs/INSTALL.md | 23 +++++++++-- 3 files changed, 96 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index fcecc1a..b668d7b 100644 --- a/README.md +++ b/README.md @@ -22,10 +22,12 @@ weight == 0 -> skip (free sparsity) The result is a full transformer that runs on **integer arithmetic only**. No CUDA required. No FP16. No matrix multiply units. Just `add` and `sub` instructions that work on any processor manufactured in the last 30 years. **This matters because:** -- A 1.33M parameter model drops from **5.3 MB to 2.0 MB** (-62%) +- A 6.85M parameter model drops from **27.4 MB to 10.1 MB** (-63%) - The C++ runtime has **zero external dependencies** -- it's a single binary - It runs on a **Raspberry Pi**, an **ESP32**, or a **2015 laptop** -- The Python training stack and C++ inference engine produce **bit-exact identical output** +- The ternary kernel is **bit-exact** against the NumPy reference, and the C++ + engine tracks the PyTorch model to float32 precision (see + [Numerical parity](#numerical-parity)) > This is experimental research software. It works, it's verified, and it's honest about what it is: a proof that useful AI inference doesn't require expensive hardware. @@ -33,15 +35,20 @@ The result is a full transformer that runs on **integer arithmetic only**. No CU ## Performance -Benchmarked on Apple M-series, single thread, sequence length 128, 50 generated tokens: +Benchmarked on Apple M-series, single thread, 50 generated tokens, using the +Quick Start config below (256d, 6 layers, 4 heads, 4096 vocab, 6.85M params): -| Metric | FP16 Baseline | Atomic-1Bit | Improvement | +| Metric | FP32 Baseline | Atomic-1Bit | Improvement | |:---|:---|:---|:---| -| **Model Size** | 5.3 MB | **2.0 MB** | **-62%** | -| **Parameters** | 1.33M | 1.33M | Same | -| **Precision** | Float16 | Ternary {-1, 0, 1} | -- | -| **Throughput (C++)** | N/A | **~160-170 TPS** | Portable runtime | -| **Throughput (Python)** | ~826 TPS | ~130 TPS | Unoptimized | +| **Model Size** | 27.4 MB | **10.1 MB** | **-63%** | +| **Parameters** | 6.85M | 6.85M | Same | +| **Precision** | Float32 | Ternary {-1, 0, 1} | -- | +| **Throughput (C++)** | N/A | **~170-185 TPS** | Portable runtime | + +The C++ runner recomputes the full prefix each step rather than keeping a +KV-cache, so throughput falls off as the sequence grows -- generation is +O(n^2) in the number of tokens. The Python path (`AtomicTransformer.generate`) +*does* cache. Reproduce with `benchmarks/run_suite.py`.
Visual benchmarks @@ -70,6 +77,14 @@ cd Atomic-1Bit pip install -r requirements.txt ``` +### Build the ternary kernel + +The Python inference path loads a shared library, so build it first: + +```bash +make -C atomic_1bit/core # produces atomic_1bit/core/libatomic.so +``` + ### Verify the kernel works This confirms the C++ ternary kernel matches the Python/NumPy reference exactly: @@ -117,10 +132,34 @@ The project has three components: 1. **Research stack** (`atomic_1bit/`) -- PyTorch training, evaluation, and model architecture. Train on TinyStories or Alpaca-cleaned datasets with thermal safety monitoring, gradient accumulation, and cosine scheduling. -2. **Bare-metal runtime** (`embedded/`) -- Standalone C++ inference engine with zero dependencies. Supports CPU, Metal (Apple Silicon), and CUDA backends through conditional compilation. Produces bit-exact output matching the Python reference. +2. **Bare-metal runtime** (`embedded/`) -- Standalone C++ inference engine with zero dependencies. Supports CPU, Metal (Apple Silicon), and CUDA backends through conditional compilation. 3. **Gist tokens** -- Pre-computed "thought vectors" that compress a system prompt into a single embedding, injected into the attention stream at zero inference cost. +### Numerical parity + +Two different guarantees are worth separating, because only one of them is exact: + +| Claim | Status | Verified by | +|:---|:---|:---| +| C++ ternary kernel == NumPy int32 reference | **Bit-exact** | `tests/test_kernel_parity.py` | +| C++ engine == PyTorch model, whole forward pass | Agrees to float32 precision | `tests/test_cpp_runtime_parity.py` | + +Whole-model bit-equality is *not* achievable by design, and the reason is worth +knowing before you go hunting for a bug. Activations are quantized with a +per-tensor scale derived from a max over float values. `round()` is +discontinuous, so a last-ULP difference anywhere upstream -- summation order in +an RMS norm, a `tanh` implementation -- can push one value across a `.5` +boundary. That shifts a single int8 by one step, which is a *discrete* jump in +the accumulator, not an epsilon. Deeper models give this more chances to happen. + +In practice a shallow model matches PyTorch to ~1e-7 relative, and deeper ones +stay within a few percent with the same next-token prediction. Greedy decoding +can still fork after enough steps once two logits are near-tied. The test suite +pins both bounds, so a genuine mismatch -- a different activation function, +scale convention, or weight layout -- fails loudly instead of hiding in the +noise. + ![Ternary Matmul Diagram](assets/diagram_ternary_matmul_1765658104682.png) For more details, see [docs/USAGE.md](docs/USAGE.md). @@ -143,7 +182,7 @@ atomic_1bit/ embedded/ Standalone C++ runner + ESP32 port guide configs/ Model presets (4K pocket to 12.5M flagship) benchmarks/ Reproducible benchmark suite vs FP16 baselines -tests/ 67 pytest tests for correctness verification +tests/ 149 pytest tests for correctness verification scripts/ Plotting, evaluation, and reproduction scripts docs/ Installation, usage, commands, benchmarking guides examples/ Runnable example scripts @@ -153,12 +192,22 @@ examples/ Runnable example scripts ## Model Configurations -| Config | Parameters | Dimensions | Use Case | -|:---|:---|:---|:---| -| [`pocket_4k`](configs/pocket_4k.yaml) | ~100K | 256d, 4L, 4H | ESP32 / microcontrollers | -| [`stories_base`](configs/stories_base.yaml) | ~1.33M | 256d, 6L, 4H | Development / testing | -| [`flagship_12m`](configs/flagship_12m.yaml) | ~12.5M | 320d, 8L, 5H | Quality demos | -| [`mixed_precision`](configs/mixed_precision.yaml) | Configurable | Hybrid 1.58/4-bit | Experimental | +| Config | Parameters | Exported size | Dimensions | Use Case | +|:---|:---|:---|:---|:---| +| [`pocket_4k`](configs/pocket_4k.yaml) | 5.28M | 8.5 MB | 256d, 4L, 4H, 4K vocab | Smallest preset | +| [`stories_base`](configs/stories_base.yaml) | 30.5M | 69.2 MB | 256d, 6L, 4H, 50257 vocab | Development / testing | +| [`flagship_12m`](configs/flagship_12m.yaml) | 12.5M | 16.7 MB | 320d, 8L, 5H, 4K vocab | Quality demos | +| [`mixed_precision`](configs/mixed_precision.yaml) | Configurable | -- | Hybrid 1.58/4-bit | Experimental | + +Parameter counts are dominated by the token embedding table, which stays float32 +(`vocab_size x dim x 4` bytes) -- so vocabulary size, not depth, drives file +size. `stories_base.yaml` ships with the 50257-entry GPT-2 vocabulary; the +Quick Start above passes `--vocab_size 4096` instead, giving a 6.85M-parameter +model that exports to 10.1 MB. Set the vocabulary deliberately for your target. + +> None of these presets fit in ESP32 SRAM as-is. See +> [embedded/ESP32_PORT_GUIDE.md](embedded/ESP32_PORT_GUIDE.md) for what actually +> fits on-device. Load any config with: @@ -191,6 +240,8 @@ config = config_to_atomic(load_config("configs/stories_base.yaml")) ## Running Tests ```bash +pip install pytest + # Run the full test suite pytest tests/ -v @@ -199,6 +250,12 @@ pytest tests/test_bitlinear.py -v pytest tests/test_kernel_parity.py -v ``` +Two modules need compiled artifacts and skip without them: +`test_kernel_parity.py` needs `make -C atomic_1bit/core`, and +`test_cpp_runtime_parity.py` needs the runner +(`cd embedded && g++ -O3 -std=c++17 atomic_runner.cpp -o runner`). Build both +before trusting a green run. + --- ## Contributing diff --git a/docs/API_REFERENCE.md b/docs/API_REFERENCE.md index b41b6ba..239dc0c 100644 --- a/docs/API_REFERENCE.md +++ b/docs/API_REFERENCE.md @@ -128,7 +128,7 @@ Core BitLinear compute kernel. Implements one linear layer of the form: quantize 2. Integer accumulation: `acc[o] += x_q[i] * w[i * out_dim + o]` (uses actual `int32_t` multiplication; no branch-on-sign optimization at this layer). 3. Dequantize: `out[o] = acc[o] / scale`. -**Known limitation:** The weight scale factor is not stored in the binary format used by `atomic_lib.h` and is therefore not applied during dequantization. The runner (`atomic_runner.cpp`) corrects this by reading and applying per-layer weight scales from the extended binary format written by `export_to_cpp.py`. +Per-tensor `float32` weight scales are read from the binary and applied during dequantization, matching the Python `BitLinear` forward. `atomic_lib.h` and `atomic_runner.cpp` now consume the identical format written by `export_to_cpp.py`. --- @@ -265,4 +265,4 @@ Each layer block contains: | `4 + dim * 4*dim` | `float32 + int8[]` | FC1 weight (scale + weights) | | `4 + 4*dim * dim` | `float32 + int8[]` | FC2 weight (scale + weights) | -Note: `atomic_lib.h` uses an older format that does not include the per-layer scale prefix for weights. The `atomic_runner.cpp` binary uses the current format with scale values. +Note: `atomic_lib.h` and `atomic_runner.cpp` read the same format, including the per-tensor `float32` scale prefix on every quantized weight tensor. diff --git a/docs/INSTALL.md b/docs/INSTALL.md index 04a385c..c01e3c9 100644 --- a/docs/INSTALL.md +++ b/docs/INSTALL.md @@ -37,11 +37,28 @@ This installs: python3 -c "import torch; import tiktoken; print('Dependencies OK')" ``` -## C++ Runtime (Optional) +## C++ Components -The C++ inference engine is optional -- you only need it if you want to run models on bare metal or embedded devices. +There are two separate C++ artifacts, and the first one is **not** optional. -### Build the engine +### 1. Ternary kernel (`libatomic.so`) -- required for Python inference + +`atomic_1bit/python/wrapper.py` loads this shared library, so +`python3 atomic_1bit/python/inference.py` and the kernel parity tests fail +without it: + +```bash +cd atomic_1bit/core +make # defaults to BACKEND=CPU +``` + +Verify: + +```bash +test -f atomic_1bit/core/libatomic.so && echo "kernel built" +``` + +### 2. Bare-metal runner -- needed to run exported models standalone ```bash cd embedded From e4936561cec85a2b95b115fe031736d412792363 Mon Sep 17 00:00:00 2001 From: Pierre Guirguis <22091706+guirguispierre@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:53:32 -0700 Subject: [PATCH 8/8] disable co-authored-by trailer Per AGENTS.md, commits carry no agent attribution. --- .claude/settings.json | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 .claude/settings.json diff --git a/.claude/settings.json b/.claude/settings.json new file mode 100644 index 0000000..bcf43f6 --- /dev/null +++ b/.claude/settings.json @@ -0,0 +1,3 @@ +{ + "includeCoAuthoredBy": false +}