nemo_automodel.components.speculative.dspark.registry

Dispatch registry mapping a target architecture string to a DSpark draft model.

Mirrors the EAGLE / DFlash registries: keyed by the target’s HF architectures string so adding a target family is a one-line append, with no recipe change.

Module Contents

Classes

Name	Description
`DraftSpec`	How to build a DSpark draft model for a particular target architecture.

Functions

Name	Description
`build_target_layer_ids`	Evenly spread `num_feature_layers` feature layers across the target depth.
`resolve_dspark_draft_spec`	Return the first registered DSpark draft spec matching `architectures`.

Data

DSPARK_DRAFT_REGISTRY

_DENSE_ARCHITECTURES

__all__

API

class nemo_automodel.components.speculative.dspark.registry.DraftSpec(
    draft_cls: type[transformers.PreTrainedModel]
)

Dataclass

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

draft_cls

type[PreTrainedModel]

nemo_automodel.components.speculative.dspark.registry.build_target_layer_ids(
    num_target_layers: int,
    num_feature_layers: int
) -> list[int]

Evenly spread num_feature_layers feature layers across the target depth.

Used as the default when target_layer_ids is not given. Returns strictly increasing ids in [1, num_target_layers - 1] (the embedding output, id -1/0, is excluded by default to match the paper’s choice of mid/late feature layers).

nemo_automodel.components.speculative.dspark.registry.resolve_dspark_draft_spec(
    architectures: list[str]
) -> nemo_automodel.components.speculative.dspark.registry.DraftSpec

Return the first registered DSpark draft spec matching architectures.

nemo_automodel.components.speculative.dspark.registry.DSPARK_DRAFT_REGISTRY: dict[str, DraftSpec] = {arch: (DraftSpec(draft_cls=Qwen3DSparkModel)) for arch in _DENSE_ARCHITECTURES}

nemo_automodel.components.speculative.dspark.registry._DENSE_ARCHITECTURES: tuple[str, ...] = ('Qwen3ForCausalLM', 'Qwen3MoeForCausalLM')

nemo_automodel.components.speculative.dspark.registry.__all__ = ['DraftSpec', 'DSPARK_DRAFT_REGISTRY', 'resolve_dspark_draft_spec', 'build_targe...