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

DFlash 2 online training wrapper.

DFlash 2 ([https://inco.ai/blog/dflash2/](https://inco.ai/blog/dflash2/)) keeps DFlash's parallel block draft and
adds a two-tap in-block convolution to the backbone plus a pairwise path selector
over each position's top-k candidates (see
`nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2`). The
convolution needs no new supervision -- it is part of the backbone and is trained
by the ordinary DFlash objective -- so this wrapper differs from
`DFlashTrainerModule` in exactly one place: it adds a second term that trains
the selector::

loss = base\_loss + selector\_loss\_weight \* selector\_loss

`base_loss` is DFlash's decay-weighted block CE over the full vocabulary, i.e.
what makes each position's candidate list good. `selector_loss` is a CE over
the `selector_top_k` candidates of that same position, scored against the
*ground-truth* predecessor (the token the walk would have committed had every
earlier position been right) and supervised with the index of the true token
inside the candidate list. Positions whose true token missed the candidate list
carry no selector signal -- there is nothing there to select -- and are excluded
from the selector term; `candidate_recall` reports how often that happens.

Both terms use the same block-position decay weights, so a position's importance
is identical in the two objectives.

## Module Contents

### Classes

| Name                                                                                                      | Description                                                                 |
| --------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`DFlash2StepMetrics`](#nemo_automodel-components-speculative-dflash-dflash2_core-DFlash2StepMetrics)     | Per-step training outputs for the DFlash 2 draft.                           |
| [`DFlash2TrainerModule`](#nemo_automodel-components-speculative-dflash-dflash2_core-DFlash2TrainerModule) | DFlash 2 online training wrapper: DFlash block CE + candidate-selection CE. |

### API

```python
class nemo_automodel.components.speculative.dflash.dflash2_core.DFlash2StepMetrics(
    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,
    base_loss: torch.Tensor,
    selector_loss: torch.Tensor,
    base_accuracy: torch.Tensor,
    base_correct_tokens: torch.Tensor,
    base_accept_len: torch.Tensor,
    base_accept_len_sum: torch.Tensor,
    candidate_recall: torch.Tensor
)
```

Dataclass

Per-step training outputs for the DFlash 2 draft.

`loss` / `accuracy` / `valid_tokens` mirror `DFlashStepMetrics` so the
shared DFlash training loop consumes them unchanged. The primary accuracy and
acceptance-length fields describe the *selector* path -- what the draft
actually emits at decode time -- and the `base_*` fields the backbone's own
top-1 picks, so the two are directly comparable on the same denominator.

**`accept_len`** `Tensor`

---

**`accept_len_sum`** `Tensor`

---

**`accuracy`** `Tensor`

---

**`base_accept_len`** `Tensor`

---

**`base_accept_len_sum`** `Tensor`

---

**`base_accuracy`** `Tensor`

---

**`base_correct_tokens`** `Tensor`

---

**`base_loss`** `Tensor`

---

**`candidate_recall`** `Tensor`

---

**`correct_tokens`** `Tensor`

---

**`loss`** `Tensor`

---

**`loss_weight`** `Tensor`

---

**`selector_loss`** `Tensor`

---

**`valid_blocks`** `Tensor`

---

**`valid_tokens`** `Tensor`

---

```python
class nemo_automodel.components.speculative.dflash.dflash2_core.DFlash2TrainerModule(
    draft_model: nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.Qwen3DFlash2DraftModel,
    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: float | None = None,
    selector_loss_weight: float = 1.0,
    sliding_window: int | None = None
)
```

**Bases:** [DFlashTrainerModule](/nemo-automodel/nemo_automodel/components/speculative/dflash/core#nemo_automodel-components-speculative-dflash-core-DFlashTrainerModule)

DFlash 2 online training wrapper: DFlash block CE + candidate-selection CE.

**`selector_loss_weight`** `= float(selector_loss_weight)`

---

```python
nemo_automodel.components.speculative.dflash.dflash2_core.DFlash2TrainerModule._depth_weights(
    mask: torch.Tensor
) -> torch.Tensor
```

Block-position decay weights for the `block_size - 1` predicted positions.

**Parameters:**

**`mask`** `torch.Tensor`

Tensor of shape \[batch, blocks, depth]; the supervised-position
mask the weights are multiplied into.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, blocks, depth] equal to `mask` scaled by

```python
nemo_automodel.components.speculative.dflash.dflash2_core.DFlash2TrainerModule._selector_scores(
    hidden: torch.Tensor,
    logits: torch.Tensor,
    target_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
```

Score each supervised position's top-k candidates against its true predecessor.

Teacher-forces the predecessor: position `k` is scored after the token at
`anchor + k - 1`, which is what the decode-time walk would have committed
if every earlier position had been accepted. Position 1's predecessor is the
anchor token itself, exactly as at decode time.

**Parameters:**

**`hidden`** `torch.Tensor`

Tensor of shape \[batch, blocks, depth, hidden]; the draft hidden
states of the predicted (non-anchor) block positions.

---

**`logits`** `torch.Tensor`

Tensor of shape \[batch, blocks, depth, vocab]; the backbone
logits at those positions.

---

**`target_ids`** `torch.Tensor`

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

---

**Returns:** `torch.Tensor`

Tuple `(scores, candidate_ids, target_index, has_target)`: `scores`

```python
nemo_automodel.components.speculative.dflash.dflash2_core.DFlash2TrainerModule.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.dflash2_core.DFlash2StepMetrics
```

Parallel block-wise training forward with the DFlash 2 path selector.

Sequence packing (`position_ids` `[B, S]` per-document reset positions,
`seq_lens` `[B, max_docs]` document lengths, `doc_remaining` `[B, S]`)
is handled by the shared DFlash prologue, which keeps every block inside one
document.

**Parameters:**

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

Long tensor of shape \[batch, sequence]; the context tokens.

---

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, layers \* hidden]; the
captured target-model context features.

---

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

Tensor of shape \[batch, sequence]; the supervised-token mask.

---

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

Long tensor of shape \[batch, sequence] with per-document
reset positions under packing, or `None`.

---

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

Long tensor of shape \[batch, max\_docs] with packed document
lengths, or `None` when unpacked.

---

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

Long tensor of shape \[batch, sequence]; real tokens left
in each position's document, or `None` when unpacked.

---

**Returns:** `DFlash2StepMetrics`

DFlash2StepMetrics for this micro-batch.