[BugFix][AMD][Quantization] Fix torch.compile issue where wvSplitKQ not being called when it should when using quantized FP8 model

[rasmith]

This PR fixes an issue where wvSplitKQ is not being called when it should be when using a quantized FP8 model. This is because during compilation, this code path is not being used so does not get called during model execution even though it should be (e.g. batch size = 1).

I tested using Llama-3.1-8B-Instruct-FP8-KV.

Before this, this kernel was not being called at all when eager mode was not enforced.

Without the fix, and running with eager mode enforced:

-------------------------------------------------------  ---------
                                                    Name  # of Calls
 -------------------------------------------------------  ---------
 void wvSplitKQ_hf_sml_<__hip_bfloat16, c10::Float8_e...      32640

After applying fix and using profiler without eager mode enforced:

 -------------------------------------------------------  ---------
                                                    Name  # of Calls
 -------------------------------------------------------  ---------
 void wvSplitKQ_hf_sml_<__hip_bfloat16, c10::Float8_e...      32640

[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 can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

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.

🚀

[gemini-code-assist]

Code Review

This pull request correctly refactors rocm_per_tensor_w8a8_scaled_mm to register a custom PyTorch operation. This ensures that the data-dependent control flow for dispatching to the wvSplitKQ kernel is correctly handled by torch.compile, which was the goal of this PR. The implementation is sound. However, I've identified a pre-existing critical issue where the bias term is dropped when the wvSplitKQ path is taken. I've included a suggested fix to ensure the bias is applied correctly in all cases.

[yewentao256]

Thanks for the work!
But in the PR description, seems without the fix and with the fix, the output is the same?

[rasmith]

Thanks for the work! But in the PR description, seems without the fix and with the fix, the output is the same?

Without the fix, and running with eager mode enforced:

[SageMoore]

Thanks for the work! But in the PR description, seems without the fix and with the fix, the output is the same?

Without the fix, and running with eager mode enforced:

I was also confused by this 🙂

[rasmith]

Thanks for the work! But in the PR description, seems without the fix and with the fix, the output is the same?

Yes, the output is the same, we just want the function called on ROCm when it is supposed to be called.

[rasmith]

Thanks for the work! But in the PR description, seems without the fix and with the fix, the output is the same?

Yes, the output is the same. We want the function to be called, which was happening before, but is not happening as it should be anymore, so this fixes the issue.