Fix C++ deployment parity, add top-p sampling, robust configs

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

CHANGELOG.md

@@ -0,0 +1,72 @@
+# Changelog
+
+All notable changes to Atomic-1Bit are documented here.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [1.4.0] - 2026-04-30
+
+### Fixed
+- **C++ deployment parity** — `embedded/atomic_lib.h` previously read 6 ints in the wrong order, skipping the `ATOM` magic and version, treating `magic` as `vocab_size`, and never consuming the per-tensor `float32` weight scale written by the exporter. Rewritten to read the full 8-int header, validate magic/version, consume weight scales, and apply them in `bit_linear` via SubLN-aware dequantization that matches the Python `BitLinear` forward exactly. (`atomic_runner.cpp` was already correct; this brings the documented header-only API in line.)
+- **Trainer crash on temperature logging** — `train_instruct.py` called a nonexistent `thermal_monitor.get_current_temp()`. Added the public method to `ThermalMonitor` and made the trainer format `None` defensively when sensors are unavailable.
+- **`chat.py` produced garbage on most checkpoints** — hardcoded `vocab=50257, dim=256, depth=4` only matched one model. Now auto-infers `vocab / dim / depth / context_length` from the checkpoint's `state_dict`, with CLI overrides.
+- **`gen_gist.py` was locked to flagship dims** — added `--dim / --depth / --heads / --vocab_size / --context_length` CLI flags.
+
+### Added
+- **Top-p (nucleus) and top-k sampling** in the C++ runner alongside the existing temperature sampling. New flags: `--top_k`, `--top_p`.
+- **Multi-token prompts** — `--prompt 1,2,3` now accepts a comma-separated starting sequence, overriding `--start_token`.
+- **Streaming controls** — `--stream` (default) prints tokens as generated; `--no-stream` buffers until end.
+- **Runner help** — `--help` / `-h` flag prints all options. Unknown flags now error instead of being silently ignored.
+- **Export/load roundtrip test** (`tests/test_export_roundtrip.py`) — parses the exported binary using the C++ loader's exact expected layout and asserts every field arrives where the loader expects it. Catches future drift between exporter and loader.
+
+### Changed
+- `chat.py`: cleaner device auto-selection (`cuda` → `mps` → `cpu`), explicit error messages on missing checkpoint, `--max_new_tokens` flag.
+- Documentation (`docs/USAGE.md`, `docs/COMMANDS.md`) updated to cover the new sampling and prompt flags.
+
+### Verification
+- 132 pytest tests pass (was 129; +3 roundtrip).
+- CPU kernel parity (NEON SIMD path on Apple Silicon) passes against NumPy reference across 9 size configurations.
+
+## [1.3.0] - 2026-03
+
+### Added
+- 12.5M flagship configuration (`flagship_12m`): 320d, 8 layers, 5 heads, 256 ctx.
+- Evaluation harness — perplexity, repetition rate, coherence, diversity metrics.
+- KV-cache for autoregressive generation.
+- YAML config system with `pocket_4k`, `stories_base`, `flagship_12m`, `mixed_precision` presets.
+- Tokenizer abstraction layer.
+- ESP32 deployment guide and reference port (`embedded/platforms/esp32/`).
+- Raspberry Pi NEON benchmark (`benchmarks/platforms/rpi/`).
+- WebAssembly port (`embedded/platforms/wasm/`).
+- CI/CD pipeline (GitHub Actions) with performance regression detection.
+
+## [1.2.0] - 2026-02
+
+### Added
+- Metal backend (Apple Silicon GPU) with verified parity vs CPU reference.
+- CUDA backend (NVIDIA GPU) with verified parity vs CPU reference.
+- Modular backend selection via `make BACKEND=CPU|METAL|CUDA`.
+- SIMD acceleration in the CPU backend: NEON (ARM) and AVX2 (x86_64) paths with scalar fallback.
+
+## [1.0.0] - 2026-01
+
+### Added
+- Bit-exact Python ↔ C++ kernel parity (zero numerical divergence vs NumPy reference).
+- Standalone C++ inference runtime (`embedded/atomic_runner.cpp`) with zero external dependencies.
+- BitLinear layer with Straight-Through Estimator (STE) for training.
+- Activation quantization (INT8) and weight quantization (ternary {-1, 0, 1}).
+- Gist tokens — pre-computed thought vectors injected at zero inference cost.
+- Training stack: TinyStories and Pocket model variants.
+- Exporter (`atomic_1bit/utils/export_to_cpp.py`) for PyTorch → `.bin` conversion.
+- Benchmark suite vs FP16 baselines.
+
+### Fixed
+- Kernel layout mismatch — Python wrapper passed `(K, N)` weights to a kernel expecting `(N, K)`. Parity now 100%.
+- Memory alignment handling in the C++ loader.
+- Export dimension validation now requires explicit flags to prevent silent shape mismatches.
+
+[1.4.0]: https://github.com/guirguispierre/Atomic-1Bit/compare/v1.3.0...v1.4.0
+[1.3.0]: https://github.com/guirguispierre/Atomic-1Bit/compare/v1.2.0...v1.3.0
+[1.2.0]: https://github.com/guirguispierre/Atomic-1Bit/compare/v1.0.0...v1.2.0
+[1.0.0]: https://github.com/guirguispierre/Atomic-1Bit/releases/tag/v1.0.0

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
@@ -209,13 +266,13 @@ We welcome contributions. See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines
 
 ## Roadmap
 
-See [ROADMAP.md](ROADMAP.md) for the full development plan.
+See [ROADMAP.md](ROADMAP.md) for the full development plan and [CHANGELOG.md](CHANGELOG.md) for per-release detail.
 
 - **v1.0** -- Parity-verified ternary inference (done)
-- **v1.2** -- Hardware-native backends: Metal, CUDA (done)
+- **v1.2** -- Hardware-native backends: Metal, CUDA, NEON/AVX2 SIMD (done)
 - **v1.3** -- Model scaling, evaluation harness, 12.5M config (done)
-- **v2.0** -- YAML configs, tokenizer abstraction, model presets (done)
-- **Next** -- SIMD acceleration (AVX2/NEON), mobile demos, mixed-precision training
+- **v1.4** -- Format-correct C++ deployment, top-p / top-k sampling, robust CLIs (current)
+- **Next** -- Mobile demos, mixed-precision training, HTTP serving wrapper
 
 ---
 

atomic_1bit/__init__.py

@@ -1 +1,3 @@
 """Atomic-1Bit: Bare-metal inference engine for 1.58-bit ternary language models."""
+
+__version__ = "1.4.0"

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