[Kernel] Fix Flashinfer Correctness

[LiuXiaoxuanPKU]

FIX #7176

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[LiuXiaoxuanPKU]

/ready

[simon-mo]

wait how did this fix it? did the profile run corrupt something?

[Yard1]

@simon-mo looks like the assumption that "prefill doesn't read KV cache" was incorrect.

[LiuXiaoxuanPKU]

wait how did this fix it? did the profile run corrupt something?

Originally, we set paged_kv_indptr to be 0s for both profile run and prefill run. We don't use flashinfer in the profile run, so it's good. But we use flashinfer in the prefill run, and the prefill run reads the kv cache. Setting paged_kv_indptr to be all 0s disable flashinfer from reading the kv cache, which corrupts the output.

[youkaichao]

test failures are not related and are fixed in the main branch.

[felixzhu555]

@LiuXiaoxuanPKU I think this line is buggy, self.block_tables here is always a tensor so is_block_tables_empty should always return False. From what I've checked this will happen to work for flashinfer v0.1.2, but not v0.1.3, which will raise RuntimeError: CHECK_EQ(paged_kv_indptr.size(0), batch_size + 1) failed. 1 vs 257 during profile run. (This is due to the if self.is_profile_run block never being run.)

I tried a quick fix of patching is_block_tables_empty by checking if block_tables has num_el == 0, which will pass the profile run fine. But this introduces logic issues for prefill, which obviously has a 0 length block table as well. Not sure what the best fix for this is, just wanted to raise this concern here.