diff --git a/README.md b/README.md index 102021b..0755811 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,57 @@ thousand parameters, so this one holds about a hundred times more. It fits becau most of the model lives in flash instead of RAM, using an idea from Google's Gemma models called Per-Layer Embeddings. +## Reproducing It + +This repo is source-first: the generated training data, checkpoints, export +artifacts, and firmware assets are not committed. To reproduce the ESP32 story +demo from a fresh checkout, you need: + +1. Python 3.12+ +2. `uv` +3. A TinyStories download and BPE/data bins from `data/prepare.py` +4. A trained checkpoint in `runs/` +5. Arduino CLI plus an ESP32 Arduino core for the firmware build + +The easiest path is the bundled wrapper: + +```bash +uv run python deploy.py +``` + +That command does the full sequence in order: prepare the data, train the deploy +checkpoint, export `firmware/model/model.bin`, generate `firmware/esp32_llm/vocab.h`, +and verify the exported binary on the host. + +If you already have the data and checkpoint, you can skip ahead with: + +```bash +uv run python deploy.py --skip-data --skip-train +``` + +For a start-to-finish run that prepares data, trains, exports, compiles, and +flashes in one command, use: + +```bash +uv run python flash.py --full-pipeline --force-train +``` + +If you want to reuse an existing checkpoint (skip retraining), run: + +```bash +uv run python flash.py --full-pipeline +``` + +If the model/checkpoint already exists and you only want build+flash, use: + +```bash +uv run python flash.py +``` + +The device sketch plays the same “Once upon a time” prompt that the firmware +README documents, so the on-chip demo is a small storyteller rather than a chat +assistant. + ## The numbers | | | diff --git a/deploy.py b/deploy.py new file mode 100644 index 0000000..90060da --- /dev/null +++ b/deploy.py @@ -0,0 +1,148 @@ +"""Single-command deploy for the ESP32 PLE demo. + +This wraps the existing flow: + 1. prepare data/tokenizer if needed + 2. train the deploy checkpoint if needed + 3. export the model binary and golden reference + 4. generate firmware vocab assets + 5. compile and run the host verifier + +The default configuration matches the deploy run documented in the repo: +vocab 32768, d_model 96, n_layers 6, ple_dim 128, target core 560000. +""" + +from __future__ import annotations + +import argparse +import subprocess +import sys +from pathlib import Path + + +ROOT = Path(__file__).resolve().parent +DATA = ROOT / "data" +RUNS = ROOT / "runs" +FIRMWARE_MODEL = ROOT / "firmware" / "model" +FIRMWARE = ROOT / "firmware" +SRC = ROOT / "src" + + +def run_step(title: str, argv: list[str]) -> None: + print(f"\n==> {title}", flush=True) + print(" " + " ".join(argv), flush=True) + subprocess.run(argv, cwd=ROOT, check=True) + + +def data_artifacts_exist(vocab: int) -> bool: + suffix = "" if vocab == 4096 else f"_v{vocab}" + return all( + (DATA / name).exists() + for name in ( + f"train{suffix}.bin", + f"val{suffix}.bin", + f"bpe{vocab}.json", + ) + ) + + +def checkpoint_path(arm: str, tag: str, seed: int) -> Path: + return RUNS / f"{arm}{('-' + tag) if tag else ''}-s{seed}.pt" + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--arm", default="ple", choices=["baseline", "ple", "ple_notable", "fatembed", "bigcore"]) + ap.add_argument("--vocab", type=int, default=32768) + ap.add_argument("--d-model", type=int, default=96) + ap.add_argument("--n-layers", type=int, default=6) + ap.add_argument("--ple-dim", type=int, default=128) + ap.add_argument("--target-core", type=int, default=560000) + ap.add_argument("--batch-size", type=int, default=16) + ap.add_argument("--seq-len", type=int, default=256) + ap.add_argument("--steps", type=int, default=5000) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--tag", default="cleandeploy") + ap.add_argument("--skip-data", action="store_true") + ap.add_argument("--skip-train", action="store_true") + ap.add_argument("--skip-export", action="store_true") + ap.add_argument("--skip-assets", action="store_true") + ap.add_argument("--skip-verify", action="store_true") + ap.add_argument("--force-train", action="store_true") + args = ap.parse_args() + + RUNS.mkdir(exist_ok=True) + FIRMWARE_MODEL.mkdir(parents=True, exist_ok=True) + + if not args.skip_data and not data_artifacts_exist(args.vocab): + run_step( + "prepare data", + [ + sys.executable, + str(ROOT / "data" / "prepare.py"), + "--vocab", + str(args.vocab), + ], + ) + elif not args.skip_data: + print("\n==> prepare data: skipped (artifacts already exist)", flush=True) + + ckpt = checkpoint_path(args.arm, args.tag, args.seed) + if not args.skip_train: + if ckpt.exists() and not args.force_train: + print(f"\n==> train: skipped (found {ckpt.relative_to(ROOT)})", flush=True) + else: + run_step( + "train deploy checkpoint", + [ + sys.executable, + str(SRC / "train.py"), + "--arm", + args.arm, + "--vocab", + str(args.vocab), + "--d-model", + str(args.d_model), + "--n-layers", + str(args.n_layers), + "--ple-dim", + str(args.ple_dim), + "--target-core", + str(args.target_core), + "--batch-size", + str(args.batch_size), + "--seq-len", + str(args.seq_len), + "--steps", + str(args.steps), + "--seed", + str(args.seed), + "--tag", + args.tag, + ], + ) + + if not ckpt.exists(): + raise FileNotFoundError(f"expected checkpoint not found: {ckpt}") + + if not args.skip_export: + run_step("export model artifact", [sys.executable, str(SRC / "export.py"), str(ckpt)]) + + if not args.skip_assets: + run_step("generate firmware vocab assets", [sys.executable, str(SRC / "gen_assets.py")]) + + if not args.skip_verify: + verify_bin = Path("/tmp/esp32-llm-verify") + run_step( + "build host verifier", + ["cc", "-O3", "-o", str(verify_bin), str(FIRMWARE / "host_verify" / "verify.c"), "-lm"], + ) + run_step( + "verify exported model", + [str(verify_bin), str(FIRMWARE_MODEL / "model.bin"), str(FIRMWARE_MODEL / "golden.txt")], + ) + + print("\nDeploy flow complete.") + + +if __name__ == "__main__": + raise SystemExit(main()) \ No newline at end of file diff --git a/firmware/esp32_llm/README.md b/firmware/esp32_llm/README.md index 42024ba..0299306 100644 --- a/firmware/esp32_llm/README.md +++ b/firmware/esp32_llm/README.md @@ -6,7 +6,13 @@ embedding/output head is staged in PSRAM at boot. ## Build and verify -Export the group-128 ragged-int4 model and verify the portable C runtime first: +If you want the full sequence (prepare/train/export/compile/flash) in one +command, run `uv run python flash.py --full-pipeline --force-train` from the +repo root. + +Export the group-128 ragged-int4 model and verify the portable C runtime first. +This repo does not ship `runs/*.pt` checkpoints or the generated `firmware/model/model.bin`, so +you need a trained checkpoint in `runs/` before `src/export.py` can produce the artifact: ```bash cd src @@ -28,24 +34,58 @@ arduino-cli compile \ ## Flash and run -Replace the port if the board enumerates under a different device name: +On Linux the board usually shows up as `/dev/ttyUSB0`; replace the port if your +system uses a different device name: + +If you already have a checkpoint/model and just want compile+flash with USB +auto-scan, run `uv run python flash.py` from the repo root. ```bash arduino-cli upload \ - -p /dev/cu.usbmodem2101 \ + -p /dev/ttyUSB0 \ --fqbn 'esp32:esp32:esp32s3:UploadSpeed=921600,USBMode=hwcdc,CDCOnBoot=cdc,UploadMode=default,CPUFreq=240,FlashMode=qio,FlashSize=16M,PartitionScheme=custom,PSRAM=opi,DebugLevel=info' \ --input-dir /tmp/esp32-llm-build \ firmware/esp32_llm -esptool.py --chip esp32s3 --port /dev/cu.usbmodem2101 --baud 921600 \ +esptool.py --chip esp32s3 --port /dev/ttyUSB0 --baud 921600 \ write_flash 0x110000 firmware/model/model.bin -arduino-cli monitor -p /dev/cu.usbmodem2101 --config baudrate=115200 +arduino-cli monitor -p /dev/ttyUSB0 --config baudrate=115200 ``` The model payload only needs reflashing after a new export. Firmware-only changes can be uploaded without rewriting the model partition. +## Current Working Story Mode (2026-07) + +The current working firmware stage runs in autonomous storytelling mode: + +1. No serial command UI is required. +2. Story text streams continuously on serial output. +3. Context rollover is automatic (keeps generating after internal reset). +4. Startup seed is randomized from a small prompt bank for more variety. +5. Throughput/telemetry lines are suppressed so serial output is story text only. + +Typical run command: + +```bash +uv run python flash.py --skip-model +stty -F /dev/ttyACM0 115200 raw -echo +cat /dev/ttyACM0 +``` + +Example serial output snippet: + +```text +the cat and wanted to go outside. +"Look, I found something big!" the little cat hopped towards it and found some flowers. +Then, a mean dog came running down the stairs to the park. +The dog jumped on a leaf and chased it with its nose. +``` + +Note: if you use `timeout ... cat`, exit code `124` is expected when timeout +expires; it does not indicate a firmware failure. + The model used for the measurements below has SHA-256: ```text diff --git a/firmware/esp32_llm/esp32_llm.ino b/firmware/esp32_llm/esp32_llm.ino index 8cc30b6..b80afd6 100644 --- a/firmware/esp32_llm/esp32_llm.ino +++ b/firmware/esp32_llm/esp32_llm.ino @@ -7,6 +7,9 @@ #include "esp_partition.h" #include "esp_heap_caps.h" #include "esp_timer.h" +#include +#include +#include #define LLM_PROFILE 1 #define LLM_PROFILE_NOW() esp_timer_get_time() #include "../common/llm.h" @@ -14,13 +17,29 @@ // Set to 1 once a GMT020-02-7P (2.0" 240x320 ST7789) is wired up — see display.h. // Leave 0 to run serial-only (no panel needed). -#define USE_DISPLAY 1 +#define USE_DISPLAY 0 #if USE_DISPLAY #include "display.h" #endif static const int PROMPT_IDS[] = {433, 447, 259, 405}; // "Once upon a time" static const int N_GENERATE = 200; +static const int DEFAULT_MORE = 32; +static const int MAX_PROMPT_IDS = 64; +static const int MAX_PROMPT_BANK = 8; +static const int SAMPLE_TOP_K = 32; +static const float SAMPLE_TEMP = 0.95f; +static const float REPEAT_PENALTY = 0.28f; +static const int RECENT_WINDOW = 96; + +static const int PROMPT_IDS_0[] = {433, 447, 259, 405}; +static const int PROMPT_IDS_1[] = {2322, 259, 1759, 1814}; +static const int PROMPT_IDS_2[] = {10938, 557, 317, 263, 1078}; +static const int PROMPT_IDS_3[] = {345, 259, 1122, 1529}; +static const int PROMPT_IDS_4[] = {2614, 282, 259, 918, 1148}; +static const int PROMPT_IDS_5[] = {1570, 263, 3051, 348, 263, 1944}; +static const int PROMPT_IDS_6[] = {44, 537, 5725, 12, 317, 259, 2397}; +static const int PROMPT_IDS_7[] = {427, 7361, 3699}; // Emit one token to every active output (serial always; TFT when enabled). static void emit(int tok) { @@ -38,6 +57,22 @@ static void emit(int tok) { Model model; Scratch s; +static bool model_ready = false; +static int cur_pos = 0; +static int cur_tok = 0; +static int decoded_total = 0; +static int prompt_bank[MAX_PROMPT_BANK][MAX_PROMPT_IDS]; +static int prompt_bank_lens[MAX_PROMPT_BANK]; +static int prompt_bank_count = 0; +static int prompt_bank_idx = 0; +static uint32_t prng_state = 1; +static int64_t window_start_us = 0; +static int64_t window_decode_us = 0; +static int window_tokens = 0; +static int recent_ids[RECENT_WINDOW]; +static int recent_len = 0; +static int recent_head = 0; +static uint8_t recent_counts[VOCAB_N]; // ---- int8 output head (SIMD-friendly) -------------------------------------- // The head is scanned in full every token and dominates runtime. We stage it as @@ -122,6 +157,441 @@ static void blink(uint8_t g) { #endif } +static uint32_t prng_next() { + prng_state = prng_state * 1664525u + 1013904223u; + return prng_state; +} + +static float prng_unit() { + // 24-bit mantissa-style fraction in [0,1). + return (float)((prng_next() >> 8) & 0x00FFFFFFu) / 16777216.0f; +} + +static void recent_clear() { + memset(recent_counts, 0, sizeof(recent_counts)); + recent_len = 0; + recent_head = 0; +} + +static void recent_note(int tok) { + if (tok < 0 || tok >= VOCAB_N) return; + + if (recent_len == RECENT_WINDOW) { + int old = recent_ids[recent_head]; + if (old >= 0 && old < VOCAB_N && recent_counts[old] > 0) recent_counts[old]--; + recent_ids[recent_head] = tok; + recent_head = (recent_head + 1) % RECENT_WINDOW; + } else { + int idx = (recent_head + recent_len) % RECENT_WINDOW; + recent_ids[idx] = tok; + recent_len++; + } + + if (recent_counts[tok] < 255) recent_counts[tok]++; +} + +static int sample_topk_token(const float *logits, int vocab, int k) { + if (k < 1) k = 1; + if (k > 16) k = 16; + + int top_ids[16] = {0}; + float top_vals[16]; + for (int i = 0; i < 16; i++) top_vals[i] = -1e30f; + + int n = 0; + for (int v = 0; v < vocab; v++) { + float val = logits[v]; + if (val <= top_vals[k - 1]) continue; + + int pos = k - 1; + while (pos > 0 && val > top_vals[pos - 1]) { + top_vals[pos] = top_vals[pos - 1]; + top_ids[pos] = top_ids[pos - 1]; + pos--; + } + top_vals[pos] = val; + top_ids[pos] = v; + if (n < k) n++; + } + + if (n <= 0) return 0; + return top_ids[prng_next() % (uint32_t)n]; +} + +static int sample_next_token(const float *logits, int vocab, int k, float temp) { + if (k < 1) k = 1; + if (k > 64) k = 64; + if (temp < 0.1f) temp = 0.1f; + + int top_ids[64] = {0}; + float top_vals[64]; + for (int i = 0; i < 64; i++) top_vals[i] = -1e30f; + + int n = 0; + for (int v = 0; v < vocab; v++) { + float val = logits[v] - REPEAT_PENALTY * (float)recent_counts[v]; + if (val <= top_vals[k - 1]) continue; + + int pos = k - 1; + while (pos > 0 && val > top_vals[pos - 1]) { + top_vals[pos] = top_vals[pos - 1]; + top_ids[pos] = top_ids[pos - 1]; + pos--; + } + top_vals[pos] = val; + top_ids[pos] = v; + if (n < k) n++; + } + + if (n <= 0) return 0; + if (n == 1) return top_ids[0]; + + float maxv = top_vals[0]; + float probs[64]; + float s = 0.0f; + for (int i = 0; i < n; i++) { + float z = (top_vals[i] - maxv) / temp; + float p = expf(z); + probs[i] = p; + s += p; + } + + float r = prng_unit() * s; + float c = 0.0f; + for (int i = 0; i < n; i++) { + c += probs[i]; + if (r <= c) return top_ids[i]; + } + return top_ids[n - 1]; +} + +static int next_token_greedy() { + int best = 0; + float bv = -1e30f; + for (int v = 0; v < VOCAB_N; v++) { + if (s.logits[v] > bv) { + bv = s.logits[v]; + best = v; + } + } + return best; +} + +// Greedy decode for n tokens, continuing from current context. +static void generate_more(int n) { + if (!model_ready) { + Serial.println("model not ready"); + return; + } + if (cur_pos >= model.c.seq_len) { + Serial.printf("context full (%d/%d). Use 'reset' or 'prompt_ids ...'\n", cur_pos, model.c.seq_len); + return; + } + if (n <= 0) n = DEFAULT_MORE; + + int can = model.c.seq_len - cur_pos; + if (n > can) n = can; + + int64_t t0 = esp_timer_get_time(); + int64_t decode_us = 0; + int decoded = 0; + + for (int step = 0; step < n; step++) { + int best = 0; + float bv = -1e30f; + for (int v = 0; v < VOCAB_N; v++) { + if (s.logits[v] > bv) { + bv = s.logits[v]; + best = v; + } + } + + cur_tok = best; + emit(cur_tok); + blink((step & 1) ? 40 : 8); + + int64_t d0 = esp_timer_get_time(); + llm_forward(&model, cur_tok, cur_pos++, &s); + decode_us += esp_timer_get_time() - d0; + decoded++; + decoded_total++; + if ((step & 7) == 0) delay(0); + } + + int64_t total_us = esp_timer_get_time() - t0; + Serial.printf("\n\n--- +%d tokens (total %d, pos %d/%d) in %.2f s ---\n", + decoded, decoded_total, cur_pos, model.c.seq_len, total_us / 1e6); + if (decoded > 0) { + Serial.printf("throughput: %.2f tok/s (%.1f ms/token)\n", + decoded * 1e6 / total_us, decode_us / 1000.0 / decoded); + } +} + +static void clear_runtime_state() { + recent_clear(); + memset(s.kcache, 0, (size_t)model.c.n_layers * model.c.seq_len * model.c.dim * sizeof(float)); + memset(s.vcache, 0, (size_t)model.c.n_layers * model.c.seq_len * model.c.dim * sizeof(float)); + memset(s.ple, 0, (size_t)model.c.n_layers * model.c.ple_dim * sizeof(float)); + memset(s.tmpP, 0, (size_t)model.c.n_layers * model.c.ple_dim * sizeof(float)); + memset(s.trow, 0, (size_t)model.c.n_layers * model.c.ple_dim * sizeof(float)); + cur_pos = 0; + cur_tok = 0; + decoded_total = 0; +} + +static void clear_cache_state() { + memset(s.kcache, 0, (size_t)model.c.n_layers * model.c.seq_len * model.c.dim * sizeof(float)); + memset(s.vcache, 0, (size_t)model.c.n_layers * model.c.seq_len * model.c.dim * sizeof(float)); + memset(s.ple, 0, (size_t)model.c.n_layers * model.c.ple_dim * sizeof(float)); + memset(s.tmpP, 0, (size_t)model.c.n_layers * model.c.ple_dim * sizeof(float)); + memset(s.trow, 0, (size_t)model.c.n_layers * model.c.ple_dim * sizeof(float)); + cur_pos = 0; +} + +static void set_prompt_slot(int slot, const int *ids, int n_ids) { + if (slot < 0 || slot >= MAX_PROMPT_BANK) return; + if (n_ids < 1) return; + if (n_ids > MAX_PROMPT_IDS) n_ids = MAX_PROMPT_IDS; + for (int i = 0; i < n_ids; i++) { + prompt_bank[slot][i] = ids[i]; + } + prompt_bank_lens[slot] = n_ids; +} + +static void init_prompt_bank() { + set_prompt_slot(0, PROMPT_IDS_0, sizeof(PROMPT_IDS_0) / sizeof(int)); + set_prompt_slot(1, PROMPT_IDS_1, sizeof(PROMPT_IDS_1) / sizeof(int)); + set_prompt_slot(2, PROMPT_IDS_2, sizeof(PROMPT_IDS_2) / sizeof(int)); + set_prompt_slot(3, PROMPT_IDS_3, sizeof(PROMPT_IDS_3) / sizeof(int)); + set_prompt_slot(4, PROMPT_IDS_4, sizeof(PROMPT_IDS_4) / sizeof(int)); + set_prompt_slot(5, PROMPT_IDS_5, sizeof(PROMPT_IDS_5) / sizeof(int)); + set_prompt_slot(6, PROMPT_IDS_6, sizeof(PROMPT_IDS_6) / sizeof(int)); + set_prompt_slot(7, PROMPT_IDS_7, sizeof(PROMPT_IDS_7) / sizeof(int)); + prompt_bank_count = 8; + prompt_bank_idx = 0; +} + +static void start_prompt_ids(const int *ids, int n_ids) { + if (!model_ready) { + Serial.println("model not ready"); + return; + } + if (n_ids <= 0) { + Serial.println("no prompt ids provided"); + return; + } + + clear_runtime_state(); + Serial.print("\n>>> "); + for (int i = 0; i < n_ids && cur_pos < model.c.seq_len; i++) { + int id = ids[i]; + if (id < 0 || id >= VOCAB_N) { + Serial.printf("\ninvalid token id %d (skipped)\n", id); + continue; + } + cur_tok = id; + emit(cur_tok); + llm_forward(&model, cur_tok, cur_pos++, &s); + recent_note(cur_tok); + } +} + +static void start_active_prompt() { + if (prompt_bank_count <= 0) { + init_prompt_bank(); + } + start_prompt_ids(prompt_bank[prompt_bank_idx], prompt_bank_lens[prompt_bank_idx]); +} + +static void print_prompt_list() { + Serial.printf("prompt bank: %d/%d entries, active=%d\n", + prompt_bank_count, MAX_PROMPT_BANK, prompt_bank_idx); + for (int i = 0; i < prompt_bank_count; i++) { + Serial.printf(" [%d]%s len=%d\n", i, (i == prompt_bank_idx ? "*" : " "), prompt_bank_lens[i]); + } +} + +static void start_default_prompt() { + prompt_bank_idx = 0; + start_active_prompt(); +} + +static void prime_story_seed() { + if (prompt_bank_count > 0) { + prompt_bank_idx = (int)(prng_next() % (uint32_t)prompt_bank_count); + } + start_active_prompt(); + int extra = 1 + (int)(prng_next() % 4u); + for (int i = 0; i < extra && cur_pos < model.c.seq_len; i++) { + int tok = sample_next_token(s.logits, VOCAB_N, 24, 1.05f); + cur_tok = tok; + emit(cur_tok); + llm_forward(&model, cur_tok, cur_pos++, &s); + recent_note(cur_tok); + } +} + +static void print_serial_help() { + Serial.println("\nserial commands:"); + Serial.println(" generate 32 more tokens (auto-resets if context is full)"); + Serial.println(" more [N] generate N more tokens"); + Serial.println(" reset reset to default prompt"); + Serial.println(" prompt_ids a,b,c reset and use token-id prompt"); + Serial.println(" prompt_add a,b,c save a prompt slot"); + Serial.println(" next_prompt rotate to next saved prompt and generate"); + Serial.println(" prompt_list show saved prompt slots"); + Serial.println(" status show token/context position"); + Serial.println(" help show this help"); +} + +static void print_cli_prompt() { + Serial.print("cmd> "); +} + +static char ascii_lower(char c) { + if (c >= 'A' && c <= 'Z') return (char)(c - 'A' + 'a'); + return c; +} + +static bool equals_ci(const char *a, const char *b) { + while (*a && *b) { + if (ascii_lower(*a) != ascii_lower(*b)) return false; + a++; b++; + } + return *a == '\0' && *b == '\0'; +} + +static bool starts_with_ci(const char *s, const char *prefix) { + while (*prefix) { + if (*s == '\0') return false; + if (ascii_lower(*s) != ascii_lower(*prefix)) return false; + s++; prefix++; + } + return true; +} + +static void handle_serial_command(const char *line_in) { + char line[192]; + strncpy(line, line_in, sizeof(line) - 1); + line[sizeof(line) - 1] = '\0'; + + // Trim leading spaces. + char *p = line; + while (*p == ' ' || *p == '\t' || *p == '\r') p++; + + // Trim trailing spaces. + size_t plen = strlen(p); + while (plen > 0 && (p[plen - 1] == ' ' || p[plen - 1] == '\t' || p[plen - 1] == '\r')) { + p[--plen] = '\0'; + } + + if (*p == '\0') { + if (cur_pos >= model.c.seq_len) { + Serial.println("context full; restarting from default prompt"); + start_default_prompt(); + } + generate_more(DEFAULT_MORE); + return; + } + + if (equals_ci(p, "help")) { + print_serial_help(); + return; + } + + if (equals_ci(p, "reset")) { + start_active_prompt(); + generate_more(DEFAULT_MORE); + return; + } + + if (equals_ci(p, "status")) { + Serial.printf("status: pos %d/%d, generated %d tokens\n", cur_pos, model.c.seq_len, decoded_total); + return; + } + + if (starts_with_ci(p, "more")) { + int n = DEFAULT_MORE; + if (p[4] != '\0') n = atoi(p + 4); + if (cur_pos >= model.c.seq_len) { + Serial.println("context full; restarting from default prompt"); + start_default_prompt(); + } + generate_more(n); + return; + } + + if (starts_with_ci(p, "prompt_ids")) { + char *args = p + 10; + while (*args == ' ' || *args == '\t') args++; + if (*args == '\0') { + Serial.println("usage: prompt_ids 433,447,259,405"); + return; + } + int ids[MAX_PROMPT_IDS]; + int n_ids = 0; + + char *tok = strtok(args, ", "); + while (tok && n_ids < MAX_PROMPT_IDS) { + ids[n_ids++] = atoi(tok); + tok = strtok(NULL, ", "); + } + set_prompt_slot(prompt_bank_idx, ids, n_ids); + start_prompt_ids(ids, n_ids); + generate_more(DEFAULT_MORE); + return; + } + + if (starts_with_ci(p, "prompt_add")) { + char *args = p + 10; + while (*args == ' ' || *args == '\t') args++; + if (*args == '\0') { + Serial.println("usage: prompt_add 433,447,259,405"); + return; + } + if (prompt_bank_count >= MAX_PROMPT_BANK) { + Serial.printf("prompt bank full (%d). overwrite with prompt_ids or reboot.\n", MAX_PROMPT_BANK); + return; + } + int ids[MAX_PROMPT_IDS]; + int n_ids = 0; + char *tok = strtok(args, ", "); + while (tok && n_ids < MAX_PROMPT_IDS) { + ids[n_ids++] = atoi(tok); + tok = strtok(NULL, ", "); + } + set_prompt_slot(prompt_bank_count, ids, n_ids); + prompt_bank_idx = prompt_bank_count; + prompt_bank_count++; + Serial.printf("saved prompt slot %d\n", prompt_bank_idx); + start_active_prompt(); + generate_more(DEFAULT_MORE); + return; + } + + if (equals_ci(p, "next_prompt")) { + if (prompt_bank_count <= 0) { + init_prompt_bank(); + } + prompt_bank_idx = (prompt_bank_idx + 1) % prompt_bank_count; + Serial.printf("switched to prompt slot %d\n", prompt_bank_idx); + start_active_prompt(); + generate_more(DEFAULT_MORE); + return; + } + + if (equals_ci(p, "prompt_list")) { + print_prompt_list(); + return; + } + + if (strchr(p, ' ') || strchr(p, '"') || strchr(p, '\'')) { + Serial.println("plain text prompts are not supported on-device; use 'prompt_ids ...' or 'reset'."); + } else { + Serial.println("unknown command. type 'help'"); + } +} + void setup() { Serial.begin(115200); delay(1500); @@ -177,43 +647,16 @@ void setup() { Serial.printf("PSRAM free after alloc: %u KB\n\n", heap_caps_get_free_size(MALLOC_CAP_SPIRAM) / 1024); - // ---- generate ---- - Serial.print(">>> "); - int n_prompt = sizeof(PROMPT_IDS) / sizeof(int); - int pos = 0, tok = 0; - int64_t t_start = 0; - int64_t decode_us = 0; - int decoded = 0; - - for (int i = 0; i < n_prompt; i++) { // prime with the prompt - tok = PROMPT_IDS[i]; - emit(tok); - llm_forward(&model, tok, pos++, &s); - } - + model_ready = true; + init_prompt_bank(); + prng_state = esp_random(); llm_profile_reset(&s); + prime_story_seed(); - t_start = esp_timer_get_time(); - for (int step = 0; step < N_GENERATE && pos < model.c.seq_len; step++) { - // greedy: argmax over the trained vocab - int best = 0; float bv = -1e30f; - for (int v = 0; v < VOCAB_N; v++) - if (s.logits[v] > bv) { bv = s.logits[v]; best = v; } - tok = best; - emit(tok); - blink((step & 1) ? 40 : 8); - - int64_t d0 = esp_timer_get_time(); - llm_forward(&model, tok, pos++, &s); - decode_us += esp_timer_get_time() - d0; - decoded++; - if ((step & 7) == 0) delay(0); // feed the task WDT ~every 8 tokens (~1.1s), near-free - } - int64_t total_us = esp_timer_get_time() - t_start; + window_start_us = esp_timer_get_time(); + window_decode_us = 0; + window_tokens = 0; - Serial.printf("\n\n--- %d tokens in %.2f s ---\n", decoded, total_us / 1e6); - Serial.printf("throughput: %.2f tok/s (%.1f ms/token)\n", - decoded * 1e6 / total_us, decode_us / 1000.0 / decoded); if (s.profile.calls) { float n = (float)s.profile.calls * 1000.f; Serial.printf("profile ms/token: input %.1f | attn %.1f | ffn %.1f | ple %.1f | head %.1f\n", @@ -223,9 +666,49 @@ void setup() { } #if USE_DISPLAY // Closing card: compute-only tok/s (the model's own speed) + ms/token. - display_stats(decoded * 1e6f / decode_us, decode_us / 1000.0f / decoded); + if (decoded_total > 0) { + display_stats(0, 0); + } #endif blink(0); } -void loop() { delay(10000); } +void loop() { + if (!model_ready) { + delay(50); + return; + } + + if (cur_pos >= model.c.seq_len) { + clear_cache_state(); + llm_forward(&model, cur_tok, cur_pos++, &s); + recent_note(cur_tok); + } + + cur_tok = sample_next_token(s.logits, VOCAB_N, SAMPLE_TOP_K, SAMPLE_TEMP); + // Prevent pathological immediate repeats such as "slide slide slide". + if (cur_pos > 0 && recent_len > 0) { + int last_idx = (recent_head + recent_len - 1) % RECENT_WINDOW; + int last_tok = recent_ids[last_idx]; + if (cur_tok == last_tok) { + cur_tok = sample_next_token(s.logits, VOCAB_N, SAMPLE_TOP_K, SAMPLE_TEMP + 0.1f); + } + } + emit(cur_tok); + blink((decoded_total & 1) ? 40 : 8); + + int64_t d0 = esp_timer_get_time(); + llm_forward(&model, cur_tok, cur_pos++, &s); + recent_note(cur_tok); + window_decode_us += esp_timer_get_time() - d0; + decoded_total++; + window_tokens++; + + if (window_tokens >= 32) { + window_start_us = esp_timer_get_time(); + window_decode_us = 0; + window_tokens = 0; + } + + if ((decoded_total & 7) == 0) delay(0); +} diff --git a/flash.py b/flash.py new file mode 100644 index 0000000..5e0ba57 --- /dev/null +++ b/flash.py @@ -0,0 +1,231 @@ +"""Build and flash the ESP32-S3 sketch and model partition. + +This script looks for the board on USB serial, preferring stable +`/dev/serial/by-id/*` entries and falling back to `/dev/ttyUSB*` and +`/dev/ttyACM*`. + +It can run the full pipeline in one command: + 1. optional deploy flow (prepare/train/export/assets/verify) + 2. compile sketch + 3. upload sketch + 4. flash model partition + +By default it uses a repo-local standalone Arduino CLI at +`.tools/arduino-cli/arduino-cli` when present, and falls back to `arduino-cli` +from PATH. +""" + +from __future__ import annotations + +import argparse +import glob +import os +import subprocess +import sys +import shutil +from pathlib import Path + + +ROOT = Path(__file__).resolve().parent +FIRMWARE = ROOT / "firmware" / "esp32_llm" +FIRMWARE_COMMON = ROOT / "firmware" / "common" +MODEL_BIN = ROOT / "firmware" / "model" / "model.bin" +BUILD_PATH = Path("/tmp/esp32-llm-build") +LOCAL_ARDUINO_CLI = ROOT / ".tools" / "arduino-cli" / "arduino-cli" +FQBN = ( + "esp32:esp32:esp32s3:" + "UploadSpeed=921600,USBMode=hwcdc,CDCOnBoot=cdc,UploadMode=default," + "CPUFreq=240,FlashMode=qio,FlashSize=16M,PartitionScheme=custom," + "PSRAM=opi,DebugLevel=info" +) + + +def run_step(title: str, argv: list[str]) -> None: + print(f"\n==> {title}", flush=True) + print(" " + " ".join(argv), flush=True) + subprocess.run(argv, cwd=ROOT, check=True) + + +def resolve_arduino_cli(explicit: str | None) -> str: + if explicit: + return explicit + if LOCAL_ARDUINO_CLI.exists(): + return str(LOCAL_ARDUINO_CLI) + path = shutil.which("arduino-cli") + if path: + return path + raise FileNotFoundError( + "arduino-cli not found. Install it or pass --arduino-cli /path/to/arduino-cli" + ) + + +def resolve_esptool() -> str: + for name in ("esptool.py", "esptool"): + path = shutil.which(name) + if path: + return path + raise FileNotFoundError( + "esptool not found in PATH. Install it (for example: `uv add esptool`) " + "or provide it on PATH as `esptool.py` or `esptool`." + ) + + +def find_usb_ports() -> list[str]: + ports: list[str] = [] + + for path in sorted(glob.glob("/dev/serial/by-id/*")): + if os.path.exists(path): + ports.append(os.path.realpath(path)) + + for pattern in ("/dev/ttyUSB*", "/dev/ttyACM*"): + for path in sorted(glob.glob(pattern)): + if path not in ports: + ports.append(path) + + return ports + + +def choose_port(explicit: str | None) -> str: + if explicit: + return explicit + + ports = find_usb_ports() + if not ports: + raise FileNotFoundError( + "no USB serial port found; pass --port /dev/ttyUSB0 or check the board connection" + ) + if len(ports) == 1: + print(f"selected USB port: {ports[0]}") + return ports[0] + + print("found multiple USB serial ports:") + for port in ports: + print(f" {port}") + print(f"using {ports[0]}") + return ports[0] + + +def run_deploy(python_exe: str, deploy_args: list[str], force_train: bool) -> None: + argv = [python_exe, str(ROOT / "deploy.py")] + argv.extend(deploy_args) + if force_train and "--force-train" not in deploy_args: + argv.append("--force-train") + run_step("full deploy flow", argv) + + +def compile_sketch(arduino_cli: str, build_path: Path) -> None: + run_step( + "compile sketch", + [ + arduino_cli, + "compile", + "--fqbn", + FQBN, + "--build-property", + "compiler.optimization_flags=-O3", + "--build-property", + f"compiler.cpp.extra_flags=-I{FIRMWARE_COMMON}", + "--build-path", + str(build_path), + str(FIRMWARE), + ], + ) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--full-pipeline", + action="store_true", + help="run deploy.py first (prepare/train/export/assets/verify), then compile+flash", + ) + ap.add_argument( + "--deploy-arg", + action="append", + default=[], + help="extra argument forwarded to deploy.py (repeatable)", + ) + ap.add_argument( + "--force-train", + action="store_true", + help="when using --full-pipeline, force a new training run", + ) + ap.add_argument("--port", help="USB serial port; auto-detected if omitted") + ap.add_argument("--build-path", default=str(BUILD_PATH)) + ap.add_argument("--arduino-cli", help="path to arduino-cli binary") + ap.add_argument("--python-exe", default=sys.executable, help="python executable for deploy.py") + ap.add_argument("--baud", type=int, default=921600) + ap.add_argument("--skip-compile", action="store_true", help="skip arduino-cli compile") + ap.add_argument("--skip-sketch", action="store_true", help="skip arduino-cli upload") + ap.add_argument("--skip-model", action="store_true", help="skip model partition flash") + ap.add_argument("--monitor", action="store_true", help="start arduino-cli monitor after flashing") + args = ap.parse_args() + + build_path = Path(args.build_path) + arduino_cli = resolve_arduino_cli(args.arduino_cli) + esptool = resolve_esptool() + + if args.full_pipeline: + run_deploy(args.python_exe, args.deploy_arg, args.force_train) + + if not args.skip_sketch and not args.skip_compile: + compile_sketch(arduino_cli, build_path) + + if not args.skip_sketch: + if not build_path.exists(): + raise FileNotFoundError( + f"build path not found: {build_path}. Run compile first or remove --skip-compile." + ) + + port: str | None = None + if not args.skip_sketch or not args.skip_model or args.monitor: + port = choose_port(args.port) + + if not args.skip_sketch: + run_step( + "upload sketch", + [ + arduino_cli, + "upload", + "-p", + port, + "--fqbn", + FQBN, + "--input-dir", + str(build_path), + str(FIRMWARE), + ], + ) + + if not args.skip_model: + if not MODEL_BIN.exists(): + raise FileNotFoundError( + f"missing model artifact: {MODEL_BIN}. Run `uv run python deploy.py` or `uv run python src/export.py` first." + ) + run_step( + "flash model partition", + [ + esptool, + "--chip", + "esp32s3", + "--port", + port, + "--baud", + str(args.baud), + "write_flash", + "0x110000", + str(MODEL_BIN), + ], + ) + + if args.monitor: + run_step( + "serial monitor", + [arduino_cli, "monitor", "-p", port, "--config", "baudrate=115200"], + ) + + print("\nFlash flow complete.") + + +if __name__ == "__main__": + raise SystemExit(main()) \ No newline at end of file