diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 7569f2b97394..fcbeb8ae1ec4 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -19,7 +19,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from packaging.version import Version from sglang.srt.batch_invariant_ops import ( is_batch_invariant_mode_enabled, @@ -70,11 +69,8 @@ from aiter import rmsnorm2d_fwd as rms_norm from aiter import rmsnorm2d_fwd_with_add as fused_add_rms_norm elif _is_hip: - import vllm from vllm._custom_ops import fused_add_rms_norm, rms_norm - _vllm_version = Version(vllm.__version__) - logger = logging.getLogger(__name__) if _is_npu: @@ -169,21 +165,8 @@ def forward_hip( # NOTE: Remove this if aiter kernel supports discontinuous input x = x.contiguous() if residual is not None: - if _vllm_version < Version("0.9"): - fused_add_rms_norm(x, residual, self.weight.data, self.variance_epsilon) - return x, residual - else: - residual_out = torch.empty_like(x) - output = torch.empty_like(x) - fused_add_rms_norm( - output, - x, - residual_out, - residual, - self.weight.data, - self.variance_epsilon, - ) - return output, residual_out + fused_add_rms_norm(x, residual, self.weight.data, self.variance_epsilon) + return x, residual out = torch.empty_like(x) rms_norm(out, x, self.weight.data, self.variance_epsilon) return out