[Refactor] Uniform PoDAttention API with Horizontal Fusion SMs Schedule

[happierpig]

Description

This PR is a follow-up to #858, which integrates the PoDAttention (arXiv link) API in a user-transparent manner. Users can now invoke PoDAttention via the same API as BatchPrefillWithPagedKVCache, without explicitly specifying whether requests are prefill or decode (example code).

Key Changes

1. Support for Non-Continuous Q/O and KV Tensor Layout
   Previously, tensor offsets were computed using indptr, assuming continuous layouts. PoDAttention requires supporting mixed prefill/decode subsets within requests, necessitating a non-continuous layout.

   • Added q_lenptr and kv_lenptr to accommodate this functionality (code link).

2. Horizontal Fusion-Style Implementation
   For improved efficiency, subsets of requests are aware of each other, enabling optimal selection of kernel hyperparameters and persistent kernel execution.

   • Current resource partitioning strategy solely depends on total KV-cache load size (scheduler code).
   • Note: This strategy is customizable based on specific workloads.

Limitations and Future Work

• CUDA Graph is currently not supported. Only FA2 is supported at this stage.
• The workload classifier (qo_len > threshold) is preliminary and requires improvement (classifier implementation).
• Performance tuning is ongoing, and correctness has only been validated on a limited set of unit tests (unit tests).
  cc @AKKamath @yzh119

[yzh119]

Some of the unittests failed, for example (test_block_sparse_attention[False-256-16-16-128-64-16-4])

RuntimeError: Error in function 'PrefillSplitQOKVIndptr' at /workspace/flashinfer/data/include/flashinfer/attention/scheduler.cuh:515: kv_len_ptr_h[0]: 0 should be positive

[Edenzzzz]

Hi, can I ask when this is planned to be merged? I made a PR to support POD Attn in SGLang using the old API and plan to get that working with CUDA graph first.
sgl-project/sglang#5169

[AKKamath]

I really like the uniform batch API that this PR presents.

I ran this on an A100 and compared it with the existing FlashInfer POD-Attention implementation. On average this performed around 10 - 15% worse, but still better than serial execution. Performance was worse for larger prefill context lengths, while for smaller context lengths the performance was more comparable.

[Edenzzzz]

Yeah this is more convenient, one issue i had during my PR is that I have to fill 2D attention mask for prefill every time, instead using page table & indices

[yzh119]

Hi @Edenzzzz @AKKamath , I'm working on another branch following this idea, it will be merged these days.

[Edenzzzz]

Will the old API be preserved? Thanks.

[Edenzzzz]

@AKKamath Btw, I wonder what was the reason for using a mask instead of page table for prefill qkv?

[AKKamath]

@yzh119 Can correct me here, but I believe the mask prefill kernel (single_prefill) had a better performance than the page table prefill because the page table prefill had a higher register usage causing register spills.