class_docstr: Ragged unembedding implementation. This implementation will gather only the last token
  of each sequence in the ragged inflight batch and calculate the logits only for those rows.
class_docstr_tok: ["Ragged","unembedding","implementation",".","This","implementation","will","gather","only","the","last","token","of","each","sequence","in","the","ragged","inflight","batch","and","calculate","the","logits","only","for","those","rows","."]
code: def __init__(
      self,
      config: DSUnembedConfig,
      implementation_config: Dict[str, Any],
  ) -> None:
      super().__init__(config, implementation_config)
  
      self._logits_gather = RaggedLogitsGather(
          config.model_dim, self._config.dtype
      )
  
      if self._config.norm_type == "layer_norm":
          self._norm = CUDAFPLN(
              self._config.model_dim, self._config.dtype
          )
      elif self._config.norm_type == "rms_norm":
          self._norm = CUDARMSNorm(
              self._config.model_dim, self._config.dtype
          )
      else:
          self._norm = None
  
      self._linear = BlasLibLinear(self._config.dtype)
      # Here the activation kernel is being used to apply bias, hence the identity activation type!
      self._act_fn = CUDABiasActivation(
          self._config.vocab_size,
          self._config.dtype,
          ActivationType.IDENTITY,
      )
  
      self._intermediate = torch.empty(
          (self._config.max_sequences, self._config.model_dim),
          dtype=self._config.dtype,
          device=get_accelerator().current_device(),
      )
  
      self._output = torch.empty(
          (self._config.max_sequences, self._config.vocab_size),
          dtype=self._config.dtype,
          device=get_accelerator().current_device(),
      )
  
file_path: DeepSpeed/deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py
imports: from typing import Any, Dict, Optional
  import torch
  from deepspeed.accelerator import get_accelerator
  from ....allocator import empty_from
  from ....inference_utils import DtypeEnum, ActivationType
  from ....kernels.core_ops import CUDAFPLN, BlasLibLinear, CUDARMSNorm, CUDABiasActivation
  from ....kernels.ragged_ops import RaggedLogitsGather
  from ....ragged import RaggedBatchWrapper
  from ...interfaces import DSUnembedBase, DSUnembedRegistry
  from ...configs import DSUnembedConfig
label_desc: Use image node_id 3 to create a new DSRaggedUnembed object from inherited base classes: DSUnembedBase with example: obj = DSRaggedUnembed(config, implementation_config)