Micro-Llama on Shrike-Lite 🔌

Overview

The goal was to push a transformer language model to the absolute edge: train a small Llama from scratch, then run it entirely bare-metal on a microcontroller with no operating system, no allocator, and no floating-point hardware. The board is the Vicharak Shrike-Lite — a dual-core RP2040 paired with a Renesas SLG47910 FPGA, sharing just 264KB of SRAM.

The full pipeline lives in one repo: JAX/Flax training on a GPU, export to a flat binary, then a hand-written C inference core that reads weights from QSPI flash via XIP and streams generated text over USB CDC.

Pipeline

TinyStories → JAX train (RunPod) → micro_llama.bin → C firmware → UF2 on Shrike-Lite

The model is trained in JAX/Flax, exported to a single micro_llama.bin checkpoint, embedded into the firmware, and flashed to the board as a UF2 — replacing MicroPython entirely.

Model Architecture

Standard Llama building blocks: GQA (2 query heads, 2 key-value heads, head_dim 32), RoPE, RMSNorm, and a SwiGLU feed-forward network with an untied output head. Special tokens: pad=0, unk=1, bos=2, eos=3.

The 264KB SRAM Constraint

Everything has to fit on a chip with 264KB of total SRAM and no FPU. The architecture is locked by that budget: weights live in QSPI flash via XIP (execute-in-place), while activations and the KV cache must fit inside SRAM at runtime. The fp32 KV cache alone occupies ~131KB — which is why the shape above is as small as it is. FPGA INT8 MAC acceleration on the SLG47910 (~1120 LUTs) is mapped out but not yet started.

Training Details

The training runtime auto-selects CUDA → Metal → CPU before importing JAX, so the same code trains on an NVIDIA GPU, an Apple Silicon Mac, or plain CPU.

Firmware Targets

The C inference core builds to two targets from the same source (-DML_TARGET=host|pico):

• host — native greedy decoder for parity testing against the JAX reference (verify_greedy.py)
• pico — Shrike-Lite UF2 with weights + tokenizer baked into flash, streaming text over USB CDC
• XIP weights — model weights and the BPE tokenizer are embedded in QSPI flash and read in place, never copied to SRAM
• Soft-float RP2040 — generation runs on the Cortex-M0+ without hardware floating point
• JAX ↔ C parity — the host decoder is verified to match greedy generation bit-for-bit

Quick Start

Train (JAX / Flax)

uv sync
uv pip install -U "jax[cuda12]"   # NVIDIA; Mac: jax-metal or CPU
cp .env.example .env              # WANDB_API_KEY, HF_TOKEN, HF_REPO_ID

uv run python train/train.py
# dry run:
uv run python train/train.py --no-wandb --no-hub --max-steps 100

Host inference + parity check

bash firmware/scripts/download_weights.sh
uv run python firmware/scripts/embed_tokenizer.py
cmake -S firmware -B firmware/build-host -DML_TARGET=host && cmake --build firmware/build-host -j
./firmware/build-host/micro_llama_host \
  --checkpoint firmware/weights/micro_llama.bin --bos --max-new-tokens 32 --text
uv run python firmware/scripts/verify_greedy.py

For the Pico target, build with -DML_TARGET=pico, then hold BOOTSEL and run flash_uf2.sh. Open the serial console before resetting so the CDC wait captures the full session.

Project Structure

• train/train.py — JAX/Flax pretraining loop with W&B logging
• train/runtime.py — device selection (CUDA → Metal → CPU) before importing JAX
• train/export.py — Flax checkpoint → flat micro_llama.bin
• train/generate.py — greedy generation in Python
• train/hf_upload.py — publish weights + tokenizer to the HuggingFace Hub
• firmware/ — C inference core (host CLI + RP2040 UF2)
• firmware/scripts/ — weight download, tokenizer embed, UF2 flash helpers

Resources

• Source Code on GitHub
• Model Weights on HuggingFace
• Weights & Biases Training Report
• Vicharak Shrike-Lite Board