Add gather_indexer_k_quant_cache kernel

[Barry-Delaney]

This PR added cp_gather_indexer_k_quant_cache for getting quantized k/k_scale from indexer k cache.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[gemini-code-assist]

Code Review

This pull request introduces a new CUDA kernel cp_gather_indexer_k_quant_cache for gathering quantized K-cache data. My review has identified some critical issues. There's a significant bug in the CUDA kernel's grid launch configuration that will lead to incorrect data gathering for longer sequences. Additionally, there are several const correctness issues and incorrect mutability annotations in the C++ code and PyTorch bindings, which misrepresent the function's contract and can cause subtle bugs. I've provided suggestions to fix these issues.

[simon-mo]

Can you integrate it and test the performance? cc @zyongye

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Barry-Delaney.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[zyongye]

I am testing this branch.
I replace the line here to

ops.cp_gather_indexer_k_quant_cache(
                kv_cache,
                k_fp8,
                k_scale,
                prefill_metadata.block_table,
                prefill_metadata.cu_seq_lens,
            )

And getting 0.58 on gsm8k 20 shots. I did rebase with the main branch. Am I doing anything wrong?

[Barry-Delaney]

I fixed the bug in the latest force push. Could you pls help on verifying again? @zyongye Thanks in advance!

[zyongye]

Retest with the new branch. Get a good result in gsm8k and gsm8k 20 shots

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9447|±  |0.0063|
|     |       |strict-match    |     5|exact_match|↑  |0.9469|±  |0.0062|

20 shots

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|    20|exact_match|↑  |0.9431|±  |0.0064|
|     |       |strict-match    |    20|exact_match|↑  |0.9439|±  |0.0063|

[heheda12345]

@zyongye I think we can merge this PR first, and then you open a new PR for kernel integration.

[heheda12345]

Thanks for this important kernel!

[zyongye]

Maybe need to add this to pass rocm build