code: def _mixtral_moe_baseline(
      hidden_states: torch.Tensor,
      gate_weight: torch.Tensor,
      mlp_w1: torch.Tensor,
      mlp_w2: torch.Tensor,
      mlp_w3: torch.Tensor,
      force_float: bool = False,
  ) -> torch.Tensor:
      """
      Baseline implementation for mixtral MoE module.
  
      Based on transformers implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py
      """
      output_dtype = hidden_states.dtype
      if force_float:
          hidden_states = hidden_states.float()
          gate_weight = gate_weight.float()
          mlp_w1 = mlp_w1.float()
          mlp_w2 = mlp_w2.float()
          mlp_w3 = mlp_w3.float()
  
      router_logits = torch.nn.functional.linear(
          hidden_states, gate_weight
      )
      routing_weights = torch.nn.functional.softmax(
          router_logits, dim=-1, dtype=torch.float
      )
      routing_weights, selected_experts = routing_weights.topk(
          k=2, dim=-1
      )
      routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
  
      # NOTE(cmikeh2): This is a difference implementation, ours will preserve the original scale
      # as float32 and perform in-kernel fused FP16->FP32->FP16 conversion.
      routing_weights = routing_weights.to(hidden_states.dtype)
  
      final_hidden_states = torch.zeros_like(hidden_states)
  
      expert_mask = torch.nn.functional.one_hot(
          selected_experts, num_classes=gate_weight.shape[0]
      ).permute(2, 1, 0)
      get_accelerator().synchronize()
  
      for expert_idx in range(gate_weight.shape[0]):
          exp_mlp_w1 = mlp_w1[expert_idx]
          exp_mlp_w2 = mlp_w2[expert_idx]
          exp_mlp_w3 = mlp_w3[expert_idx]
  
          idx, top_x = torch.where(expert_mask[expert_idx])
  
          if top_x.shape[0] == 0:
              continue
  
          top_x_list = top_x.tolist()
          idx_list = idx.tolist()
  
          current_state = hidden_states[top_x_list]
  
          linear = torch.nn.functional.linear
          intermediate = torch.nn.functional.silu(
              linear(current_state, exp_mlp_w1)
          ) * linear(current_state, exp_mlp_w3)
          output = linear(intermediate, exp_mlp_w2) * routing_weights[
              top_x_list, idx_list
          ].unsqueeze(-1)
          final_hidden_states.index_add_(
              0, top_x, output.to(final_hidden_states.dtype)
          )
  
      return final_hidden_states.to(output_dtype)
  
file_path: DeepSpeed/tests/unit/inference/v2/modules/test_cutlass_moe.py
imports: from typing import Tuple
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum
  from deepspeed.inference.v2.modules import ConfigBundle
  from deepspeed.inference.v2.modules.configs import DSMoEConfig
  from deepspeed.inference.v2.modules.interfaces import DSMoERegistry
  from ..kernels.ragged_ops.ragged_testing_utils import build_simple_batch
  from ...v2.inference_test_utils import allclose, get_dtypes
label_desc: Use image node_id 9 for calling a global function with example usage: _mixtral_moe_baseline(hidden_states, gate_weight, mlp_w1, mlp_w2, mlp_w3, force_float) and returns: final_hidden_states