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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "miniworld-engine"
-version = "2.3.6"
+version = "2.3.8"
 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.6 → miniworld_engine-2.3.8}/src/miniworld_engine/architecture/blocks/pairformer.py

@@ -116,12 +116,17 @@ class PairformerBlock(nn.Module):
         transition_pair: nn.Module,
         attention_single: nn.Module | None = None,
         transition_single: nn.Module | None = None,
+        parallel: bool = False,
     ) -> "PairformerBlock":
         """Compose checkpoint-compatible ops with the standard residual contract.
 
         Model adapters may provide different parameter layouts or mask semantics.
         Every supplied self-residual op returns its complete residual output;
         this shared block owns only ordering. All tensors retain a batch axis.
+
+        ``parallel`` (Chai-1, RoseTTAFold3): every part reads the block input, ``pair + sum_i f_i(pair)`` (the single attention reads
+        the pair entering the block, the single transition the single entering it), summed from the parts' residual outputs by
+        ``kernels.dit_variants.cuda.parallel_residual``.
         """
         if (attention_single is None) != (transition_single is None):
             raise ValueError(
@@ -142,6 +147,7 @@ class PairformerBlock(nn.Module):
         block.tri_atten_starting = attention_starting
         block.tri_atten_ending = attention_ending
         block.transition_pair = transition_pair
+        block.parallel = parallel
         if block.use_single:
             block.attention_pair_bias = attention_single
             block.transition_single = transition_single
@@ -161,6 +167,8 @@ class PairformerBlock(nn.Module):
     ]:
         """Forward pass."""
         triangle_mask = pair_mask if pair_mask is not None else mask
+        if getattr(self, "parallel", False):
+            return self._parallel_forward(pair, single, mask, triangle_mask)
         if self.use_bidirectional_trimul:
             pair = self.tri_multi(pair, triangle_mask)
         else:
@@ -179,6 +187,41 @@ class PairformerBlock(nn.Module):
             single = self.transition_single(single)
         return pair, single
 
+    def _parallel_forward(self, pair, single, mask, triangle_mask):
+        """Every part on the block input; the parts own their residuals, so the update is ``pair + sum_i (y_i - pair)``."""
+        if self.use_bidirectional_trimul:
+            outs = [self.tri_multi(pair, triangle_mask)]
+        else:
+            outs = [self.tri_multi_outgoing(pair, triangle_mask), self.tri_multi_incoming(pair, triangle_mask)]
+        if self.use_triangle_attention:
+            outs += [self.tri_atten_starting(pair, triangle_mask), self.tri_atten_ending(pair, triangle_mask)]
+        outs.append(self.transition_pair(pair))
+        new_pair = _parallel_residual(pair, outs)
+        if self.use_single:
+            if single is None:
+                msg = "single must be provided when use_single=True"
+                raise ValueError(msg)
+            single = _parallel_residual(single, [self.attention_pair_bias(single, pair, mask), self.transition_single(single)])
+        return new_pair, single
+
+
+def _parallel_residual(x: torch.Tensor, outs: list[torch.Tensor]) -> torch.Tensor:
+    """x + sum_i (y_i - x): the CUDA sum (fp32, one rounding) on H100, else the same in torch."""
+    if (
+        x.is_cuda
+        and x.dtype in (torch.bfloat16, torch.float32)
+        and torch.cuda.get_device_capability(x.device) == (9, 0)
+        and x.numel() % 4 == 0
+        and len(outs) <= 8
+    ):
+        from miniworld_engine.kernels.dit_variants import cuda
+
+        x = x.contiguous()
+        return cuda.parallel_residual(x, [y.contiguous() for y in outs])
+    acc = x.to(torch.promote_types(x.dtype, torch.float32))
+    total = acc + sum(y.to(acc.dtype) - acc for y in outs)
+    return total.to(x.dtype)
+
 
 class Pairformer(nn.Module):
     """The Pairformer stack consisting of multiple Pairformer blocks."""

{miniworld_engine-2.3.6 → miniworld_engine-2.3.8}/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.5,650b4231cfc7,clean,2026-10-21,v2.3.4-2-g650b4231,
+#provenance,2.3.7,435116ba906e,clean,2026-10-21,v2.3.6-5-g435116ba,
 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=2.50e-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.09e-07 dbeta=1.67e-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
@@ -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=4.02e-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=3.22e-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
@@ -130,24 +130,24 @@ transition_bwd_residual_sm90_cuda,cuda,transition,miniworld_engine/kernels/trans
 swa_dit_inproj_fwd_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/forward.py,ok,rel out=5.90e-03 (band 5e-02 declared),bf16
 swa_dit_softmax_fwd_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/forward.py,ok,rel out=5.90e-03 (band 5e-02 declared),bf16
 swa_dit_output_swiglu_fwd_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/forward.py,ok,rel out=5.90e-03 (band 5e-02 declared),bf16
-swa_dit_swiglu_bwd_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dmod=3.64e-03 dw_up=4.36e-03 dw_down=4.69e-03 (band 5e-02 declared),bf16
-swa_dit_swiglu_dw_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dw_up=4.36e-03 dw_down=4.69e-03 (band 5e-02 declared),bf16
-swa_dit_output_bwd_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dmod=3.64e-03 dw_out=4.61e-03 (band 5e-02 declared),bf16
-swa_dit_softmax_bwd_dq_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dq=5.60e-03 dw_qkv=6.01e-03 (band 5e-02 declared),bf16
-swa_dit_softmax_bwd_dkdv_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dq=5.60e-03 dw_qkv=6.01e-03 (band 5e-02 declared),bf16
-swa_dit_inproj_bwd_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dq=5.60e-03 dmod=3.64e-03 dw_qkv=6.01e-03 dw_gate=4.86e-03 (band 5e-02 declared),bf16
+swa_dit_swiglu_bwd_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dmod=3.40e-03 dw_up=4.71e-03 dw_down=4.70e-03 (band 5e-02 declared),bf16
+swa_dit_swiglu_dw_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dw_up=4.71e-03 dw_down=4.70e-03 (band 5e-02 declared),bf16
+swa_dit_output_bwd_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dmod=3.40e-03 dw_out=4.32e-03 (band 5e-02 declared),bf16
+swa_dit_softmax_bwd_dq_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dq=6.18e-03 dw_qkv=6.74e-03 (band 5e-02 declared),bf16
+swa_dit_softmax_bwd_dkdv_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dq=6.18e-03 dw_qkv=6.74e-03 (band 5e-02 declared),bf16
+swa_dit_inproj_bwd_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/backward.py,ok,rel dq=6.18e-03 dmod=3.40e-03 dw_qkv=6.74e-03 dw_gate=4.86e-03 (band 5e-02 declared),bf16
 swa_dit_inproj_fwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_qkvg_fwd.cu,ok,rel out=5.90e-03 (band 5e-02 declared),bf16
 swa_dit_output_swiglu_fwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_ffn_fwd.cu,ok,rel out=5.90e-03 (band 5e-02 declared),bf16
-swa_dit_swiglu_bwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_ffn_bwd.cu,ok,rel dq=5.60e-03 dmod=3.64e-03 dw_up=4.36e-03 dw_down=4.69e-03 (band 5e-02 declared),bf16
+swa_dit_swiglu_bwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_ffn_bwd.cu,ok,rel dq=6.18e-03 dmod=3.40e-03 dw_up=4.71e-03 dw_down=4.70e-03 (band 5e-02 declared),bf16
 swa_dit_softmax_fwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_attn.cu,ok,rel out=5.90e-03 (band 5e-02 declared),bf16
-swa_dit_softmax_bwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_attn.cu,ok,rel dq=5.60e-03 dmod=3.64e-03 dw_qkv=6.01e-03 (band 5e-02 declared),bf16
-swa_dit_output_bwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_slab.cu,ok,rel dq=5.60e-03 dmod=3.64e-03 dw_out=4.61e-03 dw_gate=4.86e-03 (band 5e-02 declared),bf16
-swa_dit_inproj_bwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_slab.cu,ok,rel dq=5.60e-03 dmod=3.64e-03 dw_qkv=6.01e-03 dw_gate=4.86e-03 (band 5e-02 declared),bf16
+swa_dit_softmax_bwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_attn.cu,ok,rel dq=5.60e-03 dmod=3.64e-03 dw_qkv=5.72e-03 (band 5e-02 declared),bf16
+swa_dit_output_bwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_slab.cu,ok,rel dq=6.18e-03 dmod=3.40e-03 dw_out=4.32e-03 dw_gate=4.86e-03 (band 5e-02 declared),bf16
+swa_dit_inproj_bwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_slab.cu,ok,rel dq=6.18e-03 dmod=3.40e-03 dw_qkv=6.74e-03 dw_gate=4.86e-03 (band 5e-02 declared),bf16
 swa_dit_inproj_fwd_rows_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_slab.cu,ok,rel out=5.90e-03 (band 5e-02 declared),bf16
 swa_dit_output_fwd_rows_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_slab.cu,ok,rel out=5.90e-03 (band 5e-02 declared),bf16
 swa_dit_swiglu_fwd_rows_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_slab.cu,ok,rel out=5.90e-03 (band 5e-02 declared),bf16
-swa_dit_swiglu_bwd_rows_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_slab.cu,ok,rel dq=5.60e-03 dmod=3.64e-03 dw_up=4.36e-03 dw_down=4.69e-03 (band 5e-02 declared),bf16
+swa_dit_swiglu_bwd_rows_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_slab.cu,ok,rel dq=6.18e-03 dmod=3.40e-03 dw_up=4.71e-03 dw_down=4.70e-03 (band 5e-02 declared),bf16
 swa_dit_gemm_fwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_mod_wg.cu,ok,rel mod=2.23e-03 (band 5e-02 declared),bf16
 swa_dit_gemm_bwd_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_mod.cu,ok,rel dc=4.71e-03 dw=3.08e-03 (band 5e-02 declared),bf16
-swa_dit_dw_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_wgrad_wg.cu,ok,rel dw_qkv=6.01e-03 dw_gate=4.86e-03 dw_out=4.61e-03 dw_up=4.36e-03 dw_down=4.69e-03 (band 5e-02 declared),bf16
+swa_dit_dw_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_wgrad_wg.cu,ok,rel dw_qkv=5.56e-03 dw_gate=5.62e-03 dw_out=5.42e-03 dw_up=4.62e-03 dw_down=5.23e-03 (band 5e-02 declared),bf16
 swa_dit_transpose_sm90_cuda,cuda,swa_dit,miniworld_engine/kernels/swa_dit/cuda/swa_pack.cu,ok,rel wab64=0.00e+00 wab32=0.00e+00 wdt=0.00e+00 wabt=0.00e+00 (band 5e-02 declared),bf16

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

@@ -136,6 +136,9 @@ def cases() -> list[dict]:
                         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})
+    # Boltz-2 / RoseTTAFold3 / Chai-1's variants (kernels/dit_variants, the 64-channel and output-variant TriangleAttention, the biased /
+    # grouped OuterProductMean, the parallel Pairformer sum): bf16 inference; their extensions take no shape keys, one small case covers them.
+    out.append({"kind": "model_variants", "train": False, "dtype": "bfloat16", "heavy": False})
     for i, c in enumerate(out):
         c["id"] = i
     return out
@@ -229,8 +232,11 @@ def _run_block(c: dict) -> None:
         )
         from miniworld_engine.modules.swa_atom_attention import build_attention_params
 
+        # esmfold2 runs the fused block (kernels/swa_dit), as MiniWorld does (MINIWORLD_SWA_FUSED): settings.swa_fused is off by
+        # default, and without it no engine case reached the fused block's extensions (2.3.0 - 2.3.6 got them from MiniWorld only).
         config = SWAAtomTransformer.Config(n_block=1, block_style="af3" if c["variant"].startswith("af3") else "esmfold2",
-                                           global_attn=c["variant"] == "af3_global")
+                                           global_attn=c["variant"] == "af3_global",
+                                           fused_triton=True if c["variant"] == "esmfold2" else None)
         model = SWAAtomTransformer(config).cuda().to(pdt)
         _randomize(model)
         a, s = c["samples"], c["N"]
@@ -241,12 +247,20 @@ def _run_block(c: dict) -> None:
         model.train(c["train"])
         if c["train"]:
             q = q.requires_grad_(True)
-        with torch.set_grad_enabled(c["train"]), torch.autocast("cuda", dtype=torch.bfloat16,
-                                                                 enabled=autocast or dt is torch.bfloat16):
-            out = model(q, cond, params)
-            if c["train"]:
-                out.float().sum().backward()
-        torch.cuda.synchronize()
+        # fp32 activations take the native H100 stages under "highest" and the composed TF32 ones (swa_fp32.cu) under "high",
+        # what MiniWorld's trainer sets: both run.
+        before = torch.get_float32_matmul_precision()
+        try:
+            for precision in ("highest", "high") if dt is torch.float32 else (before,):
+                torch.set_float32_matmul_precision(precision)
+                with torch.set_grad_enabled(c["train"]), torch.autocast("cuda", dtype=torch.bfloat16,
+                                                                         enabled=autocast or dt is torch.bfloat16):
+                    out = model(q, cond, params)
+                    if c["train"]:
+                        out.float().sum().backward()
+                torch.cuda.synchronize()
+        finally:
+            torch.set_float32_matmul_precision(before)
         return
     elif c["kind"] == "trimul_hidden2x":
         from miniworld_engine.modules.exceptions import ImplementationType as MI
@@ -295,6 +309,9 @@ def _run_block(c: dict) -> None:
         _randomize(model)
         _step(model, args, False, autocast, {"rope": (angles.cos(), angles.sin())})
         return
+    elif c["kind"] == "model_variants":
+        _model_variants()
+        return
     elif c["kind"] == "local_dit":
         from miniworld_engine.modules.exceptions import ImplementationType as MI
         from miniworld_engine.modules.local_dit import LocalDiTBlock
@@ -310,6 +327,58 @@ def _run_block(c: dict) -> None:
     _step(model, args, c["train"], autocast)
 
 
+def _model_variants() -> None:
+    """Every Boltz-2 / RoseTTAFold3 / Chai-1 variant path once, bf16 inference, small shapes."""
+    import torch
+
+    from miniworld_engine.architecture.blocks.pairformer import PairformerBlock
+    from miniworld_engine.modules import Transition, TriangleMultiplication
+    from miniworld_engine.modules import dit_variant as dv
+    from miniworld_engine.modules.exceptions import ImplementationType as MI
+    from miniworld_engine.modules.outer_product import OuterProductMean
+    from miniworld_engine.modules.triangle_attention import TriangleAttention
+
+    bf, dev = torch.bfloat16, "cuda"
+
+    def prep(m):
+        m = m.to(dev).eval().to(bf)
+        _randomize(m)
+        return m
+
+    with torch.no_grad():
+        n = 256
+        for opts in (dv.BOLTZ2, dv.ROSETTAFOLD3, dv.CHAI1):
+            prep(dv.DiTVariantStack(2, **opts))(torch.randn(2, n, 768, device=dev, dtype=bf),
+                                               torch.randn(n, opts["d_cond"], device=dev, dtype=bf),
+                                               torch.randn(n, n, opts["d_pair"], device=dev, dtype=bf), torch.ones(n, dtype=torch.bool, device=dev))
+        n = 512
+        for opts in (dv.BOLTZ2_ATOM, dv.ROSETTAFOLD3_ATOM, dv.CHAI1_ATOM):
+            uid = torch.arange(n, device=dev) // 40 if opts is dv.CHAI1_ATOM else None
+            prep(dv.AtomDiTVariantStack(2, **opts))(torch.randn(2, n, 128, device=dev, dtype=bf), torch.randn(n, 128, device=dev, dtype=bf),
+                                                    torch.randn(n // 32, 32, 128, 16, device=dev, dtype=bf),
+                                                    torch.arange(n, device=dev) < n - 20, uid)
+        n = 128
+        mask = torch.ones(1, n, dtype=torch.bool, device=dev)
+        for d_pair, d_hidden, starting, options in ((256, 256, True, {}), (256, 256, False, {"untransposed_output": True}),
+                                                     (256, 256, True, {"dual_output": True}), (64, 128, False, {"untransposed_output": True}),
+                                                     (128, 128, True, {"gate_bias": True, "out_bias": True}),
+                                                     (64, 256, False, {"gate_bias": True, "out_bias": True})):
+            m = prep(TriangleAttention(d_pair, 4, d_hidden=d_hidden, starting=starting, transposed_bias=False, implementation=MI.MINIWORLD,
+                                       p_drop=0.0, **options))
+            m(torch.randn(1, n, n, d_pair, device=dev, dtype=bf), mask)
+        msa_mask = torch.ones(1, 256, n, dtype=torch.bool, device=dev)
+        for d_pair, d_hidden, options in ((128, 32, {"projection_bias": True, "normalize_before_proj": False}), (256, 64, {"groups": 8})):
+            prep(OuterProductMean(64, d_pair, d_hidden, implementation=MI.MINIWORLD, **options))(
+                torch.randn(1, 256, n, 64, device=dev, dtype=bf), msa_mask)
+        block = prep(PairformerBlock.from_components(
+            triangle_outgoing=TriangleMultiplication(128, outgoing=True, implementation=MI.MINIWORLD, p_drop=0.0),
+            triangle_incoming=TriangleMultiplication(128, outgoing=False, implementation=MI.MINIWORLD, p_drop=0.0),
+            attention_starting=None, attention_ending=None, transition_pair=Transition(128, 4, implementation=MI.MINIWORLD),
+            parallel=True))
+        block(torch.randn(1, n, n, 128, device=dev, dtype=bf), None, mask)
+    torch.cuda.synchronize()
+
+
 def run(c: dict) -> dict:
     """Run one descriptor; never raises. ``status``: ok / skip (the module refused the shape) / oom / error."""
     import torch

{miniworld_engine-2.3.6 → miniworld_engine-2.3.8}/src/miniworld_engine/integrations/triattn_h100_sm90.py

@@ -43,6 +43,10 @@ def _kind(module, pair: torch.Tensor, mask: torch.Tensor | None) -> int | None:
         return None
     if not module.use_self_attention or module.use_qk_norm:
         return None
+    if any(b is not None for b in _biases(module)) and (
+        torch.is_grad_enabled() or (module.training and module.p_drop > 0.0)
+    ):
+        return None  # RoseTTAFold3's gate / out-projection biases: inference only (the backward knows neither)
     if module._bias_untransposed and (
         torch.is_grad_enabled() or (module.training and module.p_drop > 0.0)
     ):
@@ -138,14 +142,16 @@ def _infer_op(
 def _infer(kind, args, eps, n_head, head_dim, trans, out=None, flip=False):
     """The inference kernels; `out`: a contiguous tensor like the pair to write the result into (a chunk's rows of a batch's output,
     eager only: an op output may not alias an input)."""
-    pair, lnw, lnb, *weights, mask, _ = args
+    pair, lnw, lnb, *weights = args[:9]
+    mask = args[9]
+    gate_bias, out_bias = (t if t.numel() else None for t in args[11:13]) if len(args) > 11 else (None, None)
     mask = mask if mask.numel() else None
     if kind == FUSED:
         from miniworld_engine.kernels.triangle_attention.cuda.h100_sm90 import inference
 
         w4, wb16, wo16 = _packed(kind, tuple(weights))
         return inference(
-            pair, lnw.float(), lnb.float(), w4, wb16, wo16, mask, eps, trans, out, flip
+            pair, lnw.float(), lnb.float(), w4, wb16, wo16, mask, eps, trans, out, flip, gate_bias, out_bias
         )
     from miniworld_engine.kernels.triangle_attention.cuda import (
         h100_sm90_generic as gen,
@@ -153,7 +159,7 @@ def _infer(kind, args, eps, n_head, head_dim, trans, out=None, flip=False):
 
     wcat, wo_p = _packed(kind, tuple(weights))
     return gen.inference(
-        pair, lnw, lnb, wcat, wo_p, n_head, head_dim, mask, eps, trans, out, flip
+        pair, lnw, lnb, wcat, wo_p, n_head, head_dim, mask, eps, trans, out, flip, gate_bias, out_bias
     )
 
 
@@ -507,6 +513,7 @@ def forward(module, pair: torch.Tensor, mask: torch.Tensor | None) -> torch.Tens
                     *params,
                     m if m is not None else x.new_empty(0, dtype=torch.bool),
                     x.new_empty(0),
+                    *_bias_args(module, x),
                 ]
                 _infer(
                     kind,
@@ -686,7 +693,7 @@ def _forward_one(
     if drop is None and not needs_grad:
         out = _infer_op(
             kind,
-            [pair, *params, mask_t, drop_t],
+            [pair, *params, mask_t, drop_t, *_bias_args(module, pair)],
             eps,
             module.n_head,
             head_dim,
@@ -708,3 +715,13 @@ def _forward_one(
     if Lp != L:
         out = out[:, :L, :L]
     return out.contiguous()
+
+
+def _biases(module) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    """(gate bias, out-projection bias): RoseTTAFold3's TriangleAttention has both, every other model neither."""
+    return module.to_gate.bias, module.to_out.bias
+
+
+def _bias_args(module, like: torch.Tensor) -> list[torch.Tensor]:
+    """The inference op's trailing arguments: the two biases, an empty tensor for an absent one."""
+    return [b.detach() if b is not None else like.new_empty(0) for b in _biases(module)]

miniworld_engine-2.3.8/src/miniworld_engine/kernels/dit_variants/__init__.py

@@ -0,0 +1,2 @@
+"""DiT block variants (Boltz-2, RoseTTAFold3, Chai-1 token DiT): CUDA rows (``cuda/``), the H100 inference path (``interface``) and the
+PyTorch reference (``reference``)."""

miniworld_engine-2.3.8/src/miniworld_engine/kernels/dit_variants/cuda/__init__.py

@@ -0,0 +1,75 @@
+"""CUDA row kernels of the DiT block variants (``rows.cu``), sm_90a. Built on first use (``load_extension``), never at import."""
+
+from __future__ import annotations
+
+import functools
+from pathlib import Path
+
+import torch
+
+_dir = Path(__file__).parent
+
+
+@functools.lru_cache(maxsize=None)
+def ext():
+    from ..._nvcc import ensure_cuda_home, gencodes, host_flags, load_extension
+
+    ensure_cuda_home()
+    return load_extension(
+        name="dit_variants_rows_sm90a",
+        sources=[str(_dir / "rows.cu")],
+        extra_cuda_cflags=[
+            *host_flags(),
+            "-O3",
+            "-std=c++17",
+            "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+            "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+            "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+            *gencodes("90a"),
+        ],
+        extra_cflags=["-std=c++17"],
+        verbose=False,
+    )
+
+
+def ln_mod_rows(x, mods, L, eps, sig, *, xin=None):
+    """xhat = LN(x) (no affine; x = xin first when given); for each (ms, mb, out) of ``mods`` (one to three): out = xhat * f(ms) + mb,
+    f = sigmoid when ``sig`` else identity; ms / mb [L, D] bf16 table views (one row stride) indexed by row % L. D = 768 or 128."""
+    ms, mb, out = (list(t) for t in zip(*mods, strict=True))
+    ext().ln_mod_rows(xin, x, ms, mb, out, int(L), float(eps), bool(sig))
+
+
+def res_rows(x, ya, ga, L, *, yt=None, gt=None, ms=None, mb=None, out_ln=None, final=None, eps=1e-5, sig=True):
+    """x += sigmoid(ga) ya [+ sigmoid(gt) yt]; into ``final`` (bf16) when given, else back into x (fp32); then
+    out_ln = LN(x) * f(ms) + mb when ``out_ln`` is given."""
+    ext().res_rows(x, ya, ga, yt, gt, ms, mb, out_ln, final, int(L), float(eps), bool(sig))
+
+
+def qk_ln_rows(q, k, wq, bq, wk, bk, eps):
+    """In place: q <- LN(q) * wq + bq, k <- LN(k) * wk + bk over D (768 or 128) columns (bf16 row views; fp32 [D] affine)."""
+    ext().qk_ln_rows(q, k, wq, bq, wk, bk, float(eps))
+
+
+def gate_rows(o, g):
+    """In place: o <- o * sigmoid(g) (bf16 row views)."""
+    ext().gate_rows(o, g)
+
+
+def glu_rows(ab, out, up):
+    """out = silu(g) * v [* u] for ab = [g | v | u] (bf16)."""
+    ext().glu_rows(ab, out, bool(up))
+
+
+def ln_rows(z, out, eps=1e-5):
+    """out = LN(z), no affine, z / out bf16 [rows, 128 | 256]."""
+    ext().ln_rows(z, out, float(eps))
+
+
+def available() -> bool:
+    return torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0)
+
+
+def parallel_residual(x, outputs):
+    """x + sum_i (y_i - x) with fp32 sums and one rounding: a parallel block's output from the residual outputs y_i = x + f_i(x) of its
+    parts (each engine module owns its residual). x / y_i contiguous, one shape, bf16 or fp32; one to eight outputs."""
+    return ext().parallel_residual(x, list(outputs))