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 module

max.pipelines.architectures.dspark_draft

Model-agnostic DSpark draft: block nn.Module + registry placeholder.

Speculators-format DSpark drafters declare the generic HuggingFace architectures ["DSparkDraftModel"] shared across target model lines, and the draft nn.Module is likewise target-agnostic: each unified target architecture imports it and drives it inside its compiled graph. The format helpers (config parsing, d2t vocabulary map, checkpoint weight handling) live in the speculators_common package.

DSparkDraftPlaceholderModel

class max.pipelines.architectures.dspark_draft.DSparkDraftPlaceholderModel(pipeline_config, session, devices, kv_cache_config, weights, adapter, return_logits, return_hidden_states=ReturnHiddenStates.NONE, max_batch_size=1)

source

Bases: PipelineModel[TextContext]

Placeholder for the draft-only registration; never constructed.

Parameters:

• pipeline_config (PipelineConfig)
• session (InferenceSession)
• devices (list[Device])
• kv_cache_config (KVCacheConfig)
• weights (Weights)
• adapter (WeightsAdapter | None)
• return_logits (ReturnLogits)
• return_hidden_states (ReturnHiddenStates)
• max_batch_size (int)

execute()

execute(model_inputs)

source

Executes the graph with the given inputs.

Parameters:

model_inputs (ModelInputs) – The model inputs to execute, containing tensors and any other required data for model execution.

Returns:

ModelOutputs containing the pipeline’s output tensors.

Return type:

ModelOutputs

This is an abstract method that must be implemented by concrete PipelineModels to define their specific execution logic.

DSparkSpeculatorsDraft

class max.pipelines.architectures.dspark_draft.DSparkSpeculatorsDraft(*, hidden_size, num_hidden_layers, num_attention_heads, num_key_value_heads, head_dim, intermediate_size, rms_norm_eps, rope_theta, sliding_window, layer_causal, vocab_size, draft_vocab_size, markov_rank, block_size, sample_from_anchor, mask_token_id, num_context_features, max_seq_len, kv_params, devices, dtype)

source

Bases: Module

Speculators-format DSpark draft transformer.

Structure per block slot: input_layernorm -> attn -> +residual -> post_attention_layernorm -> silu MLP -> +residual (2-norm sandwich), final norm after the stack. Attention is GQA with separate K/V, per-head-dim q/k RMSNorm before full-rotary neox RoPE, and a per-layer causal/sliding mask derived from layer_causal/sliding_window.

Parameters:

• hidden_size (int)
• num_hidden_layers (int)
• num_attention_heads (int)
• num_key_value_heads (int)
• head_dim (int)
• intermediate_size (int)
• rms_norm_eps (float)
• rope_theta (float)
• sliding_window (int | None)
• layer_causal (Sequence[bool])
• vocab_size (int)
• draft_vocab_size (int)
• markov_rank (int)
• block_size (int)
• sample_from_anchor (bool)
• mask_token_id (int)
• num_context_features (int)
• max_seq_len (int)
• kv_params (KVCacheParams)
• devices (list[DeviceRef])
• dtype (DType)

forward_block()

forward_block(input_embeds, kv_collection, input_row_offsets)

source

Runs the draft stack over the block tokens.

Returns hidden states for ALL block slots (anchor slot 0 included); dropping slot 0 when sample_from_anchor is false is the unified graph’s job.

Parameters:

• input_embeds (TensorValue)
• kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue])
• input_row_offsets (TensorValue)

Return type:

TensorValue

materialize_kv()

materialize_kv(ctx_hidden, input_row_offsets, kv_collection)

source

Projects context hidden states to per-layer K/V and writes the paged KV cache.

K gets k_norm + RoPE at the context positions; V is the plain v_proj output (un-normed, un-roped), matching the reference precompute_and_store_context_kv.

Parameters:

• ctx_hidden (TensorValue)
• input_row_offsets (TensorValue)
• kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue])

Return type:

None

project_target_hidden()

project_target_hidden(target_hs_concat)

source

Projects concatenated target taps to the context-KV input, hidden_norm(fc(concat(taps))).

Parameters:

target_hs_concat (TensorValue)

Return type:

TensorValue

sample_draft_tokens()

sample_draft_tokens(base_logits, anchor_tokens)

source

Greedily samples the markov-corrected draft tokens.

Unrolled at graph-build time over the static slot axis. Per slot: bias = markov_w2(markov_w1[prev]) with prev a TARGET-vocab id (the anchor seeds the first slot), draft = argmax(base + bias) over the draft vocab, then the in-chain d2t gather maps draft to the target id that both feeds the next slot’s bias and is emitted.

Parameters:

• base_logits (TensorValue) – Base draft logits [batch, num_slots, draft_vocab] with a static slot axis (slot 0 = the first drafted position, i.e. after the caller’s anchor-slot drop).
• anchor_tokens (TensorValue) – Anchor/bonus token ids [batch], target vocab.

Returns:

The drafted TARGET-vocab token ids [batch, num_slots] (int64).

Return type:

TensorValue