miniworld-engine 2.5.0__tar.gz → 2.5.1__tar.gz

miniworld_engine-2.5.1/PKG-INFO

@@ -0,0 +1,94 @@
+Metadata-Version: 2.4
+Name: miniworld-engine
+Version: 2.5.1
+Summary: H100 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/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: miniworld-engine-kernels-sm90-cu129==2.5.1; 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
+
+PyTorch modules of AF3-style structure-prediction models (pair stack, MSA ops, token / atom diffusion transformers, sampler glue,
+ProteinMPNN, losses), each with a PyTorch reference and hand-written H100 kernels as alternative backends, plus the tooling that
+states what each module must serve and judges whether a backend is accurate and fast enough.
+
+Scope: **H100 (sm90) only.** The `cuda` backends are built for sm_90a; the launch-config tables, the prebuilder's default target and
+the judge's recorded environment are H100.
+
+## The four parts
+
+| part | where | what |
+|---|---|---|
+| requirements | `src/miniworld_engine/requirements/` | `requirements_module.csv`: every (module class, dims, options) a downstream model builds, with lengths, dtypes, modes. `requirements_kernel.csv`: every kernel a module backend launches, its family, role and human-managed `version`. |
+| modules + backends | `src/miniworld_engine/modules/` | One `BackendModule` per op. `pytorch` is the reference; `cuda` / `triton` are the engine's kernels; `anthropic` / `cueq` are external baselines. A backend states its range in `_supports_<backend>` and never falls back. |
+| kernels | `src/miniworld_engine/kernels/` | Kernel families (CUDA and Triton), per-GPU launch-config tables, the prebuilder (CUDA artifacts ahead of time) and the pretuner (launch-config tables). |
+| measure + judge | `src/miniworld_engine/measure/`, `src/miniworld_engine/judge/` | Measurers for accuracy (fp64 gold) and speed (CUDA graphs); the judge turns records into tokens and decides which cases are owed a measurement, by human-managed versions. |
+
+## Quick usage
+
+Build a module with a backend:
+
+```python
+import torch
+from miniworld_engine.modules import TriangleMultiplication, TokenDiT
+
+tri = TriangleMultiplication(128, backend="cuda").cuda()     # H100 kernels
+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)
+```
+
+A call outside a backend's range raises `BackendUnsupported` with the reason; it is never run by another backend.
+
+Prebuild the CUDA kernels (CPU node, no GPU needed):
+
+```sh
+python -m miniworld_engine.kernels._prebuild --list            # plan from the code + requirements, what is missing
+python -m miniworld_engine.kernels._prebuild --jobs 16 --cutlass $CUTLASS_PATH
+```
+
+Measure and judge (the measurers need a GPU; the judge runs on a CPU):
+
+```sh
+python -m miniworld_engine.judge owed --plan $DIR             # which cases are owed; writes configs + groups G1..G4
+# one 1-GPU, 10-minute job per group: for each line "<kind> <module> <cfg>" of $DIR/G<i>
+python -m miniworld_engine.measure.<kind> <cfg> --out $DIR/out/<kind>.<module>.jsonl
+python -m miniworld_engine.judge issue $DIR/out               # records made by current versions -> tokens
+python -m miniworld_engine.judge status
+```
+
+## Documentation
+
+[docs/README.md](docs/README.md) is the index: requirements, modules, kernels, measurement, judge, and per-family H100 kernel notes.
+
+## License
+
+MIT (copyright 2026 Sanggeun Park). The license text and third-party notices are `LICENSE` and `licenses/` at the repository root;
+the external release used by the `anthropic` backends is under `third_party/anthropic/`.

miniworld_engine-2.5.1/pyproject.toml

@@ -0,0 +1,160 @@
+[project]
+name = "miniworld-engine"
+version = "2.5.1"
+requires-python = ">=3.10"
+readme = "README.md"
+description = "H100 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/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
+  # The prebuilt H100 kernels (CUDA 12.9 / torch 2.13 / CPython 3.12), installed into miniworld_engine/_prebuilt/; elsewhere a kernel
+  # JIT-builds on first use (needs nvcc and, for the CUTLASS kernels, CUTLASS_PATH).
+  'miniworld-engine-kernels-sm90-cu129==2.5.1; 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
+
+# --------------------------------------------------------------------------------------------------------------------------- pyright
+[tool.pyright]
+include = ["src"]
+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 = "standard"
+
+# ----------------------------------------------------------------------------------------------------------------------------- pixi
+# One environment for the repo (torch 2.13 cu129 + triton 3.7, cuequivariance baseline, TE). The materialized `.pixi/` is ignored.
+# Cluster (CSSB H100, driver 575 = CUDA 12.9 max): cu12 only; resolve / download on the login node (compute nodes are offline);
+# 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"
+format = "ruff format src"
+types = "pyright"
+# 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.1/src/miniworld_engine/autotune/data/bias_only_dispatch/NVIDIA_A100_80GB_PCIe_sm80_triton3.7.1.json

@@ -0,0 +1,14 @@
+{
+  "gate|128|128|131072": {
+    "choice": "fused",
+    "ms": {}
+  },
+  "gate|128|128|262144": {
+    "choice": "fused",
+    "ms": {}
+  },
+  "gate|128|128|524288": {
+    "choice": "fused",
+    "ms": {}
+  }
+}

miniworld_engine-2.5.1/src/miniworld_engine/autotune/data/ln_bwd_dispatch/NVIDIA_A100_80GB_PCIe_sm80_triton3.7.1.json

@@ -0,0 +1,65 @@
+{
+  "128|131072": {
+    "ms": {
+      "atomic": 0.152754,
+      "persistent": 0.174959
+    },
+    "path": "atomic"
+  },
+  "128|262144": {
+    "ms": {
+      "atomic": 0.219835,
+      "persistent": 0.260306
+    },
+    "path": "atomic"
+  },
+  "128|4096": {
+    "ms": {
+      "atomic": 0.016848,
+      "persistent": 0.030278
+    },
+    "path": "atomic"
+  },
+  "128|524288": {
+    "ms": {
+      "atomic": 0.305845,
+      "persistent": 0.35656
+    },
+    "path": "atomic"
+  },
+  "16|64": {
+    "ms": {
+      "atomic": 0.011037,
+      "persistent": 0.040249
+    },
+    "path": "atomic"
+  },
+  "256|131072": {
+    "ms": {
+      "atomic": 0.431104,
+      "persistent": 0.558324
+    },
+    "path": "atomic"
+  },
+  "256|4096": {
+    "ms": {
+      "atomic": 0.024739,
+      "persistent": 0.039898
+    },
+    "path": "atomic"
+  },
+  "384|128": {
+    "ms": {
+      "atomic": 0.024792,
+      "persistent": 0.063316
+    },
+    "path": "atomic"
+  },
+  "384|256": {
+    "ms": {
+      "atomic": 0.013423,
+      "persistent": 0.027009
+    },
+    "path": "atomic"
+  }
+}