> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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`](#nemo_automodel-components-speculative-dspark-draft_minimax_m3-MiniMaxM3DSparkAttention)       | Dense, non-causal MiniMax M3 GQA attention over a `[context \| noise-block]` layout.     |
| [`MiniMaxM3DSparkDecoderLayer`](#nemo_automodel-components-speculative-dspark-draft_minimax_m3-MiniMaxM3DSparkDecoderLayer) | Pre-norm residual block: MiniMax M3 DSpark attention followed by a dense SwiGLU-OAI MLP. |
| [`MiniMaxM3DSparkMLP`](#nemo_automodel-components-speculative-dspark-draft_minimax_m3-MiniMaxM3DSparkMLP)                   | Dense SwiGLU-OAI MLP for the DSpark draft (plain `nn.Linear`, no MoE).                   |
| [`MiniMaxM3DSparkModel`](#nemo_automodel-components-speculative-dspark-draft_minimax_m3-MiniMaxM3DSparkModel)               | DSpark draft model with a MiniMax M3 attention backbone (dense, non-causal).             |

### Functions

| Name                                                                                                                      | Description                                                              |
| ------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`_minimax_m3_flex_attention`](#nemo_automodel-components-speculative-dspark-draft_minimax_m3-_minimax_m3_flex_attention) | Use fused FlexAttention on CUDA while retaining the eager CPU test path. |

### Data

[`__all__`](#nemo_automodel-components-speculative-dspark-draft_minimax_m3-__all__)

[`_minimax_m3_flex_attention_compiled`](#nemo_automodel-components-speculative-dspark-draft_minimax_m3-_minimax_m3_flex_attention_compiled)

### API

```python
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`**

---

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

```python
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: torch.Tensor | None = None,
    kwargs = {}
) -> torch.Tensor
```

```python
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`**

---

```python
nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkDecoderLayer.forward(
    target_hidden_states: torch.Tensor | None = None,
    hidden_states: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    cos: torch.Tensor | None = None,
    sin: torch.Tensor | None = None,
    kwargs = {}
) -> torch.Tensor
```

```python
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`**

---

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

```python
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)`

---

**`fc`**

---

**`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`

---

```python
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.

```python
nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel._forward_backbone(
    position_ids: torch.LongTensor,
    attention_mask: torch.Tensor | None = None,
    noise_embedding: torch.Tensor | None = None,
    target_hidden_states: torch.Tensor | None = None,
    kwargs = {}
) -> torch.Tensor
```

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

```python
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: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput
```

```python
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
)
```

```python
nemo_automodel.components.speculative.dspark.draft_minimax_m3.MiniMaxM3DSparkModel.predict_confidence_step(
    hidden_states: torch.Tensor,
    prev_token_ids: torch.Tensor | None = None
) -> torch.Tensor | None
```

```python
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: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
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: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]
```

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

```python
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.

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

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