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

{miniworld_engine-2.5.2/src/miniworld_engine.egg-info → miniworld_engine-2.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: miniworld-engine
-Version: 2.5.2
+Version: 2.5.3
 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
@@ -22,8 +22,8 @@ Requires-Dist: pydantic
 Requires-Dist: pyyaml
 Requires-Dist: cuda-bindings
 Requires-Dist: ninja
-Requires-Dist: miniworld-engine-kernels-sm90-cu129==2.5.2; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
-Requires-Dist: miniworld-engine-kernels-sm80-cu129==2.5.2; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
+Requires-Dist: miniworld-engine-kernels-sm90-cu129==2.5.3; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
+Requires-Dist: miniworld-engine-kernels-sm80-cu129==2.5.3; 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
@@ -212,7 +212,7 @@ A paper describing the engine is in preparation. Until then, please cite the rep
   author  = {Park, Sanggeun},
   title   = {miniworld-engine: GPU kernels for AlphaFold 3-style modules},
   year    = {2026},
-  version = {2.5.2},
+  version = {2.5.3},
   note    = {TODO: replace with the paper reference}
 }
 ```

{miniworld_engine-2.5.2 → miniworld_engine-2.5.3}/README.md

@@ -176,7 +176,7 @@ A paper describing the engine is in preparation. Until then, please cite the rep
   author  = {Park, Sanggeun},
   title   = {miniworld-engine: GPU kernels for AlphaFold 3-style modules},
   year    = {2026},
-  version = {2.5.2},
+  version = {2.5.3},
   note    = {TODO: replace with the paper reference}
 }
 ```

{miniworld_engine-2.5.2 → miniworld_engine-2.5.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "miniworld-engine"
-version = "2.5.2"
+version = "2.5.3"
 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."
@@ -24,8 +24,8 @@ dependencies = [
   # 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.2; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"',
-  'miniworld-engine-kernels-sm80-cu129==2.5.2; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"',
+  'miniworld-engine-kernels-sm90-cu129==2.5.3; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"',
+  'miniworld-engine-kernels-sm80-cu129==2.5.3; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"',
 ]
 
 [project.optional-dependencies]

miniworld_engine-2.5.3/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"]

{miniworld_engine-2.5.2 → miniworld_engine-2.5.3}/src/miniworld_engine/kernels/_capture.py

@@ -216,6 +216,54 @@ _FORMS: dict = {}
 _HELD_BY_CAPTURES: list = []
 
 
+# ------------------------------------------------------------------------------------------------------------- the weight-cache registry
+#: name -> store: every module-level cache whose entries are derived from weights or per-item inputs (keyed on tensor addresses / versions /
+#: identities). A cache registers itself where it is defined (``_PACKS = _capture.weight_cache("family.packs")``), so
+#: ``clear_weight_caches`` finds it without a hand list; tests/cpu/test_capture_api.py fails on a module-level dict in the kernel sources
+#: that is neither registered nor listed there as not weight-derived.
+_WEIGHT_CACHES: dict[str, dict] = {}
+
+#: Per-instance cache attributes (``module.__dict__[name]``) holding weight-derived forms: dropped by ``clear_weight_caches(*modules)``.
+_INSTANCE_CACHE_ATTRS: set[str] = set()
+
+
+def weight_cache(name: str, store: dict | None = None) -> dict:
+    """Register ``store`` (a new empty dict by default) as the weight-derived cache ``name`` and return it. Registering a name again (a module
+    reload) replaces the earlier store."""
+    store = {} if store is None else store
+    _WEIGHT_CACHES[name] = store
+    return store
+
+
+def instance_cache_attr(name: str) -> str:
+    """Register ``name`` as a per-instance cache attribute (a weight-derived form kept in a module's ``__dict__``); returns ``name``."""
+    _INSTANCE_CACHE_ATTRS.add(name)
+    return name
+
+
+def weight_caches() -> dict[str, dict]:
+    """The registered weight-derived caches by name (of the kernel modules imported so far: a module not imported has filled none)."""
+    return dict(_WEIGHT_CACHES)
+
+
+def clear_weight_caches(*modules: torch.nn.Module) -> None:
+    """Drop every weight- or input-derived cache of the engine: each registered store (``weight_cache``: packs, hoisted tables, the parameter
+    forms, 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 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 (replaying one afterwards is undefined).
+    Forms already held for live captures (``capture_reads_eager``) stay held."""
+    for store in _WEIGHT_CACHES.values():
+        store.clear()
+    _PINNED.clear()
+    for module in modules:
+        for sub in module.modules():
+            for attr in _INSTANCE_CACHE_ATTRS:
+                sub.__dict__.pop(attr, None)
+
+
+weight_cache("capture.forms", _FORMS)
+
+
 def param_forms(tag, params, build, *, inputs_scope: bool = False, capture_reads_eager: bool = False):
     """``build()`` -- the kernel-dtype forms of some parameters (casts, contiguous copies, transposes, folds) -- cached while the very same
     parameter tensors are alive and unmodified, scoped to the capture like every weight pack (``lookup``): an eager call and a captured

{miniworld_engine-2.5.2 → miniworld_engine-2.5.3}/src/miniworld_engine/kernels/_prebuild.py

@@ -101,6 +101,7 @@ RECIPES: dict[str, str] = {
     "pair_embed_adaln": "miniworld_engine.kernels.pair_embed_adaln.cuda:prebuild",
     "pair_embed_geom": "miniworld_engine.kernels.pair_embed_geom.cuda:prebuild",
     "token_pair_init": "miniworld_engine.kernels.token_pair_init.cuda:prebuild",
+    "parallel_residual": "miniworld_engine.kernels.parallel_residual.cuda:prebuild",
 }
 
 

{miniworld_engine-2.5.2 → miniworld_engine-2.5.3}/src/miniworld_engine/kernels/_tf32.py

@@ -35,6 +35,8 @@ import torch
 import triton
 import triton.language as tl
 
+from miniworld_engine.kernels import _capture
+
 
 @contextlib.contextmanager
 def tf32_matmul(on: bool = True):
@@ -159,7 +161,7 @@ def tf32_grid_many(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
     return tuple(cast("list[torch.Tensor]", out))
 
 
-_CACHE: dict[tuple, tuple] = {}
+_CACHE: dict[tuple, tuple] = _capture.weight_cache("tf32.operands")
 _CACHE_MAX = 256
 
 

{miniworld_engine-2.5.2 → miniworld_engine-2.5.3}/src/miniworld_engine/kernels/augmented_attention/cuda/__init__.py

@@ -552,6 +552,21 @@ def prebuild(target, rows):
                 }
                 for opts in sorted(optsets):
                     calls.append(ft.partial(sm90_atom.kernel_cubin, kind, tf32, opts))
+        # AtomDiTESMFold2Pair (integrations/local_dit_esm): its inference rows on the sm90_atom row kernels built with RMS, the q / k RMSNorm +
+        # rotary (qknorm_local.cu) and the 192-key rank-window attention (lattn90_esm.cu; its pair-bias backward finishes on lbias.cu)
+        for tf32 in (False, True):
+            if need(["atom_rms_rows_fwd"], "f32" if tf32 else "bf16"):
+                for kind in ("cond", "cond_split", "pre", "pre_split", "post"):
+                    calls.append(ft.partial(sm90_atom.kernel_cubin, kind, tf32, ("RMS",)))
+        if need(["atom_rms_rows_fwd"]):
+            calls.append(ft.partial(sm90_atom.kernel_cubin, "pack", False, ()))
+        if need(["qknorm_rope_fwd", "qknorm_rope_bwd"]):
+            ext("qknorm_local")
+        if need(["local_esm_attn_fwd", "local_esm_pair_bias_fwd", "local_esm_attn_bwd_dq", "local_esm_attn_bwd_dkv"]):
+            calls.append(ft.partial(sm100_atom_local.esm_cubin, True))
+        if need(["local_esm_pair_bias_bwd"]):
+            calls.append(ft.partial(sm100_atom_local.esm_cubin, True))
+            calls.append(ft.partial(sm100_atom_local.cubin, "lbias", True))
         # sm_90a halves of the sm100_atom / sm100_atom_local sources
         for stem, _, _, _, *defs in sm100_atom._HOPPER_KERNELS.values():
             if need(_ATOM_STEMS[stem]):

{miniworld_engine-2.5.2 → miniworld_engine-2.5.3}/src/miniworld_engine/kernels/augmented_attention/cuda/apb/__init__.py

@@ -152,15 +152,17 @@ def pair_bias_bwd(
     d: int = 384,
     saved: tuple[torch.Tensor, torch.Tensor] | None = None,
     grid: int = 0,
+    out: torch.Tensor | None = None,
 ):
     """(dpair [L L, dp] bf16, dWf [H, dp] fp32, per-head dbias sums [H]) of ``pair_bias`` (without the mask: masked keys
     carry zero dbias). dWf and the sums are added into ``acc`` (``acc_n(H, d, dp)`` floats, zeroed by the caller) when given; the
     returned dWf / sums are views of it. A wide pair (dp > 128) takes ``saved`` = (row stats, unmasked bias) of its forward.
-    A100: ``grid`` CTAs of the persistent pass (0: the full occupancy)."""
+    A100: ``grid`` CTAs of the persistent pass (0: the full occupancy). ``out``: a contiguous [L L, dp] buffer dpair is written into
+    (a batch row of the caller's gradient; default a new tensor)."""
     heads, dp = wf.shape
     if acc is None:
         acc = torch.zeros(acc_n(heads, d, dp), device=pair2d.device)
-    dz = torch.empty_like(pair2d)
+    dz = torch.empty_like(pair2d) if out is None else out
     if dp == DP and grid > 0:
         ext().pair_bias_bwd_grid(pair2d, dbias.contiguous(), wf, dz, acc, int(L), float(eps), int(d), int(grid))
     elif dp == DP:

miniworld_engine-2.5.3/src/miniworld_engine/kernels/augmented_attention/cuda/qknorm_local.cu

@@ -0,0 +1,237 @@
+// qknorm_local.cu -- QK-norm of the AF3 atom block's q and k (H100 mma family, kernels/integrations/local_dit_h100.py): per head (4 heads x 32) RMSNorm
+// with one shared weight per tensor, eps in the denominator as torch's rms_norm.
+//
+//   forward   qo = q * rsqrt(mean(q^2) + eps) * wq,    ko likewise              (q, k [M, 128] bf16 or fp32, weights fp32 [32])
+//   backward  with x_hat = x rstd, g = dn w:   dx = rstd (g - x_hat mean(g x_hat)),   dw += dn x_hat
+//   rotary    (qknorm_rope_fwd, ESMFold2, inference) then qo = rot(qo), ko = rot(ko) with per-atom tables cos / sin [N, 16] (row r is atom r % N):
+//             channel c < 16 of a head pairs with c + 16, (a, b) -> (a cos - b sin, b cos + a sin) -- the partner is lane ^ 4.
+//
+// A warp per row, 8 lanes per head, 4 elements a lane: the mean of squares and the backward's mean are quad-of-octet shuffles, nothing
+// touches shared memory in the forward. The backward runs in place on the gradient slices of the block's [dg | dq | dk | dv] buffer
+// (row stride ``ld``) and leaves the weight gradients as per-block partial rows [blocks, 2, 32] that the caller sums in order: the
+// rows go to the blocks in a fixed pattern and each block adds its warps in a fixed order, so the sum is bitwise repeatable.
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+namespace {
+constexpr int DS = 128, HD = 32, WARPS = 8;
+
+template <typename T> __device__ __forceinline__ float4 load4(const T* p);
+template <> __device__ __forceinline__ float4 load4<float>(const float* p) { return *reinterpret_cast<const float4*>(p); }
+template <> __device__ __forceinline__ float4 load4<__nv_bfloat16>(const __nv_bfloat16* p) {
+  const uint2 u = *reinterpret_cast<const uint2*>(p);
+  const __nv_bfloat162 a = *reinterpret_cast<const __nv_bfloat162*>(&u.x), b = *reinterpret_cast<const __nv_bfloat162*>(&u.y);
+  return make_float4(__low2float(a), __high2float(a), __low2float(b), __high2float(b));
+}
+__device__ __forceinline__ float rna(float x) {                   // round to nearest tf32: the MMA operands of the fp32 path
+  unsigned r;
+  asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(r) : "f"(x));
+  return __uint_as_float(r);
+}
+template <typename T> __device__ __forceinline__ void store4(T* p, float4 v, bool rnd);
+template <> __device__ __forceinline__ void store4<float>(float* p, float4 v, bool rnd) {
+  if (rnd) v = make_float4(rna(v.x), rna(v.y), rna(v.z), rna(v.w));
+  *reinterpret_cast<float4*>(p) = v;
+}
+template <> __device__ __forceinline__ void store4<__nv_bfloat16>(__nv_bfloat16* p, float4 v, bool) {
+  const __nv_bfloat162 a = __floats2bfloat162_rn(v.x, v.y), b = __floats2bfloat162_rn(v.z, v.w);
+  uint2 u;
+  u.x = *reinterpret_cast<const unsigned*>(&a);
+  u.y = *reinterpret_cast<const unsigned*>(&b);
+  *reinterpret_cast<uint2*>(p) = u;
+}
+
+__device__ __forceinline__ float octet_sum(float v) {          // the 8 lanes of a head
+  v += __shfl_xor_sync(0xffffffffu, v, 1);
+  v += __shfl_xor_sync(0xffffffffu, v, 2);
+  return v + __shfl_xor_sync(0xffffffffu, v, 4);
+}
+
+template <typename T>
+__global__ void __launch_bounds__(WARPS * 32)
+qknorm_fwd_k(const T* __restrict__ q, const T* __restrict__ k, const float* __restrict__ wq, const float* __restrict__ wk,
+             T* __restrict__ qo, T* __restrict__ ko, long M, float eq, float ek, bool rnd, const float* __restrict__ cs = nullptr,
+             const float* __restrict__ sn = nullptr, long natoms = 1) {
+  const long row = (long)blockIdx.x * WARPS + threadIdx.x / 32;
+  if (row >= M) return;
+  const int lane = threadIdx.x % 32, col = lane * 4, d = (lane % 8) * 4;
+  const T* src[2] = {q + row * DS + col, k + row * DS + col};
+  T* dst[2] = {qo + row * DS + col, ko + row * DS + col};
+  const float* w[2] = {wq + d, wk + d};
+  const float eps[2] = {eq, ek};
+#pragma unroll
+  for (int t = 0; t < 2; ++t) {
+    const float4 x = load4<T>(src[t]);
+    const float rstd = rsqrtf(octet_sum(x.x * x.x + x.y * x.y + x.z * x.z + x.w * x.w) / HD + eps[t]);
+    const float4 wv = *reinterpret_cast<const float4*>(w[t]);
+    float4 y = make_float4(x.x * rstd * wv.x, x.y * rstd * wv.y, x.z * rstd * wv.z, x.w * rstd * wv.w);
+    if (cs != nullptr) {                                         // the rotary of the normalised values (fp32)
+      const int r = d % 16;
+      const float4 c = *reinterpret_cast<const float4*>(cs + (row % natoms) * 16 + r), s = *reinterpret_cast<const float4*>(sn + (row % natoms) * 16 + r);
+      const float4 p = make_float4(__shfl_xor_sync(0xffffffffu, y.x, 4), __shfl_xor_sync(0xffffffffu, y.y, 4),
+                                   __shfl_xor_sync(0xffffffffu, y.z, 4), __shfl_xor_sync(0xffffffffu, y.w, 4));
+      const float sg = d < 16 ? -1.f : 1.f;                     // first half: a cos - b sin; second: b cos + a sin
+      y = make_float4(y.x * c.x + sg * p.x * s.x, y.y * c.y + sg * p.y * s.y, y.z * c.z + sg * p.z * s.z, y.w * c.w + sg * p.w * s.w);
+    }
+    store4<T>(dst[t], y, rnd);
+  }
+}
+
+template <typename T>
+__global__ void __launch_bounds__(WARPS * 32)
+qknorm_bwd_k(T* __restrict__ dq, T* __restrict__ dk, long ld, const T* __restrict__ qr, const T* __restrict__ kr,
+             const float* __restrict__ wq, const float* __restrict__ wk, float* __restrict__ part, long M, float eq, float ek, bool rnd,
+             const float* __restrict__ cs = nullptr, const float* __restrict__ sn = nullptr, long natoms = 1) {
+  __shared__ float red[WARPS][2][4][8][4];                      // [warp][q/k][head][lane in head][dim]
+  const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, col = lane * 4, d = (lane % 8) * 4, head = lane / 8;
+  T* dg[2] = {dq, dk};
+  const T* raw[2] = {qr, kr};
+  const float* w[2] = {wq + d, wk + d};
+  const float eps[2] = {eq, ek};
+  float4 acc[2] = {make_float4(0, 0, 0, 0), make_float4(0, 0, 0, 0)};
+  for (long row = (long)blockIdx.x * WARPS + warp; row < M; row += (long)gridDim.x * WARPS) {
+#pragma unroll
+    for (int t = 0; t < 2; ++t) {
+      const float4 x = load4<T>(raw[t] + row * DS + col);
+      float4 dn = load4<T>(dg[t] + row * ld + col);
+      if (cs != nullptr) {   // the gradient of the rotated values back through the rotary: the transpose of the forward's rotation
+        const int r = d % 16;
+        const float4 c = *reinterpret_cast<const float4*>(cs + (row % natoms) * 16 + r), s = *reinterpret_cast<const float4*>(sn + (row % natoms) * 16 + r);
+        const float4 p = make_float4(__shfl_xor_sync(0xffffffffu, dn.x, 4), __shfl_xor_sync(0xffffffffu, dn.y, 4),
+                                     __shfl_xor_sync(0xffffffffu, dn.z, 4), __shfl_xor_sync(0xffffffffu, dn.w, 4));
+        const float sg = d < 16 ? 1.f : -1.f;
+        dn = make_float4(dn.x * c.x + sg * p.x * s.x, dn.y * c.y + sg * p.y * s.y, dn.z * c.z + sg * p.z * s.z, dn.w * c.w + sg * p.w * s.w);
+      }
+      const float rstd = rsqrtf(octet_sum(x.x * x.x + x.y * x.y + x.z * x.z + x.w * x.w) / HD + eps[t]);
+      const float4 xh = make_float4(x.x * rstd, x.y * rstd, x.z * rstd, x.w * rstd);
+      const float4 wv = *reinterpret_cast<const float4*>(w[t]);
+      const float4 g = make_float4(dn.x * wv.x, dn.y * wv.y, dn.z * wv.z, dn.w * wv.w);
+      const float mean = octet_sum(g.x * xh.x + g.y * xh.y + g.z * xh.z + g.w * xh.w) / HD;
+      store4<T>(dg[t] + row * ld + col, make_float4(rstd * (g.x - xh.x * mean), rstd * (g.y - xh.y * mean),
+                                                    rstd * (g.z - xh.z * mean), rstd * (g.w - xh.w * mean)), rnd);
+      acc[t].x += dn.x * xh.x; acc[t].y += dn.y * xh.y; acc[t].z += dn.z * xh.z; acc[t].w += dn.w * xh.w;
+    }
+  }
+#pragma unroll
+  for (int t = 0; t < 2; ++t) *reinterpret_cast<float4*>(&red[warp][t][head][lane % 8][0]) = acc[t];
+  __syncthreads();
+  if (threadIdx.x < 64) {                                       // (q/k, dim): heads and warps added in a fixed order
+    const int t = threadIdx.x / 32, dim = threadIdx.x % 32, j = dim / 4, e = dim % 4;
+    float s = 0.f;
+    for (int h = 0; h < 4; ++h)
+      for (int wp = 0; wp < WARPS; ++wp) s += red[wp][t][h][j][e];
+    part[((long)blockIdx.x * 2 + t) * HD + dim] = s;
+  }
+}
+
+void check(const at::Tensor& t, const char* what) {
+  TORCH_CHECK(t.is_cuda() && t.dim() == 2 && t.size(1) == DS && t.stride(1) == 1, what, ": [M, 128] rows, last dim contiguous");
+  TORCH_CHECK(t.scalar_type() == at::kFloat || t.scalar_type() == at::kBFloat16, what, ": fp32 or bf16");
+}
+void check_w(const at::Tensor& w, const char* what) {
+  TORCH_CHECK(w.is_cuda() && w.scalar_type() == at::kFloat && w.is_contiguous() && w.numel() == HD, what, ": fp32 [32]");
+}
+}  // namespace
+
+void qknorm_fwd(at::Tensor q, at::Tensor k, at::Tensor wq, at::Tensor wk, double eq, double ek, at::Tensor qo, at::Tensor ko) {
+  check(q, "q"); check(k, "k"); check(qo, "qo"); check(ko, "ko"); check_w(wq, "wq"); check_w(wk, "wk");
+  TORCH_CHECK(q.stride(0) == DS && k.stride(0) == DS && qo.stride(0) == DS && ko.stride(0) == DS, "contiguous rows");
+  TORCH_CHECK(q.scalar_type() == k.scalar_type() && q.scalar_type() == qo.scalar_type() && k.sizes() == q.sizes(), "operands match");
+  const long M = q.size(0);
+  const at::cuda::CUDAGuard guard(q.device());
+  const dim3 grid((unsigned)((M + WARPS - 1) / WARPS));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (q.scalar_type() == at::kFloat)
+    qknorm_fwd_k<float><<<grid, WARPS * 32, 0, stream>>>(q.data_ptr<float>(), k.data_ptr<float>(), wq.data_ptr<float>(), wk.data_ptr<float>(),
+                                                          qo.data_ptr<float>(), ko.data_ptr<float>(), M, (float)eq, (float)ek, true);
+  else
+    qknorm_fwd_k<__nv_bfloat16><<<grid, WARPS * 32, 0, stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(q.data_ptr()), reinterpret_cast<const __nv_bfloat16*>(k.data_ptr()), wq.data_ptr<float>(),
+        wk.data_ptr<float>(), reinterpret_cast<__nv_bfloat16*>(qo.data_ptr()), reinterpret_cast<__nv_bfloat16*>(ko.data_ptr()), M, (float)eq, (float)ek, false);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+// ESMFold2 (inference): the norm, then the rotary with the per-atom tables cos, sin [natoms, 16] fp32 (rows r of q / k are atom r % natoms).
+void qknorm_rope_fwd(at::Tensor q, at::Tensor k, at::Tensor wq, at::Tensor wk, double eq, double ek, at::Tensor cs, at::Tensor sn, at::Tensor qo,
+                     at::Tensor ko) {
+  check(q, "q"); check(k, "k"); check(qo, "qo"); check(ko, "ko"); check_w(wq, "wq"); check_w(wk, "wk");
+  TORCH_CHECK(q.stride(0) == DS && k.stride(0) == DS && qo.stride(0) == DS && ko.stride(0) == DS, "contiguous rows");
+  TORCH_CHECK(q.scalar_type() == k.scalar_type() && q.scalar_type() == qo.scalar_type() && k.sizes() == q.sizes(), "operands match");
+  TORCH_CHECK(cs.is_cuda() && sn.is_cuda() && cs.scalar_type() == at::kFloat && sn.scalar_type() == at::kFloat && cs.is_contiguous() &&
+              sn.is_contiguous() && cs.dim() == 2 && cs.size(1) == HD / 2 && sn.sizes() == cs.sizes(), "rotary tables: fp32 [N, 16]");
+  const long M = q.size(0), natoms = cs.size(0);
+  TORCH_CHECK(natoms > 0 && M % natoms == 0, "rows are samples x atoms");
+  const at::cuda::CUDAGuard guard(q.device());
+  const dim3 grid((unsigned)((M + WARPS - 1) / WARPS));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (q.scalar_type() == at::kFloat)
+    qknorm_fwd_k<float><<<grid, WARPS * 32, 0, stream>>>(q.data_ptr<float>(), k.data_ptr<float>(), wq.data_ptr<float>(), wk.data_ptr<float>(),
+                                                          qo.data_ptr<float>(), ko.data_ptr<float>(), M, (float)eq, (float)ek, true,
+                                                          cs.data_ptr<float>(), sn.data_ptr<float>(), natoms);
+  else
+    qknorm_fwd_k<__nv_bfloat16><<<grid, WARPS * 32, 0, stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(q.data_ptr()), reinterpret_cast<const __nv_bfloat16*>(k.data_ptr()), wq.data_ptr<float>(),
+        wk.data_ptr<float>(), reinterpret_cast<__nv_bfloat16*>(qo.data_ptr()), reinterpret_cast<__nv_bfloat16*>(ko.data_ptr()), M, (float)eq, (float)ek, false,
+        cs.data_ptr<float>(), sn.data_ptr<float>(), natoms);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+// dq, dk: [M, 128] views into the block's gradient buffer (row stride ld), rewritten in place from the gradient of the normalised
+// tensor to that of the raw one; qr, kr: the raw q, k. Returns the weight-gradient partials [blocks, 2, 32] fp32.
+at::Tensor qknorm_bwd(at::Tensor dq, at::Tensor dk, at::Tensor qr, at::Tensor kr, at::Tensor wq, at::Tensor wk, double eq, double ek) {
+  check(dq, "dq"); check(dk, "dk"); check(qr, "q raw"); check(kr, "k raw"); check_w(wq, "wq"); check_w(wk, "wk");
+  TORCH_CHECK(dq.stride(0) == dk.stride(0) && qr.stride(0) == DS && kr.stride(0) == DS, "row strides");
+  TORCH_CHECK(dq.scalar_type() == qr.scalar_type() && dk.scalar_type() == qr.scalar_type() && kr.scalar_type() == qr.scalar_type(), "operands match");
+  const long M = qr.size(0);
+  const at::cuda::CUDAGuard guard(qr.device());
+  const int blocks = (int)std::min<long>((M + WARPS - 1) / WARPS, 1056);          // a fixed function of M: the partials' order is too
+  auto part = at::empty({blocks, 2, HD}, qr.options().dtype(at::kFloat));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (qr.scalar_type() == at::kFloat)
+    qknorm_bwd_k<float><<<blocks, WARPS * 32, 0, stream>>>(dq.data_ptr<float>(), dk.data_ptr<float>(), dq.stride(0), qr.data_ptr<float>(),
+                                                            kr.data_ptr<float>(), wq.data_ptr<float>(), wk.data_ptr<float>(), part.data_ptr<float>(), M, (float)eq, (float)ek, true);
+  else
+    qknorm_bwd_k<__nv_bfloat16><<<blocks, WARPS * 32, 0, stream>>>(
+        reinterpret_cast<__nv_bfloat16*>(dq.data_ptr()), reinterpret_cast<__nv_bfloat16*>(dk.data_ptr()), dq.stride(0),
+        reinterpret_cast<const __nv_bfloat16*>(qr.data_ptr()), reinterpret_cast<const __nv_bfloat16*>(kr.data_ptr()), wq.data_ptr<float>(),
+        wk.data_ptr<float>(), part.data_ptr<float>(), M, (float)eq, (float)ek, false);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return part;
+}
+
+// ESMFold2 (training): :func:`qknorm_bwd` for q / k that were rotated after the norm; cs, sn the forward's tables [natoms, 16].
+at::Tensor qknorm_rope_bwd(at::Tensor dq, at::Tensor dk, at::Tensor qr, at::Tensor kr, at::Tensor wq, at::Tensor wk, double eq, double ek, at::Tensor cs,
+                           at::Tensor sn) {
+  check(dq, "dq"); check(dk, "dk"); check(qr, "q raw"); check(kr, "k raw"); check_w(wq, "wq"); check_w(wk, "wk");
+  TORCH_CHECK(dq.stride(0) == dk.stride(0) && qr.stride(0) == DS && kr.stride(0) == DS, "row strides");
+  TORCH_CHECK(dq.scalar_type() == qr.scalar_type() && dk.scalar_type() == qr.scalar_type() && kr.scalar_type() == qr.scalar_type(), "operands match");
+  TORCH_CHECK(cs.is_cuda() && sn.is_cuda() && cs.scalar_type() == at::kFloat && sn.scalar_type() == at::kFloat && cs.is_contiguous() &&
+              sn.is_contiguous() && cs.dim() == 2 && cs.size(1) == HD / 2 && sn.sizes() == cs.sizes(), "rotary tables: fp32 [N, 16]");
+  const long M = qr.size(0), natoms = cs.size(0);
+  TORCH_CHECK(natoms > 0 && M % natoms == 0, "rows are samples x atoms");
+  const at::cuda::CUDAGuard guard(qr.device());
+  const int blocks = (int)std::min<long>((M + WARPS - 1) / WARPS, 1056);          // a fixed function of M: the partials' order is too
+  auto part = at::empty({blocks, 2, HD}, qr.options().dtype(at::kFloat));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (qr.scalar_type() == at::kFloat)   // fp32 ESMFold2 training: the gradients stay fp32 (its GEMMs run at the caller's matmul precision)
+    qknorm_bwd_k<float><<<blocks, WARPS * 32, 0, stream>>>(dq.data_ptr<float>(), dk.data_ptr<float>(), dq.stride(0), qr.data_ptr<float>(),
+        kr.data_ptr<float>(), wq.data_ptr<float>(), wk.data_ptr<float>(), part.data_ptr<float>(), M, (float)eq, (float)ek, false, cs.data_ptr<float>(),
+        sn.data_ptr<float>(), natoms);
+  else
+    qknorm_bwd_k<__nv_bfloat16><<<blocks, WARPS * 32, 0, stream>>>(
+        reinterpret_cast<__nv_bfloat16*>(dq.data_ptr()), reinterpret_cast<__nv_bfloat16*>(dk.data_ptr()), dq.stride(0),
+        reinterpret_cast<const __nv_bfloat16*>(qr.data_ptr()), reinterpret_cast<const __nv_bfloat16*>(kr.data_ptr()), wq.data_ptr<float>(),
+        wk.data_ptr<float>(), part.data_ptr<float>(), M, (float)eq, (float)ek, false, cs.data_ptr<float>(), sn.data_ptr<float>(), natoms);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return part;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("qknorm_rope_bwd", &qknorm_rope_bwd);
+  m.def("qknorm_fwd", &qknorm_fwd);
+  m.def("qknorm_rope_fwd", &qknorm_rope_fwd);
+  m.def("qknorm_bwd", &qknorm_bwd);
+}

{miniworld_engine-2.5.2 → miniworld_engine-2.5.3}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/__init__.py

@@ -1009,3 +1009,144 @@ def pair_bias_bwd_mma(z, gamma, wb, dbias, eps=1e-5):
     )
     _load("pb_fin", d)((1, 1, 1), (64, 1, 1), part, int(nb), wbf, gf, dgamma, dwb, 0)
     return dz, dgamma, dwb
+
+
+# ------------------------------------------------------------------------------------------------ ESMFold2's attention (lattn90_esm.cu)
+# AtomDiTESMFold2Pair (kernels/integrations/local_dit_esm.py): the 192-key slid subset with the +/- 64 rank window, mma.sync bf16 and TF32,
+# forward (row LSE optional) and backward. sm_90a: the cubin of this package (PDL); sm_80: the same source with its dependent-launch wait
+# emptied (dit_variants/cuda/sm80.cubin), so there is one source of these kernels.
+ESM_KEYS, ESM_SPAN = 192, 160
+
+#: name -> (function, dynamic shared memory) of lattn90_esm.cu
+_ESM_KERNELS = {
+    "h_pb_esm": ("h_pair_bias_esm", 0),
+    "h_esm": ("h_attn_esm", 64 * 168 * 4 + 2 * (64 + 192 + 192) * 40 * 2),
+    "h_esm32": ("h_attn_esm_tf32", 64 * 168 * 4 + (64 + 192 + 192) * 36 * 4),
+    "h_esm_dq": ("h_attn_esm_dq", 64 * 168 * 4 + (64 * 2 + 192 * 2) * 40 * 2 + 2 * 64 * 4),
+    "h_esm_dkv": ("h_attn_esm_dkv", 4 * 16 * 164 * 4 + (192 * 2 + 64 * 2) * 40 * 2 + 2 * 192 * 4),
+    "h_esm_dq32": ("h_attn_esm_dq_tf32", 64 * 168 * 4 + (64 * 2 + 192 * 2) * 36 * 4 + 2 * 64 * 4),
+    "h_esm_dkv32": ("h_attn_esm_dkv_tf32", 4 * 16 * 164 * 4 + (192 * 2 + 64 * 2) * 36 * 4 + 2 * 192 * 4),
+    "h_pbb_esm": ("h_pair_bias_esm_bwd", 0),
+}
+
+
+def esm_cubin(hopper: bool) -> str:
+    """The cubin of lattn90_esm.cu: sm_90a (``hopper``) or sm_80."""
+    if hopper:
+        return cubin("lattn90_esm", True)
+    from miniworld_engine.kernels.dit_variants.cuda import sm80
+
+    return sm80.cubin("lattn90_esm")
+
+
+@functools.lru_cache(maxsize=None)
+def _esm_load(name: str, device_index: int):
+    from miniworld_engine.kernels.augmented_attention.cuda.sm100 import driver
+    from miniworld_engine.kernels.augmented_attention.cuda.sm100_atom import _hopper
+
+    hopper = _hopper(device_index)
+    func, smem = _ESM_KERNELS[name]
+    with torch.cuda.device(device_index):
+        return driver.Kernel(esm_cubin(hopper), func, smem, pdl=hopper)
+
+
+def _pb_fin(device_index: int):
+    """The pair-bias weight finish (``local_bias_fin`` of lbias.cu): this package's sm_90a build, or sm80_atom_local's."""
+    from miniworld_engine.kernels.augmented_attention.cuda.sm100_atom import _hopper
+
+    if _hopper(device_index):
+        return _load("pb_fin", device_index)
+    from miniworld_engine.kernels.augmented_attention.cuda import sm80_atom_local
+
+    return sm80_atom_local._load("pb_fin", device_index)
+
+
+def pair_bias_esm(z, gamma, wb, kmask, nreal, n, eps=1e-5):
+    """ESMFold2's trunked pair z [nwin, 32, 192, 16] bf16 / fp32 -> bias2 [4, nwin, 32, 160] fp32 by POSITION: entry (w, i, c) is the bias of
+    the key at 32 w - 64 + c, read from the slot of AlphaFold 3's 192-key slid subset that holds it (log2 domain), -inf outside the +/- 64
+    rank window, outside the real keys (``kmask`` [n] or None) or outside the subset. ``nreal``: the real atom count, a device int32 scalar."""
+    d = _dev(z)
+    out = torch.empty(NH, z.shape[0], WQ, ESM_SPAN, device=z.device, dtype=torch.float32)
+    _esm_load("h_pb_esm", d)(
+        (min(8 * num_sms(d), (out[0].numel() + 255) // 256), 1, 1),
+        (256, 1, 1),
+        z.contiguous(),
+        gamma.float().contiguous(),
+        wb.float().contiguous(),
+        _u8(kmask),
+        nreal,
+        out,
+        int(n),
+        float(eps),
+        int(z.dtype == torch.bfloat16),
+    )
+    return out
+
+
+def attn_esm(q, k, v, bias2, tf32: bool):
+    """ESMFold2's window attention forward (inference). q, k, v [A, N, 128] (bf16, or fp32 rounded to tf32; q / k normalised and rotated),
+    bias2 from ``pair_bias_esm`` -> O [A, N, 128] in q's dtype. N a multiple of 64."""
+    A, N, _ = q.shape
+    d = _dev(q)
+    o = torch.empty_like(q)
+    spc = _chunk(A, (N // 64) * NH)
+    grid = (N // 64, NH, (A + spc - 1) // spc)
+    if tf32:
+        _esm_load("h_esm32", d)(grid, (128, 1, 1), q, k, v, bias2, o, int(N), int(A), int(spc), None)
+    else:
+        _esm_load("h_esm", d)(grid, (128, 1, 1), q, k, v, bias2, o, None, int(N), int(A), int(spc))
+    return o
+
+
+def attn_esm_fwd_lse(q, k, v, bias2):
+    """The forward of :func:`attn_esm` for training (bf16, or fp32 on TF32 by q's dtype): O [A, N, 128] and the row LSE [A, 4, N] fp32 (log2
+    units; 1e30 for a row with no key)."""
+    A, N, _ = q.shape
+    d = _dev(q)
+    o = torch.empty_like(q)
+    lse = torch.empty(A, NH, N, device=q.device, dtype=torch.float32)
+    spc = _chunk(A, (N // 64) * NH)
+    grid = (N // 64, NH, (A + spc - 1) // spc)
+    if q.dtype == torch.float32:
+        _esm_load("h_esm32", d)(grid, (128, 1, 1), q, k, v, bias2, o, int(N), int(A), int(spc), lse)
+    else:
+        _esm_load("h_esm", d)(grid, (128, 1, 1), q, k, v, bias2, o, lse, int(N), int(A), int(spc))
+    return o, lse
+
+
+def attn_esm_bwd(q, k, v, do, bias2, lse, dd, dkv):
+    """Backward of the ESMFold2 window attention (bf16, or TF32 for fp32 operands): dQ, dK, dV into ``dkv`` [A N, >= 384] (row stride
+    ``dkv.stride(0)``, columns 0 / 128 / 256 relative to its base) and the bias gradient [4, nwin, 32, 160] fp32 (natural units, by position,
+    summed over the samples). ``lse`` from :func:`attn_esm_fwd_lse`, ``dd`` [A, 4, N] = rowsum(dO * O) per head."""
+    A, N, _ = q.shape
+    d = _dev(q)
+    nwin = nwindows(N)
+    ldd = dkv.stride(0)
+    spc = _chunk(A, (N // 64) * NH)
+    sp = (A + spc - 1) // spc
+    partial = torch.empty(sp, NH, nwin, WQ, ESM_SPAN, device=q.device, dtype=torch.float32)
+    grid = (N // 64, NH, sp)
+    sfx = "32" if q.dtype == torch.float32 else ""  # fp32 operands: the TF32 kernels, fp32 gradients into an fp32 dkv
+    _esm_load("h_esm_dq" + sfx, d)(grid, (128, 1, 1), q, k, v, do, bias2, lse, dd, dkv, partial, int(ldd), int(N), int(A), int(spc))
+    _esm_load("h_esm_dkv" + sfx, d)(grid, (128, 1, 1), q, k, v, do, bias2, lse, dd, dkv, int(ldd), int(N), int(A), int(spc))
+    return partial.sum(0)
+
+
+def pair_bias_esm_bwd(z, gamma, wb, dbias, nreal, eps=1e-5):
+    """dbias [4, nwin, 32, 160] fp32 (natural units) -> dz (z's shape / dtype), dgamma [16], dWb [4, 16] (fp32); ``nreal`` as in
+    :func:`pair_bias_esm`."""
+    d = _dev(z)
+    n = dbias.shape[1] * WQ
+    z = z.contiguous()
+    rows = z.numel() // DP
+    dz = torch.empty_like(z)
+    nb = min(num_sms(d), (rows + 255) // 256)
+    part = torch.empty(nb, NH * DP, device=z.device, dtype=torch.float32)
+    dgamma = torch.empty(DP, device=z.device, dtype=torch.float32)
+    dwb = torch.empty(NH, DP, device=z.device, dtype=torch.float32)
+    wbf, gf = wb.float().contiguous(), gamma.float().contiguous()
+    _esm_load("h_pbb_esm", d)(
+        (nb, 1, 1), (256, 1, 1), z, wbf, gf, dbias.contiguous(), nreal, dz, part, int(n), float(eps), int(z.dtype == torch.bfloat16)
+    )
+    _pb_fin(d)((1, 1, 1), (64, 1, 1), part, int(nb), wbf, gf, dgamma, dwb, 0)
+    return dz, dgamma, dwb