fix layernorm once for all

vllm/hpu/ops.py

@@ -12,6 +12,16 @@
 import torch.nn.functional as F
 
 import vllm.hpu.utils as hpu_utils
+from vllm.logger import init_logger
+
+logger = init_logger()
+HPUFusedRMSNorm = None
+try:
+ from habana_frameworks.torch.hpex.normalization import FusedRMSNorm
+ HPUFusedRMSNorm = FusedRMSNorm
+except ImportError:
+ logger.warning("Could not import HPU FusedRMSNorm kernel. "
+ "vLLM will use forward_native implementation of RMSNorm.")
 
 PA_SPLIT_VALUE = (os.environ.get('PA_SPLIT_VALUE', '1') == '1')
 

vllm/model_executor/layers/layernorm.py

@@ -6,18 +6,8 @@
 
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
-from vllm.utils import is_hpu
 
 logger = init_logger(__name__)
-if is_hpu():
- try:
- from habana_frameworks.torch.hpex.normalization import (
- FusedRMSNorm as HPUFusedRMSNorm)
- except ImportError:
- logger.warning(
- "Could not import HPU FusedRMSNorm kernel. "
- "vLLM will use forward_native implementation of RMSNorm.")
- HPUFusedRMSNorm = None
 
 
 class RMSNorm(CustomOp):
@@ -85,6 +75,7 @@ def forward_hpu(
  x: torch.Tensor,
  residual: Optional[torch.Tensor] = None,
  ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ from vllm.hpu.ops import HPUFusedRMSNorm
  if HPUFusedRMSNorm is None:
  return self.forward_native(x, residual)
  if residual is not None: