diff --git a/CONTEXT.md b/CONTEXT.md index e8d31ad..86dcc10 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -38,7 +38,11 @@ On a fresh machine delete the symlinks and run `.tools/setup.sh` and `.tools/ems FLASH), Playwright smoke test (5 tests, production build). - 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. +- Phase 9 (docs, landing page, launch): built. Landing page with the animated FLASH schedule, six docs pages with live + demos, `/og` social card, generated `measured.ts`, README rewritten for the series. Deployment: see the PR and + ARCHITECTURE §14. +- Next: Phase 10 (FPGA synthesis, ECP5 for the ULX3S). Blocked here: no yosys / nextpnr / ecppack installed, and the disk + had under 1 GB free during Phase 9. ## 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. diff --git a/README.md b/README.md index 1325cd0..426510f 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,47 @@ # TinyAttention -A real, synthesizable **FlashAttention engine** in SystemVerilog, compiled to WebAssembly and running live -in your browser: online softmax with a running max and sum, SRAM tiles filling and emptying, and standard -vs FlashAttention memory traffic side by side. +A synthesizable **FlashAttention engine** in SystemVerilog, verified bit-exact against a golden model, compiled to +WebAssembly, and running live in your browser. + +**Live:** https://tiny-attention.vercel.app + +![The visualizer running one FLASH pass](docs/media/visualizer.png) Fourth project in the series after [tiny-gpu](https://github.com/deaneeth/tiny-gpu), -[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. +[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 debug port on the compiled RTL, or from the golden model the RTL is checked against. + +## What it does + +Attention computes `softmax(Q Kᵀ / √d) V`. The score matrix is N × N, and storing it is what makes standard attention +expensive in memory. FlashAttention computes the output one tile at a time with an online softmax, so the matrix never +exists in memory. This engine implements that schedule, plus the standard three-pass schedule for comparison, with +optional overlap of the K/V loads with compute. + +## Verification -**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. +| What | Result | +|---|---| +| Engine vs golden model, end to end | Output, HBM bytes, tile count, K/V loads and busy cycles match on every run. Every sequence length 4 to 32, causal and not, every block size, both schedules, overlap on and off, random HBM latency, 200 seeds. | +| Exhaustive unit coverage | exp2 is checked on every input in its range; reciprocal on every mantissa. DMA burst timing is pinned cycle by cycle. | +| Output error vs float attention | Max 0.0028 to 0.0260 across the six presets (N = 12 to 16), in float units. The recorded 1000-seed study is in `sim/golden/study_results_1000seeds.txt`. | +| HBM traffic at N = 32 | Non-causal: FLASH 2432 B (4-row blocks) or 896 B (16-row blocks), NAIVE 12672 B. Causal: FLASH 1536 B or 768 B, NAIVE 7296 B. | +| WebAssembly | The compiled module, on golden fixtures for each schedule, matches the same quantities. 83 KB gzipped. | +| Web app | Browser tests of every page, the visualizer at phone width and in both colour schemes. | +| FPGA | Not yet synthesized (Phase 10). No timing or area numbers are claimed. | + +Every row is a test in the repository, run by GitHub Actions on each pull request. The docs page +[Verification](https://tiny-attention.vercel.app/docs/verification) lists each check and its numbers. + +## Repository + +``` +rtl/ SystemVerilog: arithmetic blocks, tile core, memory, controllers, top level +sim/ cocotb unit and end-to-end tests, the golden model (sim/golden), lint and run scripts +wasm/ Verilator + Emscripten build, the C++ harness, and the Node test over golden fixtures +web/ Astro + React site: the visualizer at /app, the docs, and the Vitest and Playwright tests +docs/ Architecture, the plan, the state schema, and the media for this README +``` ## Quick start @@ -17,10 +50,32 @@ 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/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 +bash sim/lint.sh # Verilator -Wall on every module +bash sim/run_unit.sh # cocotb unit suites +SEEDS=200 bash sim/run_e2e.sh # engine vs golden, end to end +pytest sim/golden -q # golden model, fixtures, generated tables + +bash wasm/build.sh && node wasm/test.mjs # WebAssembly engine + golden fixtures + +cd web && pnpm install && pnpm dev # the site (dev server; see note below) +pnpm build && pnpm test:e2e # production build and browser tests ``` +Note: `astro dev` hangs without output in some environments. `pnpm build` followed by any static file server over +`web/dist` works. + +## Docs + +Six pages with live demos: [what attention computes](https://tiny-attention.vercel.app/docs/what-is-attention), +[the online softmax](https://tiny-attention.vercel.app/docs/online-softmax), +[the FlashAttention schedule](https://tiny-attention.vercel.app/docs/flashattention), +[fixed-point formats and error](https://tiny-attention.vercel.app/docs/fixed-point), +[the architecture](https://tiny-attention.vercel.app/docs/architecture), and +[verification](https://tiny-attention.vercel.app/docs/verification). + +The plan and the status of each phase are in [`docs/PLAN.md`](docs/PLAN.md). The architecture, with the measured +schedule and traffic models, is in [`docs/ARCHITECTURE.md`](docs/ARCHITECTURE.md). + ## License + MIT diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index d7f3655..8b90b88 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -421,3 +421,25 @@ Gaps against the plan's Phase 8 text: 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`. + +## 14. Phase 9: docs, landing page, launch + +- **Landing page** (`/`): the hero animates the FLASH schedule for N = 16 on a causal mask (pure React, no engine), with + the three claims the site makes and links to the visualizer and the docs. Respects `prefers-reduced-motion`. +- **Docs** (`/docs/*`, six pages, plus an index). Each page has a live demo of what it explains: + - *What is attention*: float attention for a chosen preset row (the reference, not the engine). + - *Online softmax*: a FLASH run of the toy sentence recorded tile by tile (`web/src/lib/record.ts`), shown by the L1 view. + - *FlashAttention*: HBM bytes and tile counts for FLASH and NAIVE from the generated traffic table. + - *Fixed point*: an INT8 quantizer for one Q row and one K row, with the folded cfg_scale in hex; the formats table; + measured error from the generated `measured.ts`. + - *Architecture*: a static block diagram of the RTL. + - *Verification*: each check and its command, with the measured tables. +- **Generated numbers.** `sim/golden/export_web_data.py` writes `web/src/lib/measured.ts`: the per-preset error (golden + model), the recorded 1000-seed study (rows parsed from `study_results_1000seeds.txt`), and the N = 32 cycle baselines + (from `sim/perf_baseline.json`). `pytest sim/golden` fails if any generated file is stale. +- **Social card and screenshots.** `/og` is a 1200 × 630 card; `web/scripts/capture-assets.mjs` captures it to + `web/public/og.png` and the visualizer to `docs/media/visualizer.png`. The card is in the page's `og:image` meta tags. +- **Tests.** Playwright covers every page (no console errors, status 200, the islands render, the quantizer and traffic + demos respond) and every page at 375 px without horizontal scroll. Tables scroll inside their own box on phones. +- **Not in this phase:** FPGA numbers (Phase 10). The README and the Verification page say so. + diff --git a/docs/PLAN.md b/docs/PLAN.md index 19bec65..411af74 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -456,7 +456,7 @@ The estimates assume one person part-time, at tiny-cpu's pace. - 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. -### Phase 9: Docs, landing page, launch (1 week) +### Phase 9: Docs, landing page, launch (1 week) ✅ DONE (see ARCHITECTURE §14; deployed as noted there) - The six docs pages, each with an embedded live mini-visualizer. Hero capture and OG image (reuse the tiny-tpu scripts). - README in the series style: live demo, verification table (bit-exact seeds, exhaustive unit coverage, measured error, HBM savings at N = 32), and FPGA numbers. diff --git a/docs/media/visualizer.png b/docs/media/visualizer.png new file mode 100644 index 0000000..9500c98 Binary files /dev/null and b/docs/media/visualizer.png differ diff --git a/sim/golden/export_web_data.py b/sim/golden/export_web_data.py index c607394..fdf882a 100644 --- a/sim/golden/export_web_data.py +++ b/sim/golden/export_web_data.py @@ -6,18 +6,23 @@ """ import json +import re from pathlib import Path import numpy as np import fixed import presets +from float_ref import attention from traffic import traffic from export_fixtures import ROOT WEB_PRESETS = ROOT / "web" / "src" / "lib" / "presets.ts" QUANT_CASES = ROOT / "wasm" / "fixtures" / "quantize_cases.json" TRAFFIC_TS = ROOT / "web" / "src" / "lib" / "traffic-table.ts" +MEASURED_TS = ROOT / "web" / "src" / "lib" / "measured.ts" +STUDY_TXT = Path(__file__).resolve().parent / "study_results_1000seeds.txt" +PERF_JSON = ROOT / "sim" / "perf_baseline.json" def r(x): @@ -101,11 +106,71 @@ def traffic_ts(entries): f"export const TRAFFIC: Record = {body};\n") +STUDY_SEQ = re.compile( + r"N=\s*(\d+) causal=(\d) flash\s+dp max\s+([\d.]+) mean\s+([\d.]+) \| tot max\s+([\d.]+) mean\s+([\d.]+)") +STUDY_VS = re.compile( + r"^(flash|naive) N=32\s+dp max\s+([\d.]+) mean\s+([\d.]+) \| tot max\s+([\d.]+) mean\s+([\d.]+)", re.M) + + +def study_numbers(): + """Rows of the recorded 1000-seed study (sim/golden/study_results_1000seeds.txt), parsed verbatim.""" + text = STUDY_TXT.read_text() + seq = [] + for m in STUDY_SEQ.finditer(text): + n, c, dpm, dpa, tm, ta = m.groups() + seq.append(dict(n=int(n), causal=c == "1", dpMax=float(dpm), dpMean=float(dpa), totMax=float(tm), totMean=float(ta))) + vs = {} + for m in STUDY_VS.finditer(text): + mode, dpm, dpa, tm, ta = m.groups() + vs[mode] = dict(dpMax=float(dpm), dpMean=float(dpa), totMax=float(tm), totMean=float(ta)) + return dict(sequence=seq, flashVsNaive=vs, raw=text) + + +PERF_KEYS = [ + ("FLASH, Q block 4 rows, no overlap", "flash/N32/c0/q4/o0/L10"), + ("FLASH, Q block 4 rows, overlap", "flash/N32/c0/q4/o1/L10"), + ("FLASH, Q block 16 rows, no overlap", "flash/N32/c0/q16/o0/L10"), + ("FLASH, Q block 16 rows, overlap", "flash/N32/c0/q16/o1/L10"), + ("NAIVE", "naive/N32/c0/L10"), +] + + +def perf_numbers(): + base = json.loads(PERF_JSON.read_text()) + return [dict(label=label, cycles=base[key]) for label, key in PERF_KEYS] + + +def preset_errors(): + """Max |engine - float| per preset, in float units, on the golden model (the same numbers the visualizer shows).""" + out = [] + for make in presets.ALL: + p = make() + q, k, v = (np.round(np.asarray(p[x], dtype=np.float64), 6) for x in ("q", "k", "v")) + n = q.shape[0] + qz = fixed.quantize(q, k, v) + res = fixed.flash(qz, n, p["causal"], qblock=min(4, n)) + out_f = fixed.dequant_out(res, qz) + ref, _ = attention(q, k, v, p["causal"]) + # six decimals: the last bits of a float64 sum depend on the BLAS build, and the freshness test must not + out.append(dict(name=p["name"], n=n, causal=bool(p["causal"]), maxErr=round(float(np.max(np.abs(out_f - ref))), 6))) + return out + + +def measured_ts(): + study = study_numbers() + body = dict(study=dict(sequence=study["sequence"], flashVsNaive=study["flashVsNaive"]), + studyRaw=study["raw"], perfN32=perf_numbers(), presetErrors=preset_errors()) + return ("// GENERATED by sim/golden/export_web_data.py from the golden model, sim/golden/study_results_1000seeds.txt\n" + "// and sim/perf_baseline.json. Do not edit.\n" + "export const MEASURED = " + json.dumps(body, indent=1) + " as const;\n") + + def main(): WEB_PRESETS.write_text(presets_ts(preset_entries())) QUANT_CASES.write_text(json.dumps(quant_cases()) + "\n") TRAFFIC_TS.write_text(traffic_ts(traffic_entries())) - print(f"wrote {WEB_PRESETS.name} ({len(presets.ALL)} presets), {QUANT_CASES.name}, {TRAFFIC_TS.name}") + MEASURED_TS.write_text(measured_ts()) + print(f"wrote {WEB_PRESETS.name} ({len(presets.ALL)} presets), {QUANT_CASES.name}, {TRAFFIC_TS.name}, {MEASURED_TS.name}") if __name__ == "__main__": diff --git a/sim/golden/test_golden.py b/sim/golden/test_golden.py index 896842f..4e00af5 100644 --- a/sim/golden/test_golden.py +++ b/sim/golden/test_golden.py @@ -278,3 +278,4 @@ def test_web_data_fresh(): assert ewd.WEB_PRESETS.read_text() == ewd.presets_ts(ewd.preset_entries()), "presets.ts is stale" assert ewd.QUANT_CASES.read_text() == json.dumps(ewd.quant_cases()) + "\n", "quantize_cases.json is stale" assert ewd.TRAFFIC_TS.read_text() == ewd.traffic_ts(ewd.traffic_entries()), "traffic-table.ts is stale" + assert ewd.MEASURED_TS.read_text() == ewd.measured_ts(), "measured.ts is stale" diff --git a/web/e2e/docs.spec.ts b/web/e2e/docs.spec.ts new file mode 100644 index 0000000..661e30f --- /dev/null +++ b/web/e2e/docs.spec.ts @@ -0,0 +1,97 @@ +import { expect, test, type Page } from "@playwright/test"; + +const pages = [ + { path: "/", h1: "FlashAttention, built in hardware" }, + { path: "/docs", h1: "Docs" }, + { path: "/docs/what-is-attention", h1: "What is attention" }, + { path: "/docs/online-softmax", h1: "Online softmax" }, + { path: "/docs/flashattention", h1: "FlashAttention" }, + { path: "/docs/fixed-point", h1: "Fixed point" }, + { path: "/docs/architecture", h1: "Architecture" }, + { path: "/docs/verification", h1: "Verification" }, +]; + +// Collects console errors and page errors so a broken island fails the test. +function watchErrors(page: Page) { + const errors: string[] = []; + page.on("pageerror", (e) => errors.push(e.message)); + page.on("console", (m) => { + if (m.type() === "error") errors.push(m.text()); + }); + return errors; +} + +for (const p of pages) { + test(`${p.path} renders with no errors`, async ({ page }) => { + const errors = watchErrors(page); + const response = await page.goto(p.path); + expect(response?.status()).toBe(200); + await expect(page.getByRole("heading", { level: 1 })).toContainText(p.h1); + await page.waitForTimeout(500); + expect(errors, errors.join("\n")).toEqual([]); + }); +} + +test("landing page links to the engine and the docs", async ({ page }) => { + await page.goto("/"); + await expect(page.getByRole("link", { name: "Open the engine" })).toHaveAttribute("href", "/app"); + await expect(page.getByRole("link", { name: "Read the docs" })).toHaveAttribute("href", "/docs"); +}); + +test("landing hero animates the tile schedule", async ({ page }) => { + await page.goto("/"); + const grid = page.getByRole("img", { name: /FLASH tile schedule/ }); + await expect(grid).toBeVisible(); + const before = await grid.evaluate((g) => g.innerHTML); + await page.waitForTimeout(700); + const after = await grid.evaluate((g) => g.innerHTML); + expect(after).not.toEqual(before); +}); + +test("what is attention: the demo's probabilities sum to 100%", async ({ page }) => { + await page.goto("/docs/what-is-attention"); + await expect(page.getByText("Token that asks (query)")).toBeVisible(); + const text = await page.locator("main").innerText(); + const shares = [...text.matchAll(/(\d+\.\d)%/g)].map((m) => Number(m[1])); + expect(shares.length).toBeGreaterThan(0); + const sum = shares.reduce((a, b) => a + b, 0); + expect(Math.abs(sum - 100)).toBeLessThan(1); +}); + +test("online softmax: the engine records a run and the row table fills", async ({ page }) => { + await page.goto("/docs/online-softmax"); + await expect(page.getByRole("cell", { name: "kv0" })).toBeVisible({ timeout: 30_000 }); + await expect(page.getByText("m engine / float")).toBeVisible(); +}); + +test("flashattention: traffic demo shows FLASH and NAIVE byte counts", async ({ page }) => { + await page.goto("/docs/flashattention"); + await expect(page.getByText(/FLASH \(Q block 4 rows\)/)).toBeVisible(); + await expect(page.getByText(/NAIVE \(N×N scores in HBM\)/)).toBeVisible(); + // at N = 32 with causal off, FLASH with 4-row blocks moves a different total than NAIVE + await page.getByLabel("N").selectOption("32"); + await expect(page.getByText(/FLASH moves [0-9.]+× fewer bytes than NAIVE here/)).toBeVisible(); +}); + +test("fixed point: the quantizer shows the INT8 values and cfg_scale in hex", async ({ page }) => { + await page.goto("/docs/fixed-point"); + await expect(page.getByText(/= \d+ = 0x[0-9a-f]{6}/)).toBeVisible(); + await page.getByLabel("q dimension 0").fill("3.5"); + await expect(page.getByText(/int8 127/).first()).toBeVisible(); +}); + +for (const p of pages) { + test(`${p.path} has no horizontal scroll at phone width`, async ({ page }) => { + await page.setViewportSize({ width: 375, height: 812 }); + await page.goto(p.path); + await page.waitForTimeout(300); + const overflow = await page.evaluate(() => document.documentElement.scrollWidth - document.documentElement.clientWidth); + expect(overflow).toBeLessThanOrEqual(0); + }); +} + +test("social metadata is present on the landing page", async ({ page }) => { + await page.goto("/"); + await expect(page.locator('meta[property="og:image"]')).toHaveCount(1); + await expect(page.locator('meta[name="twitter:card"]')).toHaveAttribute("content", "summary_large_image"); +}); diff --git a/web/public/og.png b/web/public/og.png new file mode 100644 index 0000000..08ea54e Binary files /dev/null and b/web/public/og.png differ diff --git a/web/scripts/capture-assets.mjs b/web/scripts/capture-assets.mjs new file mode 100644 index 0000000..daf5f97 --- /dev/null +++ b/web/scripts/capture-assets.mjs @@ -0,0 +1,55 @@ +// Captures the social card and a screenshot of the visualizer from the built site. +// bash: pnpm build && node scripts/capture-assets.mjs +// Writes public/og.png (1200 x 630, from /og) and docs/media/visualizer.png (from /app, for the README). Rebuild after +// running this so public/og.png is included in dist/. +import { spawn } from "node:child_process"; +import { mkdirSync } from "node:fs"; +import { dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; +import { chromium } from "@playwright/test"; + +const webRoot = join(dirname(fileURLToPath(import.meta.url)), ".."); +const repoRoot = join(webRoot, ".."); +const port = 4340; + +const server = spawn("python3", ["-m", "http.server", String(port), "--bind", "127.0.0.1"], { + cwd: join(webRoot, "dist"), + stdio: "ignore", +}); + +async function waitForServer() { + for (let i = 0; i < 50; i++) { + try { + const r = await fetch(`http://127.0.0.1:${port}/og/`); + if (r.ok) return; + } catch { + // not up yet + } + await new Promise((res) => setTimeout(res, 200)); + } + throw new Error("static server did not start"); +} + +try { + await waitForServer(); + const browser = await chromium.launch(); + try { + const og = await browser.newPage({ viewport: { width: 1200, height: 630 } }); + await og.goto(`http://127.0.0.1:${port}/og/`); + await og.waitForLoadState("load"); + await og.screenshot({ path: join(webRoot, "public", "og.png") }); + + mkdirSync(join(repoRoot, "docs", "media"), { recursive: true }); + const app = await browser.newPage({ viewport: { width: 1200, height: 900 }, colorScheme: "dark" }); + await app.goto(`http://127.0.0.1:${port}/app/`); + await app.getByRole("button", { name: "Start" }).click({ timeout: 30_000 }); + await app.getByRole("button", { name: "Run to end" }).click(); + await app.getByRole("heading", { name: "Output" }).waitFor({ timeout: 30_000 }); + await app.screenshot({ path: join(repoRoot, "docs", "media", "visualizer.png"), fullPage: false }); + console.log("wrote web/public/og.png and docs/media/visualizer.png"); + } finally { + await browser.close(); + } +} finally { + server.kill(); +} diff --git a/web/src/components/docs/AttentionDemo.tsx b/web/src/components/docs/AttentionDemo.tsx new file mode 100644 index 0000000..33518fb --- /dev/null +++ b/web/src/components/docs/AttentionDemo.tsx @@ -0,0 +1,50 @@ +// Float attention for one row of a preset: the probabilities over the keys. No engine needed. +import { useMemo, useState } from "react"; +import { PRESETS } from "@/lib/presets"; +import { attention } from "@/lib/reference"; + +export default function AttentionDemo() { + const [preset, setPreset] = useState(0); + const p = PRESETS[preset] ?? PRESETS[0]; + const n = p?.q.length ?? 0; + const [row, setRow] = useState(0); + const labels = p?.labels ?? []; + + const probs = useMemo(() => (p ? attention(p.q, p.k, p.v, p.causal).probs : []), [p]); + const weights = probs[row] ?? []; + + return ( +
+
+ + +
+
+ {weights.map((w, j) => ( +
+ {labels[j] ?? j} +
+
+
+ {(w * 100).toFixed(1)}% +
+ ))} +
+

+ Each bar is one key's share of the query's attention. The shares sum to 100%. Causal presets give later tokens no weight. +

+
+ ); +} diff --git a/web/src/components/docs/DocsLayout.astro b/web/src/components/docs/DocsLayout.astro new file mode 100644 index 0000000..438aaa2 --- /dev/null +++ b/web/src/components/docs/DocsLayout.astro @@ -0,0 +1,51 @@ +--- +import Layout from "@/layouts/Layout.astro"; + +interface Props { + title: string; + description: string; + current?: string; +} + +const { title, description, current = "" } = Astro.props; +const links = [ + { href: "/docs/what-is-attention", label: "What is attention" }, + { href: "/docs/online-softmax", label: "Online softmax" }, + { href: "/docs/flashattention", label: "FlashAttention" }, + { href: "/docs/fixed-point", label: "Fixed point" }, + { href: "/docs/architecture", label: "Architecture" }, + { href: "/docs/verification", label: "Verification" }, +]; +--- + + +
+ +
+
+

{title}

+ +
+ + + diff --git a/web/src/components/docs/Hero.tsx b/web/src/components/docs/Hero.tsx new file mode 100644 index 0000000..bd745ba --- /dev/null +++ b/web/src/components/docs/Hero.tsx @@ -0,0 +1,58 @@ +// Landing hero: the FLASH schedule for N = 16 (Q block of 8 rows, causal) sweeping the tile grid. Pure React, no engine. +import { useEffect, useMemo, useState } from "react"; +import { flashTileOrder } from "@/lib/schedule"; + +const N = 16; +const QROWS = 8; +const STEP_MS = 180; + +export default function Hero() { + const order = useMemo(() => flashTileOrder(N, QROWS, true), []); + const [step, setStep] = useState(0); + + useEffect(() => { + const reduce = window.matchMedia("(prefers-reduced-motion: reduce)").matches; + if (reduce) { + setStep(order.length); + return; + } + const id = window.setInterval(() => setStep((s) => (s + 1) % (order.length + 8)), STEP_MS); + return () => window.clearInterval(id); + }, [order.length]); + + const blocks = N / 4; + const done = new Set(); + order.slice(0, Math.min(step, order.length)).forEach((t) => done.add(`${t.qi},${t.kj}`)); + const current = step < order.length ? order[step] : undefined; + const cur = current ? `${current.qi},${current.kj}` : ""; + + return ( +
+
+
+ {Array.from({ length: blocks }, (_, kj) => ( +
kv{kj}
+ ))} + {Array.from({ length: blocks }, (_, qi) => ( +
+
q{qi}
+ {Array.from({ length: blocks }, (_, kj) => { + const masked = kj > qi; + const key = `${qi},${kj}`; + const bg = masked ? "transparent" : key === cur ? "var(--accent)" : done.has(key) ? "var(--ok)" : "var(--panel)"; + return ( +
+ ); + })} +
+ ))} +
+

+ Each cell is a 4×4 tile of the attention matrix. Blue is computing, green is done, faded tiles are masked and never + computed. The full matrix is never stored. +

+
+ ); +} diff --git a/web/src/components/docs/OnlineSoftmaxDemo.tsx b/web/src/components/docs/OnlineSoftmaxDemo.tsx new file mode 100644 index 0000000..368ffc4 --- /dev/null +++ b/web/src/components/docs/OnlineSoftmaxDemo.tsx @@ -0,0 +1,40 @@ +// The online softmax on the toy sentence: the engine records each tile's row state; SoftmaxCellView shows it. +import { useEffect, useState } from "react"; +import SoftmaxCellView from "@/components/SoftmaxCellView"; +import { loadEngine } from "@/lib/engine"; +import { PRESETS } from "@/lib/presets"; +import { recordFlashRun } from "@/lib/record"; +import type { TileSample } from "@/lib/softmax-cell"; +import type { EngineModule } from "@/lib/state-schema"; + +const preset = PRESETS[0]!; + +export default function OnlineSoftmaxDemo() { + const [mod, setMod] = useState(null); + const [samples, setSamples] = useState(null); + const [error, setError] = useState(null); + + useEffect(() => { + loadEngine().then(setMod).catch((e: unknown) => setError(e instanceof Error ? e.message : String(e))); + }, []); + + useEffect(() => { + if (!mod) return; + try { + setSamples(recordFlashRun(mod, preset.q, preset.k, preset.v, preset.causal, 4, 10)); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } + }, [mod]); + + if (error) return

{error}

; + if (!samples) return

Running the engine on the toy sentence…

; + return ( +
+

+ The toy sentence on the engine, FLASH schedule. Pick a token to see the running state its row goes through. +

+ +
+ ); +} diff --git a/web/src/components/docs/QuantizerDemo.tsx b/web/src/components/docs/QuantizerDemo.tsx new file mode 100644 index 0000000..6ac1643 --- /dev/null +++ b/web/src/components/docs/QuantizerDemo.tsx @@ -0,0 +1,63 @@ +// Host-side INT8 quantization of one Q row and one K row, the way the engine's inputs are prepared (quantize.ts mirrors +// sim/golden/fixed.py). Shows the INT8 values, the per-tensor scales, and the folded cfg_scale in hex. +import { useMemo, useState } from "react"; +import { D, LOG2E, quantize } from "@/lib/quantize"; + +const fmtNum = (x: number) => x.toFixed(4); + +export default function QuantizerDemo() { + const [qRow, setQRow] = useState([1.2, -0.7, 2.4, 0.1]); + const [kRow, setKRow] = useState([0.9, 1.8, -2.1, 0.4]); + + const result = useMemo(() => { + const vRow = [[0.5, -0.5, 0.25, 1]]; + return quantize([qRow], [kRow], vRow); + }, [qRow, kRow]); + + const setCell = (which: "q" | "k", i: number, value: number) => { + const update = (row: number[]) => row.map((x, j) => (j === i ? value : x)); + if (which === "q") setQRow((r) => update(r)); + else setKRow((r) => update(r)); + }; + + const dq = (ints: number[], s: number) => ints.map((x) => x * s); + + return ( +
+ {(["q", "k"] as const).map((which) => { + const row = which === "q" ? qRow : kRow; + const ints = which === "q" ? result.q : result.k; + const scale = which === "q" ? result.sq : result.sk; + return ( +
+
{which === "q" ? "Q row" : "K row"}
+
+ {row.map((x, i) => ( + { const v = e.target.valueAsNumber; if (Number.isFinite(v)) setCell(which, i, v); }} /> + ))} +
+
+ {ints.map((x, i) => ( +
+
int8 {x}
+
≈ {fmtNum(dq([x], scale)[0] ?? 0)}
+
+ ))} +
+
+ scale {scale.toFixed(6)} (int8 = round(x / scale), halves to even; step error ≤ scale / 2 = {(scale / 2).toFixed(4)}) +
+
+ ); + })} +
+
cfg_scale = round(s_Q · s_K · log2(e) / √{D} · 2²⁰)
+
+ = {result.cfgScale} = 0x{result.cfgScale.toString(16).padStart(6, "0")} (24 bits, UQ4.20; scales log2(e)≈{LOG2E.toFixed(4)}) +
+
+
+ ); +} diff --git a/web/src/components/docs/TrafficDemo.tsx b/web/src/components/docs/TrafficDemo.tsx new file mode 100644 index 0000000..849faa4 --- /dev/null +++ b/web/src/components/docs/TrafficDemo.tsx @@ -0,0 +1,63 @@ +// HBM traffic for FLASH and NAIVE from the golden traffic model (web/src/lib/traffic-table.ts, generated). +// The cocotb end-to-end test checks these formulas against the RTL byte counters for every N, causal setting and block size. +import { useState } from "react"; +import { TRAFFIC } from "@/lib/traffic-table"; + +const NS = [4, 8, 12, 16, 20, 24, 28, 32]; + +export default function TrafficDemo() { + const [n, setN] = useState(32); + const [causal, setCausal] = useState(false); + const [qrows, setQrows] = useState(4); + const c = causal ? 1 : 0; + const flash = TRAFFIC[`FLASH:${n}:${c}:${qrows}`]; + const naive = TRAFFIC[`NAIVE:${n}:${c}:4`]; + const flashTotal = flash ? flash.rd + flash.wr : 0; + const naiveTotal = naive ? naive.rd + naive.wr : 0; + const max = Math.max(flashTotal, naiveTotal, 1); + + return ( +
+
+ + + +
+ + {[ + { name: "NAIVE (N×N scores in HBM)", total: naiveTotal, color: "var(--err)" }, + { name: `FLASH (Q block ${qrows} rows)`, total: flashTotal, color: "var(--accent)" }, + ].map((row) => ( +
+
+ {row.name} + {row.total} B +
+
+
+
+
+ ))} + +

+ FLASH moves {flashTotal ? (naiveTotal / flashTotal).toFixed(1) : "-"}× fewer bytes than NAIVE here. + Tiles computed: FLASH {flash?.tiles ?? "-"}, NAIVE {naive?.tiles ?? "-"}. +

+
+ ); +} diff --git a/web/src/layouts/Layout.astro b/web/src/layouts/Layout.astro index 77fb92b..30df1a1 100644 --- a/web/src/layouts/Layout.astro +++ b/web/src/layouts/Layout.astro @@ -4,12 +4,15 @@ import "@/styles/global.css"; interface Props { title: string; description?: string; + image?: string; } const { title, description = "A real FlashAttention engine in SystemVerilog, compiled to WebAssembly and running live in your browser.", + image = "/og.png", } = Astro.props; +const imageUrl = new URL(image, Astro.site ?? "https://tiny-attention.vercel.app").href; --- @@ -30,6 +33,16 @@ const { })(); + + + + + + + + + + {title} diff --git a/web/src/lib/measured.ts b/web/src/lib/measured.ts new file mode 100644 index 0000000..74557f5 --- /dev/null +++ b/web/src/lib/measured.ts @@ -0,0 +1,147 @@ +// GENERATED by sim/golden/export_web_data.py from the golden model, sim/golden/study_results_1000seeds.txt +// and sim/perf_baseline.json. Do not edit. +export const MEASURED = { + "study": { + "sequence": [ + { + "n": 4, + "causal": false, + "dpMax": 0.624, + "dpMean": 0.0665, + "totMax": 6.431, + "totMean": 0.3839 + }, + { + "n": 4, + "causal": true, + "dpMax": 0.591, + "dpMean": 0.0418, + "totMax": 8.085, + "totMean": 0.3209 + }, + { + "n": 8, + "causal": false, + "dpMax": 0.849, + "dpMean": 0.064, + "totMax": 10.942, + "totMean": 0.4172 + }, + { + "n": 8, + "causal": true, + "dpMax": 0.849, + "dpMean": 0.0486, + "totMax": 8.855, + "totMean": 0.3728 + }, + { + "n": 16, + "causal": false, + "dpMax": 1.061, + "dpMean": 0.0661, + "totMax": 9.016, + "totMean": 0.4541 + }, + { + "n": 16, + "causal": true, + "dpMax": 0.981, + "dpMean": 0.0529, + "totMax": 9.383, + "totMean": 0.4065 + }, + { + "n": 32, + "causal": false, + "dpMax": 1.26, + "dpMean": 0.0691, + "totMax": 11.904, + "totMean": 0.4722 + }, + { + "n": 32, + "causal": true, + "dpMax": 1.023, + "dpMean": 0.0563, + "totMax": 13.076, + "totMean": 0.4403 + } + ], + "flashVsNaive": { + "flash": { + "dpMax": 1.26, + "dpMean": 0.0691, + "totMax": 11.904, + "totMean": 0.4722 + }, + "naive": { + "dpMax": 1.744, + "dpMean": 0.0989, + "totMax": 11.806, + "totMean": 0.4798 + } + } + }, + "studyRaw": "== P width sweep (N=32, FLASH, 1000 seeds per cell)\nsoft pf=5 dp max 11.576 mean 0.4951 | tot max 11.309 mean 0.5591\nsoft pf=6 dp max 6.080 mean 0.2603 | tot max 5.701 mean 0.3600\nsoft pf=7 dp max 3.309 mean 0.1352 | tot max 3.869 mean 0.2752\nsoft pf=8 dp max 1.572 mean 0.0701 | tot max 4.088 mean 0.2438\npeaky pf=5 dp max 8.862 mean 0.4016 | tot max 11.935 mean 0.6607\npeaky pf=6 dp max 4.469 mean 0.2251 | tot max 11.912 mean 0.5428\npeaky pf=7 dp max 2.052 mean 0.1250 | tot max 11.904 mean 0.4913\npeaky pf=8 dp max 1.260 mean 0.0691 | tot max 11.904 mean 0.4722\nsharp pf=5 dp max 6.742 mean 0.2361 | tot max 21.089 mean 0.6762\nsharp pf=6 dp max 3.850 mean 0.1341 | tot max 20.401 mean 0.6218\nsharp pf=7 dp max 2.383 mean 0.0758 | tot max 21.257 mean 0.5989\nsharp pf=8 dp max 1.046 mean 0.0422 | tot max 20.999 mean 0.5900\n== rounding knobs (peaky, N=32, exp/z/acc round vs truncate)\nall truncate (default) dp max 1.260 mean 0.0691 | tot max 11.904 mean 0.4722\nexp round dp max 1.260 mean 0.0685 | tot max 11.900 mean 0.4721\nz round dp max 1.260 mean 0.0690 | tot max 11.904 mean 0.4720\nacc round dp max 1.260 mean 0.0682 | tot max 11.896 mean 0.4720\nall round dp max 1.260 mean 0.0677 | tot max 11.896 mean 0.4718\n== recip refinement (peaky, N=32)\nnr_steps=0 dp max 1.437 mean 0.1264 | tot max 11.892 mean 0.4910\nnr_steps=1 dp max 1.260 mean 0.0691 | tot max 11.904 mean 0.4722\nnr_steps=2 dp max 1.260 mean 0.0692 | tot max 11.904 mean 0.4723\n== l/O extra fractional bits (peaky, N=32)\nlf=of=2 dp max 1.260 mean 0.0723 | tot max 11.912 mean 0.4728\nlf=of=4 dp max 1.260 mean 0.0691 | tot max 11.904 mean 0.4722\nlf=of=6 dp max 1.260 mean 0.0684 | tot max 11.900 mean 0.4721\nlf=of=8 dp max 1.260 mean 0.0683 | tot max 11.900 mean 0.4720\n== cfg_scale fractional bits sf (cfg_scale is 24 bits; N=32)\nsmall sf=12 overflow: cfg_scale 0 does not fit 24 bits at sf=12 (inputs badly scaled)\nsmall sf=16 dp max 1.143 mean 0.0631 | tot max 1.101 mean 0.0798\nsmall sf=20 dp max 0.080 mean 0.0106 | tot max 0.223 mean 0.0424\nsmall sf=24 dp max 0.056 mean 0.0095 | tot max 0.219 mean 0.0421\nsoft sf=12 dp max 19.778 mean 1.1309 | tot max 19.950 mean 1.1711\nsoft sf=16 dp max 2.414 mean 0.1072 | tot max 4.147 mean 0.2580\nsoft sf=20 dp max 1.572 mean 0.0701 | tot max 4.088 mean 0.2438\nsoft sf=24 dp max 1.572 mean 0.0700 | tot max 3.979 mean 0.2439\nsharp sf=12 dp max 2.801 mean 0.1098 | tot max 20.401 mean 0.6020\nsharp sf=16 dp max 1.046 mean 0.0430 | tot max 20.999 mean 0.5901\nsharp sf=20 dp max 1.046 mean 0.0422 | tot max 20.999 mean 0.5900\nsharp sf=24 dp max 1.046 mean 0.0422 | tot max 20.999 mean 0.5901\n== sequence length and causal (default fmt)\nN= 4 causal=0 flash dp max 0.624 mean 0.0665 | tot max 6.431 mean 0.3839\nN= 4 causal=1 flash dp max 0.591 mean 0.0418 | tot max 8.085 mean 0.3209\nN= 8 causal=0 flash dp max 0.849 mean 0.0640 | tot max 10.942 mean 0.4172\nN= 8 causal=1 flash dp max 0.849 mean 0.0486 | tot max 8.855 mean 0.3728\nN=16 causal=0 flash dp max 1.061 mean 0.0661 | tot max 9.016 mean 0.4541\nN=16 causal=1 flash dp max 0.981 mean 0.0529 | tot max 9.383 mean 0.4065\nN=32 causal=0 flash dp max 1.260 mean 0.0691 | tot max 11.904 mean 0.4722\nN=32 causal=1 flash dp max 1.023 mean 0.0563 | tot max 13.076 mean 0.4403\n== FLASH vs NAIVE (peaky, default fmt)\nflash N=32 dp max 1.260 mean 0.0691 | tot max 11.904 mean 0.4722\nnaive N=32 dp max 1.744 mean 0.0989 | tot max 11.806 mean 0.4798\n", + "perfN32": [ + { + "label": "FLASH, Q block 4 rows, no overlap", + "cycles": 4113 + }, + { + "label": "FLASH, Q block 4 rows, overlap", + "cycles": 2937 + }, + { + "label": "FLASH, Q block 16 rows, no overlap", + "cycles": 2133 + }, + { + "label": "FLASH, Q block 16 rows, overlap", + "cycles": 1657 + }, + { + "label": "NAIVE", + "cycles": 10961 + } + ], + "presetErrors": [ + { + "name": "toy_sentence", + "n": 12, + "causal": true, + "maxErr": 0.004894 + }, + { + "name": "previous_token", + "n": 16, + "causal": true, + "maxErr": 0.010519 + }, + { + "name": "attention_sink", + "n": 16, + "causal": true, + "maxErr": 0.02535 + }, + { + "name": "uniform", + "n": 16, + "causal": false, + "maxErr": 0.002773 + }, + { + "name": "self", + "n": 16, + "causal": false, + "maxErr": 0.025997 + }, + { + "name": "random", + "n": 16, + "causal": false, + "maxErr": 0.016339 + } + ] +} as const; diff --git a/web/src/lib/record.ts b/web/src/lib/record.ts new file mode 100644 index 0000000..20d2336 --- /dev/null +++ b/web/src/lib/record.ts @@ -0,0 +1,47 @@ +// Records a whole FLASH run tile by tile (the row state at each tile_done pulse), for the docs demos. Same observation +// rule as the visualizer: the tile completed on an edge is order[tiles - 1] of that run. +import { quantize } from "@/lib/quantize"; +import { flashTileOrder } from "@/lib/schedule"; +import type { EngineModule } from "@/lib/state-schema"; +import type { TileSample } from "@/lib/softmax-cell"; + +const MAX_CYCLES = 400_000; + +export function recordFlashRun( + mod: EngineModule, + q: number[][], + k: number[][], + v: number[][], + causal: boolean, + qrows: number, + latency: number, +): TileSample[] { + const n = q.length; + const qz = quantize(q, k, v); + const order = flashTileOrder(n, qrows, causal); + const sim = new mod.TinyAttnSim(); + const samples: TileSample[] = []; + const observe = (s: { tileDone: boolean; tiles: number }) => { + if (!s.tileDone) return; + const ref = order[s.tiles - 1]; + if (!ref) return; + const t = sim.getTile(); + samples.push({ qi: ref.qi, kj: ref.kj, z: t.z, m: t.m, l: t.l, alpha: t.alpha }); + }; + try { + sim.configure(n, qrows, causal, 0, false, latency, qz.cfgScale); + sim.loadQKV(qz.q, qz.k, qz.v); + sim.start(); + for (let i = 0; i < MAX_CYCLES; i++) { + const s = sim.step(); + observe(s); + if (s.done) { + observe(sim.step()); + break; + } + } + } finally { + sim.delete(); + } + return samples; +} diff --git a/web/src/pages/docs/architecture.astro b/web/src/pages/docs/architecture.astro new file mode 100644 index 0000000..9f8bade --- /dev/null +++ b/web/src/pages/docs/architecture.astro @@ -0,0 +1,90 @@ +--- +import DocsLayout from "@/components/docs/DocsLayout.astro"; +--- + + +

+ The top level is tiny_attn_top. A host loads Q, K and V into a model of off-chip memory, starts the engine, + and reads back the output. Two interchangeable controllers drive one tile core. One selects the FLASH schedule, and the + other the NAIVE schedule. +

+ + + + host + + hbm_model + words, byte counters + + dma + bursts, latency + + flash_ctrl + tiles, prefetch + + naive_ctrl + four passes + + K/V buffers + double-buffered + + attn_tile + + MM1 S = Q·Kᵀ (4×4 PEs) + + scale, mask, online softmax + + MM2 P·V (4×4 PEs) + O accumulator, normalize + + sram_bank + output staging + + perf_counters + cycles, bytes, tiles + + + + + + + + + + + + +

Memory

+
    +
  • hbm_model is a 2,048-word memory: Q, K and V rows, the output, and the NAIVE score and probability tiles. It counts the bytes read and written.
    • +
  • dma moves one burst at a time, with a configurable first-word latency (0 to 64 cycles) and one word per cycle after it.
    • +
  • K/V buffers are two K tiles and two V tiles, each 4 rows. Without overlap only the first pair is used. With overlap, the second pair is filled while the first is in use.
    • +
  • sram_bank stages output rows before the DMA writes them.
    • +
+ +

Controllers

+

+ flash_ctrl loops over Q blocks and, inside each block, over the K/V tiles. It issues the DMA loads, the + weight loads and the group tiles, skips tiles above the causal diagonal, and normalizes and stores each block. + naive_ctrl runs the four passes of the standard softmax over each row block, writing and reading the + score and probability tiles through HBM. +

+ +

Tile core

+

+ attn_tile processes one 4-row group against one K/V tile. The two 4×4 systolic arrays are the matrix + products. Between them sit the scale-and-mask unit, the row maximum, the exp2 unit and the sum. The O accumulator + rescales the running output by α and adds the new product. Five operations share the datapath: FLASH (online), + and the four NAIVE passes. +

+ +

Timing

+

+ The schedule is exact: busy cycles follow the model in sim/golden/schedule.py, and the end-to-end test + checks them for every configuration it runs. Overlap changes the schedule, so its cycle counts are pinned as + regression baselines in sim/perf_baseline.json. +

+

+ The design has not been synthesized for an FPGA yet (Phase 10). The numbers on this site come from simulation. +

+ diff --git a/web/src/pages/docs/fixed-point.astro b/web/src/pages/docs/fixed-point.astro new file mode 100644 index 0000000..1f70240 --- /dev/null +++ b/web/src/pages/docs/fixed-point.astro @@ -0,0 +1,79 @@ +--- +import DocsLayout from "@/components/docs/DocsLayout.astro"; +import QuantizerDemo from "@/components/docs/QuantizerDemo"; +import { MEASURED } from "@/lib/measured"; + +const presetRows = MEASURED.presetErrors; +--- + + +

+ The engine is integer hardware. Every value has a fixed format, and the formats were chosen by a numerics study before any RTL was written. + The golden model in sim/golden/fixed.py uses the same integers, so the RTL can be checked bit for bit. +

+ +

Formats

+ + + + + + + + + + + + + + + + +
ValueFormatNotes
Q, K, Vint8Symmetric per-tensor scale, set on the host (see below)
S = Q·Kᵀint17Signed dot product of int8 values
cfg_scaleUQ4.20, 24 bitsFolds s_Q · s_K · log₂e / √d into one multiplier
Z, mQ8.8, signedScaled scores in log₂ units, and their running maximum
Z − mQ9.8, signedAlways ≤ 0; the input of exp2
2^(Z − m), αUQ1.15exp2 output and the rescale factor
PUQ1.8, 9 bitsProbabilities, 0 to 256
P·V (one tile)int18, signedMatrix product in the second array
l (running sum)18 bits, unsignedReal value × 2¹²
O (accumulator)26 bits, signedReal value × 2⁴, in V units
ReciprocalUQ1.16 mantissa, 5-bit exponentUsed once per row, at the end
OutputQ8.8, signedDequantize with s_V: value = out / 256 · s_V
+ +

Quantization

+

+ The host picks one scale per tensor: s = max|x| / 127, then int8 = round(x / s) with ties to + even, clipped to ±127. The scales of Q and K go into cfg_scale, so the datapath never multiplies by a + scale. Try it: +

+ + +

The exp2 unit

+

+ 2^(x/256) for x ≤ 0 is split into an integer shift and a fraction. The fraction indexes a + 256-entry table of 16-bit values. The integer part is a right shift. Results that shift past 16 bits are exactly zero, + and the output is truncated, not rounded. The table is generated by sim/golden/gen_luts.py, and CI checks + that the committed table is up to date. +

+

The unit is tested exhaustively over its input range, so there is no sampling gap.

+ +

Measured error

+

+ Errors are against float attention on the same inputs, in float units. The first table is for each preset, with the + engine's output after dequantization. The second is from the recorded 1000-seed study at N = 32, in V units (one int8 + step of V). The study separates the datapath error (dp, on the INT8 inputs) from the total error + (tot, which includes input quantization). +

+ + + + {presetRows.map((r) => ( + + ))} + +
PresetNCausalMax |engine − float|
{r.name}{r.n}{r.causal ? "yes" : "no"}{r.maxErr.toFixed(4)}
+ + + + {Object.entries(MEASURED.study.flashVsNaive).map(([mode, r]) => ( + + ))} + +
Study, N = 32, peaky inputsdp maxdp meantot maxtot mean
{mode === "flash" ? "FLASH" : "NAIVE"}{r.dpMax.toFixed(3)}{r.dpMean.toFixed(4)}{r.totMax.toFixed(3)}{r.totMean.toFixed(4)}
+

+ For the default format at N = 32 (peaky inputs), the datapath error averages 0.07 of a V step for FLASH and 0.10 for + NAIVE. The largest errors are 1.26 steps for FLASH and 1.74 for NAIVE. The total error is larger because it also includes the rounding of + the inputs to INT8. The full study output is recorded in sim/golden/study_results_1000seeds.txt. +

+ diff --git a/web/src/pages/docs/flashattention.astro b/web/src/pages/docs/flashattention.astro new file mode 100644 index 0000000..a10f1ec --- /dev/null +++ b/web/src/pages/docs/flashattention.astro @@ -0,0 +1,36 @@ +--- +import DocsLayout from "@/components/docs/DocsLayout.astro"; +import TrafficDemo from "@/components/docs/TrafficDemo"; +--- + + +

+ FlashAttention computes attention in tiles. A block of query rows stays on chip. The engine streams key and value tiles + past it, one 4-row tile at a time, folds each tile into the running state of + online softmax, and writes the output once at the end. The N × N score + matrix never exists in memory. +

+

The schedule

+
    +
  • Q rows are split into blocks. A block holds cfg_qrows rows (4, 8, 16 or 32 here).
    • +
  • For each block, every key and value tile is loaded once and used by all the row groups in the block.
    • +
  • With the causal mask, tiles above the diagonal are skipped, and the K/V loads stop at the block's last needed tile.
    • +
  • The output of a block is normalized and written once.
    • +
+

The traffic

+

+ Each K or V tile is 16 bytes (4 rows of 4 int8 values), each Q tile is 16 bytes, and each O tile is 32 bytes (4 rows of + 4 int16 values). With nb = N / 4 row groups: +

+
    +
  • FLASH: reads 16·nb bytes of Q, plus 32 bytes per K/V tile load. Writes 32·nb bytes of O. A larger Q block means fewer K/V loads.
    • +
  • NAIVE: writes and reads the score tiles Z (32 bytes each) and the probability tiles P (32 bytes each) through HBM, once per pass.
    • +
+

+ The NAIVE schedule is the standard three-pass softmax: pass 1 writes the scores, pass 2 finds the row maxima, pass 3 + writes the probabilities, and pass 4 multiplies them by V. +

+

Compare the bytes

+

The tables come from the golden traffic model, which the cocotb end-to-end test checks against the RTL byte counters.

+ + diff --git a/web/src/pages/docs/index.astro b/web/src/pages/docs/index.astro new file mode 100644 index 0000000..ee9916a --- /dev/null +++ b/web/src/pages/docs/index.astro @@ -0,0 +1,15 @@ +--- +import DocsLayout from "@/components/docs/DocsLayout.astro"; +--- + + +

Six pages, each with a live demo of the part it explains.

+
    +
  • What is attention: queries, keys and values, and the score matrix.
    • +
  • Online softmax: how a softmax over a row can be computed a tile at a time.
    • +
  • FlashAttention: the tile schedule, and the memory traffic it saves.
    • +
  • Fixed point: the number formats, the exp2 unit, and the measured error.
    • +
  • Architecture: the blocks of the RTL and how they connect.
    • +
  • Verification: what is tested, how, and the results.
    • +
+ diff --git a/web/src/pages/docs/online-softmax.astro b/web/src/pages/docs/online-softmax.astro new file mode 100644 index 0000000..d783cd1 --- /dev/null +++ b/web/src/pages/docs/online-softmax.astro @@ -0,0 +1,32 @@ +--- +import DocsLayout from "@/components/docs/DocsLayout.astro"; +import OnlineSoftmaxDemo from "@/components/docs/OnlineSoftmaxDemo"; +--- + + +

+ A softmax over a row needs the row's maximum before it can exponentiate. The usual recipe makes three passes over the + row: find the maximum, sum the exponentials of the shifted values, then divide. A tiled engine sees the row a few keys + at a time, so it cannot do that without storing the whole row. +

+

The online softmax keeps two running values for each row: the maximum m so far, and the sum l of exponentials relative to it. When a new tile brings keys with a larger maximum, the old sum was relative to a smaller maximum, so it must shrink. It is multiplied by

+

α = exp(m_old − m_new)

+

+ and the same factor rescales the output accumulator O, which is a sum of exponentials times values. After + the last tile, O / l is the exact attention output for that row. Nothing about the row has to be kept + except m, l and O. +

+

In the engine

+

+ The engine works in log base 2, so it uses 2^z in place of exp. The scores are scaled to + z = s · log₂e / √d and the rescale is α = 2^(m_old − m_new). The power of two is computed by + the exp2 unit (see Fixed point), and α is stored as an unsigned 16-bit + fraction. +

+

Watch one row

+

+ The demo runs the engine on the toy sentence, records the row state after each K/V tile, and shows it next to the float + softmax of the same inputs. The difference is fixed-point rounding only. +

+ + diff --git a/web/src/pages/docs/verification.astro b/web/src/pages/docs/verification.astro new file mode 100644 index 0000000..b3698d6 --- /dev/null +++ b/web/src/pages/docs/verification.astro @@ -0,0 +1,105 @@ +--- +import DocsLayout from "@/components/docs/DocsLayout.astro"; +import { MEASURED } from "@/lib/measured"; + +const seq = MEASURED.study.sequence; +--- + + +

+ Every claim on this site is checked by a test in the repository. The golden model is the reference: a Python program + with the same integer formats, which the RTL is compared against on every run. +

+ +

What is checked

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
LayerWhat it checksHow it runs
Unit blocksEach arithmetic block against the golden model, including exhaustive sweeps of the exp2 and reciprocal inputs, and the DMA burst timing cycle by cycle.bash sim/run_unit.sh (cocotb with Verilator)
Tile coreOne tile against the golden trace, including causal masking and multi-group blocks.sim/test_attn_tile.py
Engine, end to endOutput bit for bit; HBM read and write bytes; tile count; K/V loads; busy cycles against the closed-form model. Every sequence length 4 to 32, causal and not, all block sizes, both schedules, overlap on and off, and random HBM latency. Random seeds on top.SEEDS=200 bash sim/run_e2e.sh
Golden modelQuantization, exp2 and reciprocal, both schedules against float attention, traffic formulas, presets, and that the committed LUT table matches its generator.pytest sim/golden
LintVerilator with all warnings on, for every module and the full design.bash sim/lint.sh
WebAssemblyThe compiled module, on golden fixtures for each schedule: output, bytes, tiles, cycles, the last tile's state, probabilities, and the K/V buffer in use.node wasm/test.mjs
Web appQuantization and tile order against the shared fixtures, the visualizer and each docs page in a real browser, at phone width and in both colour schemes.pnpm test, pnpm test:e2e
+

All of these run in GitHub Actions on every pull request. The workflows are rtl.yml, wasm.yml and web.yml.

+ +

Output error

+

Maximum error of the engine's output against float attention, per preset, in float units. These numbers come from the golden model, which the RTL matches bit for bit.

+ + + + {MEASURED.presetErrors.map((r) => ( + + ))} + +
PresetNCausalMax error
{r.name}{r.n}{r.causal ? "yes" : "no"}{r.maxErr.toFixed(4)}
+ +

Accuracy study

+

+ The recorded study (sim/golden/study_results_1000seeds.txt), in V units (one int8 step of V). dp is the datapath + error on the INT8 inputs. tot is the error against float attention on the original inputs. +

+ + + + {seq.map((r) => ( + + + + + + ))} + +
NCausaldp maxdp meantot maxtot mean
{r.n}{r.causal ? "yes" : "no"}{r.dpMax.toFixed(3)}{r.dpMean.toFixed(4)}{r.totMax.toFixed(3)}{r.totMean.toFixed(4)}
+ +

Cycles at N = 32

+

+ Busy cycles at N = 32, non-causal, with 10 cycles of HBM latency. The sequential FLASH rows match the closed-form + schedule exactly. The overlap and NAIVE rows are pinned baselines from the RTL, and the test fails if they change. +

+ + + + {MEASURED.perfN32.map((r) => ( + + ))} + +
ConfigurationBusy cycles
{r.label}{r.cycles}
+ +

Not verified yet

+
    +
  • Synthesis for an FPGA. The timing and area numbers are not measured yet (Phase 10).
    • +
  • Real hardware. Everything here runs in simulation or in WebAssembly.
    • +
+ diff --git a/web/src/pages/docs/what-is-attention.astro b/web/src/pages/docs/what-is-attention.astro new file mode 100644 index 0000000..f492d75 --- /dev/null +++ b/web/src/pages/docs/what-is-attention.astro @@ -0,0 +1,33 @@ +--- +import DocsLayout from "@/components/docs/DocsLayout.astro"; +import AttentionDemo from "@/components/docs/AttentionDemo"; +--- + + +

+ Attention lets each token in a sequence take information from the others. For token i it builds a + query vector q_i. Every token j offers a key k_j and a value v_j. + The score of j for i is the dot product q_i · k_j, scaled by + 1/√d. A softmax turns each row of scores into weights that sum to one. The output for token + i is the weighted sum of the values. +

+

In matrix form, with N tokens of dimension d:

+

O = softmax(Q Kᵀ / √d) V

+

+ The middle term, S = Q Kᵀ, is an N × N matrix. For a 2,048-token context that is about four + million scores for every head. Storing it is what makes naive attention expensive in memory, and it is the reason + FlashAttention exists. +

+

Causal masks

+

+ In a language model, token i must not see later tokens. A causal mask sets the scores for j > i + to minus infinity before the softmax, so their weights are zero. The engine has a causal mode; with it, half of the + score matrix is never computed. +

+

Try it

+

+ The demo uses the float reference, not the engine. Pick a preset and a query token. In the toy sentence, the pronoun "it" gives + "cat" about 0.66 of its weight and "mouse" about 0.15. +

+ + diff --git a/web/src/pages/index.astro b/web/src/pages/index.astro index bba47d0..c32f390 100644 --- a/web/src/pages/index.astro +++ b/web/src/pages/index.astro @@ -1,16 +1,52 @@ --- import Layout from "@/layouts/Layout.astro"; +import Hero from "@/components/docs/Hero"; --- - -
-

TinyAttention

-

- A real, synthesizable FlashAttention engine in SystemVerilog, compiled to WebAssembly and - running live in your browser. Under construction. -

-

- Toolchain spike → -

+ +
+
+

TinyAttention

+

FlashAttention, built in hardware and run in your browser.

+

+ A synthesizable SystemVerilog engine for attention. It computes the output one tile at a time, so the N×N score + matrix never goes to memory. The RTL is verified bit-exact against a golden model, and the same RTL runs here as + WebAssembly. +

+
+ Open the engine + Read the docs +
+
+ +
+

The schedule

+ +
+ +
+
+
Bit-exact
+

Output, HBM traffic, tile counts and cycle counts match the golden model on every run in the test sweep.

+
+
+
No N×N matrix
+

At N = 32, FLASH moves 1536 bytes against 7296 for the naive schedule, with the same result.

+
+
+
In the browser
+

The Verilator model is compiled to an 83 KB (gzip) WebAssembly module. You step it clock by clock.

+
+
+ +
+

+ Start with what attention computes, + then the online softmax that makes the tiling possible. +

+
diff --git a/web/src/pages/og.astro b/web/src/pages/og.astro new file mode 100644 index 0000000..97b8ac3 --- /dev/null +++ b/web/src/pages/og.astro @@ -0,0 +1,53 @@ +--- +// The social card, 1200 x 630. Captured to public/og.png by web/scripts/capture-assets.mjs; not linked from the site. +import { flashTileOrder } from "@/lib/schedule"; + +const N = 16; +const QROWS = 8; +const order = flashTileOrder(N, QROWS, true); +const blocks = N / 4; +const done = new Set(order.map((t) => `${t.qi},${t.kj}`)); +--- + + + + + + + TinyAttention + + + +
+
+
TinyAttention
+

FlashAttention in SystemVerilog

+

Bit-exact against a golden model. Compiled to WebAssembly. Run it in your browser.

+
+ +
+ + diff --git a/web/src/styles/global.css b/web/src/styles/global.css index 8a26f44..ad36c9f 100644 --- a/web/src/styles/global.css +++ b/web/src/styles/global.css @@ -9,6 +9,7 @@ --ok: #3fb950; --err: #f85149; --panel: #161b22; + --line: #30363d; } :root[data-theme="light"] { @@ -19,6 +20,7 @@ --ok: #1a7f37; --err: #cf222e; --panel: #f6f8fa; + --line: #d0d7de; } body {