├── .gitignore ├── .pre-commit-config.yaml ├── LICENSE ├── README.md ├── benchmark ├── benchmark_kvpacked_func.py └── benchmark_varlen_kvpacked_func.py ├── pyproject.toml ├── ring_flash_attn ├── __init__.py ├── adapters │ ├── __init__.py │ └── hf_adapter.py ├── llama3_flash_attn_varlen.py ├── ring_flash_attn.py ├── ring_flash_attn_varlen.py ├── stripe_flash_attn.py ├── triton_utils.py ├── utils.py ├── zigzag_ring_flash_attn.py └── zigzag_ring_flash_attn_varlen.py ├── setup.py └── test ├── test.sh ├── test_llama3_flash_attn_varlen_func.py ├── test_llama3_prepare_cu_seqlens.py ├── test_ring_flash_attn_func.py ├── test_ring_flash_attn_varlen_func.py ├── test_stripe_flash_attn_func.py ├── test_triton_kernels.py ├── test_zigzag_ring_flash_attn_func.py ├── test_zigzag_ring_flash_attn_varlen_func.py └── utils.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/.gitignore -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/.pre-commit-config.yaml -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/README.md -------------------------------------------------------------------------------- /benchmark/benchmark_kvpacked_func.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/benchmark/benchmark_kvpacked_func.py -------------------------------------------------------------------------------- /benchmark/benchmark_varlen_kvpacked_func.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/benchmark/benchmark_varlen_kvpacked_func.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/pyproject.toml -------------------------------------------------------------------------------- /ring_flash_attn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/__init__.py -------------------------------------------------------------------------------- /ring_flash_attn/adapters/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/adapters/__init__.py -------------------------------------------------------------------------------- /ring_flash_attn/adapters/hf_adapter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/adapters/hf_adapter.py -------------------------------------------------------------------------------- /ring_flash_attn/llama3_flash_attn_varlen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/llama3_flash_attn_varlen.py -------------------------------------------------------------------------------- /ring_flash_attn/ring_flash_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/ring_flash_attn.py -------------------------------------------------------------------------------- /ring_flash_attn/ring_flash_attn_varlen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/ring_flash_attn_varlen.py -------------------------------------------------------------------------------- /ring_flash_attn/stripe_flash_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/stripe_flash_attn.py -------------------------------------------------------------------------------- /ring_flash_attn/triton_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/triton_utils.py -------------------------------------------------------------------------------- /ring_flash_attn/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/utils.py -------------------------------------------------------------------------------- /ring_flash_attn/zigzag_ring_flash_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/zigzag_ring_flash_attn.py -------------------------------------------------------------------------------- /ring_flash_attn/zigzag_ring_flash_attn_varlen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/ring_flash_attn/zigzag_ring_flash_attn_varlen.py -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/setup.py -------------------------------------------------------------------------------- /test/test.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/test/test.sh -------------------------------------------------------------------------------- /test/test_llama3_flash_attn_varlen_func.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/test/test_llama3_flash_attn_varlen_func.py -------------------------------------------------------------------------------- /test/test_llama3_prepare_cu_seqlens.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/test/test_llama3_prepare_cu_seqlens.py -------------------------------------------------------------------------------- /test/test_ring_flash_attn_func.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/test/test_ring_flash_attn_func.py -------------------------------------------------------------------------------- /test/test_ring_flash_attn_varlen_func.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/test/test_ring_flash_attn_varlen_func.py -------------------------------------------------------------------------------- /test/test_stripe_flash_attn_func.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/test/test_stripe_flash_attn_func.py -------------------------------------------------------------------------------- /test/test_triton_kernels.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/test/test_triton_kernels.py -------------------------------------------------------------------------------- /test/test_zigzag_ring_flash_attn_func.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/test/test_zigzag_ring_flash_attn_func.py -------------------------------------------------------------------------------- /test/test_zigzag_ring_flash_attn_varlen_func.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/test/test_zigzag_ring_flash_attn_varlen_func.py -------------------------------------------------------------------------------- /test/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhuzilin/ring-flash-attention/HEAD/test/utils.py --------------------------------------------------------------------------------