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

Target-model wrapper for DFlash training.

Unlike EAGLE-3 (which captures exactly three aux layers and left-shifts the
supervision), DFlash captures an arbitrary set of decoder layers, concatenates
them along the feature dim, and feeds the result to the draft as *context*. No
shifting is applied -- the DFlash block attention mask handles anchor alignment.

## Module Contents

### Classes

| Name                                                                                              | Description                                                              |
| ------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`DFlashTargetBatch`](#nemo_automodel-components-speculative-dflash-target-DFlashTargetBatch)     | Target-model context features needed by the DFlash trainer.              |
| [`HFDFlashTargetModel`](#nemo_automodel-components-speculative-dflash-target-HFDFlashTargetModel) | Capture a set of decoder-layer hidden states from a frozen HF causal LM. |

### Functions

| Name                                                                                              | Description                                                              |
| ------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`resolve_text_config`](#nemo_automodel-components-speculative-dflash-target-resolve_text_config) | Return the decoder-stack config, unwrapping a multimodal wrapper config. |

### API

```python
class nemo_automodel.components.speculative.dflash.target.DFlashTargetBatch(
    hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    logits: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
)
```

Dataclass

Target-model context features needed by the DFlash trainer.

`position_ids` / `seq_lens` / `doc_remaining` are `None` off the
packing path and carry the (unshifted) packing metadata to the trainer on it.

**`attention_mask`** `Tensor`

---

**`doc_remaining`** `Tensor | None = None`

---

**`hidden_states`** `Tensor`

---

**`input_ids`** `Tensor`

---

**`logits`** `Tensor | None = None`

---

**`loss_mask`** `Tensor`

---

**`position_ids`** `Tensor | None = None`

---

**`seq_lens`** `Tensor | None = None`

---

```python
class nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel(
    model: torch.nn.Module,
    target_layer_ids: typing.Sequence[int],
    capture_logits: bool = False,
    cp_mesh = None
)
```

Capture a set of decoder-layer hidden states from a frozen HF causal LM.

A forward hook on decoder layer `i` captures that layer's output, which in
HuggingFace's `output_hidden_states` convention is `hidden_states[i + 1]`
\-- matching SpecForge's `extract_context_feature` (offset 1).

**`_cp_size`** `= cp_mesh.size() if cp_mesh is not None else 1`

---

**`capture_logits`** `= bool(capture_logits)`

---

**`model`** `= model.eval()`

---

**`target_layer_ids`** `= self._validate_layer_ids(target_layer_ids)`

---

```python
nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel._check_captured(
    captured: dict[int, torch.Tensor]
) -> None
```

```python
nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel._get_transformer_layers() -> list[torch.nn.Module]
```

Return decoder layers as an ordered, integer-indexable list.

```python
nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel._validate_layer_ids(
    target_layer_ids: typing.Sequence[int]
) -> list[int]
```

```python
nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel.generate_batch(
    input_ids: torch.Tensor,
    attention_mask: 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.target.DFlashTargetBatch
```

Run the target model and capture the selected layers' hidden states as context.

With `seq_lens` (`[B, max_docs]`, per-document lengths summing to `S`)
the target runs with a document-level block-causal mask and per-document
`position_ids` so the captured context hidden states do not leak across
document boundaries (SDPA/eager consume the `[B, 1, S, S]` block-causal
additive mask; FlashAttention infers boundaries from `position_ids` at batch
size 1). The packing metadata is carried through to the trainer unchanged.

```python
nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel.get_input_embeddings() -> torch.nn.Embedding
```

Return the target model input embeddings.

```python
nemo_automodel.components.speculative.dflash.target.resolve_text_config(
    config: typing.Any
) -> typing.Any
```

Return the decoder-stack config, unwrapping a multimodal wrapper config.

A `*ForConditionalGeneration` target (e.g. `Qwen3_5ForConditionalGeneration`,
which is what `Qwen/Qwen3.8-27B` ships as) keeps `num_hidden_layers`,
`hidden_size`, and `vocab_size` on a nested `text_config` and leaves the
outer config holding only the vision/text composition. Causal-LM targets keep
them at the top level, so this is the identity for them.

**Parameters:**

**`config`** `Any`

A `PretrainedConfig` for the frozen target model.

---

**Returns:** `Any`

`config.text_config` when present, otherwise `config` unchanged.