diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 8a0cc63..c3822eb 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -30,7 +30,9 @@ // Vulkan loader + software fallback; scripts/setup_nvidia_vulkan.sh adds // the NVIDIA ICD manifest, the GLVND EGL vendor manifest, the modeset // node, and fixes /dev/dri permissions. - "postCreateCommand": "sudo apt-get update && sudo apt-get install -y libvulkan1 vulkan-tools mesa-vulkan-drivers libegl1 libgl1 && bash scripts/setup_nvidia_vulkan.sh || true", + // install_hooks.sh points core.hooksPath at .githooks/ — the local CI gate + // (see scripts/ci_local.sh). git does not wire that up per clone on its own. + "postCreateCommand": "sudo apt-get update && sudo apt-get install -y libvulkan1 vulkan-tools mesa-vulkan-drivers libegl1 libgl1 && bash scripts/install_hooks.sh; bash scripts/setup_nvidia_vulkan.sh || true", // /dev/dri is re-mounted with the host's restrictive perms on every // container start (not just creation), so the DAC/GLVND fixes need to // re-run on every start too, not only once via postCreateCommand. diff --git a/.githooks/pre-commit b/.githooks/pre-commit new file mode 100755 index 0000000..2ee29c4 --- /dev/null +++ b/.githooks/pre-commit @@ -0,0 +1,18 @@ +#!/usr/bin/env bash +# Fast local CI gate: fmt, clippy, the wasm32 and forge-top builds, and the TUI +# dependency-leak assert. About 6s with a warm target/. +# +# The release test suite runs on pre-push instead (see .githooks/pre-push) — +# it takes about a minute, too slow to pay on every WIP commit or --amend. +# +# Activated by scripts/install_hooks.sh, which points core.hooksPath here. +# Bypass with `git commit --no-verify`. +set -euo pipefail + +# This gate checks the working tree, not the staged snapshot, so a partially +# staged change can commit a state that was never checked on its own. +if ! git diff --quiet; then + printf '\033[33mnote:\033[0m unstaged changes present — this gate checks the working tree, not the index\n' >&2 +fi + +exec ./scripts/ci_local.sh fast diff --git a/.githooks/pre-push b/.githooks/pre-push new file mode 100755 index 0000000..2eda9b5 --- /dev/null +++ b/.githooks/pre-push @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +# Full local CI gate, run before anything leaves the machine: everything +# pre-commit checks, plus `cargo test --release --locked` — the whole suite, +# including gpt2_e2e and kv_cache against real GPU hardware and the HF golden +# fixture. +# +# That is a superset of what .github/workflows/ci.yml used to verify: the +# GitHub runners had no GPU (software Vulkan only) and skipped the suites +# needing models/gpt2/, since the 548 MB of weights are correctly gitignored. +# +# `full` re-runs the pre-commit checks on purpose — a push can carry commits +# made with --no-verify, or fetched from another machine. +# +# Activated by scripts/install_hooks.sh. Bypass with `git push --no-verify`. +set -euo pipefail + +# git feeds one " " line per +# pushed ref on stdin, with an all-zero local sha for a deletion. A +# deletion-only push (`git push origin :branch`) changes no code, so there is +# nothing to verify. An empty stdin means the hook was invoked by hand — run. +refs=0 +updates=0 +while read -r _local_ref local_sha _remote_ref _remote_sha; do + refs=$((refs + 1)) + [[ "$local_sha" =~ ^0+$ ]] || updates=$((updates + 1)) +done +if [[ "$refs" -gt 0 && "$updates" -eq 0 ]]; then + echo "pre-push: deletions only, nothing to verify" + exit 0 +fi + +exec ./scripts/ci_local.sh full diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fba5ade..5351bbb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,16 +1,25 @@ -name: CI +name: CI (manual) +# The primary gate for this repo is local, not here: .githooks/pre-commit runs +# fmt/clippy/builds and .githooks/pre-push runs the full release test suite +# before anything is pushed. See scripts/ci_local.sh, which both hooks and the +# jobs below invoke, and CONTRIBUTING.md for the one-time install. +# +# This workflow therefore has no push or pull_request trigger — it only runs on +# manual dispatch. It is kept for the one case the local hooks cannot cover: +# a pull request from a fork, where nobody's pre-push hook ran on the code. +# # Scope note: GitHub-hosted runners have no NVIDIA GPU. The WGPU-backed tests # below therefore execute against Mesa's software Vulkan driver (lavapipe), -# which exercises the real WGSL kernel path but not real GPU hardware. -# GPU-backed parity — the published CPU<->WGPU max logit diff of 8.4e-5 and the -# Forge<->HF transformers diff of 1.75e-4 — is verified locally on an NVIDIA -# RTX A5000, not in CI. See scripts/setup_nvidia_vulkan.sh. +# which exercises the real WGSL kernel path but not real GPU hardware. Runners +# also lack the gitignored 548 MB models/gpt2/, so gpt2_e2e and kv_cache are +# excluded here. Both gaps are covered by the pre-push hook on a machine with +# an NVIDIA RTX A5000 and the weights present — the published CPU<->WGPU max +# logit diff of 8.4e-5 and the Forge<->HF transformers diff of 1.75e-4 come +# from there. See scripts/setup_nvidia_vulkan.sh. on: - push: - branches: [main] - pull_request: + workflow_dispatch: env: CARGO_TERM_COLOR: always @@ -79,6 +88,7 @@ jobs: # The TUI deps must never reach the library's dependents or the wasm # build; they are optional, behind the `tui` feature, for that reason. + # Kept in sync with assert_no_tui_deps in scripts/ci_local.sh. - name: Assert the default build pulls no TUI dependencies run: | for dep in ratatui crossterm sysinfo nvml-wrapper memmap2; do diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5928800..1eb7e70 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -75,6 +75,10 @@ add a new op, add it to both backends and add a parity case. ## Environment setup ```bash +# One-time: activate the local CI hooks (see "The local CI gate" below). +# Already run for you in the devcontainer. +./scripts/install_hooks.sh + # Fetch GPT-2 124M weights + tokenizer into models/gpt2/ # (reads HF_TOKEN from .env if set; gpt2 is public so this is optional) ./scripts/download_gpt2.sh @@ -89,6 +93,51 @@ you're running in a container with an NVIDIA GPU and `wgpu::Device` initialization fails to find a hardware adapter (falls back to software rendering via Mesa's llvmpipe otherwise, which works but is slow). +## The local CI gate + +Verification runs on your machine, in git hooks, rather than on GitHub. The +only remaining workflow that runs automatically is the Pages deploy +(`.github/workflows/pages.yml`); `.github/workflows/ci.yml` is dispatch-only, +kept for pull requests from forks where nobody's hooks ran. + +The reason is that the GitHub runners were verifying strictly less than a +developer machine can. They have no GPU, so every WGPU test ran against Mesa's +software Vulkan driver, and they don't have the gitignored 548 MB +`models/gpt2/`, so `gpt2_e2e` and `kv_cache` — the suites that check real +GPT-2 numerics against HF `transformers` — were skipped entirely. The pre-push +hook runs both, on real hardware. + +`scripts/ci_local.sh` is the single source of truth for what "green" means, +and both hooks are thin wrappers around it: + +| stage | checks | cost (warm `target/`) | hook | +| --- | --- | --- | --- | +| `fast` | `cargo fmt --check`, `cargo clippy -D warnings`, wasm32 build, `forge-top` build, TUI dependency-leak assert | ~6s | `pre-commit` | +| `full` | everything in `fast`, plus `cargo test --release --locked` (all suites) | ~1m10s | `pre-push` | + +Stages are ordered cheapest-first and stop at the first failure, so a +formatting slip doesn't cost you a minute of GPU tests. Run either by hand: + +```bash +./scripts/ci_local.sh fast +./scripts/ci_local.sh full +``` + +`full` deliberately repeats the `fast` checks — a push can carry commits made +with `--no-verify`, or fetched from another machine. + +Two things worth knowing: + +- **Activation is per clone.** git ignores `.githooks/` until + `core.hooksPath` points at it, which is what `./scripts/install_hooks.sh` + does. Undo with `git config --unset core.hooksPath`. +- **`pre-commit` checks the working tree, not the index.** With a partially + staged change it verifies a different state than the one being committed, + and warns when it notices unstaged changes. `pre-push` has no such gap. + +To commit or push a knowingly-broken WIP state, bypass with +`git commit --no-verify` / `git push --no-verify`. + ## Testing Run the full suite before sending a change: @@ -165,7 +214,7 @@ If you're touching the browser/wasm path, build and serve the demo: 2. If you add or modify an op: implement it in `backend/cpu.rs`, add/update the matching WGSL kernel in `shaders/`, wire it through `ops.rs`, and add a parity case in `tests/op_parity.rs`. -3. Run `cargo test --release` and, if relevant, the manual CPU/WGPU - generation comparison above. +3. Run `./scripts/ci_local.sh full` (or just let the `pre-push` hook do it) + and, if relevant, the manual CPU/WGPU generation comparison above. 4. Keep changes scoped to what GPT-2 needs — this project deliberately avoids generality for its own sake. diff --git a/Cargo.toml b/Cargo.toml index 3ea0395..cf07ecd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,6 +44,11 @@ sysinfo = { version = "0.39", optional = true } nvml-wrapper = { version = "0.12", optional = true } memmap2 = { version = "0.9", optional = true } +# Native only: the integration tests exercise the async inference API — the +# path the browser actually takes — so they need a blocking executor. +[target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies] +pollster = "0.4" + # Browser (wasm32): async-only device APIs + JS bindings. [target.'cfg(target_arch = "wasm32")'.dependencies] wasm-bindgen = "0.2" diff --git a/docs/src/demo.js b/docs/src/demo.js index 0bec517..0988603 100644 --- a/docs/src/demo.js +++ b/docs/src/demo.js @@ -29,6 +29,57 @@ function explain(title, body, retry) { } } +// ── The 3D stack ────────────────────────────────────────────────────────── +// Separate from the WebGPU check on purpose: WebGL and WebGPU fail +// independently, and either one missing must still leave a complete page. + +let scenePromise = null; + +/** Start the scene at most once; resolves to the controller or null. */ +function ensureScene() { + scenePromise = scenePromise || startScene(); + return scenePromise; +} + +async function startScene() { + const canvas = $("scene"); + if (!canvas) return null; + try { + // three.js is 751 KB and lives behind this call, so it is fetched when + // the section is reached rather than on first paint. + const { createStack } = await import("./scene.js"); + return createStack({ canvas, label: $("scene-label") }); + } catch (e) { + // No WebGL, or the module itself failed to load. Drop the canvas + // entirely — an empty rectangle is worse than no rectangle — and open + // the text architecture, which says the same thing in words. + console.warn("3D stack unavailable:", e); + $("scene-card")?.remove(); + $("demo-grid")?.classList.remove("md:grid-cols-2"); + const text = $("stack-text"); + if (text) text.open = true; + return null; + } +} + +const section = $("demo"); +if (section && "IntersectionObserver" in window) { + const io = new IntersectionObserver( + (entries, obs) => { + if (entries.some((e) => e.isIntersecting)) { + obs.disconnect(); + ensureScene(); + } + }, + { rootMargin: "200px" }, + ); + io.observe(section); +} else { + ensureScene(); +} + +// ── The demo itself ─────────────────────────────────────────────────────── + if (!("gpu" in navigator)) { explain( "WebGPU is not available in this browser", @@ -139,6 +190,17 @@ function wire() { } if (!checkCharset()) return; + // The visualization must describe the model that is running, not the + // defaults it was built with. + const scene = await ensureScene(); + scene?.setConfig({ + nLayer: model.n_layer(), + nHead: model.n_head(), + nEmbd: model.n_embd(), + nCtx: model.n_ctx(), + }); + scene?.reset(); + $("demo-output").textContent = ""; $("demo-stop").hidden = false; stop = false; @@ -150,24 +212,33 @@ function wire() { let count = 0; let first = null; - await model.generate( + const onText = (s) => { + // Returning false stops generation after the current token. + if (stop) return false; + if (first === null) first = performance.now(); + count += 1; + $("demo-output").textContent += s; + const dt = (performance.now() - first) / 1000; + if (dt > 0) { + status(`generating — ${(count / dt).toFixed(1)} tok/s`); + } + }; + const args = [ $("demo-prompt").value, n, topk, 0.8, BigInt(Date.now() % 100000), - (s) => { - // Returning false stops generation after the current token. - if (stop) return false; - if (first === null) first = performance.now(); - count += 1; - $("demo-output").textContent += s; - const dt = (performance.now() - first) / 1000; - if (dt > 0) { - status(`generating — ${(count / dt).toFixed(1)} tok/s`); - } - }, - ); + onText, + ]; + + // Text generation never depends on the 3D view: without it the plain + // path runs, and it does no attention readback at all. + await (scene + ? model.generate_with_attention(...args, (layer, nHead, weights) => + scene.pushAttention(layer, nHead, weights), + ) + : model.generate(...args)); const dt = (performance.now() - t0) / 1000; const decode = first === null ? dt : (performance.now() - first) / 1000; diff --git a/docs/src/index.html b/docs/src/index.html index 254364a..7cbf3f2 100644 --- a/docs/src/index.html +++ b/docs/src/index.html @@ -13,6 +13,9 @@ + @@ -33,10 +36,9 @@ WebGPU-native.

- Train a small Shakespeare model from scratch, then load real OpenAI - GPT-2 124M weights with the same code. Every operator, WGSL kernel, - and module exists because GPT-2 needs it. + Forge trains a small Shakespeare model from scratch and runs real + OpenAI GPT-2 124M weights with the same code — in Rust, on WebGPU. + It is built to be read as much as run: every operator, WGSL kernel, + and module exists because GPT-2 needs it, and each one is checked + against a plain-Rust CPU implementation.

Run it in your browser @@ -68,314 +72,242 @@

>

-
-
-
Model on this page
-
43 MB
-
-
-
WGSL kernels
-
23
-
-
-
CPU ↔ WGPU logits
-
8.4e-5
-
-
-
Forge ↔ HF
-
1.75e-4
-
-
+ +

+ verified — CPU and WGPU + logits agree to 8.4e-5, and GPT-2's + output matches HuggingFace + transformers token for token + (tests/gpt2_e2e.rs, on an RTX A5000). +

- +
-

Live demo — Shakespeare, on your GPU

+

Watch it think

- A 10.77M-parameter character-level GPT (6 layers, 6 heads, 384 - dimensions, 65-token vocabulary), trained from scratch by Forge on - Tiny Shakespeare and running here through WebGPU in your browser. - No server does the work — your GPU does. + A 10.77M-parameter character-level GPT — 6 blocks, 6 heads, 384 + dimensions, a 65-character vocabulary — trained from scratch by + Forge on Tiny Shakespeare and running here on your own GPU. The + attention beside it is read back off that GPU as each character is + produced, not an animation of one.

-
- -
- - -
- - -
-
-

Two backends, one model

-

- Model code is backend-agnostic: the same - Gpt2 runs on either device. -

-
-
-

WGPU backend

-

- The production path, via - wgpu - — Vulkan, Metal, D3D12, and browser WebGPU from one set of WGSL - kernels. Reaching an NVIDIA GPU and reaching WebGPU are the same - code path; there is no separate CUDA backend to build. -

-
-
-

CPU backend

-

- A mathematically identical reference implementation, used for - testing and verification. Every kernel is checked against it, - which is what makes the parity numbers below meaningful. -

-
-
- -
-

Why “nanoGPT-class” is not a rebrand

-

- nanoGPT is the GPT-2 architecture — pre-LN blocks, causal - self-attention, a 4× GELU MLP, learned positional embeddings, and - weight tying — and its own headline feature is loading OpenAI - GPT-2 weights. This is one framework, not two modes. -

-

- One honest caveat. - Forge stores Linear weights as - [in, out] (the HF Conv1D - convention), whereas upstream nanoGPT uses - nn.Linear's - [out, in]. Checkpoints are - not - binary-interchangeable with nanoGPT - .pt files: four keys per block would - need transposing, and Forge reads safetensors rather than PyTorch - pickles. + +

+ + The same stack, in text + +
    +
  1. + models/gpt2/ + — config, the 6 transformer blocks above, KV-cache decode, and + generation. Each block is LayerNorm → causal self-attention + (6 heads × 64) → LayerNorm → MLP (384 → 1536 → 384, GELU). +
    2. +
  2. + nn/, autograd/, optim/ + — Linear, LayerNorm, Embedding with a row-chunked token table; + a tape-based reverse mode; AdamW. +
    3. +
  3. + ops.rs + — shape-checked dispatch, one API over both backends. +
    4. +
  4. + backend/ + — WGPU (23 WGSL kernels) and CPU (a plain Rust reference). +
    5. +
+

+ Context 256 characters, 65-token vocabulary, LM head weight-tied + to the token embedding. GPT-2 124M is the same code with + 12 blocks, 12 heads and n_embd 768 — 548 MB of weights, so it + stays a local artifact: + ./scripts/build_site.sh && ./scripts/serve_web.sh.

-
-
-
+ - -
-
-

The stack, twelve blocks deep

-

- GPT-2 124M is twelve identical transformer blocks. Hover one to - highlight it; click to expand its sub-layers with real tensor - shapes. +

+ No server does the work — your GPU does. Forge needs no + SharedArrayBuffer, because + rayon is native-only, so this page is + plain static files.

- -
-
- -

-
-
- - -
-

Architecture

-
- 
-GPT-2 model            models/gpt2/  config, blocks, KV-cache, generation
-  ↓
-nn modules + autograd  nn/  Linear, LayerNorm, Embedding (row-chunked wte)
-                       autograd/  tape, backward ops     optim/  AdamW
-  ↓
-tensor ops             ops.rs  shape-checked dispatch, both backends
-  ↓
-backend abstraction    backend/
-  ↓
-WGPU  ·············  CPU
-23 WGSL kernels      plain Rust reference
-
-

- Per block: LayerNorm → causal self-attention (12 heads of 64 - dimensions) → LayerNorm → MLP (768 → 3072 → 768, GELU). Context - 1024 tokens, vocabulary 50257, LM head weight-tied to the token - embedding. -

-
- -
+ +
-

Verification

-

- Measured on this repository, not estimated. The CPU backend is the - correctness reference; HF transformers - is the external one. -

- -
- - - - - - - - - - - - - - - - - - - - - - - - - - -
- Maximum absolute logit differences -
ComparisonMax differenceSuite
CPU ↔ WGPU logits8.4e-5tests/gpt2_e2e.rs
Forge ↔ HF transformers1.75e-4tests/gpt2_e2e.rs
Every WGSL kernel ↔ CPU≤ 1e-4tests/op_parity.rs
-
+

Why Forge is built this way

+

Three decisions, and what each one buys.

-
+
-

Identical greedy output

+

Rust, not Python

- Greedy continuations match token-for-token across CPU, WGPU, and - HF transformers. + Explicit memory and dispatch, no interpreter in the loop, and + one binary to run. Training and inference are the same code on + the same tensors, so there is no second implementation to keep + honest.

-

Gradients too

+

WebGPU, not CUDA

- Analytic gradients are checked against central differences, and - CPU ↔ WGPU gradient parity is a gate — not just the forward pass. + One set of WGSL kernels reaches Vulkan, Metal, D3D12, and the + browser through + wgpu. + Reaching an NVIDIA GPU and reaching this page are the same code + path; there is no CUDA toolchain to install.

-

41 tests, 9 suites

+

A CPU reference backend

- Kernels, tokenizers, autograd, training ops, optimizer - convergence, KV-cache decode, streaming, and end-to-end GPT-2. + Every kernel has a mathematically identical plain-Rust twin to + check against. That is what makes the parity numbers mean + something — and what makes the thing teachable, since the + reference and the kernel can be read side by side.

+ +

+ “nanoGPT-class” is not a rebrand. + nanoGPT is the GPT-2 architecture — pre-LN blocks, causal + self-attention, a 4× GELU MLP, learned positional embeddings, weight + tying — and loading OpenAI's weights is its own headline feature. + One caveat: Forge stores Linear + weights as [in, out] (the HF Conv1D + convention) rather than nanoGPT's + [out, in], and reads safetensors + rather than PyTorch pickles, so checkpoints are not binary- + interchangeable with nanoGPT .pt + files. +

@@ -384,8 +316,9 @@

41 tests, 9 suites

23 WGSL kernels

- The complete compute surface. Nothing here is generic — each kernel - exists because GPT-2 inference or training needs it. + The complete compute surface, generated from + shaders/ at build time so this list + cannot drift from the code.

@@ -506,105 +439,6 @@

Run this demo locally

- - -
-
-

Roadmap

-

- Stages 1–11 are implemented and gated by tests. Stage 12 is - explicitly not a 1.0 requirement. -

- -
    -
  1. - 1 ✓ - Core types - Shape, DType, Device, Error, Arc'd contiguous Tensor. -
    2. -
  2. - 2 ✓ - CPU reference ops - Plain Rust, unit-tested against hand-computed values. -
    3. -
  3. - 3 ✓ - WebGPU runtime - Adapter, device, queue, shader loader, pipeline cache. -
    4. -
  4. - 4 ✓ - WGSL kernels - Every kernel within 1e-4 of the CPU reference. -
    5. -
  5. - 5 ✓ - Tokenizer & serialization - Byte-level BPE, character-level vocab, safetensors. -
    6. -
  6. - 6 ✓ - GPT-2 inference - Greedy and top-k sampling, weight-tied chunked LM head. -
    7. -
  7. - 7 ✓ - KV-cache decode - Token-identical to full recompute over ≥ 64 tokens. -
    8. -
  8. - 8 ✓ - Autograd - Tape-based reverse mode over the existing ops layer. -
    9. -
  9. - 9 ✓ - Training modules - Deterministic dropout, fused cross-entropy, AdamW. -
    10. -
  10. - 10 ✓ - GPT-2 training - Tiny Shakespeare from scratch, gradient accumulation, - checkpoints. -
    11. -
  11. - 11 ✓ - Browser deployment - wasm32 + WebGPU — the demo at the top of this page. -
    12. -
  12. - 12 - Optimization - Kernel fusion, buffer pooling, matmul tiling. Post-1.0. -
    13. -
-
-
@@ -632,8 +466,8 @@

Roadmap

- - + ', html, re.S) -if imap: - refs |= set(json.loads(imap.group(1))["imports"].values()) +importmap = json.loads(imap.group(1))["imports"] if imap else {} +refs |= set(importmap.values()) -demo = (dist / "demo.js").read_text() if (dist / "demo.js").exists() else "" -refs |= set(re.findall(r'import\("([^"]+)"\)', demo)) # Fetched at runtime, so they must ship even though no tag names them. refs |= {"./model/model.safetensors", "./model/config.json", "./model/vocab.json"} @@ -54,6 +56,50 @@ if not (dist / r.lstrip("./")).exists(): problems.append(f"missing asset (would 404): {r}") +# ── Module graph ────────────────────────────────────────────────────────── +# Static (`from "x"`, `import "x"`, `export … from "x"`) and dynamic +# (`import("x")`) specifiers, in every shipped module including the vendored +# and wasm-bindgen-generated ones. Resolving these is what would have caught +# the missing three.core.min.js at build time instead of in production. +# +# The specifier charset is deliberately narrow: minified three.js contains +# English strings like "…resized from ("+w+")", and a permissive pattern reads +# those as imports. +SPEC = r"""["']([A-Za-z0-9_@~./-]+)["']""" +SPECIFIERS = re.compile( + rf"""\bfrom\s*{SPEC}|\bimport\s*{SPEC}|\bimport\s*\(\s*{SPEC}""" +) + + +def specifiers(text): + """Every module specifier in `text`, from whichever alternative matched.""" + return {next(g for g in m if g) for m in SPECIFIERS.findall(text)} + +for js in sorted(dist.rglob("*.js")): + text = js.read_text(errors="replace") + for spec in sorted(specifiers(text)): + if spec.startswith(("http://", "https://", "//", "data:")): + problems.append(f"{js.relative_to(dist)}: cross-origin import: {spec}") + continue + if spec.startswith("/"): + problems.append( + f"{js.relative_to(dist)}: root-absolute import (404 under /forge/): {spec}" + ) + continue + if spec.startswith("."): + target = (js.parent / spec).resolve() + elif spec in importmap: + # Bare specifier, resolved by the page's importmap relative to the + # document, not to the importing file. + target = (dist / importmap[spec].lstrip("./")).resolve() + else: + problems.append( + f"{js.relative_to(dist)}: bare import with no importmap entry: {spec}" + ) + continue + if not target.exists(): + problems.append(f"{js.relative_to(dist)}: import would 404: {spec}") + if problems: for p in problems: print(f"error: {p}", file=sys.stderr) diff --git a/scripts/ci_local.sh b/scripts/ci_local.sh new file mode 100755 index 0000000..5176b4f --- /dev/null +++ b/scripts/ci_local.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash +# The local CI gate. This script is the single source of truth for what +# "green" means: the hooks in .githooks/ are thin wrappers around it, and +# .github/workflows/ci.yml (manual dispatch only) runs the same stages. +# +# ./scripts/ci_local.sh fast # fmt, clippy, both builds, dep-leak assert +# ./scripts/ci_local.sh full # everything in fast, plus the release tests +# +# `fast` runs on pre-commit (~6s with a warm target/), `full` on pre-push +# (~1m10s). Bypass either with `git commit --no-verify` / `git push +# --no-verify`. +# +# Stages are ordered cheapest-first and the script stops at the first failure, +# so a formatting slip does not cost you a minute of GPU tests. +set -uo pipefail +cd "$(dirname "$0")/.." + +STAGE="${1:-fast}" +case "$STAGE" in + fast|full) ;; + *) echo "usage: $0 [fast|full]" >&2; exit 2 ;; +esac + +if [[ -t 1 ]]; then + BOLD=$'\033[1m'; RED=$'\033[31m'; GREEN=$'\033[32m'; DIM=$'\033[2m'; OFF=$'\033[0m' +else + BOLD=''; RED=''; GREEN=''; DIM=''; OFF='' +fi + +step=0 +started=$SECONDS + +run() { + local name="$1"; shift + step=$((step + 1)) + printf '%s[%d] %s%s\n' "$BOLD" "$step" "$name" "$OFF" + local t0=$SECONDS + if ! "$@"; then + printf '\n%s%s✗ %s failed%s %s(%s)%s\n' "$RED" "$BOLD" "$name" "$OFF" "$DIM" "$*" "$OFF" + printf '%sfix it, or bypass this gate with --no-verify%s\n' "$DIM" "$OFF" + exit 1 + fi + printf ' %sok%s %s(%ss)%s\n' "$GREEN" "$OFF" "$DIM" "$((SECONDS - t0))" "$OFF" +} + +# The TUI deps must never reach the library's dependents or the wasm build; +# they are optional, behind the `tui` feature, for exactly that reason. +assert_no_tui_deps() { + local tree dep n + tree=$(cargo tree -e normal --locked) || return 1 + for dep in ratatui crossterm sysinfo nvml-wrapper memmap2; do + n=$(grep -c "^.*[^a-z-]$dep v" <<<"$tree" || true) + if [[ "$n" -ne 0 ]]; then + echo "$dep leaked into the default dependency tree" >&2 + return 1 + fi + done +} + +printf '%sforge local CI — %s stage%s\n\n' "$BOLD" "$STAGE" "$OFF" + +run "cargo fmt --check" cargo fmt --all --check +run "cargo clippy -D warnings" cargo clippy --all-targets --locked -- -D warnings +run "build wasm32" cargo build --release --locked --target wasm32-unknown-unknown +run "build forge-top (tui)" cargo build --release --locked --features tui --bin forge-top +run "no TUI deps in default tree" assert_no_tui_deps + +if [[ "$STAGE" == full ]]; then + # The weight-dependent suites (gpt2_e2e, kv_cache, the tokenizer's BPE cases) + # self-skip rather than fail when models/gpt2/ is absent, which would make a + # green run quietly weaker than it looks. Say so. + if [[ ! -f models/gpt2/model.safetensors ]]; then + printf '%snote: models/gpt2/ missing — gpt2_e2e and kv_cache will self-skip%s\n' "$DIM" "$OFF" + printf '%s run ./scripts/download_gpt2.sh for full coverage%s\n' "$DIM" "$OFF" + fi + run "cargo test --release" cargo test --release --locked +fi + +printf '\n%s%s✓ all %s checks passed%s %s(%ss)%s\n' \ + "$GREEN" "$BOLD" "$STAGE" "$OFF" "$DIM" "$((SECONDS - started))" "$OFF" diff --git a/scripts/gen_kernels.py b/scripts/gen_kernels.py index 09ca851..2c32774 100755 --- a/scripts/gen_kernels.py +++ b/scripts/gen_kernels.py @@ -57,9 +57,12 @@ def main(target): + "\n".join(cards) + "\n " ) + # The pattern consumes three closing tags — the last card, #kernel-list, + # and the enclosing .wrap — so the replacement must emit exactly three. + # `block` already carries the first two. html, n = re.subn( r'
\s*
\s*', - block + "\n \n ", + block + "\n ", html, count=1, flags=re.S, diff --git a/scripts/install_hooks.sh b/scripts/install_hooks.sh new file mode 100755 index 0000000..87c298b --- /dev/null +++ b/scripts/install_hooks.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# Point git at the repo's committed hooks. Idempotent; safe to re-run. +# +# The hooks live in .githooks/ rather than .git/hooks/ so they are version +# controlled and shared. git never picks that directory up on its own, so +# every clone has to set core.hooksPath once — that is all this script does. +set -euo pipefail +cd "$(dirname "$0")/.." + +git config core.hooksPath .githooks +chmod +x .githooks/* + +echo "core.hooksPath -> .githooks" +echo " pre-commit ./scripts/ci_local.sh fast fmt, clippy, wasm + forge-top builds" +echo " pre-push ./scripts/ci_local.sh full the above + cargo test --release" +echo +echo "bypass either with --no-verify; undo with 'git config --unset core.hooksPath'" diff --git a/src/backend/wgpu/mod.rs b/src/backend/wgpu/mod.rs index 0241ca7..6c2ed9e 100644 --- a/src/backend/wgpu/mod.rs +++ b/src/backend/wgpu/mod.rs @@ -235,6 +235,66 @@ impl WgpuContext { Ok(out) } + /// Read several regions back in one submit and one fence wait. + /// + /// [`WgpuContext::readback_async`] costs a submit and a wait *each*, which + /// dominates when a single logical step wants several small tensors: the + /// attention probe reads `n_layer + 1` per generated token, and one at a + /// time that cost more than the decode itself. Staged into one encoder + /// they cost one round trip regardless of how many there are. + /// + /// Regions are returned in the order given. + pub async fn readback_many_async( + &self, + regions: &[(&wgpu::Buffer, usize, usize)], + ) -> Result>> { + if regions.is_empty() { + return Ok(Vec::new()); + } + let mut encoder = self.device.create_command_encoder(&Default::default()); + let staging: Vec = regions + .iter() + .map(|(buf, offset_bytes, size_bytes)| { + let s = self.device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: *size_bytes as u64, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + encoder.copy_buffer_to_buffer(buf, *offset_bytes as u64, &s, 0, *size_bytes as u64); + s + }) + .collect(); + self.queue.submit([encoder.finish()]); + + // Every map request is issued before anything is awaited, so one poll + // services all of them. + let waits: Vec<_> = staging + .iter() + .map(|s| { + let (tx, rx) = oneshot::channel(); + s.slice(..) + .map_async(wgpu::MapMode::Read, move |r| tx.send(r)); + rx + }) + .collect(); + #[cfg(not(target_arch = "wasm32"))] + self.device + .poll(wgpu::PollType::Wait) + .map_err(|e| ForgeError::Wgpu(format!("poll: {e:?}")))?; + #[cfg(target_arch = "wasm32")] + let _ = self.device.poll(wgpu::PollType::Poll); + + let mut out = Vec::with_capacity(regions.len()); + for (rx, s) in waits.into_iter().zip(&staging) { + rx.await + .map_err(|e| ForgeError::Wgpu(format!("map_async: {e:?}")))?; + out.push(s.slice(..).get_mapped_range().to_vec()); + s.unmap(); + } + Ok(out) + } + /// Dispatch `name` with binding 0 = `params` (uniform, raw words) and /// bindings 1.. = `buffers` (storage). Each buffer entry is /// (buffer, offset_bytes, size_bytes). diff --git a/src/lib.rs b/src/lib.rs index da4f423..84a3a9e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,7 +22,7 @@ pub mod wasm; pub use device::Device; pub use dtype::DType; pub use error::{ForgeError, Result}; -pub use models::gpt2::{Gpt2, Gpt2Config, KvCache, Sampling}; +pub use models::gpt2::{AttnStep, Gpt2, Gpt2Config, KvCache, Sampling}; pub use shape::Shape; pub use tensor::Tensor; pub use tokenizer::{AnyTokenizer, CharTokenizer, Gpt2Tokenizer, Tokenizer}; diff --git a/src/models/gpt2/mod.rs b/src/models/gpt2/mod.rs index b9719ae..e4cdfa0 100644 --- a/src/models/gpt2/mod.rs +++ b/src/models/gpt2/mod.rs @@ -111,6 +111,24 @@ pub enum Sampling { }, } +/// One block's attention probabilities for one decode step, read back from +/// the device. +/// +/// `probs` is the row-major `[n_head, q_len, kv_len]` tensor the model +/// actually attended with — captured after the softmax, not recomputed — so a +/// visualization built from it shows the same arithmetic that produced the +/// text. `ops::softmax` is out-of-place, so capturing it perturbs nothing. +#[derive(Debug, Clone)] +pub struct AttnStep { + pub layer: usize, + pub n_head: usize, + /// Query positions in this step: the whole prompt on prefill, 1 per decode. + pub q_len: usize, + /// Past positions attended to, including the queries themselves. + pub kv_len: usize, + pub probs: Vec, +} + /// Preallocated per-layer K/V tensors (`[n_head, n_ctx, head_dim]`) for /// incremental decode. `len` positions are filled; new tokens append. pub struct KvCache { @@ -251,6 +269,10 @@ impl Gpt2 { /// Attention over new tokens only, appending their K/V to the cache and /// attending to all `len + t` cached positions. + /// + /// `probe`, when present, collects the post-softmax probabilities — an + /// `Arc` handle to a tensor that is computed regardless, so the probing + /// and non-probing paths run identical arithmetic. fn attention_cached( &self, block: &Block, @@ -258,6 +280,7 @@ impl Gpt2 { k_cache: &mut Tensor, v_cache: &mut Tensor, len: usize, + probe: Option<&mut Vec>, ) -> Result { let n_head = self.config.n_head; let hd = self.config.n_embd / n_head; @@ -279,6 +302,9 @@ impl Gpt2 { }, )?; // [h, t, kv_len] let att = ops::softmax(&att, true, len)?; // off = kv_len - t + if let Some(probe) = probe { + probe.push(att.clone()); + } let y = ops::matmul( &att, v_cache, @@ -320,8 +346,14 @@ impl Gpt2 { } /// Hidden states for `ids` (new tokens) continuing from `cache`; - /// appends their K/V and advances `cache.len`. - fn hidden_cached(&self, ids: &[u32], cache: &mut KvCache) -> Result { + /// appends their K/V and advances `cache.len`. `probe` is threaded through + /// to [`Gpt2::attention_cached`] and collects one tensor per block. + fn hidden_cached( + &self, + ids: &[u32], + cache: &mut KvCache, + mut probe: Option<&mut Vec>, + ) -> Result { if ids.is_empty() { return Err(ForgeError::Shape("empty token sequence".into())); } @@ -342,7 +374,14 @@ impl Gpt2 { .iter() .zip(cache.k.iter_mut().zip(cache.v.iter_mut())) { - let attn_out = self.attention_cached(block, &block.ln_1.forward(&x)?, kc, vc, pos)?; + let attn_out = self.attention_cached( + block, + &block.ln_1.forward(&x)?, + kc, + vc, + pos, + probe.as_deref_mut(), + )?; x = ops::add(&x, &attn_out)?; let mlp_in = block.ln_2.forward(&x)?; let mlp_out = block @@ -386,7 +425,7 @@ impl Gpt2 { /// Last-position logits for `ids` continuing from `cache` (incremental /// decode: pass the full prompt once, then one token at a time). pub fn logits_step(&self, ids: &[u32], cache: &mut KvCache) -> Result> { - let h = self.hidden_cached(ids, cache)?; + let h = self.hidden_cached(ids, cache, None)?; let last = h.narrow_rows(ids.len() - 1, 1)?; ops::matmul_chunked_transb(&last, &self.emb.wte_chunks, 1.0)?.to_vec_f32() } @@ -394,13 +433,58 @@ impl Gpt2 { /// Async form of [`Gpt2::logits_step`] — identical math; the readback is /// awaited so it works on wasm32 (roadmap v4, pitfall 14). pub async fn logits_step_async(&self, ids: &[u32], cache: &mut KvCache) -> Result> { - let h = self.hidden_cached(ids, cache)?; + let h = self.hidden_cached(ids, cache, None)?; let last = h.narrow_rows(ids.len() - 1, 1)?; ops::matmul_chunked_transb(&last, &self.emb.wte_chunks, 1.0)? .to_vec_f32_async() .await } + /// [`Gpt2::logits_step_async`] plus every block's attention probabilities + /// for this step, in layer order. + /// + /// The logits are identical to `logits_step_async` — the probe only reads + /// tensors the step computes anyway. The whole step is enqueued first, and + /// the logits and every block's attention come back in a *single* batched + /// readback: one at a time they cost ~2.7x the decode itself, since the + /// price is the round trip rather than the ~36 KB. + pub async fn logits_step_attn_async( + &self, + ids: &[u32], + cache: &mut KvCache, + ) -> Result<(Vec, Vec)> { + let mut probe = Vec::with_capacity(self.config.n_layer + 1); + let h = self.hidden_cached(ids, cache, Some(&mut probe))?; + let last = h.narrow_rows(ids.len() - 1, 1)?; + // Logits last, so `probe` stays in layer order and the shapes below + // line up with the tensors that produced them. + let shapes: Vec> = probe.iter().map(|t| t.shape().dims().to_vec()).collect(); + probe.push(ops::matmul_chunked_transb( + &last, + &self.emb.wte_chunks, + 1.0, + )?); + + let mut read = Tensor::to_vec_f32_batch(&probe).await?; + let logits = read.pop().expect("logits were pushed last"); + let mut steps = Vec::with_capacity(read.len()); + for (layer, (probs, dims)) in read.into_iter().zip(shapes).enumerate() { + let [n_head, q_len, kv_len] = dims[..] else { + return Err(ForgeError::Shape(format!( + "attention probe expected rank 3, got {dims:?}" + ))); + }; + steps.push(AttnStep { + layer, + n_head, + q_len, + kv_len, + probs, + }); + } + Ok((logits, steps)) + } + // ---- training (roadmap v4, Stages 8-10) ---- /// Random initialization for training from scratch: N(0, 0.02) weights, @@ -844,12 +928,42 @@ impl Gpt2 { /// [`ControlFlow::Break`] stops after the current token, which is how a /// page offers a working "stop" button. pub async fn generate_async_ctl( + &self, + tokenizer: &impl Tokenizer, + prompt: &str, + max_new_tokens: usize, + sampling: Sampling, + on_text: impl FnMut(&str) -> ControlFlow<()>, + ) -> Result { + // `None` turns off the probe entirely: no capture, no readback, and + // the same `logits_step_async` this method has always called. + self.generate_async_probe( + tokenizer, + prompt, + max_new_tokens, + sampling, + on_text, + None::, + ) + .await + } + + /// [`Gpt2::generate_async_ctl`] with an optional attention probe: + /// `on_attn`, when present, fires once per decode step with every block's + /// attention probabilities — including the prompt prefill, whose `q_len` + /// is the prompt length rather than 1. + /// + /// Opt-in because it costs one readback per block per token. Passing + /// `None` is exactly the non-probing path; the generated text is identical + /// either way, since the probe reads tensors the step already computed. + pub async fn generate_async_probe( &self, tokenizer: &impl Tokenizer, prompt: &str, max_new_tokens: usize, sampling: Sampling, mut on_text: impl FnMut(&str) -> ControlFlow<()>, + mut on_attn: Option, ) -> Result { // The streaming helper takes a plain sink, so the break request is // captured here and checked once the delta has been forwarded. A Cell @@ -869,7 +983,7 @@ impl Gpt2 { Sampling::Greedy => None, }; let mut cache = self.new_cache()?; - let mut logits = self.logits_step_async(&ids, &mut cache).await?; // prompt prefill + let mut logits = self.step_async(&ids, &mut cache, &mut on_attn).await?; // prompt prefill // Stream over the raw byte-level decode (append-only per token), // emitting only its valid-UTF-8 prefix: a multi-byte character split // across BPE tokens is held back until its trailing bytes arrive. @@ -889,10 +1003,27 @@ impl Gpt2 { if ids.len() >= self.config.n_ctx { break; } - logits = self.logits_step_async(&[next], &mut cache).await?; // single-token decode + logits = self.step_async(&[next], &mut cache, &mut on_attn).await?; // single-token decode } Ok(tokenizer.decode(&ids)) } + + /// One decode step, forwarding attention to `on_attn` when probing. + async fn step_async( + &self, + ids: &[u32], + cache: &mut KvCache, + on_attn: &mut Option, + ) -> Result> { + match on_attn { + Some(f) => { + let (logits, steps) = self.logits_step_attn_async(ids, cache).await?; + f(&steps); + Ok(logits) + } + None => self.logits_step_async(ids, cache).await, + } + } } /// Send `bytes[sent..]` to `on_text` up to the longest decodable prefix; diff --git a/src/tensor.rs b/src/tensor.rs index 6d3d966..0bd6487 100644 --- a/src/tensor.rs +++ b/src/tensor.rs @@ -175,6 +175,51 @@ impl Tensor { } } + /// Read several f32 tensors back in one GPU round trip, in the order + /// given. + /// + /// [`Tensor::to_vec_f32_async`] is a submit and a fence wait per call, so + /// a step that wants several small tensors — the attention probe reads + /// `n_layer + 1` per generated token — pays for the round trips, not the + /// bytes. Every tensor must be on the same device. + pub async fn to_vec_f32_batch(tensors: &[Tensor]) -> Result>> { + let mixed = || ForgeError::Shape("to_vec_f32_batch needs one device".into()); + let mut ctx: Option<&Arc> = None; + for t in tensors { + if t.dtype != DType::F32 { + return Err(ForgeError::Shape("to_vec_f32 on non-f32 tensor".into())); + } + if let Storage::Wgpu(s) = &t.storage { + match ctx { + None => ctx = Some(&s.ctx), + Some(c) if Arc::ptr_eq(c, &s.ctx) => {} + Some(_) => return Err(mixed()), + } + } + } + let Some(ctx) = ctx else { + // All host-side: nothing to stage, and no round trip to save. + return tensors.iter().map(Tensor::to_vec_f32).collect(); + }; + let mut regions = Vec::with_capacity(tensors.len()); + for t in tensors { + match &t.storage { + Storage::Wgpu(s) => { + regions.push((s.buf.as_ref(), s.offset * 4, t.shape.numel() * 4)) + } + // A host tensor has nothing to stage, so a mixed batch would + // misalign the results with their inputs. + Storage::Cpu(_) => return Err(mixed()), + } + } + Ok(ctx + .readback_many_async(®ions) + .await? + .iter() + .map(|b| bytemuck::pod_collect_to_vec(b)) + .collect()) + } + pub async fn to_vec_u32_async(&self) -> Result> { if self.dtype != DType::U32 { return Err(ForgeError::Shape("to_vec_u32 on non-u32 tensor".into())); diff --git a/src/wasm.rs b/src/wasm.rs index cc3931a..58e5a27 100644 --- a/src/wasm.rs +++ b/src/wasm.rs @@ -5,7 +5,7 @@ use wasm_bindgen::prelude::*; use crate::Device; -use crate::models::gpt2::{Gpt2, Gpt2Config, Sampling}; +use crate::models::gpt2::{AttnStep, Gpt2, Gpt2Config, Sampling}; use crate::tokenizer::{AnyTokenizer, CharTokenizer, Gpt2Tokenizer, Tokenizer as _}; #[wasm_bindgen(start)] @@ -80,6 +80,24 @@ impl WasmGpt2 { self.tokenizer.vocab_size() } + // The page's architecture view is sized from the model it is actually + // running, not from GPT-2 124M constants baked into JavaScript. + pub fn n_layer(&self) -> usize { + self.model.config.n_layer + } + + pub fn n_head(&self) -> usize { + self.model.config.n_head + } + + pub fn n_embd(&self) -> usize { + self.model.config.n_embd + } + + pub fn n_ctx(&self) -> usize { + self.model.config.n_ctx + } + /// Characters of `prompt` this model's vocabulary cannot represent, /// deduplicated. Empty when the prompt is fine. The char model knows only /// 65 characters, so the page checks before generating rather than @@ -132,6 +150,73 @@ impl WasmGpt2 { .map_err(js_err) } + /// [`WasmGpt2::generate`] plus a live attention feed: after every token, + /// `on_attn(layer, n_head, weights)` fires once per block with that + /// block's attention probabilities as a `Float32Array` of + /// `n_head * kv_len` — head-major, so `kv_len` is `weights.length / + /// n_head`, and `weights[h * kv_len + p]` is how much head `h` weighted + /// position `p` when producing this token. + /// + /// Only the newest query row is sent: on the prompt prefill the model + /// attends with every prompt position at once, but the row that produced + /// the next token is the last one. + /// + /// Opt-in — `generate` does no attention readback at all. + pub async fn generate_with_attention( + &self, + prompt: &str, + max_new_tokens: usize, + top_k: usize, + temperature: f32, + seed: u64, + on_text: &js_sys::Function, + on_attn: &js_sys::Function, + ) -> Result { + let sampling = if top_k == 0 { + Sampling::Greedy + } else { + Sampling::TopK { + k: top_k, + temperature, + seed, + } + }; + let this = JsValue::NULL; + self.model + .generate_async_probe( + &self.tokenizer, + prompt, + max_new_tokens, + sampling, + |s| match on_text.call1(&this, &JsValue::from_str(s)) { + Ok(v) if v.is_falsy() && !v.is_undefined() && !v.is_null() => { + std::ops::ControlFlow::Break(()) + } + _ => std::ops::ControlFlow::Continue(()), + }, + Some(|steps: &[AttnStep]| { + for s in steps { + let last = (s.q_len - 1) * s.kv_len; + let mut row = Vec::with_capacity(s.n_head * s.kv_len); + for head in 0..s.n_head { + let base = head * s.q_len * s.kv_len + last; + row.extend_from_slice(&s.probs[base..base + s.kv_len]); + } + // A failing visualization callback must not abort + // generation; the text is the point. + let _ = on_attn.call3( + &this, + &JsValue::from_f64(s.layer as f64), + &JsValue::from_f64(s.n_head as f64), + &js_sys::Float32Array::from(&row[..]), + ); + } + }), + ) + .await + .map_err(js_err) + } + /// Greedy continuation as raw token ids — used by the Stage 11 gate to /// compare browser output against native WGPU token-for-token. pub async fn greedy_ids( diff --git a/tests/char_model.rs b/tests/char_model.rs index cd5b2f9..69da13f 100644 --- a/tests/char_model.rs +++ b/tests/char_model.rs @@ -5,7 +5,9 @@ //! run on a fresh clone. They still self-skip, because the artifact is only //! present once someone has trained and shipped it. -use forge::{AnyTokenizer, Device, Gpt2, Gpt2Config, Sampling, Tokenizer as _}; +use std::ops::ControlFlow; + +use forge::{AnyTokenizer, AttnStep, Device, Gpt2, Gpt2Config, Sampling, Tokenizer as _}; const DIR: &str = "assets/shakespeare_char"; const PROMPT: &str = "ROMEO:"; @@ -83,3 +85,98 @@ fn generation_stays_inside_the_vocabulary() { "output is not mostly letters: {text:?}" ); } + +// ── Attention probe ─────────────────────────────────────────────────────── +// The website renders these numbers as a live 3D view while it generates, so +// what matters is that they are the model's own softmax output and that +// capturing them changes nothing. + +fn model_on(config: Gpt2Config, device: &Device) -> Gpt2 { + Gpt2::from_safetensors(format!("{DIR}/model.safetensors"), config, device).unwrap() +} + +#[test] +fn attention_probe_captures_the_real_softmax() { + let Some((config, tok)) = assets() else { + return; + }; + // Both backends: WGPU reads every block back in one batched round trip, + // which is a different path from the CPU backend's plain clone. + for device in [Device::Cpu, Device::wgpu().unwrap()] { + let model = model_on(config.clone(), &device); + let ids = tok.encode(PROMPT).unwrap(); + let mut cache = model.new_cache().unwrap(); + let (_, steps) = + pollster::block_on(model.logits_step_attn_async(&ids, &mut cache)).unwrap(); + + assert_eq!(steps.len(), config.n_layer, "one capture per block"); + for (layer, s) in steps.iter().enumerate() { + assert_eq!(s.layer, layer); + assert_eq!(s.n_head, config.n_head); + // Prefill attends with every prompt position at once. + assert_eq!((s.q_len, s.kv_len), (ids.len(), ids.len())); + assert_eq!(s.probs.len(), s.n_head * s.q_len * s.kv_len); + + for head in 0..s.n_head { + for q in 0..s.q_len { + let row = &s.probs[(head * s.q_len + q) * s.kv_len..][..s.kv_len]; + let sum: f32 = row.iter().sum(); + assert!( + (sum - 1.0).abs() < 1e-5, + "{}: layer {layer} head {head} row {q} sums to {sum}, not 1", + device.describe() + ); + // Causal mask: a query never attends to a later position. + for (k, &w) in row.iter().enumerate().skip(q + 1) { + assert_eq!( + w, + 0.0, + "{}: layer {layer} head {head} row {q} sees future {k}", + device.describe() + ); + } + } + } + } + + // One decode step: a single query against one more cached position. + let (_, steps) = + pollster::block_on(model.logits_step_attn_async(&[ids[0]], &mut cache)).unwrap(); + for s in &steps { + assert_eq!((s.q_len, s.kv_len), (1, ids.len() + 1)); + } + } +} + +#[test] +fn attention_probe_does_not_change_the_output() { + let Some((config, tok)) = assets() else { + return; + }; + let model = model_on(config.clone(), &Device::Cpu); + const N: usize = 24; + + let plain = pollster::block_on(model.generate_async(&tok, PROMPT, N, Sampling::Greedy, |_| {})) + .unwrap(); + + let mut seen = 0usize; + let probed = pollster::block_on(model.generate_async_probe( + &tok, + PROMPT, + N, + Sampling::Greedy, + |_| ControlFlow::Continue(()), + Some(|steps: &[AttnStep]| { + assert_eq!(steps.len(), config.n_layer); + seen += 1; + }), + )) + .unwrap(); + + assert_eq!( + plain, probed, + "the probe perturbed the computation it is meant to observe" + ); + // Prefill plus one capture per generated token. + assert_eq!(seen, N + 1, "the probe skipped or duplicated a decode step"); +}