miniworld-engine 2.3.10__tar.gz → 2.3.11__tar.gz

{miniworld_engine-2.3.10 → miniworld_engine-2.3.11}/PKG-INFO

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

@@ -1,6 +1,6 @@
 [project]
 name = "miniworld-engine"
-version = "2.3.10"
+version = "2.3.11"
 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.10 → miniworld_engine-2.3.11}/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.9,d6d1eb084d32,clean,2026-10-24,v2.3.8-9-gd6d1eb08,
+#provenance,2.3.10,acfebaea42b4,clean,2026-10-24,v2.3.9-27-gacfebaea,
 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
@@ -11,6 +11,9 @@ augmented_attention_bwd_atomic_triton,triton,augmented_attention,miniworld_engin
 augmented_attention_bwd_pre_triton,triton,augmented_attention,miniworld_engine/kernels/augmented_attention/triton/main.py,ok,rel out=8.02e-08 (band 4e-07 declared),bf16
 augmented_attention_bwd_split_triton,triton,augmented_attention,miniworld_engine/kernels/augmented_attention/triton/main.py,ok,rel dq=4.04e-03 dk=3.25e-03 dv=4.70e-03 dbias=2.58e-03 (band 2e-02 declared),bf16
 augmented_attention_fwd_triton,triton,augmented_attention,miniworld_engine/kernels/augmented_attention/triton/main.py,ok,rel out=2.17e-03 m=9.38e-08 (band 9e-03 declared),bf16
+augmented_attention_fwd_window_triton,triton,augmented_attention,miniworld_engine/kernels/augmented_attention/triton/window.py,ok,rel out=2.29e-03 lse=1.23e-07 (band 1e-02 declared),bf16
+augmented_attention_bwd_window_triton,triton,augmented_attention,miniworld_engine/kernels/augmented_attention/triton/window.py,ok,rel dq=3.33e-03 dk=3.39e-03 dv=2.44e-03 dbias=2.12e-03 (band 2e-02 declared),bf16
+augmented_attention_bwd_reduce_window_triton,triton,augmented_attention,miniworld_engine/kernels/augmented_attention/triton/window.py,ok,rel out=8.28e-08 (band 7e-07 declared),bf16
 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
@@ -32,9 +35,9 @@ 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.78e-07 dbeta=3.34e-07 (band 1e-02 declared),bf16
+layernorm_bwd_atomic_strided_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/te_style.py,ok,rel dx=2.60e-03 dgamma=4.41e-07 dbeta=1.25e-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_foldstats_triton,triton,transition,miniworld_engine/kernels/transition/triton/fused.py,ok,rel dx=2.61e-03 dgamma=2.62e-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
 layernorm_bwd_split_cuda,cuda,layernorm,miniworld_engine/kernels/layernorm/cuda/layer_norm_cuda_kernel.cu,ok,rel dx=2.03e-03 dw=2.26e-03 db=2.41e-03 (band 1e-02 declared),bf16
 layernorm_bwd_split_mmajor_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/mmajor_bwd.py,ok,rel dx=2.60e-03 dgamma=3.07e-07 dbeta=1.04e-07 (band 1e-02 declared),bf16
@@ -49,25 +52,28 @@ layernorm_fwd_strided_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/
 layernorm_linear_bwd_fp32_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/pair_bias.py,ok,rel dx=2.30e-03 dln_weight=1.74e-03 dproj_weight=2.01e-03 (band 1e-02 declared),bf16
 layernorm_linear_fwd_fp32_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/pair_bias.py,ok,rel out=2.37e-03 mean=2.08e-08 rstd=1.74e-07 (band 9e-03 declared),bf16
 layernorm_linear_fwd_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/fused.py,ok,rel out=3.51e-03 (band 2e-02 declared),bf16
-layernorm_stats_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/stats.py,ok,rel rstd=8.69e-08 c1=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
+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 dmsa_pbias=4.23e-03 dln_weight_pbias=7.72e-03 dln_bias_pbias=2.32e-03 dw_left_pbias=5.16e-03 dw_right_pbias=4.23e-03 db_left_pbias=1.97e-03 db_right_pbias=3.41e-03 dmsa_grouped=4.30e-03 dln_weight_grouped=4.27e-03 dln_bias_grouped=2.00e-03 dw_left_grouped=5.39e-03 dw_right_grouped=4.50e-03 (band 1e-02 declared),bf16
 outer_product_mean_epilogue_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel out=2.88e-03 (band 1e-02 declared),bf16
-outer_product_mean_layernorm_gemm_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel out=2.88e-03 (band 1e-02 declared),bf16
-pair_weighted_averaging_bwd_gate_gemm_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel dmsa=2.84e-03 dw_gate=4.47e-03 dw_out=3.30e-03 (band 1e-02 declared),bf16
-pair_weighted_averaging_bwd_layernorm_gemm_dw_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel dw_value=4.19e-03 (band 1e-02 declared),bf16
-pair_weighted_averaging_bwd_layernorm_gemm_dx_dlnw_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel dmsa=2.84e-03 dln_msa_weight=4.39e-03 dln_msa_bias=4.52e-03 (band 1e-02 declared),bf16
-pair_weighted_averaging_bwd_layernorm_gemm_softmax_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel dpair=4.99e-03 dln_pair_weight=3.87e-03 dw_bias=4.26e-03 (band 1e-02 declared),bf16
-pair_weighted_averaging_gate_gemm_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel out=3.02e-03 (band 1e-02 declared),bf16
-pair_weighted_averaging_layernorm_gemm_softmax_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel out=3.02e-03 (band 1e-02 declared),bf16
-pair_weighted_averaging_layernorm_gemm_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel out=3.02e-03 (band 1e-02 declared),bf16
+outer_product_mean_bwd_epilogue_ln_dw_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel dw_out_grouped=4.30e-03 (band 2e-02 declared),bf16
+outer_product_mean_bwd_epilogue_ln_dx_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel dmsa_grouped=4.71e-03 db_out_grouped=2.05e-03 dresidual_grouped=0.00e+00 dproduct_weight_grouped=6.06e-03 dproduct_bias_grouped=2.72e-03 (band 3e-02 declared),bf16
+outer_product_mean_epilogue_ln_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel out_grouped=5.06e-03 (band 2e-02 declared),bf16
+outer_product_mean_layernorm_gemm_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel out=2.88e-03 out_pbias=2.82e-03 out_grouped=5.06e-03 (band 1e-02 declared),bf16
+pair_weighted_averaging_bwd_gate_gemm_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel dmsa=2.84e-03 dw_gate=4.47e-03 dw_out=3.30e-03 dmsa_c8=2.69e-03 dw_gate_c8=4.03e-03 dw_out_c8=4.06e-03 (band 1e-02 declared),bf16
+pair_weighted_averaging_bwd_layernorm_gemm_dw_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel dw_value=4.19e-03 dw_value_c8=3.40e-03 (band 1e-02 declared),bf16
+pair_weighted_averaging_bwd_layernorm_gemm_dx_dlnw_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel dmsa=2.84e-03 dln_msa_weight=4.39e-03 dln_msa_bias=4.52e-03 dmsa_c8=2.69e-03 dln_msa_weight_c8=5.20e-03 dln_msa_bias_c8=2.51e-03 (band 1e-02 declared),bf16
+pair_weighted_averaging_bwd_layernorm_gemm_softmax_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel dpair=4.99e-03 dln_pair_weight=3.87e-03 dw_bias=4.26e-03 dpair_c8=5.66e-03 dln_pair_weight_c8=6.34e-03 dw_bias_c8=7.50e-03 (band 1e-02 declared),bf16
+pair_weighted_averaging_gate_gemm_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel out=3.02e-03 out_c8=2.99e-03 (band 1e-02 declared),bf16
+pair_weighted_averaging_layernorm_gemm_softmax_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel out=3.02e-03 out_c8=2.99e-03 (band 1e-02 declared),bf16
+pair_weighted_averaging_layernorm_gemm_triton,triton,pair_weighted_averaging,miniworld_engine/kernels/pair_weighted_averaging/triton/main.py,ok,rel out=3.02e-03 out_c8=2.99e-03 (band 1e-02 declared),bf16
 rmsnorm_adamod_bwd_triton,triton,rmsnorm_adamod,miniworld_engine/kernels/rmsnorm/triton/main.py,ok,rel dq_plain=2.56e-03 dc_plain=4.44e-03 dWsc_plain=4.10e-03 dWsh_plain=1.76e-03 dWg_plain=4.17e-03 dq_aff=1.55e-03 dc_aff=4.22e-03 dWsc_aff=4.48e-03 dWsh_aff=1.76e-03 dWg_aff=4.17e-03 dweight_aff=0.00e+00 (band 1e-02 declared),bf16
 rmsnorm_adamod_fwd_triton,triton,rmsnorm_adamod,miniworld_engine/kernels/rmsnorm/triton/main.py,ok,rel y_plain=2.89e-03 gate_plain=1.94e-03 y_aff=1.34e-03 gate_aff=1.94e-03 (band 5e-03 declared),bf16
 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
@@ -99,7 +105,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=2.41e-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=8.05e-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
@@ -127,22 +133,22 @@ mpnn_edge_dropout_bwd_packed_triton,triton,mpnn_edge_dropout,miniworld_engine/ke
 trimul_fwd_sm90_cuda,cuda,trimul_inproj,miniworld_engine/kernels/trimul_inproj/cuda/h100_inference.py,ok,rel h128_dir0=4.58e-03 h128_dir1=4.64e-03 h128_dir2=4.63e-03 h256_dir0=5.06e-03 h256_dir1=4.88e-03 h256_dir2=5.34e-03 (band 2e-02 declared),bf16
 transition_fwd_residual_sm90_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_fused_fwd_sm90a_kernel.cu,ok,rel n4_y=3.57e-03 n4_xn=3.10e-03 n4_rstd=8.69e-08 n4_c1=6.69e-08 n2_y=3.33e-03 n2_xn=3.10e-03 n2_rstd=8.69e-08 n2_c1=6.69e-08 (band 1e-02 declared),bf16
 transition_bwd_residual_sm90_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_fused_bwd_sm90a_kernel.cu,ok,rel n4_dx=5.85e-03 n4_dgamma=4.31e-03 n4_dbeta=3.03e-03 n4_dWa=3.97e-03 n4_dWb=4.61e-03 n4_dWs=4.17e-03 n2_dx=6.01e-03 n2_dgamma=4.39e-03 n2_dbeta=3.81e-03 n2_dWa=3.92e-03 n2_dWb=4.27e-03 n2_dWs=3.68e-03 (band 2e-02 declared),bf16
-swa_dit_inproj_fwd_triton,triton,swa_dit,miniworld_engine/kernels/swa_dit/triton/forward.py,ok,rel out=5.94e-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.94e-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.94e-03 (band 5e-02 declared),bf16
+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.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.66e-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=5.72e-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=5.72e-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=5.72e-03 dw_gate=4.68e-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.94e-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=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=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.66e-03 dw_gate=4.68e-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=5.72e-03 dw_gate=4.68e-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

{miniworld_engine-2.3.10 → miniworld_engine-2.3.11}/src/miniworld_engine/build/bundle_coverage.py

@@ -147,6 +147,19 @@ def cases() -> list[dict]:
     # The Anthropic OPM path (integrations/anthropic_msa) finishes with THIS repo's csrc/opm_epilogue.cu, built once under the name
     # ``miniworld_opm_epilogue`` (no variant suffix: the ``_native*`` ones belong to integrations/opm_train). Needs no payload to build.
     out.append({"kind": "anthropic_ext", "train": False, "dtype": "bfloat16", "heavy": False})
+    # La-Proteina (kernels/pair_embed_geom, pair_embed_adaln; the token DiT block at its widths: 768 / 12 heads x 64 / pair 256 / cond 256 or 128 /
+    # transition x4): the pair builders once, and the batched-pairs block (block_psb under serves_batched) inference and training at the padded
+    # token lengths, B = 5 at 384 (the TOKEN_EMPTY_UNIFORM build) and a larger B elsewhere. H100 only: the others report a skip.
+    out.append({"kind": "la_proteina_pair", "train": True, "dtype": "bfloat16", "heavy": False})
+    for train in (False, True):
+        for length in (128, 256, 384, 512):
+            for d_cond in (256, 128):
+                for qk in (True, False):
+                    if not qk and d_cond != 256:
+                        continue
+                    for b in ((5, 2) if length == 384 else (2, 12)):
+                        out.append({"kind": "la_proteina_dit", "L": length, "train": train, "dtype": "bfloat16", "d_cond": d_cond,  # noqa: PERF401
+                                    "qk": qk, "batch": b, "heavy": False})
     for i, c in enumerate(out):
         c["id"] = i
     return out
@@ -320,6 +333,12 @@ def _run_block(c: dict) -> None:
     elif c["kind"] == "model_variants":
         _model_variants()
         return
+    elif c["kind"] == "la_proteina_pair":
+        _la_proteina_pair()
+        return
+    elif c["kind"] == "la_proteina_dit":
+        _la_proteina_dit(c)
+        return
     elif c["kind"] == "sampler_ops":
         _sampler_ops()
         return
@@ -346,6 +365,65 @@ def _run_block(c: dict) -> None:
     _step(model, args, c["train"], autocast)
 
 
+def _la_proteina_pair() -> None:
+    """La-Proteina's pair builders, every extension variant once: the encoder pair (inference; training with and without the folded LayerNorm), the
+    decoder pair (training, with and without it) with the table-gradient backward, and the initial pair builder (self-conditioning on / off)."""
+    import torch
+
+    from miniworld_engine.kernels.pair_embed_adaln import pair_embed_adaln
+    from miniworld_engine.kernels.pair_embed_geom import pair_embed_ca, pair_embed_geom, pair_embed_geom_train
+    from miniworld_engine.kernels.pair_embed_geom.cuda import sm90 as geom
+
+    if torch.cuda.get_device_capability() != (9, 0):
+        raise RuntimeError("La-Proteina's pair builders are H100 (sm_90a) kernels")
+    dev = "cuda"
+    for b, n in ((48, 128), (2, 512)):
+        bb = torch.randn(b, n, 4, 3, device=dev) * 0.5
+        flags = [torch.ones(b, n, device=dev, dtype=torch.bool) for _ in range(3)]
+        idx, chains = torch.arange(1, n + 1, device=dev).expand(b, n), torch.zeros(b, n, device=dev, dtype=torch.int32)
+        w = torch.randn(256, 317, device=dev, requires_grad=True)
+        wc = torch.randn(256, 157, device=dev, requires_grad=True)
+        ca = bb[:, :, 1].contiguous()
+        with torch.no_grad():
+            pair_embed_geom(w.detach(), bb, flags[0], flags[1], flags[2], idx, chains)
+        for eps in (None, 1e-5):
+            pair_embed_geom_train(w, bb, flags[0], flags[1], flags[2], idx, chains, ln_eps=eps).float().sum().backward()
+            pair_embed_ca(wc, ca, flags[2], ln_eps=eps).float().sum().backward()
+        for impl in (0, 1):   # the shared-memory read-modify-write backward is the A/B switch MINIWORLD_PAIR_EMBED_BWD_IMPL=0
+            pair, codes = geom.forward(w.detach().t().to(torch.bfloat16).contiguous(), bb, flags[0].to(torch.uint8), flags[1].to(torch.uint8),
+                                       flags[2].to(torch.uint8), idx.to(torch.int32), chains, torch.linspace(0.1, 2, 20, device=dev),
+                                       torch.linspace(-torch.pi, torch.pi, 20, device=dev), with_codes=True)
+            geom.backward(pair, codes, 0, impl=impl)
+        ln_w, ln_b = torch.ones(256, device=dev, requires_grad=True), torch.zeros(256, device=dev, requires_grad=True)
+        wa = torch.randn(256, 187, device=dev, requires_grad=True)
+        gamma, beta = torch.rand(b, 256, device=dev, requires_grad=True), torch.randn(b, 256, device=dev, requires_grad=True)
+        for sc in (None, ca):
+            pair_embed_adaln(wa, ln_w, ln_b, gamma, beta, ca, sc, idx, flags[2]).float().sum().backward()
+    torch.cuda.synchronize()
+
+
+def _la_proteina_dit(c: dict) -> None:
+    """The token DiT block at La-Proteina's widths over ``batch`` independent proteins, each with its own pair (masked), bf16."""
+    import torch
+
+    from miniworld_engine.modules.dit import DiTBlock
+    from miniworld_engine.modules.exceptions import ImplementationType as MI
+
+    b, n, d_cond, train = c["batch"], c["L"], c["d_cond"], c["train"]
+    model = DiTBlock(d_single=768, d_cond=d_cond, d_pair=256, n_head=12, n=4, use_qk_norm=c["qk"], implementation=MI.MINIWORLD).cuda().bfloat16()
+    _randomize(model)
+    model.train(train)
+    single = torch.randn(1, b, n, 768, device="cuda", dtype=torch.bfloat16, requires_grad=train)
+    cond = torch.randn(1, b, n, d_cond, device="cuda", dtype=torch.bfloat16, requires_grad=train)
+    pair = torch.randn(b, n, n, 256, device="cuda", dtype=torch.bfloat16, requires_grad=train)
+    mask = torch.arange(n, device="cuda")[None] < torch.randint(n // 2, n + 1, (b, 1), device="cuda")
+    with torch.set_grad_enabled(train):
+        out = model(single, cond, pair, mask)
+        if train:
+            out.float().sum().backward()
+    torch.cuda.synchronize()
+
+
 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

{miniworld_engine-2.3.10 → miniworld_engine-2.3.11}/src/miniworld_engine/integrations/token_dit_train.py

@@ -473,6 +473,8 @@ def _pack(P, qk, dev, at):
             d_ // P["attention.to_bias.weight"].shape[0],
             rpb=_native_rpb(),
             dc=P["attention.ada_ln_in.to_scale.weight"].shape[1],
+            n=P["transition.expand_a.weight"].shape[0] // d_,
+            dp=P["attention.ln_pair.weight"].shape[0],
         )
         weights = [P[n] for n in names]
         W = dict(
@@ -598,7 +600,13 @@ def _own_cond(cond, fp32) -> bool:
     return fp32 and cond.dtype is torch.float32 and not _native90(cond.device)
 
 
-def _saved_like(single, cond, pair, fp32, qk=True, H=H, psb=False):
+def _nd(params, qk, D):
+    """(transition expansion n, pair channels of ln_pair) read off the block's parameters: expand_a is [n D, D], ln_pair [dp]."""
+    P = dict(zip(_names_for(qk), params, strict=False))
+    return P["transition.expand_a.weight"].shape[0] // D, P["attention.ln_pair.weight"].shape[0]
+
+
+def _saved_like(single, cond, pair, fp32, qk=True, H=H, psb=False, n=2, dp=DP):
     """Shapes / dtypes of the forward's saved activations (the fake implementation and the contract of _fwd). The fp32 path
     saves no copy of cond (_own_cond), no copy of v (QK-norm: the core reads it from qkvg) and no LN(pair) (pair_bias_bwd
     recomputes it): those slots are empty [0, n] tensors. H100: q / k / v are always copies (the q pass scales q; the TF32 core's
@@ -636,17 +644,18 @@ def _saved_like(single, cond, pair, fp32, qk=True, H=H, psb=False):
         e((M, D), dtype=f32),
         e((M, 2), dtype=f32),
         e((M, D), dtype=at),
-        e((M, 4 * D), dtype=at),
+        e((M, 2 * n * D), dtype=at),   # a | b of the SwiGLU: two [M, n D] halves
         e((M, D), dtype=at),
-        e((0 if (psb or (fp32 and not h100)) else R, DP), dtype=at),
+        e((0 if (psb or (fp32 and not h100)) else R, dp), dtype=at),
         e((0 if psb else R, 2), dtype=f32),
     ]
 
 
 def _fwd_fake(single, cond, pair, mask, params, qk, eq, ek, fp32, heads, psb=False):
+    n, dp = _nd(params, qk, single.shape[-1])
     return [
         torch.empty_like(single),
-        *_saved_like(single, cond, pair, fp32, qk, heads, psb),
+        *_saved_like(single, cond, pair, fp32, qk, heads, psb, n, dp),
     ]
 
 
@@ -680,12 +689,15 @@ def _fwd_body(single, cond, pair, mask, params, qk, eq, ek, fp32, H, psb=False):
 
     A, _, L, D = single.shape
     DH = D // H
+    nexp, dpw = _nd(params, qk, D)
     T = ext(
         D,
         DH,
         rpb=_native_rpb() if _native90(single.device) else 8,
         empty_uniform=_empty_uniform90(single.device, single.shape[0], single.shape[2]),
         dc=cond.shape[-1],
+        n=nexp,
+        dp=dpw,
     )
     names = _names_for(qk)
     P = dict(zip(names, params, strict=False))
@@ -739,21 +751,21 @@ def _fwd_body(single, cond, pair, mask, params, qk, eq, ek, fp32, H, psb=False):
         # one bias per sample, produced by the caller (head-major [H, A, L, L], log2 units): there is no pair LayerNorm and no
         # bias GEMM in the block, and nothing of either is kept for the backward
         assert h100 and not fp32 and pair.shape == (H, A, L, L), "per-sample bias: the H100 bf16 block"
-        ph = pair.new_empty(0, DP, dtype=at)
+        ph = pair.new_empty(0, dpw, dtype=at)
         pst = pair.new_empty(0, 2, dtype=torch.float32)
         bias = pair
     else:
-        p2 = pair.reshape(R, DP)
+        p2 = pair.reshape(R, dpw)
         pst = torch.empty(R, 2, device=dev)
     if psb:
         pass
     elif h100 or not fp32:
-        ph = torch.empty(R, DP, device=dev, dtype=at)
+        ph = torch.empty(R, dpw, device=dev, dtype=at)
         T.pair_ln(p2, ph, pst, EPS)
         # head-major bias, Wf = Wb diag(wp): natural units for sm_100a, log2 units (and key-permuted for the TF32 core) for sm_90a
         bias = mm(W["Wf_l2"] if h100 else W["Wf_at"], ph.t()).view(H, L, L)
     else:  # B200 fp32: LN(pair) folded straight into the bias, fp32 FMAs, key mask as -inf; LN(pair) not kept
-        ph = pair.new_empty(0, DP, dtype=at)
+        ph = pair.new_empty(0, dpw, dtype=at)
         bias = torch.empty(H, L, L, device=dev)
         T.pair_bias(
             p2,
@@ -811,13 +823,13 @@ def _fwd_body(single, cond, pair, mask, params, qk, eq, ek, fp32, H, psb=False):
     x1st = torch.empty(M, 2, device=dev)
     T.res_adaln_b(x, y, Gg, bg1, G, bs2, x1, xt, x1st, EPS)
     Wab = W["Wab"]
-    h = torch.empty(M, 2 * D, device=dev, dtype=at)
+    h = torch.empty(M, nexp * D, device=dev, dtype=at)
     gsw = _swiglu_gemm(xt, Wab)
     if gsw is not None:  # expand GEMM + SwiGLU, a | b saved on the way
-        ab = torch.empty(M, 4 * D, device=dev, dtype=at)
+        ab = torch.empty(M, 2 * nexp * D, device=dev, dtype=at)
         gsw(xt, Wab, h, ab)
     else:
-        ab = mm(xt, Wab.t())  # [M, 2 * 2D] a | b
+        ab = mm(xt, Wab.t())  # [M, 2 * n D] a | b
         rows.swiglu_rows(ab, h)
     Wsq = W["Wsq"]
     z = mm(h, Wsq.t())
@@ -945,7 +957,7 @@ def _native_attn_bwd90(T, qs, k, v, dob, bias, mask, lse, dd, A, L, fp32):
         return dq, dk, dv, dbn.to(torch.bfloat16)
     else:
         sm90._ext(
-            "attn_dq", (*(() if L % 192 == 0 else ("NWG=2",)), *bf_defs, *empty_defs)
+            "attn_dq", (*(() if L % 192 == 0 else ("NWG=2",)), *sm90.dq_defs(H, DH), *empty_defs)
         ).attn_dq(qs, k, v, dob, bias, km, lse, dd, dq)
         sm90._ext("attn_dkv", ("DB_PER_SAMPLE=1", *bf_defs, *empty_defs)).attn_dkv(
             qs, k, v, dob, bias, km, lse, dd, dk, dv, db
@@ -983,8 +995,8 @@ def _native_final90(
         dWn[3 * D :],
         dWg[D:],
         allp,
-        dWab[: 2 * D],
-        dWab[2 * D :],
+        dWab[: dWab.shape[0] // 2],
+        dWab[dWab.shape[0] // 2 :],
         dWsq,
     ]
     if qk:
@@ -1056,12 +1068,15 @@ def _bwd_body(single, cond, pair, mask, params, saved, dout, qk, eq, ek, fp32, H
 
     D = single.shape[-1]
     DH = D // H
+    nexp, dpw = _nd(params, qk, D)
     T = ext(
         D,
         DH,
         rpb=_native_rpb() if _native90(single.device) else 8,
         empty_uniform=_empty_uniform90(single.device, single.shape[0], single.shape[2]),
         dc=cond.shape[-1],
+        n=nexp,
+        dp=dpw,
     )
     h100 = _sm90(single.device)
     saved = list(saved)
@@ -1126,7 +1141,7 @@ def _bwd_body(single, cond, pair, mask, params, saved, dout, qk, eq, ek, fp32, H
     native = _native90(dev)
     c2, p2 = (
         (None if native else cond.reshape(M, cond.shape[-1]).contiguous()),
-        None if psb else pair.reshape(R, DP),
+        None if psb else pair.reshape(R, dpw),
     )  # native: cond_bwd90 reads cond as given
     if psb:
         bias = pair  # the bias is an input here, not a saved activation
@@ -1139,7 +1154,7 @@ def _bwd_body(single, cond, pair, mask, params, saved, dout, qk, eq, ek, fp32, H
     dGg = torch.empty(M, 2 * D, device=dev, dtype=at)
     dz = torch.empty(M, D, device=dev, dtype=at)
     T.res_c_bwd(dout, z, Gg, bg2, dz, dGg, part[0])
-    dh = mm(dz, Wsq)  # [M, 2D]
+    dh = mm(dz, Wsq)  # [M, n D]
     dab = torch.empty_like(ab)
     h = torch.empty_like(dh)
     T.swiglu_bwd(dh, ab, dab, h)
@@ -1218,12 +1233,12 @@ def _bwd_body(single, cond, pair, mask, params, saved, dout, qk, eq, ek, fp32, H
         T.cond_bwd90(dchat, dcg, cond, cst, dc)
         if psb:
             # the bias projection's weights belong to the producer of the biases, which differentiates them itself
-            dWf = torch.zeros(H, DP, device=dev)
+            dWf = torch.zeros(H, dpw, device=dev)
         else:
             dbv = DB.view(H, R)
             dWf = mm32(dbv, ph)
             dph = mm32(dbv.t(), W["Wf_at"])
-            dpair = torch.empty(R, DP, device=dev, dtype=pdtype)
+            dpair = torch.empty(R, dpw, device=dev, dtype=pdtype)
             T.pair_ln_bwd(dph, p2, pst, dpair)
         outs = _native_final90(
             T, names, params, W, allp, dWn, dWf, dWqkvg, dWo, dWg, dWab, dWsq, qk
@@ -1231,7 +1246,7 @@ def _bwd_body(single, cond, pair, mask, params, saved, dout, qk, eq, ek, fp32, H
         return [
             dx.view(A, 1, L, D),
             dc.view(cond.shape),
-            DB.view(pair.shape) if psb else dpair.view(1, L, L, DP),
+            DB.view(pair.shape) if psb else dpair.view(1, L, L, dpw),
             *outs,
         ]
     T.cond_bwd(dchat, dcg, c2, cst, dc)
@@ -1239,19 +1254,19 @@ def _bwd_body(single, cond, pair, mask, params, saved, dout, qk, eq, ek, fp32, H
     dw12 = torch.zeros(2, cond.shape[-1], device=dev)
     Ws1, Wb1, Ws2, Wb2 = (Wraw[i * D : (i + 1) * D] for i in range(4))
     T.unfold_lnw(dWn, Ws1, Wb1, Ws2, Wb2, w1, w2, dWu, dw12)
-    dpair = torch.empty(R, DP, device=dev, dtype=pdtype)
+    dpair = torch.empty(R, dpw, device=dev, dtype=pdtype)
     if (
         fp32 and not h100
     ):  # B200 fp32: LN(pair) recomputed: dpair and the dWf partials in one pass
-        pwf = torch.empty(T.pair_grid(R), H * DP, device=dev)
+        pwf = torch.empty(T.pair_grid(R), H * dpw, device=dev)
         T.pair_bias_bwd(DB.view(H, R), p2, pst, Wf, dpair, pwf)
-        dWf = pwf.sum(0).view(H, DP)
+        dWf = pwf.sum(0).view(H, dpw)
     else:
         dbv = DB.reshape(H, R).to(
             at
         )  # masked key columns are 0 (P = 0); H100: may be a transposed view
-        dWf = mm32(dbv, ph)  # [H, DP]
-        dph = mm32(dbv.t(), Wf.to(at))  # [R, DP]
+        dWf = mm32(dbv, ph)  # [H, dp]
+        dph = mm32(dbv.t(), Wf.to(at))  # [R, dp]
         T.pair_ln_bwd(dph, p2, pst, dpair)
     dWs1, dWb1, dWs2, dWb2 = dWu.split(D)
     grads = {
@@ -1277,8 +1292,8 @@ def _bwd_body(single, cond, pair, mask, params, saved, dout, qk, eq, ek, fp32, H
         "transition.ada_ln_in.to_bias.weight": dWb2,
         "transition.to_scale.weight": dWg[D:],
         "transition.to_scale.bias": pb[0],
-        "transition.expand_a.weight": dWab[: 2 * D],
-        "transition.expand_b.weight": dWab[2 * D :],
+        "transition.expand_a.weight": dWab[: nexp * D],
+        "transition.expand_b.weight": dWab[nexp * D :],
         "transition.squeeze.weight": dWsq,
     }
     # custom-op outputs may not alias each other: the gradients above are views of a few shared buffers, so each leaves as
@@ -1287,7 +1302,7 @@ def _bwd_body(single, cond, pair, mask, params, saved, dout, qk, eq, ek, fp32, H
     torch._foreach_copy_(
         outs, [grads[n].reshape(p.shape) for n, p in zip(names, params, strict=False)]
     )
-    return [dx.view(A, 1, L, D), dc.view(cond.shape), dpair.view(1, L, L, DP), *outs]
+    return [dx.view(A, 1, L, D), dc.view(cond.shape), dpair.view(1, L, L, dpw), *outs]
 
 
 class _Block(torch.autograd.Function):
@@ -1387,10 +1402,13 @@ def block(module, single, cond, pair, mask=None, compute_dtype=None):
     return out[:A, :, :L] if padded else out
 
 
-def serves_psb(module, single, cond, bias, mask, compute_dtype=None) -> bool:
-    """Can ``block_psb`` run this call: the H100 bf16 block with an attention bias of its own per sample, bias [H, A, L, L]
-    (log2 units, as the block's bias GEMM leaves it). Any L (``block_psb`` pads to the kernels' tile), one batch row. Without
-    autograd it runs the same forward (the saved activations are freed on return)."""
+PSB_COND = (128, 256, 384, 768)  # conditioning widths the row kernels build for (a multiple of 128)
+
+
+def _serves_psb_core(module, single, cond, mask, compute_dtype=None) -> bool:
+    """Everything ``serves_psb`` asks except the bias tensor itself: the settings, the H100 bf16 block, one batch row, the widths
+    (cond in ``PSB_COND``, transition n = 2 or 4, pair channels of ``ln_pair`` 128 or 256, a head layout of LAYOUTS), LayerNorm eps,
+    and a key mask that is one [1, L] row or a row per sample [A, L]."""
     if not settings.current().token_dit_train:
         return False
     if settings.current().force_triton_kernels:
@@ -1405,13 +1423,13 @@ def serves_psb(module, single, cond, bias, mask, compute_dtype=None) -> bool:
     if single.ndim != 4 or single.shape[1] != 1 or single.shape[2] == 0:
         return False
     A, _, L, d = single.shape
-    if tuple(cond.shape) not in ((A, 1, L, 384), (A, 1, L, 768)) or tuple(bias.shape) != (a.n_head, A, L, L):
+    if tuple(cond.shape[:3]) != (A, 1, L) or cond.shape[3] not in PSB_COND:
         return False
-    if bias.dtype is not BF or not bias.is_contiguous():
+    if (a.n_head, d) not in LAYOUTS or module.transition.expand_a.weight.shape[0] not in (2 * d, 4 * d):
         return False
-    if (a.n_head, d) not in LAYOUTS or module.transition.expand_a.weight.shape[0] != 2 * d:
+    if a.ln_pair.weight.numel() not in (128, 256):
         return False
-    if mask is not None and not (mask.ndim == 2 and tuple(mask.shape) == (1, L)):
+    if mask is not None and not (mask.ndim == 2 and tuple(mask.shape) in ((1, L), (A, L))):
         return False
     norms = (
         a.ada_ln_in.ln_in,
@@ -1422,6 +1440,16 @@ def serves_psb(module, single, cond, bias, mask, compute_dtype=None) -> bool:
     return all(n.eps == EPS for n in norms)
 
 
+def serves_psb(module, single, cond, bias, mask, compute_dtype=None) -> bool:
+    """Can ``block_psb`` run this call: the H100 bf16 block with an attention bias of its own per sample, bias [H, A, L, L]
+    (log2 units, as the block's bias GEMM leaves it). Any L (``block_psb`` pads to the kernels' tile), one batch row. Without
+    autograd it runs the same forward (the saved activations are freed on return)."""
+    if not _serves_psb_core(module, single, cond, mask, compute_dtype):
+        return False
+    A, _, L, _ = single.shape
+    return tuple(bias.shape) == (module.attention.n_head, A, L, L) and bias.dtype is BF and bias.is_contiguous()
+
+
 def block_psb(module, single, cond, bias, mask=None, compute_dtype=None):
     """One DiTBlock whose attention bias is its own per sample: bias [H, A, L, L] bf16 in log2 units (what a stacked bias GEMM
     leaves), the key mask folded in as an additive row, not into the bias. The block differentiates everything but the bias's
@@ -1438,6 +1466,7 @@ def block_psb(module, single, cond, bias, mask=None, compute_dtype=None):
     Lp = -(-L // QM) * QM
     if mask is None:
         mask = torch.ones(1, L, device=single.device, dtype=torch.bool)
+    mask = mask.to(torch.bool).contiguous()
     if Lp != L:
         F = torch.nn.functional
         rows = (0, 0, 0, Lp - L)
@@ -1450,4 +1479,47 @@ def block_psb(module, single, cond, bias, mask=None, compute_dtype=None):
     return out[:, :, :L] if Lp != L else out
 
 
-__all__ = ["block", "block_psb", "serves", "serves_psb"]
+def _batched_mask(mask, B, L):
+    """A [B, L] / [A=1, B, L] key mask of B independent proteins as the [B, L] row per sample ``block_psb`` takes, or None."""
+    if mask is None:
+        return None
+    if mask.ndim == 3 and mask.shape[0] == 1:
+        mask = mask[0]
+    return mask if mask.ndim == 2 and tuple(mask.shape) == (B, L) else False
+
+
+def serves_batched(module, single, cond, pair, mask, compute_dtype=None) -> bool:
+    """A DiTBlock call over B > 1 independent proteins, each with its own pair: ``single`` [1, B, L, d], ``cond`` [1, B, L, dc],
+    ``pair`` [B, L, L, dp], a key mask [B, L] (or none). They run as B samples of ``block_psb``, each with the attention bias its own
+    pair gives."""
+    if single.ndim != 4 or single.shape[0] != 1 or single.shape[1] < 2:
+        return False
+    B, L = single.shape[1], single.shape[2]
+    key = _batched_mask(mask, B, L)
+    if key is False or cond.ndim != 4 or tuple(cond.shape[:3]) != (1, B, L):
+        return False
+    if pair.ndim != 4 or tuple(pair.shape[:3]) != (B, L, L) or pair.shape[3] != module.attention.ln_pair.weight.numel():
+        return False
+    return _serves_psb_core(module, single.reshape(B, 1, L, -1), cond.reshape(B, 1, L, -1), key, compute_dtype)
+
+
+def block_batched(module, single, cond, pair, mask=None, compute_dtype=None):
+    """The fused block over B independent proteins (see ``serves_batched``). Each protein's attention bias is read off its pair by
+    differentiable torch ops -- LayerNorm, the head projection, log2 units, head-major bf16 -- so the pair, ``ln_pair`` and
+    ``to_bias`` get their gradients from autograd, and everything else from the fused block."""
+    import torch.nn.functional as F
+
+    a = module.attention
+    B, L = single.shape[1], single.shape[2]
+    dp = pair.shape[-1]
+    bias = F.linear(
+        F.layer_norm(pair, (dp,), a.ln_pair.weight.to(pair.dtype), None, a.ln_pair.eps), a.to_bias.weight.to(pair.dtype)
+    )
+    bias = (bias * LOG2E).permute(3, 0, 1, 2).contiguous().to(BF)  # [H, B, L, L]
+    out = block_psb(
+        module, single.reshape(B, 1, L, -1), cond.reshape(B, 1, L, -1), bias, _batched_mask(mask, B, L), compute_dtype
+    )
+    return out.reshape(1, B, L, -1)
+
+
+__all__ = ["block", "block_batched", "block_psb", "serves", "serves_batched", "serves_psb"]