facebookresearch / xformers

Hackable and optimized Transformers building blocks, supporting a composable construction.
https://facebookresearch.github.io/xformers/
Other
8.66k stars 614 forks source link

Enable complete BlockDiagonalGappyKeysMask and BlockDiagonalPaddedKey Support on CK #1108

Closed qianfengz closed 1 month ago

qianfengz commented 1 month ago

Some Change to the xformers/ops/fmha/ck.py so that correct seqstart_k and seqlen_k can be passed to the C++ codes for

And their derived attn-bias types.

The change has passed all tests in

pytest tests/test_mem_eff_attention.py::test_forward -k "not triton"  -k "not flshattF"  -k "not fa2F"