nemo_automodel.components.speculative.dspark.draft_minimax_m3

DSpark speculative-decoding draft model with a MiniMax M3 attention backbone.

Mirrors the Qwen3 DSpark draft’s control flow (standard GQA: separate q/k/v/o projections, per-head norm on Q/K, K/V built by projecting the context and the noise block separately then concatenating). The two differences from Qwen3 are MiniMax M3’s own building blocks: per-head Gemma RMSNorm (MiniMaxM3RMSNorm), SwiGLU-OAI activation (swiglu_oai), and partial RoPE via the gpt_oss rotary utilities (matching the target’s own MiniMaxM3Attention/MiniMaxM3TextModel).

The draft runs dense, non-causal attention over a [context | noise-block] layout, with visibility supplied entirely by the DFlash additive/block attention mask. There is no block-sparse indexer and no MoE here: M3’s sparse attention and routed experts belong to the target model, not this draft (the draft is always dense, matching the DeepSeek V4 DSpark draft’s convention).

Module Contents

Classes

Name	Description
`MiniMaxM3DSparkAttention`	Dense, non-causal MiniMax M3 GQA attention over a `[context \| noise-block]` layout.
`MiniMaxM3DSparkDecoderLayer`	Pre-norm residual block: MiniMax M3 DSpark attention followed by a dense SwiGLU-OAI MLP.
`MiniMaxM3DSparkMLP`	Dense SwiGLU-OAI MLP for the DSpark draft (plain `nn.Linear`, no MoE).
`MiniMaxM3DSparkModel`	DSpark draft model with a MiniMax M3 attention backbone (dense, non-causal).

Functions

Name	Description
`_minimax_m3_flex_attention`	Use fused FlexAttention on CUDA while retaining the eager CPU test path.

Data

__all__

_minimax_m3_flex_attention_compiled

API

class nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkAttention(
    config,
    layer_idx: int
)

Bases: Module

Dense, non-causal MiniMax M3 GQA attention over a [context | noise-block] layout.

Standard GQA (separate q/k/v/o projections, per-head Gemma RMSNorm on Q and K before RoPE), mirroring Qwen3DSparkAttention’s control flow. M3 has no shared K=V latent (unlike DeepSeek V4’s MLA draft), so K and V are genuinely separate projections; there is no block-sparse indexer or MoE here (target-only machinery, not the draft’s concern).

RoPE follows the gpt_oss rotary utilities’ broadcasting convention (x in [B, S, H, D] layout while cos/sin apply), matching the target’s own MiniMaxM3Attention, so RoPE is applied to Q/K before the transpose to [B, H, S, D] for attention — not after, unlike the Gemma4/Qwen3 drafts’ own RoPE helpers, which use the opposite (post-transpose) convention.

head_dim

k_norm

k_proj

layer_idx

= int(layer_idx)

num_heads

= int(config.num_attention_heads)

num_key_value_groups

= self.num_heads // self.num_key_value_heads

num_key_value_heads

= int(config.num_key_value_heads)

o_proj

q_norm

q_proj

scaling

= self.head_dim ** -0.5

v_proj

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkAttention._repeat_kv(
    hidden_states: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkAttention.forward(
    hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    attention_mask: typing.Optional[torch.Tensor] = None,
    kwargs = {}
) -> torch.Tensor

class nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkDecoderLayer(
    config,
    layer_idx: int
)

Bases: Module

Pre-norm residual block: MiniMax M3 DSpark attention followed by a dense SwiGLU-OAI MLP.

hidden_size

= config.hidden_size

input_layernorm

mlp

= MiniMaxM3DSparkMLP(config)

post_attention_layernorm

self_attn

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkDecoderLayer.forward(
    target_hidden_states: typing.Optional[torch.Tensor] = None,
    hidden_states: typing.Optional[torch.Tensor] = None,
    attention_mask: typing.Optional[torch.Tensor] = None,
    cos: typing.Optional[torch.Tensor] = None,
    sin: typing.Optional[torch.Tensor] = None,
    kwargs = {}
) -> torch.Tensor

class nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkMLP(
    config
)

Bases: Module

Dense SwiGLU-OAI MLP for the DSpark draft (plain nn.Linear, no MoE).

alpha

= float(getattr(config, 'swiglu_alpha', 1.702))

down_proj

gate_proj

limit

= float(getattr(config, 'swiglu_limit', 7.0))

up_proj

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkMLP.forward(
    x: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel(
    config
)

Bases: Module

DSpark draft model with a MiniMax M3 attention backbone (dense, non-causal).

_no_split_modules

= ['MiniMaxM3DSparkDecoderLayer']

block_size

= int(config.block_size)

confidence_head_with_markov

embed_tokens

enable_confidence_head

= bool(config.enable_confidence_head)

head_dim

hidden_norm

layers

lm_head

markov_head

= build_markov_head(config)

mask_token_id

= config.mask_token_id

norm

num_anchors

= int(config.num_anchors)

rotary_emb

target_layer_ids

= config.target_layer_ids

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel._apply(
    fn,
    recurse: bool = True
)

Keep the rotary module’s device in sync across .to() calls.

RotaryEmbedding caches its concentration/inv_freq via functools.cache keyed on the module instance, not a persistent buffer, so nn.Module._apply (which only visits parameters/buffers) never touches it. A later .to(device=...) (the training build path) would otherwise silently leave the cached frequencies on whatever device was current at construction, causing a device mismatch the first time they are used against q/k on the real device. Sync .device and drop the cache after any move so the next call recomputes on the right device.

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel._forward_backbone(
    position_ids: torch.LongTensor,
    attention_mask: typing.Optional[torch.Tensor] = None,
    noise_embedding: typing.Optional[torch.Tensor] = None,
    target_hidden_states: typing.Optional[torch.Tensor] = None,
    kwargs = {}
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel.forward(
    input_ids: torch.Tensor,
    target_hidden_states: torch.Tensor,
    loss_mask: torch.Tensor,
    target_last_hidden_states: typing.Optional[torch.Tensor] = None
) -> nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel.initialize_embeddings_and_head(
    embed_tokens: torch.nn.Module,
    lm_head: torch.nn.Module,
    freeze: bool = True
)

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel.predict_confidence_step(
    hidden_states: torch.Tensor,
    prev_token_ids: typing.Optional[torch.Tensor] = None
) -> typing.Optional[torch.Tensor]

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel.sample_draft_token_step(
    base_logits: torch.Tensor,
    prev_token_ids: torch.Tensor,
    temperature: float = 0.0,
    hidden_states: typing.Optional[torch.Tensor] = None
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel.sample_draft_tokens(
    base_logits: torch.Tensor,
    first_prev_token_ids: torch.Tensor,
    temperature: float = 0.0,
    hidden_states: typing.Optional[torch.Tensor] = None
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel.set_embedding_head_trainable(
    trainable: bool
)

nemo_automodel.components.speculative.dspark.draft_minimax_m3._minimax_m3_flex_attention(
    q,
    k,
    v,
    block_mask,
    scale
)

Use fused FlexAttention on CUDA while retaining the eager CPU test path.

nemo_automodel.components.speculative.dspark.draft_minimax_m3.__all__ = ['MiniMaxM3DSparkModel', 'MiniMaxM3DSparkAttention', 'MiniMaxM3DSparkDecoderLaye...

nemo_automodel.components.speculative.dspark.draft_minimax_m3._minimax_m3_flex_attention_compiled = torch.compile(flex_attention, mode='max-autotune-no-cudagraphs', dynamic=True)