Skip to content

Repository files navigation

LLM Weight Streaming on ESP32

Run a real Llama-2-architecture model on an ESP32 that could never hold it: the weights live on your PC and stream over WiFi, the entire model crossing the air for every single token. The ESP32 keeps only a small working buffer; RAM usage stays flat no matter the model size, and generation is bounded by network bandwidth, not flash or RAM.

heap [model ready]: free 110 KiB, largest block 68 KiB
token   9038 'Once'  (7.13 s, free heap 110 KiB, worst-ever heap 45 KiB)
        net 5.38 s (57 reqs, 94.4 ms/req, 3103 KB/s) matmul 1.58 s other 0.16 s
=======
Once
=======
token   2501 ' upon'  (6.97 s, free heap 110 KiB, worst-ever heap 45 KiB)
        net 5.22 s (57 reqs, 91.6 ms/req, 3200 KB/s) matmul 1.58 s other 0.17 s
=======
Once upon
=======
token    263 ' a'  (6.97 s, free heap 110 KiB, worst-ever heap 44 KiB)
        net 5.24 s (57 reqs, 91.8 ms/req, 3190 KB/s) matmul 1.57 s other 0.17 s
=======
Once upon a
=======
token    931 ' time'  (6.96 s, free heap 110 KiB, worst-ever heap 44 KiB)
        net 5.20 s (57 reqs, 91.3 ms/req, 3209 KB/s) matmul 1.60 s other 0.16 s
=======
Once upon a time
=======

(example output; numbers depend on your WiFi; replace with a capture from your run)

How it works

Architecture: ModelHost on the PC serves weight slices over WiFi to the ESP32's shared buffer

Per token, the ESP32 sends the whole forward pass's request recipe in one batched write: embedding row, then per layer rms_att → wq → wk → wv → wo → rms_ffn → w1 → w3 → w2, then the final norm and a classifier sweep, and the compute path drains the responses in the same order:

Per-token sequence: one batched request send, then per-layer weight bites and the streamed classifier sweep

The tricks that make it fit:

  • No framing, order is the contract: requests name a tensor plus up to two slices (LLAMA2_Layers.h); both sides derive response sizes from shapes, so the wire carries pure payload.
  • ChunkedProject: a layer matrix larger than the buffer is consumed in bites of whole rows, each matmul'd into its slice of the output. Chunking never changes the math.
  • Sliding-window KV cache: k/v stored in q8, pre-RoPE, in a sliding window with BOS pinned. Each step feeds one token.
  • Streamed logits: all 32000 classifier rows flow through the buffer; only the running argmax (or Gumbel-max sample) survives. Full logits never materialize.
  • Weights are q8 (32 int8 quants + f32 scale = 36 B per 32 elements); activations, norms and scores stay f32.

Quickstart

You need: an ESP32-S3 (or classic ESP32/WROOM) on 2.4 GHz WiFi, ESP-IDF ≥ 5.x, a C++17 compiler on the PC, Python 3 with numpy.

1. Get and quantize the model (PC)

cd models
python download.py    # stories15M.bin + tokenizer.bin (Karpathy's llama2.c artifacts)
python quantize.py    # -> stories15M-q8.bin, 3.6x smaller

2. Flash the ESP32

Put your WiFi credentials in main/wifi_config.h, then:

idf.py set-target esp32s3     # or esp32 for a WROOM board
idf.py build flash monitor

The monitor prints the board's IP and listening on port 9000.

3. Serve the weights (PC)

cmake -S host -B host/build
cmake --build host/build --config Release
./host/build/Release/ModelHost <esp32-ip>     # run from the repo root

ModelHost connects, greets with the model config, and the ESP32 starts generating a TinyStories tale token by token, one connection per story; it listens again after EOS.

Performance

target buffer tok/s notes
ESP32-S3 96 KiB ~0.2 ~17 MB downloaded per token
ESP32 (WROOM) 48 KiB ~0.15 smaller TCP window, less DRAM

The per-token log splits time into net (wire), matmul, and other, so you can see exactly what you're bound by. Tuning lives in sdkconfig.defaults (TCP window, WiFi RX buffers) with WROOM overrides in sdkconfig.defaults.esp32.

All matmuls are plain scalar loops, no SIMD yet. The ESP32-S3 has 128-bit vector instructions (PIE) that fit the q8 int8 dot products well, so there is known headroom on the matmul share whenever the network stops being the bottleneck.

Layout

main/            ESP32 app: main.cpp, WifiManager, Esp32MemoryManager
main/tensor/     Tensor.h (q8/f32), TensorOps.h, StreamingTensor.h
main/llama2/     StreamingModel.h (ESP32), Model.h (PC reference), wire protocol
host/            ModelHost.cpp: loads the model, serves weight slices
models/          download.py, quantize.py, model binaries (not in git)

Model.h runs the same model standalone on the PC (with and without KV cache) and doubles as the reference implementation the streaming side is checked against.

Credits

Model, tokenizer and checkpoint format from Andrej Karpathy's llama2.c (stories15M trained on TinyStories). PC sockets via kissnet.

License

MIT

About

A technical demo of streaming LLM weights over wifi to run a 16 MB quantized llama2 in an ESP32 with 300KB of ram

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages