Skip to content

[fix] recover auto-dispatch for rmsnorm and rope #6745

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 19 commits into from
Jun 4, 2025
Merged

[fix] recover auto-dispatch for rmsnorm and rope #6745

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
22b02c0
[fix] recover auto-dispatch for rmsnorm and rope
Alcanderian May 29, 2025
c796d35
Merge branch 'main' into auto-dispatch
zhyncs May 29, 2025
329e102
try to fix amd ci
Alcanderian Jun 2, 2025
ef58e45
re-dispatch _forward_method
Alcanderian Jun 2, 2025
002c1b9
fix _to_torch
Alcanderian Jun 2, 2025
b98da83
.
Alcanderian Jun 2, 2025
fcb5f81
fix
Alcanderian Jun 2, 2025
8e33357
Merge branch 'main' into auto-dispatch
Alcanderian Jun 3, 2025
a040f6b
fix
Alcanderian Jun 3, 2025
3ba86bb
fix enter_torch_compile and leave_torch_compile
Alcanderian Jun 3, 2025
15e6d9a
doc typo
Alcanderian Jun 3, 2025
08610bc
lint
Alcanderian Jun 3, 2025
04cf251
Update python/sglang/srt/custom_op.py
Alcanderian Jun 3, 2025
4adc232
Update python/sglang/srt/custom_op.py
Alcanderian Jun 3, 2025
dce7b18
Update custom_op.py
Alcanderian Jun 3, 2025
877b1cb
Update rotary_embedding.py
Alcanderian Jun 3, 2025
9a85a63
Update layernorm.py
Alcanderian Jun 3, 2025
ec8f2d9
Revert "Update rotary_embedding.py"
Alcanderian Jun 4, 2025
7605c4a
fix layernorm.py
Alcanderian Jun 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 2 additions & 8 deletions python/sglang/srt/layers/layernorm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,12 +52,8 @@ def __init__(
def forward(self, *args, **kwargs):
Copy link
Collaborator

@HaiShaw HaiShaw Jun 4, 2025

Choose a reason for hiding this comment

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

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.

if torch.compiler.is_compiling():
return self.forward_native(*args, **kwargs)
if _is_cuda:
return self.forward_cuda(*args, **kwargs)
elif _is_hip:
return self.forward_hip(*args, **kwargs)
else:
return self.forward_native(*args, **kwargs)
return self._forward_method(*args, **kwargs)

def forward_cuda(
self,
Expand Down Expand Up @@ -120,10 +116,8 @@ def __init__(
def forward(self, *args, **kwargs):
if torch.compiler.is_compiling():
return self.forward_native(*args, **kwargs)
if _is_cuda:
return self.forward_cuda(*args, **kwargs)
else:
return self.forward_native(*args, **kwargs)
return self._forward_method(*args, **kwargs)

def forward_native(
self,
Expand Down
14 changes: 6 additions & 8 deletions python/sglang/srt/layers/rotary_embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,12 @@ def _compute_cos_sin_cache(self) -> torch.Tensor:
cache = torch.cat((cos, sin), dim=-1)
return cache

def forward(self, *args, **kwargs):
if torch.compiler.is_compiling():
return self.forward_native(*args, **kwargs)
else:
return self._forward_method(*args, **kwargs)

def forward_native(
self,
positions: torch.Tensor,
Expand Down Expand Up @@ -653,14 +659,6 @@ def _compute_cos_sin_cache(self) -> torch.Tensor:
def forward_hip(self, *args, **kwargs):
return self.forward_native(*args, **kwargs)

def forward(self, *args, **kwargs):
if torch.compiler.is_compiling():
return self.forward_native(*args, **kwargs)
if _is_cuda:
return self.forward_cuda(*args, **kwargs)
else:
return self.forward_native(*args, **kwargs)

def forward_native(
self,
positions: torch.Tensor,
Expand Down
Loading