Skip to content

[XPU] Add fused GemmaRMSNorm path for eager execution - #53678

Merged
jikunshang merged 4 commits into
vllm-project:mainfrom
ccrhx4:feat/xpu-gemma-rms-norm
Sep 2, 2026
Merged

jikunshang merged 4 commits into
vllm-project:mainfrom
ccrhx4:feat/xpu-gemma-rms-norm

Conversation

@ccrhx4

@ccrhx4 ccrhx4 commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

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).

  • 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.
  • Depends on [Qwen3.5] feat(layernorm): add fused XPU GemmaRMSNorm kernels vllm-xpu-kernels#532
  • Also add the corresponding _custom_ops wrappers and guarded register_fake entries.

AI assistance (GitHub Copilot CLI) was used to author this change.

Test Result

Performance (enforce-eager prefill on Intel B70 input_len=8192)

Model TP eager (native fallback) eager (fused) torch.compile
Qwen3.5-9B 1 1.256 s 1.023 s 1.013 s
Qwen3.5-35B-A3B 4 0.580 s 0.419 s 0.421 s

Accuracy (GSM8K, lm-eval-harness, max_gen_toks=4096)

Model TP exact_match (strict) exact_match (flexible)
Qwen3.5-9B 1 0.8923 0.8901
Qwen3.6-27B 4 0.6748 0.6755
Qwen3.5-35B-A3B 4 0.9462 0.9538

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

Comment thread vllm/_custom_ops.py Outdated
@ccrhx4
ccrhx4 force-pushed the feat/xpu-gemma-rms-norm branch 2 times, most recently from 90e05c8 to 5e7b92a Compare August 31, 2026 02:14
@jikunshang jikunshang added the verified Run pre-commit for new contributors without triggering other tests label Aug 31, 2026
@jikunshang

Copy link
Copy Markdown
Member

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #86370 for commit 5dda576fa86c.

@jikunshang

Copy link
Copy Markdown
Member

there are some failed case related to gemma in intel/ci, please check.
Multi-Modal Models (Standard) 2: qwen3 + gemma https://buildkite.com/vllm/intel-ci/builds/10131#01a05851-ff34-4e76-9a97-2d1b2f1d1fa3

ccrhx4 and others added 2 commits September 1, 2026 01:08
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
ccrhx4 force-pushed the feat/xpu-gemma-rms-norm branch from 5dda576 to 6056c28 Compare September 1, 2026 02:11
@ccrhx4

ccrhx4 commented Sep 1, 2026

Copy link
Copy Markdown
Contributor Author

there are some failed case related to gemma in intel/ci, please check. Multi-Modal Models (Standard) 2: qwen3 + gemma https://buildkite.com/vllm/intel-ci/builds/10131#01a05851-ff34-4e76-9a97-2d1b2f1d1fa3

The CI failed is caused by the output tensor creation. Update the related code.

    # empty_like preserves x's strides, but the kernel requires a
    # contiguous out (unlike x, which it can handle non-contiguous).
    out = torch.empty(x.shape, device=x.device, dtype=x.dtype)
    ops.gemma_rms_norm(out, x, self.weight.data, self.variance_epsilon)

Now the failed CI seems unrelated to this PR. @jikunshang Would you please again kindly help? Thank you.

@jikunshang

Copy link
Copy Markdown
Member

/ci run

@github-actions

github-actions Bot commented Sep 2, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #86802 for commit 21cff155b1e5.

@jikunshang
jikunshang merged commit a0d3e5c into vllm-project:main Sep 2, 2026
105 checks passed
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

intel-gpu Related to Intel GPU verified Run pre-commit for new contributors without triggering other tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants