Extend paged attention to support query_len>1

[vanbasten23]

This PR extends the existing paged attention kernel to support query_len>1. Additionally, it upgrades the flash attention from v1 to v2.

Test plan:

• python pytorch/xla/test/test_pallas.py -v -k PallasTest.test_paged_attention_multi_queries_wrapper
• python pytorch/xla/test/test_tpu_paged_attention_kernel.py 2>&1 | tee out.txt

cc: @miladm

[vanbasten23]

hey @WoosukKwon, this is the integration point between vLLM and torch_xla. I'm thinking if vLLM can switch this flag use_kernel perhaps by using some flags. I want to use the nonkernel version as a per baseline. Do you know if it possible?

[vanbasten23]

For dynamo, it's similar. The integration point is at def multi_queries_paged_attention_xla( in the same file.

[Liyang90]

Here, we assume the input query corresponds to the last (q_len) of the input kv. For example, if the input q_len is 8, and kv_len is 24, we assume the query corresponds to the kv at index [16. 24), and applies the causal mask accordingly.

@WoosukKwon please let us know if this assumption is valid or nor for the use cases in vLLM.

[WoosukKwon]

Yes that's the desired behavior. Thanks for checking it out with me!