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

{miniworld_engine-2.5.4/src/miniworld_engine.egg-info → miniworld_engine-2.5.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: miniworld-engine
-Version: 2.5.4
+Version: 2.5.5
 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.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"
+Requires-Dist: miniworld-engine-kernels-sm90-cu129==2.5.5; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
+Requires-Dist: miniworld-engine-kernels-sm80-cu129==2.5.5; 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.4},
+  version = {2.5.5},
   note    = {TODO: replace with the paper reference}
 }
 ```

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/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.4},
+  version = {2.5.5},
   note    = {TODO: replace with the paper reference}
 }
 ```

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/pyproject.toml

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

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/_capture.py

@@ -189,6 +189,8 @@ def lookup_inputs(store: dict, key, build, limit: int = 64, valid=None, alive=No
         cid = _graph_capture_id()
         if cid == 0:
             return build()
+    if alive is not None:
+        _INPUT_STORES[id(store)] = (store, alive)
     sk = (cid, key)
     pinned = _pin_if_static_capture(store, sk)
     hit = store.get(sk)
@@ -206,6 +208,30 @@ def lookup_inputs(store: dict, key, build, limit: int = 64, valid=None, alive=No
     return out
 
 
+#: id(store) -> (the store, its ``alive(entry)``): every per-input cache ``lookup_inputs`` filled with an ``alive`` test, for
+#: ``release_dead_inputs``.
+_INPUT_STORES: dict[int, tuple[dict, object]] = {}
+
+
+def release_dead_inputs() -> int:
+    """Drop every per-input cache entry (``lookup_inputs`` / ``hoisted_table``: a hoisted pair bias, conditioning tables) whose input
+    tensor is gone; returns how many were dropped. ``lookup_inputs`` drops them only at its next call on the same store, so a table of a
+    finished item (a token DiT's hoisted pair bias: [heads x blocks, S, S], 11.3 GiB for OpenDDE's 3968 structural tokens) otherwise
+    outlives its item until the next item reaches the same module -- through the whole next trunk. A graph that read such an entry read
+    its input too, so it is no longer replayable either (as at ``lookup_inputs``'s own drop). Cheap (a weak-reference test per entry);
+    called by :class:`~miniworld_engine.modules.PairformerBlock` on every eager call, and by callers between items."""
+    dropped = 0
+    for store, alive in list(_INPUT_STORES.values()):
+        dead = [k for k, e in store.items() if not alive(e)]  # type: ignore[operator]
+        pinned = _PINNED.get(id(store))
+        for k in dead:
+            del store[k]
+            if pinned is not None:
+                pinned.discard(k)
+        dropped += len(dead)
+    return dropped
+
+
 #: (scope, (tag, ((pointer, version, dtype, shape), ...))) -> (weak references to the parameters, the forms). Also the Transition's ``Ws^T``
 #: (``transition/cuda/_weights.py``) and what FoldForge clears before it loads a model.
 _FORMS: dict = {}

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/aux_sm80.cuh

@@ -254,51 +254,101 @@ __global__ void __launch_bounds__(PB_NT) pair_bias_gen_kernel(const TI* __restri
 //   dW phase   thread (head group, channel slot) owns HB = H / 4 heads x CB = C / 64 channels of dW' and runs over the tile's elements and rows: per element one g read per head and one xh read per
 //              channel feed HB x CB FMAs; the block's partial [H, C] goes to pw (the caller sums the blocks in a fixed order).
 // C a multiple of 64 (64, 128, 256, 512: TE = 128, 64, 32, 16), H = 8, 12, 16 or 24 (a multiple of 4).
+// W' sits channel-major with a row pitch WP(H) of an odd number of float4s (H or H + 4): the dz phase's TPE threads of an element read the float4s of TPE consecutive
+// channels, which then fall in distinct bank quads (a pitch of an even number of float4s -- H = 8, 16 -- put them 2- / 4-way on the same banks).
+// 2026-11-12 (A100): the padded pitch and the next row's tile prefetched into registers during the current row's two phases (PbRaw8): 3.65 -> 1.72 ms at H 16, C 256,
+// the pair 768 x 768 (C 128: 1.07 -> 0.74; C 512 H 24, 256 x 256: 0.82 -> 0.59), the same values bitwise.
 template <int C_> struct PbBwdCfg {
   static constexpr int C = C_, TE = C_ == 64 ? 128 : C_ == 128 ? 64 : C_ == 256 ? 32 : 16, TPE = 256 / TE, P = C_ + 4, CB = C_ / 64;
   static_assert(C_ % 64 == 0 && C_ <= 512, "C: 64 .. 512, a multiple of 64");
-  static constexpr int smem(int H) { return (C_ * H + 2 * TE * P + TE * H) * 4; }
+  static constexpr int wp(int H) { return ((H / 4) & 1) ? H : H + 4; }
+  static constexpr int smem(int H) { return (C_ * wp(H) + 2 * TE * P + TE * H) * 4; }
+};
+// The raw 16-byte vectors of 8 channels (bf16: one uint4; fp32: two float4), held in registers between the prefetch and the shared-memory fill.
+template <typename T> struct PbRaw8;
+template <> struct PbRaw8<__nv_bfloat16> {
+  uint4 a;
+  __device__ __forceinline__ void load(const __nv_bfloat16* p) { a = __ldg(reinterpret_cast<const uint4*>(p)); }
+  __device__ __forceinline__ void zero() { a = make_uint4(0u, 0u, 0u, 0u); }
+  __device__ __forceinline__ void put(float* d) const {                  // the values pb_load8 gives, as two float4s
+    const uint32_t w[4] = {a.x, a.y, a.z, a.w};
+    float x[8];
+#pragma unroll
+    for (int k = 0; k < 4; ++k) { x[2 * k] = __uint_as_float(w[k] << 16); x[2 * k + 1] = __uint_as_float(w[k] & 0xffff0000u); }
+    reinterpret_cast<float4*>(d)[0] = make_float4(x[0], x[1], x[2], x[3]); reinterpret_cast<float4*>(d)[1] = make_float4(x[4], x[5], x[6], x[7]);
+  }
+};
+template <> struct PbRaw8<float> {
+  float4 a, b;
+  __device__ __forceinline__ void load(const float* p) { a = __ldg(reinterpret_cast<const float4*>(p)); b = __ldg(reinterpret_cast<const float4*>(p) + 1); }
+  __device__ __forceinline__ void zero() { a = b = make_float4(0.f, 0.f, 0.f, 0.f); }
+  __device__ __forceinline__ void put(float* d) const { reinterpret_cast<float4*>(d)[0] = a; reinterpret_cast<float4*>(d)[1] = b; }
 };
 
 template <int H, int C, typename T>
 __global__ void __launch_bounds__(256) pair_bias_gen_bwd_kernel(const T* __restrict__ z, const float* __restrict__ wcm, const float* __restrict__ db, T* __restrict__ dz, float* __restrict__ pw,
                                                                 int N, int NZ, int RPC, const unsigned char* __restrict__ kv, float eps) {
   using G = PbBwdCfg<C>;
-  constexpr int TE = G::TE, TPE = G::TPE, P = G::P, CB = G::CB, HB = H / 4;
+  constexpr int TE = G::TE, TPE = G::TPE, P = G::P, CB = G::CB, HB = H / 4, WP = G::wp(H);
+  constexpr int NT = (TE * (C / 8) + 255) / 256, NG = (TE * H + 255) / 256;  // z tasks (8 channels) / g values per thread and row
   static_assert(H % 4 == 0, "heads: a multiple of 4");
   extern __shared__ __align__(16) float bsm[];
-  float* ws = bsm;                                                       // W' channel-major [C][H]
-  float* xh = ws + C * H;                                                // [TE][P]: z, then xh
+  float* ws = bsm;                                                       // W' channel-major [C][WP]
+  float* xh = ws + C * WP;                                               // [TE][P]: z, then xh
   float* dx = xh + TE * P;                                               // [TE][P]: dxh, then dz
   float* gs = dx + TE * P;                                               // [TE][H]
   const int tid = threadIdx.x, e = tid / TPE, q = tid % TPE, hg = tid >> 6, slot = tid & 63;
-  for (int k = tid; k < C * H; k += 256) ws[k] = wcm[k];
+  for (int k = tid; k < C * H; k += 256) ws[(k / H) * WP + k % H] = wcm[k];
   const int j0 = blockIdx.x * TE;
   float acc[HB][CB];
 #pragma unroll
   for (int hb = 0; hb < HB; ++hb)
 #pragma unroll
     for (int cb = 0; cb < CB; ++cb) acc[hb][cb] = 0.f;
+  // Row i's z tile and dbias, fetched into registers one row ahead (raw: the conversions and the key mask wait for the shared-memory fill), so the loads of
+  // row i + 1 are in flight during row i's dz and dW phases.
+  PbRaw8<T> zr[NT];
+  float gr[NG];
+  unsigned char kr[NG];
+  auto fetch = [&](int i) {
+#pragma unroll
+    for (int t = 0; t < NT; ++t) {
+      const int idx = tid + 256 * t;
+      if (TE * (C / 8) % 256 != 0 && idx >= TE * (C / 8)) continue;
+      const int el = idx / (C / 8), c8 = (idx - el * (C / 8)) * 8;
+      if (j0 + el < NZ) zr[t].load(z + ((size_t)i * NZ + j0 + el) * C + c8);
+      else zr[t].zero();
+    }
+#pragma unroll
+    for (int t = 0; t < NG; ++t) {
+      const int idx = tid + 256 * t;
+      if (TE * H % 256 != 0 && idx >= TE * H) continue;
+      const int h = idx / TE, el = idx - h * TE, j = j0 + el;
+      gr[t] = j < NZ ? __ldg(db + ((size_t)h * N + i) * N + j) : 0.f;
+      kr[t] = j < NZ && (kv == nullptr || kv[j]);
+    }
+  };
+  if (blockIdx.y * RPC < NZ && RPC > 0) fetch(blockIdx.y * RPC);
   for (int r = 0; r < RPC; ++r) {
     const int i = blockIdx.y * RPC + r;
     if (i >= NZ) break;                                                  // uniform over the block
     __syncthreads();                                                     // the previous row's readers are done with xh / dx / gs
-    for (int idx = tid; idx < TE * (C / 8); idx += 256) {                // the z tile: 8 channels per task
-      const int el = idx / (C / 8), c8 = (idx - el * (C / 8)) * 8;
-      float v[8];
-      if (j0 + el < NZ) pb_load8(z + ((size_t)i * NZ + j0 + el) * C + c8, v);
-      else {
 #pragma unroll
-        for (int k = 0; k < 8; ++k) v[k] = 0.f;
-      }
-      float4* d = reinterpret_cast<float4*>(xh + el * P + c8);
-      d[0] = make_float4(v[0], v[1], v[2], v[3]); d[1] = make_float4(v[4], v[5], v[6], v[7]);
+    for (int t = 0; t < NT; ++t) {                                       // the z tile: 8 channels per task
+      const int idx = tid + 256 * t;
+      if (TE * (C / 8) % 256 != 0 && idx >= TE * (C / 8)) continue;
+      const int el = idx / (C / 8), c8 = (idx - el * (C / 8)) * 8;
+      zr[t].put(xh + el * P + c8);
     }
-    for (int idx = tid; idx < TE * H; idx += 256) {                      // g[el][h]: a masked key and a padded element have none
-      const int h = idx / TE, el = idx - h * TE, j = j0 + el;
-      gs[el * H + h] = (j < NZ && (kv == nullptr || kv[j])) ? __ldg(db + ((size_t)h * N + i) * N + j) : 0.f;
+#pragma unroll
+    for (int t = 0; t < NG; ++t) {                                       // g[el][h]: a masked key and a padded element have none
+      const int idx = tid + 256 * t;
+      if (TE * H % 256 != 0 && idx >= TE * H) continue;
+      const int h = idx / TE, el = idx - h * TE;
+      gs[el * H + h] = kr[t] ? gr[t] : 0.f;
     }
     __syncthreads();
+    if (r + 1 < RPC && i + 1 < NZ) fetch(i + 1);                         // the next row's loads overlap this row's phases
     // ---- dz phase: element e, channels q, q + TPE, ...
     float* xr = xh + e * P;
     float* dr = dx + e * P;
@@ -317,7 +367,7 @@ __global__ void __launch_bounds__(256) pair_bias_gen_bwd_kernel(const T* __restr
       const int c = q + TPE * k;
       const float xn = (xr[c] - mean) * rstd;
       xr[c] = xn;
-      const float4* wr = reinterpret_cast<const float4*>(ws + (size_t)c * H);
+      const float4* wr = reinterpret_cast<const float4*>(ws + (size_t)c * WP);
       float d = 0.f;
 #pragma unroll
       for (int hq = 0; hq < H / 4; ++hq) {

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/augmented_attention/cuda/sm80.py

@@ -762,7 +762,7 @@ def pair_bias_generic_bwd_supported(dp: int, heads: int, dtype: torch.dtype) ->
         te is not None
         and dtype in (torch.bfloat16, torch.float32)
         and heads in GEN_BWD_HEADS
-        and (dp * heads + 2 * te * (dp + 4) + te * heads) * 4 <= 160 * 1024
+        and (dp * (heads if (heads // 4) % 2 else heads + 4) + 2 * te * (dp + 4) + te * heads) * 4 <= 160 * 1024
     )
 
 

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/integrations/attention_pair_bias_b200.py

@@ -64,9 +64,9 @@ _cores: dict = {}
 EXACT_DWB = False
 
 
-def refusal(module, single: torch.Tensor, pair: torch.Tensor, mask: torch.Tensor | None, grad: bool) -> str | None:
+def refusal(module, single: torch.Tensor, pair: torch.Tensor, mask: torch.Tensor | None, grad: bool, padded: bool = False) -> str | None:
     """Why the H100 (sm_90) native path cannot take this call (``grad``: the training Function, else the inference op), or ``None``.
-    The ``cuda`` backend of ``AttentionPairBias`` states exactly this."""
+    The ``cuda`` backend of ``AttentionPairBias`` states exactly this. ``padded``: the caller pads L to a multiple of 128 (inference)."""
     if not single.is_cuda:
         return "CUDA tensors only"
     if torch.cuda.get_device_capability(single.device) != (9, 0):
@@ -98,8 +98,8 @@ def refusal(module, single: torch.Tensor, pair: torch.Tensor, mask: torch.Tensor
     if module.ln_single.weight is None or module.ln_pair.weight is None:
         return "the kernels take the affine LayerNorms"
     L = single.shape[1]
-    if L % 128:
-        return f"L {L}: a multiple of 128 only"
+    if L % 128 and not padded:
+        return f"L {L}: a multiple of 128 only (inference pads it: AttentionPairBias._forward_cuda)"
     if tuple(pair.shape) != (B, L, L, dp) or not pair.is_contiguous():
         return f"pair must be one contiguous [{B}, {L}, {L}, {dp}] tensor"
     if mask is not None and (tuple(mask.shape) != (B, L) or mask.dtype != torch.bool):

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/integrations/attention_pair_bias_sm80.py

@@ -140,14 +140,14 @@ def _sm80():
     return sm80
 
 
-def refusal(module, single: torch.Tensor, pair: torch.Tensor, mask: torch.Tensor | None, grad: bool) -> str | None:
+def refusal(module, single: torch.Tensor, pair: torch.Tensor, mask: torch.Tensor | None, grad: bool, padded: bool = False) -> str | None:
     """Why the A100 (sm_80) paths cannot take this call (``grad``: the training Function, else the inference op), or ``None``: the fused C++ step
     (``_step_refusal``), else the composed path (``_generic_refusal``: fp32 on the TF32 core, the wide pairs, ``gate_bias``). The ``cuda`` backend of
     ``AttentionPairBias`` states exactly this on an A100."""
     why = _step_refusal(module, single, pair, mask, grad)
     if why is None:
         return None
-    generic = _generic_refusal(module, single, pair, mask, grad)
+    generic = _generic_refusal(module, single, pair, mask, grad, padded)
     return None if generic is None else f"{why}; composed path: {generic}"
 
 
@@ -620,7 +620,7 @@ def _hdp(hd: int) -> int:
     return max(hd, 32)
 
 
-def _generic_refusal(module, single, pair, mask, grad: bool) -> str | None:
+def _generic_refusal(module, single, pair, mask, grad: bool, padded: bool = False) -> str | None:
     sm80 = _sm80()
     if module.use_qk_norm:
         return "per-head q / k RMSNorm"
@@ -635,7 +635,7 @@ def _generic_refusal(module, single, pair, mask, grad: bool) -> str | None:
     if single.ndim != 3 or single.shape[0] < 1:
         return "single must be [B, L, d_single] (B >= 1)"
     B, L, H = single.shape[0], single.shape[1], module.n_head
-    if L % 128:
+    if L % 128 and not padded:  # padded: the caller pads L (inference: AttentionPairBias._forward_cuda)
         return f"L {L}: a multiple of 128"
     hd = module.to_query.weight.shape[0] // H
     if hd not in GENERIC_HEAD_DIMS or module.to_query.weight.shape[0] != H * hd:
@@ -651,10 +651,11 @@ def _generic_refusal(module, single, pair, mask, grad: bool) -> str | None:
         return "the kernels take the affine LayerNorms"
     idx = single.device.index if single.device.index is not None else torch.cuda.current_device()
     hdp = _hdp(hd)
+    lc = -(-L // 128) * 128 if padded else L  # the core runs at the padded length
     if not (
-        sm80.supported_plain(torch.bfloat16, L, H * hdp, H, idx) if dt is torch.bfloat16 else sm80.supported_tf32(dt, L, H * hdp, H, idx)
+        sm80.supported_plain(torch.bfloat16, lc, H * hdp, H, idx) if dt is torch.bfloat16 else sm80.supported_tf32(dt, lc, H * hdp, H, idx)
     ):
-        return f"{H} x {hdp}, L {L}: outside the sm_80 attention core"
+        return f"{H} x {hdp}, L {lc}: outside the sm_80 attention core"
     if not _generic_loads(dt):
         return "the sm_80 attention extensions failed to build / load"
     return None
@@ -814,10 +815,13 @@ def _gfwd_row(single, pair, mask, params, heads, gate_bias, eps_s, eps_p, save,
         qkvg = torch.addmm(b, s, w.t())
         q, k, v, g = (qkvg[:, i * n : (i + 1) * n] for i in range(4))
         m = None if mask is None else mask.reshape(L)
+        nz = pair.shape[1]  # < L: the caller padded single / mask / extra to L (inference); the bias kernel masks the keys past nz
         bias = (
             _bias_split(pair[0], wcm, m, L, eps_p)
-            if f32 and pair.shape[-1] >= 512 and L <= 128 and heads % 8 == 0
-            else sm80.pair_bias_generic(pair[0], wcm, m, L, eps_p, oscale=1.0 if f32 else math.sqrt(hd), rows=_gen_rows(L))
+            if f32 and pair.shape[-1] >= 512 and L <= 128 and heads % 8 == 0 and nz == L
+            else sm80.pair_bias_generic(
+                pair[0], wcm, None if m is None else m[:nz], L, eps_p, oscale=1.0 if f32 else math.sqrt(hd), rows=_gen_rows(L)
+            )
         )
         if extra is not None:  # the bias in the core's units (bf16: x sqrt(head dim)); a masked key stays far below every real logit
             bias.add_(extra[0].to(bias.dtype), alpha=1.0 if f32 else math.sqrt(hd))

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/integrations/triattn_h100_sm90.py

@@ -195,7 +195,17 @@ def _rows_blocks(kind: int, pair: torch.Tensor, n_head: int) -> int:
     cap = INFERENCE_WORKSET_FRACTION * _total_memory(index)
     if _bytes_per_token(kind, C, n_head, False)[0] * L * L <= cap:
         return 1
-    return min(L, -(-_bytes_per_token(GENERIC, C, n_head, False)[0] * L * L // int(cap)))
+    blocks = -(-_bytes_per_token(GENERIC, C, n_head, False)[0] * L * L // int(cap))
+    if blocks >= 3:
+        # far past the cap (FoldForge OpenDDE's refiner: d_pair 384 at 2976 / 3968 tokens; the requirement sizes up to L2048 take 1 - 2
+        # blocks): blocks of at most ROWS_BLOCK_BYTES of q | k | v | g, so an allocator holding a fragmented model can place them (a 9 GiB
+        # block ran out of memory with 26 GiB reserved but free in the n = 16 cell's second fold)
+        blocks = max(blocks, -(-L * L * 4 * 32 * n_head * 2 // ROWS_BLOCK_BYTES))
+    return min(L, blocks)
+
+
+#: The largest q | k | v | g block of a row-blocked inference far past the working-set cap (``_rows_blocks``).
+ROWS_BLOCK_BYTES = 2 << 30
 
 
 @torch.autocast("cuda", enabled=False)  # the fp16 core takes the fp16 projections as made: an enclosing bf16 autocast turned the

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/integrations/triattn_sm80.py

@@ -228,7 +228,8 @@ def refusal(module, pair: torch.Tensor, mask: torch.Tensor | None, *, train: boo
     if not _serves_kind(module, pair, mask, train=train):
         return (
             f"L {pair.shape[1]}, d_pair {d_pair}, {module.n_head} heads, {'training' if train else 'inference'}: refused by the A100 kernels' "
-            "gates (L % 128 == 0, contiguous [B, L, L, d_pair], the training bias-gradient budget) or their extensions failed to build"
+            "gates (L % 128 == 0 for training and the fused d_pair 64 / 128 geometries, L > 128 for inference of the others, contiguous "
+            "[B, L, L, d_pair], the training bias-gradient budget) or their extensions failed to build"
         )
     return None
 
@@ -253,7 +254,10 @@ def forward(module, pair: torch.Tensor, mask: torch.Tensor | None) -> torch.Tens
     if kind is not None and kind[0] == "wide":
         from miniworld_engine.kernels.triangle_attention.cuda import sm80_wide
 
-        return sm80_wide.forward(
+        # L not a multiple of 128, or a whole-pair front output past sm80_wide.ROWS_BLOCKED_BYTES (FoldForge OpenDDE's refiner: 2976 / 3968
+        # tokens at d_pair 384): blocks of the attended frame's rows, at L padded to a multiple of 128
+        call = sm80_wide.forward_rows if sm80_wide.rows_blocked(pair, kind[1]) else sm80_wide.forward
+        return call(
             pair,
             kind[1],
             weights,

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/swa_dit/dispatch.py

@@ -111,6 +111,11 @@ _WGRAD_F32_WG_MAX_ROWS = 65536
 #: (bf16: ``triton/forward.py`` / ``backward.py``; fp32: ``forward_fp32.py`` / ``backward_fp32.py``) and the weight gradients cuBLAS.
 _PORTABLE = [False]
 
+#: The half window the Triton window kernels take for global attention under ``_PORTABLE`` (any is_global half window: negative or
+#: at least GLOBAL_HALF_WINDOW): every key of the row is then in reach, so the window kernels are dense attention over the row's valid keys
+#: (``j < seqused``) -- the same kernels, one constexpr specialisation. 2^30 keeps the kernels' int32 bounds (i0 +/- HW) in range.
+TRITON_GLOBAL_HW = 1 << 30
+
 
 @contextlib.contextmanager
 def _portable(on: bool):
@@ -183,11 +188,28 @@ def _composed_builds(device: torch.device, extra: str | None) -> bool:
 
 
 def _h100_global(q, half_window):
-    """Global attention on H100 always takes the native sm_90a stages (FlashAttention-4 between them): the composed path's
-    global attention is the sm_100a stages, which an H100 cannot run, whatever the per-shape defaults pick for the window."""
+    """Whether this global-attention call takes the native sm_90a stages (FlashAttention-4 between them) on H100: fp32 (and a bf16
+    call whose composed extensions do not build). A bf16 call takes the composed stages around FlashAttention-4
+    (:func:`_h100_global_composed`)."""
+    from miniworld_engine.kernels.swa_dit.cuda.loader import is_sm90
+
+    return not _PORTABLE[0] and is_global(half_window) and q.is_cuda and is_sm90(q.device) and not _h100_global_composed(q, half_window)
+
+
+def _h100_global_composed(q, half_window):
+    """bf16 global attention on H100: the composed bf16 stages (the window block's hand-CUDA qkvg / out-projection + FFN and their
+    backwards) with FlashAttention-4 in the window attention's place -- 2026-11-12: A 1 x B 48, 1024 atoms, master training 2.19 ms on
+    the native stages (six row passes and cuBLAS GEMMs per direction) against 1.997 for the "triton" backend."""
     from miniworld_engine.kernels.swa_dit.cuda.loader import is_sm90
 
-    return not _PORTABLE[0] and is_global(half_window) and q.is_cuda and is_sm90(q.device)
+    return (
+        not _PORTABLE[0]
+        and is_global(half_window)
+        and q.dtype == torch.bfloat16
+        and q.is_cuda
+        and is_sm90(q.device)
+        and _composed_cuda(q.device)
+    )
 
 
 def _native_b200_fusion(q=None, training=None, B=None):
@@ -709,6 +731,8 @@ def swa_dit_block_fwd(
 
 def _swa_dit_block_fwd(q, mod, cos, sin, seqused, wqkv, wg, wo, wu, wd, B, half_window, eps, save):
     """:func:`swa_dit_block_fwd`'s launches: the card's stages, or the Triton ones under ``_PORTABLE``."""
+    if _PORTABLE[0] and is_global(half_window):
+        half_window = TRITON_GLOBAL_HW  # global attention on the Triton window kernels (dense over the valid keys)
     ext32 = _fp32_global_cuda(q, half_window, save)
     if ext32 is not None:
         return _swa_dit_fwd_fp32_global(ext32, q, mod, cos, sin, seqused, wqkv, wg, wo, wu, wd, B, eps)
@@ -716,7 +740,8 @@ def _swa_dit_block_fwd(q, mod, cos, sin, seqused, wqkv, wg, wo, wu, wd, B, half_
         # The native stages take an fp32 modulation; a bf16 one comes from the composed bf16 hoist (global attention, which only the
         # native stages serve on H100, under a bf16 call that the composed stages otherwise take).
         return _native_h100_fwd(q, mod.float(), cos, sin, seqused, wqkv, wg, wo, wu, wd, B, half_window, eps, save)
-    if is_global(half_window):
+    fa4 = _h100_global_composed(q, half_window)  # the composed bf16 stages with FlashAttention-4 as the attention
+    if is_global(half_window) and not _PORTABLE[0] and not fa4:
         sm100 = _sm100(q, wd.shape[1], half_window, eps)
         if sm100 is None:
             raise RuntimeError(
@@ -785,7 +810,13 @@ def _swa_dit_block_fwd(q, mod, cos, sin, seqused, wqkv, wg, wo, wu, wd, B, half_
             SAVE=save,
         )
     ext = _cuda("attn", _attn_served(C, H, half_window), q.device)
-    if ext is not None:
+    if fa4:
+        from miniworld_engine.kernels.swa_dit.cuda.sm100 import attn_global_fwd
+
+        O = torch.zeros(M, C, device=q.device, dtype=torch.bfloat16)  # FA4 skips the rows past seqused: they must read zero
+        lse = torch.empty(N, H, S, device=q.device, dtype=torch.float32)
+        attn_global_fwd(Qh, Kh, Vh, seqused, O, lse)
+    elif ext is not None:
         O, lse = ext.attn_fwd(Qh, Kh, Vh, seqused, D**-0.5)
     else:
         O = torch.empty_like(qf)
@@ -1057,6 +1088,8 @@ def _swa_dit_block_bwd(
     dy, q, mod, cos, sin, seqused, wqkv, wg, wo, wu, wd, qh, kh, vh, g, o, lse, q1, x, pq, pk, att, y, ffn, B, half_window, eps
 ):
     """:func:`swa_dit_block_bwd`'s launches: the card's stages, or the Triton ones under ``_PORTABLE``."""
+    if _PORTABLE[0] and is_global(half_window):
+        half_window = TRITON_GLOBAL_HW  # as the forward
     if native_h100_stages(q) or _h100_global(q, half_window):
         grads = _native_h100_bwd(
             dy,
@@ -1089,7 +1122,7 @@ def _swa_dit_block_bwd(
         )
         grads[1] = grads[1].to(mod.dtype)  # dmod like mod (the contract), whatever the native stages computed it in
         return grads
-    if is_global(half_window):
+    if is_global(half_window) and not _PORTABLE[0] and not _h100_global_composed(q, half_window):
         sm100 = _sm100(q, wd.shape[1], half_window, eps)
         if sm100 is None:
             raise RuntimeError(
@@ -1487,7 +1520,19 @@ def _bwd_attention_side(
         if ext_q is not None:
             dP[:, 3 * C :] = dG
     heads: list[torch.Tensor] | None = None
-    if ext_a is not None:
+    if _h100_global_composed(q, half_window):
+        # global attention: FlashAttention-4's backward (its own softmax-delta pass; Dv unused). It writes dq | dk | dv straight into
+        # dP's column blocks 0 | 1 | 2 (strided [N, S, H, D] views) and leaves the rows past seqused alone: they are zeroed first.
+        from miniworld_engine.kernels.swa_dit.cuda.sm100 import attn_global_bwd
+
+        if ext_q is not None:
+            dP[:, : 3 * C].zero_()
+            blocks = [dP[:, i * C : (i + 1) * C].view(N, S, H, D).transpose(1, 2) for i in range(3)]
+            attn_global_bwd(qh, kh, vh, o, dO, lse, seqused, *blocks)
+        else:
+            heads = [torch.zeros(N, H, S, D, device=dev, dtype=bf) for _ in range(3)]
+            attn_global_bwd(qh, kh, vh, o, dO, lse, seqused, *heads)
+    elif ext_a is not None:
         if ext_q is not None:
             ext_a.attn_bwd(qh, kh, vh, dO, lse, Dv, seqused, dP, scale)
         else:

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/transition/cuda/__init__.py

@@ -43,9 +43,25 @@ def transition_sm80(x, ln_weight, ln_bias, wa, wb, ws, eps):
     limit = fused_fwd_sm80.max_rows(x.device, x.dtype)
     grad = torch.is_grad_enabled() and any(t is not None and t.requires_grad for t in (x, ln_weight, ln_bias, wa, wb, ws))
     if rows > limit and not grad:
-        # inference past one call's rows (the kernels' work lists / row grid): one call per contiguous row block (rows are independent), joined.
+        # inference past one call's rows (the kernels' work lists / row grid): one call per contiguous row block (rows are independent).
         # Training past it is refused by modules.Transition (the backward would need the same split of its weight-gradient sums).
         flat = x.reshape(-1, x.shape[-1]).contiguous()
+        d, h = x.shape[-1], wa.shape[0]
+        # balanced blocks of whole 256-row tiles (no short tail block), each written into its rows of one output: the wide chain's hand squeeze
+        # writes there directly (no join copy: at 2976^2 rows of D 384 the copy was 8 % of the call)
+        nblk = -(-rows // limit)
+        per = -(-rows // nblk)
+        step = min(-(-per // 256) * 256, limit)
+        starts = range(0, rows, step)
+        if (
+            x.dtype is torch.bfloat16
+            and fused_wide_sm80.route(flat[:step], wa, ws) == "wide"
+            and all(fused_wide_sm80.writes_into(min(step, rows - r), d, h) for r in starts)
+        ):
+            out = torch.empty_like(flat)
+            for r in starts:
+                fused_wide_sm80.forward_into(flat[r : r + step], ln_weight, ln_bias, wa, wb, ws, float(eps), out[r : r + step])
+            return out.reshape(x.shape)
         parts = [transition_sm80(flat[r : r + limit], ln_weight, ln_bias, wa, wb, ws, eps) for r in range(0, rows, limit)]
         if any(p is None for p in parts):
             return None

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/transition/cuda/fused_wide_sm80.py

@@ -263,6 +263,44 @@ def _fwd_launch(
     )
 
 
+def writes_into(m: int, d: int, h: int) -> bool:
+    """The inference forward of an ``m``-row bf16 call can write its output into a caller's buffer (``forward_into``): the chain whose squeeze is the hand
+    ``gemm_res`` (``_hand_squeeze``), not the one-kernel D 64 / 128 forward."""
+    return _hand_squeeze(m, d, True) and not fused_fwd_sm80.supports(d, h, m)
+
+
+def _fwd_into_fake(x, gamma, beta, wa, wb, ws, eps, out):
+    """Writes ``out`` only."""
+    return None
+
+
+@opaque(fake=_fwd_into_fake, name="transition_wide_fwd_into_sm80", mutates_args=("out",))
+def forward_into(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    beta: torch.Tensor,
+    wa: torch.Tensor,
+    wb: torch.Tensor,
+    ws: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    """``_fwd_launch``'s inference forward (``save`` off, ``ln`` on) of x [M, D] written into ``out`` (contiguous [M, D] bf16, a row block of a larger
+    output: ``transition_sm80`` joins the blocks of a split call without a copy).  The same launches as ``_fwd_launch`` where ``writes_into(M, D, H)``
+    (LayerNorm, dual GEMM + SwiGLU, the hand squeeze with the residual): the same output bits."""
+    if _is_fake(x, wa):
+        return
+    m, d = x.shape
+    ext = _ext()
+    gamma, beta = _f32(gamma, True), _f32(beta, True)
+    wa, wb, ws = _kernel_weights(wa, wb, ws, x.dtype, cache=True)
+    dual, res = _tiles(m, d, wa.shape[0])
+    xn, _ = ext.ln_fwd(x, gamma, beta, eps, False)
+    hid = ext.dual_swiglu(xn, wa, wb, dual)
+    del xn
+    ext.gemm_res(hid, ws, x, res, out)
+
+
 def _bwd_launch_fake(dy, x, xn, stats, gamma, wa, wb, ws, hid, ab, ln, gdt, wdt):
     """(dx like x, dgamma, dbeta [D] in ``gdt`` (empty without ``ln``), dW = [dWa; dWb] [2H, D], dWs [D, H] in ``wdt``)."""
     d = x.shape[1]

{miniworld_engine-2.5.4 → miniworld_engine-2.5.5}/src/miniworld_engine/kernels/transition/cuda/fused_wide_sm90a.py

@@ -534,6 +534,19 @@ class _WideTransitionSM90A(torch.autograd.Function):
         )
 
 
+#: Inference at D >= 384 (the two-kernel forward materializes the hidden h [M, n D]): a call whose h exceeds _H_ONE_LAUNCH_BYTES runs in
+#: launches of h <= _H_CHUNK_BYTES each. Requirement sizes (L <= 2048: h <= 12 GiB at D 384 n 4) keep one launch.
+_H_ONE_LAUNCH_BYTES = 16 << 30
+_H_CHUNK_BYTES = 2 << 30
+
+
+def _inference_chunk_rows(rows: int, d: int, hidden: int) -> int:
+    """Rows of one inference launch: all of them, or a multiple of 128 (whole tiles) keeping h within ``_H_CHUNK_BYTES``."""
+    if d < 384 or rows * hidden * 2 <= _H_ONE_LAUNCH_BYTES:
+        return rows
+    return max(128, _H_CHUNK_BYTES // (hidden * 2) // 128 * 128)
+
+
 def transition_wide_sm90a(x, gamma, beta, wa, wb, ws, eps):
     """Module-facing entry, same signature as ``fused_sm90a.transition_fused_sm90a``. Call ``supported()`` first: this
     raises rather than falling back. Inference (grad mode off, or nothing requiring grad) runs the forward without the
@@ -550,8 +563,8 @@ def transition_wide_sm90a(x, gamma, beta, wa, wb, ws, eps):
         else:
             wa, wb, ws = _inference_forms(wa, wb, ws, False)[:3]
             wst = None
-        out, _, _, _, _ = _fwd_launch(
-            x.reshape(-1, shape[-1]).contiguous(),
+        flat = x.reshape(-1, shape[-1]).contiguous()
+        args = (
             gamma.float().contiguous(),
             beta.float().contiguous(),
             wa.contiguous(),
@@ -561,6 +574,17 @@ def transition_wide_sm90a(x, gamma, beta, wa, wb, ws, eps):
             False,
             wst,
         )
+        rows = _inference_chunk_rows(flat.shape[0], shape[-1], wa.shape[0])
+        if rows >= flat.shape[0]:
+            return _fwd_launch(flat, *args)[0].reshape(shape)
+        # D >= 384 materializes the hidden h [M, n D] (bf16): at OpenDDE's structural pair (3968^2 tokens, D 384, n 2) that is one
+        # 22.5 GiB allocation next to the input and the output. Whole 128-row tiles in chunks of h <= _H_CHUNK_BYTES, one after
+        # another, each written into its rows of one output: every row is the same tile computation (bitwise the one-launch output).
+        out = torch.empty_like(flat)
+        for start in range(0, flat.shape[0], rows):
+            part = _fwd_launch(flat[start : start + rows], *args)[0]
+            out[start : start + rows] = part
+            del part
         return out.reshape(shape)
     out = _WideTransitionSM90A.apply(x, gamma, beta, wa, wb, ws, eps)
     from miniworld_engine.kernels.integrations.token_dit_train import _native_caller_output