[Kernel] Use fused rmsnorm for some models like qwen3 series #17735
Conversation
Signed-off-by: evian <eviantai@u.nus.edu>
|
👋 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 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 🚀 |
|
Please take a look for me, thx a lot! @Isotr0py |
Signed-off-by: evian <eviantai@u.nus.edu>
github-actions
bot
added
the
ready
…oject#17735) Signed-off-by: evian <eviantai@u.nus.edu> Co-authored-by: evian <eviantai@u.nus.edu>
…oject#17735) Signed-off-by: evian <eviantai@u.nus.edu> Co-authored-by: evian <eviantai@u.nus.edu> Signed-off-by: Mu Huai <tianbowen.tbw@antgroup.com>
…oject#17735) Signed-off-by: evian <eviantai@u.nus.edu> Co-authored-by: evian <eviantai@u.nus.edu>
…oject#17735) Signed-off-by: evian <eviantai@u.nus.edu> Co-authored-by: evian <eviantai@u.nus.edu> Signed-off-by: Yuqi Zhang <yuqizhang@google.com>
3 participants
Some models like qwen3 series still use forward_native method when q and k need to be normalized.
The reason why fused rmsnorm kernel is not being used is that before q norm and k norm, operations like split are being done, which makes q and k non-contiguous.
However, currently fused rmsnorm kernel does not support non-contiguous input. If non-contiguous input were given, the ops would not check and gave the wrong output.
To optimize, this PR add contiguous call inside rmsnorm kernel wrapper to enable fused rmsnorm kernel, and input contiguous checking is being done to ensure that non-contiguous input would never be used.
The decode time of qwen3-8b eager mode decreases from 14ms to 13ms using 1*H20.