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 supports_config(config: DSUnembedConfig):
      if DtypeEnum(config.dtype) not in [
          DtypeEnum.fp16,
          DtypeEnum.bf16,
          DtypeEnum.fp32,
      ]:
          return False
  
      try:
          _ = RaggedLogitsGather(config.model_dim, config.dtype)
      except ValueError:
          return False
  
      if config.norm_type == "rms_norm":
          try:
              _ = CUDARMSNorm(config.model_dim, config.dtype)
          except ValueError:
              return False
      elif config.norm_type == "layer_norm":
          try:
              _ = CUDAFPLN(config.model_dim, config.dtype)
          except ValueError:
              return False
  
      return True
  
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 2 for calling the DSRaggedUnembed obj's underlying member method code with example usage: obj.supports_config(config) and returns: True, False, False, False, False