Publish cross-API Acrobot artifact: implementations + equivalence gate + Colab runner (#1)

* Publish the cross-API Acrobot artifact: implementations, equivalence gate, Colab runner

The paper's headline cross-API numbers (720x CUDA / 172x JAX / 27x Triton
vs per-step PyTorch on T4) had no published code. This adds it:

- benchmarks/colab/acrobot_impls.py — one canonical dynamics definition
  (Gym "book" Acrobot, explicit Euler dt=0.05, matching the WebGPU shader
  lineage) with five implementations: NumPy reference, per-step PyTorch
  baseline, JAX lax.scan, hand-fused CUDA (CuPy RawKernel), Triton
- tests/acrobot_equivalence_test.py + CI job — NumPy/PyTorch/JAX must
  produce IDENTICAL fitness for shared genomes in float64 (verified: 1.0000
  match) before any timing means anything
- benchmarks/colab/acrobot_cross_api.ipynb — Colab notebook (GPU runtime)
  that clones the repo, verifies GPU kernels against the same reference
  (>=95% match required; chaotic-system caveat documented), then times all
  implementations and saves a results JSON
- Label correction: the integrator is explicit Euler, not RK4 as earlier
  table captions claimed; artifact note added to paper.tex including the
  framework-versions-are-moving-targets caveat (cross-ref to the companion
  transformer paper's v2 erratum)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* Add T4 re-measurement (Colab, 2026-07-28) with verified-equivalent kernels

Ran the notebook on a Colab Tesla T4 (python 3.12.13, torch 2.11.0+cu128,
jax 0.7.2, cupy 14.0.1). Equivalence vs the NumPy reference: torch 99.66%,
fused CUDA 99.63%, jax 99.63% (chaotic f32 flips; gate >=95%).

Timings (N=10 medians, POP=4096, 500 steps): per-step PyTorch 554.7ms
(1.8 gen/s), fused CUDA 3.15ms (317.3 gen/s, 176x), JAX lax.scan 16.5ms
(60.5 gen/s, 33.5x). v1's 720x/172x reflected a ~3x-slower torch —
qualitative result unchanged, ratios compressed. Also fixes the Triton
kernel for newer Triton (tl.math.tanh removed; manual tanh).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: abgnydn

.github/workflows/ci.yml

@@ -28,3 +28,20 @@ jobs:
       # copy-paste errors and stale numbers without needing GPU access.
       - name: Verify paper arithmetic
         run: node tests/paper_arithmetic.test.js
+
+  acrobot-equivalence:
+    name: acrobot cross-API equivalence (CPU, f64)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+      - run: pip install numpy jax torch --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
+      # NumPy, PyTorch, and JAX must produce IDENTICAL fitness for shared
+      # genomes — the property that makes the cross-API speedup table
+      # meaningful. GPU kernels (CUDA/Triton) are checked in the Colab
+      # notebook against the same reference.
+      - name: Acrobot implementation equivalence
+        run: JAX_ENABLE_X64=1 python tests/acrobot_equivalence_test.py

CHANGELOG.md

@@ -5,6 +5,26 @@ format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and the project follows [Semantic Versioning](https://semver.org/) starting
 from `0.1.0`.
 
+## [0.2.0] — 2026-07-28
+
+### Added
+
+- **Cross-API Acrobot artifact published** — the code behind the headline
+  T4 numbers now lives in `benchmarks/colab/acrobot_impls.py` (NumPy
+  reference, per-step PyTorch baseline, JAX `lax.scan`, hand-fused CUDA via
+  CuPy, Triton) with a CI-gated float64 equivalence test and a Colab
+  notebook runner.
+- **Re-measured on Colab T4 (2026-07-28)** with verified-equivalent
+  kernels: fused CUDA 176x, JAX 33.5x over per-step PyTorch (torch
+  2.11.0+cu128). v1's 720x/172x reflected a ~3x-slower torch of spring
+  2026 — ratios compress as frameworks improve; qualitative result
+  unchanged. Data in `benchmarks/results/2026-07-28_t4_colab/`.
+
+### Fixed
+
+- Acrobot tables were labeled "RK4"; the integrator is explicit Euler at
+  dt=0.05 (corrected in README, PAPER.md, paper.tex + artifact note).
+
 ## [0.1.0] — 2026-05-04
 
 First public release of the kernel-fusion benchmark suite + paper companion.

PAPER.md

@@ -117,7 +117,7 @@ WebGPU executes the **entire 1,500-step simulation as a single compute shader di
 
 ### 4.3 Sequential Workload: Acrobot-v1
 
-**Table 3: Acrobot-v1 (POP=4,096, 163 params, 500 steps, RK4)**
+**Table 3: Acrobot-v1 (POP=4,096, 163 params, 500 steps, explicit Euler — earlier versions mislabeled this RK4)**
 
 *Same hardware: Tesla T4 (CUDA)*
 

README.md

@@ -4,12 +4,13 @@
 [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](./LICENSE)
 [![Live](https://img.shields.io/badge/live-kernelfusion.dev-6ea8ff)](https://kernelfusion.dev)
 [![Paper arithmetic](https://img.shields.io/badge/paper%20arithmetic-54%20%E2%9C%93-82c98b)](./tests/paper_arithmetic.test.js)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/abgnydn/webgpu-kernel-fusion/blob/main/benchmarks/colab/acrobot_cross_api.ipynb)
 
 Fusing sequential fitness evaluations into single GPU compute shader dispatches eliminates per-step kernel launch overhead. We prove this across **4 GPU APIs on 2 hardware platforms** — the fusion advantage is **GPU-API-agnostic**.
 
 ## Key Results
 
-**Same hardware: Tesla T4 (Acrobot-v1, 500 steps, RK4)**
+**Same hardware: Tesla T4 (Acrobot-v1, 500 steps, explicit Euler dt=0.05 — earlier versions mislabeled this RK4)**
 
 | System | gen/s | vs PyTorch |
 |---|---|---|
@@ -52,6 +53,31 @@ The fusion advantage scales with dispatch overhead fraction. On compute-bound N-
 - **Paper (LaTeX):** [paper.tex](paper.tex)
 - **DOI:** [10.5281/zenodo.19342888](https://doi.org/10.5281/zenodo.19342888)
 
+### Re-measured (Tesla T4, Colab, 2026-07-28 — corrected artifact, verified-equivalent kernels)
+
+| System | gen/s | vs PyTorch |
+|---|---|---|
+| PyTorch CUDA per-step (torch 2.11.0) | 1.8 | 1x |
+| JAX lax.scan+vmap (jax 0.7.2) | 60.5 | **33.5x** |
+| **Hand-fused CUDA kernel (cupy 14.0.1)** | **317.3** | **176x** |
+
+All kernels cross-checked for output equivalence before timing (≥ 99.6% match,
+f32 chaotic-flip residue documented). PyTorch's per-step dispatch got ~3x
+faster between the v1 measurement and torch 2.11, compressing the ratios —
+the qualitative result is unchanged. Raw data:
+[`benchmarks/results/2026-07-28_t4_colab/`](benchmarks/results/2026-07-28_t4_colab/).
+
+## Reproduce the Cross-API Table (Colab)
+
+The Acrobot cross-API implementations (per-step PyTorch, hand-fused CUDA,
+JAX `lax.scan`, Triton) live in [`benchmarks/colab/acrobot_impls.py`](benchmarks/colab/acrobot_impls.py)
+with one canonical dynamics definition. Open
+[`benchmarks/colab/acrobot_cross_api.ipynb`](benchmarks/colab/acrobot_cross_api.ipynb)
+in Colab on a GPU runtime — it **verifies all implementations produce identical
+fitness on shared genomes before any timing** (CI additionally gates NumPy /
+PyTorch / JAX equivalence in float64 on every push). Speedups are tied to the
+framework versions the notebook prints; treat baselines as moving targets.
+
 ## Reproduce Every Result
 
 ### Prerequisites

benchmarks/colab/acrobot_cross_api.ipynb

@@ -0,0 +1,178 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Acrobot Cross-API Fusion Benchmark (Colab, NVIDIA GPU)\n",
+    "\n",
+    "Reproduces the cross-API table of *Single-Kernel Fusion for Sequential Fitness\n",
+    "Evaluation via WebGPU Compute Shaders* (doi:10.5281/zenodo.19331833): per-step\n",
+    "PyTorch dispatch vs hand-fused CUDA vs JAX `lax.scan` vs Triton, all evaluating\n",
+    "the **same function** (verified below before any timing).\n",
+    "\n",
+    "**Runtime \u2192 Change runtime type \u2192 GPU** (T4 matches the paper).\n",
+    "\n",
+    "All implementations live in the repo (`benchmarks/colab/acrobot_impls.py`) and are\n",
+    "CI-tested for CPU equivalence; this notebook adds the GPU implementations and timings.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!nvidia-smi -L\n",
+    "!git clone -q https://github.com/abgnydn/webgpu-kernel-fusion.git repo 2>/dev/null || (cd repo && git pull -q)\n",
+    "import sys; sys.path.insert(0, 'repo/benchmarks/colab')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install -q cupy-cuda12x 'jax[cuda12]'\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np, time, json, torch\n",
+    "import acrobot_impls as A\n",
+    "print(f'POP={A.POP} GENOME={A.GENOME_SIZE} STEPS={A.STEPS} (explicit Euler dt={A.DT})')\n",
+    "genomes = A.make_genomes()  # shared across every implementation\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1 \u2014 Equivalence before timing\n",
+    "\n",
+    "Speedups over kernels that compute different functions are meaningless, so every\n",
+    "implementation is checked against the NumPy reference on the shared genomes first.\n",
+    "The double pendulum is chaotic and the argmax policy discontinuous, so f32\n",
+    "op-order differences can flip a small fraction of episodes \u2014 f64 CPU equivalence\n",
+    "is exact (CI-gated); for f32 GPU kernels we report the match fraction and require\n",
+    "\u2265 95%.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ref32 = A.acrobot_numpy(genomes)\n",
+    "\n",
+    "checks = {}\n",
+    "checks['torch-cuda f32'] = A.acrobot_torch(genomes, device='cuda')\n",
+    "checks['fused CUDA'] = A.acrobot_cuda(genomes)\n",
+    "import jax.numpy as jnp\n",
+    "jax_eval = A.make_acrobot_jax()\n",
+    "checks['jax f32'] = np.asarray(jax_eval(jnp.asarray(genomes)))\n",
+    "try:\n",
+    "    triton_eval = A.make_acrobot_triton()\n",
+    "    checks['triton'] = triton_eval(genomes)\n",
+    "except Exception as e:\n",
+    "    triton_eval = None\n",
+    "    print(f'Triton unavailable/failed to compile: {e}')\n",
+    "\n",
+    "for name, r in checks.items():\n",
+    "    frac = A.match_fraction(ref32, r, atol=1e-4)\n",
+    "    status = 'OK' if frac >= 0.95 else 'FAIL'\n",
+    "    print(f'{name:16} match={frac:.4f}  {status}')\n",
+    "    assert frac >= 0.95, f'{name} diverged \u2014 investigate before timing'\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2 \u2014 Timings (N=10, medians, gen/s = full-population evaluations per second)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "N = 10\n",
+    "def bench(fn, warmup=3):\n",
+    "    for _ in range(warmup): fn()\n",
+    "    ts = []\n",
+    "    for _ in range(N):\n",
+    "        t0 = time.perf_counter(); fn(); ts.append(time.perf_counter() - t0)\n",
+    "    ts.sort()\n",
+    "    med = ts[len(ts)//2]\n",
+    "    return med\n",
+    "\n",
+    "g_cuda = torch.as_tensor(genomes, device='cuda')\n",
+    "results = {}\n",
+    "results['pytorch per-step (baseline)'] = bench(lambda: A.acrobot_torch(genomes, device='cuda'))\n",
+    "results['fused CUDA (cupy)'] = bench(lambda: A.acrobot_cuda(genomes))\n",
+    "gj = jnp.asarray(genomes)\n",
+    "results['jax lax.scan (jit)'] = bench(lambda: np.asarray(jax_eval(gj)))\n",
+    "if triton_eval: results['triton fused'] = bench(lambda: triton_eval(genomes))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "base = results['pytorch per-step (baseline)']\n",
+    "print(f\"{'implementation':30} {'median':>10} {'gen/s':>8} {'speedup':>9}\")\n",
+    "for name, med in results.items():\n",
+    "    print(f'{name:30} {med*1000:8.1f}ms {1/med:8.1f} {base/med:8.1f}x')\n",
+    "\n",
+    "out = {'gpu': torch.cuda.get_device_name(0), 'n': N,\n",
+    "       'pop': A.POP, 'steps': A.STEPS,\n",
+    "       'results_ms': {k: round(v*1000, 2) for k, v in results.items()},\n",
+    "       'speedups': {k: round(base/v, 1) for k, v in results.items()}}\n",
+    "print(json.dumps(out, indent=2))\n",
+    "with open('acrobot_cross_api_results.json', 'w') as f: json.dump(out, f, indent=2)\n",
+    "print('\\nSaved acrobot_cross_api_results.json \u2014 commit it under benchmarks/results/.')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Notes\n",
+    "\n",
+    "- The baseline is *per-step framework dispatch* \u2014 the pattern naive GPU RL/EC\n",
+    "  evaluation uses. It is a fair representation of that pattern, not of the best\n",
+    "  possible unfused pipeline (e.g. CUDA graphs would batch launches).\n",
+    "- Speedups are tied to the framework versions printed by the install cell \u2014\n",
+    "  they are moving targets; report versions with any numbers you publish.\n",
+    "- Dynamics: Gym's 'book' Acrobot equations, explicit Euler at dt=0.05 across\n",
+    "  ALL implementations (earlier artifacts mislabeled this as RK4).\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}