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

DFlash draft model for Llama3-family targets.

The draft is a Qwen3-style transformer (per-head Q/K RMSNorm, non-causal attention) that fuses concatenated target hidden states into its KV cache via AttentionWithRope.materialize_kv_from_hidden() and runs a single non-causal block forward over [verified_id, MASK, MASK, …] per iteration.

DFlashLlama3

class max.pipelines.architectures.dflash_llama3.DFlashLlama3(config, *, num_context_features, layer_types=None)

source

Bases: Module

DFlash draft transformer for a Llama3 target.

Builds the draft stack.

Parameters:

• config (Llama3Config) – The draft’s own Llama3-shaped config.
• num_context_features (int) – Number of target hidden-state taps the fc projection consumes. Decoupled from the draft’s layer count: the Gemma4 DFlash drafter fuses six taps into five layers.
• layer_types (Sequence[str] | None) – Per-layer "sliding_attention" / "full_attention" selection from the draft checkpoint. None applies config.sliding_window to every layer.

ctx_at_draft_cache_length

ctx_at_draft_cache_length = False

source

forward_block()

forward_block(input_embeds, kv_collections, input_row_offsets, signal_buffers)

source

Runs the block, on the draft’s one device.

Parameters:

• input_embeds (Sequence[TensorValue])
• kv_collections (Sequence[KVCacheInputsPerDevice[TensorValue, BufferValue]])
• input_row_offsets (Sequence[TensorValue])
• signal_buffers (Sequence[BufferValue])

Return type:

list[TensorValue]

mask_embeddings()

mask_embeddings(target_mask)

source

Returns the target’s embedding of the mask token, as trained.

Parameters:

target_mask (Sequence[TensorValue])

Return type:

list[TensorValue]

materialize()

materialize(target_hidden, input_row_offsets, kv_collections, signal_buffers)

source

Writes the context’s K/V, on the draft’s one device.

Parameters:

• target_hidden (Sequence[TensorValue])
• input_row_offsets (Sequence[TensorValue])
• kv_collections (Sequence[KVCacheInputsPerDevice[TensorValue, BufferValue]])
• signal_buffers (Sequence[BufferValue])

Return type:

None

supports_zero_draft_rows

supports_zero_draft_rows = True

source

DFlashLlama3ArchConfig

class max.pipelines.architectures.dflash_llama3.DFlashLlama3ArchConfig(*, hidden_size, num_attention_heads, num_key_value_heads, num_hidden_layers, rope_theta, rope_scaling_params, max_seq_len, intermediate_size, interleaved_rope_weights, vocab_size, dtype, model_quantization_encoding, quantization_config, kv_params, return_logits=ReturnLogits.LAST_TOKEN, norm_method='rms_norm', norm_dtype=None, attention_bias=False, rms_norm_eps=None, tie_word_embeddings=False, stacked_mlp=False, stacked_qkv=False, attention_multiplier, embedding_multiplier, residual_multiplier, devices, clip_qkv, quant_config=None, longrope_scaling_params=None, logits_scaling=1.0, return_hidden_states=ReturnHiddenStates.NONE, target_layer_ids=None, use_subgraphs=True, data_parallel_degree=1, sliding_window=None, quantization_encoding=None)

source

Bases: Llama3Config

The DFlash draft’s Llama 3 config, which builds the draft’s module.

Parameters:

• hidden_size (int)
• num_attention_heads (int)
• num_key_value_heads (int)
• num_hidden_layers (int)
• rope_theta (float)
• rope_scaling_params (Llama3RopeScalingParams | None)
• max_seq_len (int)
• intermediate_size (int)
• interleaved_rope_weights (bool)
• vocab_size (int)
• dtype (DType)
• model_quantization_encoding (QuantizationEncoding | None)
• quantization_config (QuantizationConfig | None)
• kv_params (KVCacheParams)
• return_logits (ReturnLogits)
• norm_method (Literal['rms_norm', 'layer_norm'])
• norm_dtype (DType | None)
• attention_bias (bool)
• rms_norm_eps (float | None)
• tie_word_embeddings (bool)
• stacked_mlp (bool)
• stacked_qkv (bool)
• attention_multiplier (float)
• embedding_multiplier (float)
• residual_multiplier (float)
• devices (list[DeviceRef])
• clip_qkv (float | None)
• quant_config (QuantConfig | None)
• longrope_scaling_params (LongRoPEScalingParams | None)
• logits_scaling (float)
• return_hidden_states (ReturnHiddenStates)
• target_layer_ids (list[int] | None)
• use_subgraphs (bool)
• data_parallel_degree (int)
• sliding_window (int | None)
• quantization_encoding (SupportedEncoding | None)

build_module()

classmethod build_module(context)

source

Builds the DFlash draft from the draft checkpoint context names.

The draft computes in its checkpoint’s own precision and has no embedding or head; it embeds and projects through its target’s.

Raises:

ValueError – If the method isn’t dflash.

Parameters:

context (ModelLoadContext)

Return type:

tuple[DFlashLlama3, DFlashLlama3Config]