EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Revert "Support aiter RMSNorm in AMD" (#5646)

Browse Source
This commit is contained in:
HAI
2025-04-22 15:20:24 -07:00
committed by GitHub
parent 6b6e748775
commit b0feda090c
1 changed files with 3 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
python/sglang/srt/layers/layernorm.py
View File

@@ -20,12 +20,9 @@ import torch
import torch.nn as nn import torch.nn as nn
from sglang.srt.custom_op import CustomOp from sglang.srt.custom_op import CustomOp
from sglang.srt.utils import is_cuda, is_hip from sglang.srt.utils import is_cuda
logger = logging.getLogger(__name__)
_is_cuda = is_cuda() _is_cuda = is_cuda()
_is_hip = is_hip()
if _is_cuda: if _is_cuda:
from sgl_kernel import ( from sgl_kernel import (
@@ -35,20 +32,8 @@ if _is_cuda:
rmsnorm, rmsnorm,
) )
if _is_hip:
from aiter.ops.rmsnorm import rms_norm, rmsnorm2d_fwd_with_add logger = logging.getLogger(__name__)
rmsnorm = rms_norm
def fused_add_rmsnorm(
x: torch.Tensor,
residual: torch.Tensor,
w: torch.Tensor,
eps: float,
) -> Tuple[torch.Tensor, torch.Tensor]:
rmsnorm2d_fwd_with_add(x, x, residual, residual, w, eps)
return x, residual
class RMSNorm(CustomOp): class RMSNorm(CustomOp):
@@ -154,7 +139,7 @@ class Gemma3RMSNorm(nn.Module):
return f"{tuple(self.weight.shape)}, eps={self.eps}" return f"{tuple(self.weight.shape)}, eps={self.eps}"
if not (_is_cuda or _is_hip): if not _is_cuda:
logger.info( logger.info(
"sgl-kernel is not available on Non-NV platforms. Fallback to other kernel libraries." "sgl-kernel is not available on Non-NV platforms. Fallback to other kernel libraries."
) )