Skip to content

Commit 2c44ad2

Browse files
authored
Revert "Support aiter RMSNorm in AMD (sgl-project#5510)"
This reverts commit 968ef51.
1 parent 2aa3f5e commit 2c44ad2

File tree

1 file changed

+3
-18
lines changed

1 file changed

+3
-18
lines changed

python/sglang/srt/layers/layernorm.py

Lines changed: 3 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,9 @@
2020
import torch.nn as nn
2121

2222
from sglang.srt.custom_op import CustomOp
23-
from sglang.srt.utils import is_cuda, is_hip
24-
25-
logger = logging.getLogger(__name__)
23+
from sglang.srt.utils import is_cuda
2624

2725
_is_cuda = is_cuda()
28-
_is_hip = is_hip()
2926

3027
if _is_cuda:
3128
from sgl_kernel import (
@@ -35,20 +32,8 @@
3532
rmsnorm,
3633
)
3734

38-
if _is_hip:
39-
40-
from aiter.ops.rmsnorm import rms_norm, rmsnorm2d_fwd_with_add
4135

42-
rmsnorm = rms_norm
43-
44-
def fused_add_rmsnorm(
45-
x: torch.Tensor,
46-
residual: torch.Tensor,
47-
w: torch.Tensor,
48-
eps: float,
49-
) -> Tuple[torch.Tensor, torch.Tensor]:
50-
rmsnorm2d_fwd_with_add(x, x, residual, residual, w, eps)
51-
return x, residual
36+
logger = logging.getLogger(__name__)
5237

5338

5439
class RMSNorm(CustomOp):
@@ -154,7 +139,7 @@ def extra_repr(self):
154139
return f"{tuple(self.weight.shape)}, eps={self.eps}"
155140

156141

157-
if not (_is_cuda or _is_hip):
142+
if not _is_cuda:
158143
logger.info(
159144
"sgl-kernel is not available on Non-NV platforms. Fallback to other kernel libraries."
160145
)

0 commit comments

Comments
 (0)