[WIP][1/N] Chunked Prefill

[rkooo567]

WIP. RFC link: #3130

[tdene]

Is the goal of this PR and #3121 not the same?

[robertgshaw2-redhat]

@rkooo567 this is an awesome initiative, thanks for kicking it off! Quick note on this PR.

There is already an attention kernel context_attention_fwd in vLLM that covers the case of prefill with kv cache. It was originally created for Manual Prefix Caching (and will be used for Automatic Prefix Caching as well). I have not looked too much into the implementation of context_attention_fwd (so I'm not sure of the performance relative to the flash attention implementation), but just wanted to make sure you were aware

Perhaps we can leverage these kernels in the Automatic Prefix caching diff

[rkooo567]

There is already an attention kernel context_attention_fwd in vLLM that covers the case of prefill with kv cache. It was originally created for Manual Prefix Caching (and will be used for Automatic Prefix Caching as well). I have not looked too much into the implementation of context_attention_fwd (so I'm not sure of the performance relative to the flash attention implementation), but just wanted to make sure you were aware

Yes! I am aware of this. I think it must be possible to use this kernel as well, but our internal benchmark shows it is pretty slow actually (I don't know much detail here though).

[robertgshaw2-redhat]

Thanks @rkooo567

[zhr2001]

After trying with the provided code, I think the implementation is not completely finished and it still leaves some bugs.

[rkooo567]

@zhr2001 actually this is stale. did you try #3884 (comment)?

[zhr2001]

I have tried.It's excellent, thanks for your contribution.
However, I just got result 7.09it/s for benchmark_serving, while without chunkedprefill the result was 3.44it/s within llama-13B,A40,tp2, max-num-batched-token and sonnet data-set.I wonder how to get the following result in the picture below.

[rkooo567]

@zhr2001 I am starting the benchmark from tomorrow. Btw, what does 7.09it/s for benchmark_serving mean (does it mean it is better or worse)? I am not familiar with this benchmark

For the result above, it is the result with our internal repo (the implementation is almost the same, but our repo has cuda-graph for prefill whereas OSS currently doesn't have it). The environment is basically sending a requests with similar length with a certain QPS