[XPU] Add fused GemmaRMSNorm path for eager execution - #53678
Merged
Merged
Conversation
jikunshang
reviewed
Aug 27, 2026
ccrhx4
force-pushed
the
feat/xpu-gemma-rms-norm
branch
2 times, most recently
from
August 31, 2026 02:14
90e05c8 to
5e7b92a
Compare
jikunshang
approved these changes
Aug 31, 2026
Member
|
/ci run |
|
✅ Triggered Buildkite CI #86370 for commit |
Member
|
there are some failed case related to gemma in intel/ci, please check. |
GemmaRMSNorm on XPU fell back to separated elementwise ops under enforce-eager, causing a large prefill latency gap versus torch.compile (which Inductor-fuses the same ops). Add GemmaRMSNorm.forward_xpu that dispatches to the fused vllm-xpu-kernels gemma_rms_norm / fused_add_gemma_rms_norm ops, passing the raw (bf16/fp16) weight so the (1 + weight) offset and fp32 multiply are folded into the kernel (matching forward_native numerics). Falls back to forward_native when the kernels are unavailable in the installed package. Also add the corresponding _custom_ops wrappers and guarded register_fake entries. AI assistance (GitHub Copilot CLI) was used to author this change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Huanxing <huanxing.shen@intel.com>
CUDA has no gemma_rms_norm kernel, so this XPU-only op should not live in the shared _custom_ops.py file. Move gemma_rms_norm and fused_add_gemma_rms_norm into _xpu_ops.py, registered as torch.ops.vllm.xpu_gemma_rms_norm and torch.ops.vllm.xpu_fused_add_gemma_rms_norm via direct_register_custom_op, guarded by hasattr(torch.ops._C, "gemma_rms_norm"), following the existing convention for other XPU-only custom ops in that file (xpu_fp8_bmm, xpu_topk_topp_sampler, gdn_attention_core_xpu, etc.). Update GemmaRMSNorm.forward_xpu to call the new torch.ops.vllm.xpu_* ops accordingly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Huanxing <huanxing.shen@intel.com>
ccrhx4
force-pushed
the
feat/xpu-gemma-rms-norm
branch
from
September 1, 2026 02:11
5dda576 to
6056c28
Compare
Contributor
Author
The CI failed is caused by the output tensor creation. Update the related code.
Now the failed CI seems unrelated to this PR. @jikunshang Would you please again kindly help? Thank you. |
Member
|
/ci run |
|
✅ Triggered Buildkite CI #86802 for commit |
mylibrar
pushed a commit
to tanyuqian/vllm
that referenced
this pull request
Sep 3, 2026
…3678) Signed-off-by: Huanxing <huanxing.shen@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
GemmaRMSNorm on XPU fell back to separated elementwise ops under enforce-eager, causing a large prefill latency gap versus torch.compile (which Inductor-fuses the same ops).
AI assistance (GitHub Copilot CLI) was used to author this change.
Test Result
Performance (enforce-eager prefill on Intel B70 input_len=8192)
Accuracy (GSM8K, lm-eval-harness, max_gen_toks=4096)