fix python/c++ parity, runner throughput, and int8 kernel overflow

.claude/settings.json

@@ -0,0 +1,3 @@
+{
+  "includeCoAuthoredBy": false
+}

.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)/'

README.md

@@ -22,26 +22,33 @@ 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.
 
 ---
 
 ## 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`.
 
 <details>
 <summary>Visual benchmarks</summary>
@@ -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

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

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)))

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.

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

embedded/atomic_lib.h

@@ -48,14 +48,15 @@ inline void softmax(vector<float> &x) {
 // Affine RMSNorm (used for the layer's ln1 / ln2 / ln_f).
 inline void rms_norm(const vector<float> &x, const vector<float> &w,
                      vector<float> &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<float> &x, const vector<int8_t> &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,30 +95,35 @@ inline void bit_linear(const vector<float> &x, const vector<int8_t> &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<int8_t> 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)
       q = -127.0f;
     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<int32_t> 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 {