flash_attn_with_kvcach return block_lse or attention_score

[NonvolatileMemory]

Hi Dao,

The flash_attn_with_kvcach function currently does not return the intermediate outputs, such as block_lse or attention_score, which are useful for further analysis or debugging. Can you add an option to return these outputs would enhance flexibility and allow users to gain more insight into the attention computation process.

[tridao]

Yeah can you add a PR?

[NonvolatileMemory]

I have a triton version but the speed is 2x slower than your cuda code.

But I don't know how to impl it with cuda code. But it should not be very hard, especially for block_lse, so can you slightly modify the code to support this func?