code: def forward(ctx, tokens, mask, dest_idx, ec):
      s = tokens.size(0)
      h = tokens.size(1)
      dtype = tokens.dtype
  
      if MOE_KERNEL is None:
          load_moe()
      if tokens.dtype != torch.float32:
          tokens = tokens.to(torch.float32)
      expert_input = MOE_KERNEL.dispatch_forward(
          s, ec, h, tokens, mask, dest_idx
      )
      if expert_input.dtype != dtype:
          expert_input = expert_input.to(dtype)
      ctx.save_for_backward(mask, dest_idx)
      ctx.s = s
      ctx.h = h
      ctx.ec = ec
      ctx.dtype = dtype
  
      return expert_input
  
imports: from typing import Any, Optional, Tuple
  import torch
  import torch.distributed
  from torch import Tensor
  from torch.cuda.amp import custom_bwd, custom_fwd
  from torch.distributed import ProcessGroup
label_desc: Use image node_id 1 for calling the MoeDispatch obj's underlying member method code with example usage: obj.forward(ctx, tokens, mask, dest_idx, ec) and returns: expert_input