miniworld-engine 2.3.2__tar.gz → 2.3.3__tar.gz

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: miniworld-engine
-Version: 2.3.2
+Version: 2.3.3
 Summary: MiniWorld GPU kernel development repo — Triton / CUDA kernels + benches for AF3-style ops (triangle multiplication, transition, attention, ...).
 License-Expression: MIT AND Apache-2.0 AND BSD-3-Clause
 Requires-Python: >=3.10
@@ -47,7 +47,7 @@ optimized in isolation: **hand-written CUDA where it exists, a Triton fallback o
 > blocks; terminal product repos assemble those into full models. Boundary rules live in
 > team-gm's `docs/ARCHITECTURE.md`.
 
-**Version 2.3.2** — CUDA + Triton only (CuTe/quack removed), torch 2.13 (CUDA 12.9 or 13.0), prebuilt H100 kernels,
+**Version 2.3.3** — CUDA + Triton only (CuTe/quack removed), torch 2.13 (CUDA 12.9 or 13.0), prebuilt H100 kernels,
 cuequivariance 0.12. [Changelog](docs/CHANGELOG.md) ·
 [attribution](licenses/THIRD_PARTY_NOTICES.md).
 

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/README.md

@@ -9,7 +9,7 @@ optimized in isolation: **hand-written CUDA where it exists, a Triton fallback o
 > blocks; terminal product repos assemble those into full models. Boundary rules live in
 > team-gm's `docs/ARCHITECTURE.md`.
 
-**Version 2.3.2** — CUDA + Triton only (CuTe/quack removed), torch 2.13 (CUDA 12.9 or 13.0), prebuilt H100 kernels,
+**Version 2.3.3** — CUDA + Triton only (CuTe/quack removed), torch 2.13 (CUDA 12.9 or 13.0), prebuilt H100 kernels,
 cuequivariance 0.12. [Changelog](docs/CHANGELOG.md) ·
 [attribution](licenses/THIRD_PARTY_NOTICES.md).
 

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "miniworld-engine"
-version = "2.3.2"
+version = "2.3.3"
 requires-python = ">=3.10"
 readme = "README.md"
 description = "MiniWorld GPU kernel development repo — Triton / CUDA kernels + benches for AF3-style ops (triangle multiplication, transition, attention, ...)."

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/src/miniworld_engine/autotune/devices.py

@@ -22,8 +22,19 @@ stays visible instead of dropping out of both sides and reading as 100% covered.
 from __future__ import annotations
 
 import csv
+import re
 from pathlib import Path
 
+#: A path into one account or one site's storage (tests/cpu/layout/test_no_machine_paths.py). A failure detail is an
+#: exception message, and a JIT failure names its build directory under the home of whoever ran it: the manifest is
+#: committed and shipped, so the path is written as `<path>`.
+_MACHINE_PATH = re.compile(r"/(?:home|users|Users|public_data|public_data02|scratch|lustre|gpfs)/[^\s,:;'\")]*")
+
+
+def _portable(detail: str) -> str:
+    return _MACHINE_PATH.sub("<path>", detail)
+
+
 _PKG = Path(__file__).resolve().parent.parent
 _REGISTRY = _PKG / "kernels" / "registry" / "registry.csv"
 #: Per-card manifests, beside the tuned caches in `autotune/data/`. They used to be resolved as
@@ -176,7 +187,7 @@ def record(gpu_key: str, results: dict[str, tuple[bool, str]],
         rows.append({
             "kernel": entry["kernel"], "backend": entry["backend"],
             "family": entry["family"], "file": entry["file"],
-            "status": status, "detail": detail, "dtype": dtype,
+            "status": status, "detail": _portable(detail), "dtype": dtype,
         })
     # Every OTHER precision's rows, carried through -- but only for precisions the registry still
     # declares. A row for a precision a kernel no longer runs at is evidence about a configuration

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/src/miniworld_engine/autotune/manifests/NVIDIA H100 80GB HBM3 (sm90).csv

@@ -1,5 +1,5 @@
 kernel,backend,family,file,status,detail,dtype
-#provenance,2.3.1,813f5c1080b0,clean,2026-10-20,archive/norm-cuda-20260930-335-g813f5c10,
+#provenance,2.3.2,230ef3679ec8,clean,2026-10-20,v2.3.1-5-g230ef367,
 adaln_bwd_dx_dlnw_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/training.py,ok,rel DCond=2.79e-03 DLnW=2.04e-03 (band 1e-02 declared),bf16
 adaln_bwd_pre_dx_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/training.py,ok,rel D=1.51e-03 DX=3.13e-03 (band 1e-02 declared),bf16
 adaln_epilogue_saveact_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/training.py,ok,rel Y=2.33e-03 Mean=2.23e-08 Rstd=1.86e-07 Gate=1.95e-03 (band 1e-02 declared),bf16
@@ -32,7 +32,7 @@ gated_projection_gate_dropres_triton,triton,trimul_inproj,miniworld_engine/kerne
 gated_projection_gate_flat_triton,triton,gated_projection,miniworld_engine/kernels/gated_projection/triton/main.py,ok,rel out=2.15e-03 (band 8e-03 declared),bf16
 gated_projection_gate_gemm_triton,triton,bias_only_attention,miniworld_engine/kernels/bias_only_attention/triton/gate_out.py,ok,rel out=3.79e-03 (band 1e-02 declared),bf16
 gated_projection_gate_triton,triton,gated_projection,miniworld_engine/kernels/gated_projection/triton/main.py,ok,rel out=2.15e-03 (band 8e-03 declared),bf16
-layernorm_bwd_atomic_strided_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/te_style.py,ok,rel dx=2.60e-03 dgamma=4.72e-07 dbeta=1.67e-07 (band 1e-02 declared),bf16
+layernorm_bwd_atomic_strided_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/te_style.py,ok,rel dx=2.60e-03 dgamma=4.41e-07 dbeta=2.50e-07 (band 1e-02 declared),bf16
 layernorm_bwd_atomic_triton,triton,layernorm,miniworld_engine/kernels/layernorm/triton/main.py,ok,rel dx=2.03e-03 dw=2.26e-03 db=2.41e-03 (band 1e-02 declared),bf16
 layernorm_bwd_foldstats_triton,triton,transition,miniworld_engine/kernels/transition/triton/fused.py,ok,rel dx=2.61e-03 dgamma=2.40e-07 dbeta=8.45e-08 (band 9e-03 declared),bf16
 layernorm_bwd_reduce_cuda,cuda,layernorm,miniworld_engine/kernels/layernorm/cuda/layer_norm_cuda_kernel.cu,ok,rel dw=2.26e-03 db=2.41e-03 (band 9e-03 declared),bf16
@@ -49,7 +49,7 @@ layernorm_fwd_strided_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/
 layernorm_linear_bwd_fp32_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/pair_bias.py,ok,rel dx=2.30e-03 dln_weight=1.74e-03 dproj_weight=2.01e-03 (band 1e-02 declared),bf16
 layernorm_linear_fwd_fp32_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/pair_bias.py,ok,rel out=2.37e-03 mean=2.08e-08 rstd=1.74e-07 (band 9e-03 declared),bf16
 layernorm_linear_fwd_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/fused.py,ok,rel out=3.51e-03 (band 2e-02 declared),bf16
-layernorm_stats_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/stats.py,ok,rel rstd=8.69e-08 c1=1.63e-07 (band 6e-07 declared),bf16
+layernorm_stats_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/stats.py,ok,rel rstd=8.69e-08 c1=8.17e-08 (band 6e-07 declared),bf16
 outer_product_mean_bwd_epilogue_dw_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel dw_out=4.55e-03 (band 1e-02 declared),bf16
 outer_product_mean_bwd_epilogue_dx_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel dmsa=5.84e-03 db_out=2.05e-03 dresidual=0.00e+00 (band 1e-02 declared),bf16
 outer_product_mean_bwd_layernorm_gemm_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel dmsa=5.84e-03 dln_weight=4.26e-03 dln_bias=3.54e-03 dw_left=4.97e-03 dw_right=6.01e-03 (band 1e-02 declared),bf16
@@ -69,12 +69,12 @@ rmsnorm_fwd_triton,triton,rmsnorm,miniworld_engine/kernels/rmsnorm/triton/main.p
 rope_fwd_triton,triton,rope,miniworld_engine/kernels/rope/triton/main.py,ok,rel y=5.14e-04 dx=7.23e-04 (band 3e-03 declared),bf16
 transition_bwd_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,failed,RuntimeError: Error building extension 'transition_cuda_ext_v2': [1/2] c++ -MMD -MF transition_cuda.o.d -DTORCH_EXTENSION_NAME=transit | c++: error: unrecognized command line option ‘-std=c++20’; did you mean ‘-std=c++2a’?,bf16
 transition_bwd_swiglu_recompute_triton,triton,transition,miniworld_engine/kernels/transition/triton/fused.py,ok,rel h=2.18e-03 dAB=2.12e-03 xn=7.02e-04 (band 1e-02 declared),bf16
-transition_cast_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,failed,ImportError: /home/psk6950/miniworld-engine-scratch/jit-runall2/miniworld_engine_jit/torch_extensions/transition_cuda_ext_v2-ee6c7a7d,bf16
+transition_cast_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,failed,ImportError: <path>,bf16
 transition_expand_swiglu_triton,triton,transition,miniworld_engine/kernels/transition/triton/main.py,ok,rel out=2.81e-03 (band 1e-02 declared),bf16
 transition_fwd_b2b_ktiled_triton,triton,transition,miniworld_engine/kernels/transition/triton/fused.py,ok,rel out=1.96e-03 (band 9e-03 declared),bf16
 transition_fwd_b2b_triton,triton,transition,miniworld_engine/kernels/transition/triton/fused.py,ok,rel out=2.48e-03 (band 1e-02 declared),bf16
 transition_layernorm_expand_swiglu_triton,triton,transition,miniworld_engine/kernels/transition/triton/fused.py,ok,rel out=2.18e-03 (band 1e-02 declared),bf16
-transition_swiglu_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,failed,ImportError: /home/psk6950/miniworld-engine-scratch/jit-runall2/miniworld_engine_jit/torch_extensions/transition_cuda_ext_v2-ee6c7a7d,bf16
+transition_swiglu_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,failed,ImportError: <path>,bf16
 triangle_attention_bwd_atomic_triton,triton,triangle_attention,miniworld_engine/kernels/triangle_attention/triton/atomic.py,ok,rel dq=1.54e-02 dk=1.29e-02 dv=4.85e-03 dbias=4.05e-03 (band 5e-02),bf16
 triangle_attention_bwd_dkdv_triton,triton,triangle_attention,miniworld_engine/kernels/triangle_attention/triton/main.py,ok,rel dk=4.44e-03 dv=2.76e-03 dbias=3.27e-03 (band 2e-02 declared),bf16
 triangle_attention_bwd_dq_triton,triton,triangle_attention,miniworld_engine/kernels/triangle_attention/triton/main.py,ok,rel dq=6.77e-03 (band 2e-02 declared),bf16
@@ -99,7 +99,7 @@ transition_fwd_b2b_sm90_cuda,cuda,transition,miniworld_engine/kernels/transition
 transition_expand_gate_sm90_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_expand_gate_kernel.cu,ok,rel out=1.94e-03 (band 8e-03 declared),bf16
 transition_bwd_gate_sm90_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_gatebwd_kernel.cu,ok,rel h=1.94e-03 dAB=2.97e-03 xn=5.27e-04 (band 1e-02 declared),bf16
 trimul_output_f567_train_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/output_fused.py,ok,rel y=4.27e-03 proj=2.95e-03 gate=2.72e-03 (band 1e-02 declared),bf16
-trimul_input_ln_residual_bwd_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/backward_fused.py,ok,rel dx=2.29e-03 dw=3.22e-07 db=8.17e-08 (band 1e-02 declared),bf16
+trimul_input_ln_residual_bwd_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/backward_fused.py,ok,rel dx=2.29e-03 dw=4.02e-07 db=8.17e-08 (band 1e-02 declared),bf16
 trimul_input_dual_bwd_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/backward_fused.py,ok,rel dx=2.34e-03 (band 1e-02 declared),bf16
 transition_squeeze_residual_triton,triton,transition,miniworld_engine/kernels/transition/triton/residual.py,ok,rel out=2.14e-03 (band 1e-02 declared),bf16
 mpnn_edge_tail_fwd_gemm_gather_saveact_triton,triton,mpnn_edge_tail,miniworld_engine/kernels/mpnn_edge_tail/triton/compute.py,ok,rel out=6.01e-03 (band 5e-02),bf16

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/src/miniworld_engine/build/audit.py

@@ -85,6 +85,11 @@ def import_all_kernels() -> list[tuple[str, str]]:
 
     failed = []
     for m in pkgutil.walk_packages(pkg.__path__, prefix=pkg.__name__ + "."):
+        # `build_*.py` are rebuild scripts run by hand (triangle_attention/cuda): importing one compiles its extension
+        # and overwrites the committed .so and *_manifest.json next to it (every CPU run with CUTLASS_PATH set did).
+        # Not edited to guard themselves: the runtime checks those files' hashes against the manifests.
+        if m.name.rpartition(".")[2].startswith("build_"):
+            continue
         try:
             importlib.import_module(m.name)
         except BaseException as exc:  # noqa: PERF203 -- a broken module is a finding, not a crash

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/src/miniworld_engine/build/bundle_coverage.py

@@ -9,7 +9,8 @@ the number of batched samples. This driver runs, on the current GPU:
   128 ... 2560, atom rungs 1024 ... 49152) for inference, the training crops (384, 768 tokens; the registry's atom
   ladder) for training, in each declared dtype and compute dtype, plus two batched samples on the pair stream;
 * the architecture blocks the model is built from (Pairformer with one- and two-directional TriMul, the token
-  DiffusionTransformer with its pair-bias / core-precision options and batched pairs, the SWA atom transformer in its
+  DiffusionTransformer with its pair-bias / core-precision options and batched pairs, the bias-only token DiT (one
+  pair shared by the samples or one per sample, masked, hoisted, d_cond 384 / 768), the SWA atom transformer in its
   three styles and the block-local atom DiT), inference and training,
 
 every one in bf16, fp32, bf16-mixed (fp32 parameters under bf16 autocast, MiniWorld's precision) and, for the blocks,
@@ -101,6 +102,17 @@ def cases() -> list[dict]:
                     for b in ((1, BATCHED) if length in (384, 768) else (1,)):
                         out.append({"kind": "diffusion", "L": length, "train": train, "dtype": prec, "options": opts,  # noqa: PERF401
                                     "samples": 48 if train else 5, "batch": b, "heavy": length >= HEAVY_TOKENS})
+                # The bias-only token DiT: one pair shared by the samples, or a pair of its own per sample (training
+                # block_psb); masked or not; trained stacks also from one hoisted pair_bias_all.
+                for d_cond in (384, 768):
+                    for pair in ("shared", "per_sample"):
+                        if pair == "per_sample" and length > TRAIN_TOKENS[-1]:
+                            continue
+                        for masked, hoist in ((False, False), (True, False), *(((True, True),) if train else ())):
+                            out.append({"kind": "bias_only_dit", "L": length, "train": train, "dtype": prec,  # noqa: PERF401
+                                        "d_cond": d_cond, "pair": pair, "masked": masked, "hoist": hoist,
+                                        "samples": 48 if train else 5,
+                                        "heavy": length >= HEAVY_TOKENS or (pair == "per_sample" and length >= 640)})
             atoms = (1024, 2048, 3072, 4096, 6144) if train else ATOM_RUNGS
             for n in atoms:
                 for variant in ("esmfold2", "af3", "af3_global"):
@@ -220,6 +232,31 @@ def _run_block(c: dict) -> None:
                 out.float().sum().backward()
         torch.cuda.synchronize()
         return
+    elif c["kind"] == "bias_only_dit":
+        from miniworld_engine.modules.bias_only_dit import BiasOnlyDiTBlock, pair_bias_all
+        from miniworld_engine.modules.exceptions import ImplementationType as MI
+
+        blocks = [BiasOnlyDiTBlock(d_cond=c["d_cond"], implementation=MI.MINIWORLD).cuda().to(pdt)
+                  for _ in range(2 if c["hoist"] else 1)]
+        for b in blocks:
+            _randomize(b)
+            b.train(c["train"])
+        a, n, per = c["samples"], c["L"], c["pair"] == "per_sample"
+        lead, nb = ((1, a), a) if per else ((a, 1), 1)  # per sample: [1, A, L, d] reading pair [A, L, L, 128]
+        x = torch.randn(*lead, n, 768, device="cuda", dtype=dt)
+        cond = torch.randn(*lead, n, c["d_cond"], device="cuda", dtype=dt)
+        pair = torch.randn(nb, n, n, 128, device="cuda", dtype=dt)
+        mask = torch.rand(nb, n, device="cuda") > 0.15 if c["masked"] else None
+        if c["train"]:
+            x, cond, pair = (t.requires_grad_(True) for t in (x, cond, pair))
+        with torch.set_grad_enabled(c["train"]), torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast):
+            biases = pair_bias_all(blocks, pair) if c["hoist"] else [None] * len(blocks)
+            for b, bias in zip(blocks, biases, strict=True):
+                x = b(x, cond, None if c["hoist"] else pair, mask, bias=bias)
+            if c["train"]:
+                x.float().sum().backward()
+        torch.cuda.synchronize()
+        return
     elif c["kind"] == "local_dit":
         from miniworld_engine.modules.exceptions import ImplementationType as MI
         from miniworld_engine.modules.local_dit import LocalDiTBlock

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/src/miniworld_engine/integrations/bias_only_dit.py

@@ -204,7 +204,8 @@ def _infer(
         else:
             P = runner.hoist(pair.contiguous(), mask)
             if pkey is not None:
-                runner._p_cache = (pkey, P)
+                # the entry keeps its key tensors alive: a freed address cannot alias another tensor
+                runner._p_cache = (pkey, P, (pair, mask))
         return runner.step(single, cond, P)
 
 

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/src/miniworld_engine/integrations/bias_only_dit_train.py

@@ -338,7 +338,7 @@ def _pack(P, dev):
     if W is not None:
         if slot is not None:
             _capture.prune(_PACKS)
-            _PACKS[slot] = (key, W)
+            _PACKS[slot] = (key, W, tuple(P.values()))  # keeps the key tensors alive
         return W
     g = lambda n: P[n]
     w1, w2 = (
@@ -407,7 +407,7 @@ def _pack(P, dev):
         W.update(mt)
     if slot is not None:
         _capture.prune(_PACKS)
-        _PACKS[slot] = (key, W)
+        _PACKS[slot] = (key, W, tuple(P.values()))  # keeps the key tensors alive
     return W
 
 
@@ -1588,7 +1588,7 @@ def _pack32(P, dev):
     }
     if slot is not None:
         _capture.prune(_PACKS32)
-        _PACKS32[slot] = (key, W)
+        _PACKS32[slot] = (key, W, tuple(P.values()))  # keeps the key tensors alive
     return W
 
 

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/src/miniworld_engine/integrations/local_dit.py

@@ -328,11 +328,18 @@ def _fp32_call(module, single, cond, pair) -> bool:
     return all(p.dtype == F32 for p in module.parameters())
 
 
+def _af3_variant(module) -> bool:
+    """AlphaFold 3's block: the K/V AdaLN beside the first one (``kv_norm="parallel"``) or the slid key window (``key_window="slide"``)."""
+    return getattr(module, "kv_norm", "chained") != "chained" or getattr(module, "key_window", "pad") != "pad"
+
+
 def serves(module, single, cond, pair, mask) -> bool:
     from miniworld_engine.integrations import local_dit_h100
 
     if local_dit_h100.matches(module, single, cond, pair, mask):
         return True
+    if _af3_variant(module):
+        return False  # only the H100 mma kernels know AlphaFold 3's K/V norm and key window; the others keep the composed path
     qualified = _serves_hopper_training(module, single, cond, pair, mask)
     if qualified is not None:
         return qualified
@@ -694,9 +701,13 @@ def _kv_pack(weights, packed_pre=None):
         weights[WEIGHTS.index("attention.to_key.weight")],
         weights[WEIGHTS.index("attention.to_value.weight")],
     ]
+    # the entry keeps the weights alive: a freed address cannot alias another tensor
     return _capture.lookup(
-        _KV, (hopper, *((t.data_ptr(), t._version) for t in tail)), build, limit=16
-    )
+        _KV,
+        (hopper, *((t.data_ptr(), t._version) for t in tail)),
+        lambda: (build(), tuple(tail)),
+        limit=16,
+    )[0]
 
 
 def _cond_tables(

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/src/miniworld_engine/integrations/local_dit_h100.py

@@ -39,16 +39,20 @@ def matches(module, single, cond, pair, mask) -> bool:
         torch.is_autocast_enabled("cuda") and single.dtype != BF
     ):
         return False
+    if _ld._af3_variant(module) and not _af3_servable(module, single, cond, pair, mask):
+        return False
     a, tr = module.attention, module.transition
-    if a.use_qk_norm or settings.current().engine_backend == "triton":
+    if settings.current().engine_backend == "triton":
         return False
     from miniworld_engine.modules.dispatch import KernelBackend
 
     if a._backend != KernelBackend.TRITON:
         return False
-    params = [module.get_parameter(n) for n in (*_ld.WEIGHTS, *_ld.CROSS_WEIGHTS)]
+    params = _weights_of(module)
     if any(p.device != single.device or not p.is_contiguous() for p in params):
         return False
+    if a.use_qk_norm and (a.norm_query.eps is not None or a.norm_key.eps is not None):
+        return False   # the kernels take torch's default eps (the float32 one for bf16 and fp32 alike)
     dt = (single.dtype, cond.dtype, pair.dtype)
     if dt == (F32, F32, F32):
         if (
@@ -112,6 +116,47 @@ def matches(module, single, cond, pair, mask) -> bool:
     return all(norm.eps == _ld.EPS for norm in norms)
 
 
+def _af3_servable(module, single, cond, pair, mask) -> bool:
+    """AlphaFold 3's K/V norm and slid key window are inference-only here (the backward kernels know neither). The slid window needs no
+    host read of the real atom count: it is a device tensor, so the block also runs inside a CUDA graph capture."""
+    params = [module.get_parameter(n) for n in (*_ld.WEIGHTS, *_ld.CROSS_WEIGHTS)]
+    return not (
+        torch.is_grad_enabled()
+        and any(t.requires_grad for t in (single, cond, pair, *params))
+    )
+
+
+def _slid_edges(o, q, k, v, z, gamma, wb, a, n, kmask, n0):
+    """Overwrite the attention output ``o`` [A, N, 128] of the windows whose slid key window is not the clipped one with the attention
+    over their slid key windows (torch, fp32). Those are windows 0 and 1 and the last three real windows (a window ends at
+    ``32 w + 80``, so at most three end past the last real atom): five windows chosen on the device from the real atom count, so there
+    is no host read and the call can be captured. A window chosen twice is computed twice with the same result."""
+    import torch.nn.functional as F
+
+    h, d = NH, DS // NH
+    nwin = n // 32
+    n_real = (
+        kmask.sum()
+        if kmask is not None
+        else torch.full((), n0, device=o.device, dtype=torch.int64)
+    )
+    last = (n_real - 1) // 32
+    idx = torch.stack([last * 0, last * 0 + 1, last - 2, last - 1, last]).clamp(0, nwin - 1)  # [5]
+    starts = torch.minimum((32 * idx - 48).clamp(min=0), (n_real - DS).clamp(min=0))
+    keys = starts[:, None] + torch.arange(DS, device=o.device)  # [5, 128]
+    valid = keys < n_real  # a system under 128 real atoms has keys past them
+    qs = q.view(a, nwin, 32, h, d)[:, idx].float()
+    ks = k.view(a, n, h, d)[:, keys].float()
+    vs = v.view(a, n, h, d)[:, keys].float()
+    zz = F.layer_norm(z[idx].float(), (DP,), gamma.float(), None, _ld.EPS)
+    bias = F.linear(zz, wb.float()).permute(0, 3, 1, 2)  # [5, h, 32, 128]
+    logits = torch.einsum("awqhd,awkhd->awhqk", qs, ks) * d**-0.5 + bias[None]
+    logits = logits.masked_fill(~valid[None, :, None, None, :], float("-inf"))
+    probs = torch.nan_to_num(torch.softmax(logits, dim=-1))
+    out = torch.einsum("awhqk,awkhd->awqhd", probs, vs).reshape(a, 5, 32, DS)
+    o.view(a, nwin, 32, DS).index_copy_(1, idx, out.to(o.dtype))
+
+
 # --------------------------------------------------------------------------------------------------------------------- the ops
 def _shared(cond) -> bool:
     """One conditioning for every sample: an expanded view (stride 0 over the samples) of [1, 1, N, 128]. The modulation tables are made
@@ -130,7 +175,15 @@ def _dims(single):
     return a, n0, n, a * n
 
 
-def _fwd_fake(single, cond, pair, kmask, weights, save):
+_NBASE = len(_ld.WEIGHTS) + len(_ld.CROSS_WEIGHTS)
+_QK_EPS = 1.1920928955078125e-07   # torch.finfo(torch.float32).eps: what the engine's RMSNorm uses for bf16 and fp32 inputs alike
+
+
+def _qk(weights) -> bool:
+    return len(weights) > _NBASE
+
+
+def _fwd_fake(single, cond, pair, kmask, weights, save, parallel_kv=False, slide=False):
     out = torch.empty_like(single)
     if not save:
         return [out]
@@ -154,6 +207,7 @@ def _fwd_fake(single, cond, pair, kmask, weights, save):
         t(m, DS),
         t(m, DS),
         t(m, DS),
+        *((t(m, DS), t(m, DS)) if _qk(weights) else ()),
     ]
 
 
@@ -165,8 +219,11 @@ def _fwd(
     kmask: torch.Tensor | None,
     weights: list[torch.Tensor],
     save: bool,
+    parallel_kv: bool = False,
+    slide: bool = False,
 ) -> list[torch.Tensor]:
-    """The block output [A, 1, N, 128] and, when ``save``, the tensors the backward reads (bias2, mod, mkv, q, k, v, sg, x1, xkv, o, lse, u, a2,
+    """``parallel_kv``: AlphaFold 3's K/V AdaLN of the raw atoms; ``slide``: its slid key window (inference only).
+    The block output [A, 1, N, 128] and, when ``save``, the tensors the backward reads (bias2, mod, mkv, q, k, v, sg, x1, xkv, o, lse, u, a2,
     x2, t): fresh tensors."""
     from miniworld_engine.kernels.augmented_attention.cuda import sm90_atom as k90
     from miniworld_engine.kernels.augmented_attention.cuda import (
@@ -176,6 +233,8 @@ def _fwd(
     tf32 = single.dtype == F32
     a, _n0, n, m = _dims(single)
     attn = local.attn_fwd_mma32 if tf32 else local.attn_fwd_mma
+    qkw = [w.float().contiguous() for w in weights[_NBASE:]]
+    weights = weights[:_NBASE]
     with torch.cuda.device(single.device):
         p = k90.pack(weights, tf32)
         s2 = _ld._padded(single, n).reshape(m, DS).contiguous()
@@ -187,7 +246,11 @@ def _fwd(
         ).contiguous()
         nmod = n if shared else m
         hoist = (not save) and settings.current().local_dit_hoist
-        wkey = tuple((w.data_ptr(), w._version) for w in weights) if hoist else None
+        wkey = (
+            (tuple((w.data_ptr(), w._version) for w in weights), parallel_kv)
+            if hoist
+            else None
+        )
 
         def tables():
             return k90.cond_fwd(c2, p, tf32)
@@ -208,10 +271,21 @@ def _fwd(
         if save:
             q, sg, k, v, x1, xkv = k90.pre_fwd(
                 s2, mod, mkv, p, tf32, save=True, nmod=nmod
-            )
+            )  # the training path is the chained K/V norm only (``matches``)
         else:
-            q, sg, k, v = k90.pre_fwd(s2, mod, mkv, p, tf32, nmod=nmod)
+            q, sg, k, v = k90.pre_fwd(
+                s2, mod, mkv, p, tf32, nmod=nmod, parallel_kv=parallel_kv
+            )
+        qraw = kraw = None
+        if qkw:   # the attention reads the normalised q / k; the backward keeps the raw ones
+            qraw, kraw = q, k
+            q, k = torch.empty_like(qraw), torch.empty_like(kraw)
+            _qk_kernels().qknorm_fwd(qraw, kraw, qkw[0], qkw[1], _QK_EPS, _QK_EPS, q, k)
+            if not save:
+                del qraw, kraw
         o, lse = attn(q.view(a, n, DS), k.view(a, n, DS), v.view(a, n, DS), bias2)
+        if slide:
+            _slid_edges(o, q, k, v, z, weights[_ld._GAMMA], weights[_ld._WBIAS], a, n, kmask, _n0)
         o = o.view(m, DS)
         if not save:
             return [
@@ -238,9 +312,16 @@ def _fwd(
         a2,
         x2,
         t,
+        *((qraw, kraw) if qkw else ()),
     ]
 
 
+def _qk_kernels():
+    from miniworld_engine.kernels.augmented_attention import cuda as sm90
+
+    return sm90._ext("qknorm_local")
+
+
 def _bwd_fake(dy, single, cond, pair, kmask, weights, saved):
     return [
         torch.empty_like(single),
@@ -266,7 +347,10 @@ def _bwd(
         sm100_atom_local as local,
     )
 
-    bias2, mod, mkv, q, k, v, sg, x1, xkv, o, lse, u, a2, x2, t = saved
+    bias2, mod, mkv, q, k, v, sg, x1, xkv, o, lse, u, a2, x2, t, *raw = saved
+    wq_dtype = weights[_NBASE].dtype if len(weights) > _NBASE else None
+    qkw = [w.float().contiguous() for w in weights[_NBASE:]]
+    weights = weights[:_NBASE]
     tf32 = single.dtype == F32
     a, _n0, n, m = _dims(single)
     dev = single.device
@@ -303,6 +387,9 @@ def _bwd(
             dp[:, DS:],
             tf32,
         )
+        if qkw:   # dq / dk leave the attention as gradients of the normalised tensors: back to the raw projections, in place
+            part = _qk_kernels().qknorm_bwd(dp[:, DS : 2 * DS], dp[:, 2 * DS : 3 * DS], raw[0], raw[1], qkw[0], qkw[1], _QK_EPS, _QK_EPS)
+            dqkw = part.sum(0)   # [2, 32]: the per-block partials in order
         dz, dgamma_pair, dwb = local.pair_bias_bwd_mma(z, p["gamma"], p["wbias"], dbias)
         dx, dmkv, cs_sc1, cs_bs, cs_bq = k90.pre_bwd(
             dp, x1, s2, mod, mkv, da2, pb, p, dmod, tf32, nmod=nmod, shared=sh
@@ -360,6 +447,8 @@ def _bwd(
             dev.index if dev.index is not None else torch.cuda.current_device(),
         )
     out = outs
+    if qkw:
+        out = [*outs, dqkw[0].to(wq_dtype).clone(), dqkw[1].to(wq_dtype).clone()]   # outputs of an opaque op may not alias
     dsingle = _ld._unpadded(dx.view(a, 1, n, DS), single.shape)
     if shared:  # the gradient of the one shared row set, split evenly over the expanded view's A copies (autograd sums them back)
         dcond = _ld._unpadded(
@@ -474,7 +563,10 @@ class _H100Block(torch.autograd.Function):
 
 
 def _weights_of(module):
-    return [module.get_parameter(n) for n in (*_ld.WEIGHTS, *_ld.CROSS_WEIGHTS)]
+    names = (*_ld.WEIGHTS, *_ld.CROSS_WEIGHTS)
+    if module.attention.use_qk_norm:
+        names = (*names, "attention.norm_query.weight", "attention.norm_key.weight")
+    return [module.get_parameter(n) for n in names]
 
 
 def run_block(module, single, cond, pair, mask=None):
@@ -486,4 +578,7 @@ def run_block(module, single, cond, pair, mask=None):
         t.requires_grad for t in (single, cond, pair, *weights)
     ):
         return _H100Block.apply(kmask, single, cond, pair, *weights)
-    return _fwd(single, cond, pair, kmask, weights, False)[0]
+    parallel = module.kv_norm == "parallel"
+    return _fwd(
+        single, cond, pair, kmask, weights, False, parallel, module.key_window == "slide"
+    )[0]

{miniworld_engine-2.3.2 → miniworld_engine-2.3.3}/src/miniworld_engine/integrations/token_dit.py

@@ -386,6 +386,7 @@ def _infer(
         assert (
             mask is not None
         )  # off sm_90 ``update`` fills a missing mask with all-True
+        pinned = (pair, mask)  # the cache entry below keeps the caller's tensors alive
         bkey = _capture.scoped(
             (
                 pair.data_ptr(),
@@ -410,7 +411,7 @@ def _infer(
                 )
             bias = runner.hoist(pair.contiguous(), mask)
             if bkey is not None:
-                runner._bias_cache = (bkey, bias)
+                runner._bias_cache = (bkey, bias, pinned)
         if L8 == L:
             return runner.step(single, cond, bias)
         # padded query rows: zero single / cond rows, computed and dropped (a shared cond stays one table)

miniworld_engine-2.3.3/src/miniworld_engine/integrations/token_dit_pair_bias.py

@@ -0,0 +1,89 @@
+"""The attention biases of a stack of token DiT blocks, from one per-sample pair, in one pass.
+
+Every block's bias is ``to_bias(LayerNorm(pair))`` with the block's own LayerNorm affine and projection, and only the affine and
+the projection differ from block to block: the LayerNorm statistics are the same. So the pair is normalised ONCE (a warp per
+row, ``pair_ln``), the blocks' ``to_bias.weight diag(ln_pair.weight) log2 e`` are stacked into one [blocks H, 128] matrix, and one
+GEMM leaves [blocks H, A L L] -- block ``b``'s bias is the b-th group of H rows, which is exactly the head-major per-sample
+layout ``token_dit_train.block_psb`` and the ``PSB`` attention cores read. Reading the pair once instead of once per block is the
+point: the pair is [A, L, L, 128] (A = 48 augmentations), 1.8 GB at L = 384 in bf16, and a per-block LayerNorm + projection read
+it twelve times each way.
+
+The backward concatenates the blocks' bias gradients, so ``dpair`` and the blocks' ``ln_pair`` / ``to_bias`` weight gradients are
+two GEMMs (dLN(pair) = dB^T Wf, dWf = dB LN(pair)) and the LayerNorm backward.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+
+from miniworld_engine.integrations.token_dit_train import BF, EPS, _mm32
+
+LOG2E = 1.0 / math.log(2.0)
+D, DH = 768, 48
+DP = 128
+
+
+def _rows():
+    from miniworld_engine.kernels.conditioned_transition.cuda.train import ext
+
+    return ext(D, DH)
+
+
+def _stacked(ln_weights, to_bias_weights):
+    """[blocks H, 128] fp32: row (b, h) = to_bias_b[h] ln_pair_b log2 e."""
+    return torch.cat([w.float() * l.float()[None, :] for l, w in zip(ln_weights, to_bias_weights, strict=True)]) * LOG2E
+
+
+class _PairBiases(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, pair, nblocks, *weights):  # noqa: D102
+        ln_weights, to_bias_weights = weights[:nblocks], weights[nblocks:]
+        A, L = pair.shape[0], pair.shape[1]
+        R, T = A * L * L, _rows()
+        heads = to_bias_weights[0].shape[0]
+        p2 = pair.reshape(R, DP)
+        ph = torch.empty(R, DP, device=pair.device, dtype=BF)
+        pst = torch.empty(R, 2, device=pair.device, dtype=torch.float32)
+        T.pair_ln(p2, ph, pst, EPS)
+        wf = _stacked(ln_weights, to_bias_weights).to(BF)
+        bias = torch.mm(wf, ph.t())                                        # [blocks H, R]
+        ctx.save_for_backward(pair, ph, pst, *weights)
+        ctx.meta = (nblocks, heads, A, L)
+        return tuple(bias[b * heads : (b + 1) * heads].view(heads, A, L, L) for b in range(nblocks))
+
+    @staticmethod
+    def backward(ctx, *grads):  # noqa: D102
+        pair, ph, pst, *weights = ctx.saved_tensors
+        nblocks, heads, A, L = ctx.meta
+        ln_weights, to_bias_weights = weights[:nblocks], weights[nblocks:]
+        R, T = A * L * L, _rows()
+        db = torch.empty(nblocks * heads, R, device=pair.device, dtype=BF)
+        for b, g in enumerate(grads):
+            rows = db[b * heads : (b + 1) * heads]
+            if g is None:
+                rows.zero_()
+            else:
+                rows.copy_(g.reshape(heads, R))
+        wf = _stacked(ln_weights, to_bias_weights).to(BF)
+        dph = _mm32(db.t(), wf)                                            # [R, 128] fp32
+        dwf = _mm32(db, ph)                                                # [blocks H, 128] fp32
+        dpair = torch.empty(R, DP, device=pair.device, dtype=pair.dtype)
+        T.pair_ln_bwd(dph, pair.reshape(R, DP), pst, dpair)
+        d_ln, d_w = [], []
+        for b in range(nblocks):
+            g = dwf[b * heads : (b + 1) * heads]
+            ln, w = ln_weights[b].float(), to_bias_weights[b].float()
+            d_w.append((g * ln[None, :] * LOG2E).to(to_bias_weights[b].dtype))
+            d_ln.append(((g * w).sum(0) * LOG2E).to(ln_weights[b].dtype))
+        return (dpair.view(pair.shape), None, *d_ln, *d_w)
+
+
+def pair_biases(pair, ln_weights, to_bias_weights):
+    """Every block's attention bias from the per-sample pair [A, L, L, 128]: a tuple of [H, A, L, L] bf16 tensors in log2 units
+    (``bias_b = (to_bias_b LN(pair) ln_pair_b) log2 e``, no key mask), differentiable in the pair and in each block's
+    ``ln_pair.weight`` / ``to_bias.weight``. The tuple's tensors are views of one [blocks H, A L L] buffer."""
+    return _PairBiases.apply(pair, len(ln_weights), *ln_weights, *to_bias_weights)
+
+
+__all__ = ["pair_biases"]