Metadata-Version: 2.4
Name: kasahare
Version: 0.5.0
Summary: A self-contained LLaMA-style inference engine with a Python binding.
Author: simple-llm
License: MIT
Keywords: llm,inference,transformer,llama,c++
Classifier: Programming Language :: C++
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: POSIX :: Linux
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.20
Provides-Extra: test
Requires-Dist: pytest>=7; extra == "test"
Dynamic: requires-python

# simple-llm / kasahare

`kasahare` (which means "speak fast" in Twi) is the PyPI distribution name for simple-llm.
A self-contained C++17 inference **and training** engine for transformer
language models. No external dependencies beyond the C++ standard library
and a recent `g++` / `clang++`. Designed to be read end-to-end in an
afternoon, then hacked on.

## What's in the box

### Inference

| file | purpose |
|---|---|
| `include/simplellm/simd.h`        | AVX2 + FMA kernels (matmul, RMSNorm, SwiGLU, vec ops) |
| `include/simplellm/tensor.h`      | Lightweight matrix view + matmul / linear kernels |
| `include/simplellm/mmap_loader.h` | Zero-copy mmap weight loader and INT4 block dequantization |
| `include/simplellm/config.h`      | Static model description |
| `include/simplellm/tokenizer.h`   | Byte-level BPE tokenizer (no HF tokenizers dep) |
| `include/simplellm/sampler.h`     | Greedy / temperature / top-k / top-p sampler |
| `include/simplellm/model_data.h`  | On-disk model format reader / writer |
| `include/simplellm/model.h`       | Public surface: `Model::load`, `forward`, `generate` |
| `src/model.cpp`                   | LLaMA block: RMSNorm, GQA attention, RoPE, SwiGLU |
| `src/main.cpp`                    | `sl-llm` CLI |
| `tools/create_tiny_model.cpp`     | Synthesises a tiny model on disk for smoke testing |
| `tools/bench.cpp`                 | Micro-benchmark: `make bench` |
| `tests/test_basic.cpp`            | End-to-end smoke test: `make test` |

### Training (Native C++ Engine)

| file | purpose |
|---|---|
| `include/simplellm/train_model.h` | `TrainModel`, `TrainContext`, `ModelGradients`, `LayerActivations` |
| `include/simplellm/optimizer.h`   | `AdamW` optimizer + `OptimizerConfig` |
| `src/train_model.cpp`             | Full-sequence forward pass with fused Flash Attention |
| `src/backward.cpp`                | Reverse-mode differentiation with on-the-fly math for Flash Attention |
| `src/optimizer.cpp`               | Native AdamW parameter updates (no PyTorch) |

## Architecture

### Inference Forward Pass

For each block (LLaMA-2 style):

```
x_n              = x_{n-1}                                  # residual
x_n              = attn_norm(x_n)
q, k, v          = Wq x_n, Wk x_n, Wv x_n                  # no biases
q, k             = RoPE(q, k, pos)
attn             = softmax(q . K^T / sqrt(d_head)) . V
x_n              = Wo(attn) + x_n                           # residual
x_n              = ffn_norm(x_n)
x_n              = W2( silu(W1 x_n) * (W3 x_n) ) + x_n      # SwiGLU, residual
```

### Training Backward Pass

The native C++ engine implements exact reverse-mode differentiation for every layer:

| Layer | Gradient |
|---|---|
| Cross-entropy loss | Numerically stable log-softmax backward |
| SwiGLU FFN | `dW2`, `dW1`, `dW3` via SiLU gate chain rule |
| RMSNorm | Per-element scale derivative |
| Scaled dot-product attention | `dQ`, `dK`, `dV`, `dWo` with causal mask |
| RoPE | Exact inverse rotation (negated sin/cos phase) |

Activations are cached at every layer during the sequence-level forward pass (`TrainContext`) and consumed during the backward pass — no re-computation needed.

### Optimisations

* `matmul` and `linear` are tiled over the inner dim with 8-wide AVX2 + FMA, and effortlessly scale across CPU threads using OpenMP parallelization.
* `rmsnorm`, `vec_add`, `vec_mul`, `silu_mul` are SIMD-ised.
* A pre-allocated KV cache (`KvSlot` per head per layer) keeps inference `forward()` O(1) allocation cost.
* Tied input / output embeddings  — one embedding matrix.
* AdamW uses per-parameter `m`/`v` moment vectors; no framework overhead.

The block uses **GQA** (grouped query attention): `n_kv_heads < n_heads` (or equal for MHA, 1 for MQA).

## Build

```sh
make            # builds build/bin/{sl-llm, create_tiny_model, bench}
make test       # builds and runs the smoke test
make bench      # runs the benchmark
make clean
```

Requirements: `g++` ≥ 9, AVX2 + FMA support. Override SIMD flags with
`make CFLAGS="-O3 -march=native"` to autodetect, or `make CFLAGS="-O3 -msse4.2"` to disable AVX2.

## Quick start (Inference CLI)

```sh
# 1) Generate a tiny model + tokenizer (~110 KB total)
./build/bin/create_tiny_model --out build

# 2) Run it
./build/bin/sl-llm \
    --model build/tiny.sllm \
    --tokenizer build/tiny.tok.txt \
    --prompt "hi" \
    --max-new 16 \
    --temperature 0.8

# 3) Show timing
./build/bin/sl-llm \
    --model build/tiny.sllm \
    --tokenizer build/tiny.tok.txt \
    --prompt "" \
    --max-new 64 --temperature 0 --timing
```

## Quick start (Trainer CLI)

The `cpp_train_infer` binary exposes dedicated `--train` and `--infer` subcommands to encapsulate distinct pipelines safely. During training, a real-time ASCII Sparkline UI calculates dynamic loss curves in your terminal!

```sh
# 1) Train a model natively inside C++
./build/cpp_train_infer --train --doc examples/trainingdocs/health/stg_chatml_finetune.txt \
    --output my_model_output \
    --batch 16 \
    --epochs 5 \
    --lr 5e-4
    
# 2) Test inference natively
./build/cpp_train_infer --infer --output my_model_output --prompt "hello" --max_new 32
```

## CLI flags
*Note: For a fully maintained reference of all C++ CLI flags available on `sl-llm` and `cpp_train_infer`, see [docs/cli_reference.md](docs/cli_reference.md).*

```
--model <file.sllm>        model file
--tokenizer <file.txt>     tokenizer file
--prompt <text>            input text
--max-new <n>              number of tokens to generate (default 32)
--temperature <f>          sampling temperature, 0 = greedy (default 0.8)
--top-k <n>                top-k filtering (default 40)
--top-p <f>                top-p filtering (default 0.95)
--seed <n>                 PRNG seed (0 = random)
--show-config              print model config and exit
--no-print-prompt          don't echo the prompt before the generated tail
--timing                   print per-step timing to stderr
--stream                   flush stdout after each generated token
--help                     this message
```

## Python binding

The engine is `pip install`-able as `kasahare` (on PyPI) or installable from source:

```sh
pip install kasahare

# From source (development):
pip install -e .
```

### Inference API

```python
import simplellm

m = simplellm.Model.load("model.sllm", "model.tok.txt")

# Tokenize
ids = m.tokenizer.encode("hello world", add_bos=True)

# Single-token forward (returns np.ndarray of shape (vocab_size,))
logits = m.forward(ids[0], reset_kv=True)

# Generate using probability sampling
out = m.generate(ids, max_new=32,
                  options=simplellm.SampleOptions(temperature=0.7, top_k=40))
print(m.tokenizer.decode(out))

# Generate using Beam Search
out_beam = m.generate_beam(ids, max_new=32, beam_width=3)
print(m.tokenizer.decode(out_beam))

# Streaming
def on_token(tid): print(m.tokenizer.decode_token(tid), end="", flush=True)
m.reset_kv()
m.generate(ids, max_new=64, options=simplellm.SampleOptions(temperature=0.8), stream=on_token)

# In-memory (no file paths)
m2 = simplellm.Model.from_bytes(open("model.sllm","rb").read(),
                                 open("model.tok.txt","rb").read())

# Tracing — per-layer stats at zero overhead normally
trace = m.forward_traced(ids[0], reset_kv=True)
print("Max logit:", trace["logit_max"], "Entropy:", trace["logit_entropy"])
print("RMS values, layer 0:", trace["layers"][0]["attn_norm_rms"])
```

### Training API (Native C++ Engine)

The training API bypasses PyTorch entirely. Gradients are computed by the native
reverse-mode differentiation engine in C++.

```python
import simplellm

# Load model weights
m = simplellm.Model.load("model.sllm", "model.tok.txt")

# Create trainer + optimizer bound to the same weights
trainer = simplellm.TrainModel(m.weights)
optimizer = simplellm.AdamW(m.weights, simplellm.OptimizerConfig(
    lr=1e-3,
    beta1=0.9,
    beta2=0.999,
    eps=1e-8,
    weight_decay=0.01,
))

# Training loop
tokens = m.tokenizer.encode("The quick brown fox", add_bos=True)
targets = tokens[1:]       # next-token targets
inputs  = tokens[:-1]

ctx = simplellm.TrainContext()
trainer.forward(inputs, ctx)          # full-sequence forward + cache activations
loss = trainer.backward(ctx, targets) # reverse-mode diff, accumulate gradients
optimizer.step()                       # AdamW parameter update
optimizer.zero_grad()                  # reset gradient accumulators

print(f"Loss: {loss:.4f}")
```

**What the C++ engine handles:**
- Full-sequence forward pass storing all intermediate activations
- Numerical cross-entropy loss with log-softmax stability
- Exact gradients for: SwiGLU W1/W2/W3, RMSNorm (attn + FFN), Attention Q/K/V/O, inverse RoPE rotation
- AdamW with momentum, variance, weight decay, and bias correction
- Zero-allocation hot path (no heap alloc per training step after warm-up)

## On-disk model format (`.sllm`)

```
+---------------------------------------------------+
| magic   (u32 LE, 0x534C4C4D == "SLLM")             |
| version (u32 LE, currently 1)                      |
| blen    (u32 LE)                                   |
| config blob  (blen bytes, key=value text)          |
| weights blob (u32 × N floats, little-endian)       |
+---------------------------------------------------+
```

Weights are stored in a single contiguous float32 array:

```
token_emb                (vocab_size * dim)
for each block:
    attn_norm            (dim)
    wq                   (dim * dim)
    wk                   (kv_dim * dim)
    wv                   (kv_dim * dim)
    wo                   (dim * dim)
    ffn_norm             (dim)
    w1                   (hidden * dim)   # SwiGLU gate
    w3                   (hidden * dim)   # SwiGLU up
    w2                   (dim * hidden)   # SwiGLU down
output_norm              (dim)
```
If `dtype=1` (Q4_0), the weights blob contains `BlockQ4_0` structures (20 bytes per 32 floats) rather than raw float32 values. Both `Model::load` and `MMapLoader` support seamless automatic conversion of these structures back to the float32 formats required by the math pipeline natively on load.

`tied_embeddings=1` reuses `token_emb` for the output projection (LLaMA style).

## Tokenizer format

```
# simple-llm tokenizer
vocab_size <N>
bos <id>
eos <id>
token <id> "<csv-escaped-payload>"
merge <rank> "<csv-escaped-left>" "<csv-escaped-right>"
```

Payloads are byte strings. The CSV-escape parser handles `"`, `\`, control chars,
embedded spaces, and arbitrary bytes via `\xHH`. The BPE merge rank determines
priority. A tokenizer with zero merges falls back to longest-match — sufficient
for byte-level pre-tokenization.

## Performance

### Inference
*(single CPU core, g++ 12.2, AVX2 + FMA, FP32)*

| config | tokens/s |
|---|---|
| dim=256,  layers=4,  heads=4,  hidden=512,  seq=128  |   ~380 |
| dim=384,  layers=6,  heads=6,  hidden=1024, seq=256  |   ~60  |
| dim=512,  layers=8,  heads=8,  hidden=1536, seq=256  |   ~24  |

The decode loop is O(seq_len) per token. Most time is in `linear()` calls for Q/K/V and FFN.

### Training Benchmark
Wall-clock step time measured via `test_backward.py` comparing the native C++ engine against PyTorch `loss.backward()` + `AdamW` (dim=64, layers=2, SEQ=16).

| Engine | Mean (ms) | Min (ms) | Max (ms) |
|---|---|---|---|
| PyTorch CPU | 2.32 | 2.20 | 3.47 |
| Native C++ | 1.70 | 1.62 | 1.93 |

**Result:** The C++ engine shows a **~1.37× speedup** over PyTorch natively on the CPU, benefiting from zero per-step allocation in the hot path. *(Note: The forward pass cache is currently incomplete; fully projecting Q/K/V will slightly reduce this speedup).*

## Limitations

* Engine compute is FP32 natively. Models can be loaded aggressively out of space-saving INT4 (Q4_0) or flat FP32 configurations securely.
* GPU training is not supported. All compute defaults purely to the CPU backend, scaling horizontally across available cores effortlessly via OpenMP limits.
* On-disk format is little-endian only.

## Where to go from here

* Add a HuggingFace `tokenizer.json` reader.

## License

Do whatever you want with it. No warranty — see "Limitations" above.
