> 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.models.deepseek_v4.mtp

DeepSeek V4 Multi-Token Prediction (MTP) blocks.

The released DSV4-Flash checkpoint stores MTP under `mtp.{depth}.*`.  Each
MTP depth mirrors the reference `MTPBlock`:

* fuse the future-token embedding and the backbone HC stream with
  `e_proj(embed) + h_proj(hidden)`;
* run one HC-enabled DSV4 attention + MoE block;
* collapse the HC stream with an MTP-local `hc_head` and `norm` before
  the shared LM head computes the auxiliary CE loss.

## Module Contents

### Classes

| Name                                                                                           | Description                                                           |
| ---------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- |
| [`DeepseekV4MTPBlock`](#nemo_automodel-components-models-deepseek_v4-mtp-DeepseekV4MTPBlock)   | One DSV4 MTP depth.                                                   |
| [`DeepseekV4MTPModule`](#nemo_automodel-components-models-deepseek_v4-mtp-DeepseekV4MTPModule) | DSV4 MTP stack, one :class:`DeepseekV4MTPBlock` per prediction depth. |

### Functions

| Name                                                                                                     | Description                                 |
| -------------------------------------------------------------------------------------------------------- | ------------------------------------------- |
| [`build_deepseek_v4_mtp`](#nemo_automodel-components-models-deepseek_v4-mtp-build_deepseek_v4_mtp)       | Construct DSV4 MTP blocks.                  |
| [`build_mtp_config_from_hf`](#nemo_automodel-components-models-deepseek_v4-mtp-build_mtp_config_from_hf) | Build an MTPConfig from a DeepseekV4Config. |

### API

```python
class nemo_automodel.components.models.deepseek_v4.mtp.DeepseekV4MTPBlock(
    config,
    layer_idx: int,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype,
    rotary_emb,
    rotary_emb_compress
)
```

**Bases:** `Module`

One DSV4 MTP depth.

**Parameters:**

Main DSV4 config.

Global layer index used by the attention implementation.

Shared MoE config.

BackendConfig for kernels/modules.

Model dtype.

Shared main rotary embedding module.

Shared compressor rotary embedding module.

```python
nemo_automodel.components.models.deepseek_v4.mtp.DeepseekV4MTPBlock.forward(
    hidden_states: torch.Tensor,
    embed_input: torch.Tensor,
    input_ids: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    attn_kwargs = {}
) -> tuple[torch.Tensor, torch.Tensor]
```

Run one MTP depth.

**Parameters:**

HC stream `[B, S, hc_mult, H]`.

Future-token embeddings `[B, S, H]`.

**Returns:** `torch.Tensor`

Tuple of `(next_hc_stream, prediction_hidden)` where

```python
nemo_automodel.components.models.deepseek_v4.mtp.DeepseekV4MTPBlock.init_weights(
    buffer_device: torch.device | None = None
) -> None
```

```python
class nemo_automodel.components.models.deepseek_v4.mtp.DeepseekV4MTPModule(
    config,
    mtp_config: nemo_automodel.components.models.common.mtp.MTPConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    dtype: torch.dtype,
    rotary_emb,
    rotary_emb_compress
)
```

**Bases:** `Module`

DSV4 MTP stack, one :class:`DeepseekV4MTPBlock` per prediction depth.

```python
nemo_automodel.components.models.deepseek_v4.mtp.DeepseekV4MTPModule.forward(
    hidden_states: torch.Tensor,
    input_ids: torch.LongTensor | None = None,
    embed_fn = None,
    embed_inputs: tuple[torch.Tensor, ...] | list[torch.Tensor] | None = None,
    position_ids: torch.LongTensor | None = None,
    block_kwargs = {}
) -> list[torch.Tensor]
```

```python
nemo_automodel.components.models.deepseek_v4.mtp.build_deepseek_v4_mtp(
    config,
    mtp_config: nemo_automodel.components.models.common.mtp.MTPConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    dtype: torch.dtype,
    rotary_emb,
    rotary_emb_compress
) -> nemo_automodel.components.models.deepseek_v4.mtp.DeepseekV4MTPModule
```

Construct DSV4 MTP blocks.

```python
nemo_automodel.components.models.deepseek_v4.mtp.build_mtp_config_from_hf(
    config,
    loss_scaling_factor: float = 0.1
) -> nemo_automodel.components.models.common.mtp.MTPConfig
```

Build an MTPConfig from a DeepseekV4Config.