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

ConvStateCache

class max.pipelines.architectures.lfm2.ConvStateCache(num_conv_layers, hidden_size, conv_kernel, dtype, max_slots, device)

source

Bases: object

Parameters:

• num_conv_layers (int)
• hidden_size (int)
• conv_kernel (int)
• dtype (DType)
• max_slots (int)
• device (Device)

claim()

claim(request_id)

source

Parameters:

request_id (RequestID)

Return type:

None

get_states()

get_states(request_ids)

source

Return one [N, hidden, kernel] buffer per conv layer.

For N == 1 this is zero-copy (returns the slot’s buffer directly). For N > 1 per-slot buffers are concatenated along the leading batch dim via numpy round-trip — the conv state is small (hidden * kernel per slot), so this is acceptable.

Parameters:

request_ids (list[RequestID])

Return type:

list[Buffer]

release()

release(request_id)

source

Parameters:

request_id (RequestID)

Return type:

None

update_states()

update_states(request_ids, new_states)

source

Store updated per-layer states back into their request slots.

For N == 1 the buffer reference is stored directly. For N > 1 the leading batch dim is split and each slice is copied into the matching slot.

Parameters:

• request_ids (list[RequestID])
• new_states (list[Buffer])

Return type:

None

LFM2Config

class max.pipelines.architectures.lfm2.LFM2Config(*, 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, lora_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, layer_types=<factory>, conv_L_cache=3, conv_bias=False, norm_eps=1e-05)

source

Bases: Llama3Config

Model configuration for LFM2 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)
• 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)
• lora_config (LoRAConfig | 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)
• layer_types (list[str])
• conv_L_cache (int)
• conv_bias (bool)
• norm_eps (float)

conv_L_cache

conv_L_cache: int = 3

source

conv_bias

conv_bias: bool = False

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

initialize_from_config()

classmethod initialize_from_config(pipeline_config, huggingface_config, model_config=None)

source

Parameters:

• pipeline_config (Any)
• huggingface_config (AutoConfig)
• model_config (Any)

Return type:

LFM2Config

layer_types

layer_types: list[str]

source

norm_eps

norm_eps: float = 1e-05

source

LFM2Inputs

class max.pipelines.architectures.lfm2.LFM2Inputs(tokens: 'Buffer', input_row_offsets: 'Buffer', signal_buffers: 'list[Buffer]', return_n_logits: 'Buffer', data_parallel_splits: 'Buffer | Sequence[Sequence[int]] | None' = None, conv_states: 'list[Buffer]' = <factory>, request_ids: 'list[RequestID]' = <factory>, *, kv_cache_inputs: 'KVCacheInputs[Buffer, Buffer] | None' = None, lora: 'LoRAInputs | None' = None, hidden_states: 'Buffer | list[Buffer] | None' = None)

source

Bases: Llama3Inputs

Parameters:

• tokens (Buffer)
• input_row_offsets (Buffer)
• signal_buffers (list[Buffer])
• return_n_logits (Buffer)
• data_parallel_splits (Buffer | Sequence[Sequence[int]] | None)
• conv_states (list[Buffer])
• request_ids (list[RequestID])
• kv_cache_inputs (KVCacheInputs[Buffer, Buffer] | None)
• lora (LoRAInputs | None)
• hidden_states (Buffer | list[Buffer] | None)

buffers

property buffers: tuple[Buffer, ...]

source

Returns positional Buffer inputs for model ABI calls.

conv_states

conv_states: list[Buffer]

source

request_ids

request_ids: list[RequestID]

source

LFM2Model

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

source

Bases: LlamaModelBase

LFM2 hybrid (full-attention + conv) pipeline model.

Parameters:

• pipeline_config (PipelineConfig) – The configuration for this pipeline.
• session (InferenceSession) – The container for the runtime for this model.
• 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

Whether to use attention bias.

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.

model_config_cls

model_config_cls

source

alias of LFM2Config

norm_method

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

source

Normalization layer.

prepare_initial_token_inputs()

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

source

Prepare the inputs for the first pass in multistep execution.

Parameters:

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

Return type:

LFM2Inputs

release()

release(request_id)

source

Parameters:

request_id (RequestID)

Return type:

None

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