[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