IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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.unified_dflash_llama3

DFlash speculative decoding for Llama3 with unified graph compilation.

DflashDraftHFConfig​

class max.pipelines.architectures.unified_dflash_llama3.DflashDraftHFConfig(mask_token_id, target_layer_ids, block_size=None, num_target_layers=None)

source

Bases: object

Parsed DFlash fields from a draft HuggingFace config.

Parameters:

  • mask_token_id (int)
  • target_layer_ids (list[int])
  • block_size (int | None)
  • num_target_layers (int | None)

block_size​

block_size: int | None = None

source

draft_width()​

draft_width(speculative, *, warn=True)

source

Returns the draft width, which is block_size - 1.

The drafter only works at its trained block size, so a width that disagrees is replaced with a warning. A checkpoint with no block_size needs an explicit width.

Parameters:

Return type:

int

mask_token_id​

mask_token_id: int

source

num_target_layers​

num_target_layers: int | None = None

source

target_layer_ids​

target_layer_ids: list[int]

source

PersistentInputBuffers​

class max.pipelines.architectures.unified_dflash_llama3.PersistentInputBuffers(tokens, input_row_offsets)

source

Bases: object

Pinned-host buffers reused across unified spec-decode batch steps.

Parameters:

alloc()​

classmethod alloc(max_batch_size, max_batch_input_tokens, device)

source

Allocates persistent token and row-offset buffers for spec-decode batching.

Parameters:

  • max_batch_size (int)
  • max_batch_input_tokens (int)
  • device (Device)

Return type:

PersistentInputBuffers

input_row_offsets​

input_row_offsets: Buffer

source

tokens​

tokens: Buffer

source

UnifiedDflashLlama3Config​

class max.pipelines.architectures.unified_dflash_llama3.UnifiedDflashLlama3Config(*, target: 'Llama3Config', draft: 'Llama3Config', speculative_config: 'SpeculativeConfig', target_layer_ids: 'list[int]' = <factory>, mask_token_id: 'int' = 0, block_size: 'int' = 0, quantization_encoding: 'SupportedEncoding | None' = None, resolved_num_speculative_tokens: 'int | None' = None)

source

Bases: ArchConfigWithKVCache

Parameters:

DEFAULT_ENCODING​

DEFAULT_ENCODING: ClassVar[SupportedEncoding] = 'bfloat16'

source

SUPPORTED_ENCODINGS​

SUPPORTED_ENCODINGS: ClassVar[set[SupportedEncoding]] = {'bfloat16', 'float32'}

source

block_size​

block_size: int = 0

source

calculate_max_seq_len()​

classmethod calculate_max_seq_len(huggingface_config, model_config)

source

Returns the resolved maximum sequence length.

Bounds or defaults the user’s model_config.max_length with the model’s own limits. Construction runs this once and stores the result on model_config.max_length; memory planning may lower it further, but only on the memory plan.

Parameters:

  • huggingface_config (AutoConfig) – The HuggingFace config to read model bounds from.
  • model_config (MAXModelConfig) – The model config whose max_length carries the user’s setting.

Return type:

int

draft​

draft: Llama3Config

source

get_kv_params()​

get_kv_params()

source

KV cache parameters to use when running the model.

Return type:

KVCacheParamInterface

get_max_seq_len()​

get_max_seq_len()

source

Returns the effective maximum sequence length for the model.

For configs that store a deployment length, this is the value initialize received; for metadata-only configs it derives from the checkpoint.

Return type:

int

initialize()​

classmethod initialize(pipeline_config, model_config=None, *, max_seq_len)

source

Initialize the config from a PipelineConfig.

Parameters:

  • pipeline_config (PipelineConfig) – The pipeline configuration.
  • model_config (MAXModelConfig | None) – The model configuration to read from. When None (the default), pipeline_config.model is used. Pass an explicit config (e.g. pipeline_config.draft_model) to initialize the arch config for a different model.
  • max_seq_len (int) – The effective maximum sequence length to store on the config. The value is received, never derived here: the pipeline model passes the memory plan’s VRAM-clamped length, while memory planning (which runs before a plan exists) passes the construction-resolved model_config.max_length. Configs whose sequence length is pure model metadata (e.g. diffusion components) ignore it.

Return type:

Self

mask_token_id​

mask_token_id: int = 0

source

quantization_encoding​

quantization_encoding: SupportedEncoding | None = None

source

resolve_block_size()​

resolve_block_size(*, default=None)

source

Parameters:

default (int | None)

Return type:

int

resolved_num_speculative_tokens​

resolved_num_speculative_tokens: int | None = None

source

explicit value if set, else the trained width.

Type:

Per-step draft count

speculative_config​

speculative_config: SpeculativeConfig

source

target​

target: Llama3Config

source

target_layer_ids​

target_layer_ids: list[int]

source

validate_dflash_fields()​

validate_dflash_fields()

source

Strict validation run from UnifiedDflashLlama3Model.load_model once the DFlash-specific fields have been populated from the draft HF config — __post_init__ accepts the empty-placeholder config produced by initialize() so we can’t enforce these there.

Return type:

None

UnifiedDflashLlama3Inputs​

class max.pipelines.architectures.unified_dflash_llama3.UnifiedDflashLlama3Inputs(tokens, input_row_offsets, return_n_logits, *, kv_cache_inputs=None, lora_buffers=(), vision_embeddings=<factory>, vision_scatter_indices=<factory>, hidden_states=None, draft_tokens=None, draft_probs_full=None, seed=None, temperature=None, top_k=None, max_k=None, top_p=None, min_top_p=None, in_thinking_phase=None, pinned_bitmask=None, wait_payload=None, device_bitmask_scratch=None, structured_output=False, sampled_draft_proposal=False)

source

Bases: UnifiedSpecDecodeInputs

Inputs for the unified DFlash Llama3 graph.

The spec-decode fields and trailing buffer packing come from UnifiedSpecDecodeInputs; tokens / input_row_offsets / return_n_logits plus the KV cache form this single-device graph’s prefix. The DFlash graph does not bind in_thinking_phase.

Parameters:

buffers​

property buffers: tuple[Buffer, ...]

source

Returns positional Buffer inputs for model ABI calls.

input_row_offsets​

input_row_offsets: Buffer

source

return_n_logits​

return_n_logits: Buffer

source

tokens​

tokens: Buffer

source

UnifiedDflashLlama3Model​

class max.pipelines.architectures.unified_dflash_llama3.UnifiedDflashLlama3Model(pipeline_config, session, devices, kv_cache_config, weights, *, memory_plan, adapter=None, return_logits=ReturnLogits.LAST_TOKEN, return_hidden_states=ReturnHiddenStates.NONE, max_batch_size=1)

source

Bases: _UnifiedSpecDecodeModelMixin, GraphPipelineModelWithKVCache[TextContext]

Unified DFlash Llama3: target + draft in one compiled graph.

Parameters:

batch_processor_cls​

batch_processor_cls

source

alias of UnifiedDflashLlama3BatchProcessor

get_kv_params()​

classmethod get_kv_params(huggingface_config, pipeline_config, devices, kv_cache_config, cache_dtype)

source

Returns the KV cache params for the pipeline model.

Delegates to model_config_cls.construct_kv_params(...). Subclasses with custom KV behavior should override this method.

Parameters:

Return type:

KVCacheParams

model​

model: Model

source

model_config_cls​

model_config_cls

source

alias of UnifiedDflashLlama3Config

Models

Cookbook

Puzzles

Blog

Community

Contact