From 7f7660da25aa5eb0db41944af161228aef02f312 Mon Sep 17 00:00:00 2001 From: Dineth Hettiarachchi Date: Sun, 25 Oct 2026 14:42:43 +0530 Subject: [PATCH] feat: Phase 6 WebAssembly harness and golden fixtures TinyAttnSim (one model instance with its own VerilatedContext, so instances run side by side), embind bindings, build script (81 KB gzipped), Node test over three golden fixtures (FLASH sequential, FLASH overlap causal, NAIVE causal) with interleaved instances, exact O, bytes, tiles, K/V loads and cycles. State schema documented; fixture freshness pinned in pytest. --- .github/workflows/wasm.yml | 13 ++ .gitignore | 1 - CONTEXT.md | 4 +- README.md | 5 +- docs/ARCHITECTURE.md | 16 +++ docs/PLAN.md | 2 +- docs/STATE_SCHEMA.md | 72 +++++++++++ sim/golden/export_fixtures.py | 88 ++++++++++++++ sim/golden/test_golden.py | 10 ++ wasm/README.md | 14 ++- wasm/build.sh | 43 +++++++ wasm/fixtures/flash_overlap_causal.json | 1 + wasm/fixtures/flash_seq.json | 1 + wasm/fixtures/naive_causal.json | 1 + wasm/harness/bindings.cpp | 54 +++++++++ wasm/harness/harness.cpp | 155 ++++++++++++++++++++++++ wasm/harness/harness.h | 60 +++++++++ wasm/test.mjs | 82 +++++++++++++ 18 files changed, 613 insertions(+), 9 deletions(-) create mode 100644 docs/STATE_SCHEMA.md create mode 100644 sim/golden/export_fixtures.py create mode 100755 wasm/build.sh create mode 100644 wasm/fixtures/flash_overlap_causal.json create mode 100644 wasm/fixtures/flash_seq.json create mode 100644 wasm/fixtures/naive_causal.json create mode 100644 wasm/harness/bindings.cpp create mode 100644 wasm/harness/harness.cpp create mode 100644 wasm/harness/harness.h create mode 100644 wasm/test.mjs diff --git a/.github/workflows/wasm.yml b/.github/workflows/wasm.yml index d9e24b8..3099e57 100644 --- a/.github/workflows/wasm.yml +++ b/.github/workflows/wasm.yml @@ -53,3 +53,16 @@ jobs: run: | test -f web/public/spike/counter.wasm echo "spike wasm: $(du -h web/public/spike/counter.wasm | cut -f1)" + + # The engine: Verilate tiny_attn_top, build the module, run the golden fixtures in Node, check the size budget. + - name: Build engine WASM + run: bash wasm/build.sh + + - name: Engine golden fixtures (Node) + run: node wasm/test.mjs + + - name: Size budget (500 KB gzipped) + run: | + GZ=$(gzip -9 -c web/public/tiny_attn.wasm | wc -c) + echo "engine wasm gzip: $GZ bytes" + test "$GZ" -le 512000 diff --git a/.gitignore b/.gitignore index 41526b2..c9d103e 100644 --- a/.gitignore +++ b/.gitignore @@ -27,7 +27,6 @@ web/public/tiny_attn.mjs web/public/spike/ *.o *.wasm -!web/public/tiny_attn.wasm # Toolchain checkouts (built by .tools/*.sh, or symlinked to a sibling project's) .tools/* diff --git a/CONTEXT.md b/CONTEXT.md index 423b34b..28ab821 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -31,7 +31,9 @@ On a fresh machine delete the symlinks and run `.tools/setup.sh` and `.tools/ems - Phase 5 (overlap, NAIVE, perf pins): DONE. `naive_ctrl` (4 passes, block-by-block), FLASH double-buffered K/V prefetch (`cfg_overlap`), tile operations OP_NZ / OP_NM / OP_NP / OP_NPV, HBM 2048 words. Full e2e (200 seeds) and unit suites pass; perf baseline in `sim/perf_baseline.json` (re-record with `PERF_RECORD=1 COCOTB_TEST_FILTER=perf_grid bash sim/run_e2e.sh`). -- Next: Phase 6 (WASM harness: dual models, bindings, state schema, Node test against the golden fixtures). +- Phase 6 (WASM harness): DONE. `wasm/build.sh` builds `web/public/tiny_attn.{mjs,wasm}` (81 KB gzipped); `node wasm/test.mjs` + checks three golden fixtures on interleaved instances, exactly. JS API in `docs/STATE_SCHEMA.md`. +- Next: Phase 7 (web UI). First: wide-signal access so the per-clock state carries per-PE values and softmax lanes. ## Gotchas found in Phase 5 - The NAIVE tile operations reuse the online-softmax unit: OP_NM uses its m path, OP_NP uses its P and l path with alpha = 1 diff --git a/README.md b/README.md index bfc4ee7..bb2902f 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ Fourth project in the series after [tiny-gpu](https://github.com/deaneeth/tiny-g [tiny-tpu](https://github.com/deaneeth/tiny-tpu) and tiny-cpu. Same rule as the others: nothing on screen is faked; every value comes from a `dbg_*` port on the compiled RTL. -**Status:** Phases 0-5 done: the FLASH engine, the NAIVE baseline and K/V overlap are bit-exact against the golden model, with byte counts, tile counts and cycle counts verified (FLASH sequential exactly against a closed-form model; overlap and NAIVE against pinned baselines). See [`docs/PLAN.md`](docs/PLAN.md) for the roadmap. +**Status:** Phases 0-6 done: the engine (FLASH, NAIVE, K/V overlap) is bit-exact against the golden model and compiled to WebAssembly (81 KB gzipped), with a Node test over golden fixtures. Next is the web UI. See [`docs/PLAN.md`](docs/PLAN.md) for the roadmap. ## Quick start @@ -17,7 +17,8 @@ bash .tools/setup.sh && bash .tools/emsdk.sh # Verilator 5.036 + Emscripten python3 -m venv .venv && .venv/bin/pip install -r sim/requirements.txt source scripts/env.sh -bash wasm/spike/build.sh && node wasm/spike/test.mjs # WASM spike +bash wasm/build.sh && node wasm/test.mjs # engine WASM + golden fixtures +bash wasm/spike/build.sh && node wasm/spike/test.mjs # toolchain spike cd web && pnpm install && pnpm dev # open /spike ``` diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 5d0222e..656149c 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -339,3 +339,19 @@ Deviations from the first plan, recorded here so they are not forgotten: - NAIVE runs block by block, not pass by pass over all blocks. Values and traffic totals are identical (the per-row computations are independent), but the sequencing differs from `fixed.naive`'s loop order. - Overlap applies to FLASH only. NAIVE is sequential. + +## 11. Phase 6: WebAssembly harness + +- `wasm/harness/harness.{h,cpp}`: `TinyAttnSim`, one instance of `tiny_attn_top` with its own `VerilatedContext`. Drives + the engine as a host would: load Q/K/V through the HBM load port, one-cycle start, read O through the readback port. +- `wasm/harness/bindings.cpp`: embind bindings. The JavaScript API and the per-clock snapshot are in `docs/STATE_SCHEMA.md`. +- `wasm/build.sh`: Verilator `--prefix Vtiny_attn_top`, `-O3`, then `em++`. `-Wno-UNOPTFLAT` is needed for the + combinational paths through the staging buffer and the DMA word mux (no logic loop: the loop is through registers). +- Size: 206 KB raw, 81 KB gzipped. The budget is 500 KB gzipped, enforced by `wasm.yml`. + +Results (`node wasm/test.mjs`, three golden fixtures interleaved on separate instances): O, HBM bytes, tiles, K/V loads and +busy cycles equal the golden values for FLASH sequential, FLASH with overlap and causal masking, and NAIVE with causal +masking. A fresh instance reproduces the run, so no state leaks between instances. + +Deviation from the plan: the per-clock state exported to JavaScript is a subset (controller state, counters, byte counts). +Per-PE values, softmax lanes and SRAM tiles need wide-signal access and are the first item of Phase 7. diff --git a/docs/PLAN.md b/docs/PLAN.md index e0bb212..ffabe8d 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -443,7 +443,7 @@ The estimates assume one person part-time, at tiny-cpu's pace. - **Exit:** the whole random matrix (§5 layer 3) is green; `perf_baseline.json` is pinned; the accuracy test is green. *At this point the engine is fully verified hardware and worth publishing on its own.* -### Phase 6: WASM harness (1 week) +### Phase 6: WASM harness (1 week) ✅ DONE - `harness.cpp`, `bindings.cpp`, `build.sh`, `STATE_SCHEMA.md`, `test.mjs` with golden fixtures. - **Exit:** `wasm.yml` is green; the size budget is met; two instances run independently in Node. diff --git a/docs/STATE_SCHEMA.md b/docs/STATE_SCHEMA.md new file mode 100644 index 0000000..e633d51 --- /dev/null +++ b/docs/STATE_SCHEMA.md @@ -0,0 +1,72 @@ +# TinyAttention: JavaScript state schema (Phase 6) + +What the WebAssembly module exposes per clock, and how to drive it. The schema mirrors the C++ structs in +`wasm/harness/harness.h`; if one changes, the other must change in the same commit. + +Phase 6 exports a **subset** of the planned per-cycle state. The full `CycleState` (per-PE values, softmax lanes, SRAM +tags, tile matrices) needs wide-signal access through Verilator and is the first item of Phase 7. The RTL already has +the ports for it (`dbg_*` in `tiny_attn_top` and `attn_tile`); only the bindings are missing. + +## Snapshot (returned by `step()`, `snapshot()`) + +| Field | Type | Meaning | +|---|---|---| +| `cycle` | number | clock edges since reset | +| `ctrlState` | number | `dbg_ctrl_state`: the active controller's state enum (`flash_ctrl` or `naive_ctrl`, see their source) | +| `busy` | boolean | controller not idle | +| `done` | boolean | one-cycle pulse: the DONE state | +| `hbmRdBytes` | number | DMA read bytes since start (4 per word) | +| `hbmWrBytes` | number | DMA write bytes since start | +| `cycles` | number | busy-cycle counter (state != idle, DONE included) | +| `tiles` | number | tile computations (FLASH: group tiles; NAIVE: pass-1 tile pairs) | +| `kvLoads` | number | K/V tile loads (FLASH: K loads incl. prefetch; NAIVE: one per tile pair) | + +**Counters are final one edge after `done`.** `done` is high during the DONE cycle, and the busy counter increments at the +end of that cycle. `wasm/test.mjs` takes one more `step()` before reading. + +## Methods + +```ts +class TinyAttnSim { + reset(): void; + configure(n: number, qrows: number, causal: boolean, mode: 0 | 1, overlap: boolean, + latency: number, scale: number): void; // mode 0 = FLASH, 1 = NAIVE; scale < 2^24 + loadQKV(q: number[], k: number[], v: number[]): void; // n*4 int8 values each, row-major; n*3 cycles + start(): void; // one-cycle start pulse; the next step() is the first cycle of the run + step(): Snapshot; // one clock + runToDone(maxCycles: number): number; // edges taken until done, -1 on timeout + readOutput(): number[]; // n*4 int16 values (O, Q8.8 V-LSB), row-major + readHbm(addr: number, len: number): number[]; // raw 32-bit words (Z/P tiles are at 256 and 768) + snapshot(): Snapshot; + isDone(): boolean; + delete(): void; +} +``` + +Input validation throws: `n` must be a multiple of 4 in 4..32, `qrows` one of 4, 8, 16, 32 and not above `n`, latency 0..64, +`scale` below 2^24. + +## Run sequence + +1. `new mod.TinyAttnSim()`, then `configure(...)`. +2. `loadQKV(q, k, v)`: quantized int8 values from `sim/golden/fixed.py` (`quantize`), and `scale = cfgScale`. +3. `start()`, then `step()` until `done`, or `runToDone(max)`. +4. `step()` once more, then `readOutput()`. + +Multiple instances are independent: each owns its Verilator context. Steps may be interleaved across instances. + +## Address map of the HBM model (words) + +| Region | Base | Size | +|---|---|---| +| Q | 0 | 32 rows (one word each) | +| K | 32 | 32 rows | +| V | 64 | 32 rows | +| O | 96 | 64 words (two per row) | +| Z tiles (NAIVE) | 256 | 8 words per tile, tile (qi, kj) at `256 + 64·qi + 8·kj` | +| P tiles (NAIVE) | 768 | same layout | + +## Fixtures + +`wasm/fixtures/*.json` are exported by `sim/golden/export_fixtures.py` and checked for freshness by `pytest sim/golden`. +Each holds the quantized inputs, the expected O, bytes, tiles, K/V loads and cycles. diff --git a/sim/golden/export_fixtures.py b/sim/golden/export_fixtures.py new file mode 100644 index 0000000..eed5513 --- /dev/null +++ b/sim/golden/export_fixtures.py @@ -0,0 +1,88 @@ +"""Export golden fixtures for the WASM Node test (wasm/test.mjs). + +Each fixture holds one engine run: the quantized inputs, the config, the expected O (int16, Q8.8 V-LSB), the expected +HBM bytes, tile count and K/V loads (traffic.py), and the expected busy cycles. Cycles come from the closed-form model +for sequential FLASH, and from sim/perf_baseline.json for overlap and NAIVE (those are pinned baselines). + + python sim/golden/export_fixtures.py # writes wasm/fixtures/*.json +""" + +import json +import random +from pathlib import Path + +import numpy as np + +import fixed +from schedule import cycles as sched_cycles +from traffic import kv_tile_loads, traffic + +ROOT = Path(__file__).resolve().parents[2] +OUT = ROOT / "wasm" / "fixtures" +BASELINE = json.loads((ROOT / "sim" / "perf_baseline.json").read_text()) + +CASES = [ + dict(name="flash_seq", mode="FLASH", n=16, causal=False, qrows=8, overlap=False, latency=10, seed=7), + dict(name="flash_overlap_causal", mode="FLASH", n=16, causal=True, qrows=4, overlap=True, latency=10, seed=8), + dict(name="naive_causal", mode="NAIVE", n=16, causal=True, qrows=4, overlap=False, latency=10, seed=9), +] + + +def case_inputs(seed, n): + rng = np.random.default_rng(seed) + qf = rng.normal(0, 2.5, (n, 4)) + kf = rng.normal(0, 2.5, (n, 4)) + vf = rng.normal(0, 1.0, (n, 4)) + return fixed.quantize(qf, kf, vf) + + +def expected_cycles(c): + if c["mode"] == "FLASH" and not c["overlap"]: + return sched_cycles(c["n"], c["causal"], c["qrows"], c["latency"]) + key = (f"naive/N{c['n']}/c{int(c['causal'])}/L{c['latency']}" if c["mode"] == "NAIVE" + else f"flash/N{c['n']}/c{int(c['causal'])}/q{c['qrows']}/o{int(c['overlap'])}/L{c['latency']}") + return BASELINE[key] + + +def build_fixture(c): + qz = case_inputs(c["seed"], c["n"]) + naive = c["mode"] == "NAIVE" + qblock = 4 if naive else c["qrows"] + res = fixed.naive(qz, c["n"], c["causal"]) if naive else fixed.flash(qz, c["n"], c["causal"], qblock=qblock) + t = traffic(c["n"], c["causal"], c["mode"], qblock=qblock) + return dict( + name=c["name"], + mode=c["mode"], + n=c["n"], + causal=c["causal"], + qrows=qblock, + overlap=c["overlap"] and not naive, + latency=c["latency"], + cfgScale=int(qz.cfg_scale), + q=qz.q.reshape(-1).tolist(), + k=qz.k.reshape(-1).tolist(), + v=qz.v.reshape(-1).tolist(), + out=res.out.reshape(-1).astype(int).tolist(), + hbmRdBytes=t["rd"], + hbmWrBytes=t["wr"], + tiles=t["tiles"], + kvLoads=t["tiles"] if naive else kv_tile_loads(c["n"], c["causal"], qblock), + cycles=expected_cycles(c), + ) + + +def fixture_text(c): + return json.dumps(build_fixture(c)) + "\n" + + +def main(): + OUT.mkdir(parents=True, exist_ok=True) + for c in CASES: + text = fixture_text(c) + (OUT / f"{c['name']}.json").write_text(text) + f = json.loads(text) + print(f"{c['name']}: cycles {f['cycles']} bytes rd {f['hbmRdBytes']} wr {f['hbmWrBytes']} tiles {f['tiles']}") + + +if __name__ == "__main__": + main() diff --git a/sim/golden/test_golden.py b/sim/golden/test_golden.py index 3f4ce63..456dde0 100644 --- a/sim/golden/test_golden.py +++ b/sim/golden/test_golden.py @@ -257,3 +257,13 @@ def test_bit_width_bounds_worst_case(): assert np.abs(t["O"]).max() < 1 << 25, name # 26-bit signed assert t["P"].max() <= 256 # 9-bit unsigned assert np.abs(res.out).max() <= 127 * 256 + + +# ------------------------------------------------------------------ WASM fixtures +def test_wasm_fixtures_fresh(): + """wasm/fixtures must be exactly what export_fixtures.py produces (the Node test trusts them).""" + import export_fixtures as ef + + for c in ef.CASES: + path = ef.OUT / f"{c['name']}.json" + assert path.read_text() == ef.fixture_text(c), f"{path.name} is stale: run sim/golden/export_fixtures.py" diff --git a/wasm/README.md b/wasm/README.md index 814491f..c028b36 100644 --- a/wasm/README.md +++ b/wasm/README.md @@ -1,10 +1,16 @@ # wasm/ -`spike/` is the Phase 0 proof that Verilator + Emscripten work here, including **two independent -instances of one Verilated model** (each with its own `VerilatedContext`) in one module. Phase 6 adds -`harness.cpp`, `bindings.cpp`, `build.sh`, `test.mjs`. +- `build.sh`: Verilates `rtl/` (top `tiny_attn_top`) and links the engine into `web/public/tiny_attn.{mjs,wasm}`. + Output: 206 KB raw, 81 KB gzipped (budget 500 KB gzipped, enforced by `wasm.yml`). +- `harness/`: `TinyAttnSim` (one model instance with its own Verilator context) and the embind bindings. +- `test.mjs`: Node test. Runs the golden fixtures in `fixtures/` with two or more instances interleaved, checks O, bytes, + tiles, K/V loads and cycles exactly, and checks that a fresh instance reproduces a run. +- `fixtures/`: golden runs exported by `sim/golden/export_fixtures.py`. +- `spike/`: the Phase 0 toolchain spike (two counters in one module), kept as a regression check. ```bash source scripts/env.sh -bash wasm/spike/build.sh && node wasm/spike/test.mjs +bash wasm/build.sh && node wasm/test.mjs ``` + +The JavaScript API is in `docs/STATE_SCHEMA.md`. diff --git a/wasm/build.sh b/wasm/build.sh new file mode 100755 index 0000000..f5ac9ef --- /dev/null +++ b/wasm/build.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +# Build the TinyAttn engine as a WebAssembly module: Verilate rtl/ (top: tiny_attn_top), compile the C++ harness +# and the Verilator runtime with Emscripten, link into web/public/tiny_attn.{mjs,wasm}. +# Requires: verilator >= 5.036 and Emscripten (`source scripts/env.sh` first). +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +OBJ="$ROOT/wasm/build/obj" +OUT="$ROOT/web/public" + +VROOT="$(verilator --getenv VERILATOR_ROOT)" +[ -f "$VROOT/include/verilated.h" ] || { echo "VERILATOR_ROOT not found: $VROOT" >&2; exit 1; } + +RTL_TOP=tiny_attn_top +SRCS=("$ROOT/rtl/lut_pkg.sv" "$ROOT/rtl/tiny_attn_pkg.sv") +for f in "$ROOT"/rtl/*.sv; do + case "$f" in */lut_pkg.sv|*/tiny_attn_pkg.sv) ;; *) SRCS+=("$f") ;; esac +done + +echo "[1/2] Verilating ${RTL_TOP}..." +rm -rf "$OBJ" +mkdir -p "$OBJ" +verilator --cc "${SRCS[@]}" --top-module "$RTL_TOP" --prefix Vtiny_attn_top --Mdir "$OBJ" \ + -O3 -Wno-fatal -Wno-UNOPTFLAT >"$OBJ/verilator.log" 2>&1 || { tail -40 "$OBJ/verilator.log"; exit 1; } + +echo "[2/2] Compiling to WASM..." +mkdir -p "$OUT" +em++ -O2 -std=c++17 -DVL_IGNORE_UNKNOWN_ARCH \ + -I "$OBJ" -I "$VROOT/include" -I "$VROOT/include/vltstd" -I "$ROOT/wasm/harness" \ + "$ROOT/wasm/harness/harness.cpp" "$ROOT/wasm/harness/bindings.cpp" \ + "$OBJ"/Vtiny_attn_top*.cpp \ + "$VROOT/include/verilated.cpp" \ + "$VROOT/include/verilated_threads.cpp" \ + "$ROOT/wasm/emscripten_stubs.cpp" \ + -lembind \ + -s MODULARIZE=1 -s EXPORT_ES6=1 -s ALLOW_MEMORY_GROWTH=1 \ + -s EXPORT_NAME=createTinyAttn -s ENVIRONMENT=web,node \ + -o "$OUT/tiny_attn.mjs" + +WASM="$OUT/tiny_attn.wasm" +RAW=$(stat -c %s "$WASM") +GZ=$(gzip -9 -c "$WASM" | wc -c) +echo "Done: $OUT/tiny_attn.mjs + tiny_attn.wasm (raw $RAW bytes, gzip $GZ bytes)" diff --git a/wasm/fixtures/flash_overlap_causal.json b/wasm/fixtures/flash_overlap_causal.json new file mode 100644 index 0000000..d0c9228 --- /dev/null +++ b/wasm/fixtures/flash_overlap_causal.json @@ -0,0 +1 @@ +{"name": "flash_overlap_causal", "mode": "FLASH", "n": 16, "causal": true, "qrows": 4, "overlap": true, "latency": 10, "cfgScale": 1780, "q": [-87, -67, -68, -18, -116, -9, -48, 45, 48, 70, 38, -3, 43, 75, -33, 31, -2, 72, -42, -15, 18, 13, -82, 18, -6, -12, -8, 11, -91, 78, -43, -112, -4, 73, -26, 78, 78, -43, -123, -62, 59, -41, -76, -67, 0, -1, 44, 49, -47, -8, -57, 4, -58, -60, 106, 2, 32, 127, 39, -6, 3, -37, 55, -17], "k": [94, -18, -19, -22, 33, 10, 73, -29, 15, 103, -13, 28, 83, 29, -34, -109, 127, 23, 67, -25, 97, 28, 3, -48, 20, -26, -12, -86, -26, 8, -66, 4, 35, -16, -33, -33, 50, -81, -16, -22, 8, -4, -59, 108, 9, -71, -47, 5, 33, -71, 38, -73, -42, 45, 90, -4, -16, 67, -57, 15, 12, 54, 16, -14], "v": [25, -45, 33, -27, -37, -15, 2, 29, -9, 30, 39, -4, -40, 37, 11, -85, -40, 6, -55, -27, -69, 16, 10, 7, 58, -14, -12, 81, -52, 29, -53, 10, 85, 11, -31, 96, -22, -8, 12, -54, 45, 65, 31, 13, -6, 36, -127, -5, 10, -80, 17, -5, -1, 4, 23, 18, -42, -5, -7, 2, -24, 39, 53, 61], "out": [6400, -11519, 8447, -6912, -8216, -4448, 1140, 6290, -2766, 6817, 9326, -464, -2304, 7680, 9983, -1024, -2687, 7766, 9637, -2026, -2846, 3410, 7471, -7763, 1010, -2319, 3089, 1805, -8336, 8409, 2015, -17765, -2304, 7680, 9983, -1024, -10111, 9309, 2860, -21645, -10061, 9296, 2816, -21507, 2772, 9415, 2634, 2012, -8948, 8810, -11825, 2467, -253, 418, 5807, 4474, -2477, 7514, 9419, -1137, 675, -16832, 3075, -612], "hbmRdBytes": 384, "hbmWrBytes": 128, "tiles": 10, "kvLoads": 10, "cycles": 611} diff --git a/wasm/fixtures/flash_seq.json b/wasm/fixtures/flash_seq.json new file mode 100644 index 0000000..d9d36bc --- /dev/null +++ b/wasm/fixtures/flash_seq.json @@ -0,0 +1 @@ +{"name": "flash_seq", "mode": "FLASH", "n": 16, "causal": false, "qrows": 8, "overlap": false, "latency": 10, "cfgScale": 1656, "q": [0, 15, -14, -45, -23, -50, 3, 68, -25, -31, 25, 18, 5, -47, -1, 35, -68, -23, -96, -65, -93, -12, -64, 14, 8, -9, -127, -27, -2, 6, -77, -24, -49, -41, 54, -41, -2, 45, -29, -6, 6, 3, -62, 4, 69, -78, 43, 6, -32, 101, 38, -61, 4, 29, -10, 34, -3, 34, 73, -34, 10, -23, 6, -60], "k": [-33, -11, 51, 65, -75, -45, 37, -113, -26, -6, 71, 39, -19, -21, -14, 86, -24, -17, 20, -7, -11, -63, -1, -25, 66, 37, -1, 38, -19, 60, 0, 33, -73, 20, -96, -115, -17, -51, 9, 127, -47, -35, 12, 28, -10, -12, 40, 29, -58, -4, 2, -60, 15, -49, 55, 11, 5, -33, -7, -113, -64, 21, -120, 48], "v": [-126, 55, -61, 56, 9, -111, 90, 104, -5, -20, -12, -70, 79, -39, -4, -57, -45, -92, 91, -11, 70, 1, -50, -24, -40, 1, -27, -22, -99, -58, 119, -48, -76, 24, 101, -105, -15, -46, -127, 53, -2, 5, -54, 33, -39, -10, -80, -88, 96, -37, 21, -2, -32, -37, 45, -22, -11, 2, 85, 49, 28, -41, -100, 68], "out": [-16809, 4267, 25180, -21357, -3681, -11710, -32198, 13316, -6644, -5700, -17150, 5840, -3052, -11228, -30137, 11818, -19455, 6144, 25854, -26878, 6177, -9878, -23686, 15760, -14265, 2900, 15824, -18245, -17001, 4610, 21112, -22797, 2304, -28414, 23039, 26622, -8871, -1900, 5151, -7939, 4955, -9115, -21290, 13686, -7432, -9175, 9662, -5171, -14326, -57, 24554, -17481, -3808, -8679, -9204, 2405, 1231, -22508, 17578, 17145, -1175, -11168, 22263, 16684], "hbmRdBytes": 320, "hbmWrBytes": 128, "tiles": 16, "kvLoads": 8, "cycles": 749} diff --git a/wasm/fixtures/naive_causal.json b/wasm/fixtures/naive_causal.json new file mode 100644 index 0000000..ebcdc23 --- /dev/null +++ b/wasm/fixtures/naive_causal.json @@ -0,0 +1 @@ +{"name": "naive_causal", "mode": "NAIVE", "n": 16, "causal": true, "qrows": 4, "overlap": false, "latency": 10, "cfgScale": 1754, "q": [-40, 12, -83, 33, 58, -23, 22, 13, -20, -43, -102, 71, -2, 127, 42, 14, -33, 70, -25, 79, -20, 9, -77, 118, -5, -19, 41, -45, 39, -59, 27, -53, -92, -30, -74, 28, 1, 26, 5, -18, 1, 53, -9, -38, -49, -11, -38, 95, -59, -54, -15, 33, -13, -48, -41, -18, 49, -49, -5, -31, 16, 20, -74, 120], "k": [10, -27, -48, -14, 2, 32, 0, 26, -86, -40, -54, 9, 25, -78, -9, 39, 49, -60, 30, 18, -23, 105, -60, -15, 24, 13, -9, 74, 77, 49, 54, -3, -62, 117, -51, 119, -7, 12, -9, 51, -27, 99, 29, 9, -31, 15, 0, 102, 50, -83, 82, -127, 82, -101, 0, 43, 0, -59, 66, -95, -78, 74, -46, 13], "v": [36, -5, 12, 5, 10, 53, 42, 14, 41, -24, -34, -45, 35, 72, 4, -88, -23, 21, 13, 4, 44, 41, -60, 19, -33, 10, 55, 82, 31, -23, -16, 24, -44, -3, 27, 41, -13, 16, 61, 14, -77, 10, 11, -19, 2, -66, 56, 19, -24, 11, 6, -39, -119, -127, 90, -50, -13, 25, -16, 53, -42, 25, -93, 32], "out": [9215, -1280, 3072, 1280, 6355, 5102, 6373, 2270, 10491, -6125, -8658, -11470, 2560, 13567, 10751, 3584, 3188, 12008, 9212, 2389, 10301, -3593, -8188, -10745, -3058, 5700, 2834, -1029, -4296, 6241, 3029, -871, 10327, -6102, -8583, -11349, 8381, 6241, -9540, 5392, 8534, 9667, -13660, 4068, -11127, -955, 6998, 10430, 10369, -6079, -8365, -11416, 7007, -5315, -4396, -10996, -7559, 759, 2776, -10063, -11263, -768, 6912, 10495], "hbmRdBytes": 1344, "hbmWrBytes": 768, "tiles": 10, "kvLoads": 10, "cycles": 1807} diff --git a/wasm/harness/bindings.cpp b/wasm/harness/bindings.cpp new file mode 100644 index 0000000..13044a6 --- /dev/null +++ b/wasm/harness/bindings.cpp @@ -0,0 +1,54 @@ +// Embind bindings: the JavaScript API of the TinyAttn WebAssembly module. +#include +#include + +#include +#include + +#include "harness.h" + +using namespace emscripten; + +namespace { +std::vector toI32(const val& arr) { return vecFromJSArray(arr); } + +val snapshotToJs(const Snapshot& s) { + val o = val::object(); + o.set("cycle", s.cycle); + o.set("ctrlState", s.ctrlState); + o.set("busy", s.busy); + o.set("done", s.done); + o.set("hbmRdBytes", s.hbmRdBytes); + o.set("hbmWrBytes", s.hbmWrBytes); + o.set("cycles", s.cycles); + o.set("tiles", s.tiles); + o.set("kvLoads", s.kvLoads); + return o; +} +} // namespace + +EMSCRIPTEN_BINDINGS(tiny_attn) { + register_vector("VectorInt32"); + register_vector("VectorUint32"); + + class_("TinyAttnSim") + .constructor<>() + .function("reset", &TinyAttnSim::reset) + .function("configure", &TinyAttnSim::configure) + .function("loadQKV", optional_override([](TinyAttnSim& self, const val& q, const val& k, const val& v) { + self.loadQKV(toI32(q), toI32(k), toI32(v)); + })) + .function("start", &TinyAttnSim::start) + .function("step", optional_override([](TinyAttnSim& self) { return snapshotToJs(self.step()); })) + .function("runToDone", &TinyAttnSim::runToDone) + .function("readOutput", optional_override([](TinyAttnSim& self) { + auto out = self.readOutput(); + return val::array(out.begin(), out.end()); + })) + .function("readHbm", optional_override([](TinyAttnSim& self, int addr, int len) { + auto words = self.readHbm(addr, len); + return val::array(words.begin(), words.end()); + })) + .function("snapshot", optional_override([](TinyAttnSim& self) { return snapshotToJs(self.snapshot()); })) + .function("isDone", &TinyAttnSim::isDone); +} diff --git a/wasm/harness/harness.cpp b/wasm/harness/harness.cpp new file mode 100644 index 0000000..0e04126 --- /dev/null +++ b/wasm/harness/harness.cpp @@ -0,0 +1,155 @@ +#include "harness.h" + +#include + +namespace { +// HBM layout (rtl/tiny_attn_pkg.sv) +constexpr uint32_t Q_BASE = 0, K_BASE = 32, V_BASE = 64, O_BASE = 96; +} // namespace + +TinyAttnSim::TinyAttnSim() { + ctx_ = new VerilatedContext; + top_ = new Vtiny_attn_top(ctx_, "TOP"); + reset(); +} + +TinyAttnSim::~TinyAttnSim() { + top_->final(); + delete top_; + delete ctx_; +} + +void TinyAttnSim::evalComb() { top_->eval(); } + +void TinyAttnSim::clockEdge() { + top_->clk = 0; + evalComb(); + top_->clk = 1; + evalComb(); + ++cycle_; +} + +void TinyAttnSim::reset() { + top_->start = 0; + top_->ld_en = 0; + top_->ld_addr = 0; + top_->ld_word = 0; + top_->o_rd_addr = 0; + top_->cfg_nb = 1; + top_->cfg_gg = 1; + top_->cfg_causal = 0; + top_->cfg_naive = 0; + top_->cfg_overlap = 0; + top_->cfg_hbm_latency = 0; + top_->cfg_scale = 0; + top_->clk = 0; + top_->rst_n = 0; + clockEdge(); + clockEdge(); + top_->rst_n = 1; + clockEdge(); + cycle_ = 0; +} + +void TinyAttnSim::configure(int n, int qrows, bool causal, int mode, bool overlap, int latency, uint32_t scale) { + if (n < 4 || n > 32 || n % 4 != 0) throw std::invalid_argument("n must be a multiple of 4 in 4..32"); + if (qrows != 4 && qrows != 8 && qrows != 16 && qrows != 32) throw std::invalid_argument("qrows must be 4, 8, 16 or 32"); + if (qrows > n) throw std::invalid_argument("qrows must not exceed n"); + if (latency < 0 || latency > 64) throw std::invalid_argument("latency must be 0..64"); + if (scale >= (1u << 24)) throw std::invalid_argument("scale must fit 24 bits"); + n_ = n; + qrows_ = qrows; + top_->cfg_nb = static_cast(n / 4); + top_->cfg_gg = static_cast(qrows / 4); + top_->cfg_causal = causal ? 1 : 0; + top_->cfg_naive = mode == 1 ? 1 : 0; + top_->cfg_overlap = (overlap && mode == 0) ? 1 : 0; + top_->cfg_hbm_latency = static_cast(latency); + top_->cfg_scale = scale; +} + +void TinyAttnSim::loadQKV(const std::vector& q, const std::vector& k, const std::vector& v) { + const size_t want = static_cast(n_) * 4; + if (q.size() != want || k.size() != want || v.size() != want) throw std::invalid_argument("q, k, v must be n*4 values"); + auto pack = [](const std::vector& m, size_t row) { + uint32_t w = 0; + for (int c = 0; c < 4; ++c) { + uint32_t lane = static_cast(m[row * 4 + c]) & 0xFFu; + w |= lane << (8 * c); + } + return w; + }; + auto write = [&](uint32_t addr, uint32_t word) { + top_->ld_en = 1; + top_->ld_addr = addr; + top_->ld_word = word; + clockEdge(); + top_->ld_en = 0; + }; + for (int i = 0; i < n_; ++i) write(Q_BASE + i, pack(q, i)); + for (int i = 0; i < n_; ++i) write(K_BASE + i, pack(k, i)); + for (int i = 0; i < n_; ++i) write(V_BASE + i, pack(v, i)); +} + +void TinyAttnSim::start() { + top_->start = 1; + clockEdge(); + top_->start = 0; +} + +Snapshot TinyAttnSim::step() { + clockEdge(); + return snapshot(); +} + +int TinyAttnSim::runToDone(int maxCycles) { + for (int i = 0; i < maxCycles; ++i) { + if (top_->done) return i; + clockEdge(); + } + return top_->done ? maxCycles : -1; +} + +bool TinyAttnSim::isDone() const { return top_->done != 0; } + +Snapshot TinyAttnSim::snapshot() const { + Snapshot s; + s.cycle = cycle_; + s.ctrlState = top_->dbg_ctrl_state; + s.busy = top_->busy != 0; + s.done = top_->done != 0; + s.hbmRdBytes = top_->hbm_rd_bytes; + s.hbmWrBytes = top_->hbm_wr_bytes; + s.cycles = top_->cycles; + s.tiles = top_->tiles; + s.kvLoads = top_->kv_loads; + return s; +} + +std::vector TinyAttnSim::readOutput() { + std::vector out(static_cast(n_) * 4); + for (int i = 0; i < n_; ++i) { + for (int half = 0; half < 2; ++half) { + top_->o_rd_addr = O_BASE + 2 * i + half; + evalComb(); + uint32_t w = top_->o_rd_word; + for (int lane = 0; lane < 2; ++lane) { + int32_t v = static_cast((w >> (16 * lane)) & 0xFFFFu); + if (v >= 32768) v -= 65536; + out[static_cast(i) * 4 + 2 * half + lane] = v; + } + } + } + return out; +} + +std::vector TinyAttnSim::readHbm(int addr, int len) { + std::vector words; + words.reserve(len); + for (int i = 0; i < len; ++i) { + top_->o_rd_addr = addr + i; + evalComb(); + words.push_back(top_->o_rd_word); + } + return words; +} diff --git a/wasm/harness/harness.h b/wasm/harness/harness.h new file mode 100644 index 0000000..fba0d6a --- /dev/null +++ b/wasm/harness/harness.h @@ -0,0 +1,60 @@ +// TinyAttnSim: one instance of the tiny_attn_top model with its own VerilatedContext, so several instances can run in +// one WebAssembly module side by side (L4 of the visualizer runs NAIVE and FLASH at once). +// +// The engine is driven exactly as a host would drive it: Q/K/V are written through the HBM load port, start is a +// one-cycle pulse, O is read back through the readback port. Every step() is one rising clock edge. +#pragma once + +#include +#include + +#include "Vtiny_attn_top.h" +#include "verilated.h" + +struct Snapshot { + uint32_t cycle; // clock edges since reset + uint32_t ctrlState; // flash_ctrl / naive_ctrl state (dbg_ctrl_state) + bool busy; + bool done; + uint32_t hbmRdBytes; + uint32_t hbmWrBytes; + uint32_t cycles; // busy cycles (perf counter) + uint32_t tiles; + uint32_t kvLoads; +}; + +class TinyAttnSim { +public: + TinyAttnSim(); + ~TinyAttnSim(); + TinyAttnSim(const TinyAttnSim&) = delete; + TinyAttnSim& operator=(const TinyAttnSim&) = delete; + + void reset(); + // n: sequence length (multiple of 4, 4..32); qrows: rows per Q block (4..32, <= n); mode: 0 FLASH, 1 NAIVE. + void configure(int n, int qrows, bool causal, int mode, bool overlap, int latency, uint32_t scale); + // q, k, v: n*4 int8 values, row-major. Written through the HBM load port (takes n*3 cycles). + void loadQKV(const std::vector& q, const std::vector& k, const std::vector& v); + // Starts the engine (one-cycle start pulse). The next step() is the first cycle of the run. + void start(); + // One clock. Returns the state after the edge. + Snapshot step(); + // Steps until done or maxCycles edges. Returns the number of edges taken, or -1 on timeout. + int runToDone(int maxCycles); + // O for the configured n: n*4 int16 values (sign-extended), row-major. + std::vector readOutput(); + // Raw HBM words [addr, addr+len). + std::vector readHbm(int addr, int len); + Snapshot snapshot() const; + bool isDone() const; + +private: + void clockEdge(); + void evalComb(); + + VerilatedContext* ctx_; + Vtiny_attn_top* top_; + int n_ = 4; + int qrows_ = 4; + uint32_t cycle_ = 0; +}; diff --git a/wasm/test.mjs b/wasm/test.mjs new file mode 100644 index 0000000..289135a --- /dev/null +++ b/wasm/test.mjs @@ -0,0 +1,82 @@ +// Headless test of the TinyAttn WebAssembly module against the golden fixtures (wasm/fixtures, exported by +// sim/golden/export_fixtures.py). Two instances run side by side with their steps interleaved: any state shared +// between instances would change an output or a cycle count. Per run it checks: +// O (every value) == golden, HBM bytes == traffic.py, tiles and K/V loads == traffic.py, busy cycles == fixture. +// Usage: node wasm/test.mjs (after bash wasm/build.sh) +import { readdirSync, readFileSync } from "node:fs"; +import { dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; + +const root = join(dirname(fileURLToPath(import.meta.url)), ".."); +const { default: createTinyAttn } = await import(join(root, "web/public/tiny_attn.mjs")); +const mod = await createTinyAttn(); + +const fixtures = readdirSync(join(root, "wasm/fixtures")) + .filter((f) => f.endsWith(".json")) + .sort() + .map((f) => JSON.parse(readFileSync(join(root, "wasm/fixtures", f), "utf8"))); + +let failures = 0; +const check = (ok, msg) => { + if (!ok) { + failures++; + console.log(` FAIL ${msg}`); + } +}; + +function load(fx) { + const sim = new mod.TinyAttnSim(); + sim.configure(fx.n, fx.qrows, fx.causal, fx.mode === "NAIVE" ? 1 : 0, fx.overlap, fx.latency, fx.cfgScale); + sim.loadQKV(fx.q, fx.k, fx.v); + sim.start(); + return sim; +} + +// Run all fixtures together, one clock at a time, round-robin over the instances. +const sims = fixtures.map((fx) => ({ fx, sim: load(fx), last: null, finished: false })); +const maxCycles = 400_000; +for (let t = 0; t < maxCycles && sims.some((s) => !s.finished); t++) { + for (const s of sims) { + if (s.finished) continue; + s.last = s.sim.step(); + if (s.last.done) { + // the DONE cycle is counted by the busy-cycle counter at its end: one more edge makes the counters final + s.last = s.sim.step(); + s.finished = true; + } + } +} + +for (const { fx, sim, last, finished } of sims) { + console.log(`${fx.name} (${fx.mode} N=${fx.n} causal=${fx.causal} qrows=${fx.qrows} overlap=${fx.overlap} L=${fx.latency})`); + check(finished, "engine did not finish"); + const out = Array.from(sim.readOutput()); + const bad = out.findIndex((v, i) => v !== fx.out[i]); + check(bad < 0, `O differs at ${bad}: expected ${fx.out[bad]}, got ${out[bad]}`); + check(last.hbmRdBytes === fx.hbmRdBytes, `read bytes ${last.hbmRdBytes} != ${fx.hbmRdBytes}`); + check(last.hbmWrBytes === fx.hbmWrBytes, `write bytes ${last.hbmWrBytes} != ${fx.hbmWrBytes}`); + check(last.tiles === fx.tiles, `tiles ${last.tiles} != ${fx.tiles}`); + check(last.kvLoads === fx.kvLoads, `K/V loads ${last.kvLoads} != ${fx.kvLoads}`); + check(last.cycles === fx.cycles, `busy cycles ${last.cycles} != ${fx.cycles}`); + console.log(` cycles ${last.cycles} read ${last.hbmRdBytes} B write ${last.hbmWrBytes} B tiles ${last.tiles}`); + sim.delete(); +} + +// A fresh instance must reproduce the same run (no state left over from earlier instances). +{ + const fx = fixtures[0]; + const sim = load(fx); + let snap; + for (let t = 0; t < maxCycles; t++) { + snap = sim.step(); + if (snap.done) break; + } + snap = sim.step(); + check(Array.from(sim.readOutput()).every((v, i) => v === fx.out[i]), "fresh instance output differs"); + check(snap.cycles === fx.cycles, "fresh instance cycles differ"); + sim.delete(); + console.log(`fresh instance (${fx.name}): checked`); +} + +console.log(failures === 0 ? "PASS" : `FAILED (${failures} checks)`); +process.exit(failures === 0 ? 0 : 1);