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

## Module Contents

### Classes

| Name                                                                                           | Description                                                                    |
| ---------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`GatedMarkovHead`](#nemo_automodel-components-speculative-dspark-markov_head-GatedMarkovHead) | -                                                                              |
| [`RNNHead`](#nemo_automodel-components-speculative-dspark-markov_head-RNNHead)                 | RNN-based head that maintains recurrent state across positions within a block. |
| [`VanillaMarkov`](#nemo_automodel-components-speculative-dspark-markov_head-VanillaMarkov)     | -                                                                              |

### Functions

| Name                                                                                               | Description |
| -------------------------------------------------------------------------------------------------- | ----------- |
| [`build_markov_head`](#nemo_automodel-components-speculative-dspark-markov_head-build_markov_head) | -           |

### Data

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

### API

```python
class nemo_automodel.components.speculative.dspark.markov_head.GatedMarkovHead(
    vocab_size: int,
    markov_rank: int,
    hidden_size: int
)
```

**Bases:** [VanillaMarkov](#nemo_automodel-components-speculative-dspark-markov_head-VanillaMarkov)

```python
nemo_automodel.components.speculative.dspark.markov_head.GatedMarkovHead.compute_gate(
    token_ids: torch.Tensor,
    hidden_states: typing.Optional[torch.Tensor]
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dspark.markov_head.GatedMarkovHead.compute_step_bias(
    token_ids: torch.Tensor,
    hidden_states: typing.Optional[torch.Tensor]
) -> torch.Tensor
```

```python
class nemo_automodel.components.speculative.dspark.markov_head.RNNHead(
    vocab_size: int,
    markov_rank: int,
    hidden_size: int
)
```

**Bases:** [VanillaMarkov](#nemo_automodel-components-speculative-dspark-markov_head-VanillaMarkov)

RNN-based head that maintains recurrent state across positions within a block.

Unlike the memoryless Markov heads, position k can access the full prefix
history x\_\{\<k} through a GRU-like recurrent state.

```python
nemo_automodel.components.speculative.dspark.markov_head.RNNHead._rnn_step(
    state: torch.Tensor,
    prev_embeddings: torch.Tensor,
    hidden_states: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Single RNN step.

**Parameters:**

\[\*, r] previous recurrent state

\[\*, r] W1\[x\_\{k-1}]

\[\*, d] backbone hidden at step k

**Returns:** `torch.Tensor`

\[\*, r]

```python
nemo_automodel.components.speculative.dspark.markov_head.RNNHead.apply_block_logits(
    base_logits: torch.Tensor,
    token_ids: torch.Tensor,
    hidden_states: typing.Optional[torch.Tensor]
) -> torch.Tensor
```

Apply RNN bias during training (teacher-forced, unrolled over block\_size).

**Parameters:**

\[B, num\_blocks, block\_size, V]

\[B, num\_blocks, block\_size] - prev token ids for each position

\[B, num\_blocks, block\_size, d]

```python
nemo_automodel.components.speculative.dspark.markov_head.RNNHead.compute_step_bias(
    token_ids: torch.Tensor,
    hidden_states: typing.Optional[torch.Tensor]
) -> torch.Tensor
```

Stateless single-step bias (state initialized to zero).

This is used for compatibility but does NOT carry state across steps.
For full RNN behavior, use apply\_block\_logits or sample\_block\_tokens.

```python
nemo_automodel.components.speculative.dspark.markov_head.RNNHead.sample_block_tokens(
    base_logits: torch.Tensor,
    first_prev_token_ids: torch.Tensor,
    hidden_states: typing.Optional[torch.Tensor],
    temperature: float = 0.0
) -> tuple[torch.Tensor, torch.Tensor]
```

Autoregressive sampling with RNN state.

**Parameters:**

\[batch, proposal\_len, vocab]

\[batch] - token preceding the first draft position

\[batch, proposal\_len, d]

sampling temperature

**Returns:** `torch.Tensor`

\[batch, proposal\_len]

```python
class nemo_automodel.components.speculative.dspark.markov_head.VanillaMarkov(
    vocab_size: int,
    markov_rank: int
)
```

**Bases:** `Module`

```python
nemo_automodel.components.speculative.dspark.markov_head.VanillaMarkov.apply_block_logits(
    base_logits: torch.Tensor,
    token_ids: torch.Tensor,
    hidden_states: typing.Optional[torch.Tensor]
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dspark.markov_head.VanillaMarkov.apply_step_logits(
    logits: torch.Tensor,
    token_ids: torch.Tensor,
    hidden_states: typing.Optional[torch.Tensor]
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dspark.markov_head.VanillaMarkov.compute_step_bias(
    token_ids: torch.Tensor,
    hidden_states: typing.Optional[torch.Tensor]
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dspark.markov_head.VanillaMarkov.get_prev_embeddings(
    token_ids: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dspark.markov_head.VanillaMarkov.project_bias(
    latent_states: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dspark.markov_head.VanillaMarkov.sample_block_tokens(
    base_logits: torch.Tensor,
    first_prev_token_ids: torch.Tensor,
    hidden_states: typing.Optional[torch.Tensor],
    temperature: float = 0.0
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.speculative.dspark.markov_head.build_markov_head(
    config
) -> torch.nn.Module | None
```

```python
nemo_automodel.components.speculative.dspark.markov_head.__all__ = ['VanillaMarkov', 'GatedMarkovHead', 'RNNHead', 'build_markov_head']
```