Skip to content

[XPU] Add forward_xpu to Gemma4RMSNorm by reusing existing kernel ops#26262

Closed
CaoE wants to merge 3 commits into
sgl-project:mainfrom
CaoE:gemma4_rmsnorm
Closed

[XPU] Add forward_xpu to Gemma4RMSNorm by reusing existing kernel ops#26262
CaoE wants to merge 3 commits into
sgl-project:mainfrom
CaoE:gemma4_rmsnorm

Conversation

@CaoE
Copy link
Copy Markdown
Contributor

@CaoE CaoE commented May 25, 2026

Motivation

Add forward_xpu to Gemma4RMSNorm by reusing ops already available in sgl-kernel for XPU.

Modifications

Refactor forward_cuda into a shared _forward_impl helper with CUDA path remaining unchanged.
Add forward_xpu that leverages the XPU kernel's native 3D stride support.
Note: The XPU gemma_rmsnorm/ rmsnorm kernels already handle non-contiguous 3D inputs via stride-aware row offset computation, so we can skip the flatten-to-2D reshape when the tensor dim is 3 with contiguous last dim.

Generated by my agent.

Accuracy Tests

Covered by sgl-project/sgl-kernel-xpu#216.

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): 🚫 Run #26428790998
Latest PR Test (Extra): ❌ Run #26428790936

@gemini-code-assist
Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@CaoE CaoE marked this pull request as ready for review May 25, 2026 15:06
@gemini-code-assist
Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@CaoE
Copy link
Copy Markdown
Contributor Author

CaoE commented May 25, 2026

/tag-run-ci-label

@CaoE
Copy link
Copy Markdown
Contributor Author

CaoE commented May 27, 2026

This is coverd by #23280.

@CaoE CaoE closed this May 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants