Fix the website issues

.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.

.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

.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 "<local ref> <local sha> <remote ref> <remote sha>" 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

.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

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.

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"

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;