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`	How to build a DFlash draft model for a particular target architecture.

Functions

Name	Description
`_no_target_kwargs`	Default: a target needs no architecture-specific `from_pretrained` kwargs.
`resolve_dflash_draft_spec`	Return the first registered DFlash draft spec matching any architecture in the list.

Data

DFLASH_DRAFT_REGISTRY

_KIMI_K3_ARCHITECTURES

_QWEN3_ARCHITECTURES

API

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.

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

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

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.

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

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

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