[Feature]: Support flashdecoding++ async unified softmax

[chizhang118]

flashdecoding++ paper: https://arxiv.org/abs/2311.01282

• Q3 Collaboration Plan of Infra and IaaS Labs: https://bytedance.us.larkoffice.com/docx/HKXfdRh1noMrbAxcgL2ureGasdQ

• FlashDecoding++ Summary: https://bytedance.larkoffice.com/wiki/WbqXwRL3qi0x18kJVkAcU0HZnte, including Asynchronized Softmax, Double Buffering, Heuristic Dataflow, and more profiling-related optimizations.

• The new kv cache memory layout and performance benchmarks is also needed.

• Inference Improvement Weekly Status and Progress: https://bytedance.us.larkoffice.com/docx/RGnPdj5gfoBN3YxuY3yuLtlQsjb

[chizhang118]

Primarily related method paged_attention_v2_reduce_kernel: https://github.com/vllm-project/vllm/blob/main/csrc/attention/attention_kernels.cu#L567

[chizhang118]

[chizhang118]

[JackChuang]

I'm taking this issue!! Please assign this to me. Thanks.

[JackChuang]

Check different block sizes and max qk sizes for a 7b model. Try to change the implementation of the qk calculation.

• figure out if the reduce kernel requires the same modification.

[JackChuang]

Confirm some implementation-related details:

• How to check overflow, how does the overflow impact, e.g., precision or crash?
• Check the usage of red_smem. Why is qk_max saved here?
• Check partition>1 + 99.99% again.
• Check why adding a space at the beginning of the prompts changes the range of the qk_max distribution.