Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsSpeculativeDsparknemo_automodel.components.speculative.dspark.registry

nemo_automodel.components.speculative.dspark.registry

View as MarkdownOpen in Claude

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

NameDescription
DraftSpecHow to build a DSpark draft model for a particular target architecture.
ModelOwnedDSparkProviderTarget config that owns construction of its architecture-specific draft.
_DSparkDraftOptionsRecipe options consumed by model-owned DSpark draft constructors.

Functions

NameDescription
build_target_layer_idsEvenly spread num_feature_layers feature layers across the target depth.
resolve_dspark_draft_specReturn the first registered DSpark draft spec matching architectures.

Data

DSPARK_DRAFT_REGISTRY

_DENSE_ARCHITECTURES

API

class nemo_automodel.components.speculative.dspark.registry.DraftSpec(
    draft_cls: type[torch.nn.Module]
)
Dataclass

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

draft_cls
type[Module]
class nemo_automodel.components.speculative.dspark.registry.ModelOwnedDSparkProvider()
Protocol

Target config that owns construction of its architecture-specific draft.

nemo_automodel.components.speculative.dspark.registry.ModelOwnedDSparkProvider.build_dspark_draft(
    options: nemo_automodel.components.speculative.dspark.registry._DSparkDraftOptions
) -> torch.nn.Module

Build a draft from the generic recipe’s declarative options.

class nemo_automodel.components.speculative.dspark.registry._DSparkDraftOptions()
Protocol

Recipe options consumed by model-owned DSpark draft constructors.

block_size
int
confidence_head_alpha
float
confidence_head_stop_gradient
bool
confidence_head_with_markov
bool
markov_head_type
str
markov_rank
int
mask_token_id
int
num_anchors
int
num_draft_layers
int
target_layer_ids
list[int]
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')