├── .github └── workflows │ └── code-check.yml ├── .gitignore ├── .pre-commit-config.yaml ├── LICENSE ├── README.md ├── README_cn.md ├── assets ├── headdim128-causal-A100.png ├── headdim64-causal-A100.png ├── logo │ ├── baai-flagopen.jpeg │ ├── horizontal-blue.png │ ├── horizontal-dark.png │ ├── horizontal-light.png │ ├── vertical-blue.png │ ├── vertical-dark.png │ └── vertical-light.png ├── piecewise_attention.png ├── piecewise_attention_interface.png └── v0.2 │ ├── flash_attention.png │ ├── flash_attention_d64.png │ └── piecewise_attention.png ├── benchmark ├── flash_benchmark.py ├── flash_decoding_benchmark.py ├── paged_benchmark.py └── piecewise_benchmark.py ├── examples ├── flash_attention_example.py ├── flash_attention_with_aux_outputs.py ├── paged_example.py ├── piecewise_example.py └── use_cutom_config_func.py ├── pyproject.toml ├── src └── flag_attn │ ├── __init__.py │ ├── dropout.py │ ├── flash.py │ ├── paged.py │ ├── piecewise.py │ ├── split_kv.py │ ├── testing │ ├── __init__.py │ ├── dropout.py │ ├── flash.py │ ├── paged.py │ └── piecewise.py │ └── total.py └── tests └── flag_attn ├── test_dropout.py ├── test_flash_attention.py ├── test_paged_attention.py └── test_piecewise_attention.py /.github/workflows/code-check.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/.github/workflows/code-check.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/.gitignore -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/.pre-commit-config.yaml -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/README.md -------------------------------------------------------------------------------- /README_cn.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/README_cn.md -------------------------------------------------------------------------------- /assets/headdim128-causal-A100.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/headdim128-causal-A100.png -------------------------------------------------------------------------------- /assets/headdim64-causal-A100.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/headdim64-causal-A100.png -------------------------------------------------------------------------------- /assets/logo/baai-flagopen.jpeg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/logo/baai-flagopen.jpeg -------------------------------------------------------------------------------- /assets/logo/horizontal-blue.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/logo/horizontal-blue.png -------------------------------------------------------------------------------- /assets/logo/horizontal-dark.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/logo/horizontal-dark.png -------------------------------------------------------------------------------- /assets/logo/horizontal-light.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/logo/horizontal-light.png -------------------------------------------------------------------------------- /assets/logo/vertical-blue.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/logo/vertical-blue.png -------------------------------------------------------------------------------- /assets/logo/vertical-dark.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/logo/vertical-dark.png -------------------------------------------------------------------------------- /assets/logo/vertical-light.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/logo/vertical-light.png -------------------------------------------------------------------------------- /assets/piecewise_attention.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/piecewise_attention.png -------------------------------------------------------------------------------- /assets/piecewise_attention_interface.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/piecewise_attention_interface.png -------------------------------------------------------------------------------- /assets/v0.2/flash_attention.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/v0.2/flash_attention.png -------------------------------------------------------------------------------- /assets/v0.2/flash_attention_d64.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/v0.2/flash_attention_d64.png -------------------------------------------------------------------------------- /assets/v0.2/piecewise_attention.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/assets/v0.2/piecewise_attention.png -------------------------------------------------------------------------------- /benchmark/flash_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/benchmark/flash_benchmark.py -------------------------------------------------------------------------------- /benchmark/flash_decoding_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/benchmark/flash_decoding_benchmark.py -------------------------------------------------------------------------------- /benchmark/paged_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/benchmark/paged_benchmark.py -------------------------------------------------------------------------------- /benchmark/piecewise_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/benchmark/piecewise_benchmark.py -------------------------------------------------------------------------------- /examples/flash_attention_example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/examples/flash_attention_example.py -------------------------------------------------------------------------------- /examples/flash_attention_with_aux_outputs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/examples/flash_attention_with_aux_outputs.py -------------------------------------------------------------------------------- /examples/paged_example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/examples/paged_example.py -------------------------------------------------------------------------------- /examples/piecewise_example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/examples/piecewise_example.py -------------------------------------------------------------------------------- /examples/use_cutom_config_func.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/examples/use_cutom_config_func.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/pyproject.toml -------------------------------------------------------------------------------- /src/flag_attn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/__init__.py -------------------------------------------------------------------------------- /src/flag_attn/dropout.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/dropout.py -------------------------------------------------------------------------------- /src/flag_attn/flash.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/flash.py -------------------------------------------------------------------------------- /src/flag_attn/paged.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/paged.py -------------------------------------------------------------------------------- /src/flag_attn/piecewise.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/piecewise.py -------------------------------------------------------------------------------- /src/flag_attn/split_kv.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/split_kv.py -------------------------------------------------------------------------------- /src/flag_attn/testing/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/testing/__init__.py -------------------------------------------------------------------------------- /src/flag_attn/testing/dropout.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/testing/dropout.py -------------------------------------------------------------------------------- /src/flag_attn/testing/flash.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/testing/flash.py -------------------------------------------------------------------------------- /src/flag_attn/testing/paged.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/testing/paged.py -------------------------------------------------------------------------------- /src/flag_attn/testing/piecewise.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/testing/piecewise.py -------------------------------------------------------------------------------- /src/flag_attn/total.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/src/flag_attn/total.py -------------------------------------------------------------------------------- /tests/flag_attn/test_dropout.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/tests/flag_attn/test_dropout.py -------------------------------------------------------------------------------- /tests/flag_attn/test_flash_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/tests/flag_attn/test_flash_attention.py -------------------------------------------------------------------------------- /tests/flag_attn/test_paged_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/tests/flag_attn/test_paged_attention.py -------------------------------------------------------------------------------- /tests/flag_attn/test_piecewise_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/flagos-ai/FlagAttention/HEAD/tests/flag_attn/test_piecewise_attention.py --------------------------------------------------------------------------------