[fix] recover auto-dispatch for rmsnorm and rope

[Alcanderian]

Motivation

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Hello @Alcanderian, 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!

Summary of Changes

Hello, Gemini here! I've reviewed this pull request by Alcanderian. The primary goal, as indicated by the title, is to restore or improve the auto-dispatch mechanism for the RMSNorm, LayerNorm, and Rotary Embedding layers within the SRT (SGLang Runtime) component. This is achieved by simplifying the main forward methods in these layers, removing explicit checks for CUDA or HIP devices, and instead relying on a delegated method (_forward_method) to handle the actual computation and potentially the device-specific dispatch. This change aims to make the layer implementations cleaner and potentially leverage PyTorch's internal dispatch mechanisms more effectively.

Highlights

• Auto-dispatch Logic Refinement: The explicit device checks (_is_cuda, _is_hip) have been removed from the main forward methods in layernorm.py and rotary_embedding.py. The logic now primarily checks for torch.compiler.is_compiling() and otherwise delegates the call to self._forward_method.
• LayerNorm/RMSNorm Forward Simplification: The forward methods in the RMSNorm and LayerNorm classes within python/sglang/srt/layers/layernorm.py have been updated to remove the direct if _is_cuda: and elif _is_hip: branches, instead calling self._forward_method in the default case.
• Rotary Embedding Forward Update: The RotaryEmbedding class in python/sglang/srt/layers/rotary_embedding.py has had its forward method modified/replaced. The new structure checks for torch.compiler.is_compiling() and then calls self._forward_method, removing the previous explicit CUDA check.

Changelog

• python/sglang/srt/layers/layernorm.py
  • Modified the forward method in the RMSNorm class (around line 52) to remove explicit CUDA/HIP checks and delegate to _forward_method.
  • Modified the forward method in the LayerNorm class (around line 120) to remove explicit CUDA checks and delegate to _forward_method.
• python/sglang/srt/layers/rotary_embedding.py
  • Added a new forward method (around line 119) that checks for compilation and delegates to _forward_method.
  • Removed an older forward method (around line 656) that contained explicit CUDA checks.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 is currently in preview and 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 to provide feedback.

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.

---

Layers normalize,
Rotary shifts the context,
Dispatch finds the path.

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 introduces a beneficial refactoring to the forward methods in RMSNorm, LayerNorm, and RotaryEmbedding classes. By leveraging self._forward_method (presumably set by the CustomOp base class), the dispatch logic for CUDA, HIP, and native backends is now centralized and standardized.

Key improvements:

• Maintainability: The dispatch logic is no longer duplicated in each subclass, making it easier to manage and extend in CustomOp.dispatch_forward().
• Clarity: The forward methods in the subclasses are now simpler and more focused on their specific logic (like the torch.compiler.is_compiling() check).
• Correctness for RotaryEmbedding: The previous forward method in RotaryEmbedding did not explicitly handle the HIP backend (forward_hip). The new approach correctly allows CustomOp to dispatch to self.forward_hip if _is_hip is true and _is_cuda is false. This aligns with the PR title regarding "recovering" auto-dispatch.

Overall, these changes enhance code quality and consistency. Well done!

A minor suggestion for future PRs: it would be helpful to fill out the "Motivation" and "Modifications" sections in the PR description, even for refactorings, to provide context for reviewers.

Summary of Findings

• Code Quality and Refactoring: The changes successfully refactor the dispatch logic in RMSNorm, LayerNorm, and RotaryEmbedding, centralizing it likely within the CustomOp base class. This improves maintainability and code clarity.
• Dispatch Correctness for RotaryEmbedding: The refactoring in RotaryEmbedding also corrects the dispatch mechanism to properly include forward_hip when appropriate, which was missing in its previous forward method implementation.
• Overall Assessment: No issues of medium, high, or critical severity were identified. The pull request introduces positive changes.

Merge Readiness

The changes in this pull request are well-implemented and improve the codebase. The refactoring centralizes dispatch logic effectively and, in the case of RotaryEmbedding, corrects previous dispatch behavior for HIP backends. The code appears to be in good shape for merging. As always, ensure all CI checks and tests pass before merging. I am unable to approve the pull request myself, so please have another reviewer approve it.

[Alcanderian]

@ispobock @fzyzcjy It works fine now! I would really appreciate it if you could take another look!

[fzyzcjy]

Only some nits (it is ok to ignore all of them) and a curious question

[fzyzcjy]

LGTM (as long as all tests pass)

[HaiShaw]

@fzyzcjy @Alcanderian @zhyncs
some question to layernorm changes

[HaiShaw]

we would still need this custom dispatch in this context - over-general abstraction as _is_hip to forward_hip isn't enough - we will need to add forward_aiter specifically in this block.