Improve the mixed chunk prefill by lanuch two kernels

[libratiger]

Motivation

Improve the mixed chunk prefill performance. see #2273

launch two kernels: one prefill attention kernel for prefill requests and one decode attention kernel for decode requests.

Modifications

Checklist

• Format your code according to the Contributor Guide.
• Add unit tests as outlined in the Contributor Guide.
• Update documentation as needed, including docstrings or example tutorials.

[libratiger]

related: #2273

[libratiger]

cc @merrymercy for review.
I will give more performance detail later.

[merrymercy]

@libratiger Could you share some perf numbers?
@hnyls2002 Could you review this?

[libratiger]

@libratiger Could you share some perf numbers? @hnyls2002 Could you review this?

Yes, I will add this as soon as possible.
But I want to ensure the implementation is correct and reliable first. Any feedback from the review would be greatly appreciated.

I added a new property from ScheduleBatch to ForwardBatch to distinguish between prefill and decode.
Now there are enough properties. If you have any optimization suggestions, I will improve it. @merrymercy @hnyls2002

[xiezhq-hermann]

Thank you for your contribution. I'm also curious about the performance. From my understanding, the advantage of mixing prefill and decode lies in their complementary resource usage: prefill is compute-bound, while decoding can utilize some remaining memory bandwidth. In my opinion, a specialized kernel for mixed computation might be an ideal solution. Alternatively, have you tried launching the two kernels concurrently?

[libratiger]

Thank you for your contribution. I'm also curious about the performance. From my understanding, the advantage of mixing prefill and decode lies in their complementary resource usage: prefill is compute-bound, while decoding can utilize some remaining memory bandwidth. In my opinion, a specialized kernel for mixed computation might be an ideal solution. Alternatively, have you tried launching the two kernels concurrently?

Your comment shows deep insight!
theoretically using a single kernel would have lower overhead. This is can be implemented by similar to the flash_attn_varlen_func in FlashAttention. However, currently FlashInfer doesn't seem to implement such an interface.

The current implementation uses two kernels, primarily due to two considerations:

1. To maintain consistency with the previous PR, making it easier to review and capture partial performance benefits
2. Other attention backends can be optimized similarly by reusing this logic, without needing to dive into low-level kernel modifications

Regarding the performance benefits of launching two kernels, I don't have sufficient performance data yet.

However, after completing my implementation, I looked into other frameworks' implementations to address this concern and found they also use two kernels, suggesting that this approach should still yield some performance benefits.

We can further optimize it to use a single kernel once we have more performance data, which theoretically would reduce some upper-layer overhead.

[xiezhq-hermann]

Thank you for your contribution. I'm also curious about the performance. From my understanding, the advantage of mixing prefill and decode lies in their complementary resource usage: prefill is compute-bound, while decoding can utilize some remaining memory bandwidth. In my opinion, a specialized kernel for mixed computation might be an ideal solution. Alternatively, have you tried launching the two kernels concurrently?

Your comment shows deep insight! theoretically using a single kernel would have lower overhead. This is can be implemented by similar to the flash_attn_varlen_func in FlashAttention. However, currently FlashInfer doesn't seem to implement such an interface.

The current implementation uses two kernels, primarily due to two considerations:

1. To maintain consistency with the previous PR, making it easier to review and capture partial performance benefits
2. Other attention backends can be optimized similarly by reusing this logic, without needing to dive into low-level kernel modifications

Regarding the performance benefits of launching two kernels, I don't have sufficient performance data yet.

However, after completing my implementation, I looked into other frameworks' implementations to address this concern and found they also use two kernels, suggesting that this approach should still yield some performance benefits.

We can further optimize it to use a single kernel once we have more performance data, which theoretically would reduce some upper-layer overhead.

Thanks for the reply, keep me posted with the performance number and we can merge the two kernel version first if the benefit justifies. Consider to benchmark concurrent launch version as well which might also be interesting.