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

Native DeepSeek V4.1 DSpark draft backbone.

The released draft is stored under `mtp.0` through `mtp.2` but is not the
autoregressive MTP objective used by earlier DeepSeek models. It is a parallel
five-position drafter whose three blocks use sliding-window MLA, routed MoE and
single-pass mHC. This module implements the trainable, cache-free backbone; the
shared frozen embedding/LM head and anchor sampling remain owned by the generic
DSpark trainer.

## Module Contents

### Classes

| Name                                                                                                                         | Description                                                                   |
| ---------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`DeepseekV41DSparkBackbone`](#nemo_automodel-components-models-deepseek_v41-dspark-DeepseekV41DSparkBackbone)               | Three-stage native drafter operating on prepared target and noise tensors.    |
| [`DeepseekV41DSparkBackboneOutput`](#nemo_automodel-components-models-deepseek_v41-dspark-DeepseekV41DSparkBackboneOutput)   | Native draft states consumed by the released output heads.                    |
| [`DeepseekV41DSparkModel`](#nemo_automodel-components-models-deepseek_v41-dspark-DeepseekV41DSparkModel)                     | Model-owned adapter from the native draft to the shared DSpark loss contract. |
| [`_DeepseekV41DSparkAttention`](#nemo_automodel-components-models-deepseek_v41-dspark-_DeepseekV41DSparkAttention)           | Cache-free DSpark attention over target context and parallel draft blocks.    |
| [`_DeepseekV41DSparkBlock`](#nemo_automodel-components-models-deepseek_v41-dspark-_DeepseekV41DSparkBlock)                   | One released DSpark stage with MLA, MoE, mHC and stage-owned heads.           |
| [`_DeepseekV41DSparkConfidenceHead`](#nemo_automodel-components-models-deepseek_v41-dspark-_DeepseekV41DSparkConfidenceHead) | Predict conditional acceptance logits in FP32.                                |
| [`_DeepseekV41DSparkMarkovHead`](#nemo_automodel-components-models-deepseek_v41-dspark-_DeepseekV41DSparkMarkovHead)         | Rank-factorized first-order token-transition bias.                            |
| [`_DeepseekV41DSparkStageOutput`](#nemo_automodel-components-models-deepseek_v41-dspark-_DeepseekV41DSparkStageOutput)       | Internal stage state kept within each stage's FSDP forward boundary.          |

### Data

[`__all__`](#nemo_automodel-components-models-deepseek_v41-dspark-__all__)

### API

```python
class nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackbone(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None
)
```

**Bases:** `Module`

Three-stage native drafter operating on prepared target and noise tensors.

**`moe_config`**

---

**`mtp`**

---

```python
nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackbone.build_attention_mask(
    anchor_positions: torch.Tensor,
    block_keep_mask: torch.Tensor,
    context_sequence: int,
    dtype: torch.dtype
) -> torch.Tensor
```

Build the released SWA-128 multi-anchor training mask.

Every query sees the target window ending immediately before its anchor,
plus every parallel input in its own draft block. It cannot see another
anchor's block. Invalid padding blocks retain their own in-block keys so
no attention row is fully masked; their losses are discarded later.

**Parameters:**

**`anchor_positions`** `torch.Tensor`

Integer tensor \[batch, num\_anchors].

---

**`block_keep_mask`** `torch.Tensor`

Boolean tensor \[batch, num\_anchors].

---

**`context_sequence`** `int`

Number of target-context tokens.

---

**`dtype`** `torch.dtype`

Floating dtype of the returned additive mask.

---

**Returns:** `torch.Tensor`

Additive tensor \[batch, 1, num\_anchors \* block\_size,

```python
nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackbone.build_position_ids(
    anchor_positions: torch.Tensor,
    context_sequence: int
) -> torch.Tensor
```

Build official target-context and draft-query positions.

**Parameters:**

**`anchor_positions`** `torch.Tensor`

Integer tensor \[batch, num\_anchors] containing the
target token that seeds each draft block.

---

**`context_sequence`** `int`

Number of target-context tokens.

---

**Returns:** `torch.Tensor`

Integer tensor \[batch, context\_sequence + num\_anchors \* block\_size].

```python
nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackbone.forward(
    noise_embeddings: torch.Tensor,
    target_hidden_states: torch.Tensor,
    position_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    previous_token_ids: torch.Tensor | None = None,
    enable_confidence_head: bool = True,
    confidence_head_stop_gradient: bool = False
) -> nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackboneOutput
```

Run the cache-free draft backbone for sampled anchors.

**Parameters:**

**`noise_embeddings`** `torch.Tensor`

Tensor of shape \[batch, draft\_sequence, hidden],
containing an anchor embedding followed by noise embeddings in
each fixed-width block.

---

**`target_hidden_states`** `torch.Tensor`

Concatenated target features of shape \[batch,
context\_sequence, target\_layers \* hidden].

---

**`position_ids`** `torch.Tensor`

Integer tensor of shape \[batch, context\_sequence + draft\_sequence].

---

**`attention_mask`** `torch.Tensor`

Additive tensor of shape \[batch, 1, draft\_sequence,
context\_sequence + draft\_sequence].

---

**`previous_token_ids`** `torch.Tensor | None` — default: None

Optional integer tensor of shape \[batch,
draft\_sequence] used by the final Markov head.

---

**`enable_confidence_head`** `bool` — default: True

Whether the final stage computes confidence.

---

**`confidence_head_stop_gradient`** `bool` — default: False

Whether the confidence head reads
detached inputs.

---

**Returns:** `DeepseekV41DSparkBackboneOutput`

Draft backbone output containing normalized states of shape \[batch,

```python
nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackbone.initialize_weights(
    buffer_device: torch.device | None = None
) -> None
```

Initialize every draft parameter for checkpoint-free training.

**Parameters:**

**`buffer_device`** `torch.device | None` — default: None

Device used by grouped expert initialization. Defaults
to the first attention projection's device.

---

```python
class nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackboneOutput(
    normalized_hidden_states: torch.Tensor,
    transition_logits: torch.Tensor | None = None,
    confidence_pred: torch.Tensor | None = None
)
```

Dataclass

Native draft states consumed by the released output heads.

**`confidence_pred`** `Tensor | None = None`

---

**`normalized_hidden_states`** `Tensor`

---

**`transition_logits`** `Tensor | None = None`

---

```python
class nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkModel(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    num_anchors: int,
    enable_confidence_head: bool,
    confidence_head_stop_gradient: bool = False
)
```

**Bases:** [DeepseekV41DSparkBackbone](#nemo_automodel-components-models-deepseek_v41-dspark-DeepseekV41DSparkBackbone)

Model-owned adapter from the native draft to the shared DSpark loss contract.

**`_keep_in_fp32_modules_strict`**

---

**`_no_split_modules`** `= ['_DeepseekV41DSparkBlock']`

---

**`confidence_head_stop_gradient`** `= bool(confidence_head_stop_gradient)`

---

**`embed_tokens`**

---

**`enable_confidence_head`** `= bool(enable_confidence_head)`

---

**`layers`** `ModuleList`

Expose native MTP stages to the shared AC/FSDP helpers.

---

**`lm_head`**

---

**`num_anchors`** `= int(num_anchors)`

---

**`state_dict_adapter`**

---

```python
nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkModel.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Project hidden states through the frozen vocabulary head in FP32.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

---

**Returns:** `torch.Tensor`

FP32 tensor of shape \[batch, sequence, vocab].

```python
nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkModel.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
```

Run native V4.1 DSpark training for sampled anchors.

**Parameters:**

**`input_ids`** `torch.Tensor`

Target-token tensor of shape \[batch, sequence].

---

**`target_hidden_states`** `torch.Tensor`

Concatenated target-feature tensor of shape
\[batch, sequence, target\_layers \* hidden].

---

**`loss_mask`** `torch.Tensor`

Supervision tensor of shape \[batch, sequence].

---

**`target_last_hidden_states`** `torch.Tensor | None` — default: None

Optional frozen target-state tensor of
shape \[batch, sequence, hidden] used by the probability-distance loss.

---

**Returns:** `DSparkForwardOutput`

Shared DSpark loss inputs. Draft logits have shape \[batch,

```python
nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkModel.initialize_embeddings_and_head(
    embed_tokens: torch.nn.Module,
    lm_head: torch.nn.Module,
    freeze: bool = True
) -> None
```

Copy the target's embedding and vocabulary projection.

**Parameters:**

**`embed_tokens`** `nn.Module`

Target embedding with weight of shape \[vocab, hidden].

---

**`lm_head`** `nn.Module`

Target output projection with weight of shape \[vocab, hidden].

---

**`freeze`** `bool` — default: True

Disable gradients for both copied modules when true.

---

```python
nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkModel.set_embedding_head_trainable(
    trainable: bool
) -> None
```

Set whether the copied embedding and LM head receive gradients.

```python
class nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkAttention(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** [DeepseekV41Attention](/nemo-automodel/nemo_automodel/components/models/deepseek_v41/attention#nemo_automodel-components-models-deepseek_v41-attention-DeepseekV41Attention)

Cache-free DSpark attention over target context and parallel draft blocks.

```python
nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkAttention.forward(
    hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    position_ids: torch.Tensor,
    attention_mask: torch.Tensor
) -> torch.Tensor
```

Attend draft queries to target context and their own parallel block.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Draft tensor of shape \[batch, draft\_sequence, hidden].

---

**`target_hidden_states`** `torch.Tensor`

Projected target tensor of shape
\[batch, context\_sequence, hidden].

---

**`position_ids`** `torch.Tensor`

Integer tensor of shape \[batch, context\_sequence +
draft\_sequence], containing absolute positions for both regions.

---

**`attention_mask`** `torch.Tensor`

Additive tensor broadcastable to shape \[batch, heads,
draft\_sequence, context\_sequence + draft\_sequence], with zero for
visible keys and negative infinity for masked keys.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, draft\_sequence, hidden]. Inputs are not mutated.

```python
class nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkBlock(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    stage_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig
)
```

**Bases:** `Module`

One released DSpark stage with MLA, MoE, mHC and stage-owned heads.

**`attn`**

---

**`attn_hc`**

---

**`attn_norm`**

---

**`confidence_head`**

---

**`ffn`** `= MoE(moe_config, backend)`

---

**`ffn_hc`**

---

**`ffn_norm`**

---

**`main_norm`**

---

**`main_proj`**

---

**`markov_head`**

---

**`mlp`** `MoE`

Expose the shared parallelizer's MoE interface without duplicate registration.

---

**`norm`**

---

```python
nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkBlock.forward(
    hidden_states: torch.Tensor,
    pre_mix: torch.Tensor,
    target_hidden_states: torch.Tensor,
    position_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    previous_token_ids: torch.Tensor | None = None,
    enable_confidence_head: bool = True,
    confidence_head_stop_gradient: bool = False
) -> nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkStageOutput
```

Apply one native DSpark stage.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, draft\_sequence, streams, hidden].

---

**`pre_mix`** `torch.Tensor`

FP32 tensor of shape \[batch, draft\_sequence, streams].

---

**`target_hidden_states`** `torch.Tensor`

Tensor of shape \[batch, context\_sequence, hidden].

---

**`position_ids`** `torch.Tensor`

Integer tensor of shape \[batch, context\_sequence + draft\_sequence].

---

**`attention_mask`** `torch.Tensor`

Additive tensor of shape \[batch, 1, draft\_sequence,
context\_sequence + draft\_sequence].

---

**`previous_token_ids`** `torch.Tensor | None` — default: None

Optional integer tensor of shape \[batch,
draft\_sequence] used by the final Markov head.

---

**`enable_confidence_head`** `bool` — default: True

Whether the final stage computes confidence.

---

**`confidence_head_stop_gradient`** `bool` — default: False

Whether the confidence head reads
detached inputs, so its loss trains only `confidence_head`.

---

**Returns:** `_DeepseekV41DSparkStageOutput`

Stage state containing updated streams, target features, and any

```python
class nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkConfidenceHead(
    hidden_size: int,
    markov_rank: int
)
```

**Bases:** `Module`

Predict conditional acceptance logits in FP32.

AutoModel-trained drafts feed the final RMSNorm output to this head. Serving
these checkpoints must use the same input instead of the released raw residual.

**`proj`**

---

```python
nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkConfidenceHead.forward(
    hidden_states: torch.Tensor,
    markov_embeddings: torch.Tensor
) -> torch.Tensor
```

Predict a conditional acceptance logit for every draft position.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Final RMSNorm output of shape \[..., hidden].

---

**`markov_embeddings`** `torch.Tensor`

Tensor of shape \[..., markov\_rank] with matching
leading dimensions.

---

**Returns:** `torch.Tensor`

FP32 tensor of shape \[...] containing uncalibrated confidence logits.

```python
class nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkMarkovHead(
    vocab_size: int,
    rank: int,
    dtype: torch.dtype
)
```

**Bases:** `Module`

Rank-factorized first-order token-transition bias.

**`embed`** `= nn.Embedding(vocab_size, rank, dtype=dtype)`

---

**`head`**

---

```python
nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkMarkovHead.forward(
    token_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Compute the transition bias and its conditioning embedding.

**Parameters:**

**`token_ids`** `torch.Tensor`

Integer tensor of shape \[...] containing preceding tokens.

---

**Returns:** `torch.Tensor`

Transition logits of shape \[..., vocab] and embeddings of shape

```python
class nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkStageOutput()
```

**Bases:** `NamedTuple`

Internal stage state kept within each stage's FSDP forward boundary.

**`confidence_pred`** `Tensor | None = None`

---

**`normalized_hidden_states`** `Tensor | None = None`

---

**`pre_mix`** `Tensor`

---

**`streams`** `Tensor`

---

**`target_hidden_states`** `Tensor`

---

**`transition_logits`** `Tensor | None = None`

---

```python
nemo_automodel.components.models.deepseek_v41.dspark.__all__ = ['DeepseekV41DSparkModel']
```