Metadata-Version: 2.4
Name: miniworld-engine
Version: 2.3.1
Summary: MiniWorld GPU kernel development repo — Triton / CUDA kernels + benches for AF3-style ops (triangle multiplication, transition, attention, ...).
License-Expression: MIT AND Apache-2.0 AND BSD-3-Clause
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: licenses/THIRD_PARTY_NOTICES.md
License-File: licenses/Anthropic-Apache-2.0.txt
License-File: licenses/Anthropic-NOTICE
License-File: licenses/CUTLASS-LICENSE.txt
License-File: licenses/FlashAttention-LICENSE.txt
Requires-Dist: torch==2.13.*
Requires-Dist: triton>=3.3
Requires-Dist: einops
Requires-Dist: jaxtyping
Requires-Dist: numpy
Requires-Dist: scipy
Requires-Dist: pydantic
Provides-Extra: baselines
Requires-Dist: cuequivariance-torch<0.13,>=0.12; extra == "baselines"
Provides-Extra: bench
Requires-Dist: hydra-core; extra == "bench"
Requires-Dist: lightning; extra == "bench"
Requires-Dist: omegaconf; extra == "bench"
Requires-Dist: pydantic; extra == "bench"
Requires-Dist: matplotlib; extra == "bench"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: ruff==0.16.9; extra == "dev"
Requires-Dist: ty==0.0.84; extra == "dev"
Provides-Extra: flash
Requires-Dist: flash-attn-4; extra == "flash"
Provides-Extra: flash2
Requires-Dist: flash-attn<3,>=2.8; extra == "flash2"
Dynamic: license-file

# miniworld-engine

GPU kernel library for MiniWorld / AF3-style ops. Each op is cut out of the model and
optimized in isolation: **hand-written CUDA where it exists, a Triton fallback on every other GPU**
(H100 and B200 have none, on purpose: they fall back to PyTorch), and a PyTorch reference that defines what "correct" means.

> **Where this fits.** miniworld-engine is the bottom layer of a three-layer stack: it owns
> the fused kernels and building-block ops; **team-gm** composes them into representative AF3
> blocks; terminal product repos assemble those into full models. Boundary rules live in
> team-gm's `docs/ARCHITECTURE.md`.

**Version 2.3.1** — CUDA + Triton only (CuTe/quack removed), torch 2.13 (CUDA 12.9 or 13.0), prebuilt H100 kernels,
cuequivariance 0.12. [Changelog](docs/CHANGELOG.md) ·
[attribution](licenses/THIRD_PARTY_NOTICES.md).

| I want to… | go to |
|---|---|
| run on a GPU, or see what is finished there for which shapes | [docs/gpus/](docs/gpus/README.md) |
| find a dated measurement or audit (never edited after the fact) | [docs/records/](docs/records/) |
| read everything else | [docs/README.md](docs/README.md) |
| work on this repo as an agent | [.claude/CLAUDE.md](.claude/CLAUDE.md) |

## Quickstart

Four steps, and the first three need no GPU. Every command in this section is executed by
`tests/cpu/layout/test_quickstart_runs.py`, so it cannot drift from what the code does.

**1. Install.** The library is a normal wheel; `[bench]` and `[baselines]` are extras you do not need
to read a number.

```bash
# cpu
python -m pip install -e .
```

**2. Check what you have.** Prints the version and the config set every triton kernel will search
if you do nothing else.

```bash
# cpu
python -c "import miniworld_engine as m; print(m.__version__)"
python -c "from miniworld_engine.autotune import configs; print(configs.default_config_dir())"
```

**3. See what the library declares.** One row per kernel: which backend, which arch it requires,
which tolerance it is held to.

```bash
# cpu
python -c "
import csv, collections, miniworld_engine.kernels as k, pathlib
reg = pathlib.Path(k.__file__).parent / 'registry' / 'registry.csv'
rows = list(csv.DictReader(reg.open()))
print(len(rows), 'kernels;', dict(collections.Counter(r['backend'] for r in rows)))"
```

**4. Run them, on a GPU.** This launches every kernel that has a driver and compares each against
its torch reference. It is the fastest way to find out whether this library works on your card,
and it writes `autotune/manifests/<your card>.csv` recording what it found.

```bash
# gpu
python -m miniworld_engine.autotune.run_all
```

The summary reports `declared`, `driven`, `ok`, `failed`, and `skipped` counts for the
current registry and device. `skipped` distinguishes unsupported architecture/dtype cases;
missing drivers and execution failures are reported separately. See
[docs/gpus/supported.md](docs/gpus/supported.md) for what has actually been run, and
[docs/gpus/troubleshooting.md](docs/gpus/troubleshooting.md) when a step does not do this.

**Then:** using the kernels means `from miniworld_engine import ops` — eight whole-op entry points
that take the same arguments as their torch equivalents. Getting them *fast* on your card means
building a tuned cache with `miniworld-engine build all`. Its work depends on the GPU,
current source and existing valid measurements; use `dev coverage` to inspect a specific GPU cache.
## Prebuilt CUDA kernels

The hand-written CUDA kernels compile on first use (tens of minutes for a new environment) unless the installed engine
is a binary wheel: release wheels (`miniworld_engine-<version>-cp3XX-cp3XX-manylinux_2_28_x86_64.whl`) carry the
compiled kernels for H100 (sm_90a) in `miniworld_engine/_prebuilt/` for both CUDA builds of torch 2.13 (the default
PyPI one, CUDA 13.0, and the cu129 one) and pin `torch==2.13.*`, so the first call loads binaries and never runs nvcc:

```bash
pip install miniworld-engine     # with the torch it pins; an sdist / git install keeps the JIT path
# torch's CUDA 12.9 build instead (NVIDIA driver 575-579, which cannot run CUDA 13):
pip install miniworld-engine --index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://pypi.org/simple
```

Shipped kernels that do not match the running torch / CUDA or GPU are ignored, and anything they do not cover still
compiles (`MINIWORLD_ENGINE_JIT_FALLBACK=0` makes that an error). How kernels are keyed by their include graph, how
the release builds and checks them (`miniworld-engine dev bundle-kernels`, `dev build-wheel`), and the release
checklist: [docs/kernels/prebuilt-kernels.md](docs/kernels/prebuilt-kernels.md).

## Layout

```
src/miniworld_engine/
├── kernels/            fusion units — one folder per family
│   ├── <family>/       reference.py · interface.py · [dispatch.py] · [whole_op.py]
│   │                   cuda/ · triton/ backends · notes/ (dated optimization logs)
│   ├── drivers/  checks/   autotune-capture harness: one module per registry family
│   └── registry/       registry.csv (declared kernels) · registry_module.csv (model shapes)
│                       · registry_kernel.csv (derived plan) · exemption / evidence lists
├── modules/            model ops: connect kernels, no backend folders
├── integrations/       hand-CUDA entry points wired into modules (trimul_h100, token_dit, …)
├── autotune/           tuner, cache reader/builder · configs/{default,grid,ab/} · data/ (per-GPU caches)
├── build/              cache-build matrix and audits
├── ops/                whole-op public API (`from miniworld_engine import ops`)
├── tools/              maintenance commands behind `miniworld-engine dev …` (import-anthropic, …)
├── viz/  cli.py  settings.py
benchmarks/             runners/bench.py + {kernels,modules}/<target>/configs · results/<gpu>/
tests/                  CPU contracts; `@pytest.mark.gpu` for device tests
docs/                   see docs/README.md
third_party/  licenses/ upstream provenance and license texts
```

Every kernel family has the same shape, enforced by `tests/cpu/layout/test_kernel_layout.py`:
`reference.py` (the torch definition checks compare against), `interface.py` (the family's one
public door), optional `dispatch.py` (a choice among implementations) and `whole_op.py`
(a whole layer with weights as arguments), and backend packages `cuda/` / `triton/`.

## Benchmarking

One entry point: `benchmarks/runners/bench.py` with a target's `configs/bench.yaml`
(`target=<name> level=module|kernel`). Final numbers are compiled and CUDA-graph timed; the
CSV is the source of truth and plots are rendered from it. Conventions:
[docs/benchmarks/](benchmarks/README.md) · traps: [cautions](benchmarks/cautions.md) ·
cluster commands per GPU: [docs/gpus/](docs/gpus/README.md).

## torch.compile

Every kernel entry point is registered as an opaque `torch.library` op, so a compiled model is
ONE graph rather than one per kernel — a pairformer block traces to 1 graph / 0 breaks instead of
27 / 26. That is `settings.compile_wrap="custom_op"`, the default.

```bash
MINIWORLD_COMPILE_WRAP=disable   # the other mode: a graph break at every kernel entry
```

`disable` is kept for A/B and as the escape hatch: it is the mode that needs no `fake`, so it
still works if one is ever wrong. It has to come from the environment because the value is read
when the kernel modules IMPORT — `settings.configure()` from a parent process is too late.

The default matters beyond fusion. Under `disable`, inductor's cudagraph-trees
(`mode="reduce-overhead"`) bail on the breaks and end up SLOWER than eager, and a manual
`torch.cuda.graph` capture over a compiled module dies with `cudaErrorStreamCaptureInvalidated`.
Numbers, and the scripts that produced them: `benchmarks/compile_wrap/`.

## Supported hardware

Every kernel declares its minimum architecture in `kernels/registry/registry.csv`'s `arch` column, and the
table below is checked against that column by `tests/cpu/registry/test_hardware_support.py` — so it cannot drift
from the code.

<!-- BEGIN GENERATED: hardware-support -->
| arch | GPUs | kernels | backends |
|---|---|---|---|
| **sm80+** | A100, A5000, A6000, RTX 4090 | 132 | triton 126, cuda 6 |
| **sm90+** | H100 | 34 | cuda 30, triton 4 |
<!-- END GENERATED: hardware-support -->

**The minimum supported architecture is sm80.** The generated table above states each
registered kernel's minimum architecture; it is not a numerical qualification for every
listed GPU. Portable Triton paths and higher-architecture Triton/CUDA alternatives are
selected by `modules/dispatch.py`. `implementation='triton'` explicitly selects the
Triton backend; native CUDA variants can require a newer architecture.

One extension is **not** in the table because it is not in the registry: `transition_b2b_cuda`,
which the `Transition` module builds on demand, is compiled for `sm_90a` and fails to build on
sm_86 ("Error building extension"). `autotune/builder.py` excludes `cuda` from that case's
implementations for exactly this reason.

The B200 (sm_100a) kernels are not in the table either: like the H100 hand-CUDA paths they are
integrations wired into the modules (`integrations/`, gated on compute capability 10.0), not
registry kernels. What runs there, for which shapes: [docs/gpus/b200/b200.md](docs/gpus/b200/b200.md).
## CLI

`miniworld-engine` (installed by the package; `python -m miniworld_engine.cli` works too):

```bash
miniworld-engine build all            # verify the current model plan, then fill its cache gaps
miniworld-engine build all --resume   # default: reuse compatible completed measurements
miniworld-engine dev coverage --arch sm86    # default cache key: NVIDIA RTX A6000 (sm86)
miniworld-engine dev audit            # registry, tuning and build-system contract checks
```

`build all` runs only the configured FoldForge/MiniWorld module shapes, then checks
required cache coverage after merging. Unreachable diagnostic kernels are not appended to
the default build; use explicit `--per-op` for a separate kernel experiment.
Model shape policy (`archive/records-20260928:docs/records/model-shape-cleanup-20260916.md`). Declared invocations, selected work and cache keys
are different counts; the command prints them for the current source and GPU.
A claim file alone does not prove completion: resume requires reusable measurement shards
and matching provenance. `--no-resume` disables completed-shard reuse.

`build` writes to the installed package's `autotune/data/` directory
(`src/miniworld_engine/autotune/data/` in a checkout). It requires a writable installation
or checkout and does not automatically commit results. A successful cache build does not
replace module numerical tests or benchmarks. Full policy:
[dispatch-cache.md](docs/kernels/autotune-dispatch-cache.md).
## Research history

Research capsules and older docs were removed from the tree; only the fastest variant of each
kernel lives in `src/`. Everything is still in git:

| where | contains |
|---|---|
| tag `archive/experiments-20260928` | the last tree with `experiments/`: `transition_fused`, `trimul_b7b12`, `trimul_k1k3_inference`, `token_dit_fused`, `token_dit_overlap`, `token_dit_train` |
| tag `archive/wip-main-20260928` | the Sept 27–28 local-H100/Vast research: `trimul_large_d_vast`, `transition_wide_fusion`, `transition_shapes_vast`, `triattn_baseline_20260927`, `triattn_compare_20260927`, `triattn_local_d128`, `trimul_config_audit_20260927`, `trimul_d128_config_20260927` |
| tag `archive/pre-tidy-20260927` | `trimul_training_v2` (2.0.0 training capsule, 6,625 run files), `legacy_branches`, `legacy_h100_runtime`, retired cache blobs |
| tag `archive/docs-20260928`, `archive/records-20260928` | removed docs: dated records, verdicts, design proposals, analyses, release pages |
| tag `archive/scripts-20260928` | the Sept H100 bring-up scripts |
| tag `archive/a100-sm80-branch-20260928` | the A100 branch as merged: `experiments/a100_trimul_fwd` (sm_80 CUDA TriMul, not ported), `a100_transition_*`, `a100_msa_fwd`, `a100_anthropic_baseline`, `a100_trimul_triton`, `msa_triton` |

```sh
git show archive/experiments-20260928:experiments/transition_fused/README.md
git archive archive/experiments-20260928 experiments/trimul_k1k3_inference | tar -x -C /tmp/restore
git show archive/wip-main-20260928:experiments/trimul_large_d_vast/LATENCY_TABLE.md
```

Where each capsule ended up in `src/`:

| capsule | production code |
|---|---|
| `transition_fused`, `transition_shapes_vast` | `kernels/transition/cuda/fused_sm90a.py`, `fused_wide_sm90a.py` |
| `trimul_k1k3_inference` | `kernels/trimul_inproj/cuda/h100_inference.py` (byte-identical K1/K3 overlay sources) |
| `trimul_b7b12`, `trimul_training_v2` | `kernels/trimul_inproj/cuda/h100_training.py` (D128 B1/B7) |
| `trimul_large_d_vast` | `kernels/trimul_inproj/cuda/h100_wide_training.py` (D256/384/512 training) |
| `token_dit_train` | `kernels/augmented_attention/cuda/` |
| `token_dit_fused`, `token_dit_overlap` | not yet ported (quack-GEMM dependent); see the v2.2.0 pending list in `docs/CHANGELOG.md` |
| `transition_wide_fusion` | not ported (explicit 1.03x D384/512 candidate) |

## Toolchain

One-time, per clone — git will not let a repository point itself at its own hooks:

```bash
git config core.hooksPath .githooks   # refuses to commit tuned cache data with code
```

```bash
pixi run ruff-check     # lint  (src tests benchmarks)
pixi run types          # ty    (src tests benchmarks) -- gates CI, no findings allowed
pixi run test           # all tests not marked gpu, on an allocated CPU node
pixi run test-gpu       # all tests marked gpu, on an allocated GPU node
pixi run ci             # all three, in CI's order
```

`ty`, not pyright: pyright cannot parse jaxtyping shape strings
(`Float[torch.Tensor, "N S 3"]`) and reported 144 parse errors and no real findings, so
`[tool.pyright]` turns it off and exists only to keep an editor quiet.
