miniworld-engine 2.3.7__tar.gz → 2.3.9__tar.gz

{miniworld_engine-2.3.7 → miniworld_engine-2.3.9}/PKG-INFO

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

@@ -1,6 +1,6 @@
 [project]
 name = "miniworld-engine"
-version = "2.3.7"
+version = "2.3.9"
 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.7 → miniworld_engine-2.3.9}/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.7 → miniworld_engine-2.3.9}/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.6,9bb9bd5abb7c,clean,2026-10-21,v2.3.5-9-g9bb9bd5a,
-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.8,7bd3ec119a29,clean,2026-10-22,v2.3.7-11-g7bd3ec11,
+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=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=2.99e-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
@@ -45,7 +45,7 @@ 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
@@ -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/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_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: <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
@@ -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=5.63e-07 db=1.63e-07 (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=1.63e-07 (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
@@ -115,7 +115,7 @@ mpnn_edge_tail_bwd_dx_recompute_triton,triton,mpnn_edge_tail,miniworld_engine/ke
 mpnn_node_message_fwd_gemm_triton,triton,mpnn_node_message,miniworld_engine/kernels/mpnn_node_message/triton/main.py,ok,rel out=1.49e-03 (band 5e-02),bf16
 mpnn_node_message_bwd_recompute_triton,triton,mpnn_node_message,miniworld_engine/kernels/mpnn_node_message/triton/main.py,failed,WRONG NUMBERS: rel edge_states=1.10e-01 query_projection=1.02e-02 neighbor_projection=3.51e-02 edge_weight=2.20e-02 hidden_weight=6.86e-03 hidden_bias=3.82e-03 (band 5e-02),bf16
 mpnn_node_message_bwd_dx_triton,triton,mpnn_node_message,miniworld_engine/kernels/mpnn_node_message/triton/main.py,failed,WRONG NUMBERS: rel edge_states=1.10e-01 query_projection=1.02e-02 neighbor_projection=3.51e-02 edge_weight=2.20e-02 hidden_weight=6.86e-03 hidden_bias=3.82e-03 (band 5e-02),bf16
-mpnn_relative_position_bwd_reduce_triton,triton,mpnn_relative_position,miniworld_engine/kernels/mpnn_relative_position/triton/main.py,ok,rel grad_table=5.77e-08 grad_bias=7.24e-08 (band 5e-02),bf16
+mpnn_relative_position_bwd_reduce_triton,triton,mpnn_relative_position,miniworld_engine/kernels/mpnn_relative_position/triton/main.py,ok,rel grad_table=1.15e-07 grad_bias=7.24e-08 (band 5e-02),bf16
 mpnn_message_fwd_gemm_triton,triton,mpnn_message,miniworld_engine/kernels/mpnn_message/triton/main.py,ok,rel out=1.07e-03 (band 5e-02),bf16
 mpnn_message_fwd_gelu_reduce_triton,triton,mpnn_message,miniworld_engine/kernels/mpnn_message/triton/main.py,ok,rel out=1.07e-03 (band 5e-02),bf16
 mpnn_message_bwd_dx_triton,triton,mpnn_message,miniworld_engine/kernels/mpnn_message/triton/main.py,ok,rel preactivation=1.01e-02 weight=4.23e-03 bias=3.97e-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.7 → miniworld_engine-2.3.9}/src/miniworld_engine/build/bundle_coverage.py

@@ -136,6 +136,14 @@ 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})
+    # The sampler's glue (kernels/sampler_ops: random rigid motion, Kabsch + Euler update, encoder pooling): one extension, no shape keys.
+    out.append({"kind": "sampler_ops", "train": False, "dtype": "float32", "heavy": False})
+    # Every hand-CUDA stage extension of the fused SWA atom block (kernels/swa_dit/cuda/loader.py): a stage the dispatch tries first (modwg)
+    # and falls back from is reached by no model shape, and would JIT-compile on a first call.
+    out.append({"kind": "swa_stages", "train": False, "dtype": "bfloat16", "heavy": False})
     for i, c in enumerate(out):
         c["id"] = i
     return out
@@ -306,6 +314,15 @@ 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"] == "sampler_ops":
+        _sampler_ops()
+        return
+    elif c["kind"] == "swa_stages":
+        _swa_stages()
+        return
     elif c["kind"] == "local_dit":
         from miniworld_engine.modules.exceptions import ImplementationType as MI
         from miniworld_engine.modules.local_dit import LocalDiTBlock
@@ -321,6 +338,92 @@ def _run_block(c: dict) -> None:
     _step(model, args, c["train"], autocast)
 
 
+def _swa_stages() -> None:
+    """Load every stage extension of the fused SWA atom block once (a stage that fails to build is an error, not a silent fallback)."""
+    from miniworld_engine.kernels.swa_dit.cuda import loader
+
+    failed = {}
+    for stage in loader.SOURCES:
+        if loader.extension(stage) is None:
+            failed[stage] = loader.ERRORS.get(stage, "?")
+    if failed:
+        raise RuntimeError(f"swa_dit stage extensions that did not build: {failed}")
+
+
+def _sampler_ops() -> None:
+    """Every sampler-glue kernel once: the rigid motion + re-noising, the Kabsch + Euler update, the encoder pooling."""
+    import torch
+
+    from miniworld_engine.kernels.sampler_ops import interface
+
+    dev = "cuda"
+    b, n = 2, 1536
+    coords, noise = torch.randn(b, n, 3, device=dev), torch.randn(b, n, 3, device=dev)
+    mask = torch.ones(n, device=dev)
+    sigma, sigma_next = torch.tensor(2.0, device=dev), torch.tensor(1.5, device=dev)
+    noisy, sc = interface.augment_noise(
+        coords, mask, torch.randn(b, 2, 3, device=dev), torch.randn(b, 3, device=dev), noise, sigma, sigma_next,
+        window=False, gamma=0.8, gamma_min=1.0, sigma_min=0.0, sigma_max=float("inf"), noise_scale=1.003, variance_floor=0.0, clamp=False)
+    interface.align_euler(noisy, coords, mask, sc, sigma_next, 1.5)
+    tokens, atoms = 64, 24
+    interface.pool_atoms(torch.randn(b, tokens * atoms, 768, device=dev, dtype=torch.bfloat16),
+                         torch.randint(0, tokens * atoms, (tokens, atoms), device=dev), torch.ones(tokens, atoms, device=dev, dtype=torch.bool),
+                         torch.ones(tokens, atoms, device=dev))
+    torch.cuda.synchronize()
+
+
+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.7 → miniworld_engine-2.3.9}/src/miniworld_engine/integrations/attention_pair_bias_b200.py

@@ -220,7 +220,7 @@ def _inference_packs(params, leaves):
         params,
     )
     return _capture.lookup(
-        _packs, tuple((t.data_ptr(), t._version) for t in params), build, limit=8
+        _packs, tuple((t.data_ptr(), t._version) for t in params), build, limit=_capture.BLOCK_CACHE_LIMIT
     )[:5]
 
 

{miniworld_engine-2.3.7 → miniworld_engine-2.3.9}/src/miniworld_engine/integrations/csrc/opm_epilogue.cu

@@ -915,7 +915,7 @@ __global__ __launch_bounds__(512, 1) void opm_prologue_tma_wide(
     const __grid_constant__ CUtensorMap xm, const __grid_constant__ CUtensorMap mm,
     const __grid_constant__ CUtensorMap wam, const __grid_constant__ CUtensorMap wbm,
     const __grid_constant__ CUtensorMap am, const __grid_constant__ CUtensorMap bm,
-    const float* gamma, const float* beta, __nv_bfloat16* mask16, float* stats, unsigned* bits,
+    const float* gamma, const float* beta, __nv_bfloat16* mask16, float* stats, unsigned* bits, const float* pbias,
     int S, int N, int tiles, float eps) {
   using namespace opm_pro_tma;
   using C = opm_pro_wide::Cfg<CMT>;
@@ -1025,6 +1025,11 @@ __global__ __launch_bounds__(512, 1) void opm_prologue_tma_wide(
       for (int k = 0; k < 16; ++k) {
         int col = k * 8 + (lane % 4) * 2, ch = cw * 16 + lane / 4, h = col / 64, cc = col % 64;
         __nv_bfloat16* out = st + b * TT + h * 4096;
+        if (pbias != nullptr) {   // RoseTTAFold3's projection biases [a | b], masked like the projections (mask16 of this tile is written)
+          const float k0 = __bfloat162float(mask16[(long)(s0 + col) * N + i]), k1 = __bfloat162float(mask16[(long)(s0 + col + 1) * N + i]);
+          acc[k * 4] += pbias[ch] * k0; acc[k * 4 + 1] += pbias[ch] * k1;
+          acc[k * 4 + 2] += pbias[ch + 8] * k0; acc[k * 4 + 3] += pbias[ch + 8] * k1;
+        }
         *reinterpret_cast<__nv_bfloat162*>(out + wg::off_n(64, 0, ch, cc / 8) + cc % 8) = __floats2bfloat162_rn(acc[k * 4], acc[k * 4 + 1]);
         *reinterpret_cast<__nv_bfloat162*>(out + wg::off_n(64, 0, ch + 8, cc / 8) + cc % 8) = __floats2bfloat162_rn(acc[k * 4 + 2], acc[k * 4 + 3]);
       }
@@ -2386,7 +2391,8 @@ std::vector<torch::Tensor> opm_native_prologue(torch::Tensor m,torch::Tensor mas
 // A2 [(i,c), s] / BT [(j,e), s] layouts, the LayerNorm statistics, the bf16 mask, and the exact mask count [N, N] (popcount of
 // the bit mask the prologue wrote).  Parameters arrive already as gamma / beta fp32 [CM] and Wa / Wb bf16 [32, CM].
 std::vector<torch::Tensor> opm_prologue_wide(torch::Tensor m, torch::Tensor mask, torch::Tensor gamma, torch::Tensor beta,
-                                             torch::Tensor wa, torch::Tensor wb, double eps, bool save_stats) {
+                                             torch::Tensor wa, torch::Tensor wb, double eps, bool save_stats,
+                                             c10::optional<torch::Tensor> pbias) {
   using Cf = opm_pro_wide::Cfg<CM_MSA>;
   TORCH_CHECK(m.is_cuda() && m.is_contiguous() && m.scalar_type() == torch::kBFloat16 && m.dim() == 3 && m.size(2) == CM_MSA,
               "wide OPM prologue: m is contiguous bf16 [S, N, ", CM_MSA, "]");
@@ -2431,8 +2437,11 @@ std::vector<torch::Tensor> opm_prologue_wide(torch::Tensor m, torch::Tensor mask
   __nv_bfloat16* mp = reinterpret_cast<__nv_bfloat16*>(mk.data_ptr());
   float* sp = save_stats ? stats.data_ptr<float>() : nullptr;
   unsigned* bitp = reinterpret_cast<unsigned*>(bits.data_ptr());
+  TORCH_CHECK(!pbias || (pbias->is_cuda() && pbias->scalar_type() == torch::kFloat32 && pbias->is_contiguous() && pbias->numel() == 2 * CH),
+              "wide OPM prologue: pbias fp32 [2 x 32] (the left | right projection biases)");
+  const float* pbp = pbias ? pbias->data_ptr<float>() : nullptr;
   int tiles = N * (S / 128), s_arg = S, n_arg = N; float ep = float(eps);
-  void* args[] = {&xm, &mm, &wam, &wbm, &am, &bm, &gp, &bp, &mp, &sp, &bitp, &s_arg, &n_arg, &tiles, &ep};
+  void* args[] = {&xm, &mm, &wam, &wbm, &am, &bm, &gp, &bp, &mp, &sp, &bitp, &pbp, &s_arg, &n_arg, &tiles, &ep};
   int grid = std::min(tiles, at::cuda::getCurrentDeviceProperties()->multiProcessorCount);
   C10_CUDA_CHECK(cudaLaunchKernel(reinterpret_cast<const void*>(kernel), dim3(grid), dim3(512), args, Cf::SMEM,
                                   at::cuda::getCurrentCUDAStream()));
@@ -2478,7 +2487,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("eps"),py::arg("bwp"),py::arg("N"),py::arg("S"),py::arg("native_finish")=false,
         py::arg("raw_weights")=std::vector<torch::Tensor>{});
   m.def("opm_prologue_wide", &opm_prologue_wide, "persistent TMA OPM prologue at this build's d_msa (64 / 128) + exact mask count",
-        py::arg("m"), py::arg("mask"), py::arg("gamma"), py::arg("beta"), py::arg("wa"), py::arg("wb"), py::arg("eps"), py::arg("save_stats"));
+        py::arg("m"), py::arg("mask"), py::arg("gamma"), py::arg("beta"), py::arg("wa"), py::arg("wb"), py::arg("eps"), py::arg("save_stats"),
+        py::arg("pbias") = py::none());
   m.def("opm_native_pack",&opm_native_pack);
   m.def("opm_native_prologue",&opm_native_prologue);
   m.def("opm_native_prologue_tma",&opm_native_prologue_tma);

{miniworld_engine-2.3.7 → miniworld_engine-2.3.9}/src/miniworld_engine/integrations/csrc/pair3.cu

@@ -451,7 +451,7 @@ torch::Tensor pwa_wgv(torch::Tensor g,torch::Tensor v) {
 torch::Tensor pair_fwd3(torch::Tensor z, torch::Tensor mask, torch::Tensor ln_w, double eps, torch::Tensor wb) {
   TORCH_CHECK(z.is_cuda() && z.scalar_type() == torch::kBFloat16 && z.is_contiguous() && z.dim() == 3 && z.size(2) == DZ && z.size(0) == z.size(1), "z: [N, N, 128] bf16");
   const int N = (int)z.size(0);
-  TORCH_CHECK(N % GJ == 0 && N <= 1024, "N must be a multiple of 16 (and <= 1024)");
+  TORCH_CHECK(N % GJ == 0 && N <= 2048, "N must be a multiple of 16 (and <= 2048; the forward holds H rows of N logits in shared memory)");
   TORCH_CHECK(mask.scalar_type() == torch::kBFloat16 && mask.is_contiguous() && mask.numel() == (long)N * N, "mask: [N, N] bf16");
   TORCH_CHECK(ln_w.scalar_type() == torch::kFloat && ln_w.is_contiguous() && ln_w.numel() == DZ, "LN scale fp32[128]");
   auto wb16 = wb.to(torch::kBFloat16).contiguous();
@@ -460,7 +460,7 @@ torch::Tensor pair_fwd3(torch::Tensor z, torch::Tensor mask, torch::Tensor ln_w,
   const int NP = N + 4;
   const int smem = NWARP * 2 * TILEB + (H * NP + DZ) * 4;
   static bool attr = false;
-  if (!attr) { C10_CUDA_CHECK(cudaFuncSetAttribute(pair_fwd_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, NWARP * 2 * TILEB + (H * 1028 + DZ) * 4)); attr = true; }
+  if (!attr) { C10_CUDA_CHECK(cudaFuncSetAttribute(pair_fwd_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, NWARP * 2 * TILEB + (H * 2052 + DZ) * 4)); attr = true; }
   pair_fwd_kernel<<<N, THREADS, smem, at::cuda::getCurrentCUDAStream()>>>(
       reinterpret_cast<const __nv_bfloat16*>(z.data_ptr<at::BFloat16>()), reinterpret_cast<const __nv_bfloat16*>(mask.data_ptr<at::BFloat16>()),
       ln_w.data_ptr<float>(), reinterpret_cast<const __nv_bfloat16*>(wb16.data_ptr<at::BFloat16>()),

{miniworld_engine-2.3.7 → miniworld_engine-2.3.9}/src/miniworld_engine/integrations/local_dit.py

@@ -657,16 +657,23 @@ def _device_of(mask):
 
 def _hoisted(store, tensor, key, build):
     """The table ``build()`` makes from ``tensor``, once per tensor (pointer, version, ``key``): the entry keeps a weak reference to the tensor, so
-    a new tensor that reuses a freed address is not served the old one's table, and entries of freed tensors are dropped."""
+    a new tensor that reuses a freed address is not served the old one's table, and entries of freed tensors are dropped.
+
+    A sample-expanded conditioning (stride 0 over the samples: a sampler step's view of the fold's one conditioning) is a NEW tensor object at
+    every call over the same small storage, so no weak reference to it outlives the call and the table would be remade at every call -- in a
+    captured step in every replay, for the whole life of the sampler. Such an entry holds the tensor itself (its storage is the fold's one
+    conditioning, which keeps its address from being reused while the entry stands); a view is otherwise named by the tensor it views."""
+    held = any(st == 0 and n > 1 for st, n in zip(tensor.stride(), tensor.shape, strict=True))
+    root = tensor._base if tensor._base is not None else tensor
     entry = _capture.lookup_inputs(
         store,
         (_tkey(tensor), key),
-        lambda: (weakref.ref(tensor), build()),
+        lambda: (tensor if held else weakref.ref(root), held, build()),
         limit=64,
-        valid=lambda e: e[0]() is tensor,
-        alive=lambda e: e[0]() is not None,
+        valid=lambda e: e[1] or e[0]() is root,
+        alive=lambda e: e[1] or e[0]() is not None,
     )
-    return entry[1]
+    return entry[2]
 
 
 def _tkey(t):

{miniworld_engine-2.3.7 → miniworld_engine-2.3.9}/src/miniworld_engine/integrations/local_dit_h100.py

@@ -23,6 +23,7 @@ F32 = torch.float32
 BF = torch.bfloat16
 _HOIST_COND: dict = {}
 _HOIST_BIAS: dict = {}
+_HOIST_EDGES: dict = {}  # (pair tensor, LN / bias weights) -> the slid edge windows' indices, key mask and pair bias
 
 
 def _hopper(t) -> bool:
@@ -128,32 +129,42 @@ def _af3_servable(module, single, cond, pair, mask) -> bool:
     )
 
 
-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."""
+def _edge_tables(z, gamma, wb, n, kmask, n0):
+    """What the slid edge windows read that does not change between the steps of a fold: the five window indices, their key positions and
+    validity, and their pair bias [5, h, 32, 128] (the pair LayerNorm and projection of those windows only). It depends on the real atom
+    count, the pair and the block's LN / bias weights, so it is built once per pair (``_HOIST_EDGES``) and read at every step."""
     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)
+        else torch.full((), n0, device=z.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]
+    keys = starts[:, None] + torch.arange(DS, device=z.device)  # [5, 128]
     valid = keys < n_real  # a system under 128 real atoms has keys past them
+    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]
+    mask = ~valid[None, :, None, None, :]
+    return idx, keys, mask, bias[None].contiguous()
+
+
+def _slid_edges(o, q, k, v, tables, a, n):
+    """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. ``tables``: ``_edge_tables``."""
+    idx, keys, mask, bias = tables
+    h, d = NH, DS // NH
+    nwin = n // 32
     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"))
+    logits = torch.einsum("awqhd,awkhd->awhqk", qs, ks) * d**-0.5 + bias
+    logits = logits.masked_fill(mask, 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))
@@ -287,7 +298,18 @@ def _fwd(
                 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)
+            def edges():
+                return _edge_tables(z, weights[_ld._GAMMA], weights[_ld._WBIAS], n, kmask, _n0)
+
+            _slid_edges(
+                o,
+                q,
+                k,
+                v,
+                _ld._hoisted(_HOIST_EDGES, pair, (wkey, n, mk), edges) if hoist else edges(),
+                a,
+                n,
+            )
         o = o.view(m, DS)
         if not save:
             return [

{miniworld_engine-2.3.7 → miniworld_engine-2.3.9}/src/miniworld_engine/integrations/pwa_train.py

@@ -48,6 +48,9 @@ _NATIVE_TRAINING_SCOPE: ContextVar[bool] = ContextVar(
 _SPLIT_MIN_N = 512
 
 
+NMAX_INFER = 2048  # longest sequence of the inference path (pair_fwd3's shared-memory row of logits)
+
+
 def _split_contraction(n: int) -> bool:
     """N >= 512: the contraction runs on cuBLAS (head-major o) and the gate / out pass in pwa_gateout; below, the fused kernel wins
     (the o round trip and the extra launch cost more than the faster contraction saves). o is rounded to bf16 BEFORE the gate here, as in
@@ -140,8 +143,10 @@ def refusal(
         native = native_aux_enabled(msa) or (
             not torch.is_grad_enabled() and native_inference_enabled(msa.device)
         )
-        if cap == (9, 0) and native and (n == 0 or s == 0 or n > 1024):
-            return f"the CUDA-only H100 path needs positive S and N <= 1024, got N={n}, S={s}"
+        # the forward pair kernel holds H rows of N logits in shared memory (N <= 2048); the backward one is built for N <= 1024
+        n_max = 1024 if torch.is_grad_enabled() else NMAX_INFER
+        if cap == (9, 0) and native and (n == 0 or s == 0 or n > n_max):
+            return f"the CUDA-only H100 path needs positive S and N <= {n_max}, got N={n}, S={s}"
         if wide:  # the wide path pads the length to a multiple of 16 itself; any depth
             if cap != (9, 0):
                 return "the wide kernels are built for sm_90a"
@@ -885,7 +890,7 @@ def _inference_math(
     else:
         lnz_w = module.ln_pair.weight.detach().float().contiguous()
         wb16 = module.to_bias.weight.detach()
-    if k["pair3"] is not None and n % 16 == 0 and n <= 1024:
+    if k["pair3"] is not None and n % 16 == 0 and n <= NMAX_INFER:
         w16 = k["pair3"].pair_fwd3(z, pm, lnz_w, eps_z, wb16)
     else:
         w16 = pair_fwd(z, pm, lnz_w, eps_z, module.to_bias.weight.detach(), H, BJ=64)