| """Flash Attention CUTE (CUDA Template Engine) implementation.""" | |
| from importlib.metadata import PackageNotFoundError, version | |
| # Update when syncing again. | |
| __version__ = "4.0.0.beta4" | |
| import cutlass.cute as cute | |
| from .interface import ( | |
| flash_attn_func, | |
| flash_attn_varlen_func, | |
| ) | |
| from .cute_dsl_utils import cute_compile_patched | |
| # Patch cute.compile to optionally dump SASS | |
| cute.compile = cute_compile_patched | |
| __all__ = [ | |
| "flash_attn_func", | |
| "flash_attn_varlen_func", | |
| ] | |