add paged-attention

[kaiyuanm]

1. Add paged attention operation commonly used in large language model inference. About paged attention view https://arxiv.org/pdf/2309.06180.pdf.
2. This implementation requires triton 2.2.0.

There is some performance data between triton and vllm.

vllm_paged_attention-B32-G8-D128-bs16-v2:
   context_len     triton  vllm-0.3.0
0        512.0   6.618643    3.582180
1       1024.0   7.798634    3.925020
2       2048.0   8.560921    4.286430
3       4096.0  12.241113    4.460762
4       8192.0  13.589705    4.558541
5      16384.0  14.094599    4.609148

Note:

• Input layout is different, performance is for reference only.
• Some input shapes is still being optimized.

[tongxin]

else assert num_splits == 1

[kaiyuanm]

Refer to flash_attn flash_attn_with_kvcache definition.

num_splits: int. If > 1, split the key/value into this many chunks along the sequence.
If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic
to automatically determine the number of splits.

[tongxin]

Yes, here you need to guard num_splits against negatives

[tongxin]

Is it possible to omit kv_mask here and the mask the output?

[kaiyuanm]

Looks like ok. What do you think? @iclementine

[iclementine]

Wehen loading, the mask is obligatory if it would access illegal memory otherwise. However you can use the modulo trick to avoid masking. Whether it is more efficient than masking depends on the cost .

[iclementine]

lgtm