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);