> This page is for version 0.7.0 · 26.10.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

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

DSpark speculative-decoding draft model with a DeepSeek V4 attention backbone.

This mirrors the Qwen3 DSpark draft (`draft_qwen3.py`) structurally; only the
attention internals, MLP, and rotary plumbing are swapped for DeepSeek V4. The
draft runs dense, non-causal attention over a `[context | noise-block]` layout,
with visibility supplied entirely by the DFlash additive attention mask. There is
no compressor / indexer / sparse-attention path (the draft is always dense; V4's
sparse machinery belongs to the target model, not this draft).

## Module Contents

### Classes

| Name                                                                                                                           | Description                                                                   |
| ------------------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------- |
| [`DeepseekV4DSparkAttention`](#nemo_automodel-components-speculative-dspark-draft_deepseek_v4-DeepseekV4DSparkAttention)       | Dense, non-causal V4 attention over a `[context \| noise-block]` layout.      |
| [`DeepseekV4DSparkDecoderLayer`](#nemo_automodel-components-speculative-dspark-draft_deepseek_v4-DeepseekV4DSparkDecoderLayer) | Pre-norm residual block: V4 DSpark attention followed by a dense SwiGLU MLP.  |
| [`DeepseekV4DSparkMLP`](#nemo_automodel-components-speculative-dspark-draft_deepseek_v4-DeepseekV4DSparkMLP)                   | Dense SwiGLU MLP for the DSpark draft.                                        |
| [`DeepseekV4DSparkModel`](#nemo_automodel-components-speculative-dspark-draft_deepseek_v4-DeepseekV4DSparkModel)               | DSpark draft model with a DeepSeek V4 attention backbone (dense, non-causal). |

### API

```python
class nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkAttention(
    config,
    layer_idx: int
)
```

**Bases:** `Module`

Dense, non-causal V4 attention over a `[context | noise-block]` layout.

Reuses DeepSeek V4's Q-LoRA, single shared K=V latent, grouped O-LoRA, and the
per-head attention sink. Queries come from the noise block only, while the
shared latent spans context plus noise. Visibility is set entirely by the
DFlash additive `attention_mask` (there is no causal flag, and no compressor
or indexer, since the draft is always dense).

**`attention_dropout`**

---

**`head_dim`** `= int(config.head_dim)`

---

**`kv_norm`**

---

**`layer_idx`** `= int(layer_idx)`

---

**`num_heads`** `= int(config.num_attention_heads)`

---

**`num_key_value_groups`** `= int(config.num_attention_heads)`

---

**`q_norm`**

---

**`rope_head_dim`** `= int(config.qk_rope_head_dim)`

---

**`scaling`** `= self.head_dim ** -0.5`

---

**`sinks_param`**

---

**`wkv`**

---

**`wo_a`**

---

**`wo_b`**

---

**`wq_a`**

---

**`wq_b`**

---

```python
nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkAttention.forward(
    hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    kwargs = {}
) -> tuple[torch.Tensor, torch.Tensor | None]
```

```python
class nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkDecoderLayer(
    config,
    layer_idx: int
)
```

**Bases:** `Module`

Pre-norm residual block: V4 DSpark attention followed by a dense SwiGLU MLP.

**`hidden_size`** `= config.hidden_size`

---

**`input_layernorm`**

---

**`mlp`** `= DeepseekV4DSparkMLP(config)`

---

**`post_attention_layernorm`**

---

**`self_attn`**

---

```python
nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkDecoderLayer.forward(
    target_hidden_states: torch.Tensor | None = None,
    hidden_states: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_value: object | None = None,
    use_cache: bool | None = False,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    kwargs = {}
) -> torch.Tensor
```

```python
class nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkMLP(
    config
)
```

**Bases:** `Module`

Dense SwiGLU MLP for the DSpark draft.

This intentionally avoids V4's MoE and Hyper-Connection paths; the draft keeps
a plain 3D hidden state and a single dense feed-forward block.

**`act_fn`** `= ACT2FN[config.hidden_act]`

---

**`down_proj`**

---

**`gate_proj`**

---

**`up_proj`**

---

```python
nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkMLP.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

```python
class nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel(
    config
)
```

**Bases:** `Module`

DSpark draft model with a DeepSeek V4 attention backbone (dense, non-causal).

**`_no_split_modules`** `= ['DeepseekV4DSparkDecoderLayer']`

---

**`block_size`** `= int(config.block_size)`

---

**`embed_tokens`**

---

**`enable_confidence_head`** `= bool(config.enable_confidence_head)`

---

**`fc`**

---

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

---

**`partial_rotary_factor`**

---

**`rotary_emb`**

---

**`target_layer_ids`** `= config.target_layer_ids`

---

```python
nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel._apply(
    fn,
    recurse = True
)
```

Keep the rotary `inv_freq` buffer in fp32 across dtype casts.

`model.to(bfloat16)` (the training build path) would otherwise round
`inv_freq` to bf16 and dephase RoPE with absolute position, eroding draft
acceptance (the mismatch grows with position, and a bf16 round-trip cannot
be undone by upcasting). Snapshot the fp32 frequencies before the cast and
restore them after (the Fp32Safe rotary idiom used elsewhere in the repo),
so the buffer never makes a bf16 round-trip.

```python
nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel._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,
    past_key_values: object | None = None,
    use_cache: bool = False,
    kwargs = {}
) -> torch.Tensor
```

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

```python
nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel.forward(
    input_ids: torch.Tensor,
    target_hidden_states: torch.Tensor,
    loss_mask: torch.Tensor,
    target_last_hidden_states: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput
```

Run one DSpark training forward.

Packing keeps every block inside its anchor's document: the anchor's first
target is in the document, the block sees only that document's context, its
RoPE uses document-local positions, and supervision stops at the document end.

**Parameters:**

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

Tensor of shape \[batch, sequence] with target tokens.

---

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

Tensor of shape \[batch, sequence, target_layers \* hidden]
with the concatenated target features.

---

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

Tensor of shape \[batch, sequence] selecting supervised tokens.

---

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

Optional tensor of shape \[batch, sequence, hidden]
with the target's final normalized states.

---

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

Optional tensor of shape \[batch, sequence] with document-local
positions of packed rows.

---

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

Optional tensor of shape \[batch, documents] with packed document
lengths; zero entries pad the document axis.

---

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

Optional tensor of shape \[batch, sequence] with the tokens left
in each position's packed document.

---

**Returns:** `DSparkForwardOutput`

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

```python
nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel.initialize_embeddings_and_head(
    embed_tokens: torch.nn.Module,
    lm_head: torch.nn.Module,
    freeze: bool = True
)
```

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

```python
nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel.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_deepseek_v4.DeepseekV4DSparkModel.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_deepseek_v4.DeepseekV4DSparkModel.set_embedding_head_trainable(
    trainable: bool
)
```