[BugFix] Fix UB in per_token_group_quant.cu

[rivos-shreeasish]

Purpose

The GroupReduceMax function causes Undefined Behavior. Existing mask specifies a half-warp causing UB for the upper 16 half of the threads in a warp.

Use threadIdx.x to select the correct mask at runtime.

This change follows a patch to SGLang for the same kernel.

Test Plan

Existing tests cover this kernel.

Test Result

Existing tests including CI should continue passing.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[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 addresses an undefined behavior bug in the GroupReduceMax CUDA function. The original implementation used a fixed mask for __shfl_xor_sync that was only valid for threads in the lower half of a warp, causing UB for threads in the upper half. The fix correctly computes the mask at runtime based on the thread's lane ID within its warp, ensuring that the shuffle operation is always valid. The change also removes an unused parameter from the function signature. The fix is correct and effectively resolves the reported issue. The changes are well-targeted and I have no further suggestions.