[Feature] Can BatchDecodeWithPagedKVCacheWrapper return attention scores to all tokens, not just logsumexp?

[yawnzh]

I'm using flashinfer for a text-to-speech model, and I need the attention score to get the alignment between the output(audio) to the input(text). I'm curious if it is possible to get the attention score to all tokens during decoding, not just the logsumexp?

[yzh119]

Yes that's feasible by defining your own attention variant:

https://github.com/flashinfer-ai/flashinfer/blob/main/tests/test_jit_example.py#L161-L216

But then you might lose the benefit of flashattention algorithm because of the O(n^2) write to global memory.