compatible with flashinfer v0.2

[zhyncs]

Motivation

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling.

[zhyncs]

This implementation is also compatible with v0.1.6, except for the additional cost of 256 padding. All CIs have currently passed, it is safe to merge now.

[zhyncs]

ref https://github.com/flashinfer-ai/flashinfer/blob/e5a3befbe3e63025f0158bc96b218a9c5f402ac7/tests/test_hopper.py#L178-L179

[zhyncs]

cc @qingquansong

[zhyncs]

FYI This PR only solves the compatibility issue and does not address the performance issues of v0.2 as expected. I will discuss and resolve this with @yzh119 offline.