flash_attn/flash_attn_interface.py · 4e8d60069f65573aca0626d1425513a7597655c2 · gaoqiong / flash-attention

Jianwei Dong authored Jul 08, 2024

Add the return_softmax_lse parameter to the flash_attn_with_kvcache function to allow returning the logsumexp of the attention scores. (#989)

4e8d6006

flash_attn_interface.py 44.7 KB