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

Llama 3 transformer architecture for text generation.

`Llama3Config`

class max.pipelines.architectures.llama3_modulev3.Llama3Config(*, 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, kv_params, return_logits=ReturnLogits.LAST_TOKEN, norm_method='rms_norm', 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=None, norm_elementwise_affine=True, longrope_scaling_params=None, logits_scaling=1.0, return_hidden_states=ReturnHiddenStates.NONE)

source

Bases: ArchConfigWithStoredKVParams, ArchConfigWithKVCache

Model configuration for Llama3 graph construction/execution.

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)
kv_params (KVCacheParams)
return_logits (ReturnLogits)
norm_method (Literal['rms_norm', 'layer_norm'])
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)
norm_elementwise_affine (bool)
longrope_scaling_params (LongRoPEScalingParams | None)
logits_scaling (float)
return_hidden_states (ReturnHiddenStates)

`attention_bias`

attention_bias: bool = False

source

`attention_multiplier`

attention_multiplier: float

source

`calculate_attention_multiplier()`

static calculate_attention_multiplier(huggingface_config)

source

Parameters:: huggingface_config (AutoConfig)
Return type:: float

`clip_qkv`

clip_qkv: float | None = None

source

`construct_kv_params()`

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

source

Grouped-attention KV with EAGLE draft-token count when speculative is on.

Parameters:

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

Return type:

KVCacheParams

`devices`

devices: list[DeviceRef]

source

`dtype`

dtype: DType

source

`embedding_multiplier`

embedding_multiplier: float

source

`finalize()`

finalize(huggingface_config, state_dict, return_logits, return_hidden_states=ReturnHiddenStates.NONE, norm_method='rms_norm', attention_bias=False)

source

Define parameters that can’t be determined just from the pipeline config.

Parameters:

huggingface_config (AutoConfig)
state_dict (dict[str, WeightData])
return_logits (ReturnLogits)
return_hidden_states (ReturnHiddenStates)
norm_method (Literal['rms_norm', 'layer_norm'])
attention_bias (bool)

Return type:

None

`get_head_dim_from_config()`

static get_head_dim_from_config(config)

source

Parameters:: config (Llama3Config)
Return type:: int

`get_num_layers()`

static get_num_layers(huggingface_config)

source

Layer count for the decoder stack (override when HF uses a different field).

Parameters:: huggingface_config (AutoConfig)
Return type:: int

`hidden_size`

hidden_size: int

source

`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

`interleaved_rope_weights`

interleaved_rope_weights: bool

source

`intermediate_size`

intermediate_size: int

source

`kv_params`

kv_params: KVCacheParams

source

`logits_scaling`

logits_scaling: float = 1.0

source

`longrope_scaling_params`

longrope_scaling_params: LongRoPEScalingParams | None = None

source

`max_seq_len`

max_seq_len: int

source

`norm_elementwise_affine`

norm_elementwise_affine: bool = True

source

`norm_method`

norm_method: Literal['rms_norm', 'layer_norm'] = 'rms_norm'

source

`num_attention_heads`

num_attention_heads: int

source

`num_hidden_layers`

num_hidden_layers: int

source

`num_key_value_heads`

num_key_value_heads: int

source

`residual_multiplier`

residual_multiplier: float

source

`return_hidden_states`

return_hidden_states: ReturnHiddenStates = 'none'

source

`return_logits`

return_logits: ReturnLogits = 'last_token'

source

`rms_norm_eps`

rms_norm_eps: float | None = None

source

`rope_scaling_params`

rope_scaling_params: Llama3RopeScalingParams | None

source

`rope_theta`

rope_theta: float

source

`stacked_mlp`

stacked_mlp: bool = False

source

`stacked_qkv`

stacked_qkv: bool = False

source

`tie_word_embeddings`

tie_word_embeddings: bool = False

source

`vocab_size`

vocab_size: int

source

`Llama3Inputs`

class max.pipelines.architectures.llama3_modulev3.Llama3Inputs(tokens, input_row_offsets, return_n_logits, *, kv_cache_inputs=None, lora_ids=None, lora_ranks=None, hidden_states=None)

source

Bases: ModelInputs

A class representing inputs for the Llama3 model.

Parameters:

tokens (Buffer)
input_row_offsets (Buffer)
return_n_logits (Buffer)
kv_cache_inputs (KVCacheInputs[Buffer, Buffer] | None)
lora_ids (Buffer | None)
lora_ranks (Buffer | None)
hidden_states (Buffer | list[Buffer] | None)

`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

`Llama3Model`

class max.pipelines.architectures.llama3_modulev3.Llama3Model(pipeline_config, session, devices, kv_cache_config, weights, adapter=None, return_logits=ReturnLogits.LAST_TOKEN, return_hidden_states=ReturnHiddenStates.NONE)

source

Bases: LogProbabilitiesMixin, PipelineModelWithKVCache[TextContext]

Llama3 pipeline model using the ModuleV3 API.

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)

`attention_bias`

attention_bias: bool = False

source

`config_class`

config_class

source

alias of Llama3Config

`execute()`

execute(model_inputs)

source

Executes the graph with the given inputs.

Parameters:: model_inputs (ModelInputs) – The model inputs to execute, containing tensors and any other required data for model execution.
Returns:: ModelOutputs containing the pipeline’s output tensors.
Return type:: ModelOutputs

This is an abstract method that must be implemented by concrete PipelineModels to define their specific execution logic.

`load_model()`

load_model()

source

Return type:: Callable[[…], Any]

`model_config_cls`

model_config_cls

source

alias of Llama3Config

`norm_method`

norm_method: Literal['rms_norm'] | Literal['layer_norm'] = 'rms_norm'

source

`prepare_initial_token_inputs()`

prepare_initial_token_inputs(replica_batches, kv_cache_inputs=None, return_n_logits=1)

source

Prepares the initial inputs to be passed to execute().

The inputs and functionality can vary per model. For example, model inputs could include encoded tensors, unique IDs per tensor when using a KV cache manager, and kv_cache_inputs (or None if the model does not use KV cache). This method typically batches encoded tensors, claims a KV cache slot if needed, and returns the inputs and caches.

Parameters:

replica_batches (Sequence[Sequence[TextContext]])
kv_cache_inputs (KVCacheInputs[Buffer, Buffer] | None)
return_n_logits (int)

Return type:

ModelInputs

Llama3Config​

attention_bias​

attention_multiplier​

calculate_attention_multiplier()​

clip_qkv​

construct_kv_params()​

devices​

dtype​

embedding_multiplier​

finalize()​

get_head_dim_from_config()​

get_num_layers()​

hidden_size​

initialize()​

interleaved_rope_weights​

intermediate_size​

kv_params​

logits_scaling​

longrope_scaling_params​

max_seq_len​

norm_elementwise_affine​

norm_method​

num_attention_heads​

num_hidden_layers​

num_key_value_heads​

residual_multiplier​

return_hidden_states​

return_logits​

rms_norm_eps​

rope_scaling_params​

rope_theta​

stacked_mlp​

stacked_qkv​

tie_word_embeddings​

vocab_size​

Llama3Inputs​

buffers​

input_row_offsets​

return_n_logits​

tokens​

Llama3Model​

attention_bias​

config_class​

execute()​

load_model()​

model_config_cls​

norm_method​

prepare_initial_token_inputs()​

`Llama3Config`

`attention_bias`

`attention_multiplier`

`calculate_attention_multiplier()`

`clip_qkv`

`construct_kv_params()`

`devices`

`dtype`

`embedding_multiplier`

`finalize()`

`get_head_dim_from_config()`

`get_num_layers()`

`hidden_size`

`initialize()`

`interleaved_rope_weights`

`intermediate_size`

`kv_params`

`logits_scaling`

`longrope_scaling_params`

`max_seq_len`

`norm_elementwise_affine`

`norm_method`

`num_attention_heads`

`num_hidden_layers`

`num_key_value_heads`

`residual_multiplier`

`return_hidden_states`

`return_logits`

`rms_norm_eps`

`rope_scaling_params`

`rope_theta`

`stacked_mlp`

`stacked_qkv`

`tie_word_embeddings`

`vocab_size`

`Llama3Inputs`

`buffers`

`input_row_offsets`

`return_n_logits`

`tokens`

`Llama3Model`

`attention_bias`

`config_class`

`execute()`

`load_model()`

`model_config_cls`

`norm_method`

`prepare_initial_token_inputs()`