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')

build_draft_config

Callable[..., PretrainedConfig] | None = None

build_target_kwargs

Callable[[Any], dict] = _no_target_kwargs

draft2_cls

type[Module] | None = None

draft_cls

type[Module]

supports_context_parallel

bool = True

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