Class MiniMaxAttention

java.lang.Object

smile.llm.model.minimax.MiniMaxAttention

All Implemented Interfaces:

Attention

public class MiniMaxAttention
extends Object
implements Attention

MiniMax-M3 grouped-query attention: per-head Gemma QK norm, partial RoPE, no output gate.

On MSA layers the index projections (index_q_proj, index_k_proj and their norms) are registered so the checkpoint loads, but this reference implementation attends over the full causal history (dense). With topk_blocks covering every visible block, MSA selects all blocks and equals dense attention, so short prompts are exact; block selection lands in the sparse-attention phase.

Uses the torch-native SDPA kernel over the paged KV pool; mixed positions are supported for single-token decode through a per-row additive mask.

Constructor Summary

Constructors

Constructor	Description
MiniMaxAttention(MiniMaxModelArgs args, int layerId)	Constructor.

Method Summary

Modifier and Type	Method	Description
Tensor	forward(Tensor x, int[] positions, Tensor cos, Tensor sin, Tensor mask)	Forward pass.
Tensor	forward(Tensor x, int startPos, Tensor cis, Tensor mask)	Forward pass through the attention module.
MemorySegment	module()	Returns pytorch module.

Methods inherited from class Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Methods inherited from interface Attention

apply, apply, forward, repeatKV

Constructor Details

MiniMaxAttention

public MiniMaxAttention(MiniMaxModelArgs args,
 int layerId)

Constructor.

Parameters:

args - model hyperparameters.

layerId - decoder layer index (also the KV pool layer).

Method Details

module

public MemorySegment module()

Description copied from interface: Attention

Returns pytorch module.

Specified by:

module in interface Attention

Returns:

pytorch module.

forward

public Tensor forward(Tensor x,
 int startPos,
 Tensor cis,
 Tensor mask)

Description copied from interface: Attention

Forward pass through the attention module.

Specified by:

forward in interface Attention

Parameters:

x - the input tensor.

startPos - the starting position for attention caching.

cis - the precomputed frequency tensor.

mask - the attention mask tensor.

Returns:

the output tensor.

forward

public Tensor forward(Tensor x,
 int[] positions,
 Tensor cos,
 Tensor sin,
 Tensor mask)

Forward pass.

Parameters:

x - hidden states [B, S, D].

positions - KV write position per row; for S > 1 all rows must share one position.

cos - RoPE cosines [S, R] (prefill) or per-row [B, 1, R] (decode).

sin - RoPE sines with the same layout.

mask - causal mask [S, startPos + S] for S > 1; ignored for decode.

Returns:

attention output [B, S, D].