miniworld-engine 2.5.0__tar.gz → 2.5.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- miniworld_engine-2.5.2/PKG-INFO +218 -0
- miniworld_engine-2.5.2/README.md +182 -0
- miniworld_engine-2.5.2/licenses/Anthropic-NOTICE +292 -0
- miniworld_engine-2.5.2/licenses/THIRD_PARTY_NOTICES.md +95 -0
- miniworld_engine-2.5.2/licenses/anthropic-trimul-native-v5-UPSTREAM.json +15 -0
- miniworld_engine-2.5.2/pyproject.toml +187 -0
- miniworld_engine-2.5.2/src/miniworld_engine/judge/__init__.py +479 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_capture.py +290 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_cast.py +81 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_cast_cuda.py +75 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_codegraph.py +525 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_compile.py +188 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_device.py +47 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_dot_precision.py +21 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_fastmath.py +24 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_h100_master.py +28 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_headers.py +220 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_launch_config.py +259 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_nvcc.py +1019 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_prebuild.py +694 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_prebuild_cuda.py +247 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_pretune.py +1158 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_shape_key.py +110 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_tf32.py +191 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/_tiles.py +103 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/__init__.py +4 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80/adaln_atom_bwd.cuh +286 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80/adaln_atom_fwd.cuh +189 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80/adaln_atom_fwd_tf32.cuh +166 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80/adaln_common.cuh +76 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80/adaln_finish.cuh +84 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80/adaln_launch.cuh +21 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80/adaln_mma.cuh +101 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80/adaln_rows.cuh +179 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80/adaln_tf32.cuh +59 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80/ops.cu +281 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/cuda/sm80.py +275 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/triton/inference.py +823 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/triton/ln_strided.py +150 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/adaln/triton/training.py +1041 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/configs/a100_sm80.json +2145 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/configs/h100_sm90.json +2419 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/__init__.py +616 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/apb/__init__.py +177 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/apb/apb_pair_bias_sm80.cuh +601 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/apb/apb_rows.cu +2412 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/apb/apb_rows_sm80.cu +2055 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/apb/step_sm80.cuh +297 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/atom_output_fused_sm80.py +35 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/atom_output_tiled_sm80.py +35 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/atom_qkvg_resident_sm80.py +39 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/atom_qkvg_tiled_sm80.py +39 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/atom_sample_uniform_sm80.py +55 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/prebuild_sm80.py +47 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/shared_mask_uniform_sm80.py +55 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/shared_score_core_sm80.py +82 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/shared_uniform_core_sm80.py +82 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/__init__.py +791 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/driver.py +187 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/__init__.py +1284 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/__init__.py +1011 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_output_fused.cu +93 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_output_tiled.cu +104 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_output_tiled_rolled.cu +104 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_pair_reduce_bwd.cu +101 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_pair_reduce_vec.cu +121 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_pair_uniform_no_db.cu +64 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_post_grad_delta.cu +146 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_post_grad_fused.cu +121 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_post_grad_rolled.cu +122 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_qkvg.cu +62 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_qkvg_lean.cu +63 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_qkvg_resident.cu +68 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_qkvg_tiled.cu +91 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/atom_sample_uniform.cu +78 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/attn_bwd_dkv_sm80.cuh +262 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/attn_bwd_dkv_tf32_sm80.cuh +228 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/attn_bwd_dq_sm80.cuh +233 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/attn_bwd_dq_tf32_sm80.cuh +198 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/attn_fwd_sm80.cuh +369 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/attn_fwd_tf32_sm80.cuh +265 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/aux_sm80.cuh +398 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/glue_ops.cu +123 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/glue_sm80.cuh +134 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/ops.cu +535 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/ops_tf32.cu +206 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/shared_mask_uniform.cu +78 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/sm80_common.cuh +109 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80/token_grad_ops.cu +72 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80.py +1092 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local/__init__.py +284 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local/glue.cu +56 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local/lattn_dkv.cu +199 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local/lattn_dq.cu +185 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local/lattn_fwd.cu +184 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local/lbias.cu +173 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local/local.cuh +62 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local_post.cu +95 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local_post.py +66 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local_projection.cu +30 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_local_projection.py +59 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/__init__.py +227 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/cond_bwd.cuh +243 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/cond_fwd.cuh +210 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/ops.cu +310 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/post_bwd.cuh +364 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/post_fwd.cuh +262 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/pre_bwd.cuh +240 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/pre_fwd.cuh +240 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/rows_common.cuh +124 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/rows_layout.cuh +131 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/tr_bwd.cuh +209 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_atom_rows/wgrad.cuh +122 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_core_order237/atom_core_order_dkv.cuh +271 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_core_order237/atom_core_order_dq.cuh +244 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_core_order237/ops.cu +201 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_core_qresident248/atom_core_qresident_dkv.cuh +271 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_core_qresident248/atom_core_qresident_dq.cuh +255 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_core_qresident248/ops.cu +201 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_fwd_qresident257/atom_fwd_qresident.cuh +372 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_fwd_qresident257/ops.cu +109 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_score/attn_bwd_dkv_sm80.cuh +261 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_score/attn_bwd_dq_sm80.cuh +232 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_score/attn_fwd_sm80.cuh +376 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_score/aux_sm80.cuh +398 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_score/ops.cu +468 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_score/sm80_common.cuh +109 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_uniform/attn_bwd_dkv_sm80.cuh +261 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_uniform/attn_bwd_dq_sm80.cuh +232 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_uniform/attn_fwd_sm80.cuh +397 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_uniform/aux_sm80.cuh +398 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_uniform/ops.cu +468 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_shared_uniform/sm80_common.cuh +109 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_grad_rne/attn_bwd_dkv_tf32_sm80.cuh +228 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_grad_rne/attn_bwd_dq_tf32_sm80.cuh +198 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_grad_rne/attn_fwd_sm80.cuh +363 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_grad_rne/attn_fwd_tf32_sm80.cuh +265 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_grad_rne/ops_tf32.cu +204 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_grad_rne/sm80_common.cuh +109 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_rne/attn_bwd_dkv_tf32_sm80.cuh +228 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_rne/attn_bwd_dq_tf32_sm80.cuh +198 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_rne/attn_fwd_sm80.cuh +363 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_rne/attn_fwd_tf32_sm80.cuh +265 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_rne/ops_tf32.cu +204 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_qk_rne/sm80_common.cuh +109 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne/attn_bwd_dkv_tf32_sm80.cuh +228 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne/attn_bwd_dq_tf32_sm80.cuh +198 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne/attn_fwd_sm80.cuh +363 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne/attn_fwd_tf32_sm80.cuh +259 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne/ops_tf32.cu +204 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne/sm80_common.cuh +109 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne_hoisted/attn_bwd_dkv_tf32_sm80.cuh +232 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne_hoisted/attn_bwd_dq_tf32_sm80.cuh +202 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne_hoisted/attn_fwd_sm80.cuh +363 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne_hoisted/attn_fwd_tf32_sm80.cuh +263 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne_hoisted/ops_tf32.cu +204 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm80_tf32_rne_hoisted/sm80_common.cuh +109 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/__init__.py +850 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/cond90.cu +124 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/cond_bwd90.cu +141 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/cond_bwd_sh90.cu +119 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/kqln_bwd90.cu +56 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/pack90.cu +143 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/post90.cu +190 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/post_bwd90.cu +274 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/pre90.cu +209 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/pre_bwd90.cu +227 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/tr_bwd90.cu +126 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/tf32_qk_grad_rne_sm80.py +24 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/tf32_qk_rne_sm80.py +24 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/tf32_rne_hoisted_sm80.py +24 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/tf32_rne_sm80.py +24 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/token_grad_sm80.py +50 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/token_shared_scale_sm80.cu +53 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/cuda/token_shared_scale_sm80.py +27 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/triton/main.py +973 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/triton/memory_efficient.py +470 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/augmented_attention/triton/window.py +570 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_attention/cuda/__init__.py +1 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_attention/cuda/sm80.py +139 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_attention/triton/core_t.py +148 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_attention/triton/gate_out.py +319 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_attention/triton/main.py +464 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_attention/triton/pair_bias.py +197 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/__init__.py +1022 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/bias_only_dit_rows.cu +692 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/bwd_fused.py +649 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/hoist.py +71 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/inf3_bf16.py +421 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/pack_sm80.cu +107 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/runner.py +734 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/sm80/__init__.py +230 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/sm80/dpb_sm80.cuh +160 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/sm80/dpb_tf32_sm80.cuh +123 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/sm80/ops.cu +281 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/sm80/pv_gate_sm80.cuh +124 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/sm80/pv_gate_tf32_sm80.cuh +134 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/sm80/sm80_common.cuh +121 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/tf32.py +977 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/bias_only_dit/cuda/train.py +52 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/__init__.py +653 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_cond_wgrad_fused_sm80.cu +189 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_cond_wgrad_fused_sm80.py +58 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_narrow_adaln.cuh +193 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_narrow_sm80.cu +36 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_narrow_sm80.py +34 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_narrow_tail.cuh +199 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_small_adaln.cuh +190 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_small_sm80.cu +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_small_sm80.py +34 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_small_tail.cuh +199 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_split_adaln.cuh +193 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_split_sm80.cu +49 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_split_sm80.py +34 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_split_tail.cuh +185 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_wgrad_split_sm80.cu +178 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_wgrad_split_sm80.py +38 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_wgrad_wide_sm80.cu +184 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/ct_atom_wgrad_wide_sm80.py +41 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/gemm_swiglu.py +192 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/sm80/ct_atom_bwd.cuh +318 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/sm80/ct_atom_fwd.cuh +198 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/sm80/ct_atom_fwd_tf32.cuh +274 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/sm80/ct_rows.cuh +80 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/sm80/ops.cu +222 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/sm80.py +227 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/token_dit_composed.py +131 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/token_dit_pack.cu +146 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/token_dit_rows.cu +899 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/token_dit_train_rows.cu +1453 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/token_train_pack_alias_sm80.cu +93 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/token_train_pack_alias_sm80.py +51 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/token_train_pack_sm80.cu +87 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/token_train_pack_sm80.py +51 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/cuda/train.py +179 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/triton/composed.py +354 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/triton/fwd_saveact.py +302 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/triton/inference.py +219 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/triton/token_dit_attn.py +301 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/triton/token_dit_kernels.py +197 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/triton/token_dit_runner.py +446 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/conditioned_transition/triton/training.py +870 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/dit_variants/cuda/__init__.py +59 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/dit_variants/cuda/sm80.py +228 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/dit_variants/cuda/train.py +767 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/dit_variants/cuda/train_rows.cu +693 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/dit_variants/triton/portable.py +232 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/gated_projection/cuda/sm80/gp_kernels.cuh +458 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/gated_projection/cuda/sm80/gp_ops.cu +191 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/gated_projection/cuda/sm80.py +245 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/gated_projection/dispatch.py +56 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/gated_projection/triton/main.py +320 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/gated_projection/triton/residual.py +115 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/gated_projection/triton/swa.py +115 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/__init__.py +182 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/adaln_sm80.py +431 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/attention_pair_bias_b200.py +920 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/attention_pair_bias_sm80.py +909 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/augattn_sm80.py +1356 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/bias_only_dit.py +239 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/bias_only_dit_hoist.py +233 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/bias_only_dit_sm80.py +319 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/bias_only_dit_train.py +1729 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/bidir_triattn_h100.py +207 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/conditioned_transition_sm80.py +494 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/csrc/opm_epilogue.cu +2901 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/csrc/pair_wide.cu +427 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/csrc/pwa_wide.cu +1173 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/local_dit.py +100 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/local_dit_h100.py +739 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/local_dit_sm80.py +262 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/local_dit_sm80_block.py +546 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/mpnn_edge_sm80.py +816 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/mpnn_msg_sm80.py +494 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/opm_sm80.py +528 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/opm_train.py +1115 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/pwa_sm80.py +671 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/pwa_train.py +1034 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/pwa_wide.py +337 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/pwa_wide_sm80.py +198 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/sigmoid_gate_sm80.py +107 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/token_dit.py +470 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/token_dit_train.py +1403 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/triattn_h100_sm90.py +830 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/triattn_sm80.py +328 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/trimul_h100.py +841 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/integrations/trimul_sm80.py +122 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/compile_native.py +129 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/configs/a100_sm80.json +2017 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/configs/h100_sm90.json +4065 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/cuda/__init__.py +42 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/cuda/general/__init__.py +232 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/cuda/general/bind.cpp +14 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/cuda/general/ln_rows.cu +523 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/cuda/pipe/__init__.py +148 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/cuda/pipe/ln_pipe.cu +1091 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/cuda/sm80/norm_common.cuh +113 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/cuda/sm80/norm_rows.cu +1270 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/cuda/sm80.py +293 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/triton/lowreg.py +60 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/triton/main.py +436 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/triton/persistent.py +297 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm/triton/transpose.py +150 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/autograd.py +168 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/configs/a100_sm80.json +457 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/cuda/__init__.py +1 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/cuda/sm80/lnl_ops.cu +112 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/cuda/sm80/lnl_sm80.cuh +492 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/cuda/sm80.py +176 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/dispatch.py +28 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/triton/fused.py +422 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/triton/mmajor_bwd.py +301 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/triton/pair_bias.py +476 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/triton/recompute.py +123 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/triton/stats.py +117 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear/triton/te_style.py +410 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_dropout/triton/main.py +179 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_layernorm/triton/main.py +144 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_mlp/triton/compute.py +255 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_mlp/triton/main.py +315 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/configs/a100_sm80.json +211 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/configs/h100_sm90.json +211 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/__init__.py +5 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/sm80/common.cuh +155 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/sm80/dw.cuh +207 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/sm80/mlp.cuh +259 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/sm80/ops.cu +396 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/sm80/pack.cuh +60 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/sm80/rows.cuh +177 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/sm80/scatter.cuh +104 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/sm80/tail_bwd.cuh +237 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/sm80/tail_fwd.cuh +278 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/cuda/sm80.py +42 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/triton/compute.py +1078 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_edge_tail/triton/main.py +1047 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/configs/a100_sm80.json +137 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/configs/h100_sm90.json +137 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/cuda/__init__.py +1 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/cuda/sm80/mpnn_common.cuh +176 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/cuda/sm80/msg_bwd_fused_sm80.cuh +296 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/cuda/sm80/msg_bwd_sm80.cuh +322 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/cuda/sm80/msg_fwd_sm80.cuh +150 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/cuda/sm80/ops.cu +118 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/cuda/sm80.py +125 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/triton/inference.py +198 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_message/triton/main.py +883 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_node_message/configs/a100_sm80.json +57 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_node_message/configs/h100_sm90.json +57 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_node_message/cuda/__init__.py +1 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_node_message/cuda/sm80/node_bwd_sm80.cuh +435 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_node_message/cuda/sm80/node_fwd_sm80.cuh +267 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_node_message/cuda/sm80/ops.cu +99 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_node_message/cuda/sm80.py +131 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_node_message/triton/main.py +830 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_relative_position/configs/a100_sm80.json +33 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_relative_position/configs/h100_sm90.json +33 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_relative_position/cuda/__init__.py +1 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_relative_position/cuda/sm80/ops.cu +40 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_relative_position/cuda/sm80/relpos_sm80.cuh +190 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_relative_position/cuda/sm80.py +64 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/mpnn_relative_position/triton/main.py +343 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/chai.py +232 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/composed.py +135 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/configs/a100_sm80.json +1067 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/__init__.py +87 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/chai_sm80.py +58 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/finalize_typed_sm80.py +26 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/mask_count_ballot_sm80.py +27 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/mask_count_sm80.py +32 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/mask_epilogue_sm80.py +36 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/ogemm_sm80.py +40 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/pipeline_native128_sm80.py +71 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/pipeline_sm80.py +71 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/rect_gemm_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/chai_sm80.cu +533 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/finalize_typed.cu +30 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/mask_count.cu +54 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/mask_count_ballot.cu +63 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/mask_epilogue.cu +69 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/ogemm_cutlass.cu +77 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/opm_dgrad_sm80.cuh +179 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/opm_dwo_sm80.cuh +121 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/opm_epilogue_mask_sm80.cuh +195 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/opm_epilogue_sm80.cuh +179 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/opm_lt_gemm.cuh +92 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/opm_lt_table.inc +39 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/opm_prologue_bwd_sm80.cuh +526 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/opm_prologue_sm80.cuh +186 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/ops.cu +289 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/pipeline.cu +65 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/pipeline_finalize.cuh +30 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/pipeline_mask_epilogue.cuh +69 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/pipeline_mask_epilogue_math.cuh +196 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/pipeline_native128.cu +80 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/pipeline_ops.cuh +254 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/rect_gemm.cu +200 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/rect_gemm_core.cuh +200 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80/sm80_common.cuh +89 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/cuda/sm80.py +232 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean/triton/main.py +1112 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_adaln/configs/a100_sm80.json +87 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_adaln/configs/h100_sm90.json +87 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_adaln/cuda/__init__.py +15 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_adaln/cuda/function.py +99 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_adaln/cuda/sm80/__init__.py +48 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_adaln/cuda/sm80/pair_embed_adaln_sm80.cu +24 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_adaln/cuda/sm90/__init__.py +100 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_adaln/cuda/sm90/pair_embed_adaln.cu +426 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_adaln/triton/function.py +32 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_adaln/triton/main.py +327 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/configs/a100_sm80.json +217 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/configs/h100_sm90.json +217 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/cuda/__init__.py +15 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/cuda/function.py +166 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/cuda/sm80/__init__.py +64 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/cuda/sm80/pair_embed_geom_sm80.cu +74 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/cuda/sm80/pe_sm80.cuh +63 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/cuda/sm90/__init__.py +77 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/cuda/sm90/pair_embed_geom.cu +995 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/triton/function.py +54 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_embed_geom/triton/main.py +400 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/configs/a100_sm80.json +1383 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/configs/h100_sm90.json +1389 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/__init__.py +3 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/broadcast_capture_sm80.py +32 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/broadcast_hybrid_sm80.py +34 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/broadcast_native_sm80.py +34 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/broadcast_pipeline_sm80.py +34 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/broadcast_warp_sm80.py +34 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/contract_gate_fused128_sm80.py +31 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/contract_gate_fused_sm80.py +31 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/contract_gate_narrow_sm80.py +31 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/dw_cooperative_sm80.py +28 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/dw_early_sm80.py +28 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/dw_fold_sm80.py +28 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/dw_group_sm80.py +31 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/dw_native_sm80.py +36 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/glue_tuned_sm80.py +58 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/grad_partials_sm80.py +63 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/grad_reduce_fused_sm80.py +27 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/grad_reduce_vector_sm80.py +27 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/pair_partial_sm80.py +31 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/broadcast_bmm.cu +40 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/broadcast_capture_old.cu +145 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/broadcast_capture_pipeline.cu +168 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/broadcast_hybrid.cu +155 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/broadcast_native.cu +144 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/broadcast_pipeline.cu +149 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/broadcast_warp.cu +146 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/contract_gate_fused.cu +198 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/contract_gate_fused128.cu +209 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/contract_gate_narrow.cu +209 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/dw_cooperative.cu +163 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/dw_early.cu +148 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/dw_fold.cu +137 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/dw_group.cu +110 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/dw_native.cu +99 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/glue_tuned_ops.cu +30 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/grad_reduce_fused.cu +104 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/grad_reduce_vector.cu +137 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/ops.cu +330 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pair_partial.cu +106 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_dgv_bwd_sm80.cuh +300 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_finalize_sm80.cuh +49 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_gate_out_sm80.cuh +237 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_glue_sm80.cuh +304 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_glue_tuned.cuh +305 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_glue_v2.cuh +306 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_ln_v_sm80.cuh +175 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_lt_gemm.cuh +136 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_lt_table.inc +145 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_pair_bwd_sm80.cuh +268 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_pair_fwd_sm80.cuh +132 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/pwa_pair_partial_sm80.cuh +288 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/sm80_common.cuh +117 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80/wide_ops.cu +348 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/sm80.py +305 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/cuda/wide_sm80.py +24 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging/triton/main.py +779 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/rmsnorm/cuda/sm80/adamod_sm80.cu +419 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/rmsnorm/cuda/sm80.py +387 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/rmsnorm/triton/main.py +867 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/rope/cuda/__init__.py +1 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/rope/cuda/sm80/rope_sm80.cu +257 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/rope/cuda/sm80.py +220 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/rope/triton/main.py +146 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/rope/triton/qk_norm.py +378 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/sampler_ops/configs/a100_sm80.json +155 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/sampler_ops/configs/h100_sm90.json +155 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/sampler_ops/cuda/__init__.py +33 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/sampler_ops/cuda/function.py +106 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/sampler_ops/cuda/sampler_ops.cu +665 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/sampler_ops/cuda/sampler_ops_sm80.cu +144 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/sampler_ops/cuda/sm80.py +34 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/sampler_ops/triton/__init__.py +10 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/sampler_ops/triton/function.py +119 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/sampler_ops/triton/main.py +533 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/smooth_lddt/configs/a100_sm80.json +133 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/smooth_lddt/configs/h100_sm90.json +133 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/smooth_lddt/cuda/__init__.py +12 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/smooth_lddt/cuda/function.py +82 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/smooth_lddt/cuda/sm100/__init__.py +177 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/smooth_lddt/cuda/sm100/smooth_lddt.cu +368 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/smooth_lddt/triton/__init__.py +0 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/smooth_lddt/triton/function.py +39 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/smooth_lddt/triton/main.py +289 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/autograd.py +77 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/__init__.py +53 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/ffn_fwd_tuned_sm80.py +43 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/grad_finalize_sm80.cu +34 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/grad_finalize_sm80.py +32 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_ffn_sm80_a100.cu +66 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_ffn_sm80_a100.cuh +296 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_qkvg_headgroup_common_sm80.cuh +121 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_qkvg_headgroup_sm80.cuh +225 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_qkvg_singlehead_common_sm80.cuh +121 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_qkvg_singlehead_sm80.cuh +225 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_qkvg_sm80_a100.cu +146 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_rowgrid_ffn_sm80_a100.cu +69 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_rowgrid_ffn_sm80_a100.cuh +296 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_rowgrid_ffn_sm80_a100.py +52 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_rowgrid_sm80_a100.py +61 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_rt8_ffn_sm80_a100.cu +66 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_rt8_ffn_sm80_a100.cuh +300 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_rt8_ffn_sm80_a100.py +94 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/inference_sm80_a100.py +149 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/loader.py +129 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm100/__init__.py +1185 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm100/tf32_fwd.py +482 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/__init__.py +493 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/adaln_bwd_sm80.cuh +87 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/attn_bwd_sm80.cuh +312 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/attn_fwd_sm80.cuh +170 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/bwd_rows_sm80.cuh +91 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/composed.py +393 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/ffn_bwd_sm80.cuh +257 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/ffn_fwd_sm80.cuh +285 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/ffn_fwd_tuned.cu +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/ffn_fwd_tuned.cuh +288 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/fp32_fwd_sm80.cuh +363 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/mod_sm80.cuh +481 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/oproj_bwd_sm80.cuh +179 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/ops.cu +528 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/qkvg_bwd_sm80.cuh +143 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/qkvg_fwd_sm80.cuh +236 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/rows_sm80.cuh +86 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/sm80_common.cuh +121 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/wgrad_grouped.cu +186 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/wgrad_grouped5.cu +217 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/wgrad_grouped8.cu +187 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/wgrad_grouped8_alloc.cu +203 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/wgrad_qkvg.cu +176 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/sm80/wgrad_split.cu +178 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/swa_fp32.cu +1826 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/wgrad_grouped5_sm80.py +43 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/wgrad_grouped8_alloc_sm80.py +41 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/wgrad_grouped8_sm80.py +47 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/wgrad_grouped_sm80.py +47 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/wgrad_qkvg_sm80.py +38 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/cuda/wgrad_split_sm80.py +39 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/dispatch.py +1712 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/triton/backward.py +525 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/triton/backward_fp32.py +298 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/triton/forward.py +247 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/swa_dit/triton/forward_fp32.py +184 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/tm1/cuda/__init__.py +0 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/tm1/cuda/sm80.py +87 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/tm1/triton/main.py +403 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/tm2/cuda/__init__.py +0 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/tm2/cuda/sm80.py +84 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/tm2/triton/main.py +292 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/token_pair_init/configs/a100_sm80.json +105 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/token_pair_init/configs/h100_sm90.json +105 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/token_pair_init/cuda/__init__.py +15 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/token_pair_init/cuda/function.py +95 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/token_pair_init/cuda/sm80/__init__.py +51 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/token_pair_init/cuda/sm80/token_pair_init_sm80.cu +70 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/token_pair_init/cuda/sm90/__init__.py +64 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/token_pair_init/triton/function.py +31 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/token_pair_init/triton/main.py +244 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/configs/a100_sm80.json +3757 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/configs/h100_sm90.json +4396 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/__init__.py +186 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/fused_bwd_sm80.py +127 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/fused_fwd_sm80.py +195 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/fused_generic_sm90a.py +523 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/fused_sm100a.py +355 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/fused_sm80.py +350 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/fused_sm90a.py +428 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/fused_wide_sm100a.py +723 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/fused_wide_sm80.py +759 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/fused_wide_sm90a.py +858 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/gemm_epilogue_sm80.py +200 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/generic_sm90a/gen_wgmma.py +27 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/parameter_cast_pack_sm80.cu +46 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/parameter_cast_pack_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/parameter_cast_sm80.cu +43 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/parameter_cast_sm80.py +33 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/rect_dx_m64_sm80.py +28 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/rect_dx_sm80.py +28 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/rect_dx_gemm.cu +96 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/rect_dx_m64.cu +96 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/sm80_common.cuh +121 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/tr_bwd_g_sm80.cuh +230 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/tr_bwd_sm80.cuh +532 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/tr_fwd_sm80.cuh +567 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/transition_bwd_sm80.cu +153 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/transition_fwd_sm80.cu +156 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/transition_wide_sm80.cu +636 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/tw_gemm_sm80.cuh +264 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/tw_kernels_sm80.cuh +900 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/weight_grad_lean.cu +155 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/sm80/weight_grad_splitk.cu +151 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/weight_grad_lean_sm80.py +36 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/weight_grad_splitk_sm80.py +36 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/wide/f32io.cu +359 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/wide/kernels/transition_bwd_f32dw.cu +186 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/wide/kernels/transition_bwd_f32dx.cu +267 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/wide/kernels/transition_fwd_f32io.cu +271 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/wide/kernels/transition_fwd_f32xs.cu +235 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/cuda/wide/kernels/transition_lnsg_f32.cu +192 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/triton/fused.py +1852 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/triton/main.py +365 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/triton/residual.py +250 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/transition/triton/training.py +226 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/configs/a100_sm80.json +591 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/__init__.py +138 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/affine263.cu +42 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/affine263.cuh +179 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/affine_training_sm80.py +109 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/b1_finalize_exact_sm80.py +81 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/b1_gamma_exact_sm80.py +69 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/b200_triattn.py +315 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/dq_rows_candidate_sm80.py +69 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/h100_bidir.py +23 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/h100_dh.py +156 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sm90.py +379 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sm90_generic.py +795 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/front_fused3_sm90.cu +269 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/triattn_bidir_sm90.cu +83 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/triattn_dh_sm90.cu +638 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/triattn_h16_sm90.cu +567 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/ln_gamma_exact_sm80.py +40 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_db_nn_sm80.py +36 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_front_cache_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_front_ln_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_front_reuse_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_front_shift_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_front_stream2_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_front_stream_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_nt_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_out_fused_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_tn_aug_sm80.py +36 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_tn_sm80.py +36 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/native_wo_tn_sm80.py +40 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/attn2_bwd_sm80.cuh +468 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/attn2_fwd_sm80.cuh +360 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/attn2_ops.cu +266 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/attn_bwd_dkv_sm80.cuh +224 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/attn_bwd_dq_sm80.cuh +193 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/attn_fwd_sm80.cuh +370 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/b1_finalize_exact.cu +37 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/b1_finalize_exact.cuh +224 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/b1_gamma_exact.cu +25 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/b1_gamma_exact.cuh +206 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/b1_sm80.cuh +184 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/b3_sm80.cuh +144 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/dq_rows_candidate.cu +55 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/dq_rows_candidate.cuh +196 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/f1_sm80.cuh +253 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/f3_sm80.cuh +133 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/fused_sm80.cuh +615 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/ln_gamma_exact.cu +24 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/ln_gamma_exact.cuh +177 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_db_nn.cu +189 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_front_cache.cu +222 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_front_ln.cu +153 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_front_reuse.cu +159 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_front_shift.cu +193 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_front_stream.cu +172 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_front_stream2.cu +173 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_nt_gemm.cu +195 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_out_fused.cu +206 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_tn_aug.cu +191 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_tn_split.cu +170 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/native_wo_tn.cu +178 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/ops.cu +624 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/rows_ops.cu +511 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/rows_sm80.cuh +471 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/sm80_common.cuh +109 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80/wgrad_sm80.cuh +86 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80.py +1145 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80_bidir.py +36 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80_core2.py +230 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80_projected.py +230 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/cuda/sm80_wide.py +1416 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/triton/atomic.py +667 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/triton/bidirectional.py +296 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/triton/main.py +689 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/triangle_attention/triton/variants.py +106 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/configs/a100_sm80.json +2712 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/configs/h100_sm90.json +3526 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/__init__.py +265 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_bmm.py +290 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_driver.py +726 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_infer_kernel.py +964 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_launch.py +872 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_manifest.py +164 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_pack.py +624 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_runtime.py +447 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_infer.py +167 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_train.py +395 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/contract_fragment_lean_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/contract_lean_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/contract_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/contract_streamed_sm80.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b1.py +181 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style.py +708 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_aux.py +135 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_d384_tmasave.py +31 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_direct.py +88 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_dpl_pipeline.py +78 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_emitxn_tma.py +57 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_k1_prefetch.py +46 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_reorder.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_scalar_tmasave.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_ss.py +88 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_ss_regtune.py +97 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_ss_saved.py +90 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_static.py +51 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_tmasave.py +29 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b7.py +192 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_d64_b200gate.py +479 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_inference.py +257 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_native.py +184 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_output.py +188 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_single.py +274 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_single_b7.py +168 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_single_output.py +175 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/inference/common/tmn_math.cuh +105 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/inference/tmn_k3_wide.cuh +543 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/inference/tmn_kernels.cuh +1281 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/inference/tmn_sm90_unit.cuh +238 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_training.py +330 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_uni_d64_b200gate.py +427 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_uni_wide_inference.py +221 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_uni_wide_training.py +512 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_wide_inference.py +287 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_wide_training.py +1149 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/b1_sm80.cuh +397 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/b1g_sm80.cuh +383 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/b7_sm80.cuh +211 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/b7j_sm80.cuh +440 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/b8_sm80.cuh +233 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/contract_fragment_lean.cu +198 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/contract_lean.cu +189 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/contract_sm80.cuh +179 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/contract_streamed.cu +185 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/contract_tuned.cu +197 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/finalize_sm80.cuh +77 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/k1_sm80.cuh +366 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/k3_sm80.cuh +355 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/ops.cu +398 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/pack_sm80.cuh +93 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_bwd.cuh +309 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_elt.cuh +43 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_fin.cuh +58 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_gemm.cuh +283 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_gemm32.cuh +172 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_k1.cuh +366 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_k1b.cuh +104 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_k3.cuh +179 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_k32.cuh +262 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_ops.cu +419 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_rows.cuh +203 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80/wide_train512_ops.cu +392 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80.py +597 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80_gate.py +112 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80_wide.py +524 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/cuda/sm80_wide512.py +499 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/triton/back.py +308 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/triton/back_fused.py +207 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/triton/backward_fused.py +388 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/triton/bidirectional.py +423 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/triton/contract.py +61 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/triton/gate_elem.py +311 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/triton/output_fused.py +218 -0
- miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj/triton/unidirectional.py +271 -0
- miniworld_engine-2.5.2/src/miniworld_engine/measure/accuracy.py +304 -0
- miniworld_engine-2.5.2/src/miniworld_engine/measure/config.py +478 -0
- miniworld_engine-2.5.2/src/miniworld_engine/measure/plot.py +527 -0
- miniworld_engine-2.5.2/src/miniworld_engine/measure/speed.py +443 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/__init__.py +140 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/_cueq.py +65 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/atom_dit.py +1559 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/attention.py +750 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/dit_variant_cuda.py +358 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/embedder.py +273 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/functional.py +372 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/init.py +546 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/init_presets/af3.yaml +59 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/init_presets/engine.yaml +6 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/init_presets/miniworld.yaml +132 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/mpnn/__init__.py +24 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/mpnn/_functional.py +237 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/mpnn/features.py +231 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/mpnn/layers.py +776 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/mpnn/loss.py +159 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/mpnn/module.py +435 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/msa.py +945 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/norm.py +526 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/pair_embed.py +507 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/pairformer.py +286 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/primitives.py +489 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/sampler_ops.py +370 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/smooth_lddt.py +116 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/token_dit.py +1075 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/token_pair_init.py +184 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/transition.py +526 -0
- miniworld_engine-2.5.2/src/miniworld_engine/modules/triangle.py +1681 -0
- miniworld_engine-2.5.2/src/miniworld_engine/requirements/requirements_kernel.csv +691 -0
- miniworld_engine-2.5.2/src/miniworld_engine/requirements/requirements_module.csv +147 -0
- miniworld_engine-2.5.2/src/miniworld_engine.egg-info/PKG-INFO +218 -0
- miniworld_engine-2.5.2/src/miniworld_engine.egg-info/SOURCES.txt +1192 -0
- miniworld_engine-2.5.2/src/miniworld_engine.egg-info/requires.txt +25 -0
- miniworld_engine-2.5.0/PKG-INFO +0 -93
- miniworld_engine-2.5.0/README.md +0 -61
- miniworld_engine-2.5.0/licenses/Anthropic-NOTICE +0 -292
- miniworld_engine-2.5.0/licenses/THIRD_PARTY_NOTICES.md +0 -46
- miniworld_engine-2.5.0/pyproject.toml +0 -155
- miniworld_engine-2.5.0/src/miniworld_engine/judge/__init__.py +0 -372
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_capture.py +0 -294
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_cast.py +0 -79
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_cast_cuda.py +0 -72
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_codegraph.py +0 -515
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_compile.py +0 -185
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_device.py +0 -47
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_dot_precision.py +0 -20
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_fastmath.py +0 -23
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_h100_master.py +0 -28
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_launch_config.py +0 -245
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_nvcc.py +0 -887
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_prebuild.py +0 -611
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_prebuild_cuda.py +0 -229
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_pretune.py +0 -1077
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_shape_key.py +0 -104
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_tf32.py +0 -175
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/_tiles.py +0 -100
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/adaln/triton/inference.py +0 -695
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/adaln/triton/ln_strided.py +0 -132
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/adaln/triton/training.py +0 -858
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/configs/h100_sm90.json +0 -2419
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/__init__.py +0 -562
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/apb/__init__.py +0 -83
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/apb/apb_rows.cu +0 -1868
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/__init__.py +0 -854
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/driver.py +0 -135
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/__init__.py +0 -1361
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/__init__.py +0 -1064
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/__init__.py +0 -805
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/cond90.cu +0 -119
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/cond_bwd90.cu +0 -134
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/cond_bwd_sh90.cu +0 -112
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/pack90.cu +0 -131
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/post90.cu +0 -132
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/post_bwd90.cu +0 -235
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/pre90.cu +0 -164
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/pre_bwd90.cu +0 -197
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/tr_bwd90.cu +0 -101
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/triton/main.py +0 -922
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/triton/memory_efficient.py +0 -460
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/augmented_attention/triton/window.py +0 -356
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_attention/triton/core_t.py +0 -147
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_attention/triton/gate_out.py +0 -279
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_attention/triton/main.py +0 -503
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_attention/triton/pair_bias.py +0 -192
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_dit/cuda/__init__.py +0 -1065
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_dit/cuda/bias_only_dit_rows.cu +0 -691
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_dit/cuda/bwd_fused.py +0 -757
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_dit/cuda/hoist.py +0 -62
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_dit/cuda/inf3_bf16.py +0 -466
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_dit/cuda/runner.py +0 -788
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_dit/cuda/tf32.py +0 -1115
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/bias_only_dit/cuda/train.py +0 -44
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/cuda/__init__.py +0 -525
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/cuda/gemm_swiglu.py +0 -200
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/cuda/token_dit_composed.py +0 -166
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/cuda/token_dit_rows.cu +0 -737
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/cuda/token_dit_train_rows.cu +0 -1347
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/cuda/train.py +0 -185
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/triton/composed.py +0 -317
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/triton/fwd_saveact.py +0 -265
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/triton/inference.py +0 -205
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/triton/token_dit_attn.py +0 -221
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/triton/token_dit_kernels.py +0 -165
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/triton/token_dit_runner.py +0 -503
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/conditioned_transition/triton/training.py +0 -793
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/dit_variants/cuda/__init__.py +0 -129
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/dit_variants/cuda/infer32.py +0 -372
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/dit_variants/cuda/rows.cu +0 -345
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/dit_variants/cuda/train.py +0 -740
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/dit_variants/cuda/train_rows.cu +0 -794
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/dit_variants/interface.py +0 -336
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/dit_variants/triton/portable.py +0 -226
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/gated_projection/triton/main.py +0 -317
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/gated_projection/triton/residual.py +0 -112
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/gated_projection/triton/swa.py +0 -121
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/__init__.py +0 -81
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/attention_pair_bias_b200.py +0 -657
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/bias_only_dit.py +0 -260
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/bias_only_dit_hoist.py +0 -269
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/bias_only_dit_train.py +0 -1908
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/csrc/opm_epilogue.cu +0 -2833
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/csrc/pair_wide.cu +0 -390
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/csrc/pwa_wide.cu +0 -700
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/local_dit.py +0 -106
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/local_dit_h100.py +0 -479
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/opm_train.py +0 -1159
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/pwa_train.py +0 -1126
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/pwa_wide.py +0 -322
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/token_dit.py +0 -408
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/token_dit_train.py +0 -1410
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/triattn_h100_sm90.py +0 -813
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/integrations/trimul_h100.py +0 -757
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/compile_native.py +0 -130
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/configs/h100_sm90.json +0 -4065
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/cuda/__init__.py +0 -14
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/cuda/general/__init__.py +0 -157
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/cuda/general/bind.cpp +0 -9
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/cuda/pipe/__init__.py +0 -148
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/cuda/pipe/ln_pipe.cu +0 -996
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/triton/lowreg.py +0 -59
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/triton/main.py +0 -427
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/triton/persistent.py +0 -269
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm/triton/transpose.py +0 -145
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm_linear/autograd.py +0 -150
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm_linear/triton/fused.py +0 -356
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm_linear/triton/mmajor_bwd.py +0 -269
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm_linear/triton/pair_bias.py +0 -472
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm_linear/triton/recompute.py +0 -109
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm_linear/triton/stats.py +0 -104
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm_linear/triton/te_style.py +0 -359
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_edge_dropout/triton/main.py +0 -185
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_edge_layernorm/triton/main.py +0 -149
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_edge_mlp/triton/compute.py +0 -260
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_edge_mlp/triton/main.py +0 -318
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_edge_tail/configs/h100_sm90.json +0 -211
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_edge_tail/triton/compute.py +0 -794
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_edge_tail/triton/main.py +0 -1115
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_message/configs/h100_sm90.json +0 -137
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_message/triton/inference.py +0 -198
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_message/triton/main.py +0 -911
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_node_message/configs/h100_sm90.json +0 -57
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_node_message/triton/main.py +0 -816
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_relative_position/configs/h100_sm90.json +0 -33
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/mpnn_relative_position/triton/main.py +0 -348
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/outer_product_mean/chai.py +0 -170
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/outer_product_mean/composed.py +0 -116
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/outer_product_mean/cuda/__init__.py +0 -81
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/outer_product_mean/triton/main.py +0 -783
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_adaln/configs/h100_sm90.json +0 -87
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_adaln/cuda/__init__.py +0 -11
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_adaln/cuda/function.py +0 -70
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_adaln/cuda/sm90/__init__.py +0 -56
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_adaln/cuda/sm90/pair_embed_adaln.cu +0 -417
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_adaln/triton/function.py +0 -31
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_adaln/triton/main.py +0 -241
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_geom/configs/h100_sm90.json +0 -217
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_geom/cuda/__init__.py +0 -11
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_geom/cuda/function.py +0 -121
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_geom/cuda/sm90/__init__.py +0 -60
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_geom/cuda/sm90/pair_embed_geom.cu +0 -813
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_geom/triton/function.py +0 -53
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_embed_geom/triton/main.py +0 -401
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_weighted_averaging/configs/h100_sm90.json +0 -1389
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_weighted_averaging/triton/main.py +0 -532
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/rmsnorm/triton/main.py +0 -763
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/rope/triton/main.py +0 -131
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/rope/triton/qk_norm.py +0 -410
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/sampler_ops/configs/h100_sm90.json +0 -155
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/sampler_ops/cuda/__init__.py +0 -29
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/sampler_ops/cuda/function.py +0 -71
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/sampler_ops/cuda/sampler_ops.cu +0 -663
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/sampler_ops/triton/__init__.py +0 -9
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/sampler_ops/triton/function.py +0 -71
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/sampler_ops/triton/main.py +0 -391
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/smooth_lddt/configs/h100_sm90.json +0 -133
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/smooth_lddt/cuda/__init__.py +0 -12
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/smooth_lddt/cuda/function.py +0 -76
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/smooth_lddt/cuda/sm100/__init__.py +0 -139
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/smooth_lddt/cuda/sm100/smooth_lddt.cu +0 -338
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/smooth_lddt/triton/function.py +0 -38
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/smooth_lddt/triton/main.py +0 -268
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/autograd.py +0 -75
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/cuda/__init__.py +0 -38
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/cuda/loader.py +0 -124
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/cuda/sm100/__init__.py +0 -1291
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/cuda/sm100/tf32_fwd.py +0 -508
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/cuda/swa_fp32.cu +0 -1767
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/dispatch.py +0 -1495
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/triton/backward.py +0 -352
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/triton/backward_fp32.py +0 -210
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/triton/forward.py +0 -173
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/swa_dit/triton/forward_fp32.py +0 -133
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/tm1/triton/main.py +0 -405
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/tm2/triton/main.py +0 -300
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/token_pair_init/configs/h100_sm90.json +0 -105
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/token_pair_init/cuda/__init__.py +0 -11
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/token_pair_init/cuda/function.py +0 -81
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/token_pair_init/cuda/sm90/__init__.py +0 -60
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/token_pair_init/triton/function.py +0 -30
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/token_pair_init/triton/main.py +0 -214
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/configs/h100_sm90.json +0 -4396
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/cuda/__init__.py +0 -127
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/cuda/fused_generic_sm90a.py +0 -561
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/cuda/fused_sm100a.py +0 -390
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/cuda/fused_sm80.py +0 -385
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/cuda/fused_sm90a.py +0 -471
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/cuda/fused_wide_sm100a.py +0 -771
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/cuda/fused_wide_sm90a.py +0 -587
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/cuda/generic_sm90a/gen_wgmma.py +0 -35
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/cuda/sm80/tr_bwd_sm80.cuh +0 -504
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/cuda/sm80/tr_fwd_sm80.cuh +0 -461
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/triton/fused.py +0 -1603
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/triton/main.py +0 -333
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/triton/residual.py +0 -189
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition/triton/training.py +0 -176
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/cuda/__init__.py +0 -54
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/cuda/b200_triattn.py +0 -258
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/cuda/h100_dh.py +0 -124
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sm90.py +0 -401
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sm90_generic.py +0 -667
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/front_fused3_sm90.cu +0 -253
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/triattn_dh_sm90.cu +0 -233
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/triattn_h16_sm90.cu +0 -544
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/triton/atomic.py +0 -674
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/triton/bidirectional.py +0 -218
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/triton/main.py +0 -488
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/triangle_attention/triton/variants.py +0 -83
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/configs/h100_sm90.json +0 -3526
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/__init__.py +0 -245
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_bmm.py +0 -265
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_driver.py +0 -782
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_infer_kernel.py +0 -1040
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_launch.py +0 -891
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_manifest.py +0 -171
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_pack.py +0 -562
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/_h100_runtime.py +0 -449
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_infer.py +0 -151
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_train.py +0 -373
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b1.py +0 -218
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style.py +0 -733
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_aux.py +0 -153
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_d384_tmasave.py +0 -31
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_direct.py +0 -90
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_dpl_pipeline.py +0 -82
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_emitxn_tma.py +0 -59
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_k1_prefetch.py +0 -56
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_reorder.py +0 -31
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_scalar_tmasave.py +0 -31
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_ss.py +0 -99
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_ss_regtune.py +0 -106
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_ss_saved.py +0 -101
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_static.py +0 -53
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_tmasave.py +0 -31
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b7.py +0 -198
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_d64_b200gate.py +0 -368
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_inference.py +0 -249
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_native.py +0 -215
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_output.py +0 -198
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_single.py +0 -273
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_single_b7.py +0 -179
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_single_output.py +0 -189
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/inference/common/tmn_math.cuh +0 -104
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/inference/tmn_k3_wide.cuh +0 -542
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/inference/tmn_kernels.cuh +0 -1280
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/inference/tmn_sm90_unit.cuh +0 -237
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_training.py +0 -302
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_uni_d64_b200gate.py +0 -314
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_uni_wide_inference.py +0 -194
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_uni_wide_training.py +0 -419
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_wide_inference.py +0 -256
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_wide_training.py +0 -1049
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/triton/back.py +0 -284
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/triton/back_fused.py +0 -213
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/triton/backward_fused.py +0 -318
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/triton/bidirectional.py +0 -396
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/triton/contract.py +0 -59
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/triton/gate_elem.py +0 -285
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/triton/output_fused.py +0 -184
- miniworld_engine-2.5.0/src/miniworld_engine/kernels/trimul_inproj/triton/unidirectional.py +0 -266
- miniworld_engine-2.5.0/src/miniworld_engine/measure/accuracy.py +0 -287
- miniworld_engine-2.5.0/src/miniworld_engine/measure/config.py +0 -414
- miniworld_engine-2.5.0/src/miniworld_engine/measure/speed.py +0 -405
- miniworld_engine-2.5.0/src/miniworld_engine/modules/__init__.py +0 -136
- miniworld_engine-2.5.0/src/miniworld_engine/modules/_cueq.py +0 -65
- miniworld_engine-2.5.0/src/miniworld_engine/modules/atom_dit.py +0 -1099
- miniworld_engine-2.5.0/src/miniworld_engine/modules/attention.py +0 -642
- miniworld_engine-2.5.0/src/miniworld_engine/modules/dit_variant_cuda.py +0 -283
- miniworld_engine-2.5.0/src/miniworld_engine/modules/embedder.py +0 -264
- miniworld_engine-2.5.0/src/miniworld_engine/modules/functional.py +0 -371
- miniworld_engine-2.5.0/src/miniworld_engine/modules/mpnn/__init__.py +0 -24
- miniworld_engine-2.5.0/src/miniworld_engine/modules/mpnn/_functional.py +0 -208
- miniworld_engine-2.5.0/src/miniworld_engine/modules/mpnn/features.py +0 -212
- miniworld_engine-2.5.0/src/miniworld_engine/modules/mpnn/layers.py +0 -635
- miniworld_engine-2.5.0/src/miniworld_engine/modules/mpnn/loss.py +0 -173
- miniworld_engine-2.5.0/src/miniworld_engine/modules/mpnn/module.py +0 -396
- miniworld_engine-2.5.0/src/miniworld_engine/modules/msa.py +0 -446
- miniworld_engine-2.5.0/src/miniworld_engine/modules/norm.py +0 -351
- miniworld_engine-2.5.0/src/miniworld_engine/modules/pair_embed.py +0 -425
- miniworld_engine-2.5.0/src/miniworld_engine/modules/pairformer.py +0 -243
- miniworld_engine-2.5.0/src/miniworld_engine/modules/primitives.py +0 -423
- miniworld_engine-2.5.0/src/miniworld_engine/modules/sampler_ops.py +0 -266
- miniworld_engine-2.5.0/src/miniworld_engine/modules/smooth_lddt.py +0 -112
- miniworld_engine-2.5.0/src/miniworld_engine/modules/token_dit.py +0 -677
- miniworld_engine-2.5.0/src/miniworld_engine/modules/token_pair_init.py +0 -152
- miniworld_engine-2.5.0/src/miniworld_engine/modules/transition.py +0 -303
- miniworld_engine-2.5.0/src/miniworld_engine/modules/triangle.py +0 -1179
- miniworld_engine-2.5.0/src/miniworld_engine/requirements/requirements_kernel.csv +0 -491
- miniworld_engine-2.5.0/src/miniworld_engine/requirements/requirements_module.csv +0 -147
- miniworld_engine-2.5.0/src/miniworld_engine.egg-info/PKG-INFO +0 -93
- miniworld_engine-2.5.0/src/miniworld_engine.egg-info/SOURCES.txt +0 -665
- miniworld_engine-2.5.0/src/miniworld_engine.egg-info/requires.txt +0 -22
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/LICENSE +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/licenses/Anthropic-Apache-2.0.txt +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/licenses/CUTLASS-LICENSE.txt +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/licenses/FlashAttention-LICENSE.txt +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/setup.cfg +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/judge/__main__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/_cast_cuda.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/adaln/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/adaln/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/adaln/cuda/row_math.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/adaln/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/attn_dkv.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/attn_dq.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/attn_dqb.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/attn_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/attn_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/pair_rows.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/attn_dkv.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/attn_dkv_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/attn_dqb.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/attn_dqb_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/attn_fwd2.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/attn_fwd_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/attn_inf.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/attn_inf_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/glue.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100/sm100.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/cond_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/cond_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/dw_atom.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/fused_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/gemm_rg_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/gemm_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/post_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/post_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/pre_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/pre_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/rows_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/sm100.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom/tr_bwd_gate.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lattn90.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lattn90_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lattn_dkv.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lattn_dkv_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lattn_dq.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lattn_dq_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lattn_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lattn_fwd_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lbias.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lbias_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/lcross.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm100_atom_local/local.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/finalize90.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/rms90.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/row90.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/cuda/sm90_atom/wg90.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/augmented_attention/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_attention/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_attention/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_attention/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bias_only_dit_f32_rows.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bias_only_dit_hoist_rows.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bias_only_dit_train_rows.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_bwd_mid.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_bwd_tail.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_front_bf16.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_front_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_gemm_sm90.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_ptx.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_pvdpb.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_rows.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_tail2_bf16.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_tail2_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_tail_bf16.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_tail_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/bo_wgrad.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/cond_tables_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/dpb32_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/dpb_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/dpbx2_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/gemm_glu_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/gemm_resln_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/pv_gate_inf.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/cuda/pv_gate_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/bias_only_dit/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/conditioned_transition/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/conditioned_transition/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/conditioned_transition/cuda/gemm_swiglu2_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/conditioned_transition/cuda/gemm_swiglu_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/conditioned_transition/cuda/gemm_swiglu_sm90.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/conditioned_transition/cuda/token_dit_common.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/conditioned_transition/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/dit_variants/cuda/edge.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/dit_variants/reference.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/dit_variants/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/gated_projection/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/gated_projection/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0/src/miniworld_engine/kernels/layernorm_linear/triton → miniworld_engine-2.5.2/src/miniworld_engine/kernels/gated_projection/cuda}/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/gated_projection/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/dgv_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/ln_vg.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/pair3.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/pwa_ctr.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/pwa_fwd2.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/pwa_fwd3.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/pwa_gateout.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/pwa_glue3.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/sm100/opm_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/sm100/pwa_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/sm100/sm100.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/sm90_ptx.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/integrations/csrc/wgmma_n64.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/layernorm/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/layernorm/cuda/general/ln_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/layernorm/cuda/general/ln_common.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/layernorm/cuda/general/ln_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/layernorm/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/layernorm_linear/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/layernorm_linear/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0/src/miniworld_engine/kernels/outer_product_mean → miniworld_engine-2.5.2/src/miniworld_engine/kernels/layernorm_linear}/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_edge_dropout/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_edge_dropout/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_edge_layernorm/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_edge_mlp/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_edge_mlp/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_edge_tail/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_edge_tail/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_message/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_message/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_message/triton/_policy.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_node_message/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_node_message/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_relative_position/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/mpnn_relative_position/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/outer_product_mean/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/outer_product_mean/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/outer_product_mean/cuda/chai.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/outer_product_mean/cuda/chai_helpers.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/outer_product_mean/cuda/gather.cu +0 -0
- {miniworld_engine-2.5.0/src/miniworld_engine/kernels/pair_weighted_averaging → miniworld_engine-2.5.2/src/miniworld_engine/kernels/outer_product_mean}/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/pair_embed_adaln/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/pair_embed_adaln/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/pair_embed_geom/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/pair_embed_geom/cuda/sm90/pe_sm90.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/pair_embed_geom/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/pair_weighted_averaging/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0/src/miniworld_engine/kernels/rmsnorm → miniworld_engine-2.5.2/src/miniworld_engine/kernels/pair_weighted_averaging}/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/rmsnorm/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/rmsnorm/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/rmsnorm/cuda/__init__.py +0 -0
- {miniworld_engine-2.5.0/src/miniworld_engine/kernels/rope → miniworld_engine-2.5.2/src/miniworld_engine/kernels/rmsnorm}/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/rope/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/rope/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0/src/miniworld_engine/kernels/smooth_lddt → miniworld_engine-2.5.2/src/miniworld_engine/kernels/rope}/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/sampler_ops/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/smooth_lddt/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/_common.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/gen/wgmma_bf16.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/gen/wgmma_tf32.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/pdl.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/attn_dkv.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/attn_dkvq.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/attn_dq.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/attn_fwd1h.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/attn_fwd3.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/attn_fwd_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/ffn_bwd_dy.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/ffn_bwd_gate.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/ffn_fwd2.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/ffn_fwd_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/mod_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/mod_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/mod_fwd_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/oproj_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/qkvg_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/qkvg_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/qkvg_fwd2.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/qkvg_fwd_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/sm100.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/sm100/wprep_tf32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/swa_attn.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/swa_ffn_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/swa_ffn_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/swa_mod.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/swa_ptx.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/swa_qkvg_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/swa_slab.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/swa_wgrad.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/swa_wgrad_wg.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/cuda/swa_wgrad_wg_f32.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/swa_dit/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/tm1/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/tm1/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/tm1/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/tm2/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/tm2/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/tm2/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/token_pair_init/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/token_pair_init/cuda/sm100/token_pair_init.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/token_pair_init/cuda/sm90/token_pair_init.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/token_pair_init/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/_weights.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/anthropic_v5/common/tmn_math.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/anthropic_v5/tmn_kernels.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/anthropic_v5/tmn_ptx.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/generic_sm90a/bind.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/generic_sm90a/tr90.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/generic_sm90a/wgmma_ops.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/sm100.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/tr_bwd_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/tr_fwd_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/transition_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tbwd_d384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tbwd_d64.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tfwd_d256.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tfwd_d384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tfwd_d64.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tgate_ab.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tgate_w.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tgemm_nd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tln_w.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tlnbwd_w.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tsmall_w.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm100/widths/tswiglu_w.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm80/tr_bwd_dw_sm80.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm80/tr_bwd_pw_sm80.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm80/tr_bwd_w_sm80.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/sm80/transition_sm80.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/transition_fused_bwd_sm90a_kernel.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/transition_fused_fwd_sm90a_kernel.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/transition_fused_sm90a.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/transition_wgmma.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/bind.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/d256_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/d64_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/d64_fwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/d64_fwd_small_f2.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/d64_fwd_small_f3.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/dxln.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/gate.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/gate_noh.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/kernels/dxn_lnbwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/kernels/gate_gemm2.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/kernels/ln_swiglu_gemm.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/kernels/squeeze_gemm.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/kernels/transition_bwd_d64.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/kernels/transition_fwd_d256.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/kernels/transition_fwd_d64.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/kernels/wgmma256rs.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/kernels/wgmma_n.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/launch_util.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/lnsg.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/cuda/wide/squeeze.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/transition/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/configs/h100_sm90.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/b200_sources/sm100.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/b200_sources/triattn_mod_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/b200_sources/triattn_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/b200_sources/triattn_wide_sm100.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/fa3_utils.h +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/pack.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/sm90.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/triattn_bwd_sm90.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/triattn_generic_sm90.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/triattn_mod_sm90.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/h100_sources/triattn_sm90.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/cuda/h100_triattn.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/triangle_attention/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/configs/candidates.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/b1g.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/b1s.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/b7g.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/b7m.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/bind_b200.cpp +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/k1w.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/k1wb.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/k3g.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/k3w.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/sm100.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/tmap.h +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/wide_aux.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/b200_sources/wide_bwd.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_b200style_gemm.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b1/b1_fused.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b1/b1_lowreg.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b1/b1_pipeline_math.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b1/b1_stream_ln.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b1/configs.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b1/lowreg_stats.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b1/xhat_read.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b200style/k1ss.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b200style/k1ssregtune.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b200style/k1sssave.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b200style/k1wb_epilogue.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b200style/k3.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b200style/k3d384_tmasave.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b200style/k3reorder.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b200style/k3scalar_tmasave.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b200style/k3tmasave.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b7_384/joint.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b7_384/single_wg.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b7_768/joint.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/b7_768/single_wg.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/common/front_mn_primitives.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/common/front_primitives.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/common/warp_primitives.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/common_b7/common_recompute.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/common_b7/ln_recompute.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/common_b7/single_wg.inc +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/d64_train/d64_bwd_wg.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/d64_train/d64_wg.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/d64_train/save_k3_d64.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/d64_train/uni_bwd_wg.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/inference/defines.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/output/save_k3.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/single_b7/joint.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/single_output/kernel.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/single_output/selection.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/uni_wide/output.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/uni_wide/selection.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/upstream/csrc/tmn_kernels.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_forward/front_shared.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_forward/output.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_forward/output_stream.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_forward/selection.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_forward/wgmma_n.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/compact_projection_d512.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/compact_projection_fallback_d512.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/contract_d256_l384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/contract_gp_d384_l384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/contract_gp_d512_l384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/contract_gp_l768.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/delta_norm_d512.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/delta_norm_fallback_d512.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/forward_norm_d384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/front_d256.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/front_pre_d384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/front_pre_d512.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/front_pre_ss.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/gate.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/gate_affine.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/input_ln_d384_l384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/input_ln_l384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/input_reduce_d512_l768.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/input_reduce_l768.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/ln_in_d512.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/lt_selection.json +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/output_dense.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/output_ln_d256_l384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/output_ln_d384_l384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/output_ln_d512_l384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/output_ln_d512_l768.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/output_ln_l768.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/source_d256_l384.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/source_d256_l768.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/weight_reduce_d512_l768.cu +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/cuda/h100_sources/wide_train/wgmma_off.cuh +0 -0
- {miniworld_engine-2.5.0/src/miniworld_engine/kernels/transition → miniworld_engine-2.5.2/src/miniworld_engine/kernels/trimul_inproj}/cuda/sm80/sm80_common.cuh +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/kernels/trimul_inproj/triton/__init__.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/measure/configs/accuracy.yaml +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/measure/configs/all.yaml +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/measure/configs/speed.yaml +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine/modules/_anthropic.py +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine.egg-info/dependency_links.txt +0 -0
- {miniworld_engine-2.5.0 → miniworld_engine-2.5.2}/src/miniworld_engine.egg-info/top_level.txt +0 -0
|
@@ -0,0 +1,218 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: miniworld-engine
|
|
3
|
+
Version: 2.5.2
|
|
4
|
+
Summary: H100 + A100 CUDA / Triton kernels for AlphaFold3-style modules, with a PyTorch reference backend per module, accuracy and speed measurers, and a version-based judge.
|
|
5
|
+
License-Expression: MIT AND Apache-2.0 AND BSD-3-Clause
|
|
6
|
+
Requires-Python: >=3.10
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
License-File: licenses/THIRD_PARTY_NOTICES.md
|
|
10
|
+
License-File: licenses/Anthropic-Apache-2.0.txt
|
|
11
|
+
License-File: licenses/Anthropic-NOTICE
|
|
12
|
+
License-File: licenses/anthropic-trimul-native-v5-UPSTREAM.json
|
|
13
|
+
License-File: licenses/CUTLASS-LICENSE.txt
|
|
14
|
+
License-File: licenses/FlashAttention-LICENSE.txt
|
|
15
|
+
Requires-Dist: torch>=2.8
|
|
16
|
+
Requires-Dist: triton>=3.3
|
|
17
|
+
Requires-Dist: einops
|
|
18
|
+
Requires-Dist: jaxtyping
|
|
19
|
+
Requires-Dist: numpy
|
|
20
|
+
Requires-Dist: scipy
|
|
21
|
+
Requires-Dist: pydantic
|
|
22
|
+
Requires-Dist: pyyaml
|
|
23
|
+
Requires-Dist: cuda-bindings
|
|
24
|
+
Requires-Dist: ninja
|
|
25
|
+
Requires-Dist: miniworld-engine-kernels-sm90-cu129==2.5.2; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
|
|
26
|
+
Requires-Dist: miniworld-engine-kernels-sm80-cu129==2.5.2; sys_platform == "linux" and platform_machine == "x86_64" and python_version == "3.12"
|
|
27
|
+
Provides-Extra: baselines
|
|
28
|
+
Requires-Dist: cuequivariance-torch<0.13,>=0.12; extra == "baselines"
|
|
29
|
+
Provides-Extra: flash
|
|
30
|
+
Requires-Dist: flash-attn-4; extra == "flash"
|
|
31
|
+
Provides-Extra: dev
|
|
32
|
+
Requires-Dist: ruff==0.16.9; extra == "dev"
|
|
33
|
+
Requires-Dist: pyright; extra == "dev"
|
|
34
|
+
Requires-Dist: matplotlib; extra == "dev"
|
|
35
|
+
Dynamic: license-file
|
|
36
|
+
|
|
37
|
+
# miniworld-engine
|
|
38
|
+
|
|
39
|
+
GPU kernels for the building blocks of AlphaFold 3-style structure-prediction models: the pair stack (triangle multiplication and
|
|
40
|
+
attention, outer product mean, MSA pair-weighted averaging, transitions, layer norms), the token and atom diffusion transformers of
|
|
41
|
+
several model families, the diffusion sampler's glue, the smooth-lDDT loss, and ProteinMPNN. Each op is a PyTorch module with a
|
|
42
|
+
**PyTorch reference** and interchangeable backends: the engine's own hand-written **CUDA** kernels, its portable **Triton** kernels,
|
|
43
|
+
and two external baselines (Anthropic's released kernels, NVIDIA cuEquivariance). The repository also contains the tooling that
|
|
44
|
+
says what each module must serve (requirements tables), measures accuracy against a float64 gold and speed under CUDA graphs, and
|
|
45
|
+
decides, from human-managed version numbers, which results are current (the judge).
|
|
46
|
+
|
|
47
|
+
## Supported cards
|
|
48
|
+
|
|
49
|
+
| card | key | `cuda` backend | status |
|
|
50
|
+
|---|---|---|---|
|
|
51
|
+
| NVIDIA H100 | `h100_sm90` | sm_90a kernels (wgmma, TMA, CUTLASS / cuBLASDx where used) | supported: every requirement case has a current result (judge: 0 owed) ([benchmarks](benchmarks/h100_sm90/README.md)) |
|
|
52
|
+
| NVIDIA A100 | `a100_sm80` | sm_80 kernels (mma.sync) | supported: every requirement case has a current result (judge: 0 owed) ([benchmarks](benchmarks/a100_sm80/README.md)) |
|
|
53
|
+
| NVIDIA B200 | `b200_sm100` | sm_100a sources exist in the tree | **not supported**: almost every `cuda` backend refuses the card; no launch-config tables, tokens or benchmarks |
|
|
54
|
+
| other CUDA GPUs | — | refused (`BackendUnsupported`) | the `triton` and `pytorch` backends run; Triton launch configs are tuned only for H100 / A100 |
|
|
55
|
+
|
|
56
|
+
## Modules and the models that use them
|
|
57
|
+
|
|
58
|
+
Every class is in `miniworld_engine.modules`. The models are the ones whose constructors produce the class's requirement rows
|
|
59
|
+
(`source` column of [requirements_module.csv](src/miniworld_engine/requirements/requirements_module.csv)): AF3-family models built by
|
|
60
|
+
FoldForge, and MiniWorld. `engine-own` rows are shapes the engine keeps although no current downstream constructor builds them.
|
|
61
|
+
|
|
62
|
+
| group | classes | used by (requirement `source`) |
|
|
63
|
+
|---|---|---|
|
|
64
|
+
| triangle multiplication | `TriangleMultiplication` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix (v1, v2), RoseTTAFold3 |
|
|
65
|
+
| | `TriangleMultiplicationRF3` | RoseTTAFold3 |
|
|
66
|
+
| | `BidirectionalTriangleMultiplication` | MiniWorld; engine-own |
|
|
67
|
+
| triangle attention | `TriangleAttention` | Boltz-2, OpenDDE, OpenFold3-preview2, Protenix (v1, v2) |
|
|
68
|
+
| | `TriangleAttentionAF3` | AF3, IntelliFold2, OpenFold3 |
|
|
69
|
+
| | `TriangleAttentionRF3` | RoseTTAFold3 |
|
|
70
|
+
| | `TriangleAttentionChai1`, `TriangleAttentionChai1Confidence` | Chai-1 |
|
|
71
|
+
| | `BidirectionalTriangleAttention` | engine-own |
|
|
72
|
+
| MSA ops | `OuterProductMean`, `MSAPairWeightedAveraging` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, RoseTTAFold3, MiniWorld |
|
|
73
|
+
| transition, norms | `Transition`, `LayerNorm` | all of the above models and MiniWorld |
|
|
74
|
+
| attention with pair bias | `AttentionPairBias` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, RoseTTAFold3, MiniWorld |
|
|
75
|
+
| pairformer block | `PairformerBlock` | engine-own (MiniWorld block layout) |
|
|
76
|
+
| token diffusion transformer | `TokenDiTProtenix` (selector `TokenDiT(model=...)`) | AF3, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, MiniWorld |
|
|
77
|
+
| | `TokenDiTBoltz2`, `TokenDiTRoseTTAFold3`, `TokenDiTChai1` | Boltz-2, RoseTTAFold3, Chai-1 |
|
|
78
|
+
| | `TokenDiTMiniWorld`, `TokenDiTMiniWorldBiasOnly` | MiniWorld |
|
|
79
|
+
| atom diffusion transformer | `AtomDiTOpenFold3` (selector `AtomDiT(model=...)`) | AF3, IntelliFold2, OpenFold3 |
|
|
80
|
+
| | `AtomDiTProtenix` | OpenDDE, Protenix |
|
|
81
|
+
| | `AtomDiTBoltz2`, `AtomDiTRoseTTAFold3`, `AtomDiTChai1` | Boltz-2, RoseTTAFold3, Chai-1 |
|
|
82
|
+
| | `AtomDiTESMFold2` | ESMFold2 |
|
|
83
|
+
| | `SWAAtomTransformer`, `InputFeatureEmbedder` | MiniWorld |
|
|
84
|
+
| diffusion sampler, loss | `AugmentNoise`, `AlignEuler`, `PoolAtoms`, `SmoothLddtLoss` | engine-own (AF3-style sampler glue and smooth-lDDT loss) |
|
|
85
|
+
| La-Proteina pair embeddings | `PairEmbedAdaLN`, `PairEmbedGeom`, `PairEmbedCA`, `TokenPairInit` | engine-own |
|
|
86
|
+
| inverse folding | `ProteinMPNN` | engine-own |
|
|
87
|
+
|
|
88
|
+
38 classes, 146 requirement rows. "MiniWorld" covers its confidence, diffusion and distogram models.
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
## Backends
|
|
92
|
+
|
|
93
|
+
| backend | what | accuracy rule |
|
|
94
|
+
|---|---|---|
|
|
95
|
+
| `pytorch` | the reference: the single source of truth of each op's math | the reference |
|
|
96
|
+
| `cuda` | the engine's hand-written kernels for the input's card (H100 or A100) | binding: GREEN = within `max(err_pytorch, fixed[dtype])` of the fp64 gold |
|
|
97
|
+
| `triton` | the engine's Triton kernels (no TMA), any CUDA card | binding |
|
|
98
|
+
| `anthropic` | Anthropic's [uplifting-biomolecular-modeling](https://github.com/anthropics/uplifting-biomolecular-modeling) release, connected as it ships | informational (baseline) |
|
|
99
|
+
| `cueq` | NVIDIA cuEquivariance 0.12 (`cuequivariance_torch`), connected as it ships | informational (baseline) |
|
|
100
|
+
|
|
101
|
+
A call outside a backend's range raises `BackendUnsupported` with the reason; it is never run by another backend. A kernel that fails
|
|
102
|
+
to build raises `KernelBuildError`; there is no silent fallback anywhere. Which classes have which backends: the `backends` tuple of
|
|
103
|
+
each class ([docs/modules.md](docs/modules.md)).
|
|
104
|
+
|
|
105
|
+
## Results
|
|
106
|
+
|
|
107
|
+
Per-module speed, geometric mean over each class's requirement cases of the speedup vs the `pytorch` reference (the engine's
|
|
108
|
+
reference implementation, eager PyTorch kernels in the same CUDA graph, TF32 allowed, not `torch.compile`d), from the speed tokens.
|
|
109
|
+
Each card's README has the full table, the AF3-only table and per-class L / D sweeps.
|
|
110
|
+
|
|
111
|
+
| | H100 ([summary](benchmarks/h100_sm90/README.md)) | A100 ([summary](benchmarks/a100_sm80/README.md)) |
|
|
112
|
+
|---|---|---|
|
|
113
|
+
| all classes |  |  |
|
|
114
|
+
|
|
115
|
+
A few rows (eval / train, Engine CUDA vs PyTorch; in brackets the Anthropic and cuEquivariance eval speedups vs PyTorch):
|
|
116
|
+
|
|
117
|
+
| class | H100 eval | H100 train | A100 eval | A100 train |
|
|
118
|
+
|---|---|---|---|---|
|
|
119
|
+
| TriangleMultiplication | 9.76x (4.64x, 3.39x) | 8.67x | 7.63x (5.44x, 3.74x) | 6.21x |
|
|
120
|
+
| TriangleAttention | 12.57x (6.53x, 3.11x) | 6.59x | 10.36x (5.49x, 3.19x) | 5.85x |
|
|
121
|
+
| Transition | 4.37x (2.62x, —) | 3.14x | 3.56x (2.47x, —) | 2.52x |
|
|
122
|
+
| AttentionPairBias | 4.54x (1.55x, 0.93x) | 3.41x | 4.33x (1.57x, 0.94x) | 2.28x |
|
|
123
|
+
| TokenDiTProtenix | 3.56x (1.22x, 1.02x) | 2.50x | 2.54x (1.19x, 0.84x) | 1.73x |
|
|
124
|
+
| AtomDiTOpenFold3 | 5.29x (1.19x, —) | 2.78x | 4.74x (1.21x, —) | 2.82x |
|
|
125
|
+
| PairformerBlock | 12.83x (8.24x, —) | 7.84x | 9.57x (6.92x, —) | 6.19x |
|
|
126
|
+
|
|
127
|
+
These are module-level numbers on the requirement shapes, not end-to-end model speedups. Accuracy: every engine-backend accuracy token
|
|
128
|
+
of both cards is in `src/miniworld_engine/judge/tokens/<card>/` (`kind: accuracy`, `light`, `err_backend`, `err_pytorch`, `tol`).
|
|
129
|
+
|
|
130
|
+
## Installation
|
|
131
|
+
|
|
132
|
+
See [docs/install.md](docs/install.md). In short:
|
|
133
|
+
|
|
134
|
+
1. Python >= 3.10 with PyTorch >= 2.8 and Triton >= 3.3; `pip install -e .` (the extra `baselines` adds cuEquivariance for `cueq`).
|
|
135
|
+
2. On Linux x86-64 with CPython 3.12 the prebuilt H100 kernels come as a dependency wheel; elsewhere each CUDA kernel is compiled on
|
|
136
|
+
first use (needs an `nvcc` matching torch's CUDA and, for the CUTLASS kernels, `CUTLASS_PATH`) or ahead of time with the prebuilder.
|
|
137
|
+
3. The `anthropic` backend needs the Anthropic release at `third_party/anthropic/upstream` (revision in `third_party/anthropic/UPSTREAM.json`)
|
|
138
|
+
or `MINIWORLD_ANTHROPIC_ROOT`.
|
|
139
|
+
|
|
140
|
+
## Quick usage
|
|
141
|
+
|
|
142
|
+
```python
|
|
143
|
+
import torch
|
|
144
|
+
from miniworld_engine.modules import TriangleMultiplication, TokenDiT
|
|
145
|
+
from miniworld_engine.modules.primitives import BackendUnsupported
|
|
146
|
+
|
|
147
|
+
tri = TriangleMultiplication(128, backend="cuda").cuda() # H100 or A100 kernels, picked from the input's device
|
|
148
|
+
ref = TriangleMultiplication(128).cuda() # backend="pytorch", the reference
|
|
149
|
+
ref.load_state_dict(tri.state_dict()) # same parameter names for every backend
|
|
150
|
+
tri.set_backend("triton") # switch later; ValueError if the class has no such backend
|
|
151
|
+
|
|
152
|
+
dit = TokenDiT(model="boltz2", ...) # selector: builds the model's own class (TokenDiTBoltz2)
|
|
153
|
+
```
|
|
154
|
+
|
|
155
|
+
Mixed precision as the models run it (`master` in the requirements): fp32 parameters, bf16 activations, forward under
|
|
156
|
+
`torch.autocast("cuda", torch.bfloat16)`; call `set_master_weights()` on a `MasterWeights` module after a bulk bf16 cast.
|
|
157
|
+
|
|
158
|
+
## Building, measuring, judging
|
|
159
|
+
|
|
160
|
+
Every Python step below runs on a compute node (the measurers need a GPU of the card; the prebuilder and the judge run on a CPU).
|
|
161
|
+
|
|
162
|
+
```sh
|
|
163
|
+
# ahead-of-time CUDA builds for a card, planned from the code and the requirements
|
|
164
|
+
python -m miniworld_engine.kernels._prebuild --list # H100 (sm90a): plan and what is missing
|
|
165
|
+
python -m miniworld_engine.kernels._prebuild --arch sm80 --jobs 16 --cutlass $CUTLASS_PATH
|
|
166
|
+
|
|
167
|
+
# which requirement cases lack a current result, as Slurm work groups
|
|
168
|
+
python -m miniworld_engine.judge --gpu h100_sm90 owed --plan $DIR
|
|
169
|
+
# for each line "<kind> <Module> <cfg>" of $DIR/G<i>, on the card:
|
|
170
|
+
python -m miniworld_engine.measure.<kind> <cfg> --out $DIR/out/<kind>.<Module>.jsonl
|
|
171
|
+
python -m miniworld_engine.judge --gpu h100_sm90 issue $DIR/out # records made by the current versions -> tokens
|
|
172
|
+
python -m miniworld_engine.judge --gpu h100_sm90 status
|
|
173
|
+
|
|
174
|
+
# figures and tables from the tokens (never measures)
|
|
175
|
+
python -m miniworld_engine.measure.plot --gpu a100_sm80
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
`tools/slurm/cycle.sh PLAN_DIR` runs one plan → prebuild → measure → issue cycle on a Slurm cluster (`CARD=a100_sm80 ARCH=sm80` for
|
|
179
|
+
A100). Details: [docs/judge.md](docs/judge.md), [docs/measure.md](docs/measure.md), [docs/kernels.md](docs/kernels.md).
|
|
180
|
+
|
|
181
|
+
## Repository layout
|
|
182
|
+
|
|
183
|
+
| path | what |
|
|
184
|
+
|---|---|
|
|
185
|
+
| `src/miniworld_engine/requirements/` | `requirements_module.csv` (what to serve), `requirements_kernel.csv` (which kernels serve it, with versions), the rules |
|
|
186
|
+
| `src/miniworld_engine/modules/` | the `BackendModule` classes, one file per op family; loaders of the external baselines |
|
|
187
|
+
| `src/miniworld_engine/kernels/` | ~40 kernel families (`cuda/`, `triton/`, `configs/` per card), `integrations/` (module-level compositions per card), the build, prebuild and pretune infrastructure |
|
|
188
|
+
| `src/miniworld_engine/measure/` | accuracy and speed measurers, their configs, the benchmark plotter |
|
|
189
|
+
| `src/miniworld_engine/judge/` | the judge, the tokens per card, the recorded answers |
|
|
190
|
+
| `benchmarks/<card>/` | generated speed figures and tables |
|
|
191
|
+
| `docs/` | documentation ([index](docs/README.md), [architecture](docs/architecture.md)); `docs/h100/` per-family H100 design notes |
|
|
192
|
+
| `tools/slurm/` | the measurement cycle script |
|
|
193
|
+
| `third_party/anthropic/` | provenance of the Anthropic release used by the `anthropic` backends |
|
|
194
|
+
|
|
195
|
+
Start with [docs/architecture.md](docs/architecture.md) for how a module call reaches the card's kernels and how the pieces fit.
|
|
196
|
+
|
|
197
|
+
## License
|
|
198
|
+
|
|
199
|
+
`MIT AND Apache-2.0 AND BSD-3-Clause`. MiniWorld Engine's own code is MIT (copyright 2026 Sanggeun Park, `LICENSE`); the code derived
|
|
200
|
+
from Anthropic's uplifting-biomolecular-modeling release, and any file whose header says `SPDX-License-Identifier: Apache-2.0`, is
|
|
201
|
+
Apache-2.0 (with Anthropic's `NOTICE`), and the FlashAttention-3 portions and the
|
|
202
|
+
CUTLASS-instantiated object code in the prebuilt kernels are BSD-3-Clause. Details, per-file attributions and the licence texts are in
|
|
203
|
+
[`licenses/THIRD_PARTY_NOTICES.md`](licenses/THIRD_PARTY_NOTICES.md). The external release used by the `anthropic` backends is cloned
|
|
204
|
+
locally under `third_party/anthropic/upstream/` (git-ignored, not redistributed).
|
|
205
|
+
|
|
206
|
+
## Citation
|
|
207
|
+
|
|
208
|
+
A paper describing the engine is in preparation. Until then, please cite the repository:
|
|
209
|
+
|
|
210
|
+
```bibtex
|
|
211
|
+
@software{miniworld_engine,
|
|
212
|
+
author = {Park, Sanggeun},
|
|
213
|
+
title = {miniworld-engine: GPU kernels for AlphaFold 3-style modules},
|
|
214
|
+
year = {2026},
|
|
215
|
+
version = {2.5.2},
|
|
216
|
+
note = {TODO: replace with the paper reference}
|
|
217
|
+
}
|
|
218
|
+
```
|
|
@@ -0,0 +1,182 @@
|
|
|
1
|
+
# miniworld-engine
|
|
2
|
+
|
|
3
|
+
GPU kernels for the building blocks of AlphaFold 3-style structure-prediction models: the pair stack (triangle multiplication and
|
|
4
|
+
attention, outer product mean, MSA pair-weighted averaging, transitions, layer norms), the token and atom diffusion transformers of
|
|
5
|
+
several model families, the diffusion sampler's glue, the smooth-lDDT loss, and ProteinMPNN. Each op is a PyTorch module with a
|
|
6
|
+
**PyTorch reference** and interchangeable backends: the engine's own hand-written **CUDA** kernels, its portable **Triton** kernels,
|
|
7
|
+
and two external baselines (Anthropic's released kernels, NVIDIA cuEquivariance). The repository also contains the tooling that
|
|
8
|
+
says what each module must serve (requirements tables), measures accuracy against a float64 gold and speed under CUDA graphs, and
|
|
9
|
+
decides, from human-managed version numbers, which results are current (the judge).
|
|
10
|
+
|
|
11
|
+
## Supported cards
|
|
12
|
+
|
|
13
|
+
| card | key | `cuda` backend | status |
|
|
14
|
+
|---|---|---|---|
|
|
15
|
+
| NVIDIA H100 | `h100_sm90` | sm_90a kernels (wgmma, TMA, CUTLASS / cuBLASDx where used) | supported: every requirement case has a current result (judge: 0 owed) ([benchmarks](benchmarks/h100_sm90/README.md)) |
|
|
16
|
+
| NVIDIA A100 | `a100_sm80` | sm_80 kernels (mma.sync) | supported: every requirement case has a current result (judge: 0 owed) ([benchmarks](benchmarks/a100_sm80/README.md)) |
|
|
17
|
+
| NVIDIA B200 | `b200_sm100` | sm_100a sources exist in the tree | **not supported**: almost every `cuda` backend refuses the card; no launch-config tables, tokens or benchmarks |
|
|
18
|
+
| other CUDA GPUs | — | refused (`BackendUnsupported`) | the `triton` and `pytorch` backends run; Triton launch configs are tuned only for H100 / A100 |
|
|
19
|
+
|
|
20
|
+
## Modules and the models that use them
|
|
21
|
+
|
|
22
|
+
Every class is in `miniworld_engine.modules`. The models are the ones whose constructors produce the class's requirement rows
|
|
23
|
+
(`source` column of [requirements_module.csv](src/miniworld_engine/requirements/requirements_module.csv)): AF3-family models built by
|
|
24
|
+
FoldForge, and MiniWorld. `engine-own` rows are shapes the engine keeps although no current downstream constructor builds them.
|
|
25
|
+
|
|
26
|
+
| group | classes | used by (requirement `source`) |
|
|
27
|
+
|---|---|---|
|
|
28
|
+
| triangle multiplication | `TriangleMultiplication` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix (v1, v2), RoseTTAFold3 |
|
|
29
|
+
| | `TriangleMultiplicationRF3` | RoseTTAFold3 |
|
|
30
|
+
| | `BidirectionalTriangleMultiplication` | MiniWorld; engine-own |
|
|
31
|
+
| triangle attention | `TriangleAttention` | Boltz-2, OpenDDE, OpenFold3-preview2, Protenix (v1, v2) |
|
|
32
|
+
| | `TriangleAttentionAF3` | AF3, IntelliFold2, OpenFold3 |
|
|
33
|
+
| | `TriangleAttentionRF3` | RoseTTAFold3 |
|
|
34
|
+
| | `TriangleAttentionChai1`, `TriangleAttentionChai1Confidence` | Chai-1 |
|
|
35
|
+
| | `BidirectionalTriangleAttention` | engine-own |
|
|
36
|
+
| MSA ops | `OuterProductMean`, `MSAPairWeightedAveraging` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, RoseTTAFold3, MiniWorld |
|
|
37
|
+
| transition, norms | `Transition`, `LayerNorm` | all of the above models and MiniWorld |
|
|
38
|
+
| attention with pair bias | `AttentionPairBias` | AF3, Boltz-2, Chai-1, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, RoseTTAFold3, MiniWorld |
|
|
39
|
+
| pairformer block | `PairformerBlock` | engine-own (MiniWorld block layout) |
|
|
40
|
+
| token diffusion transformer | `TokenDiTProtenix` (selector `TokenDiT(model=...)`) | AF3, ESMFold2, IntelliFold2, OpenDDE, OpenFold3, Protenix, MiniWorld |
|
|
41
|
+
| | `TokenDiTBoltz2`, `TokenDiTRoseTTAFold3`, `TokenDiTChai1` | Boltz-2, RoseTTAFold3, Chai-1 |
|
|
42
|
+
| | `TokenDiTMiniWorld`, `TokenDiTMiniWorldBiasOnly` | MiniWorld |
|
|
43
|
+
| atom diffusion transformer | `AtomDiTOpenFold3` (selector `AtomDiT(model=...)`) | AF3, IntelliFold2, OpenFold3 |
|
|
44
|
+
| | `AtomDiTProtenix` | OpenDDE, Protenix |
|
|
45
|
+
| | `AtomDiTBoltz2`, `AtomDiTRoseTTAFold3`, `AtomDiTChai1` | Boltz-2, RoseTTAFold3, Chai-1 |
|
|
46
|
+
| | `AtomDiTESMFold2` | ESMFold2 |
|
|
47
|
+
| | `SWAAtomTransformer`, `InputFeatureEmbedder` | MiniWorld |
|
|
48
|
+
| diffusion sampler, loss | `AugmentNoise`, `AlignEuler`, `PoolAtoms`, `SmoothLddtLoss` | engine-own (AF3-style sampler glue and smooth-lDDT loss) |
|
|
49
|
+
| La-Proteina pair embeddings | `PairEmbedAdaLN`, `PairEmbedGeom`, `PairEmbedCA`, `TokenPairInit` | engine-own |
|
|
50
|
+
| inverse folding | `ProteinMPNN` | engine-own |
|
|
51
|
+
|
|
52
|
+
38 classes, 146 requirement rows. "MiniWorld" covers its confidence, diffusion and distogram models.
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
## Backends
|
|
56
|
+
|
|
57
|
+
| backend | what | accuracy rule |
|
|
58
|
+
|---|---|---|
|
|
59
|
+
| `pytorch` | the reference: the single source of truth of each op's math | the reference |
|
|
60
|
+
| `cuda` | the engine's hand-written kernels for the input's card (H100 or A100) | binding: GREEN = within `max(err_pytorch, fixed[dtype])` of the fp64 gold |
|
|
61
|
+
| `triton` | the engine's Triton kernels (no TMA), any CUDA card | binding |
|
|
62
|
+
| `anthropic` | Anthropic's [uplifting-biomolecular-modeling](https://github.com/anthropics/uplifting-biomolecular-modeling) release, connected as it ships | informational (baseline) |
|
|
63
|
+
| `cueq` | NVIDIA cuEquivariance 0.12 (`cuequivariance_torch`), connected as it ships | informational (baseline) |
|
|
64
|
+
|
|
65
|
+
A call outside a backend's range raises `BackendUnsupported` with the reason; it is never run by another backend. A kernel that fails
|
|
66
|
+
to build raises `KernelBuildError`; there is no silent fallback anywhere. Which classes have which backends: the `backends` tuple of
|
|
67
|
+
each class ([docs/modules.md](docs/modules.md)).
|
|
68
|
+
|
|
69
|
+
## Results
|
|
70
|
+
|
|
71
|
+
Per-module speed, geometric mean over each class's requirement cases of the speedup vs the `pytorch` reference (the engine's
|
|
72
|
+
reference implementation, eager PyTorch kernels in the same CUDA graph, TF32 allowed, not `torch.compile`d), from the speed tokens.
|
|
73
|
+
Each card's README has the full table, the AF3-only table and per-class L / D sweeps.
|
|
74
|
+
|
|
75
|
+
| | H100 ([summary](benchmarks/h100_sm90/README.md)) | A100 ([summary](benchmarks/a100_sm80/README.md)) |
|
|
76
|
+
|---|---|---|
|
|
77
|
+
| all classes |  |  |
|
|
78
|
+
|
|
79
|
+
A few rows (eval / train, Engine CUDA vs PyTorch; in brackets the Anthropic and cuEquivariance eval speedups vs PyTorch):
|
|
80
|
+
|
|
81
|
+
| class | H100 eval | H100 train | A100 eval | A100 train |
|
|
82
|
+
|---|---|---|---|---|
|
|
83
|
+
| TriangleMultiplication | 9.76x (4.64x, 3.39x) | 8.67x | 7.63x (5.44x, 3.74x) | 6.21x |
|
|
84
|
+
| TriangleAttention | 12.57x (6.53x, 3.11x) | 6.59x | 10.36x (5.49x, 3.19x) | 5.85x |
|
|
85
|
+
| Transition | 4.37x (2.62x, —) | 3.14x | 3.56x (2.47x, —) | 2.52x |
|
|
86
|
+
| AttentionPairBias | 4.54x (1.55x, 0.93x) | 3.41x | 4.33x (1.57x, 0.94x) | 2.28x |
|
|
87
|
+
| TokenDiTProtenix | 3.56x (1.22x, 1.02x) | 2.50x | 2.54x (1.19x, 0.84x) | 1.73x |
|
|
88
|
+
| AtomDiTOpenFold3 | 5.29x (1.19x, —) | 2.78x | 4.74x (1.21x, —) | 2.82x |
|
|
89
|
+
| PairformerBlock | 12.83x (8.24x, —) | 7.84x | 9.57x (6.92x, —) | 6.19x |
|
|
90
|
+
|
|
91
|
+
These are module-level numbers on the requirement shapes, not end-to-end model speedups. Accuracy: every engine-backend accuracy token
|
|
92
|
+
of both cards is in `src/miniworld_engine/judge/tokens/<card>/` (`kind: accuracy`, `light`, `err_backend`, `err_pytorch`, `tol`).
|
|
93
|
+
|
|
94
|
+
## Installation
|
|
95
|
+
|
|
96
|
+
See [docs/install.md](docs/install.md). In short:
|
|
97
|
+
|
|
98
|
+
1. Python >= 3.10 with PyTorch >= 2.8 and Triton >= 3.3; `pip install -e .` (the extra `baselines` adds cuEquivariance for `cueq`).
|
|
99
|
+
2. On Linux x86-64 with CPython 3.12 the prebuilt H100 kernels come as a dependency wheel; elsewhere each CUDA kernel is compiled on
|
|
100
|
+
first use (needs an `nvcc` matching torch's CUDA and, for the CUTLASS kernels, `CUTLASS_PATH`) or ahead of time with the prebuilder.
|
|
101
|
+
3. The `anthropic` backend needs the Anthropic release at `third_party/anthropic/upstream` (revision in `third_party/anthropic/UPSTREAM.json`)
|
|
102
|
+
or `MINIWORLD_ANTHROPIC_ROOT`.
|
|
103
|
+
|
|
104
|
+
## Quick usage
|
|
105
|
+
|
|
106
|
+
```python
|
|
107
|
+
import torch
|
|
108
|
+
from miniworld_engine.modules import TriangleMultiplication, TokenDiT
|
|
109
|
+
from miniworld_engine.modules.primitives import BackendUnsupported
|
|
110
|
+
|
|
111
|
+
tri = TriangleMultiplication(128, backend="cuda").cuda() # H100 or A100 kernels, picked from the input's device
|
|
112
|
+
ref = TriangleMultiplication(128).cuda() # backend="pytorch", the reference
|
|
113
|
+
ref.load_state_dict(tri.state_dict()) # same parameter names for every backend
|
|
114
|
+
tri.set_backend("triton") # switch later; ValueError if the class has no such backend
|
|
115
|
+
|
|
116
|
+
dit = TokenDiT(model="boltz2", ...) # selector: builds the model's own class (TokenDiTBoltz2)
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
Mixed precision as the models run it (`master` in the requirements): fp32 parameters, bf16 activations, forward under
|
|
120
|
+
`torch.autocast("cuda", torch.bfloat16)`; call `set_master_weights()` on a `MasterWeights` module after a bulk bf16 cast.
|
|
121
|
+
|
|
122
|
+
## Building, measuring, judging
|
|
123
|
+
|
|
124
|
+
Every Python step below runs on a compute node (the measurers need a GPU of the card; the prebuilder and the judge run on a CPU).
|
|
125
|
+
|
|
126
|
+
```sh
|
|
127
|
+
# ahead-of-time CUDA builds for a card, planned from the code and the requirements
|
|
128
|
+
python -m miniworld_engine.kernels._prebuild --list # H100 (sm90a): plan and what is missing
|
|
129
|
+
python -m miniworld_engine.kernels._prebuild --arch sm80 --jobs 16 --cutlass $CUTLASS_PATH
|
|
130
|
+
|
|
131
|
+
# which requirement cases lack a current result, as Slurm work groups
|
|
132
|
+
python -m miniworld_engine.judge --gpu h100_sm90 owed --plan $DIR
|
|
133
|
+
# for each line "<kind> <Module> <cfg>" of $DIR/G<i>, on the card:
|
|
134
|
+
python -m miniworld_engine.measure.<kind> <cfg> --out $DIR/out/<kind>.<Module>.jsonl
|
|
135
|
+
python -m miniworld_engine.judge --gpu h100_sm90 issue $DIR/out # records made by the current versions -> tokens
|
|
136
|
+
python -m miniworld_engine.judge --gpu h100_sm90 status
|
|
137
|
+
|
|
138
|
+
# figures and tables from the tokens (never measures)
|
|
139
|
+
python -m miniworld_engine.measure.plot --gpu a100_sm80
|
|
140
|
+
```
|
|
141
|
+
|
|
142
|
+
`tools/slurm/cycle.sh PLAN_DIR` runs one plan → prebuild → measure → issue cycle on a Slurm cluster (`CARD=a100_sm80 ARCH=sm80` for
|
|
143
|
+
A100). Details: [docs/judge.md](docs/judge.md), [docs/measure.md](docs/measure.md), [docs/kernels.md](docs/kernels.md).
|
|
144
|
+
|
|
145
|
+
## Repository layout
|
|
146
|
+
|
|
147
|
+
| path | what |
|
|
148
|
+
|---|---|
|
|
149
|
+
| `src/miniworld_engine/requirements/` | `requirements_module.csv` (what to serve), `requirements_kernel.csv` (which kernels serve it, with versions), the rules |
|
|
150
|
+
| `src/miniworld_engine/modules/` | the `BackendModule` classes, one file per op family; loaders of the external baselines |
|
|
151
|
+
| `src/miniworld_engine/kernels/` | ~40 kernel families (`cuda/`, `triton/`, `configs/` per card), `integrations/` (module-level compositions per card), the build, prebuild and pretune infrastructure |
|
|
152
|
+
| `src/miniworld_engine/measure/` | accuracy and speed measurers, their configs, the benchmark plotter |
|
|
153
|
+
| `src/miniworld_engine/judge/` | the judge, the tokens per card, the recorded answers |
|
|
154
|
+
| `benchmarks/<card>/` | generated speed figures and tables |
|
|
155
|
+
| `docs/` | documentation ([index](docs/README.md), [architecture](docs/architecture.md)); `docs/h100/` per-family H100 design notes |
|
|
156
|
+
| `tools/slurm/` | the measurement cycle script |
|
|
157
|
+
| `third_party/anthropic/` | provenance of the Anthropic release used by the `anthropic` backends |
|
|
158
|
+
|
|
159
|
+
Start with [docs/architecture.md](docs/architecture.md) for how a module call reaches the card's kernels and how the pieces fit.
|
|
160
|
+
|
|
161
|
+
## License
|
|
162
|
+
|
|
163
|
+
`MIT AND Apache-2.0 AND BSD-3-Clause`. MiniWorld Engine's own code is MIT (copyright 2026 Sanggeun Park, `LICENSE`); the code derived
|
|
164
|
+
from Anthropic's uplifting-biomolecular-modeling release, and any file whose header says `SPDX-License-Identifier: Apache-2.0`, is
|
|
165
|
+
Apache-2.0 (with Anthropic's `NOTICE`), and the FlashAttention-3 portions and the
|
|
166
|
+
CUTLASS-instantiated object code in the prebuilt kernels are BSD-3-Clause. Details, per-file attributions and the licence texts are in
|
|
167
|
+
[`licenses/THIRD_PARTY_NOTICES.md`](licenses/THIRD_PARTY_NOTICES.md). The external release used by the `anthropic` backends is cloned
|
|
168
|
+
locally under `third_party/anthropic/upstream/` (git-ignored, not redistributed).
|
|
169
|
+
|
|
170
|
+
## Citation
|
|
171
|
+
|
|
172
|
+
A paper describing the engine is in preparation. Until then, please cite the repository:
|
|
173
|
+
|
|
174
|
+
```bibtex
|
|
175
|
+
@software{miniworld_engine,
|
|
176
|
+
author = {Park, Sanggeun},
|
|
177
|
+
title = {miniworld-engine: GPU kernels for AlphaFold 3-style modules},
|
|
178
|
+
year = {2026},
|
|
179
|
+
version = {2.5.2},
|
|
180
|
+
note = {TODO: replace with the paper reference}
|
|
181
|
+
}
|
|
182
|
+
```
|