For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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)

source

Bases: ArchConfigWithKVCache

Parameters:

• target (Llama3Config)
• draft (Llama3Config)
• speculative_config (SpeculativeConfig)
• enable_structured_output (bool)

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 default maximum sequence length for the model.

Subclasses should determine whether this value can be overridden by setting the --max-length (pipeline_config.model.max_length) flag.

Return type:

int

initialize()

classmethod initialize(pipeline_config, model_config=None)

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.

Return type:

Self

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

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 (LoRAInputs | None)
• hidden_states (Buffer | list[Buffer] | None)
• draft_tokens (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)

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, 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)
• 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

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit