├── .DS_Store ├── .gitignore ├── LICENSE ├── Licenses ├── KERNL ├── STABLE-FAST ├── TRIDENT └── TRITON ├── README.md ├── cheat_sheet.sh ├── implementations ├── ComfyUI │ ├── README.md │ └── example.py ├── Diffusers │ ├── README.md │ ├── load_sdxl_pipeline.py │ ├── load_unet.py │ └── unet_pt.py └── sgm_ │ ├── README.md │ ├── config.yaml │ └── example.py ├── optimizations.txt ├── requirements.txt ├── setup.cfg ├── setup.py └── src └── stabletriton ├── __init__.py ├── kernels ├── Conv_Kernels │ ├── README.md │ ├── conv_implicit_gemm.py │ ├── conv_torch_inductor.py │ ├── conv_trident.py │ └── conv_utils.py ├── __init__.py ├── attention_fa1.py ├── attention_fa2.py ├── attention_proj.py ├── dropout.py ├── geglu.py ├── groupnorm.py ├── layer_norm.py ├── linear.py ├── timestep.py └── timestep_bench.py ├── optimization.py ├── optimizers ├── __init__.py ├── cuda │ ├── cuda_graph.py │ ├── graphs.py │ ├── pool_cuda_graphs.py │ └── test_cuda_graphs.py ├── remove_dropout.py ├── replace_attention.py ├── replace_geglu.py ├── replace_groupnorm.py ├── replace_layernorm.py ├── replace_linear.py ├── replace_timesteps.py ├── unet_pt.py └── utils │ ├── fx.py │ └── util.py └── tests ├── flash_fft.py ├── test_attention.py ├── test_fft_conv.py ├── test_forward_overwrite.py └── torch_tut.py /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/.DS_Store -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/LICENSE -------------------------------------------------------------------------------- /Licenses/KERNL: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/Licenses/KERNL -------------------------------------------------------------------------------- /Licenses/STABLE-FAST: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/Licenses/STABLE-FAST -------------------------------------------------------------------------------- /Licenses/TRIDENT: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/Licenses/TRIDENT -------------------------------------------------------------------------------- /Licenses/TRITON: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/Licenses/TRITON -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/README.md -------------------------------------------------------------------------------- /cheat_sheet.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/cheat_sheet.sh -------------------------------------------------------------------------------- /implementations/ComfyUI/README.md: -------------------------------------------------------------------------------- 1 | # Example usage with ComfyUI (WIP) 2 | -------------------------------------------------------------------------------- /implementations/ComfyUI/example.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /implementations/Diffusers/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/implementations/Diffusers/README.md -------------------------------------------------------------------------------- /implementations/Diffusers/load_sdxl_pipeline.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/implementations/Diffusers/load_sdxl_pipeline.py -------------------------------------------------------------------------------- /implementations/Diffusers/load_unet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/implementations/Diffusers/load_unet.py -------------------------------------------------------------------------------- /implementations/Diffusers/unet_pt.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/implementations/Diffusers/unet_pt.py -------------------------------------------------------------------------------- /implementations/sgm_/README.md: -------------------------------------------------------------------------------- 1 | # Example usage with Automatic1111 (WIP) 2 | -------------------------------------------------------------------------------- /implementations/sgm_/config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/implementations/sgm_/config.yaml -------------------------------------------------------------------------------- /implementations/sgm_/example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/implementations/sgm_/example.py -------------------------------------------------------------------------------- /optimizations.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/optimizations.txt -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/requirements.txt -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/setup.cfg -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/setup.py -------------------------------------------------------------------------------- /src/stabletriton/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/stabletriton/kernels/Conv_Kernels/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/Conv_Kernels/README.md -------------------------------------------------------------------------------- /src/stabletriton/kernels/Conv_Kernels/conv_implicit_gemm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/Conv_Kernels/conv_implicit_gemm.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/Conv_Kernels/conv_torch_inductor.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/Conv_Kernels/conv_torch_inductor.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/Conv_Kernels/conv_trident.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/Conv_Kernels/conv_trident.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/Conv_Kernels/conv_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/Conv_Kernels/conv_utils.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/stabletriton/kernels/attention_fa1.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/attention_fa1.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/attention_fa2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/attention_fa2.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/attention_proj.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/attention_proj.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/dropout.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/dropout.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/geglu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/geglu.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/groupnorm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/groupnorm.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/layer_norm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/layer_norm.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/linear.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/linear.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/timestep.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/timestep.py -------------------------------------------------------------------------------- /src/stabletriton/kernels/timestep_bench.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/kernels/timestep_bench.py -------------------------------------------------------------------------------- /src/stabletriton/optimization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimization.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/__init__.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/cuda/cuda_graph.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/cuda/cuda_graph.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/cuda/graphs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/cuda/graphs.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/cuda/pool_cuda_graphs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/cuda/pool_cuda_graphs.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/cuda/test_cuda_graphs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/cuda/test_cuda_graphs.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/remove_dropout.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/remove_dropout.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/replace_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/replace_attention.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/replace_geglu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/replace_geglu.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/replace_groupnorm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/replace_groupnorm.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/replace_layernorm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/replace_layernorm.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/replace_linear.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/replace_linear.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/replace_timesteps.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/replace_timesteps.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/unet_pt.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/unet_pt.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/utils/fx.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/utils/fx.py -------------------------------------------------------------------------------- /src/stabletriton/optimizers/utils/util.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/optimizers/utils/util.py -------------------------------------------------------------------------------- /src/stabletriton/tests/flash_fft.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/tests/flash_fft.py -------------------------------------------------------------------------------- /src/stabletriton/tests/test_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/tests/test_attention.py -------------------------------------------------------------------------------- /src/stabletriton/tests/test_fft_conv.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/tests/test_fft_conv.py -------------------------------------------------------------------------------- /src/stabletriton/tests/test_forward_overwrite.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/tests/test_forward_overwrite.py -------------------------------------------------------------------------------- /src/stabletriton/tests/torch_tut.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/arnavdantuluri/StableTriton/HEAD/src/stabletriton/tests/torch_tut.py --------------------------------------------------------------------------------