Extend FlashAttention Prefill with KV cache

[min-jean-cho]

Moved to #331.

This extends FlashAttention prefill with cached KV in addition to current KV (blue box in the below figure). Both causal and non-causal are supported.

[pengzhao-intel]

Great Job Min Jean!

[sunjiweiswift]

Refer to sglang/test/srt/test_triton_attention_kernels.py,
qo_indptr, kv_indptr, kv_indices, and other parameters are also needed

[pengzhao-intel]

Refer to sglang/test/srt/test_triton_attention_kernels.py, qo_indptr, kv_indptr, kv_indices, and other parameters are also needed

@sunjiweiswift Thanks for good suggestion.

This is non-contiguous input of cached KV for this feature and the major part of code is same.
I suggest to review and merge this PR first and we can add this feature in next PR (or only leave it in sglang because this is the cutlass example, we want to keep it as simple as possible in general).

[mehdi-goli]

Refer to sglang/test/srt/test_triton_attention_kernels.py, qo_indptr, kv_indptr, kv_indices, and other parameters are also needed

@sunjiweiswift Thanks for good suggestion.

This is non-contiguous input of cached KV for this feature and the major part of code is same. I suggest to review and merge this PR first and we can add this feature in next PR (or only leave it in sglang because this is the cutlass example, we want to keep it as simple as possible in general).

Yes agree with you @pengzhao-intel . We need to keep the example as simple as possible and leave the rest in sglang. If there is aby performance regression there, we would be able to offer help.

[mehdi-goli]

@min-jean-cho I put the prefetch for the cached version, I also fixed some index/strides as well

[min-jean-cho]

@mehdi-goli, looks good. Thanks for the update!

[min-jean-cho]

Hi @mehdi-goli, any further comments on this? Thanks.

[pengzhao-intel]

@mehdi-goli @muhammad-tanvir-1211 please help approve and merge this PR and we have further works based on this :)

[pengzhao-intel]

LGTM! Thanks for the contribution!

[muhammad-tanvir-1211]

Do we consider the cached key-value pairs to be the same across all batches? My understanding is that each batch would have it's seq_len for the cached keys and values, which would mean that seq_len_kv_cache would also be of Variable Length type (same as seq_len_qo and seq_len_kv). This code would potentially give out of bound access because it is missing the multiplication with l_coord (if we want to keep seq_len_kv_cache fixed length), or a multiplication with kv_cache_cumulative_length[l_coord] (if we want to change the type to Variable Length)

[min-jean-cho]

Moved to #331. As discussed offline, separating without KV cache vs. with KV cache into separate pipelines. Created a new PR rather than update in here to keep the difference clear.