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_eagle_llama3

EAGLE speculative decoding draft model for Llama 3 with unified graph compilation.

PersistentInputBuffers

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

source

Bases: object

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

Parameters:

• tokens (Buffer)
• input_row_offsets (Buffer)

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

UnifiedEagleLlama3Config

class max.pipelines.architectures.unified_eagle_llama3.UnifiedEagleLlama3Config(*, target: 'Llama3Config', draft: 'Llama3Config', speculative_config: 'SpeculativeConfig', enable_structured_output: 'bool' = False, quantization_encoding: 'SupportedEncoding | None' = None)

source

Bases: ArchConfigWithKVCache

Parameters:

• target (Llama3Config)
• draft (Llama3Config)
• speculative_config (SpeculativeConfig)
• enable_structured_output (bool)
• quantization_encoding (SupportedEncoding | None)

DEFAULT_ENCODING

DEFAULT_ENCODING: ClassVar[SupportedEncoding] = 'bfloat16'

source

SUPPORTED_ENCODINGS

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

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

enable_structured_output

enable_structured_output: bool = False

source

When True, the graph accepts a bitmask input for grammar-constrained decoding.

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

quantization_encoding

quantization_encoding: SupportedEncoding | None = None

source

speculative_config

speculative_config: SpeculativeConfig

source

target

target: Llama3Config

source

UnifiedEagleLlama3Inputs

class max.pipelines.architectures.unified_eagle_llama3.UnifiedEagleLlama3Inputs(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 EAGLE Llama3 model.

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 unified_eagle_llama3 graph does not bind in_thinking_phase.

Parameters:

• tokens (Buffer)
• input_row_offsets (Buffer)
• return_n_logits (Buffer)
• kv_cache_inputs (KVCacheInputsInterface[Buffer, Buffer] | None)
• lora_buffers (tuple[Buffer, ...])
• vision_embeddings (list[Buffer])
• vision_scatter_indices (list[Buffer])
• hidden_states (Buffer | list[Buffer] | None)
• draft_tokens (Buffer | None)
• draft_probs_full (Buffer | None)
• seed (Buffer | None)
• temperature (Buffer | None)
• top_k (Buffer | None)
• max_k (Buffer | None)
• top_p (Buffer | None)
• min_top_p (Buffer | None)
• in_thinking_phase (Buffer | None)
• pinned_bitmask (Buffer | None)
• wait_payload (Buffer | None)
• device_bitmask_scratch (Buffer | None)
• structured_output (bool)
• sampled_draft_proposal (bool)

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

UnifiedEagleLlama3Model

class max.pipelines.architectures.unified_eagle_llama3.UnifiedEagleLlama3Model(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 EAGLE Llama3: target + draft in one compiled graph.

Parameters:

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

batch_processor_cls

batch_processor_cls

source

alias of UnifiedEagleLlama3BatchProcessor

get_kv_params()

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

source

Target KV params for memory planning; load_model upgrades to multi-KV.

Parameters:

• huggingface_config (PreTrainedConfig)
• pipeline_config (PipelineConfig)
• devices (list[DeviceRef])
• kv_cache_config (KVCacheConfig)
• cache_dtype (DType)

Return type:

KVCacheParams

model

model: Model

source

model_config_cls

model_config_cls

source

alias of UnifiedEagleLlama3Config