> This page is for version 0.5.0 · 26.06.
> 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.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

DSpark speculative-decoding draft model and training objective.

A semi-autoregressive parallel drafter: a parallel backbone produces every
position of a block in one pass, a lightweight serial Markov head injects
intra-block token dependency, and a confidence head predicts per-position
acceptance probability for scheduled verification.

## Submodules

* **[`nemo_automodel.components.speculative.dspark._sampling`](/nemo-automodel/nemo_automodel/components/speculative/dspark/_sampling)**
* **[`nemo_automodel.components.speculative.dspark.common`](/nemo-automodel/nemo_automodel/components/speculative/dspark/common)**
* **[`nemo_automodel.components.speculative.dspark.config`](/nemo-automodel/nemo_automodel/components/speculative/dspark/config)**
* **[`nemo_automodel.components.speculative.dspark.core`](/nemo-automodel/nemo_automodel/components/speculative/dspark/core)**
* **[`nemo_automodel.components.speculative.dspark.draft_deepseek_v4`](/nemo-automodel/nemo_automodel/components/speculative/dspark/draft_deepseek_v4)**
* **[`nemo_automodel.components.speculative.dspark.draft_gemma4`](/nemo-automodel/nemo_automodel/components/speculative/dspark/draft_gemma4)**
* **[`nemo_automodel.components.speculative.dspark.draft_glm_5_2`](/nemo-automodel/nemo_automodel/components/speculative/dspark/draft_glm_5_2)**
* **[`nemo_automodel.components.speculative.dspark.draft_kimi_k3`](/nemo-automodel/nemo_automodel/components/speculative/dspark/draft_kimi_k3)**
* **[`nemo_automodel.components.speculative.dspark.draft_minimax_m3`](/nemo-automodel/nemo_automodel/components/speculative/dspark/draft_minimax_m3)**
* **[`nemo_automodel.components.speculative.dspark.draft_qwen3`](/nemo-automodel/nemo_automodel/components/speculative/dspark/draft_qwen3)**
* **[`nemo_automodel.components.speculative.dspark.loss`](/nemo-automodel/nemo_automodel/components/speculative/dspark/loss)**
* **[`nemo_automodel.components.speculative.dspark.markov_head`](/nemo-automodel/nemo_automodel/components/speculative/dspark/markov_head)**
* **[`nemo_automodel.components.speculative.dspark.registry`](/nemo-automodel/nemo_automodel/components/speculative/dspark/registry)**
* **[`nemo_automodel.components.speculative.dspark.target`](/nemo-automodel/nemo_automodel/components/speculative/dspark/target)**
* **[`nemo_automodel.components.speculative.dspark.target_utils`](/nemo-automodel/nemo_automodel/components/speculative/dspark/target_utils)**

## Package Contents

### Classes

| Name                                                                                              | Description                              |
| ------------------------------------------------------------------------------------------------- | ---------------------------------------- |
| [`DSparkForwardOutput`](#nemo_automodel-components-speculative-dspark-common-DSparkForwardOutput) | Outputs for one DSpark training forward. |
| [`Qwen3DSparkModel`](#nemo_automodel-components-speculative-dspark-draft_qwen3-Qwen3DSparkModel)  | -                                        |

### Functions

| Name                                                                                            | Description |
| ----------------------------------------------------------------------------------------------- | ----------- |
| [`build_draft_config`](#nemo_automodel-components-speculative-dspark-config-build_draft_config) | -           |
| [`compute_dspark_loss`](#nemo_automodel-components-speculative-dspark-loss-compute_dspark_loss) | -           |

### API

```python
class nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput(
    draft_logits: torch.Tensor,
    target_ids: torch.Tensor,
    eval_mask: torch.Tensor,
    block_keep_mask: torch.Tensor,
    confidence_pred: torch.Tensor | None = None,
    aligned_target_logits: torch.Tensor | None = None
)
```

Dataclass

Outputs for one DSpark training forward.

Shape symbols:
batch\_size: number of samples in the batch
seq\_len: source sequence length
num\_anchors: sampled anchor blocks per sample
block\_size: number of draft positions per anchor
vocab\_size: vocabulary size

The sampler keeps anchors whose first draft target is enabled by
`loss_mask`. Later slots are supervised only while they remain inside
`seq_len` and form a contiguous enabled prefix. Dummy anchors can still
appear when a sample has too few valid anchors; they are masked out by
`block_keep_mask` and `eval_mask`.

**`aligned_target_logits`** `Tensor | None = None`

---

**`block_keep_mask`** `Tensor`

---

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

---

**`draft_logits`** `Tensor`

---

**`eval_mask`** `Tensor`

---

**`target_ids`** `Tensor`

---

```python
class nemo_automodel.components.speculative.dspark.draft_qwen3.Qwen3DSparkModel(
    config
)
```

**Bases:** `Qwen3PreTrainedModel`

**`_no_split_modules`** `= ['Qwen3DSparkDecoderLayer']`

---

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

---

**`rotary_emb`** `= Qwen3RotaryEmbedding(config)`

---

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

---

```python
nemo_automodel.components.speculative.dspark.draft_qwen3.Qwen3DSparkModel._apply(
    fn,
    recurse: bool = True
)
```

Keep the RoPE `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 (see `pin_rope_inv_freq_fp32`).

```python
nemo_automodel.components.speculative.dspark.draft_qwen3.Qwen3DSparkModel._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: transformers.cache_utils.Cache | None = None,
    use_cache: bool = False,
    kwargs = {}
) -> torch.Tensor
```

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

```python
nemo_automodel.components.speculative.dspark.draft_qwen3.Qwen3DSparkModel.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.

Sequence packing (`position_ids` `[B, S]` per-document reset positions,
`seq_lens` `[B, max_docs]`, `doc_remaining` `[B, S]`) keeps every block
inside its anchor's document: the anchor's first target must be in-document,
the block's context prefix and supervision are restricted to that document,
and the draft's RoPE uses the per-document positions.

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

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

```python
nemo_automodel.components.speculative.dspark.draft_qwen3.Qwen3DSparkModel.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_qwen3.Qwen3DSparkModel.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_qwen3.Qwen3DSparkModel.set_embedding_head_trainable(
    trainable: bool
)
```

```python
nemo_automodel.components.speculative.dspark.config.build_draft_config(
    target_config,
    model_args
)
```

```python
nemo_automodel.components.speculative.dspark.loss.compute_dspark_loss(
    outputs: nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput,
    loss_decay_gamma: float | None,
    ce_loss_alpha: float,
    l1_loss_alpha: float,
    confidence_head_alpha: float,
    return_terms: bool = False
)
```