For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

DraftRegistration

DraftRegistration

class max.pipelines.lib.registration_draft.DraftRegistration(method, configure, create_proposer=None, target_hidden_states=<function default_target_hidden_states>, shared_modules=(), optional_weights=frozenset({}), drafter_export=None, strict_weights=True)

source

Bases: Generic[_ModuleT, _ConfigT]

The construction and proposer contract of a draft.

The config factory only finalizes configuration. The framework loads the checkpoint, constructs the registered module and drives it through its proposer with the method’s driver. Other drafts retain their existing builders until migrated.

Parameters:

• method (SpeculativeMethod)
• configure (Callable[[ModelLoadContext, Mapping[str, WeightData]], _ConfigT])
• create_proposer (DraftProposerFactory[_ModuleT, _ConfigT] | None)
• target_hidden_states (Callable[[SpeculativeConfig, _ConfigT], TargetHiddenStates])
• shared_modules (tuple[Literal['embed_tokens', 'lm_head'], ...])
• optional_weights (frozenset[str])
• drafter_export (Callable[[ModelLoadContext, _ConfigT], object] | None)
• strict_weights (bool)

bind_shared_modules()

bind_shared_modules(target, draft)

source

Aliases declared target modules onto the draft before tracing.

Raises:

ValueError – If either component lacks a declared module.

Parameters:

• target (object)
• draft (object)

Return type:

None

configure

configure: Callable[[ModelLoadContext, Mapping[str, WeightData]], _ConfigT]

source

Finalizes the draft’s config from its checkpoint, before the module is constructed.

create_proposer

create_proposer: DraftProposerFactory[_ModuleT, _ConfigT] | None = None

source

Returns the draft’s proposer, for a draft the method’s own proposer can’t drive. It is a fallback: a target adapter’s proposer, when the target’s architecture has one, runs instead. None runs the method’s own, as sequential_proposer() and block_proposer() pick it.

drafter_export

drafter_export: Callable[[ModelLoadContext, _ConfigT], object] | None = None

source

Describes what the graph bakes in of the draft, for an export’s sidecar.

method

method: SpeculativeMethod

source

optional_weights

optional_weights: frozenset[str] = frozenset({})

source

shared_modules

shared_modules: tuple[Literal['embed_tokens', 'lm_head'], ...] = ()

source

strict_weights

strict_weights: bool = True

source

Rejects unexpected owned checkpoint weights as well as missing ones.

Missing required weights are always rejected. Shared target modules and absent optional weights remain valid. Set this to False only when the draft intentionally accepts unused checkpoint weights; those keys are logged as warnings instead.

target_hidden_states()

target_hidden_states(draft_config)

source

Returns the hidden states the target returns for the draft, and the layers it captures for them.

Parameters:

• speculative (SpeculativeConfig)
• draft_config (object)

Return type:

TargetHiddenStates