class_docstr: CUDA implementation of top-1 gating. This will perform a softmax on the logits,
  and return the scale as well as its idx within that expert's allocation.
docstr: Perform the ragged top_k_gating.
  
  Arguments:
      expert_counts (torch.Tensor): Tensor of 0s of shape [n_experts] to be filled with
          number of tokens assigned to each expert. This must be filled with 0s else
          the copy kernel will buffer overflow. In order to minimize the zero-fill cost,
          it is recommended to write to 0 during the MoE output remapping.
      scores (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place expert scaling
          value.
      expert_assignment (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place
          which expert a token has been assigned to.
      expert_offset (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place which
          offset within an experts group a token is.
      logits (torch.Tensor): Raw logits of gating function.
      batch (RaggedBatchWrapper): Batch information for ragged tensor.
  
  Returns:
      tuple of (expert_counts, scores, expert_assignment, expert_offset)
docstr_tok: ["Perform","the","ragged","top_k_gating",".","Arguments",":","expert_counts","(","torch.Tensor",")",":","Tensor","of","0s","of","shape","[","n_experts","]","to","be","filled","with","number","of","tokens","assigned","to","each","expert",".","This","must","be","filled","with","0s","else","the","copy","kernel","will","buffer","overflow",".","In","order","to","minimize","the","zero-fill","cost",",","it","is","recommended","to","write","to","0","during","the","MoE","output","remapping",".","scores","(","torch.Tensor",")",":","Preallocated","output","of","shape","[","n_tokens",",","n_top_k","]","to","place","expert","scaling","value",".","expert_assignment","(","torch.Tensor",")",":","Preallocated","output","of","shape","[","n_tokens",",","n_top_k","]","to","place","which","expert","a","token","has","been","assigned","to",".","expert_offset","(","torch.Tensor",")",":","Preallocated","output","of","shape","[","n_tokens",",","n_top_k","]","to","place","which","offset","within","an","experts","group","a","token","is",".","logits","(","torch.Tensor",")",":","Raw","logits","of","gating","function",".","batch","(","RaggedBatchWrapper",")",":","Batch","information","for","ragged","tensor",".","Returns",":","tuple","of","(","expert_counts",",","scores",",","expert_assignment",",","expert_offset",")"]
code: def __call__(
      self,
      expert_counts: torch.Tensor,
      scores: torch.Tensor,
      assignments: torch.Tensor,
      offsets: torch.Tensor,
      logits: torch.Tensor,
      batch: RaggedBatchWrapper,
  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
      """
      Perform the ragged top_k_gating.
  
      Arguments:
          expert_counts (torch.Tensor): Tensor of 0s of shape [n_experts] to be filled with
              number of tokens assigned to each expert. This must be filled with 0s else
              the copy kernel will buffer overflow. In order to minimize the zero-fill cost,
              it is recommended to write to 0 during the MoE output remapping.
          scores (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place expert scaling
              value.
          expert_assignment (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place
              which expert a token has been assigned to.
          expert_offset (torch.Tensor): Preallocated output of shape [n_tokens, n_top_k] to place which
              offset within an experts group a token is.
          logits (torch.Tensor): Raw logits of gating function.
          batch (RaggedBatchWrapper): Batch information for ragged tensor.
  
      Returns:
          tuple of (expert_counts, scores, expert_assignment, expert_offset)
      """
      self.kernel(
          expert_counts,
          scores,
          assignments,
          offsets,
          logits,
          batch.batch_metadata_buffer(),
      )
      return expert_counts, scores, assignments, offsets
  
file_path: deepspeed/deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.py
imports: import torch
  from typing import Tuple
  from ...None import DSKernelBase
  from ....inference_utils import DtypeEnum
  from ....ragged import RaggedBatchWrapper
  from deepspeed.ops.op_builder import RaggedOpsBuilder
label_desc: Use image node_id 2 for calling the RaggedTopKGating obj's underlying member method code with example usage: obj.__call__(expert_counts, scores, assignments, offsets, logits, batch) and returns: expert_counts,  scores,  assignments,  offsets