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

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.

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, 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)
• 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 proposer’s job.

Parameters:

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

Return type:

TensorValue

has_own_embed_tokens

has_own_embed_tokens = True

source

The draft embeds through its own copy, never the target’s scaled one.

has_own_lm_head

has_own_lm_head = True

source

The draft scores its own pruned vocabulary.

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, sampler=None)

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, that is, after the caller’s anchor-slot drop).
• anchor_tokens (TensorValue) – Anchor/bonus token ids [batch], target vocab.
• sampler (Callable[[TensorValue, int], TensorValue] | None) – Draws step k’s [batch] DRAFT-vocab id from its [batch, draft_vocab] corrected logits in place of the argmax.

Returns:

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

Return type:

TensorValue