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

## Module Contents

### Classes

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

### Functions

| Name                                                                                                              | Description                                                                    |
| ----------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`build_anchor_candidate_mask`](#nemo_automodel-components-speculative-dspark-common-build_anchor_candidate_mask) | -                                                                              |
| [`build_eval_mask`](#nemo_automodel-components-speculative-dspark-common-build_eval_mask)                         | -                                                                              |
| [`context_doc_ids`](#nemo_automodel-components-speculative-dspark-common-context_doc_ids)                         | Per-context-token document id `[B, S]` from packed `seq_lens` `[B, max_docs]`. |
| [`create_noise_embed`](#nemo_automodel-components-speculative-dspark-common-create_noise_embed)                   | -                                                                              |
| [`create_position_ids`](#nemo_automodel-components-speculative-dspark-common-create_position_ids)                 | Position ids for the parallel draft blocks (`base position + offset`).         |
| [`extract_context_feature`](#nemo_automodel-components-speculative-dspark-common-extract_context_feature)         | -                                                                              |
| [`pin_rope_inv_freq_fp32`](#nemo_automodel-components-speculative-dspark-common-pin_rope_inv_freq_fp32)           | Keep a RoPE module's `inv_freq` buffers in fp32 after a dtype cast.            |
| [`sample_anchor_positions`](#nemo_automodel-components-speculative-dspark-common-sample_anchor_positions)         | -                                                                              |
| [`validate_target_layer_ids`](#nemo_automodel-components-speculative-dspark-common-validate_target_layer_ids)     | -                                                                              |

### Data

[`__all__`](#nemo_automodel-components-speculative-dspark-common-__all__)

[`logger`](#nemo_automodel-components-speculative-dspark-common-logger)

### API

```python
class nemo_automodel.components.speculative.dspark.common.AcceptRatePredictor(
    input_dim: int
)
```

**Bases:** `Module`

```python
nemo_automodel.components.speculative.dspark.common.AcceptRatePredictor.forward(
    features
)
```

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

Dataclass

Outputs for one DSpark training forward.

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

```python
nemo_automodel.components.speculative.dspark.common.build_anchor_candidate_mask(
    seq_len: int,
    loss_mask: torch.Tensor,
    doc_remaining: typing.Optional[torch.Tensor] = None
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dspark.common.build_eval_mask(
    seq_len: int,
    loss_mask: torch.Tensor,
    label_indices: torch.Tensor,
    safe_label_indices: torch.Tensor,
    block_keep_mask: torch.Tensor,
    doc_remaining: typing.Optional[torch.Tensor] = None,
    anchor_positions: typing.Optional[torch.Tensor] = None
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dspark.common.context_doc_ids(
    seq_lens: torch.Tensor,
    seq_len: int,
    device: torch.device
) -> torch.Tensor
```

Per-context-token document id `[B, S]` from packed `seq_lens` `[B, max_docs]`.

Mirrors the `doc_id` construction in `build_block_causal_additive_mask`: a
token's id is the number of document boundaries at or before its position, so
0-length padding entries never split a real document.

```python
nemo_automodel.components.speculative.dspark.common.create_noise_embed(
    embed_tokens: torch.nn.Module,
    input_ids: torch.Tensor,
    anchor_positions: torch.Tensor,
    block_keep_mask: torch.Tensor,
    mask_token_id: int,
    block_size: int
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dspark.common.create_position_ids(
    anchor_positions: torch.Tensor,
    block_size: int,
    context_position_ids: typing.Optional[torch.Tensor] = None
) -> torch.Tensor
```

Position ids for the parallel draft blocks (`base position + offset`).

Without packing the anchor's position equals its row index, so the block
positions are `anchor + offset`. Under packing `context_position_ids`
`[B, S]` holds per-document reset positions, so the block's base position is
gathered from it at the anchor to keep the draft's RoPE phase document-local.

```python
nemo_automodel.components.speculative.dspark.common.extract_context_feature(
    hidden_states,
    layer_ids
)
```

```python
nemo_automodel.components.speculative.dspark.common.pin_rope_inv_freq_fp32(
    rotary_emb: typing.Optional[torch.nn.Module]
) -> None
```

Keep a RoPE module's `inv_freq` buffers in fp32 after a dtype cast.

`module.to(bfloat16)` (the training build path) rounds the rotary
`inv_freq` buffers to bf16. The rounded frequencies dephase with absolute
position, so the train/inference RoPE diverges (worse with longer context)
and draft acceptance erodes, while the serving runtime keeps an fp32 RoPE
cache. A bf16 round-trip cannot be undone by upcasting, so recompute fresh
fp32 frequencies from a fresh rotary module built off the same config (the
same values HF derives on the fp32 paths) and copy them back in.

Works for both the single-buffer layout (`inv_freq` / `original_inv_freq`,
e.g. Qwen3) and the per-layer-type layout where each frequency buffer is named
`<layer_type>_inv_freq` (e.g. Gemma4). No-op when every frequency buffer is
already fp32 or on a meta device.

```python
nemo_automodel.components.speculative.dspark.common.sample_anchor_positions(
    seq_len: int,
    loss_mask: torch.Tensor,
    num_anchors: int,
    device: torch.device,
    doc_remaining: typing.Optional[torch.Tensor] = None
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.speculative.dspark.common.validate_target_layer_ids(
    layer_ids,
    num_target_layers: int
)
```

```python
nemo_automodel.components.speculative.dspark.common.__all__ = ['DSparkForwardOutput', 'AcceptRatePredictor', 'extract_context_feature', 'valid...
```

```python
nemo_automodel.components.speculative.dspark.common.logger = logging.getLogger(__name__)
```