miniworld-engine 2.5.2__tar.gz → 2.5.4__tar.gz

miniworld_engine-2.5.4/PKG-INFO

@@ -0,0 +1,218 @@
+Metadata-Version: 2.4
+Name: miniworld-engine
+Version: 2.5.4
+Summary: H100 + A100 CUDA / Triton kernels for AlphaFold3-style modules, with a PyTorch reference backend per module, accuracy and speed measurers, and a version-based judge.
+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/anthropic-trimul-native-v5-UPSTREAM.json
+License-File: licenses/CUTLASS-LICENSE.txt
+License-File: licenses/FlashAttention-LICENSE.txt
+Requires-Dist: torch>=2.8
+Requires-Dist: triton>=3.3
+Requires-Dist: einops
+Requires-Dist: jaxtyping
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: pydantic
+Requires-Dist: pyyaml
+Requires-Dist: cuda-bindings
+Requires-Dist: ninja
+Requires-Dist: miniworld-engine-kernels-sm90-cu129==2.5.4; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
+Requires-Dist: miniworld-engine-kernels-sm80-cu129==2.5.4; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
+Provides-Extra: baselines
+Requires-Dist: cuequivariance-torch<0.13,>=0.12; extra == "baselines"
+Provides-Extra: flash
+Requires-Dist: flash-attn-4; extra == "flash"
+Provides-Extra: dev
+Requires-Dist: ruff==0.16.9; extra == "dev"
+Requires-Dist: pyright; extra == "dev"
+Requires-Dist: matplotlib; extra == "dev"
+Dynamic: license-file
+
+# miniworld-engine
+
+GPU kernels for the building blocks of AlphaFold 3-style structure-prediction models: the pair stack (triangle multiplication and
+attention, outer product mean, MSA pair-weighted averaging, transitions, layer norms), the token and atom diffusion transformers of
+several model families, the diffusion sampler's glue, the smooth-lDDT loss, and ProteinMPNN. Each op is a PyTorch module with a
+**PyTorch reference** and interchangeable backends: the engine's own hand-written **CUDA** kernels, its portable **Triton** kernels,
+and two external baselines (Anthropic's released kernels, NVIDIA cuEquivariance). The repository also contains the tooling that
+says what each module must serve (requirements tables), measures accuracy against a float64 gold and speed under CUDA graphs, and
+decides, from human-managed version numbers, which results are current (the judge).
+
+## Supported cards
+
+| card | key | `cuda` backend | status |
+|---|---|---|---|
+| NVIDIA H100 | `h100_sm90` | sm_90a kernels (wgmma, TMA, CUTLASS / cuBLASDx where used) | supported: every requirement case has a current result (judge: 0 owed) ([benchmarks](benchmarks/h100_sm90/README.md)) |
+| NVIDIA A100 | `a100_sm80` | sm_80 kernels (mma.sync) | supported: every requirement case has a current result (judge: 0 owed) ([benchmarks](benchmarks/a100_sm80/README.md)) |
+| NVIDIA B200 | `b200_sm100` | sm_100a sources exist in the tree | **not supported**: almost every `cuda` backend refuses the card; no launch-config tables, tokens or benchmarks |
+| other CUDA GPUs | — | refused (`BackendUnsupported`) | the `triton` and `pytorch` backends run; Triton launch configs are tuned only for H100 / A100 |
+
+## Modules and the models that use them
+
+Every class is in `miniworld_engine.modules`. The models are the ones whose constructors produce the class's requirement rows
+(`source` column of [requirements_module.csv](src/miniworld_engine/requirements/requirements_module.csv)): AF3-family models built by
+FoldForge, and MiniWorld. `engine-own` rows are shapes the engine keeps although no current downstream constructor builds them.
+
+| group | classes | used by (requirement `source`) |
+|---|---|---|
+| triangle multiplication | `TriangleMultiplication` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix (v1, v2), RoseTTAFold3 |
+| | `TriangleMultiplicationRF3` | RoseTTAFold3 |
+| | `BidirectionalTriangleMultiplication` | MiniWorld; engine-own |
+| triangle attention | `TriangleAttention` | Boltz-2, OpenDDE, OpenFold3-preview2, Protenix (v1, v2) |
+| | `TriangleAttentionAF3` | AF3, IntelliFold2, OpenFold3 |
+| | `TriangleAttentionRF3` | RoseTTAFold3 |
+| | `TriangleAttentionChai1`, `TriangleAttentionChai1Confidence` | Chai-1 |
+| | `BidirectionalTriangleAttention` | engine-own |
+| MSA ops | `OuterProductMean`, `MSAPairWeightedAveraging` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, RoseTTAFold3, MiniWorld |
+| transition, norms | `Transition`, `LayerNorm` | all of the above models and MiniWorld |
+| attention with pair bias | `AttentionPairBias` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, RoseTTAFold3, MiniWorld |
+| pairformer block | `PairformerBlock` | engine-own (MiniWorld block layout) |
+| token diffusion transformer | `TokenDiTProtenix` (selector `TokenDiT(model=...)`) | AF3, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, MiniWorld |
+| | `TokenDiTBoltz2`, `TokenDiTRoseTTAFold3`, `TokenDiTChai1` | Boltz-2, RoseTTAFold3, Chai-1 |
+| | `TokenDiTMiniWorld`, `TokenDiTMiniWorldBiasOnly` | MiniWorld |
+| atom diffusion transformer | `AtomDiTOpenFold3` (selector `AtomDiT(model=...)`) | AF3, IntelliFold2, OpenFold3 |
+| | `AtomDiTProtenix` | OpenDDE, Protenix |
+| | `AtomDiTBoltz2`, `AtomDiTRoseTTAFold3`, `AtomDiTChai1` | Boltz-2, RoseTTAFold3, Chai-1 |
+| | `AtomDiTESMFold2` | ESMFold2 |
+| | `SWAAtomTransformer`, `InputFeatureEmbedder` | MiniWorld |
+| diffusion sampler, loss | `AugmentNoise`, `AlignEuler`, `PoolAtoms`, `SmoothLddtLoss` | engine-own (AF3-style sampler glue and smooth-lDDT loss) |
+| La-Proteina pair embeddings | `PairEmbedAdaLN`, `PairEmbedGeom`, `PairEmbedCA`, `TokenPairInit` | engine-own |
+| inverse folding | `ProteinMPNN` | engine-own |
+
+38 classes, 146 requirement rows. "MiniWorld" covers its confidence, diffusion and distogram models.
+
+
+## Backends
+
+| backend | what | accuracy rule |
+|---|---|---|
+| `pytorch` | the reference: the single source of truth of each op's math | the reference |
+| `cuda` | the engine's hand-written kernels for the input's card (H100 or A100) | binding: GREEN = within `max(err_pytorch, fixed[dtype])` of the fp64 gold |
+| `triton` | the engine's Triton kernels (no TMA), any CUDA card | binding |
+| `anthropic` | Anthropic's [uplifting-biomolecular-modeling](https://github.com/anthropics/uplifting-biomolecular-modeling) release, connected as it ships | informational (baseline) |
+| `cueq` | NVIDIA cuEquivariance 0.12 (`cuequivariance_torch`), connected as it ships | informational (baseline) |
+
+A call outside a backend's range raises `BackendUnsupported` with the reason; it is never run by another backend. A kernel that fails
+to build raises `KernelBuildError`; there is no silent fallback anywhere. Which classes have which backends: the `backends` tuple of
+each class ([docs/modules.md](docs/modules.md)).
+
+## Results
+
+Per-module speed, geometric mean over each class's requirement cases of the speedup vs the `pytorch` reference (the engine's
+reference implementation, eager PyTorch kernels in the same CUDA graph, TF32 allowed, not `torch.compile`d), from the speed tokens.
+Each card's README has the full table, the AF3-only table and per-class L / D sweeps.
+
+| | H100 ([summary](benchmarks/h100_sm90/README.md)) | A100 ([summary](benchmarks/a100_sm80/README.md)) |
+|---|---|---|
+| all classes | ![H100 summary](benchmarks/h100_sm90/summary.svg) | ![A100 summary](benchmarks/a100_sm80/summary.svg) |
+
+A few rows (eval / train, Engine CUDA vs PyTorch; in brackets the Anthropic and cuEquivariance eval speedups vs PyTorch):
+
+| class | H100 eval | H100 train | A100 eval | A100 train |
+|---|---|---|---|---|
+| TriangleMultiplication | 9.76x (4.64x, 3.39x) | 8.67x | 7.63x (5.44x, 3.74x) | 6.21x |
+| TriangleAttention | 12.57x (6.53x, 3.11x) | 6.59x | 10.36x (5.49x, 3.19x) | 5.85x |
+| Transition | 4.37x (2.62x, —) | 3.14x | 3.56x (2.47x, —) | 2.52x |
+| AttentionPairBias | 4.54x (1.55x, 0.93x) | 3.41x | 4.33x (1.57x, 0.94x) | 2.28x |
+| TokenDiTProtenix | 3.56x (1.22x, 1.02x) | 2.50x | 2.54x (1.19x, 0.84x) | 1.73x |
+| AtomDiTOpenFold3 | 5.29x (1.19x, —) | 2.78x | 4.74x (1.21x, —) | 2.82x |
+| PairformerBlock | 12.83x (8.24x, —) | 7.84x | 9.57x (6.92x, —) | 6.19x |
+
+These are module-level numbers on the requirement shapes, not end-to-end model speedups. Accuracy: every engine-backend accuracy token
+of both cards is in `src/miniworld_engine/judge/tokens/<card>/` (`kind: accuracy`, `light`, `err_backend`, `err_pytorch`, `tol`).
+
+## Installation
+
+See [docs/install.md](docs/install.md). In short:
+
+1. Python >= 3.10 with PyTorch >= 2.8 and Triton >= 3.3; `pip install -e .` (the extra `baselines` adds cuEquivariance for `cueq`).
+2. On Linux x86-64 with CPython 3.12 the prebuilt H100 kernels come as a dependency wheel; elsewhere each CUDA kernel is compiled on
+   first use (needs an `nvcc` matching torch's CUDA and, for the CUTLASS kernels, `CUTLASS_PATH`) or ahead of time with the prebuilder.
+3. The `anthropic` backend needs the Anthropic release at `third_party/anthropic/upstream` (revision in `third_party/anthropic/UPSTREAM.json`)
+   or `MINIWORLD_ANTHROPIC_ROOT`.
+
+## Quick usage
+
+```python
+import torch
+from miniworld_engine.modules import TriangleMultiplication, TokenDiT
+from miniworld_engine.modules.primitives import BackendUnsupported
+
+tri = TriangleMultiplication(128, backend="cuda").cuda()     # H100 or A100 kernels, picked from the input's device
+ref = TriangleMultiplication(128).cuda()                     # backend="pytorch", the reference
+ref.load_state_dict(tri.state_dict())                        # same parameter names for every backend
+tri.set_backend("triton")                                    # switch later; ValueError if the class has no such backend
+
+dit = TokenDiT(model="boltz2", ...)                          # selector: builds the model's own class (TokenDiTBoltz2)
+```
+
+Mixed precision as the models run it (`master` in the requirements): fp32 parameters, bf16 activations, forward under
+`torch.autocast("cuda", torch.bfloat16)`; call `set_master_weights()` on a `MasterWeights` module after a bulk bf16 cast.
+
+## Building, measuring, judging
+
+Every Python step below runs on a compute node (the measurers need a GPU of the card; the prebuilder and the judge run on a CPU).
+
+```sh
+# ahead-of-time CUDA builds for a card, planned from the code and the requirements
+python -m miniworld_engine.kernels._prebuild --list                 # H100 (sm90a): plan and what is missing
+python -m miniworld_engine.kernels._prebuild --arch sm80 --jobs 16 --cutlass $CUTLASS_PATH
+
+# which requirement cases lack a current result, as Slurm work groups
+python -m miniworld_engine.judge --gpu h100_sm90 owed --plan $DIR
+# for each line "<kind> <Module> <cfg>" of $DIR/G<i>, on the card:
+python -m miniworld_engine.measure.<kind> <cfg> --out $DIR/out/<kind>.<Module>.jsonl
+python -m miniworld_engine.judge --gpu h100_sm90 issue $DIR/out   # records made by the current versions -> tokens
+python -m miniworld_engine.judge --gpu h100_sm90 status
+
+# figures and tables from the tokens (never measures)
+python -m miniworld_engine.measure.plot --gpu a100_sm80
+```
+
+`tools/slurm/cycle.sh PLAN_DIR` runs one plan → prebuild → measure → issue cycle on a Slurm cluster (`CARD=a100_sm80 ARCH=sm80` for
+A100). Details: [docs/judge.md](docs/judge.md), [docs/measure.md](docs/measure.md), [docs/kernels.md](docs/kernels.md).
+
+## Repository layout
+
+| path | what |
+|---|---|
+| `src/miniworld_engine/requirements/` | `requirements_module.csv` (what to serve), `requirements_kernel.csv` (which kernels serve it, with versions), the rules |
+| `src/miniworld_engine/modules/` | the `BackendModule` classes, one file per op family; loaders of the external baselines |
+| `src/miniworld_engine/kernels/` | ~40 kernel families (`cuda/`, `triton/`, `configs/` per card), `integrations/` (module-level compositions per card), the build, prebuild and pretune infrastructure |
+| `src/miniworld_engine/measure/` | accuracy and speed measurers, their configs, the benchmark plotter |
+| `src/miniworld_engine/judge/` | the judge, the tokens per card, the recorded answers |
+| `benchmarks/<card>/` | generated speed figures and tables |
+| `docs/` | documentation ([index](docs/README.md), [architecture](docs/architecture.md)); `docs/h100/` per-family H100 design notes |
+| `tools/slurm/` | the measurement cycle script |
+| `third_party/anthropic/` | provenance of the Anthropic release used by the `anthropic` backends |
+
+Start with [docs/architecture.md](docs/architecture.md) for how a module call reaches the card's kernels and how the pieces fit.
+
+## License
+
+`MIT AND Apache-2.0 AND BSD-3-Clause`. MiniWorld Engine's own code is MIT (copyright 2026 Sanggeun Park, `LICENSE`); the code derived
+from Anthropic's uplifting-biomolecular-modeling release, and any file whose header says `SPDX-License-Identifier: Apache-2.0`, is
+Apache-2.0 (with Anthropic's `NOTICE`), and the FlashAttention-3 portions and the
+CUTLASS-instantiated object code in the prebuilt kernels are BSD-3-Clause. Details, per-file attributions and the licence texts are in
+[`licenses/THIRD_PARTY_NOTICES.md`](licenses/THIRD_PARTY_NOTICES.md). The external release used by the `anthropic` backends is cloned
+locally under `third_party/anthropic/upstream/` (git-ignored, not redistributed).
+
+## Citation
+
+A paper describing the engine is in preparation. Until then, please cite the repository:
+
+```bibtex
+@software{miniworld_engine,
+  author  = {Park, Sanggeun},
+  title   = {miniworld-engine: GPU kernels for AlphaFold 3-style modules},
+  year    = {2026},
+  version = {2.5.4},
+  note    = {TODO: replace with the paper reference}
+}
+```

miniworld_engine-2.5.4/README.md

@@ -0,0 +1,182 @@
+# miniworld-engine
+
+GPU kernels for the building blocks of AlphaFold 3-style structure-prediction models: the pair stack (triangle multiplication and
+attention, outer product mean, MSA pair-weighted averaging, transitions, layer norms), the token and atom diffusion transformers of
+several model families, the diffusion sampler's glue, the smooth-lDDT loss, and ProteinMPNN. Each op is a PyTorch module with a
+**PyTorch reference** and interchangeable backends: the engine's own hand-written **CUDA** kernels, its portable **Triton** kernels,
+and two external baselines (Anthropic's released kernels, NVIDIA cuEquivariance). The repository also contains the tooling that
+says what each module must serve (requirements tables), measures accuracy against a float64 gold and speed under CUDA graphs, and
+decides, from human-managed version numbers, which results are current (the judge).
+
+## Supported cards
+
+| card | key | `cuda` backend | status |
+|---|---|---|---|
+| NVIDIA H100 | `h100_sm90` | sm_90a kernels (wgmma, TMA, CUTLASS / cuBLASDx where used) | supported: every requirement case has a current result (judge: 0 owed) ([benchmarks](benchmarks/h100_sm90/README.md)) |
+| NVIDIA A100 | `a100_sm80` | sm_80 kernels (mma.sync) | supported: every requirement case has a current result (judge: 0 owed) ([benchmarks](benchmarks/a100_sm80/README.md)) |
+| NVIDIA B200 | `b200_sm100` | sm_100a sources exist in the tree | **not supported**: almost every `cuda` backend refuses the card; no launch-config tables, tokens or benchmarks |
+| other CUDA GPUs | — | refused (`BackendUnsupported`) | the `triton` and `pytorch` backends run; Triton launch configs are tuned only for H100 / A100 |
+
+## Modules and the models that use them
+
+Every class is in `miniworld_engine.modules`. The models are the ones whose constructors produce the class's requirement rows
+(`source` column of [requirements_module.csv](src/miniworld_engine/requirements/requirements_module.csv)): AF3-family models built by
+FoldForge, and MiniWorld. `engine-own` rows are shapes the engine keeps although no current downstream constructor builds them.
+
+| group | classes | used by (requirement `source`) |
+|---|---|---|
+| triangle multiplication | `TriangleMultiplication` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix (v1, v2), RoseTTAFold3 |
+| | `TriangleMultiplicationRF3` | RoseTTAFold3 |
+| | `BidirectionalTriangleMultiplication` | MiniWorld; engine-own |
+| triangle attention | `TriangleAttention` | Boltz-2, OpenDDE, OpenFold3-preview2, Protenix (v1, v2) |
+| | `TriangleAttentionAF3` | AF3, IntelliFold2, OpenFold3 |
+| | `TriangleAttentionRF3` | RoseTTAFold3 |
+| | `TriangleAttentionChai1`, `TriangleAttentionChai1Confidence` | Chai-1 |
+| | `BidirectionalTriangleAttention` | engine-own |
+| MSA ops | `OuterProductMean`, `MSAPairWeightedAveraging` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, RoseTTAFold3, MiniWorld |
+| transition, norms | `Transition`, `LayerNorm` | all of the above models and MiniWorld |
+| attention with pair bias | `AttentionPairBias` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, RoseTTAFold3, MiniWorld |
+| pairformer block | `PairformerBlock` | engine-own (MiniWorld block layout) |
+| token diffusion transformer | `TokenDiTProtenix` (selector `TokenDiT(model=...)`) | AF3, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, MiniWorld |
+| | `TokenDiTBoltz2`, `TokenDiTRoseTTAFold3`, `TokenDiTChai1` | Boltz-2, RoseTTAFold3, Chai-1 |
+| | `TokenDiTMiniWorld`, `TokenDiTMiniWorldBiasOnly` | MiniWorld |
+| atom diffusion transformer | `AtomDiTOpenFold3` (selector `AtomDiT(model=...)`) | AF3, IntelliFold2, OpenFold3 |
+| | `AtomDiTProtenix` | OpenDDE, Protenix |
+| | `AtomDiTBoltz2`, `AtomDiTRoseTTAFold3`, `AtomDiTChai1` | Boltz-2, RoseTTAFold3, Chai-1 |
+| | `AtomDiTESMFold2` | ESMFold2 |
+| | `SWAAtomTransformer`, `InputFeatureEmbedder` | MiniWorld |
+| diffusion sampler, loss | `AugmentNoise`, `AlignEuler`, `PoolAtoms`, `SmoothLddtLoss` | engine-own (AF3-style sampler glue and smooth-lDDT loss) |
+| La-Proteina pair embeddings | `PairEmbedAdaLN`, `PairEmbedGeom`, `PairEmbedCA`, `TokenPairInit` | engine-own |
+| inverse folding | `ProteinMPNN` | engine-own |
+
+38 classes, 146 requirement rows. "MiniWorld" covers its confidence, diffusion and distogram models.
+
+
+## Backends
+
+| backend | what | accuracy rule |
+|---|---|---|
+| `pytorch` | the reference: the single source of truth of each op's math | the reference |
+| `cuda` | the engine's hand-written kernels for the input's card (H100 or A100) | binding: GREEN = within `max(err_pytorch, fixed[dtype])` of the fp64 gold |
+| `triton` | the engine's Triton kernels (no TMA), any CUDA card | binding |
+| `anthropic` | Anthropic's [uplifting-biomolecular-modeling](https://github.com/anthropics/uplifting-biomolecular-modeling) release, connected as it ships | informational (baseline) |
+| `cueq` | NVIDIA cuEquivariance 0.12 (`cuequivariance_torch`), connected as it ships | informational (baseline) |
+
+A call outside a backend's range raises `BackendUnsupported` with the reason; it is never run by another backend. A kernel that fails
+to build raises `KernelBuildError`; there is no silent fallback anywhere. Which classes have which backends: the `backends` tuple of
+each class ([docs/modules.md](docs/modules.md)).
+
+## Results
+
+Per-module speed, geometric mean over each class's requirement cases of the speedup vs the `pytorch` reference (the engine's
+reference implementation, eager PyTorch kernels in the same CUDA graph, TF32 allowed, not `torch.compile`d), from the speed tokens.
+Each card's README has the full table, the AF3-only table and per-class L / D sweeps.
+
+| | H100 ([summary](benchmarks/h100_sm90/README.md)) | A100 ([summary](benchmarks/a100_sm80/README.md)) |
+|---|---|---|
+| all classes | ![H100 summary](benchmarks/h100_sm90/summary.svg) | ![A100 summary](benchmarks/a100_sm80/summary.svg) |
+
+A few rows (eval / train, Engine CUDA vs PyTorch; in brackets the Anthropic and cuEquivariance eval speedups vs PyTorch):
+
+| class | H100 eval | H100 train | A100 eval | A100 train |
+|---|---|---|---|---|
+| TriangleMultiplication | 9.76x (4.64x, 3.39x) | 8.67x | 7.63x (5.44x, 3.74x) | 6.21x |
+| TriangleAttention | 12.57x (6.53x, 3.11x) | 6.59x | 10.36x (5.49x, 3.19x) | 5.85x |
+| Transition | 4.37x (2.62x, —) | 3.14x | 3.56x (2.47x, —) | 2.52x |
+| AttentionPairBias | 4.54x (1.55x, 0.93x) | 3.41x | 4.33x (1.57x, 0.94x) | 2.28x |
+| TokenDiTProtenix | 3.56x (1.22x, 1.02x) | 2.50x | 2.54x (1.19x, 0.84x) | 1.73x |
+| AtomDiTOpenFold3 | 5.29x (1.19x, —) | 2.78x | 4.74x (1.21x, —) | 2.82x |
+| PairformerBlock | 12.83x (8.24x, —) | 7.84x | 9.57x (6.92x, —) | 6.19x |
+
+These are module-level numbers on the requirement shapes, not end-to-end model speedups. Accuracy: every engine-backend accuracy token
+of both cards is in `src/miniworld_engine/judge/tokens/<card>/` (`kind: accuracy`, `light`, `err_backend`, `err_pytorch`, `tol`).
+
+## Installation
+
+See [docs/install.md](docs/install.md). In short:
+
+1. Python >= 3.10 with PyTorch >= 2.8 and Triton >= 3.3; `pip install -e .` (the extra `baselines` adds cuEquivariance for `cueq`).
+2. On Linux x86-64 with CPython 3.12 the prebuilt H100 kernels come as a dependency wheel; elsewhere each CUDA kernel is compiled on
+   first use (needs an `nvcc` matching torch's CUDA and, for the CUTLASS kernels, `CUTLASS_PATH`) or ahead of time with the prebuilder.
+3. The `anthropic` backend needs the Anthropic release at `third_party/anthropic/upstream` (revision in `third_party/anthropic/UPSTREAM.json`)
+   or `MINIWORLD_ANTHROPIC_ROOT`.
+
+## Quick usage
+
+```python
+import torch
+from miniworld_engine.modules import TriangleMultiplication, TokenDiT
+from miniworld_engine.modules.primitives import BackendUnsupported
+
+tri = TriangleMultiplication(128, backend="cuda").cuda()     # H100 or A100 kernels, picked from the input's device
+ref = TriangleMultiplication(128).cuda()                     # backend="pytorch", the reference
+ref.load_state_dict(tri.state_dict())                        # same parameter names for every backend
+tri.set_backend("triton")                                    # switch later; ValueError if the class has no such backend
+
+dit = TokenDiT(model="boltz2", ...)                          # selector: builds the model's own class (TokenDiTBoltz2)
+```
+
+Mixed precision as the models run it (`master` in the requirements): fp32 parameters, bf16 activations, forward under
+`torch.autocast("cuda", torch.bfloat16)`; call `set_master_weights()` on a `MasterWeights` module after a bulk bf16 cast.
+
+## Building, measuring, judging
+
+Every Python step below runs on a compute node (the measurers need a GPU of the card; the prebuilder and the judge run on a CPU).
+
+```sh
+# ahead-of-time CUDA builds for a card, planned from the code and the requirements
+python -m miniworld_engine.kernels._prebuild --list                 # H100 (sm90a): plan and what is missing
+python -m miniworld_engine.kernels._prebuild --arch sm80 --jobs 16 --cutlass $CUTLASS_PATH
+
+# which requirement cases lack a current result, as Slurm work groups
+python -m miniworld_engine.judge --gpu h100_sm90 owed --plan $DIR
+# for each line "<kind> <Module> <cfg>" of $DIR/G<i>, on the card:
+python -m miniworld_engine.measure.<kind> <cfg> --out $DIR/out/<kind>.<Module>.jsonl
+python -m miniworld_engine.judge --gpu h100_sm90 issue $DIR/out   # records made by the current versions -> tokens
+python -m miniworld_engine.judge --gpu h100_sm90 status
+
+# figures and tables from the tokens (never measures)
+python -m miniworld_engine.measure.plot --gpu a100_sm80
+```
+
+`tools/slurm/cycle.sh PLAN_DIR` runs one plan → prebuild → measure → issue cycle on a Slurm cluster (`CARD=a100_sm80 ARCH=sm80` for
+A100). Details: [docs/judge.md](docs/judge.md), [docs/measure.md](docs/measure.md), [docs/kernels.md](docs/kernels.md).
+
+## Repository layout
+
+| path | what |
+|---|---|
+| `src/miniworld_engine/requirements/` | `requirements_module.csv` (what to serve), `requirements_kernel.csv` (which kernels serve it, with versions), the rules |
+| `src/miniworld_engine/modules/` | the `BackendModule` classes, one file per op family; loaders of the external baselines |
+| `src/miniworld_engine/kernels/` | ~40 kernel families (`cuda/`, `triton/`, `configs/` per card), `integrations/` (module-level compositions per card), the build, prebuild and pretune infrastructure |
+| `src/miniworld_engine/measure/` | accuracy and speed measurers, their configs, the benchmark plotter |
+| `src/miniworld_engine/judge/` | the judge, the tokens per card, the recorded answers |
+| `benchmarks/<card>/` | generated speed figures and tables |
+| `docs/` | documentation ([index](docs/README.md), [architecture](docs/architecture.md)); `docs/h100/` per-family H100 design notes |
+| `tools/slurm/` | the measurement cycle script |
+| `third_party/anthropic/` | provenance of the Anthropic release used by the `anthropic` backends |
+
+Start with [docs/architecture.md](docs/architecture.md) for how a module call reaches the card's kernels and how the pieces fit.
+
+## License
+
+`MIT AND Apache-2.0 AND BSD-3-Clause`. MiniWorld Engine's own code is MIT (copyright 2026 Sanggeun Park, `LICENSE`); the code derived
+from Anthropic's uplifting-biomolecular-modeling release, and any file whose header says `SPDX-License-Identifier: Apache-2.0`, is
+Apache-2.0 (with Anthropic's `NOTICE`), and the FlashAttention-3 portions and the
+CUTLASS-instantiated object code in the prebuilt kernels are BSD-3-Clause. Details, per-file attributions and the licence texts are in
+[`licenses/THIRD_PARTY_NOTICES.md`](licenses/THIRD_PARTY_NOTICES.md). The external release used by the `anthropic` backends is cloned
+locally under `third_party/anthropic/upstream/` (git-ignored, not redistributed).
+
+## Citation
+
+A paper describing the engine is in preparation. Until then, please cite the repository:
+
+```bibtex
+@software{miniworld_engine,
+  author  = {Park, Sanggeun},
+  title   = {miniworld-engine: GPU kernels for AlphaFold 3-style modules},
+  year    = {2026},
+  version = {2.5.4},
+  note    = {TODO: replace with the paper reference}
+}
+```

miniworld_engine-2.5.4/pyproject.toml

@@ -0,0 +1,187 @@
+[project]
+name = "miniworld-engine"
+version = "2.5.4"
+requires-python = ">=3.10"
+readme = "README.md"
+description = "H100 + A100 CUDA / Triton kernels for AlphaFold3-style modules, with a PyTorch reference backend per module, accuracy and speed measurers, and a version-based judge."
+license = "MIT AND Apache-2.0 AND BSD-3-Clause"
+license-files = ["LICENSE", "licenses/THIRD_PARTY_NOTICES.md", "licenses/Anthropic-Apache-2.0.txt", "licenses/Anthropic-NOTICE",
+                 "licenses/anthropic-trimul-native-v5-UPSTREAM.json",
+                 "licenses/CUTLASS-LICENSE.txt", "licenses/FlashAttention-LICENSE.txt"]
+# The runtime: what `miniworld_engine.modules` and `miniworld_engine.kernels` import. Measurement (measure/, judge/) adds pyyaml only.
+dependencies = [
+  "torch>=2.8",
+  "triton>=3.3",
+  "einops",
+  "jaxtyping",
+  "numpy",
+  "scipy",          # modules/primitives.py: truncated-normal parameter init
+  "pydantic",       # modules/atom_dit.py: SWA atom block configs
+  "pyyaml",         # measure/config.py: measurement configs
+  "cuda-bindings",  # kernels/trimul_inproj/cuda: the H100 TriMul driver API
+  "ninja",          # torch.utils.cpp_extension: every JIT build of a CUDA kernel (A100, or a Python / torch without prebuilt kernels)
+  # The prebuilt H100 (sm90) and A100 (sm80) kernels (CUDA 12.9 / torch 2.13 / CPython 3.12), both installed into
+  # miniworld_engine/_prebuilt/ (content-keyed paths; the few arch-independent artifacts, multi-arch fatbins / host-only glue, share a
+  # key and are shipped by both, either copy serves both GPUs); elsewhere a kernel
+  # JIT-builds on first use (needs a CUDA toolkit; the CUTLASS / MathDx headers are fetched by kernels/_headers.py, see docs/install.md).
+  'miniworld-engine-kernels-sm90-cu129==2.5.4; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"',
+  'miniworld-engine-kernels-sm80-cu129==2.5.4; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"',
+]
+
+[project.optional-dependencies]
+# External baselines (backends "cueq" / "anthropic" connect to them; never required by the engine's own backends).
+#   (the transformers fork Anthropic's esmfold2 kit needs -- its ef2_msa / ef2_w4 levers import transformers.models.esmfold2 -- is a git
+#   dependency, which PyPI metadata may not carry: the pixi env pins it below; pip users see docs/install.md)
+baselines = [
+  "cuequivariance-torch>=0.12,<0.13",
+]
+# FlashAttention-4 (CuTe DSL, sm90+): the global-attention path of the SWA atom transformer's cuda backend. Prerelease wheel:
+#   uv pip install --prerelease=allow "flash-attn-4[cu13]"
+flash = ["flash-attn-4"]
+dev = ["ruff==0.16.9", "pyright", "matplotlib"]   # matplotlib: measure/plot.py (benchmarks/)
+
+[build-system]
+requires = ["setuptools>=77"]  # >=77 for PEP 639 (SPDX `license` + `license-files`)
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+# kernel families are namespace packages (no __init__.py); the generator script of swa_dit's wgmma_bf16.inc is build-time only
+namespaces = true
+exclude = ["*.swa_dit.cuda.gen"]
+
+[tool.setuptools.package-data]
+# Non-Python files the runtime reads in a wheel install:
+#   **/*.cu / *.cuh / *.cpp / *.h / *.inc   CUDA sources the prebuilder compiles (and a JIT build falls back to)
+#   **/configs/*.json                       launch-config tables per GPU and candidate lists (kernels/<family>/configs/)
+#   requirements/*.csv                      the requirement lists the measurers, the judge and the prebuild plan read
+#   measure/configs/*.yaml                  the canonical measurement configs
+# Not shipped: judge/tokens/ and judge/answers.jsonl (the judge's measurement state, kept in the repository, not runtime).
+miniworld_engine = ["**/*.cu", "**/*.cuh", "**/*.cpp", "**/*.h", "**/*.inc", "**/*.json", "**/*.csv", "**/*.yaml"]
+[tool.setuptools.exclude-package-data]
+miniworld_engine = ["judge/tokens/**", "judge/answers.jsonl"]
+
+# ----------------------------------------------------------------------------------------------------------------------------- ruff
+[tool.ruff]
+line-length = 140
+target-version = "py310"
+src = ["src"]
+
+[tool.ruff.lint]
+# Correctness and modernity, one import style, one string style. What is off and why: ANN (kernel code is unannotated where
+# triton hands back Unknown; pyright is the type gate), D (docstrings are prose), N (math notation: L, D, Q, K, V, Wp),
+# S / PL / SLF / ARG / T20 (subprocess builds, lazy imports that avoid triton compilation at import, reaching into triton
+# internals on purpose, uniform launcher signatures, CLI prints), TC (jaxtyping resolves annotations at runtime).
+select = [
+  "E", "F", "W", "I", "UP", "B",
+  "C4", "PIE", "RET", "SIM", "PERF", "FURB",
+  "TID", "ICN", "FA",
+  "Q", "ISC",
+  "LOG", "G",
+  "SLOT", "INT", "FLY", "NPY", "DTZ", "EXE", "YTT", "ASYNC",
+  "RUF",
+]
+ignore = [
+  "E501",    # line length is the formatter's target, not a hard wrap for tables and prose
+  "E701", "E702",
+  "E402",
+  "E731",    # one-line kernel adapters
+  "E741",    # `l`, `L`, `I` are dimensions
+  "F722",    # jaxtyping shape strings
+  "COM812",
+  "RUF002", "RUF003",   # full-width characters in Korean comments
+]
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+[tool.ruff.lint.per-file-ignores]
+# Kernel bodies are faithful copies (team-gm branches, FlashAttentionBias, Anthropic kits): not linted, as the originals are not.
+"src/miniworld_engine/kernels/**/triton/*.py" = ["ALL"]
+"src/miniworld_engine/kernels/**/cuda/*.py" = ["ALL"]
+"src/miniworld_engine/kernels/**/cuda/gen/*.py" = ["ALL"]
+
+[tool.ruff.format]
+docstring-code-format = false
+
+# ---------------------------------------------------------------------------------------------------------------------------- pytest
+[tool.pytest.ini_options]
+# tests/cpu run anywhere (no GPU); tests/gpu need a CUDA device and skip cleanly without one (or without the card a marker names).
+#   CPU:  pytest -m "not gpu"        GPU (one H100 job):  pytest tests/gpu -m gpu
+testpaths = ["tests"]
+pythonpath = ["src"]             # the checkout under test, ahead of any installed miniworld-engine
+addopts = "-ra --strict-markers --import-mode=importlib -p no:cacheprovider"
+markers = [
+  "gpu: needs a CUDA device (skipped without one)",
+  "h100: needs an H100 (sm_90); skipped on other cards",
+  "a100: needs an A100 (sm_80); skipped on other cards",
+  "slow: takes more than a few seconds (deselect with -m \"not slow\")",
+]
+
+# --------------------------------------------------------------------------------------------------------------------------- pyright
+[tool.pyright]
+# The type gate (tools/check.sh, 0 errors): basic mode over the package and the tests. Out: the kernel bodies ruff leaves alone
+# (kernels/**/triton, kernels/**/cuda: faithful copies of team-gm / FlashAttentionBias / Anthropic kits). Checked, they report ~1000
+# errors in ~90 files, nearly all from idioms the stubs do not model (tl.constexpr parameters given plain values, torch.empty(*shape)
+# unpacking, opaque launch tuples); fixing them means editing @triton.jit sources, whose text keys the Triton cache.
+include = ["src", "tests"]
+exclude = [
+  ".pixi", "**/.pixi", "**/__pycache__",
+  # the same vendored kernel bodies ruff leaves alone
+  "src/miniworld_engine/kernels/**/triton/**",
+  "src/miniworld_engine/kernels/**/cuda/**",
+]
+pythonVersion = "3.10"
+typeCheckingMode = "basic"
+
+# ----------------------------------------------------------------------------------------------------------------------------- pixi
+# One environment for the repo (torch 2.13 cu129 + triton 3.7, cuequivariance baseline, TE). The materialized `.pixi/` is ignored.
+# Pinned to CUDA 12.9 (driver >= 575): cu12 only; on a cluster whose compute nodes are offline, resolve / download on the login node;
+# at runtime export LD_LIBRARY_PATH=$CONDA_PREFIX/lib.
+[tool.pixi.workspace]
+channels = ["conda-forge", "nvidia"]
+platforms = ["linux-64"]
+
+[tool.pixi.dependencies]
+python = "3.12.*"
+pip = "*"
+cuda-toolkit = "12.9.*"
+cudnn = "*"
+nccl = "*"
+cuda-nvtx-dev = "12.9.*"
+ninja = "*"
+cmake = "*"
+setuptools = "*"
+wheel = "*"
+pybind11 = "*"
+
+[tool.pixi.pypi-options]
+extra-index-urls = ["https://download.pytorch.org/whl/cu129"]
+no-build-isolation = ["transformer-engine-torch"]
+
+[tool.pixi.pypi-dependencies]
+miniworld-engine = { path = ".", editable = true, extras = ["baselines", "dev"] }
+transformers = { git = "https://github.com/huggingface/transformers.git", rev = "ef32577f55da19a4989cd7b22e004dc43a4998cb" }   # Anthropic esmfold2 kit (baselines)
+torch = { version = "==2.13.0", index = "https://download.pytorch.org/whl/cu129" }
+transformer-engine = { version = "==2.19.*", extras = ["pytorch", "core-cu12"] }
+transformer-engine-cu12 = "==2.19.*"
+cuequivariance-ops-torch-cu12 = "==0.12.*"
+
+[tool.pixi.activation.env]
+PYTHONNOUSERSITE = "1"
+CUDA_HOME = "$CONDA_PREFIX"
+
+[tool.pixi.tasks]
+# transformer-engine-torch's sdist declares the cu13 core; reinstall the cu12 one last. Run after every `pixi install`.
+fix-te-cu12 = "pip install --no-deps --force-reinstall transformer_engine_cu12==2.19.0"
+lint = "ruff check src tools tests"
+format = "ruff format src"
+types = "pyright"
+# the gate (tools/check.sh: ruff, pyright, CPU tests) and the GPU tests; on the cluster inside a Slurm job, never on the login node
+check = "tools/check.sh"
+test = "pytest tests/cpu -m 'not gpu'"
+test-gpu = "pytest tests/gpu -m gpu"
+# Everything below needs Slurm (never the login node): prebuild on a CPU node, pretune / measure on an H100.
+prebuild = "python -m miniworld_engine.kernels._prebuild"
+pretune = "python -m miniworld_engine.kernels._pretune"
+judge = "python -m miniworld_engine.judge"

miniworld_engine-2.5.4/src/miniworld_engine/capture.py

@@ -0,0 +1,38 @@
+"""The public CUDA-graph capture and weight-cache API.
+
+``miniworld_engine`` is a namespace package (no ``__init__.py``: the prebuilt kernel wheels install into it), so nothing is exported from
+the package root; this module is the public home of what a caller that captures CUDA graphs or replaces weights needs::
+
+    from miniworld_engine.capture import clear_weight_caches, static_inputs, static_weights
+
+    with static_weights(), static_inputs():      # an inference graph over one item: weights and inputs fixed between replays
+        graph = capture(model_step)
+
+    clear_weight_caches(model)                   # before loading a second set of weights (drop old graphs first)
+
+``static_weights()``
+    Weights stay fixed between graph replays (inference graphs only): a capture inside it serves from, and fills, the eager weight-pack
+    caches, so a replay runs none of the packing / casting kernels. An in-place weight update after the capture is then NOT seen by
+    replays: never around a training graph.
+``static_inputs()``
+    Stronger: the conditioning / pair inputs of the hoisted per-item tables (a block's conditioning table, a pair's bias) also keep their
+    contents for the life of the graph, so a replay re-runs none of their producing kernels. A sampler that refills those tensors in
+    place between replays must not set it.
+``clear_weight_caches(*modules)``
+    Drop every weight- or input-derived cache of the engine (each store registered with ``kernels._capture.weight_cache``: the
+    parameter forms, weight packs, hoisted tables, the TF32 operand cache) and, for the given modules and their submodules, the
+    per-instance cache attributes. Call it before weights are replaced by new tensors (a second model, a test that builds many): a new
+    tensor allocated at the address of a freed one must not be served the old one's form. Graphs captured before the call may read the
+    dropped entries: drop them first.
+``weight_caches()``
+    The registered stores by name (inspection; a kernel module not imported yet has registered nothing and filled nothing).
+
+Semantics are those of ``miniworld_engine.kernels._capture`` (its module docstring: capture scopes, the fused-optimizer
+``increment_version`` rule); this module is the stable name for them.
+"""
+
+from __future__ import annotations
+
+from miniworld_engine.kernels._capture import clear_weight_caches, static_inputs, static_weights, weight_caches
+
+__all__ = ["clear_weight_caches", "static_inputs", "static_weights", "weight_caches"]