[WIP] support paged attention - #710

Draft

cao lei (jslhcl) wants to merge 18 commits into

mainfrom

leca/pagedAttention

Commits on May 6, 2024

support paged attention
Lei Cao
committed
0927cca
Merge branch 'main' of https://github.com/microsoft/onnxruntime-extensions into leca/pagedAttention
Lei Cao
committed
cf10813

Commits on May 7, 2024

add kernel functions and build successfully
Lei Cao
committed
6931a52

Commits on May 8, 2024

move checkInput() to paged_attention.h
Lei Cao
committed
77ae3b0
call flash attention code for prompt mode
Lei Cao
committed
9ac5f18

Commits on May 20, 2024

runtime error in RunMultiHeadAttention
cao lei (jslhcl)
committed
90ee6a6

Commits on May 21, 2024

UT can run now
cao lei (jslhcl)
committed
17c2c41

Commits on May 29, 2024

new UT with q,k and v whose shape is 381x512, and not use GetScratchBufferUnderMultiStream
cao lei (jslhcl)
committed
75b2d04

Commits on Jun 2, 2024

make context_lens CPU input, change scale in test, change batch_size and seq_len
cao lei (jslhcl)
committed
68dac5d

Commits on Jun 4, 2024

change window_size and causal to make ut pass
cao lei (jslhcl)
committed
e8bf9ea

Commits on Jun 7, 2024

enhance flash attention lib to support page attention, and UT
cao lei (jslhcl)
committed
8352a13

Commits on Jun 10, 2024

runtime error in gpu_data_transfer.cc when copying result back to cpu for test test_cuda_paged_attention_prompt_decoding
cao lei (jslhcl)
committed
4bfcfad

Commits on Jun 19, 2024

sync ort-genai changes
cao lei (jslhcl)
committed
05c0ba2

Commits on Jun 22, 2024

sync flash attn 2.5.9 kernel code. test_cuda_paged_attention3 is good, test_cuda_paged_attention_decoding runtime error
cao lei (jslhcl)
committed
7191ea5

Commits on Jun 23, 2024

fix runtime error and add new test case
cao lei (jslhcl)
committed
654cabf

Commits on Jun 25, 2024

UT for paged attention decoding case works
cao lei (jslhcl)
committed
e374fb4

Commits on Jun 26, 2024

fix block_tables and context_lens to avoid unknown random error
cao lei (jslhcl)
committed
17cf303

Commits on Jul 1, 2024

sync with ort-genai and add document
cao lei (jslhcl)
committed
cb65793