class_docstr: CUTLASS implementation of MoE GEMM.
docstr: Performs a MoE GEMM. Note that the stride between token inputs must be even (the distance between byte 1 of token 0 and token 1 must be the same as the distance between byte 1 of token 1 and token 2).
  
  Arguments:
      ordered_output (torch.Tensor): The output of the MoE GEMM of shape [n_tokens, out_neurons].
      ordered_input (torch.Tensor): The direct input for the MoE GEMM of shape [n_tokens, in_neurons].
      weights (torch.Tensor): The weights of shape [n_experts, in_neurons, out_neurons]. These weights must be contiguous.
      scales (torch.Tensor): The scales of shape [n_experts, out_neurons]. These scales must be contiguous.
      total_rows_before_expert (torch.Tensor): The total number of rows before each expert of shape [n_experts].
      biases (torch.Tensor): The biases of shape [n_experts, out_neurons]. These biases must be contiguous.
  
  Returns:
      ordered_output
docstr_tok: ["Performs","a","MoE","GEMM",".","Note","that","the","stride","between","token","inputs","must","be","even","(","the","distance","between","byte","1","of","token","0","and","token","1","must","be","the","same","as","the","distance","between","byte","1","of","token","1","and","token","2",")",".","Arguments",":","ordered_output","(","torch.Tensor",")",":","The","output","of","the","MoE","GEMM","of","shape","[","n_tokens",",","out_neurons","]",".","ordered_input","(","torch.Tensor",")",":","The","direct","input","for","the","MoE","GEMM","of","shape","[","n_tokens",",","in_neurons","]",".","weights","(","torch.Tensor",")",":","The","weights","of","shape","[","n_experts",",","in_neurons",",","out_neurons","]",".","These","weights","must","be","contiguous",".","scales","(","torch.Tensor",")",":","The","scales","of","shape","[","n_experts",",","out_neurons","]",".","These","scales","must","be","contiguous",".","total_rows_before_expert","(","torch.Tensor",")",":","The","total","number","of","rows","before","each","expert","of","shape","[","n_experts","]",".","biases","(","torch.Tensor",")",":","The","biases","of","shape","[","n_experts",",","out_neurons","]",".","These","biases","must","be","contiguous",".","Returns",":","ordered_output"]
code: def __call__(
      self,
      ordered_output: torch.Tensor,
      ordered_input: torch.Tensor,
      weights: torch.Tensor,
      scales: torch.Tensor,
      total_rows_before_expert: torch.Tensor,
      biases: Optional[torch.Tensor] = None,
  ) -> None:
      """
      Performs a MoE GEMM. Note that the stride between token inputs must be even (the distance between byte 1 of token 0 and token 1 must be the same as the distance between byte 1 of token 1 and token 2).
  
      Arguments:
          ordered_output (torch.Tensor): The output of the MoE GEMM of shape [n_tokens, out_neurons].
          ordered_input (torch.Tensor): The direct input for the MoE GEMM of shape [n_tokens, in_neurons].
          weights (torch.Tensor): The weights of shape [n_experts, in_neurons, out_neurons]. These weights must be contiguous.
          scales (torch.Tensor): The scales of shape [n_experts, out_neurons]. These scales must be contiguous.
          total_rows_before_expert (torch.Tensor): The total number of rows before each expert of shape [n_experts].
          biases (torch.Tensor): The biases of shape [n_experts, out_neurons]. These biases must be contiguous.
  
      Returns:
          ordered_output
      """
      self.kernel(
          ordered_output,
          ordered_input,
          weights,
          scales,
          biases,
          total_rows_before_expert,
          self.num_bits,
          self.act_fn,
      )
      return ordered_output
  
file_path: deepspeed/deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/mixed_moe_gemm.py
imports: import torch
  from ...None import DSKernelBase
  from ....inference_utils import ActivationType, DtypeEnum
  from deepspeed.ops.op_builder import InferenceCutlassBuilder
  from typing import Optional
label_desc: Use image node_id 2 for calling the MixedMoEGEMM obj's underlying member method code with example usage: obj.__call__(ordered_output, ordered_input, weights, scales, total_rows_before_expert, biases) and returns: ordered_output