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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: miniworld-engine
-Version: 2.3.0
+Version: 2.3.2
 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
@@ -17,6 +17,7 @@ Requires-Dist: einops
 Requires-Dist: jaxtyping
 Requires-Dist: numpy
 Requires-Dist: scipy
+Requires-Dist: pydantic
 Provides-Extra: baselines
 Requires-Dist: cuequivariance-torch<0.13,>=0.12; extra == "baselines"
 Provides-Extra: bench
@@ -46,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.2.0** — CUDA + Triton only (CuTe/quack removed), torch 2.13 / cu129,
+**Version 2.3.2** — 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).
 
@@ -119,6 +120,8 @@ PyPI one, CUDA 13.0, and the cu129 one) and pin `torch==2.13.*`, so the first ca
 
 ```bash
 pip install miniworld-engine     # with the torch it pins; an sdist / git install keeps the JIT path
+# torch's CUDA 12.9 build instead (NVIDIA driver 575-579, which cannot run CUDA 13):
+pip install miniworld-engine --index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://pypi.org/simple
 ```
 
 Shipped kernels that do not match the running torch / CUDA or GPU are ignored, and anything they do not cover still

{miniworld_engine-2.3.0 → miniworld_engine-2.3.2}/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.2.0** — CUDA + Triton only (CuTe/quack removed), torch 2.13 / cu129,
+**Version 2.3.2** — 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).
 
@@ -82,6 +82,8 @@ PyPI one, CUDA 13.0, and the cu129 one) and pin `torch==2.13.*`, so the first ca
 
 ```bash
 pip install miniworld-engine     # with the torch it pins; an sdist / git install keeps the JIT path
+# torch's CUDA 12.9 build instead (NVIDIA driver 575-579, which cannot run CUDA 13):
+pip install miniworld-engine --index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://pypi.org/simple
 ```
 
 Shipped kernels that do not match the running torch / CUDA or GPU are ignored, and anything they do not cover still

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

@@ -1,6 +1,6 @@
 [project]
 name = "miniworld-engine"
-version = "2.3.0"
+version = "2.3.2"
 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, ...)."
@@ -24,6 +24,7 @@ dependencies = [
   "jaxtyping",
   "numpy",
   "scipy",  # module parameter initialization uses scipy.stats.truncnorm
+  "pydantic",  # miniworld_engine.architecture (block / diffusion configs are pydantic models)
 ]
 
 [project.optional-dependencies]

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

@@ -1,12 +1,12 @@
 kernel,backend,family,file,status,detail,dtype
-#provenance,2.2.0,e0ada497fd78,clean,2026-10-17,archive/norm-cuda-20260930-310-ge0ada497,
-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
-adaln_fwd_gate_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/training.py,ok,rel Y=2.14e-03 Gate=1.95e-03 (band 1e-02 declared),bf16
+#provenance,2.3.1,813f5c1080b0,clean,2026-10-20,archive/norm-cuda-20260930-335-g813f5c10,
+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 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
 augmented_attention_bwd_atomic_triton,triton,augmented_attention,miniworld_engine/kernels/augmented_attention/triton/memory_efficient.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_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
@@ -14,25 +14,25 @@ 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
-gated_projection_bwd_gate_recompute_flat_triton,triton,triangle_multiplication,miniworld_engine/modules/triangle_multiplication/baseline_dtv1.py,ok,rel d_gate=3.02e-03 d_proj=2.73e-03 (band 1e-02 declared),bf16
+gated_projection_bwd_gate_recompute_flat_triton,triton,triangle_multiplication,miniworld_engine/modules/triangle_multiplication/baseline_dtv1.py,failed,ModuleNotFoundError: No module named 'cuequivariance_ops',bf16
 gated_projection_bwd_gate_triton,triton,gated_projection,miniworld_engine/kernels/gated_projection/triton/main.py,ok,rel dgate=2.69e-03 drep=2.22e-03 (band 1e-02 declared),bf16
 gated_projection_gate_dropres_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/gate_elem.py,ok,rel y=2.94e-03 gate=2.73e-03 (band 1e-02 declared),bf16
 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=5.04e-07 dbeta=2.50e-07 (band 1e-02 declared),bf16
+layernorm_bwd_atomic_strided_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/te_style.py,ok,rel dx=2.60e-03 dgamma=4.72e-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,11 +45,11 @@ layernorm_fwd_recompute_noaffine_triton,triton,layernorm_linear,miniworld_engine
 layernorm_fwd_recompute_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/recompute.py,ok,rel out=2.11e-03 (band 1e-02 declared),bf16
 layernorm_fwd_saveact_strided_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/te_style.py,ok,rel x_normed=2.10e-03 mean=4.71e-08 rstd=1.82e-07 (band 1e-02 declared),bf16
 layernorm_fwd_saveact_triton,triton,layernorm,miniworld_engine/kernels/layernorm/triton/main.py,ok,rel out=2.11e-03 (band 1e-02 declared),bf16
-layernorm_fwd_strided_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/ln_strided.py,ok,rel out=2.56e-03 (band 9e-03 declared),bf16
+layernorm_fwd_strided_triton,triton,adaln,miniworld_engine/kernels/adaln/triton/ln_strided.py,ok,rel out=2.64e-03 (band 9e-03 declared),bf16
 layernorm_linear_bwd_fp32_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/pair_bias.py,ok,rel dx=2.30e-03 dln_weight=1.74e-03 dproj_weight=2.01e-03 (band 1e-02 declared),bf16
 layernorm_linear_fwd_fp32_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/pair_bias.py,ok,rel out=2.37e-03 mean=2.08e-08 rstd=1.74e-07 (band 9e-03 declared),bf16
 layernorm_linear_fwd_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/fused.py,ok,rel out=3.51e-03 (band 2e-02 declared),bf16
-layernorm_stats_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/stats.py,ok,rel rstd=8.69e-08 c1=8.17e-08 (band 6e-07 declared),bf16
+layernorm_stats_triton,triton,layernorm_linear,miniworld_engine/kernels/layernorm_linear/triton/stats.py,ok,rel rstd=8.69e-08 c1=1.63e-07 (band 6e-07 declared),bf16
 outer_product_mean_bwd_epilogue_dw_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel dw_out=4.55e-03 (band 1e-02 declared),bf16
 outer_product_mean_bwd_epilogue_dx_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel dmsa=5.84e-03 db_out=2.05e-03 dresidual=0.00e+00 (band 1e-02 declared),bf16
 outer_product_mean_bwd_layernorm_gemm_triton,triton,outer_product_mean,miniworld_engine/kernels/outer_product_mean/triton/main.py,ok,rel dmsa=5.84e-03 dln_weight=4.26e-03 dln_bias=3.54e-03 dw_left=4.97e-03 dw_right=6.01e-03 (band 1e-02 declared),bf16
@@ -67,14 +67,14 @@ 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,ok,rel dx=4.67e-03 dwa=3.88e-03 dwb=3.58e-03 dws=4.18e-03 (band 2e-02 declared),bf16
+transition_bwd_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,failed,RuntimeError: Error building extension 'transition_cuda_ext_v2': [1/2] c++ -MMD -MF transition_cuda.o.d -DTORCH_EXTENSION_NAME=transit | c++: error: unrecognized command line option ‘-std=c++20’; did you mean ‘-std=c++2a’?,bf16
 transition_bwd_swiglu_recompute_triton,triton,transition,miniworld_engine/kernels/transition/triton/fused.py,ok,rel h=2.18e-03 dAB=2.12e-03 xn=7.02e-04 (band 1e-02 declared),bf16
-transition_cast_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,ok,rel y=4.13e-03 (band 1e-02 declared),bf16
+transition_cast_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,failed,ImportError: /home/psk6950/miniworld-engine-scratch/jit-runall2/miniworld_engine_jit/torch_extensions/transition_cuda_ext_v2-ee6c7a7d,bf16
 transition_expand_swiglu_triton,triton,transition,miniworld_engine/kernels/transition/triton/main.py,ok,rel out=2.81e-03 (band 1e-02 declared),bf16
 transition_fwd_b2b_ktiled_triton,triton,transition,miniworld_engine/kernels/transition/triton/fused.py,ok,rel out=1.96e-03 (band 9e-03 declared),bf16
 transition_fwd_b2b_triton,triton,transition,miniworld_engine/kernels/transition/triton/fused.py,ok,rel out=2.48e-03 (band 1e-02 declared),bf16
 transition_layernorm_expand_swiglu_triton,triton,transition,miniworld_engine/kernels/transition/triton/fused.py,ok,rel out=2.18e-03 (band 1e-02 declared),bf16
-transition_swiglu_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,ok,rel y=4.13e-03 (band 2e-02 declared),bf16
+transition_swiglu_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_cuda_kernel.cu,failed,ImportError: /home/psk6950/miniworld-engine-scratch/jit-runall2/miniworld_engine_jit/torch_extensions/transition_cuda_ext_v2-ee6c7a7d,bf16
 triangle_attention_bwd_atomic_triton,triton,triangle_attention,miniworld_engine/kernels/triangle_attention/triton/atomic.py,ok,rel dq=1.54e-02 dk=1.29e-02 dv=4.85e-03 dbias=4.05e-03 (band 5e-02),bf16
 triangle_attention_bwd_dkdv_triton,triton,triangle_attention,miniworld_engine/kernels/triangle_attention/triton/main.py,ok,rel dk=4.44e-03 dv=2.76e-03 dbias=3.27e-03 (band 2e-02 declared),bf16
 triangle_attention_bwd_dq_triton,triton,triangle_attention,miniworld_engine/kernels/triangle_attention/triton/main.py,ok,rel dq=6.77e-03 (band 2e-02 declared),bf16
@@ -86,10 +86,10 @@ trimul_bwd_gate_packed_recompute_triton,triton,trimul_inproj,miniworld_engine/ke
 trimul_bwd_gate_packed_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/back_fused.py,ok,rel d_gLlog=3.04e-03 d_pL=1.96e-03 d_gRlog=2.97e-03 d_pR=2.02e-03 (band 1e-02 declared),bf16
 trimul_bwd_gate_recompute_triton,triton,tm1,miniworld_engine/kernels/tm1/triton/main.py,ok,rel dLA=2.54e-03 dLB=2.07e-03 dRA=2.84e-03 dRB=2.07e-03 (band 1e-02 declared),bf16
 trimul_gemm_gate_mmajor_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/bidirectional.py,ok,rel left=3.30e-03 right=1.85e-03 preact_left=2.64e-03 preact_right=2.73e-03 (band 1e-02 declared),bf16
-trimul_gemm_gate_saveact_triton,triton,triangle_multiplication,miniworld_engine/modules/triangle_multiplication/baseline_dtv1.py,ok,rel ab=2.09e-03 sig_m=1.91e-06 (band 1e-02 declared),bf16
+trimul_gemm_gate_saveact_triton,triton,triangle_multiplication,miniworld_engine/modules/triangle_multiplication/baseline_dtv1.py,failed,ModuleNotFoundError: No module named 'cuequivariance_ops',bf16
 trimul_gemm_gate_triton,triton,tm1,miniworld_engine/kernels/tm1/triton/main.py,ok,rel left=3.19e-03 right=1.84e-03 (band 8e-03 declared),bf16
 trimul_outproj_bwd_gate_recompute_triton,triton,tm2,miniworld_engine/kernels/tm2/triton/main.py,ok,rel dA=2.81e-03 dB=2.25e-03 (band 1e-02 declared),bf16
-trimul_outproj_gemm_gate_saveact_triton,triton,triangle_multiplication,miniworld_engine/modules/triangle_multiplication/baseline_dtv1.py,ok,rel ab=3.74e-03 sig=1.76e-06 (band 9e-03 declared),bf16
+trimul_outproj_gemm_gate_saveact_triton,triton,triangle_multiplication,miniworld_engine/modules/triangle_multiplication/baseline_dtv1.py,failed,ModuleNotFoundError: No module named 'cuequivariance_ops',bf16
 trimul_outproj_gemm_gate_triton,triton,tm2,miniworld_engine/kernels/tm2/triton/main.py,ok,rel out=1.88e-03 (band 8e-03 declared),bf16
 trimul_outproj_layernorm_gemm_gate_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/back.py,ok,rel y_residual=2.21e-03 (band 1e-02 declared),bf16
 qk_norm_rope_fwd_triton,triton,rope,miniworld_engine/kernels/rope/triton/qk_norm.py,ok,rel q=0.00e+00 k=0.00e+00 (band 8e-03 declared),bf16
@@ -99,7 +99,7 @@ transition_fwd_b2b_sm90_cuda,cuda,transition,miniworld_engine/kernels/transition
 transition_expand_gate_sm90_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_expand_gate_kernel.cu,ok,rel out=1.94e-03 (band 8e-03 declared),bf16
 transition_bwd_gate_sm90_cuda,cuda,transition,miniworld_engine/kernels/transition/cuda/transition_gatebwd_kernel.cu,ok,rel h=1.94e-03 dAB=2.97e-03 xn=5.27e-04 (band 1e-02 declared),bf16
 trimul_output_f567_train_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/output_fused.py,ok,rel y=4.27e-03 proj=2.95e-03 gate=2.72e-03 (band 1e-02 declared),bf16
-trimul_input_ln_residual_bwd_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/backward_fused.py,ok,rel dx=2.29e-03 dw=1.05e-06 db=8.17e-08 (band 1e-02 declared),bf16
+trimul_input_ln_residual_bwd_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/backward_fused.py,ok,rel dx=2.29e-03 dw=3.22e-07 db=8.17e-08 (band 1e-02 declared),bf16
 trimul_input_dual_bwd_triton,triton,trimul_inproj,miniworld_engine/kernels/trimul_inproj/triton/backward_fused.py,ok,rel dx=2.34e-03 (band 1e-02 declared),bf16
 transition_squeeze_residual_triton,triton,transition,miniworld_engine/kernels/transition/triton/residual.py,ok,rel out=2.14e-03 (band 1e-02 declared),bf16
 mpnn_edge_tail_fwd_gemm_gather_saveact_triton,triton,mpnn_edge_tail,miniworld_engine/kernels/mpnn_edge_tail/triton/compute.py,ok,rel out=6.01e-03 (band 5e-02),bf16
@@ -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=1.15e-07 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=5.77e-08 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

{miniworld_engine-2.3.0 → miniworld_engine-2.3.2}/src/miniworld_engine/build/kernel_bundle.py

@@ -397,15 +397,39 @@ def update_prebuilt(entries: dict[str, dict], root: Path, arch: str, *, keep: bo
 
 
 # --------------------------------------------------------------------------------------------------- replay
+def keyed_recipe(entry: dict) -> dict:
+    """``entry``'s recipe with the ``extra_ldflags`` its key was computed from. torch's ``load`` appends its own link
+    flags (``-L<torch>/lib -lc10 ... -lcudart``) to the list it is given, in place, and records up to 2.3.1 stored
+    that list after the call: replayed as recorded, every extension with ``extra_ldflags`` (the SWA atom DiT, the
+    token DiT's owned GEMM) keyed differently from the one a model asks for. The key material holds the flags as
+    they were passed."""
+    recipe = entry["recipe"]
+    flags = (entry.get("material") or {}).get("flags") or {}
+    if (recipe.get("fn") != "ext" or "extra_ldflags" not in recipe.get("kwargs", {})
+            or any("<" in f for f in flags.get("ldflags", ["<"]))):  # absent, or machine-free placeholders
+        return recipe
+    return {**recipe, "kwargs": {**recipe["kwargs"], "extra_ldflags": list(flags["ldflags"])}}
+
+
+def rebase_recipe(recipe: dict, recorded_jit: str | None, jit: str | None, package: Path = PKG_ROOT) -> dict:
+    """``recipe`` with the recording's checkout and JIT root replaced by this one's: a replay compiles THIS tree's
+    sources (and writes generated files under this JIT root), so it keys them as this tree's own builds would. (A
+    recipe still carries the generated TriMul header and ``compile_text`` texts of the recording: refresh the records
+    with a coverage run after editing the sources they are generated from.)"""
+    text = json.dumps(recipe)
+    m = re.search(r'(?:"|-I)(/[^"]*?/miniworld_engine)/(?:kernels|integrations|modules)/', text)
+    for old, new in ((recipe.get("jit") or recorded_jit, jit), (m.group(1) if m else None, str(package))):
+        if old and new and old != new:
+            text = text.replace(old, new)
+    return json.loads(text)
+
+
 def _replay_one(task: dict) -> dict:
     """Rebuild one recorded artifact here (a worker process). Returns {old, key, path} or {old, error}."""
     from miniworld_engine.kernels import _codegraph as cg
     from miniworld_engine.kernels import _nvcc
 
-    r = task["recipe"]
-    old_jit, new_jit = r.get("jit") or task.get("jit"), cg.context().jit
-    if old_jit and new_jit and old_jit != new_jit:  # generated headers / sources lived under the recording's JIT root
-        r = json.loads(json.dumps(r).replace(old_jit, new_jit))
+    r = rebase_recipe(task["recipe"], task.get("jit"), cg.context().jit)
     try:
         if r["fn"] == "derived":
             base = _replay_one({"recipe": task["base"], "old": r["base_key"], "jit": task.get("jit")})
@@ -440,6 +464,8 @@ def replay(records: list[Path], out: Path, root: Path, arch: str, jobs: int, onl
     if only:
         recipes = {k: e for k, e in recipes.items() if only in e["kind"] or only in e["name"]}
     record = out / "records" / f"replay-{_abi_tag(sys.implementation.cache_tag)}.jsonl"
+    if record.exists():  # one replay, one record: never pack an earlier run's artifacts
+        record.rename(record.with_suffix(f".{int(time.time())}.prev"))
     env = driver_env(out, arch, record)
     os.environ.update({k: env[k] for k in (prebuilt.RECORD_ENV, "MINIWORLD_ENGINE_BUNDLE_ARCH", "XDG_CACHE_HOME",
                                            "MINIWORLD_ENGINE_JIT_FALLBACK")})
@@ -465,9 +491,9 @@ def replay(records: list[Path], out: Path, root: Path, arch: str, jobs: int, onl
             if base is None:
                 log(f"  skip {e['name']}: its base cubin was not recorded")
                 continue
-            tasks.append({"old": key, "recipe": r, "base": base["recipe"], "jit": jit_of(base)})
+            tasks.append({"old": key, "recipe": keyed_recipe(e), "base": keyed_recipe(base), "jit": jit_of(base)})
         else:
-            tasks.append({"old": key, "recipe": r, "jit": jit_of(e)})
+            tasks.append({"old": key, "recipe": keyed_recipe(e), "jit": jit_of(e)})
     tasks.sort(key=lambda t: t["recipe"]["fn"] == "derived")
     log(f"replaying {len(tasks)} recipes with {jobs} processes")
     started, errors = time.monotonic(), []
@@ -502,7 +528,7 @@ def bundle_kernels(args) -> int:
         rec = replay([Path(p) for p in args.replay], out, root, arch, args.jobs, args.only)
         timings["replay"] = time.monotonic() - t0
         records = [rec]
-        keep = True
+        keep = args.keep  # artifacts of this ABI (and cubins) the replay did not produce are pruned
     else:
         device_arch = cg.target_arch()
         if device_arch not in ("cpu", arch):

{miniworld_engine-2.3.0 → miniworld_engine-2.3.2}/src/miniworld_engine/integrations/bias_only_dit_hoist.py

@@ -51,8 +51,10 @@ F32 = torch.float32
 SINK_ATTR = "_bias_only_dit_hoist_sink"
 
 
-def serves(implementation, attns, pair) -> bool:
-    """The fused hoist serves ``pair`` and the attention modules ``attns`` (eps / head count / no offset checked by the caller)."""
+def serves(implementation, attns, pair, per_sample=False) -> bool:
+    """The fused hoist serves ``pair`` and the attention modules ``attns`` (eps / head count / no offset checked by the caller).
+    ``per_sample``: ``pair`` is [A, L, L, 128] with A > 1 (a pair of its own per sample, the samples in the batch axis) instead of
+    [1, L, L, 128]."""
     from miniworld_engine.modules.exceptions import ImplementationType
 
     if not settings.current().bias_only_dit_hoist:
@@ -65,7 +67,7 @@ def serves(implementation, attns, pair) -> bool:
         return False
     if (
         pair.ndim != 4
-        or pair.shape[0] != 1
+        or (pair.shape[0] < 2 if per_sample else pair.shape[0] != 1)
         or pair.shape[1] != pair.shape[2]
         or pair.shape[-1] != DP
         or pair.shape[1] == 0
@@ -134,6 +136,9 @@ def _hoist_bwd(
 
     N, R = dball.shape
     f32 = pair2d.dtype is F32
+    if settings.current().bias_only_dit_hoist_bwd_fused and HK.fused_bwd_served(pair2d, wf):
+        dpair, dwf = HK.hoist_bwd_fused(pair2d, wf, dball, eps)  # one pass over the pair (bo_hoist_bwd.cu)
+        return [dpair, dwf]
     K = HK.operand_cols(pair2d.dtype)
     with _gemms(pair2d.dtype):
         dln = (
@@ -165,11 +170,11 @@ class _Sink:
     time); a gradient that does not sit in its slot (autograd summed it with another, a reentrant checkpoint copied the input) is
     copied in, a missing one is zero."""
 
-    __slots__ = ("buf", "device", "dtype", "given", "h", "nb", "r")
+    __slots__ = ("buf", "device", "dtype", "given", "h", "nb", "psb", "r")
 
-    def __init__(self, nb, h, r, dtype, device):
+    def __init__(self, nb, h, r, dtype, device, psb=False):
         self.nb, self.h, self.r, self.dtype, self.device = nb, h, r, dtype, device
-        self.buf, self.given = None, set()
+        self.buf, self.given, self.psb = None, set(), psb
 
     def slot(self, i, dtype, device):
         """Block i's [H, R] slice of the buffer, or None (another dtype / device, or the slot already handed out)."""
@@ -197,7 +202,8 @@ class _Sink:
             s = buf[i * self.h : (i + 1) * self.h]
             if g is None:
                 s.zero_()
-            elif not (
+                continue
+            if not (
                 g.data_ptr() == s.data_ptr()
                 and g.dtype == s.dtype
                 and g.is_contiguous()
@@ -221,13 +227,16 @@ class _Hoist(torch.autograd.Function):
     def forward(ctx, pair, eps, nb, sink, *params):  # params: ln_pair.weight, to_bias.weight of each block
         L, dp = pair.shape[1], pair.shape[-1]
         H = params[1].shape[0]
+        A = pair.shape[0] if sink.psb else 1
         gam, wts, wf = _fold(params, pair.dtype)
         ctx.set_materialize_grads(
             False
         )  # an unused block's bias: None, not a zero tensor
-        ball = _hoist_fwd(pair.reshape(L * L, dp), wf, eps)
+        ball = _hoist_fwd(pair.reshape(A * L * L, dp), wf, eps)
         ctx.save_for_backward(pair, wf, gam, wts)
         ctx.meta = (eps, nb, H, sink, [p.dtype for p in params])
+        if sink.psb:  # per-sample: block b's slice is the head-major [H, A, L, L]
+            return tuple(ball[i * H : (i + 1) * H].view(H, A, L, L) for i in range(nb))
         return tuple(ball[i * H : (i + 1) * H].view(1, H, L, L) for i in range(nb))
 
     @staticmethod
@@ -235,8 +244,9 @@ class _Hoist(torch.autograd.Function):
         pair, wf, gam, wts = ctx.saved_tensors
         eps, nb, H, sink, pdts = ctx.meta
         L, dp = pair.shape[1], pair.shape[-1]
+        R = (pair.shape[0] if sink.psb else 1) * L * L
         dball = sink.take(grads)
-        dpair, dwf = _hoist_bwd(pair.reshape(L * L, dp), wf, dball, eps, _split(L * L))
+        dpair, dwf = _hoist_bwd(pair.reshape(R, dp), wf, dball, eps, _split(R))
         del dball
         dwf = dwf.view(nb, H, dp)
         dW = dwf * gam[:, None, :]  # dW_b = dW'_b diag(g_b)
@@ -272,4 +282,21 @@ def pair_bias_all(attns, pair) -> tuple[torch.Tensor, ...]:
     return tuple(outs)
 
 
-__all__ = ["SINK_ATTR", "pair_bias_all", "serves"]
+def pair_bias_all_psb(attns, pair) -> tuple[torch.Tensor, ...]:
+    """Each attention module's bias [A, H, L, L] for a pair of its own per sample ``pair`` [A, L, L, 128] (A > 1): the ``[A, H, L, L]``
+    view of block b's head-major [H, A, L, L] slice of ONE buffer (one LayerNorm pass and one GEMM for every block); differentiable in
+    the pair and in every ``ln_pair.weight`` / ``to_bias.weight``. The fused per-sample block reads the head-major buffer without a
+    copy and writes its bias gradient into the same slot of the shared gradient buffer. Call ``serves(..., per_sample=True)`` first."""
+    nb, H, A, L = len(attns), attns[0].to_bias.weight.shape[0], pair.shape[0], pair.shape[1]
+    sink = _Sink(nb, H, A * L * L, pair.dtype, pair.device, psb=True)
+    params = [p for a in attns for p in (a.ln_pair.weight, a.to_bias.weight)]
+    outs = _Hoist.apply(pair.contiguous(), float(attns[0].ln_pair.eps), nb, sink, *params)
+    res = []
+    for i, o in enumerate(outs):
+        v = o.permute(1, 0, 2, 3)
+        setattr(v, SINK_ATTR, (sink, i))
+        res.append(v)
+    return tuple(res)
+
+
+__all__ = ["SINK_ATTR", "pair_bias_all", "pair_bias_all_psb", "serves"]