class_docstr: CUDA Kernel implementation for gather the hidden states of the final token
  of each sequence. This is used to reduce the cost of the performing the unembedding.
docstr: Parameters:
      fp_dtype (torch.dtype): Data type for the input/output. Supported values
          are torch.float16, torch.bfloat16, and torch.float32.
docstr_tok: ["Parameters",":","fp_dtype","(","torch.dtype",")",":","Data","type","for","the","input\/output",".","Supported","values","are","torch.float16",",","torch.bfloat16",",","and","torch.float32","."]
code: def __init__(self, model_dim: int, fp_dtype: torch.dtype):
      """
      Parameters:
          fp_dtype (torch.dtype): Data type for the input/output. Supported values
              are torch.float16, torch.bfloat16, and torch.float32.
      """
      if fp_dtype not in RaggedLogitsGather.supported_dtypes:
          raise ValueError(
              "Unsupported data type: {}, supported_dtypes are {}".format(
                  fp_dtype, RaggedLogitsGather.supported_dtypes
              )
          )
  
      if elem_size(fp_dtype) * model_dim % 16 != 0:
          raise ValueError(
              "Embedding dimension must be aligned to 16 bytes, got {}".format(
                  model_dim
              )
          )
  
      inf_module = RaggedOpsBuilder().load()
      self.kernel = inf_module.gather_for_logits
  
file_path: deepspeed/deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather.py
imports: import torch
  from ...None import DSKernelBase
  from deepspeed.ops.op_builder import RaggedOpsBuilder
  from ....inference_utils import elem_size
  from ....ragged import RaggedBatchWrapper
label_desc: Use image node_id 1 to create a new RaggedLogitsGather object from inherited base classes: DSKernelBase with example: obj = RaggedLogitsGather(model_dim, fp_dtype)