Use FlashAttention for multi_query_kv_attention

[WoosukKwon]

This PR is to use FlashAttention kernels for multi_query_kv_attention, which performs masked attention for the prompt inputs.

Pros

• FlashAttention is fast and memory-efficient.
• FlashAttention supports 1D inputs and only invokes a single kernel to handle multiple sequences with variable lengths.

Cons

• FlashAttention does NOT support FP32.
• FlashAttention does not support head_size > 128. (This is fine for all models except GPT-J).
  • Ref: Support for 256 head dim Dao-AILab/flash-attention#67
• FlashAttention does not support attention bias (GPT-J, BLOOM, LLaMA).

Besides, note that FlashAttention does not support cached KV, which is required for interactive generation.

Tested models:

• OPT-125M
• OPT-350M
• OPT-1.3B
• OPT-2.7B
• OPT-6.7B
• OPT-13B

Tested GPUs:

• A100