ColossalAI/extensions/pybind/flash_attention
..
__init__.py
flash_attention_dao_cuda.py
flash_attention_npu.py
flash_attention_sdpa_cuda.py