code: def __init__(
      self, config: DSNormConfig, implementation_config: Dict[str, Any]
  ):
      super().__init__(config, implementation_config)
      self._fp_rms = CUDARMSNorm(
          self._config.channels,
          self._config.residual_dtype,
          epsilon=self._config.eps,
      )
      self._fp_rms_pre = CUDARMSPreNorm(
          self._config.channels,
          self._config.residual_dtype,
          epsilon=self._config.eps,
      )
  
      # Buffers for both the hidden and residual outputs
      self._hidden_output = torch.empty(
          (config.max_tokens, config.channels),
          dtype=config.output_dtype,
          device=get_accelerator().current_device(),
      )
      self._residual_output = torch.empty(
          (config.max_tokens, config.channels),
          dtype=config.output_dtype,
          device=get_accelerator().current_device(),
      )
  
file_path: DeepSpeed/deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_rms.py
imports: from typing import Any, Dict, Optional, Tuple
  import torch
  from deepspeed.accelerator import get_accelerator
  from ...interfaces import DSPreNormBase, DSPreNormRegistry
  from ...configs import DSNormConfig, NormTypeEnum
  from ....kernels.core_ops import CUDARMSNorm, CUDARMSPreNorm
  from ....allocator import empty_from
  from ....inference_parameter import InferenceParameter
label_desc: Use image node_id 3 to create a new DSPreRMSCUDAModule object from inherited base classes: DSPreNormBase with example: obj = DSPreRMSCUDAModule(config, implementation_config)