Class MiniMaxModel

java.lang.Object
smile.deep.layer.LayerBlock
smile.llm.model.minimax.MiniMaxModel
All Implemented Interfaces:
Function<Tensor,Tensor>, Layer
public class MiniMaxModel extends LayerBlock
MiniMax-M3 text model: embeddings, decoder blocks, final norm and LM head.

Construct on CPU, load weights, call to(Device) to place parameters and the RoPE tables, then install the KV pool with setKvCachePool(KvCachePool, boolean). This first version is the eager reference path (dense attention, host-routed MoE): no CUDA graphs, no tensor parallelism.

  • Constructor Details

    • MiniMaxModel

      public MiniMaxModel(MiniMaxModelArgs args)
      Constructs the module graph on CPU.
      Parameters:
      args - hyperparameters.

    • MiniMaxModel

      public MiniMaxModel(MiniMaxModelArgs args, MiniMaxShard shard, TensorParallelGroup group)
      Expert-parallel constructor: one model replica per rank, each owning a slice of the experts.
      Parameters:
      args - hyperparameters.
      shard - this rank's expert slice (its expert weights are allocated on shard.device()).
      group - all-reduce group, or null for a single rank.

  • Method Details

    • params

      public MiniMaxModelArgs params()
      Returns the hyperparameters.
      Returns:
      model args.

    • shard

      public MiniMaxShard shard()
      Returns this rank's expert slice.
      Returns:
      the shard.

    • quantLinears

      public Map<String, MxFp8Layer> quantLinears()
      Returns the FP8-resident (Marlin) linear layers by their SMILE weight names, e.g. layers.3.self_attn.q_proj.weight; empty unless the model was built in an MXFP8 weight mode. Routed experts are loaded through MiniMaxMoE instead.
      Returns:
      name to layer map.

    • layers

      public List<MiniMaxBlock> layers()
      Returns the decoder blocks.
      Returns:
      blocks in layer order.

    • kvCachePool

      public KvCachePool kvCachePool()
      Returns the installed KV pool.
      Returns:
      KV pool, or null before setKvCachePool(KvCachePool, boolean).

    • to

      public MiniMaxModel to(Device device)
      Description copied from interface: Layer
      Moves the layer block to a device.
      Specified by:
      to in interface Layer
      Overrides:
      to in class LayerBlock
      Parameters:
      device - the compute device.
      Returns:
      this layer.

    • to

      public MiniMaxModel to(Device device, ScalarType dtype)
      Description copied from interface: Layer
      Moves the layer block to a device.
      Specified by:
      to in interface Layer
      Overrides:
      to in class LayerBlock
      Parameters:
      device - the compute device.
      dtype - the data type.
      Returns:
      this layer.

    • fuseProjections

      public void fuseProjections()
      Merges each layer's attention projections into one GEMM where the weights allow it; call after loading.

    • setKvCachePool

      public void setKvCachePool(KvCachePool pool, boolean closePrevious)
      Installs the KV pool in every attention layer.
      Parameters:
      pool - new KV pool (must not be null).
      closePrevious - true to close the previous pool when replaced.

    • forward

      public Tensor forward(Tensor tokens, int startPos, boolean allTokenLogits)
      Prefill / multi-token window forward at a uniform cache position.
      Parameters:
      tokens - token ids [B, S].
      startPos - cache start position.
      allTokenLogits - score every position when true, otherwise only the last.
      Returns:
      float32 logits.

    • forward

      public Tensor forward(Tensor tokens, int startPos, boolean allTokenLogits, long[] rows, Tensor replacement)
      Prefill window whose embeddings are partly replaced, for multimodal prompts: rows of the token embeddings at rows are overwritten by replacement (vision embeddings) before the first layer.
      Parameters:
      tokens - token ids [1, S].
      startPos - cache start position.
      allTokenLogits - score every position when true, otherwise only the last.
      rows - window positions to overwrite, or null.
      replacement - [rows.length, hidden] embeddings in the compute dtype, on this device.
      Returns:
      float32 logits.

    • forward

      public Tensor forward(Tensor tokens, int[] positions)
      Single-token decode with per-row absolute positions.
      Parameters:
      tokens - token ids [B, 1].
      positions - write position per batch row (also the RoPE position).
      Returns:
      float32 logits [B, 1, V].

    • closeGraphs

      public void closeGraphs()
      Frees every captured graph (the KV pool or the weights changed under them).

    • decodeGraph

      public Tensor decodeGraph(long[] realTokens, int[] realPositions)
      Single-token decode through a captured CUDA graph.

      The first step of a (batch, window) shape runs eagerly twice (library handles and kernel attributes must exist before capture), then captures the whole forward pass and replays it; later steps only refill the static metadata buffer and replay. Nothing is built on the host while capturing: slots, positions and the attention mask come from the static buffers (see MiniMaxStep.External).

      Parameters:
      realTokens - last token of each row.
      realPositions - KV write position of each row.
      Returns:
      float32 logits [B, V] (a fresh tensor), or null when graphs are off or not applicable and the caller must run the eager path.

    • forward

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