[AMD] Fix RotaryEmbedding crash on AMD/ROCm (regression from #17934)

[michaelzhang-ai]

Motivation

#17934 (@BBuf) changed the HIP rotary embedding fallback from sgl_kernel.rotary_embedding to sglang.jit_kernel.pos_enc.rotary_embedding, which uses tvm_ffi JIT compilation that requires nvidia-smi / CUDA_HOME — unavailable on AMD GPUs.

This breaks all models on ROCm since _is_cuda=False on HIP always routes through the JIT fallback path.

Failing nightlies:

• 2-GPU accuracy: https://github.com/sgl-project/sglang/actions/runs/22026607130/job/63644169624
• 1-GPU unit: https://github.com/sgl-project/sglang/actions/runs/22045758633/job/63694174433
• 2-GPU VLM accuracy: https://github.com/sgl-project/sglang/actions/runs/22071313560/job/63788023338
• 2-GPU VLM perf: https://github.com/sgl-project/sglang/actions/runs/22071313560/job/63788023366

Fix: Add forward_hip() override to RotaryEmbedding that routes to forward_native() (pure PyTorch), bypassing the NVIDIA-specific JIT path. Uses *args, **kwargs because subclasses (MRotaryEmbedding, DeepseekScalingRotaryEmbedding, etc.) have different forward_native() signatures.

Pls help review @yctseng0211 @bingxche

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. 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.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello @michaelzhang-ai, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a significant regression that led to RotaryEmbedding JIT kernel failures on AMD/ROCm platforms. The change introduces a dedicated forward_hip method within the RotaryEmbedding class, redirecting AMD GPU operations to a pure PyTorch implementation. This effectively circumvents an incompatible NVIDIA-specific JIT compilation path, restoring full functionality and stability for models deployed on AMD hardware.

Highlights

• Critical Bug Fix for AMD/ROCm: This pull request resolves a critical regression introduced in PR [Kernel] Add JIT rotary_embedding_kernel #17934, which caused the RotaryEmbedding JIT kernel to crash on AMD/ROCm GPUs. The previous change incorrectly routed HIP rotary embedding fallbacks through an NVIDIA-specific JIT compilation path.
• Compatibility Restoration: The fix ensures that models running on AMD GPUs with ROCm are no longer broken by bypassing the NVIDIA-specific JIT path that relies on nvidia-smi and CUDA_HOME, which are unavailable on AMD systems.
• Implementation of forward_hip: A new forward_hip() method has been added to the RotaryEmbedding class. This method explicitly routes the computation to the forward_native() (pure PyTorch) implementation, providing a compatible solution for HIP/ROCm environments.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• python/sglang/srt/layers/rotary_embedding.py
  • Added a forward_hip method to provide a ROCm-specific implementation for rotary embedding, routing to the native PyTorch fallback.

Ignored Files

• Ignored by pattern: .github/workflows/** (2)
  • .github/workflows/nightly-test-amd-rocm720.yml
  • .github/workflows/nightly-test-amd.yml

Activity

• The pull request was opened by michaelzhang-ai to fix a regression causing crashes on AMD/ROCm GPUs.
• The motivation section clearly explains the issue, linking it to a previous PR ([Kernel] Add JIT rotary_embedding_kernel #17934) and detailing why the JIT fallback path was problematic for AMD.
• Failing nightly build links are provided as evidence of the regression.
• The PR includes a standard checklist for code formatting, unit tests, documentation, and benchmarking, along with a review process outline.
• No specific review comments or approvals have been recorded yet in the provided context.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request addresses a regression that caused crashes on AMD/ROCm platforms by introducing a forward_hip method in the RotaryEmbedding class. This new method correctly bypasses the NVIDIA-specific JIT kernel path by dispatching to the pure PyTorch forward_native implementation. The change is well-contained, correct, and effectively resolves the issue for ROCm users. The accompanying docstring clearly explains the rationale for the fix. The implementation looks solid.