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","."]
docstr: Return final model logits.
  
  Args:
      hidden_states (torch.Tensor): The hidden states from the model. This is the output of the
          final layer of the model.
      vocab_embedding (torch.Tensor): The vocab embedding table.
      raged_metadata (RaggedBatchWrapper): The ragged batch metadata.
      gamma (Optional[torch.Tensor]): The gamma tensor for normalization.
      beta (Optional[torch.Tensor]): The beta tensor for normalization.
docstr_tok: ["Return","final","model","logits",".","Args",":","hidden_states","(","torch.Tensor",")",":","The","hidden","states","from","the","model",".","This","is","the","output","of","the","final","layer","of","the","model",".","vocab_embedding","(","torch.Tensor",")",":","The","vocab","embedding","table",".","raged_metadata","(","RaggedBatchWrapper",")",":","The","ragged","batch","metadata",".","gamma","(","Optional","[","torch.Tensor","]",")",":","The","gamma","tensor","for","normalization",".","beta","(","Optional","[","torch.Tensor","]",")",":","The","beta","tensor","for","normalization","."]
code: def forward(
      self,
      hidden_states: torch.Tensor,
      vocab_embedding: torch.Tensor,
      ragged_metadata: RaggedBatchWrapper,
      bias: Optional[torch.Tensor] = None,
      gamma: Optional[torch.Tensor] = None,
      beta: Optional[torch.Tensor] = None,
  ) -> torch.Tensor:
      """
      Return final model logits.
  
      Args:
          hidden_states (torch.Tensor): The hidden states from the model. This is the output of the
              final layer of the model.
          vocab_embedding (torch.Tensor): The vocab embedding table.
          raged_metadata (RaggedBatchWrapper): The ragged batch metadata.
          gamma (Optional[torch.Tensor]): The gamma tensor for normalization.
          beta (Optional[torch.Tensor]): The beta tensor for normalization.
      """
  
      cut_down_hidden_states = empty_from(
          self._intermediate,
          (ragged_metadata.current_sequences, self._config.model_dim),
      )
      self._logits_gather(
          cut_down_hidden_states, hidden_states, ragged_metadata
      )
  
      if self._config.norm_type == "rms_norm":
          if gamma is None:
              raise ValueError(
                  "RMS Normalization enabled but gamma not provided."
              )
          self._norm(
              cut_down_hidden_states, cut_down_hidden_states, gamma
          )
      elif self._config.norm_type == "layer_norm":
          if gamma is None or beta is None:
              raise ValueError(
                  "Normalization enabled but gamma and/or beta not provided."
              )
          self._norm(
              cut_down_hidden_states,
              cut_down_hidden_states,
              gamma,
              beta,
          )
  
      output = empty_from(
          self._output,
          (ragged_metadata.current_sequences, self._config.vocab_size),
      )
      self._linear(output, cut_down_hidden_states, vocab_embedding)
      if bias is not None:
          self._act_fn(output, bias)
  
      return output
  
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 5 for calling the DSRaggedUnembed obj's underlying member method code with example usage: obj.forward(hidden_states, vocab_embedding, ragged_metadata, bias, gamma, beta) and returns: output