diff --git a/CONTEXT.md b/CONTEXT.md index 9f83467..e8d31ad 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -36,7 +36,18 @@ On a fresh machine delete the symlinks and run `.tools/setup.sh` and `.tools/ems - Phase 7 (web visualizer `/app`): DONE. Presets, editable Q/K/V (quantized in the browser), schedule, tile values, K/V buffers (read from the RTL), HBM meter against the golden traffic model, heatmap (engine P for NAIVE, float reference for FLASH), Playwright smoke test (5 tests, production build). -- Next: Phase 8 (L1 softmax cell, L2 tile dataflow, L4 naive vs flash with two instances, N sweep, error panel). +- Phase 8 (L1 softmax cell, L2 tile dataflow, L4 naive vs flash, N sweep, error panel): DONE. Engine readouts `tileDone` + and `getDatapath`. 12 Playwright tests, 15 Vitest, Node WASM test, one-seed cocotb. Light-mode screenshots not reviewed. +- Next: Phase 9 (docs, landing page, launch). Also open: the `astro dev` hang. + +## Gotchas found in Phase 8 +- The debug datapath wires are idle between tiles. The L2 view keeps the last busy clock's values, or the panel shows zeros. +- A datapath or buffer value is only meaningful while the engine is inside a tile; say so in the UI. +- Verilator's wide ports are word arrays; the `wordsOf` template reads any of them. +- A full run in the browser is cheap, but the pages overflowed at 375 px only after a run (long decimals in the output grid, + then a 1 px overhang in the tile grids). The phone check must run a computation first. +- `astro preview` hung (process state D) during Phase 8; a static file server over `dist/` works. Restart it after each build, + because the build removes its working directory. ## Gotchas found in Phase 7 - `wasm/fixtures/` holds engine runs and `quantize_cases.json`; the Node test must skip the quantization file (it did not for diff --git a/README.md b/README.md index abb9cbf..1325cd0 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-7 done: the engine (FLASH, NAIVE, K/V overlap) is bit-exact against the golden model, compiled to WebAssembly, and shown clock by clock in a browser visualizer at `/app`, with editable inputs, K/V buffers, HBM traffic against the model, and a Playwright smoke test. Next: the L1, L2 and L4 views. See [`docs/PLAN.md`](docs/PLAN.md) for the roadmap. +**Status:** Phases 0-8 done: the engine (FLASH, NAIVE, K/V overlap) is bit-exact against the golden model and compiled to WebAssembly. The browser visualizer at `/app` has four views (softmax cell, tile dataflow, schedule, naive vs flash with an N sweep) with editable inputs, K/V buffers, HBM traffic against the model, and 12 Playwright tests. Next: the landing page and the documentation pages. See [`docs/PLAN.md`](docs/PLAN.md) for the roadmap. ## Quick start diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 868d4cf..d7f3655 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -391,3 +391,33 @@ browser checks use `astro build` + `astro preview`. needs the files in git. Rebuild with `bash wasm/build.sh` after any RTL or harness change and commit them with it. CI does not yet check that the committed files match a fresh build; the wasm workflow's `node wasm/test.mjs` checks behavior, not the files. + +## 13. Phase 8: the other levels (L1, L2, L4) + +The `/app` page has four views over one engine: L1 softmax cell, L2 tile dataflow, L3 schedule (Phase 7), L4 naive vs +flash. + +- **Engine readouts added.** `tileDone` (a one-cycle pulse, from the tile core's `tile_done`) and `getDatapath()` (the + MM1 and MM2 weights, activations and partial sums, live). Checked in `wasm/test.mjs`: one completion per FLASH tile, and + the running row max never decreases within a row group (keyed on the global row group; the in-block index repeats across + Q blocks). +- **L1 (`SoftmaxCellView`).** The visualizer records the engine's row state at every tile completion. For a chosen row, + the table shows each K/V tile's four logits (engine and float), m, α and l. The float reference (`softmax-cell.ts`) + uses the same INT8 inputs dequantized, so the difference is fixed-point rounding. A FLASH run is required; NAIVE has no + running max. +- **L2 (`TileView`).** Live wires of MM1 and MM2 and the tile's S, Z, P, m, α, l, O. The values are from the last busy + clock of the run so far; between tiles the wires are idle. Stepping through a tile shows the activations streaming. +- **L4 (`CompareView`).** Two engine instances on the same inputs, each freed after its run. Bars for HBM bytes, busy + cycles, tiles and K/V loads (measured). The on-chip footprint, the full-S-matrix size and the tile-core occupancy are + formulas, labelled as such. The error panel shows |engine − float| per element for both schedules on one colour scale. + The N sweep runs N = 4 to 32 on seeded inputs and plots HBM bytes against N for both schedules. + +Gaps against the plan's Phase 8 text: +- The N sweep uses `runToDone` (one call per run), not a separate `runFast`; the cost is the same. +- The L2 view is a grid of values, not an animated wire diagram. Each block shows its live values. +- L1's exp2 split is described in text; the view does not draw the integer and LUT parts separately. +- "Every level works in light and dark": the tests check that both colour schemes render with a background. Nobody has + reviewed the light-mode screenshots yet. + +Checks: `pnpm test:e2e` (12 tests, including all four views at 375 px and the light/dark background), `pnpm test` (15), +`node wasm/test.mjs`, `SEEDS=1 bash sim/run_e2e.sh`, `bash sim/lint.sh`, `pnpm lint`, `pnpm typecheck`, `pnpm build`. diff --git a/docs/PLAN.md b/docs/PLAN.md index f008036..19bec65 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -452,7 +452,7 @@ The estimates assume one person part-time, at tiny-cpu's pace. `HbmMeter`, and **L3 Schedule first**. - **Exit:** pick a preset, step forward and back, and watch tiles, buffers and bytes update; the Playwright smoke test passes. -### Phase 8: Remaining levels (1.5 weeks) +### Phase 8: Remaining levels (1.5 weeks) ✅ DONE (see ARCHITECTURE §13) - L1 Softmax cell, L2 Tile dataflow, L4 Naive vs Flash (two instances, N sweep, error panel). - **Exit:** every level works on desktop and at 375 px, in light and dark themes. diff --git a/docs/STATE_SCHEMA.md b/docs/STATE_SCHEMA.md index 9666bfb..bcb8416 100644 --- a/docs/STATE_SCHEMA.md +++ b/docs/STATE_SCHEMA.md @@ -20,6 +20,7 @@ the ports for it (`dbg_*` in `tiny_attn_top` and `attn_tile`); only the bindings | `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) | +| `tileDone` | boolean | one-cycle pulse: a tile completed on the last edge (FLASH: one per tile; NAIVE: one per tile operation) | **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. @@ -41,6 +42,7 @@ class TinyAttnSim { isDone(): boolean; getTile(): TileState; // tile-core values of the most recent tile (S, Z, P, m, l, alpha, O) getKVBuffers(): KVBuffers; // K and V tile buffers (both), and the buffer in use + getDatapath(): Datapath; // MM1 and MM2 weights, activations and partial sums (live wires; 16 values each) readProbabilities(): number[]; // NAIVE: P from HBM, n*n values (UQ1.8) delete(): void; } @@ -58,6 +60,12 @@ Input validation throws: `n` must be a multiple of 4 in 4..32, `qrows` one of 4, Multiple instances are independent: each owns its Verilator context. Steps may be interleaved across instances. +## Datapath (`getDatapath`) + +Index `r*4 + c` is PE (r, c). `mm1W` (Kᵀ weights) and `mm1A` (Q activations) are int8; `mm1Psum` (S partial sums) is int17. +`mm2W` (V weights) is int8, `mm2A` (P activations) is uint9, `mm2Psum` (P·V partial sums) is int18. These are the wires at +the current clock. + ## K/V buffers (`getKVBuffers`) `k[(buf*4 + row)*4 + col]` is K[row][col] of buffer `buf` (int8, sign-extended); `v` the same for V. `sel` is the buffer diff --git a/rtl/tiny_attn_top.sv b/rtl/tiny_attn_top.sv index 3a4305a..5d54174 100644 --- a/rtl/tiny_attn_top.sv +++ b/rtl/tiny_attn_top.sv @@ -58,7 +58,8 @@ module tiny_attn_top // K/V buffers: word (buf*4 + row) holds row `row` of buffer `buf` (4 int8 lanes); dbg_kv_sel is the buffer in use output logic [2*4*32-1:0] dbg_kbuf, output logic [2*4*32-1:0] dbg_vbuf, - output logic dbg_kv_sel + output logic dbg_kv_sel, + output logic dbg_tile_done // one-cycle pulse: a group tile (or NAIVE row group) completed ); // ---- controllers (only the one selected by cfg_naive is started) logic f_busy, f_done, n_busy, n_done; @@ -117,6 +118,7 @@ module tiny_attn_top end end assign dbg_kv_sel = kv_sel; + assign dbg_tile_done = tile_done; logic [31:0] tin_words [8]; logic q_load_valid; diff --git a/wasm/harness/bindings.cpp b/wasm/harness/bindings.cpp index 3b1290e..560ed19 100644 --- a/wasm/harness/bindings.cpp +++ b/wasm/harness/bindings.cpp @@ -12,6 +12,17 @@ using namespace emscripten; namespace { std::vector toI32(const val& arr) { return vecFromJSArray(arr); } +val datapathToJs(const Datapath& d) { + val o = val::object(); + o.set("mm1W", val::array(d.mm1W.begin(), d.mm1W.end())); + o.set("mm1A", val::array(d.mm1A.begin(), d.mm1A.end())); + o.set("mm1Psum", val::array(d.mm1Psum.begin(), d.mm1Psum.end())); + o.set("mm2W", val::array(d.mm2W.begin(), d.mm2W.end())); + o.set("mm2A", val::array(d.mm2A.begin(), d.mm2A.end())); + o.set("mm2Psum", val::array(d.mm2Psum.begin(), d.mm2Psum.end())); + return o; +} + val kvToJs(const KVBuffers& kv) { val o = val::object(); o.set("sel", kv.sel); @@ -45,6 +56,7 @@ val snapshotToJs(const Snapshot& s) { o.set("cycles", s.cycles); o.set("tiles", s.tiles); o.set("kvLoads", s.kvLoads); + o.set("tileDone", s.tileDone); return o; } } // namespace @@ -73,6 +85,7 @@ EMSCRIPTEN_BINDINGS(tiny_attn) { })) .function("snapshot", optional_override([](TinyAttnSim& self) { return snapshotToJs(self.snapshot()); })) .function("isDone", &TinyAttnSim::isDone) + .function("getDatapath", optional_override([](TinyAttnSim& self) { return datapathToJs(self.getDatapath()); })) .function("getKVBuffers", optional_override([](TinyAttnSim& self) { return kvToJs(self.getKVBuffers()); })) .function("getTile", optional_override([](TinyAttnSim& self) { return tileToJs(self.getTile()); })) .function("readProbabilities", optional_override([](TinyAttnSim& self) { diff --git a/wasm/harness/harness.cpp b/wasm/harness/harness.cpp index c261864..01666ab 100644 --- a/wasm/harness/harness.cpp +++ b/wasm/harness/harness.cpp @@ -149,6 +149,7 @@ Snapshot TinyAttnSim::snapshot() const { s.cycles = top_->cycles; s.tiles = top_->tiles; s.kvLoads = top_->kv_loads; + s.tileDone = top_->dbg_tile_done != 0; return s; } @@ -243,3 +244,30 @@ KVBuffers TinyAttnSim::getKVBuffers() const { } return kv; } + +namespace { +// Signed or unsigned field of `width` bits at bit `lsb` of a word array, as int32. +int32_t field(const uint32_t* words, int lsb, int width, bool isSigned) { + uint64_t v = bitsAt(words, lsb, width); + return static_cast(isSigned ? signExtend(v, width) : static_cast(v)); +} +} // namespace + +Datapath TinyAttnSim::getDatapath() const { + Datapath d; + const uint32_t* w1 = wordsOf(top_->dbg_mm1_w); + const uint32_t* a1 = wordsOf(top_->dbg_mm1_a); + const uint32_t* p1 = wordsOf(top_->dbg_mm1_psum); + const uint32_t* w2 = wordsOf(top_->dbg_mm2_w); + const uint32_t* a2 = wordsOf(top_->dbg_mm2_a); + const uint32_t* p2 = wordsOf(top_->dbg_mm2_psum); + for (int i = 0; i < 16; ++i) { + d.mm1W.push_back(field(w1, i * 8, 8, true)); + d.mm1A.push_back(field(a1, i * 8, 8, true)); + d.mm1Psum.push_back(field(p1, i * 17, 17, true)); + d.mm2W.push_back(field(w2, i * 8, 8, true)); + d.mm2A.push_back(field(a2, i * 9, 9, false)); + d.mm2Psum.push_back(field(p2, i * 18, 18, true)); + } + return d; +} diff --git a/wasm/harness/harness.h b/wasm/harness/harness.h index 3243752..7a63ef9 100644 --- a/wasm/harness/harness.h +++ b/wasm/harness/harness.h @@ -21,6 +21,7 @@ struct Snapshot { uint32_t cycles; // busy cycles (perf counter) uint32_t tiles; uint32_t kvLoads; + bool tileDone; // one-cycle pulse: a tile completed on the last edge }; // Tile-core state of the most recent tile (see attn_tile.sv dbg_* ports). S, Z, P, m, l, alpha and O are for the @@ -39,6 +40,13 @@ struct TileState { // K and V tile buffers. k[(buf*4 + row)*4 + col] is K[row][col] of buffer buf (int8, sign-extended); same for v. // sel is the buffer the engine is using for the current tile. +// Datapath wires of the two systolic arrays, live. Index r*4 + c is PE (r, c). MM1 computes S = Q . Kt (weights Kt, +// activations Q rows); MM2 computes P . V (weights V, activations P). Values are signed except mm2A (P, unsigned). +struct Datapath { + std::vector mm1W, mm1A, mm1Psum; // 16 each: int8, int8, int17 + std::vector mm2W, mm2A, mm2Psum; // 16 each: int8, uint9, int18 +}; + struct KVBuffers { int sel; std::vector k; // 32 @@ -71,6 +79,7 @@ class TinyAttnSim { bool isDone() const; TileState getTile() const; KVBuffers getKVBuffers() const; + Datapath getDatapath() const; // P tiles written by a NAIVE run (HBM words 768..): n x n probabilities (UQ1.8), row-major. Requires n known. std::vector readProbabilities(); diff --git a/wasm/test.mjs b/wasm/test.mjs index bffdd99..c8c5c65 100644 --- a/wasm/test.mjs +++ b/wasm/test.mjs @@ -25,6 +25,25 @@ const check = (ok, msg) => { } }; +// FLASH tile order (web/src/lib/schedule.ts): Q blocks outer, K/V tiles, row groups; causal skips masked tiles. +function flashOrder(n, qrows, causal) { + const nb = n / 4; + const g = qrows / 4; + const order = []; + for (let gb = 0; gb < nb; gb += g) { + const bg = Math.min(g, nb - gb); + const last = causal ? gb + bg - 1 : nb - 1; + for (let kj = 0; kj <= last; kj++) { + for (let gi = 0; gi < bg; gi++) { + const qi = gb + gi; + if (causal && kj > qi) continue; + order.push({ qi, kj }); + } + } + } + return order; +} + 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); @@ -40,6 +59,23 @@ 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.tileDone && s.fx.mode === "FLASH") { + // one completion per tile; the running row max m never decreases within a row group + s.tileDones = (s.tileDones ?? 0) + 1; + const t = s.sim.getTile(); + s.order = s.order ?? flashOrder(s.fx.n, s.fx.qrows, s.fx.causal); + const key = s.order[s.tileDones - 1].qi; // global row group of this completion + s.prevM = s.prevM ?? new Map(); + const prev = s.prevM.get(key); + if (prev) { + for (let r = 0; r < 4; r++) { + if (t.m[r] < prev[r]) { + check(false, `${s.fx.name}: m decreased for group ${key} row ${r}: ${prev[r]} -> ${t.m[r]}`); + } + } + } + s.prevM.set(key, t.m.slice()); + } 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(); @@ -48,7 +84,7 @@ for (let t = 0; t < maxCycles && sims.some((s) => !s.finished); t++) { } } -for (const { fx, sim, last, finished } of sims) { +for (const { fx, sim, last, finished, tileDones } 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()); @@ -57,6 +93,7 @@ for (const { fx, sim, last, finished } of sims) { 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}`); + if (fx.mode === "FLASH") check(tileDones === fx.tiles, `tile completions ${tileDones} != ${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}`); if (fx.lastTile) { diff --git a/web/e2e/levels.spec.ts b/web/e2e/levels.spec.ts new file mode 100644 index 0000000..49e91cc --- /dev/null +++ b/web/e2e/levels.spec.ts @@ -0,0 +1,85 @@ +import { expect, test, type Page } from "@playwright/test"; + +async function openApp(page: Page) { + await page.goto("/app"); + await expect(page.getByRole("button", { name: "Start" })).toBeVisible({ timeout: 30_000 }); +} + +test("L1 softmax cell: running max and sum for the chosen row, against the float reference", async ({ page }) => { + await openApp(page); + await page.getByRole("button", { name: "Start" }).click(); + await page.getByRole("button", { name: "Run to end" }).click(); + await page.getByRole("tab", { name: "L1 Softmax cell" }).click(); + + await expect(page.getByText("m engine / float")).toBeVisible(); + await expect(page.getByRole("cell", { name: "kv0" })).toBeVisible(); + const text = await page.locator("main").innerText(); + const m = text.match(/m engine \/ float\n([0-9.-]+) \/ ([0-9.-]+)/); + expect(m, "m engine / float line").not.toBeNull(); + expect(Math.abs(Number(m?.[1]) - Number(m?.[2]))).toBeLessThan(0.05); +}); + +test("L1 says so for NAIVE", async ({ page }) => { + await openApp(page); + await page.getByRole("combobox", { name: /^Schedule/ }).selectOption("1"); + await page.getByRole("tab", { name: "L1 Softmax cell" }).click(); + await expect(page.getByText("NAIVE has no running max")).toBeVisible(); +}); + +test("L2 tile dataflow shows the MM1 and MM2 arrays with live values", async ({ page }) => { + await openApp(page); + await page.getByRole("button", { name: "Start" }).click(); + await page.getByRole("button", { name: "Run to end" }).click(); + await page.getByRole("tab", { name: "L2 Tile dataflow" }).click(); + + await expect(page.getByText("MM1: S = Q · Kᵀ (array)")).toBeVisible(); + await expect(page.getByText("MM2: P · V (array)")).toBeVisible(); + await expect(page.getByText("O accumulator (rescaled by α)")).toBeVisible(); +}); + +test("L4 runs both schedules on the same inputs and sweeps N", async ({ page }) => { + await openApp(page); + await page.getByRole("tab", { name: "L4 Naive vs Flash" }).click(); + await page.getByRole("button", { name: "Run both on these inputs" }).click(); + await expect(page.getByText("Same inputs, both schedules")).toBeVisible(); + await expect(page.getByText(/FLASH, max [0-9.]+, mean [0-9.]+/)).toBeVisible(); + + await page.getByRole("button", { name: "Sweep N = 4 to 32" }).click(); + await expect(page.getByText("HBM bytes against N (seeded inputs, same settings)")).toBeVisible(); + // header row plus one row per N: 4, 8, ..., 32 + const sweep = page.locator("table").last(); + await expect(sweep.locator("tbody tr")).toHaveCount(8); +}); + +test("light and dark colour schemes both render the visualizer", async ({ browser }) => { + for (const scheme of ["light", "dark"] as const) { + const context = await browser.newContext({ colorScheme: scheme }); + const page = await context.newPage(); + await openApp(page); + const bg = await page.evaluate(() => getComputedStyle(document.body).backgroundColor); + expect(bg).not.toBe("rgba(0, 0, 0, 0)"); + await context.close(); + } +}); + +test("L4 at phone width: no horizontal scroll", async ({ page }) => { + await page.setViewportSize({ width: 375, height: 812 }); + await openApp(page); + await page.getByRole("tab", { name: "L4 Naive vs Flash" }).click(); + await page.getByRole("button", { name: "Run both on these inputs" }).click(); + await expect(page.getByText("Same inputs, both schedules")).toBeVisible(); + const overflow = await page.evaluate(() => document.documentElement.scrollWidth - document.documentElement.clientWidth); + expect(overflow).toBeLessThanOrEqual(0); +}); + +test("every view has no horizontal scroll at phone width", async ({ page }) => { + await page.setViewportSize({ width: 375, height: 812 }); + await openApp(page); + await page.getByRole("button", { name: "Start" }).click(); + await page.getByRole("button", { name: "Run to end" }).click(); + for (const name of ["L1 Softmax cell", "L2 Tile dataflow", "L3 Schedule", "L4 Naive vs Flash"]) { + await page.getByRole("tab", { name }).click(); + const overflow = await page.evaluate(() => document.documentElement.scrollWidth - document.documentElement.clientWidth); + expect(overflow, name).toBeLessThanOrEqual(0); + } +}); diff --git a/web/e2e/smoke.spec.ts b/web/e2e/smoke.spec.ts index 5466029..28bc024 100644 --- a/web/e2e/smoke.spec.ts +++ b/web/e2e/smoke.spec.ts @@ -36,7 +36,7 @@ test("step forward and back moves the clock", async ({ page }) => { test("NAIVE run shows the engine's probability matrix from HBM", async ({ page }) => { await openApp(page); - await page.getByLabel("Schedule").selectOption("1"); + await page.getByRole("combobox", { name: /^Schedule/ }).selectOption("1"); await page.getByRole("button", { name: "Start" }).click(); await page.getByRole("button", { name: "Run to end" }).click(); diff --git a/web/public/tiny_attn.wasm b/web/public/tiny_attn.wasm index 3b8f791..6f65f01 100755 Binary files a/web/public/tiny_attn.wasm and b/web/public/tiny_attn.wasm differ diff --git a/web/src/components/CompareView.tsx b/web/src/components/CompareView.tsx new file mode 100644 index 0000000..0b5492d --- /dev/null +++ b/web/src/components/CompareView.tsx @@ -0,0 +1,203 @@ +// L4: NAIVE and FLASH on the same inputs. Two engine instances; each run is its own instance, so the two never share state. +// Bars: HBM bytes, cycles, tiles. Formulas (not measurements): on-chip footprint and tile-core occupancy. Error panel: +// |engine − float attention| per element. N sweep: HBM bytes against N on seeded inputs, both schedules. +import { useMemo, useState } from "react"; +import { attention } from "@/lib/reference"; +import { runEngine, sweep, type RunSummary, type SweepRow } from "@/lib/engine-run"; +import type { EngineModule } from "@/lib/state-schema"; + +interface Props { + mod: EngineModule; + q: number[][]; + k: number[][]; + v: number[][]; + causal: boolean; + qrows: number; + latency: number; + overlap: boolean; +} + +const muted = { color: "var(--muted)" } as const; +const panel = { background: "var(--panel)" } as const; +const TILE_CORE_CYCLES = 21; // one group tile: tile_start to the group's tile_done, plus the weight load share + +function Bar({ label, a, b, fmtv }: { label: string; a: number; b: number; fmtv: (x: number) => string }) { + const max = Math.max(a, b, 1); + return ( +
+
{label}
+ {[{ name: "FLASH", v: a, c: "var(--accent)" }, { name: "NAIVE", v: b, c: "var(--err)" }].map((row) => ( +
+ {row.name} +
+
+
+ {fmtv(row.v)} +
+ ))} +
+ ); +} + +function ErrorGrid({ title, err, scale }: { title: string; err: number[][]; scale: number }) { + return ( +
+
{title}
+
+ {err.flatMap((row, i) => + row.map((e, c) => ( +
+ {e.toFixed(3)} +
+ )), + )} +
+
+ ); +} + +function SweepChart({ rows }: { rows: SweepRow[] }) { + const W = 360; + const H = 150; + const maxB = Math.max(...rows.flatMap((r) => [r.flashBytes, r.naiveBytes]), 1); + const x = (n: number) => 10 + ((n - 4) / 28) * (W - 20); + const y = (b: number) => H - 14 - (b / maxB) * (H - 28); + const line = (key: "flashBytes" | "naiveBytes") => rows.map((r) => `${x(r.n)},${y(r[key])}`).join(" "); + return ( +
+ + + + {rows.map((r) => ( + {r.n} + ))} + +
+ ━ FLASH HBM bytes + ━ NAIVE HBM bytes + x-axis: N +
+
+ + + + + + + + + {rows.map((r) => ( + + + + + ))} + +
NFLASH bytesNAIVE bytesFLASH cyclesNAIVE cycles
{r.n}{r.flashBytes}{r.naiveBytes}{r.flashCycles}{r.naiveCycles}
+
+
+ ); +} + +export default function CompareView({ mod, q, k, v, causal, qrows, latency, overlap }: Props) { + const [both, setBoth] = useState<{ flash: RunSummary; naive: RunSummary } | null>(null); + const [sweepRows, setSweepRows] = useState(null); + const [error, setError] = useState(null); + const n = q.length; + const flashQrows = Math.min(qrows, n); + + const reference = useMemo(() => attention(q, k, v, causal).out, [q, k, v, causal]); + + const runBoth = () => { + setError(null); + try { + const base = { n, causal, latency }; + const flash = runEngine(mod, q, k, v, { ...base, qrows: flashQrows, overlap, naive: false }); + const naive = runEngine(mod, q, k, v, { ...base, qrows: 4, overlap: false, naive: true }); + setBoth({ flash, naive }); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } + }; + + const runSweep = () => { + setError(null); + try { + setSweepRows(sweep(mod, qrows, causal, latency, 1234)); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } + }; + + const errs = (out: number[][]) => out.map((row, i) => row.map((x, c) => Math.abs(x - (reference[i]?.[c] ?? 0)))); + const flashErr = both ? errs(both.flash.out) : null; + const naiveErr = both ? errs(both.naive.out) : null; + const flatMean = (e: number[][]) => e.flat().reduce((a, b) => a + b, 0) / e.flat().length; + const flatMax = (e: number[][]) => Math.max(...e.flat()); + const heatScale = Math.max(0.02, flashErr && naiveErr ? Math.max(flatMax(flashErr), flatMax(naiveErr)) : 0.02); + + return ( +
+
+ + + + FLASH: Q block {flashQrows} rows{overlap ? ", overlap on" : ""}. NAIVE: 4-row blocks. Causal {causal ? "on" : "off"}. + +
+ {error &&

{error}

} + + {both && ( +
+

Same inputs, both schedules

+
+ `${x} B`} /> + `${x}`} /> + `${x}`} /> + `${x}`} /> +
+
+ + + + + + + + + + + + + + + + +
MeasureFLASHNAIVEBasis
On-chip footprint{flashQrows * 4 + 64 + flashQrows * 13} B32 B (one Z tile)formula: Q block + K/V double buffer + O accumulators (FLASH); one 4×4 Z tile (NAIVE)
Full S matrixnever stored{n * n * 2} B per tensor, in HBMformula: N² × 2 B for Z (NAIVE writes and reads it; P is the same size)
Tile-core occupancy{Math.round((both.flash.tiles * TILE_CORE_CYCLES / both.flash.cycles) * 100)}%n/aestimate: tiles × {TILE_CORE_CYCLES} / busy cycles; NAIVE per-pass cost is not counted separately
+
+ +
+

Error against float attention (same inputs, unquantized float reference)

+
+ + +
+

+ Both errors are dominated by the INT8 quantization of Q, K and V. NAIVE quantizes P against the global max; FLASH + rescales as it goes. The grids share one colour scale. +

+
+
+ )} + + {sweepRows && ( +
+

HBM bytes against N (seeded inputs, same settings)

+ +
+ )} +
+ ); +} diff --git a/web/src/components/LevelTabs.tsx b/web/src/components/LevelTabs.tsx new file mode 100644 index 0000000..312c015 --- /dev/null +++ b/web/src/components/LevelTabs.tsx @@ -0,0 +1,39 @@ +// Tab strip for the views (L1 to L4). Same look as tiny-cpu's tabs. +export interface LevelDef { + id: T; + label: string; +} + +interface Props { + levels: LevelDef[]; + value: T; + onChange: (id: T) => void; +} + +export default function LevelTabs({ levels, value, onChange }: Props) { + return ( +
+ {levels.map((l) => { + const selected = l.id === value; + return ( + + ); + })} +
+ ); +} diff --git a/web/src/components/SoftmaxCellView.tsx b/web/src/components/SoftmaxCellView.tsx new file mode 100644 index 0000000..9db3959 --- /dev/null +++ b/web/src/components/SoftmaxCellView.tsx @@ -0,0 +1,126 @@ +// L1: one row lane of the online softmax. Each K/V tile adds four logits; the running max m rises, the rescale α scales +// the old sum down, and the sum l grows. Values are the engine's, read at each tile completion. The reference column is +// the float softmax of the same INT8 inputs (rowReference), so the differences are fixed-point rounding only. +import { useMemo, useState } from "react"; +import { rowReference, rowSteps, type TileSample } from "@/lib/softmax-cell"; + +interface Props { + samples: TileSample[]; + q: number[][]; + k: number[][]; + causal: boolean; + labels: string[]; + naive: boolean; +} + +const muted = { color: "var(--muted)" } as const; +const fmt = (x: number, d = 2) => (Number.isFinite(x) ? x.toFixed(d) : "-"); + +export default function SoftmaxCellView({ samples, q, k, causal, labels, naive }: Props) { + const [row, setRow] = useState(0); + const n = q.length; + const steps = useMemo(() => rowSteps(samples, row), [samples, row]); + const ref = useMemo(() => rowReference(q, k, row, causal), [q, k, row, causal]); + + if (naive) { + return ( +

+ NAIVE has no running max: the softmax of every row is computed in three passes over HBM. Switch to FLASH to see the + online cell. +

+ ); + } + if (steps.length === 0) { + return

Run the engine (Start, then Step or Run to end) to fill this view.

; + } + + const last = steps[steps.length - 1]; + + // polyline of m and l across the row's tiles, each scaled to its own range + const W = 320; + const H = 110; + const pts = (vals: number[]) => { + const lo = Math.min(...vals); + const hi = Math.max(...vals); + const span = hi - lo || 1; + return vals.map((v, i) => [ + (i / Math.max(1, vals.length - 1)) * (W - 20) + 10, + H - 10 - ((v - lo) / span) * (H - 20), + ] as const); + }; + const polyline = (vals: number[]) => pts(vals).map(([x, y]) => `${x},${y}`).join(" "); + const mVals = steps.map((s) => s.m); + const lVals = steps.map((s) => s.l); + + return ( +
+ + +
+ + + + + + + + + + + + + {steps.map((s, i) => ( + + + + + + + + + ))} + +
K/V tilelogits, engine (log2)logits, floatmαl
kv{s.kj} + {s.z.map((x) => (x < -127 ? "masked" : fmt(x))).join(" ")} + + {[0, 1, 2, 3].map((c) => fmt(ref.z[s.kj * 4 + c] ?? Number.NaN)).join(" ")} + {fmt(s.m)}{fmt(s.alpha, 3)}{fmt(s.l, 3)}
+
+ + + + + {pts(mVals).map(([x, y], i) => )} + {pts(lVals).map(([x, y], i) => )} + +
+ ━ m (running max, log2) + ┅ l (running sum, scaled per line) +
+ +
+
After all tiles
+
+ m engine / float + {fmt(last?.m ?? Number.NaN, 3)} / {fmt(ref.m, 3)} + l engine / float + {fmt(last?.l ?? Number.NaN, 3)} / {fmt(ref.l, 3)} + |Δm|, |Δl| + {fmt(Math.abs((last?.m ?? 0) - ref.m), 4)}, {fmt(Math.abs((last?.l ?? 0) - ref.l), 4)} +
+

+ Each logit is split into an integer shift and a LUT fraction inside exp2 (`2^z = 2^⌊z⌋ · 2^frac(z)`). The float + column is the same INT8 inputs in full precision. +

+
+
+ ); +} diff --git a/web/src/components/TileView.tsx b/web/src/components/TileView.tsx new file mode 100644 index 0000000..6ced87a --- /dev/null +++ b/web/src/components/TileView.tsx @@ -0,0 +1,82 @@ +// L2: the datapath of one (qi, kj) step, with the live wire values at the current clock. The two arrays are PE grids: +// PE (r, c) holds the value index r*4 + c. Values are the engine's; the grids show what the wires carry now. +import type { ReactNode } from "react"; +import type { Datapath, TileState } from "@/lib/state-schema"; + +interface Props { + datapath: Datapath | null; + tile: TileState | null; + naive: boolean; +} + +const muted = { color: "var(--muted)" } as const; + +function Grid({ title, values, scale = 1, digits = 0 }: { title: string; values: number[]; scale?: number; digits?: number }) { + const cells = values.slice(0, 16); + return ( +
+
{title}
+
+
+ {cells.map((v, i) => {(v * scale).toFixed(digits)})} +
+
+
+ ); +} + +function Stage({ name, children }: { name: string; children: ReactNode }) { + return ( +
+
{name}
+
{children}
+
+ ); +} + +export default function TileView({ datapath, tile, naive }: Props) { + if (naive) { + return ( +

+ NAIVE runs its passes through the same arrays, but its Z and P tiles go out to HBM between passes. Switch to FLASH + to see the single-tile datapath. +

+ ); + } + if (!datapath || !tile) { + return

Run the engine to fill this view.

; + } + return ( +
+

+ Wires at the last busy clock of the run so far (between tiles they are idle). The S, Z, P and O grids are for the row group of the most recent tile. +

+
+ + + + + + + + + +
+ + + +
+ + + + + + + + + + +
+
+ ); +} diff --git a/web/src/components/Visualizer.tsx b/web/src/components/Visualizer.tsx index ca0ca80..3093c5b 100644 --- a/web/src/components/Visualizer.tsx +++ b/web/src/components/Visualizer.tsx @@ -1,11 +1,16 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import AttentionHeatmap from "@/components/AttentionHeatmap"; +import CompareView from "@/components/CompareView"; +import LevelTabs, { type LevelDef } from "@/components/LevelTabs"; +import SoftmaxCellView from "@/components/SoftmaxCellView"; +import TileView from "@/components/TileView"; import HbmMeter from "@/components/HbmMeter"; import MatrixEditor from "@/components/MatrixEditor"; import SramPanel from "@/components/SramPanel"; import { loadEngine } from "@/lib/engine"; import { PRESETS } from "@/lib/presets"; import { quantize } from "@/lib/quantize"; +import type { TileSample } from "@/lib/softmax-cell"; import { attention } from "@/lib/reference"; import { flashTileOrder, naiveTileOrder, type TileRef } from "@/lib/schedule"; import { TRAFFIC } from "@/lib/traffic-table"; @@ -13,6 +18,7 @@ import { FLASH_STATES, NAIVE_STATES, type EngineInstance, + type Datapath, type EngineModule, type KVBuffers, type Snapshot, @@ -21,6 +27,15 @@ import { const MAX_STEPS = 400_000; +type Level = "softmax" | "tile" | "schedule" | "compare"; + +const LEVELS: LevelDef[] = [ + { id: "softmax", label: "L1 Softmax cell" }, + { id: "tile", label: "L2 Tile dataflow" }, + { id: "schedule", label: "L3 Schedule" }, + { id: "compare", label: "L4 Naive vs Flash" }, +]; + interface Settings { preset: number; mode: 0 | 1; @@ -78,6 +93,14 @@ export default function Visualizer() { const [finished, setFinished] = useState(false); const [result, setResult] = useState(null); const simRef = useRef(null); + const [level, setLevel] = useState("schedule"); + const [datapath, setDatapath] = useState(null); + const [samples, setSamples] = useState([]); + // completed tiles of the current run, in order (filled at each tile_done pulse; cleared when a run restarts) + const historyRef = useRef([]); + const orderRef = useRef([]); + // datapath wires at the last busy clock of the run so far (between tiles the wires are idle) + const lastBusyDatapath = useRef(null); const preset = PRESETS[settings.preset] ?? PRESETS[0]; const n = preset?.q.length ?? 4; @@ -95,6 +118,10 @@ export default function Visualizer() { setSnap(null); setTile(null); setKv(null); + setDatapath(null); + setSamples([]); + historyRef.current = []; + lastBusyDatapath.current = null; setFinished(false); setResult(null); }, []); @@ -112,18 +139,35 @@ export default function Visualizer() { [n, settings.mode, settings.qrows, settings.causal], ); + orderRef.current = order; + + /** Record a completed tile: its global row group and K/V tile, and the engine's row state at that point. */ + const observe = useCallback((sim: EngineInstance, s: Snapshot) => { + if (s.busy) lastBusyDatapath.current = sim.getDatapath(); + if (!s.tileDone) return; + const ref = orderRef.current[s.tiles - 1]; + if (!ref) return; + const t = sim.getTile(); + historyRef.current.push({ qi: ref.qi, kj: ref.kj, z: t.z, m: t.m, l: t.l, alpha: t.alpha }); + }, []); + /** Advance one clock. The DONE edge is followed by one finalizing edge (counters are final after it). */ - const advance = useCallback((sim: EngineInstance, state: { count: number; finished: boolean }) => { - if (state.finished) return null; - let s = sim.step(); - state.count++; - if (s.done) { - s = sim.step(); + const advance = useCallback( + (sim: EngineInstance, state: { count: number; finished: boolean }) => { + if (state.finished) return null; + let s = sim.step(); state.count++; - state.finished = true; - } - return s; - }, []); + observe(sim, s); + if (s.done) { + s = sim.step(); + state.count++; + observe(sim, s); + state.finished = true; + } + return s; + }, + [observe], + ); /** Publish the engine state to the view. Builds the result once the run has finished. */ const commit = useCallback( @@ -133,6 +177,8 @@ export default function Visualizer() { setFinished(done); setTile(sim.getTile()); setKv(sim.getKVBuffers()); + setDatapath(lastBusyDatapath.current); + setSamples(historyRef.current.slice()); if (!done) { setResult(null); return; @@ -156,6 +202,8 @@ export default function Visualizer() { (target: number) => { if (!mod) return; simRef.current?.delete(); + historyRef.current = []; + lastBusyDatapath.current = null; const qz = quantize(inputs.q, inputs.k, inputs.v); const sim = new mod.TinyAttnSim(); sim.configure(n, settings.qrows, settings.causal, settings.mode, settings.overlap, settings.latency, qz.cfgScale); @@ -233,7 +281,6 @@ export default function Visualizer() { return { bg: "var(--panel)", label: "" }; }; - const fmt = (v: number, digits = 0) => (Number.isFinite(v) ? v.toFixed(digits) : "-"); if (error) return

{error}

; if (!mod) return

Loading engine…

; @@ -337,7 +384,12 @@ export default function Visualizer() { -
+ + +
+ {level === "schedule" && ( + <> +
{/* schedule */}

@@ -368,48 +420,24 @@ export default function Visualizer() { {settings.mode === 1 ? " NAIVE passes 2 to 4 reuse these tiles." : ""}

- - {/* current tile */} -
-

Current tile values

- {tile && tilesStarted > 0 ? ( -
- {settings.mode === 0 ? ( - <> -

Row group {tile.curGrp} (rows {tile.curGrp * 4}–{tile.curGrp * 4 + 3})

- - - -
- fmt(v)} /> - fmt(v)} /> - fmt(v)} /> -
- - ) : ( - <> -

- NAIVE tiles go through HBM, not the tile core, so the per-element S, Z and P are not shown here. - The row max m and sum l are shown below; the final probabilities are in the heatmap after the run. -

-
- fmt(v)} /> - fmt(v)} /> -
- - )} -
- ) : ( -

Press Start, then Step.

- )} -
-

On-chip K/V buffers

+ + )} + {level === "softmax" && ( + + )} + {level === "tile" && } + {level === "compare" && ( + + )} +
{/* result */} {result && ( @@ -432,18 +460,13 @@ function Grid({ title, values, cols }: { title: string; values: number[]; cols: return (
{title}
-
- {rows.flatMap((row, i) => row.map((v, j) => {v}))} +
+
+ {rows.flatMap((row, i) => + row.map((v, j) => {Number.isInteger(v) ? v : v.toFixed(4)}), + )} +
); } - -function Vec({ title, values, fmt }: { title: string; values: number[]; fmt: (v: number) => string }) { - return ( -
-
{title}
-
{values.map((v) => fmt(v)).join(" ")}
-
- ); -} diff --git a/web/src/lib/engine-run.ts b/web/src/lib/engine-run.ts new file mode 100644 index 0000000..6f2acce --- /dev/null +++ b/web/src/lib/engine-run.ts @@ -0,0 +1,76 @@ +// Whole runs of the engine, for the comparison view and the N sweep. Each run uses its own instance, freed afterwards. +import { quantize } from "@/lib/quantize"; +import type { EngineModule } from "@/lib/state-schema"; + +const MAX_CYCLES = 400_000; + +export interface RunConfig { + n: number; + qrows: number; + causal: boolean; + naive: boolean; + overlap: boolean; + latency: number; +} + +export interface RunSummary { + out: number[][]; // n rows x 4, dequantized (float units) + hbmRd: number; + hbmWr: number; + cycles: number; + tiles: number; + kvLoads: number; +} + +/** Run to DONE and read everything. The counters are final one edge after DONE, so one more edge is taken. */ +export function runEngine(mod: EngineModule, q: number[][], k: number[][], v: number[][], cfg: RunConfig): RunSummary { + const qz = quantize(q, k, v); + const sim = new mod.TinyAttnSim(); + try { + sim.configure(cfg.n, cfg.qrows, cfg.causal, cfg.naive ? 1 : 0, cfg.overlap && !cfg.naive, cfg.latency, qz.cfgScale); + sim.loadQKV(qz.q, qz.k, qz.v); + sim.start(); + if (sim.runToDone(MAX_CYCLES) < 0) throw new Error("engine did not finish"); + const last = sim.step(); + const flat = sim.readOutput(); + const out: number[][] = []; + for (let i = 0; i < cfg.n; i++) { + out.push([0, 1, 2, 3].map((c) => ((flat[i * 4 + c] ?? 0) / 256) * qz.sv)); + } + return { out, hbmRd: last.hbmRdBytes, hbmWr: last.hbmWrBytes, cycles: last.cycles, tiles: last.tiles, kvLoads: last.kvLoads }; + } finally { + sim.delete(); + } +} + +/** Deterministic inputs for sweeps (not the preset): uniform Q and K in [-2, 2), V in [-1, 1). */ +export function seededInputs(n: number, seed: number): { q: number[][]; k: number[][]; v: number[][] } { + let s = seed >>> 0; + const next = () => { + s = (Math.imul(s, 1664525) + 1013904223) >>> 0; + return s / 4294967296; + }; + const mat = (amp: number) => Array.from({ length: n }, () => [0, 1, 2, 3].map(() => (next() * 2 - 1) * amp)); + return { q: mat(2), k: mat(2), v: mat(1) }; +} + +export interface SweepRow { + n: number; + flashBytes: number; + naiveBytes: number; + flashCycles: number; + naiveCycles: number; +} + +/** HBM bytes and cycles for N = 4, 8, ..., 32 under both schedules, on seeded inputs. */ +export function sweep(mod: EngineModule, qrows: number, causal: boolean, latency: number, seed: number): SweepRow[] { + const rows: SweepRow[] = []; + for (let n = 4; n <= 32; n += 4) { + const { q, k, v } = seededInputs(n, seed + n); + const base = { n, qrows: Math.min(qrows, n), causal, overlap: false, latency }; + const f = runEngine(mod, q, k, v, { ...base, naive: false }); + const nv = runEngine(mod, q, k, v, { ...base, qrows: 4, naive: true }); + rows.push({ n, flashBytes: f.hbmRd + f.hbmWr, naiveBytes: nv.hbmRd + nv.hbmWr, flashCycles: f.cycles, naiveCycles: nv.cycles }); + } + return rows; +} diff --git a/web/src/lib/softmax-cell.test.ts b/web/src/lib/softmax-cell.test.ts new file mode 100644 index 0000000..913ef46 --- /dev/null +++ b/web/src/lib/softmax-cell.test.ts @@ -0,0 +1,60 @@ +import { describe, expect, it } from "vitest"; +import { rowReference, rowSteps, type TileSample } from "@/lib/softmax-cell"; +import { seededInputs } from "@/lib/engine-run"; + +const q = seededInputs(8, 7).q; +const k = seededInputs(8, 7).k; + +describe("rowReference (float, same INT8 inputs)", () => { + it("causal row 0 sees only key 0: m is its logit and l is 1", () => { + const r = rowReference(q, k, 0, true); + expect(r.z.slice(1).every((x) => x === -Infinity)).toBe(true); + expect(r.m).toBe(r.z[0]); + expect(r.l).toBeCloseTo(1, 12); + }); + + it("non-causal: l is the sum of 2^(z - m) and is at least 1", () => { + const r = rowReference(q, k, 3, false); + expect(r.z.every((x) => Number.isFinite(x))).toBe(true); + expect(r.l).toBeGreaterThanOrEqual(1); + const direct = r.z.reduce((acc, x) => acc + 2 ** (x - r.m), 0); + expect(r.l).toBeCloseTo(direct, 12); + }); +}); + +describe("rowSteps", () => { + const sample = (qi: number, kj: number, base: number): TileSample => ({ + qi, + kj, + z: Array.from({ length: 16 }, (_, i) => (base + i) * 256), + m: [base * 256, 0, 0, 0], + l: [base * 4096, 0, 0, 0], + alpha: [32768, 0, 0, 0], + }); + + it("keeps only the tiles of the row's group and reads its lane", () => { + const samples = [sample(0, 0, 1), sample(1, 0, 5), sample(0, 1, 2)]; + const steps = rowSteps(samples, 0); + expect(steps.map((s) => s.kj)).toEqual([0, 1]); + expect(steps[0]?.z).toEqual([1, 2, 3, 4]); + expect(steps[1]?.m).toBe(2); + expect(steps[0]?.alpha).toBe(1); + }); + + it("row 5 is lane 1 of group 1", () => { + const steps = rowSteps([sample(1, 0, 5)], 5); + expect(steps[0]?.z).toEqual([5 + 4, 5 + 5, 5 + 6, 5 + 7]); + }); +}); + +describe("seededInputs", () => { + it("is deterministic and has the requested shape and bounds", () => { + const a = seededInputs(12, 99); + const b = seededInputs(12, 99); + expect(a).toEqual(b); + expect(a.q.length).toBe(12); + expect(a.q[0]?.length).toBe(4); + expect(a.q.flat().every((x) => Math.abs(x) <= 2)).toBe(true); + expect(a.v.flat().every((x) => Math.abs(x) <= 1)).toBe(true); + }); +}); diff --git a/web/src/lib/softmax-cell.ts b/web/src/lib/softmax-cell.ts new file mode 100644 index 0000000..3c0cfee --- /dev/null +++ b/web/src/lib/softmax-cell.ts @@ -0,0 +1,65 @@ +// One softmax row, in the engine's units, and the float reference for the same INT8 inputs. +// +// The engine works in log2 units: z = S · log2(e) / sqrt(d), with S = q·k computed on INT8 values. The float reference +// here uses the same INT8 values dequantized (q_int · s_q), so the only difference left is the engine's fixed-point +// rounding, not the quantization. +import { D, LOG2E, quantize } from "@/lib/quantize"; + +export interface RowReference { + z: number[]; // log2-domain logit per key; -Infinity where the causal mask removes the key + m: number; // max over the allowed keys + l: number; // sum of 2^(z - m) over the allowed keys +} + +export function rowReference(q: number[][], k: number[][], row: number, causal: boolean): RowReference { + const qz = quantize(q, k, q); + const n = q.length; + const qd = (i: number, c: number) => (qz.q[i * D + c] ?? 0) * qz.sq; + const kd = (j: number, c: number) => (qz.k[j * D + c] ?? 0) * qz.sk; + const z: number[] = []; + for (let j = 0; j < n; j++) { + if (causal && j > row) { + z.push(-Infinity); + continue; + } + let dot = 0; + for (let c = 0; c < D; c++) dot += qd(row, c) * kd(j, c); + z.push((dot / Math.sqrt(D)) * LOG2E); + } + const m = Math.max(...z.filter((x) => Number.isFinite(x))); + const l = z.reduce((acc, x) => (Number.isFinite(x) ? acc + 2 ** (x - m) : acc), 0); + return { z, m, l }; +} + +/** A completed tile as the engine reported it (recorded at the tile_done pulse). */ +export interface TileSample { + qi: number; // global row group + kj: number; // K/V tile + z: number[]; // 16 logits, Q8.8 log2 units (row-major: row, key column) + m: number[]; // 4 rows, Q8.8 + l: number[]; // 4 rows, real × 2^12 + alpha: number[]; // 4 rows, UQ1.15 +} + +export interface RowStep { + kj: number; + z: number[]; // 4 logits of this tile for the row, real units + m: number; // running max after the tile + alpha: number; // rescale factor applied to the old sum (0..1) + l: number; // running sum after the tile +} + +/** The steps row `row` went through, in tile order. */ +export function rowSteps(samples: TileSample[], row: number): RowStep[] { + const g = Math.floor(row / 4); + const r = row % 4; + return samples + .filter((s) => s.qi === g) + .map((s) => ({ + kj: s.kj, + z: [0, 1, 2, 3].map((c) => (s.z[r * 4 + c] ?? 0) / 256), + m: (s.m[r] ?? 0) / 256, + alpha: (s.alpha[r] ?? 0) / 32768, + l: (s.l[r] ?? 0) / 4096, + })); +} diff --git a/web/src/lib/state-schema.ts b/web/src/lib/state-schema.ts index 70a5009..22165c5 100644 --- a/web/src/lib/state-schema.ts +++ b/web/src/lib/state-schema.ts @@ -11,6 +11,7 @@ export interface Snapshot { cycles: number; tiles: number; kvLoads: number; + tileDone: boolean; // one-cycle pulse: a tile completed on the last edge } /** Tile-core state of the most recent tile (4 rows of one row group). */ @@ -26,6 +27,16 @@ export interface TileState { o: number[]; // 16 values, row-major (4 rows x 4 columns) } +/** Datapath wires of the two systolic arrays. Index r*4 + c is PE (r, c). */ +export interface Datapath { + mm1W: number[]; // Kᵀ weights, int8 + mm1A: number[]; // Q activations, int8 + mm1Psum: number[]; // S partial sums, int17 + mm2W: number[]; // V weights, int8 + mm2A: number[]; // P activations, uint9 + mm2Psum: number[]; // P·V partial sums, int18 +} + /** K and V tile buffers. k[(buf*4 + row)*4 + col] is K[row][col] of buffer buf (int8). sel: buffer in use. */ export interface KVBuffers { sel: number; @@ -47,6 +58,7 @@ export interface EngineInstance { isDone(): boolean; getTile(): TileState; getKVBuffers(): KVBuffers; + getDatapath(): Datapath; delete(): void; }