Do you know of any code framework that supports fast attention score calculation similar to flash attention?

[qiyuxinlin]

I am in a scenario with limited memory. I found that flash attention does not occupy VRAM and the calculation speed is very fast, but it does not support the output softmax (qk). Do you know that there is similar code that can achieve qk calculation in low memory scenarios?