> 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.dflash.core

DFlash online training wrapper.

Ported from SpecForge's `specforge/core/dflash.py`. `DFlashTrainerModule`
samples a set of anchor positions per sequence, builds one parallel draft block
per anchor (the block's first token is the real anchor token, the rest are
`MASK`), runs the draft model under a bespoke block attention mask, and
computes a block-wise cross-entropy loss against the ground-truth continuation
of each anchor.

Two training objectives are supported via `loss_type`:

* `"dflash"` (default): the DFlash paper's fixed-anchor objective. Only block
  position 0 is a real token; positions `1..block_size-1` are supervised with
  the decay-weighted CE of Eq. 4 (`w_k = exp(-(k-1)/gamma)`).
* `"variable_prefix"`: the D2SD VP-Drafter objective (arXiv:2606.04446). Each
  block draws a visible-prefix length `l` from a truncated geometric prior
  (`Pr(l) ~ prefix_weight_base ** l`), positions `< l` are filled with the
  real tokens, and only the masked positions `>= l` are supervised, with the
  decay re-anchored at the prefix boundary (`w_k = exp(-(k-l)/gamma)`). This
  trains the draft to re-draft block suffixes behind a variable accepted
  prefix, the regime a variable-prefix drafter sees at inference time.

## Module Contents

### Classes

| Name                                                                                            | Description                                                           |
| ----------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- |
| [`DFlashStepMetrics`](#nemo_automodel-components-speculative-dflash-core-DFlashStepMetrics)     | Per-step training outputs for the DFlash draft.                       |
| [`DFlashTrainerModule`](#nemo_automodel-components-speculative-dflash-core-DFlashTrainerModule) | DFlash online training wrapper with block-wise CE loss.               |
| [`NoValidAnchorsError`](#nemo_automodel-components-speculative-dflash-core-NoValidAnchorsError) | Raised when a batch has no sample long enough to form a DFlash block. |

### Functions

| Name                                                                                                      | Description                                                                    |
| --------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`_context_doc_ids`](#nemo_automodel-components-speculative-dflash-core-_context_doc_ids)                 | Per-context-token document id `[B, S]` from packed `seq_lens` `[B, max_docs]`. |
| [`_to_full_tensor`](#nemo_automodel-components-speculative-dflash-core-_to_full_tensor)                   | Materialise a (possibly tensor-parallel) tensor as a plain local tensor.       |
| [`compute_accept_len`](#nemo_automodel-components-speculative-dflash-core-compute_accept_len)             | Count consecutive accepted predictions for every draft block.                  |
| [`compute_acceptance_stats`](#nemo_automodel-components-speculative-dflash-core-compute_acceptance_stats) | Return mean, sum, and count for bonus-token-inclusive acceptance length.       |

### Data

[`_DFLASH_LOSS_TYPES`](#nemo_automodel-components-speculative-dflash-core-_DFLASH_LOSS_TYPES)

### API

```python
class nemo_automodel.components.speculative.dflash.core.DFlashStepMetrics(
    loss: torch.Tensor,
    loss_weight: torch.Tensor,
    accuracy: torch.Tensor,
    valid_tokens: torch.Tensor,
    correct_tokens: torch.Tensor,
    accept_len: torch.Tensor,
    accept_len_sum: torch.Tensor,
    valid_blocks: torch.Tensor
)
```

Dataclass

Per-step training outputs for the DFlash draft.

The count and sum fields retain the additive statistics needed for
token-weighted, distributed validation. Averaging `accuracy` or
`accept_len` per micro-batch would bias batches with fewer valid tokens or blocks.

```python
class nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule(
    draft_model: nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashDraftModel,
    target_lm_head: torch.nn.Module,
    target_embed_tokens: torch.nn.Module,
    mask_token_id: int,
    block_size: int = 16,
    attention_backend: str = 'flex_attention',
    num_anchors: int = 512,
    loss_decay_gamma: typing.Optional[float] = None,
    loss_type: str = 'dflash',
    prefix_weight_base: float = 0.9
)
```

**Bases:** `Module`

DFlash online training wrapper with block-wise CE loss.

```python
nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._build_block_targets(
    input_ids: torch.Tensor,
    loss_mask: torch.Tensor,
    anchor_positions: torch.Tensor,
    block_keep_mask: torch.Tensor,
    seq_len: int,
    label_start: int = 0,
    doc_remaining: torch.Tensor | None = None
) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Per-block ground-truth tokens and the supervised-position mask.

Returns `(label_indices, target_ids, block_mask)` each of shape
`[B, N, block_size]`. `label_indices[..., k]` is the sequence position
block position `k` predicts (`anchor + label_start + k`); `block_mask`
is the product of block validity, in-bounds, and the gathered loss mask.
Shared by the block-wise trainers (DFlash, JetSpec, and Domino, which passes
`label_start=1` for `shift_label`) so the label/mask gathering lives in
one place. Under packing, `doc_remaining` `[B, S]` truncates labels at
the anchor's document boundary: anchor sampling only keeps offsets up to
`block_size - 1` inside the anchor's document, and a shifted label window
(`label_start > 0`) reaches one past that guarantee.

```python
nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._create_noise_embed(
    input_ids,
    anchor_positions,
    block_keep_mask
)
```

Embed each block as `[anchor_token, MASK, MASK, ...]` (invalid blocks all MASK).

```python
nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._create_position_ids(
    anchor_positions: torch.Tensor,
    context_position_ids: typing.Optional[torch.Tensor] = None
) -> torch.Tensor
```

Position ids for the parallel draft blocks (anchor 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 (`context_position_ids[anchor] + offset`)
to keep the draft's RoPE phase document-local.

```python
nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._create_vp_noise_embed(
    input_ids,
    anchor_positions,
    block_keep_mask,
    prefix_lengths
)
```

Embed the draft blocks with a visible prefix of real tokens, then `MASK`.

The variable-prefix analogue of :meth:`_create_noise_embed`: block
positions `< prefix_lengths` hold the real sequence tokens, the rest
(and every position of an invalid block) hold `MASK`.

**Parameters:**

Long tensor of shape `[batch, sequence]`.

Long tensor of shape `[batch, blocks]`; each block's
start position in the sequence.

Bool tensor of shape `[batch, blocks]`; invalid
(padding) blocks are embedded as all `MASK`.

Long tensor of shape `[batch, blocks]`; this block's
visible-prefix length.

**Returns:**

Tensor of shape `[batch, blocks * block_size, hidden]`.

```python
nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._prepare_block_inputs(
    input_ids: torch.Tensor,
    loss_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None,
    causal: bool = False
) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, 'torch.Tensor | BlockMask', typing.Optional[torch.Tensor]]
```

Shared block-drafting prologue: anchors, noise embedding, positions, mask.

Centralises the sequence-packing handling for every block-wise trainer
(DFlash and its Domino / JetSpec subclasses): anchors are sampled so each
block stays inside one document, the block's context prefix is restricted
to its anchor's document, and the draft RoPE uses per-document positions.

Under `loss_type="variable_prefix"` the block embedding shows a sampled
real-token prefix (`_create_vp_noise_embed`) instead of the single anchor
token, and the sampled `prefix_lengths` are returned for the loss; every
other `loss_type` returns `prefix_lengths=None`.

**Parameters:**

`[B, S]` context token ids (long).

`[B, S]` supervised-token mask.

`[B, S]` per-document reset positions (packing), or `None`.

`[B, max_docs]` packed document lengths, or `None` (unpacked).

`[B, S]` remaining real tokens of each position's document.

When True, build the in-block-causal (JetSpec) mask instead
of the bidirectional (DFlash / Domino) one.

**Returns:** `torch.Tensor`

\`\`(anchor\_positions \[B, N], block\_keep\_mask \[B, N], noise\_embedding

```python
nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._sample_anchor_positions(
    seq_len: int,
    loss_mask: torch.Tensor,
    device: torch.device,
    doc_remaining: typing.Optional[torch.Tensor] = None
) -> typing.Tuple[torch.Tensor, torch.Tensor]
```

Randomly sample anchor positions per sample; returns `(anchors, keep_mask)`.

`doc_remaining` `[B, S]` (sequence packing) restricts anchors so the
whole block stays inside one document (`doc_remaining >= block_size - 1`),
the per-document analogue of the `anchor <= seq_len - block_size` bound.
This is required for correctness -- `_build_block_targets` gathers labels
by absolute offset and does not encode document boundaries, so a block that
crossed one would be supervised on the next document's tokens. A side effect
is that a packed document shorter than `block_size` yields no anchors (the
unpacked path still supervises such a short sequence's partial block); pack
with documents at least `block_size` long to avoid dropping their signal.

```python
nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._sample_prefix_lengths(
    bsz: int,
    n_blocks: int,
    device: torch.device
) -> torch.Tensor
```

Sample a visible-prefix length per block for variable-prefix training.

A prefix length `l` means block positions `[0, l)` are visible real
tokens and positions `[l, block_size)` are masked prediction targets.
Lengths follow D2SD's truncated geometric prior `Pr(l) ~ base ** l` over
`[min(2, block_size - 1), block_size - 1]`; a base below 1 biases toward
short prefixes. The lower bound skips the degenerate fixed-anchor DFlash
case, and the upper bound keeps at least one masked target per block.

**Returns:** `torch.Tensor`

Long tensor of shape `[batch, blocks]`.

```python
nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._variable_prefix_loss(
    logits: torch.Tensor,
    target_ids: torch.Tensor,
    block_mask: torch.Tensor,
    prefix_lengths: torch.Tensor
) -> nemo_automodel.components.speculative.dflash.core.DFlashStepMetrics
```

Decay-weighted CE over each block's masked suffix (D2SD Eq. for L\_VP).

Only positions at or past the sampled visible prefix are supervised, and
the exponential decay restarts at the prefix boundary:
`w_k = exp(-(k - l) / gamma)` for block position `k >= l`
(`loss_decay_gamma=None` disables decay). The loss is the weighted mean
`sum(nll * w) / sum(w)`, mirroring the fixed-anchor path's
`normalize="mean"`. Assumes every `prefix_lengths` entry is at least
`min(2, block_size - 1)` (what :meth:`_sample_prefix_lengths` produces),
so the leading always-visible positions can be sliced off before the CE.

**Parameters:**

Tensor of shape `[batch, blocks, block_size, vocab]`.

Long tensor of shape `[batch, blocks, block_size]`; the
ground-truth token at `anchor + k` for block position `k`.

Tensor of shape `[batch, blocks, block_size]`; 0/1
product of block validity, in-bounds, and the loss mask.

Long tensor of shape `[batch, blocks]`; this block's
visible-prefix length.

**Returns:** `DFlashStepMetrics`

DFlashStepMetrics with the weighted-mean loss, the argmax accuracy

```python
nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule.forward(
    input_ids: torch.Tensor,
    hidden_states: torch.Tensor,
    loss_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.dflash.core.DFlashStepMetrics
```

Parallel block-wise training forward pass.

Sequence packing (`position_ids` `[B, S]` per-document reset positions,
`seq_lens` `[B, max_docs]` document lengths, `doc_remaining` `[B, S]`)
keeps every block inside one document: anchors are constrained so the block
does not cross a boundary, the block's context prefix attends only within the
anchor's document, and the draft's RoPE uses the per-document positions.

```python
class nemo_automodel.components.speculative.dflash.core.NoValidAnchorsError()
```

**Bases:** `ValueError`

Raised when a batch has no sample long enough to form a DFlash block.

A DFlash anchor needs at least `block_size + 1` supervised tokens (the
anchor plus its block). Datasets always contain some short conversations;
the training loop catches this and skips the offending micro-batch rather
than aborting the run.

```python
nemo_automodel.components.speculative.dflash.core._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.dflash.core._to_full_tensor(
    tensor: torch.Tensor
) -> torch.Tensor
```

Materialise a (possibly tensor-parallel) tensor as a plain local tensor.

Under tensor parallelism the target's column-parallel `lm_head` and
vocab-parallel `embed_tokens` return `DTensor` outputs. The draft and the
block-wise loss consume plain tensors, so gather the full tensor. A no-op for
an already-plain (unsharded / replicated) tensor.

```python
nemo_automodel.components.speculative.dflash.core.compute_accept_len(
    pred_ids_4d: torch.Tensor,
    target_ids_4d: torch.Tensor,
    valid_mask_4d: torch.Tensor
) -> torch.Tensor
```

Count consecutive accepted predictions for every draft block.

**Parameters:**

Long tensor of shape `[batch, blocks, depth]`.

Long tensor of shape `[batch, blocks, depth]`.

Bool tensor of shape `[batch, blocks, depth]`.

**Returns:** `torch.Tensor`

Float tensor of shape `[batch, blocks]` containing the accepted draft

```python
nemo_automodel.components.speculative.dflash.core.compute_acceptance_stats(
    pred_ids_4d: torch.Tensor,
    target_ids_4d: torch.Tensor,
    valid_mask_4d: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Return mean, sum, and count for bonus-token-inclusive acceptance length.

**Parameters:**

Long tensor of shape `[batch, blocks, depth]`.

Long tensor of shape `[batch, blocks, depth]`.

Bool tensor of shape `[batch, blocks, depth]`.

**Returns:** `torch.Tensor`

Three scalar tensors: mean acceptance length, its additive sum, and the

```python
nemo_automodel.components.speculative.dflash.core._DFLASH_LOSS_TYPES = ('dflash', 'variable_prefix')
```