miniworld-engine 2.3.5__tar.gz → 2.3.6__tar.gz

{miniworld_engine-2.3.5 → miniworld_engine-2.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: miniworld-engine
-Version: 2.3.5
+Version: 2.3.6
 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.5** — CUDA + Triton only (CuTe/quack removed), torch 2.13 (CUDA 12.9 or 13.0), prebuilt H100 kernels,
+**Version 2.3.6** — 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.5 → miniworld_engine-2.3.6}/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.5** — CUDA + Triton only (CuTe/quack removed), torch 2.13 (CUDA 12.9 or 13.0), prebuilt H100 kernels,
+**Version 2.3.6** — 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.5 → miniworld_engine-2.3.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "miniworld-engine"
-version = "2.3.5"
+version = "2.3.6"
 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.5 → miniworld_engine-2.3.6}/src/miniworld_engine/autotune/manifests/NVIDIA H100 80GB HBM3 (sm90).csv

@@ -1,9 +1,9 @@
 kernel,backend,family,file,status,detail,dtype
-#provenance,2.3.4,ca7d86bb8dd4,clean,2026-10-21,v2.3.3-2-gca7d86bb,
-adaln_bwd_dx_dlnw_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/training.py,ok,rel DCond=2.15e-03 DLnW=2.62e-03 (band 1e-02 declared),bf16
-adaln_bwd_pre_dx_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/training.py,ok,rel D=1.49e-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.65e-03 Mean=2.23e-08 Rstd=1.86e-07 Gate=1.96e-03 (band 1e-02 declared),bf16
-adaln_epilogue_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/inference.py,ok,rel out=2.69e-03 (band 1e-02 declared),bf16
+#provenance,2.3.5,650b4231cfc7,clean,2026-10-21,v2.3.4-2-g650b4231,
+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
+adaln_epilogue_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/inference.py,ok,rel out=2.48e-03 (band 1e-02 declared),bf16
 adaln_fwd_gate_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/training.py,ok,rel Y=2.14e-03 Gate=1.95e-03 (band 8e-03 declared),bf16
 adaln_fwd_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/inference.py,ok,rel out=3.32e-03 (band 1e-02 declared),bf16
 adaln_gemm_gate_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/inference.py,ok,rel out=2.14e-03 (band 8e-03 declared),bf16
@@ -14,15 +14,15 @@ augmented_attention_fwd_triton,triton,augmented_attention,miniworld_engine/kerne
 bias_only_attention_bwd_pre_triton,triton,bias_only_attention,miniworld_engine/kernels/bias_only_attention/triton/main.py,ok,rel out=6.62e-08 (band 3e-07 declared),bf16
 bias_only_attention_bwd_triton,triton,bias_only_attention,miniworld_engine/kernels/bias_only_attention/triton/main.py,ok,rel dv=7.22e-03 dbias=5.84e-03 (band 5e-02 declared),bf16
 bias_only_attention_fwd_triton,triton,bias_only_attention,miniworld_engine/kernels/bias_only_attention/triton/main.py,ok,rel out=8.88e-03 m=8.11e-04 (band 4e-02 declared),bf16
-cond_transition_bwd_gemm_swiglu_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/training.py,ok,rel out=2.49e-03 (band 1e-02 declared),bf16
-cond_transition_bwd_swiglu_flat_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/training.py,ok,rel out=2.69e-03 (band 8e-03 declared),bf16
-cond_transition_expand_swiglu_saveact_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/fwd_saveact.py,ok,rel H=1.97e-03 AB=2.23e-03 (band 1e-02 declared),bf16
+cond_transition_bwd_gemm_swiglu_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/training.py,ok,rel out=2.98e-03 (band 1e-02 declared),bf16
+cond_transition_bwd_swiglu_flat_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/training.py,ok,rel out=2.37e-03 (band 8e-03 declared),bf16
+cond_transition_expand_swiglu_saveact_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/fwd_saveact.py,ok,rel H=1.97e-03 AB=2.17e-03 (band 1e-02 declared),bf16
 cond_transition_expand_swiglu_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/composed.py,ok,rel out=1.97e-03 (band 1e-02 declared),bf16
-cond_transition_fwd_b2b_saveact_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/training.py,ok,rel Y=1.88e-03 AB=2.23e-03 H=1.97e-03 Out=1.88e-03 Scale=2.06e-03 (band 3e-03 declared),bf16
-cond_transition_fwd_b2b_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/inference.py,ok,rel out=3.42e-03 (band 1e-02 declared),bf16
-cond_transition_squeeze_gate_saveact_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/fwd_saveact.py,ok,rel Y=3.43e-03 Out=3.15e-03 Scale=2.06e-03 (band 1e-02 declared),bf16
-cond_transition_squeeze_gate_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/composed.py,ok,rel out=3.43e-03 (band 8e-03 declared),bf16
-cond_transition_swiglu_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/training.py,ok,rel out=5.71e-03 (band 1e-02 declared),bf16
+cond_transition_fwd_b2b_saveact_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/training.py,ok,rel Y=1.77e-03 AB=2.17e-03 H=1.97e-03 Out=1.78e-03 Scale=2.06e-03 (band 3e-03 declared),bf16
+cond_transition_fwd_b2b_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/inference.py,ok,rel out=3.23e-03 (band 1e-02 declared),bf16
+cond_transition_squeeze_gate_saveact_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/fwd_saveact.py,ok,rel Y=3.22e-03 Out=2.96e-03 Scale=2.06e-03 (band 1e-02 declared),bf16
+cond_transition_squeeze_gate_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/composed.py,ok,rel out=3.22e-03 (band 8e-03 declared),bf16
+cond_transition_swiglu_triton,triton,conditioned_transition,miniworld_engine/kernels/conditioned_transition/triton/training.py,ok,rel out=3.98e-03 (band 1e-02 declared),bf16
 gated_projection_bwd_dx_triton,triton,bias_only_attention,miniworld_engine/kernels/bias_only_attention/triton/gate_out.py,ok,rel d_out_r=2.66e-03 d_gate=2.40e-03 gated=2.06e-03 (band 1e-02 declared),bf16
 gated_projection_bwd_gate_dropres_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/gate_elem.py,ok,rel d_proj=3.32e-03 d_glogit=2.39e-03 d_proj_drop=2.17e-03 d_glogit_drop=1.79e-03 d_proj_preact=3.57e-03 d_glogit_preact=2.40e-03 (band 1e-02 declared),bf16
 gated_projection_bwd_gate_flat_triton,triton,gated_projection,miniworld_engine/kernels/gated_projection/triton/main.py,ok,rel dg=2.69e-03 do=2.22e-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=3.46e-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.72e-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
@@ -45,11 +45,11 @@ layernorm_fwd_recompute_noaffine_triton,triton,layernorm_linear,miniworld_engine
 layernorm_fwd_recompute_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/recompute.py,ok,rel out=2.11e-03 (band 1e-02 declared),bf16
 layernorm_fwd_saveact_strided_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/te_style.py,ok,rel x_normed=2.10e-03 mean=4.71e-08 rstd=1.82e-07 (band 1e-02 declared),bf16
 layernorm_fwd_saveact_triton,triton,layernorm,miniworld_engine/kernels/layernorm/triton/main.py,ok,rel out=2.11e-03 (band 1e-02 declared),bf16
-layernorm_fwd_strided_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/ln_strided.py,ok,rel out=2.56e-03 (band 9e-03 declared),bf16
+layernorm_fwd_strided_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/ln_strided.py,ok,rel out=2.64e-03 (band 9e-03 declared),bf16
 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=8.17e-08 (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=1.63e-07 (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
@@ -67,7 +67,7 @@ rmsnorm_adamod_fwd_triton,triton,rmsnorm_adamod,miniworld_engine/kernels/rmsnorm
 rmsnorm_bwd_triton,triton,rmsnorm,miniworld_engine/kernels/rmsnorm/triton/main.py,ok,rel dx_aff=2.03e-03 dweight_aff=2.73e-03 dx_plain=2.54e-03 (band 4e-03 declared),bf16
 rmsnorm_fwd_triton,triton,rmsnorm,miniworld_engine/kernels/rmsnorm/triton/main.py,ok,rel y_aff=2.37e-03 y_plain=3.80e-03 (band 4e-03 declared),bf16
 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_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,failed,RuntimeError: Error building extension 'transition_cuda_ext_v2': [1/3] 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: <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

{miniworld_engine-2.3.5 → miniworld_engine-2.3.6}/src/miniworld_engine/build/bundle_coverage.py

@@ -11,7 +11,9 @@ the number of batched samples. This driver runs, on the current GPU:
 * 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 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,
+  three styles, the block-local atom DiT (AF3, and ESMFold2's variant for inference) and the template-stack TriMul (one
+  direction, hidden twice the pair width)),
+  inference and training,
 
 every one in bf16, fp32, bf16-mixed (fp32 parameters under bf16 autocast, MiniWorld's precision) and
 fp32 parameters fed bf16 activations (``PRECISIONS``).
@@ -122,6 +124,18 @@ def cases() -> list[dict]:
                                 "samples": 48 if train else 5, "heavy": n >= HEAVY_ATOMS})
                 out.append({"kind": "local_dit", "N": n, "train": train, "dtype": prec,
                             "samples": 48 if train else 5, "heavy": n >= HEAVY_ATOMS})
+                if not train:   # ESMFold2's atom block (RMS / SiLU AdaLN, raw gates, rotary q / k, 192-key +/- 64 window): inference
+                    out.append({"kind": "local_dit_esmfold2", "N": n, "train": False, "dtype": prec,  # noqa: PERF401
+                                "samples": 5, "heavy": n >= HEAVY_ATOMS})
+    # Template-stack TriMul: one direction, hidden twice the pair width (pair 64 / hidden 128, 128 / 256), B 1 and 4 templates.
+    for train in (False, True):
+        for prec in PRECISIONS:
+            for d in (64, 128):
+                for length in (TRAIN_TOKENS if train else TOKEN_RUNGS):
+                    for outgoing in (True, False):
+                        for b in (1, 4):
+                            out.append({"kind": "trimul_hidden2x", "L": length, "train": train, "dtype": prec,  # noqa: PERF401
+                                        "d": d, "outgoing": outgoing, "batch": b, "heavy": length >= HEAVY_TOKENS})
     for i, c in enumerate(out):
         c["id"] = i
     return out
@@ -234,6 +248,14 @@ def _run_block(c: dict) -> None:
                 out.float().sum().backward()
         torch.cuda.synchronize()
         return
+    elif c["kind"] == "trimul_hidden2x":
+        from miniworld_engine.modules.exceptions import ImplementationType as MI
+        from miniworld_engine.modules.triangle_multiplication import TriangleMultiplication
+
+        d, b, n = c["d"], c["batch"], c["L"]
+        model = TriangleMultiplication(d, d_hidden=2 * d, outgoing=c["outgoing"], implementation=MI.MINIWORLD,
+                                       p_drop=0.25 if c["train"] else 0.0)
+        args = (torch.randn(b, n, n, d, device="cuda", dtype=dt), torch.ones(b, n, device="cuda", dtype=torch.bool))
     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
@@ -259,6 +281,20 @@ def _run_block(c: dict) -> None:
                 x.float().sum().backward()
         torch.cuda.synchronize()
         return
+    elif c["kind"] == "local_dit_esmfold2":
+        from miniworld_engine.modules.exceptions import ImplementationType as MI
+        from miniworld_engine.modules.local_dit import ESMFOLD2, LocalDiTBlock
+
+        model = LocalDiTBlock(128, 128, 16, 4, n=2, implementation=MI.TRITON, **ESMFOLD2)
+        a, n = c["samples"], c["N"]
+        angles = torch.randn(1, n, 16, device="cuda")
+        mask = torch.ones(1, n, device="cuda", dtype=torch.bool)
+        args = (torch.randn(a, 1, n, 128, device="cuda", dtype=dt), torch.randn(a, 1, n, 128, device="cuda", dtype=dt),
+                torch.randn(1, (n + 31) // 32, 32, 192, 16, device="cuda", dtype=dt), mask)
+        model = model.cuda().to(pdt)
+        _randomize(model)
+        _step(model, args, False, autocast, {"rope": (angles.cos(), angles.sin())})
+        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.5 → miniworld_engine-2.3.6}/src/miniworld_engine/integrations/local_dit.py

@@ -329,14 +329,22 @@ def _fp32_call(module, single, cond, pair) -> bool:
 
 
 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"
+    """AlphaFold 3's block: the K/V AdaLN beside the first one (``kv_norm="parallel"``) or the slid key window (``key_window="slide"``);
+    or ESMFold2's options on top of it (192 keys, the rank window, RMS / SiLU conditioning, rotary)."""
+    return (getattr(module, "kv_norm", "chained") != "chained" or getattr(module, "key_window", "pad") != "pad"
+            or _esm_variant(module))
 
 
-def serves(module, single, cond, pair, mask) -> bool:
+def _esm_variant(module) -> bool:
+    """Any of ESMFold2's options (``modules.local_dit.ESMFOLD2``): only the H100 kernels know them."""
+    return (getattr(module, "keys", KEYS) != KEYS or getattr(module, "window_half", None) is not None
+            or getattr(module, "conditioning", "layernorm") != "layernorm" or getattr(module, "rope", False))
+
+
+def serves(module, single, cond, pair, mask, rope=None) -> bool:
     from miniworld_engine.integrations import local_dit_h100
 
-    if local_dit_h100.matches(module, single, cond, pair, mask):
+    if local_dit_h100.matches(module, single, cond, pair, mask, rope):
         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
@@ -1252,12 +1260,12 @@ def _infer(single, cond, pair, mask, weights):
     return _block_fwd(single, cond, pair, mask, list(weights), False)[0]
 
 
-def block(module, single, cond, pair, mask=None):
+def block(module, single, cond, pair, mask=None, rope=None):
     """The block's output stream (both residuals included), on the sm_100a kernels; ``serves()`` must have accepted the call."""
     from miniworld_engine.integrations import local_dit_h100
 
-    if local_dit_h100.matches(module, single, cond, pair, mask):
-        return local_dit_h100.run_block(module, single, cond, pair, mask)
+    if local_dit_h100.matches(module, single, cond, pair, mask, rope):
+        return local_dit_h100.run_block(module, single, cond, pair, mask, rope)
     weights = [
         module.get_parameter(name)
         for name in (*WEIGHTS, *(CROSS_WEIGHTS if module.cross_attention else ()))

{miniworld_engine-2.3.5 → miniworld_engine-2.3.6}/src/miniworld_engine/integrations/local_dit_h100.py

@@ -29,8 +29,10 @@ def _hopper(t) -> bool:
     return t.is_cuda and torch.cuda.get_device_capability(t.device) == (9, 0)
 
 
-def matches(module, single, cond, pair, mask) -> bool:
+def matches(module, single, cond, pair, mask, rope=None) -> bool:
     """H100, cross-attention AF3 atom block at the atom widths, fp32 inference (no gradient recorded)."""
+    if _ld._esm_variant(module):
+        return _esm_matches(module, single, cond, pair, mask, rope)
     if not (single.is_cuda and _hopper(single)):
         return False
     if not settings.current().local_dit_h100_mma:
@@ -569,8 +571,10 @@ def _weights_of(module):
     return [module.get_parameter(n) for n in names]
 
 
-def run_block(module, single, cond, pair, mask=None):
+def run_block(module, single, cond, pair, mask=None, rope=None):
     """The block's output stream on the H100 warp-level kernels; ``matches`` must have accepted the call."""
+    if _ld._esm_variant(module):
+        return _run_esm(module, single, cond, pair, mask, rope)
     weights = _weights_of(module)
     n0 = single.shape[2]
     kmask = _ld._keys(mask, n0, _ld._padded_length(n0))
@@ -582,3 +586,183 @@ def run_block(module, single, cond, pair, mask=None):
     return _fwd(
         single, cond, pair, kmask, weights, False, parallel, module.key_window == "slide"
     )[0]
+
+
+# --------------------------------------------------------------------------------------------------------------------- ESMFold2
+# ``LocalDiTBlock(**ESMFOLD2)`` for inference, on the same row kernels built with -DRMS=1 (RMSNorm AdaLNs, tables of silu(c), (1 + s) scales,
+# raw gates), the q / k RMSNorm + rotary of qknorm_local.cu, and the 192-key rank-window attention of lattn90_esm.cu:
+#
+#     cond90 (hoisted) -> pre90 -> qknorm_rope -> pair_bias_esm (hoisted) -> attn_esm -> post90
+#
+# The row kernels take the AF3 block's 27 parameters; ESMFold2 has no conditioning LayerNorm and no projection biases, so those slots hold
+# constants: the LN weights 1, the scale biases 1 (the identity-centred scale), the gate biases 0.
+ESM_KEYS = 192
+_ESM_NAMES = {
+    "sc1_w": "attention.ada_ln_in.to_scale.weight",
+    "bi1_w": "attention.ada_ln_in.to_bias.weight",
+    "wq": "attention.to_query.weight",
+    "bq": "attention.to_query.bias",
+    "wk": "attention.to_key.weight",
+    "wv": "attention.to_value.weight",
+    "ln_pair": "attention.ln_pair.weight",
+    "wb": "attention.to_bias.weight",
+    "wg": "attention.to_gate.weight",
+    "wo": "attention.to_out.weight",
+    "so_w": "attention.to_scale.weight",
+    "sc2_w": "transition.ada_ln_in.to_scale.weight",
+    "bi2_w": "transition.ada_ln_in.to_bias.weight",
+    "wa": "transition.expand_a.weight",
+    "wb2": "transition.expand_b.weight",
+    "wd": "transition.squeeze.weight",
+    "st_w": "transition.to_scale.weight",
+    "kv_sc_w": "attention.ada_ln_kv.to_scale.weight",
+    "kv_bi_w": "attention.ada_ln_kv.to_bias.weight",
+}
+_ESM_ONES = ("g1", "sc1_b", "g2", "sc2_b", "gkv", "kv_sc_b")
+_ESM_ZEROS = ("so_b", "st_b")
+_ESM_CONST: dict = {}
+
+
+def _esm_const(device, value: float):
+    """A persistent fp32 constant vector [128] per device (made eagerly and kept; inside a capture with none kept yet, a fresh one that is not kept)."""
+    key = (device, value)
+    c = _ESM_CONST.get(key)
+    if c is None:
+        c = torch.full((DS,), value, device=device, dtype=F32)
+        if not torch.cuda.is_current_stream_capturing():
+            _ESM_CONST[key] = c
+    return c
+
+
+def _esm_weights(module):
+    """The module's parameters in the row kernels' 27 slots (constants where ESMFold2 has none), then the q / k RMSNorm weights."""
+    dev = module.attention.to_query.weight.device
+    out = []
+    for name in _NAMES:
+        if name in _ESM_NAMES:
+            out.append(module.get_parameter(_ESM_NAMES[name]))
+        else:
+            out.append(_esm_const(dev, 1.0 if name in _ESM_ONES else 0.0))
+    return out, [module.attention.norm_query.weight, module.attention.norm_key.weight]
+
+
+def _esm_matches(module, single, cond, pair, mask, rope) -> bool:
+    """H100, exactly ESMFold2's option set (``modules.local_dit.ESMFOLD2``) at the atom widths, inference (no gradient recorded), bf16 or fp32
+    single / cond / pair, the rotary tables [1, N, 16]."""
+    from miniworld_engine.modules.adaptive_layernorm.module import RMS_EPS
+    from miniworld_engine.modules.dispatch import KernelBackend
+
+    if not (single.is_cuda and _hopper(single)) or not settings.current().local_dit_h100_mma:
+        return False
+    if settings.current().engine_backend == "triton" or module.attention._backend != KernelBackend.TRITON:
+        return False
+    a, tr = module.attention, module.transition
+    options = (module.cross_attention, module.kv_norm, module.key_window, module.keys, module.window_half, module.conditioning,
+               module.rope, a.use_qk_norm)
+    if options != (True, "parallel", "slide", ESM_KEYS, 64, "rms_silu", True, True):
+        return False
+    if a.norm_query.eps != RMS_EPS or a.norm_key.eps != RMS_EPS:
+        return False
+    params = [*(module.get_parameter(n) for n in _ESM_NAMES.values()), a.norm_query.weight, a.norm_key.weight]
+    if torch.is_grad_enabled() and any(t.requires_grad for t in (single, cond, pair, *params)):
+        return False   # inference only: the backward kernels know neither the conditioning nor the window
+    if any(p.device != single.device or not p.is_contiguous() for p in params):
+        return False
+    if torch.is_autocast_enabled("cuda") and single.dtype != BF:
+        return False
+    dt = (single.dtype, cond.dtype, pair.dtype)
+    if dt == (F32, F32, F32):
+        if any(p.dtype != F32 for p in params) or not settings.current().local_dit_h100_mma_fp32:
+            return False
+    elif dt == (BF, BF, BF):
+        if any(p.dtype not in (F32, BF) for p in params) or not settings.current().local_dit_h100_mma_bf16:
+            return False
+    else:
+        return False
+    if single.ndim != 4 or single.shape[1] != 1 or single.shape[3] != DS or single.shape[2] == 0:
+        return False
+    n = single.shape[2]
+    if tuple(cond.shape) != (*single.shape[:3], DC) or tuple(pair.shape) != (1, (n + 31) // 32, 32, ESM_KEYS, DP):
+        return False
+    if not (single.is_contiguous() and pair.is_contiguous() and (cond.is_contiguous() or _shared(cond))):
+        return False
+    if mask is not None and (mask.dtype != torch.bool or tuple(mask.shape) != (1, n) or mask.device != single.device):
+        return False
+    if rope is None or len(rope) != 2 or any(
+        t.device != single.device or tuple(t.shape) != (1, n, DS // NH // 2) or not t.is_floating_point() for t in rope
+    ):
+        return False
+    if (a.n_head, a.to_query.weight.shape[0], a.to_bias.weight.shape[1], tr.expand_a.weight.shape[0]) != (NH, DS, DP, 2 * DS):
+        return False
+    return a.ln_pair.eps == _ld.EPS
+
+
+def _esm_fwd_fake(single, cond, pair, kmask, weights, qkw, cos, sin):
+    return torch.empty_like(single)
+
+
+@opaque(fake=_esm_fwd_fake, name="local_dit_h100_esm_fwd")
+def _esm_fwd(
+    single: torch.Tensor,
+    cond: torch.Tensor,
+    pair: torch.Tensor,
+    kmask: torch.Tensor | None,
+    weights: list[torch.Tensor],
+    qkw: list[torch.Tensor],
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    """ESMFold2's block output [A, 1, N, 128] (inference): the RMS row kernels, the q / k RMSNorm + rotary and the rank-window attention.
+    ``cos`` / ``sin``: the rotary tables, fp32 [n, 16] at the padded length."""
+    from miniworld_engine.kernels.augmented_attention.cuda import sm90_atom as k90
+    from miniworld_engine.kernels.augmented_attention.cuda import (
+        sm100_atom_local as local,
+    )
+    from miniworld_engine.modules.adaptive_layernorm.module import RMS_EPS
+
+    tf32 = single.dtype == F32
+    a, n0, n, m = _dims(single)
+    with torch.cuda.device(single.device):
+        p = k90.pack(weights, tf32)
+        s2 = _ld._padded(single, n).reshape(m, DS).contiguous()
+        shared = _shared(cond)
+        c2 = (_ld._padded(cond[:1], n).reshape(n, DC) if shared else _ld._padded(cond, n).reshape(m, DC)).contiguous()
+        nmod = n if shared else m
+        hoist = settings.current().local_dit_hoist
+        wkey = (tuple((w.data_ptr(), w._version) for w in weights), "esm") if hoist else None
+
+        def tables():
+            return k90.cond_fwd(c2, p, tf32, rms=True)
+
+        mod, mkv = _ld._hoisted(_HOIST_COND, cond, (wkey, n), tables) if hoist else tables()
+        z = _ld._trunk(pair, n).contiguous()
+
+        def pair_bias():
+            nreal = (kmask.sum(dtype=torch.int32) if kmask is not None
+                     else torch.full((), n0, device=single.device, dtype=torch.int32))
+            return local.pair_bias_esm(z, p["gamma"], p["wbias"], kmask, nreal, n)
+
+        if hoist:
+            mk = None if kmask is None else _ld._tkey(kmask)
+            bias2 = _ld._hoisted(_HOIST_BIAS, pair, (wkey, n, mk), pair_bias)
+        else:
+            bias2 = pair_bias()
+        qraw, sg, kraw, v = k90.pre_fwd(s2, mod, mkv, p, tf32, eps=RMS_EPS, nmod=nmod, parallel_kv=True, rms=True)
+        q, k = torch.empty_like(qraw), torch.empty_like(kraw)
+        _qk_kernels().qknorm_rope_fwd(qraw, kraw, qkw[0].float().contiguous(), qkw[1].float().contiguous(), RMS_EPS, RMS_EPS, cos, sin, q, k)
+        del qraw, kraw
+        o = local.attn_esm(q.view(a, n, DS), k.view(a, n, DS), v.view(a, n, DS), bias2, tf32).view(m, DS)
+        out = k90.post_fwd(s2, o, sg, mod, p, tf32, eps=RMS_EPS, nmod=nmod, rms=True)
+    return _ld._unpadded(out.view(a, 1, n, DS), single.shape)
+
+
+def _run_esm(module, single, cond, pair, mask, rope):
+    weights, qkw = _esm_weights(module)
+    n0 = single.shape[2]
+    n = _ld._padded_length(n0)
+    kmask = _ld._keys(mask, n0, n)
+    # fp32 [n, 16] at the padded length (no copy when the caller's tables already are)
+    cos, sin = (
+        (t[0].float() if n == n0 else torch.nn.functional.pad(t[0].float(), (0, 0, 0, n - n0))).contiguous() for t in rope
+    )
+    return _esm_fwd(single, cond, pair, kmask, weights, qkw, cos, sin)

{miniworld_engine-2.3.5 → miniworld_engine-2.3.6}/src/miniworld_engine/integrations/trimul_h100.py

@@ -1,4 +1,5 @@
-"""Automatic dispatch to the packaged, selected bidirectional H100 training path."""
+"""Automatic dispatch to the packaged, selected bidirectional H100 training path. The hidden width equals the pair width, or is
+twice it in one direction at D64 / D128 (``_hidden_twice``: AF3 / Protenix template blocks, pair 64, hidden 128)."""
 
 from __future__ import annotations
 
@@ -74,6 +75,21 @@ def _total_memory(index: int) -> int:
     return torch.cuda.get_device_properties(index).total_memory
 
 
+def _hidden_twice(module, pair) -> bool:
+    """A single-direction module whose hidden width is twice the pair width, at D64 / D128 (AF3 / Protenix's template TriMul: pair
+    64, hidden 128). Per token it has exactly the bidirectional block's shapes -- K1 writes 4 D planes, K3 normalises a 2 D-channel
+    contraction output and projects it back to D, the same B1 / B2 / B7 backward -- only the cuBLAS contraction between the kernels
+    pairs the planes differently (all 2 D channels in the module's one direction instead of D per direction). So the bidirectional
+    D64 / D128 paths serve it with ``direction`` 1 / 2 (``h100_d64_b200gate``, ``h100_training``), and inference is the K1 / K3
+    unit at hidden 2 D (the bidirectional one) with the one-direction contraction."""
+    return hasattr(module, "outgoing") and pair.shape[-1] in (64, 128) and module.d_hidden == 2 * pair.shape[-1]
+
+
+def _hidden_ok(module, pair) -> bool:
+    """The hidden widths the D64 / D128 native paths take: the pair width, or twice it in one direction (``_hidden_twice``)."""
+    return module.d_hidden == pair.shape[-1] or _hidden_twice(module, pair)
+
+
 def _plan(module, pair, inference):
     """(chunk, checkpoint): the samples one launch takes and whether training chunks are checkpointed. chunk == B: one call."""
     bsz = pair.shape[0]
@@ -83,8 +99,9 @@ def _plan(module, pair, inference):
     granule = TRAINING_GRANULE if d >= 256 and inference else 16
     n = -(-n // granule) * granule                    # the padded length one launch runs (``_padded_inference``)
     bidirectional = not hasattr(module, "outgoing")
+    planes4 = bidirectional or _hidden_twice(module, pair)  # hidden 2 D in one direction: the bidirectional working set
     tokens = d * n * n                                 # bytes per sample = coefficient x tokens
-    peak = _WORKING_SET_PER_CHANNEL[(not inference, bidirectional, d <= 128)] * 1.05
+    peak = _WORKING_SET_PER_CHANNEL[(not inference, planes4, d <= 128)] * 1.05
     budget = CHUNK_MEMORY_FRACTION * _total_memory(pair.device.index if pair.device.index is not None else 0)
     limit = max(1, min(max_batch(d), (2 ** 31 - 1) // (n * n)))
     if (not inference and d >= 256 and bsz <= PER_SAMPLE_LEGACY_TRAINING.get((d, pair.shape[1], bidirectional), 0)
@@ -102,7 +119,7 @@ def _plan(module, pair, inference):
         free = budget - bsz * 4.0 * tokens
         return max(1, min(limit, int(free // (peak * tokens)))), False
     caller = bsz * _CALLER_PER_CHANNEL * tokens
-    saved = (_HELD_PER_CHANNEL[(bidirectional, d <= 128)] * 1.05 - 2.0) * tokens      # y is the caller's
+    saved = (_HELD_PER_CHANNEL[(planes4, d <= 128)] * 1.05 - 2.0) * tokens      # y is the caller's
     work = (peak - _CALLER_PER_CHANNEL) * tokens - saved                              # one chunk's backward workspace
     free = budget - caller - bsz * saved
     if free >= work:                                   # every chunk's saved state fits next to one chunk's backward workspace
@@ -296,7 +313,7 @@ def _serves_one(module, pair: torch.Tensor) -> bool:
     ):
         return False
     if (
-        module.d_hidden != pair.shape[-1]
+        not _hidden_ok(module, pair)
         or module.ln_pair.eps != 1e-5
         or module.ln_out.eps != 1e-5
     ):
@@ -441,9 +458,12 @@ def _update(module, pair, mask, dropscale, *, bidirectional=True):
     from miniworld_engine.kernels.trimul_inproj.cuda._h100_runtime import (
         preserve_training,
     )
+    # hidden 2 D in one direction: the bidirectional D64 / D128 kernels (``_hidden_twice``), whatever the bucket policy
+    planes4 = bidirectional or _hidden_twice(module, pair)
     preserve = preserve_training(d) or _selected_native_training(module, pair)
     # (B samples take the preserved D128 kernels below: the general buckets run one sample)
-    if _native_buckets_enabled(pair) and d != 64 and not preserve and pair.shape[0] == 1:
+    if (_native_buckets_enabled(pair) and d != 64 and not preserve and pair.shape[0] == 1
+            and (bidirectional or not planes4)):
         from miniworld_engine.kernels.trimul_inproj.cuda._h100_pack import (
             training_unit_scale,
         )
@@ -461,9 +481,9 @@ def _update(module, pair, mask, dropscale, *, bidirectional=True):
         if not bidirectional:
             return single_trimul(module.outgoing, *args)
         return native_buckets(True, True, *args)
-    # Bidirectional D128 takes the token mask [n]: its front K1 and B7 form m[i] & m[j] themselves.
-    # The other paths read a bf16 pair mask in forward and backward, so it is built once here.
-    token_path = bidirectional and d == 128 and (mask is None or mask.ndim == 2)
+    # Bidirectional D128 (and hidden 2 D in one direction, its kernels) takes the token mask [n]: its front K1 and B7 form m[i] & m[j]
+    # themselves. The other paths read a bf16 pair mask in forward and backward, so it is built once here.
+    token_path = planes4 and d == 128 and (mask is None or mask.ndim == 2)
     if preserve and not token_path:
         from miniworld_engine.kernels.trimul_inproj.cuda._h100_pack import (
             owned_training_mask,
@@ -504,12 +524,15 @@ def _update(module, pair, mask, dropscale, *, bidirectional=True):
         module.to_gate.weight,
         module.to_out.weight,
     )
-    if bidirectional and d == 64:
-        from miniworld_engine.kernels.trimul_inproj.cuda.h100_d64_b200gate import (
-            bidirectional_trimul as apply,
-        )
-    elif bidirectional:
-        apply = bidirectional_trimul
+    if planes4:
+        if d == 64:
+            from miniworld_engine.kernels.trimul_inproj.cuda.h100_d64_b200gate import (
+                bidirectional_trimul as apply,
+            )
+        else:
+            apply = bidirectional_trimul
+        if not bidirectional:   # hidden 2 D in one direction: the contraction pairs all 2 D channels in it
+            apply = functools.partial(apply, direction=1 if module.outgoing else 2)
     elif d in (64, 256, 384, 512):
         if d == 64:
             from miniworld_engine.kernels.trimul_inproj.cuda.h100_uni_d64_b200gate import (
@@ -595,7 +618,7 @@ def _serves_inference_one(module, pair, *, bidirectional, dropscale=None):
         return False
     if ((pair.shape[-1] in (64, 128) or (not bidirectional and _native_residual_bucket(module, pair)))
             and (pair.shape[1] < MIN_INFERENCE_LENGTH
-                 or module.d_hidden != pair.shape[-1] or module.ln_pair.eps != 1e-5
+                 or not _hidden_ok(module, pair) or module.ln_pair.eps != 1e-5
                  or any(p.dtype != torch.float32 for p in (
                      module.ln_pair.weight, module.ln_pair.bias,
                      module.ln_out.weight, module.ln_out.bias)))):

{miniworld_engine-2.3.5 → miniworld_engine-2.3.6}/src/miniworld_engine/kernels/augmented_attention/cuda/qknorm_local.cu

@@ -3,6 +3,8 @@
 //
 //   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
@@ -51,7 +53,8 @@ __device__ __forceinline__ float octet_sum(float v) {          // the 8 lanes of
 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) {
+             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;
@@ -64,7 +67,16 @@ qknorm_fwd_k(const T* __restrict__ q, const T* __restrict__ k, const float* __re
     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]);
-    store4<T>(dst[t], make_float4(x.x * rstd * wv.x, x.y * rstd * wv.y, x.z * rstd * wv.z, x.w * rstd * wv.w), rnd);
+    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);
   }
 }
 
@@ -132,6 +144,31 @@ void qknorm_fwd(at::Tensor q, at::Tensor k, at::Tensor wq, at::Tensor wk, double
   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) {
@@ -157,5 +194,6 @@ at::Tensor qknorm_bwd(at::Tensor dq, at::Tensor dk, at::Tensor qr, at::Tensor kr
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("qknorm_fwd", &qknorm_fwd);
+  m.def("qknorm_rope_fwd", &qknorm_rope_fwd);
   m.def("qknorm_bwd", &qknorm_bwd);
 }