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

```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: typing.Optional[torch.Tensor] = 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.

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

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

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

```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: typing.Optional[torch.Tensor] = None,
    noise_embedding: typing.Optional[torch.Tensor] = None,
    target_hidden_states: typing.Optional[torch.Tensor] = 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: typing.Optional[torch.Tensor] = 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: typing.Optional[torch.Tensor] = None
) -> typing.Optional[torch.Tensor]
```

```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: typing.Optional[torch.Tensor] = 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: typing.Optional[torch.Tensor] = 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)
```