Revert "Fix RMSNorm API CALL mismatch issue. (#10032)" (#13727)

This commit is contained in:
ErsongWang
2025-11-23 06:25:52 +08:00
committed by GitHub
parent cad7878964
commit 047935080d

View File

@@ -19,7 +19,6 @@ from typing import Optional, Tuple, Union
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 @@ if _use_aiter:
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 @@ class RMSNorm(CustomOp):
# 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