> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

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

Dispatch registry mapping target architecture -> DFlash draft model.

Mirrors the EAGLE registry (`components/speculative/eagle/registry.py`). The
Qwen3 DFlash draft is a non-causal Qwen3-style stack and is config-driven, so
adding a Qwen3-shaped architecture is a one-line append; a target whose backbone
differs (Kimi K3's MLA) registers its own draft class.

## Module Contents

### Classes

| Name                                                                                        | Description                                                             |
| ------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`DFlashDraftSpec`](#nemo_automodel-components-speculative-dflash-registry-DFlashDraftSpec) | How to build a DFlash draft model for a particular target architecture. |

### Functions

| Name                                                                                                            | Description                                                                          |
| --------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------ |
| [`_no_target_kwargs`](#nemo_automodel-components-speculative-dflash-registry-_no_target_kwargs)                 | Default: a target needs no architecture-specific `from_pretrained` kwargs.           |
| [`resolve_dflash_draft_spec`](#nemo_automodel-components-speculative-dflash-registry-resolve_dflash_draft_spec) | Return the first registered DFlash draft spec matching any architecture in the list. |

### Data

[`DFLASH_DRAFT_REGISTRY`](#nemo_automodel-components-speculative-dflash-registry-DFLASH_DRAFT_REGISTRY)

[`_KIMI_K3_ARCHITECTURES`](#nemo_automodel-components-speculative-dflash-registry-_KIMI_K3_ARCHITECTURES)

[`_QWEN3_ARCHITECTURES`](#nemo_automodel-components-speculative-dflash-registry-_QWEN3_ARCHITECTURES)

### API

```python
class nemo_automodel.components.speculative.dflash.registry.DFlashDraftSpec(
    draft_cls: type[torch.nn.Module],
    build_draft_config: typing.Callable[..., transformers.PretrainedConfig] | None = None,
    draft2_cls: type[torch.nn.Module] | None = None,
    build_target_kwargs: typing.Callable[[Any], dict] = _no_target_kwargs,
    attention_backends: tuple[str, ...] = ('flex_attention', 'sdpa', ...,
    supports_context_parallel: bool = True
)
```

Dataclass

How to build a DFlash draft model for a particular target architecture.

**`attention_backends`** `tuple[str, ...] = ('flex_attention', 'sdpa', 'eager')`

The attention backends this draft can run. The
trainer's mask format is driven by the same knob and must agree (a
flex `BlockMask` only works with the flex attention function, a
dense additive mask only with sdpa / eager).

---

**`build_draft_config`** `Callable[..., PretrainedConfig] | None = None`

Builds the draft config from the target's text
config; called with the keyword arguments `num_draft_layers`,
`num_target_layers`, `block_size`, `dflash_config`, and
`attention_backend`.

---

**`build_target_kwargs`** `Callable[[Any], dict] = _no_target_kwargs`

Extra `from_pretrained` keyword arguments for the
frozen target, derived from the recipe's `recipe_args`.

---

**`draft2_cls`** `type[Module] | None = None`

The DFlash 2 draft class for this target: the same backbone
plus the in-block convolutions and the candidate selector. `None`
when the family has no DFlash 2 draft, which makes the DFlash 2
recipe reject the target rather than silently training the plain
DFlash architecture under a DFlash 2 config.

---

**`draft_cls`** `type[Module]`

The draft model class.

---

**`supports_context_parallel`** `bool = True`

Whether the frozen target can be sharded by
the DFlash context-parallel path, which installs a key/value-gather
hook on the target's SDPA call. False for a target that shards the
sequence itself (it declares `_owns_cp_attention` and never routes
through that hook). This is declared here rather than read off the
loaded target because the recipe's other CP gates -- which force the
target onto HuggingFace SDPA -- run before the target exists, and
would otherwise report a misleading error first.

---

```python
nemo_automodel.components.speculative.dflash.registry._no_target_kwargs(
    recipe_cfg
) -> dict
```

Default: a target needs no architecture-specific `from_pretrained` kwargs.

```python
nemo_automodel.components.speculative.dflash.registry.resolve_dflash_draft_spec(
    architectures: list[str]
) -> nemo_automodel.components.speculative.dflash.registry.DFlashDraftSpec
```

Return the first registered DFlash draft spec matching any architecture in the list.

```python
nemo_automodel.components.speculative.dflash.registry.DFLASH_DRAFT_REGISTRY: dict[str, DFlashDraftSpec] = {None: {arch: (DFlashDraftSpec(draft_cls=Qwen3DFlashDraftModel, draft2_cls=Qwen3...
```

```python
nemo_automodel.components.speculative.dflash.registry._KIMI_K3_ARCHITECTURES: tuple[str, ...] = ('KimiK3ForCausalLM', 'KimiK3ForConditionalGeneration')
```

```python
nemo_automodel.components.speculative.dflash.registry._QWEN3_ARCHITECTURES: tuple[str, ...] = ('Qwen3ForCausalLM', 'Qwen3MoeForCausalLM', 'Qwen3_5ForCausalLM', 'Qwen3_5ForCon...
```