Class MiniMaxMoE

java.lang.Object
smile.deep.layer.LayerBlock
smile.llm.model.minimax.MiniMaxMoE
All Implemented Interfaces:
Function<Tensor,Tensor>, Layer
public class MiniMaxMoE extends LayerBlock
MiniMax-M3 sparse mixture of experts: sigmoid router with a selection bias, top-k routed experts and an always-on shared expert.

Routing: scores = sigmoid(x Wg^T); experts are chosen by the top-k of scores + e_score_correction_bias, but the combine weights are the unbiased scores, renormalized over the chosen experts and multiplied by routed_scaling_factor. The router runs in float32.

Routed experts are stored stacked, one tensor per projection (experts_w1, experts_w3: [local, I, D]; experts_w2: [local, D, I]), so a rank allocates only the experts it owns and no per-expert modules exist. Under expert parallelism every rank routes all tokens identically, computes only the tokens assigned to its own experts, and the partial outputs are summed with an all-reduce; the shared expert is replicated and added afterwards.

This is the reference (eager) implementation: routing decisions are read back to the host and each active expert runs on its own gathered rows, so it is not CUDA-graph capturable.

  • Constructor Details

    • MiniMaxMoE

      public MiniMaxMoE(MiniMaxModelArgs args)
      Constructor for the unsharded layout.
      Parameters:
      args - model hyperparameters.

    • MiniMaxMoE

      public MiniMaxMoE(MiniMaxModelArgs args, MiniMaxShard shard, TensorParallelGroup group)
      Constructor.
      Parameters:
      args - model hyperparameters.
      shard - this rank's slice of the routed experts; their weights are allocated on its device.
      group - all-reduce group for shard.size() > 1, otherwise null.

  • Method Details

    • setRouterFp32

      public void setRouterFp32(Tensor weight, Tensor bias)
      Installs float32 router parameters so routing is not rounded to the model's compute dtype.
      Parameters:
      weight - router weight [numExperts, dim], float32, on the model's device.
      bias - selection bias [numExperts], float32, on the model's device.

    • sharedExpert

      public MiniMaxMlp sharedExpert()
      Returns the shared expert.
      Returns:
      the shared expert, or null when none is configured.

    • expertStart

      public int expertStart()
      Returns the first global expert index owned by this rank.
      Returns:
      first local expert.

    • localExperts

      public int localExperts()
      Returns the number of experts owned by this rank.
      Returns:
      local expert count.

    • expertWeights

      public Tensor expertWeights(int projection)
      Returns the stacked weight tensor of a projection.
      Parameters:
      projection - 1 (gate), 2 (down) or 3 (up).
      Returns:
      the stacked tensor [localExperts, ...].

    • forward

      public Tensor forward(Tensor input)
      Description copied from interface: Layer
      Forward propagation (or forward pass) through the layer.
      Parameters:
      input - the input tensor.
      Returns:
      the output tensor.