Python module

pipeline

Hugging Face Token Generation Pipeline.

KVCacheMixin

class max.pipelines.lib.pipeline.KVCacheMixin(*args, **kwargs)

estimate_kv_cache_size()

abstract classmethod estimate_kv_cache_size(pipeline_config: PipelineConfig, available_cache_memory: int, devices: list[Device], huggingface_config: AutoConfig, kv_cache_config: KVCacheConfig, cache_dtype: DType) → int

Estimates the size of the kv cache in bytes.

load_kv_manager()

load_kv_manager(session: InferenceSession, available_cache_memory: int | None) → KVCacheManager

Provided a PipelineConfig and InferenceSession, loads the KV manager.

• Parameters:

  • session – Inference session to compile and init the KV cache.
  • available_cache_memory – Amount of memory available to the KV cache, in bytes.

• Returns:

  one per input modality.

• Return type:

  Either a single KV cache manager or a tuple of KV cache managers

ModelInputs

class max.pipelines.lib.pipeline.ModelInputs

Base class for model inputs. Use this class to encapsulate inputs for your model. You may store any number of dataclass fields

The following example demonstrates how to create a custom inputs class for a model:

class ReplitInputs(ModelInputs):
    tokens: Tensor
    input_row_offsets: Tensor

    def __init__(self, tokens: Tensor, input_row_offsets: Tensor):
        self.tokens = tokens
        self.input_row_offsets = input_row_offsets

tokens = Tensor.zeros((1, 2, 3), DType.int64)
input_row_offsets = Tensor.zeros((1, 1, 1), DType.int64)

# Initialize inputs
inputs = ReplitInputs(tokens=tokens, input_row_offsets=input_row_offsets)

# Access tensors
list(inputs) == [tokens, input_row_offsets]  # Output: True

class ReplitInputs(ModelInputs):
    tokens: Tensor
    input_row_offsets: Tensor

    def __init__(self, tokens: Tensor, input_row_offsets: Tensor):
        self.tokens = tokens
        self.input_row_offsets = input_row_offsets

tokens = Tensor.zeros((1, 2, 3), DType.int64)
input_row_offsets = Tensor.zeros((1, 1, 1), DType.int64)

# Initialize inputs
inputs = ReplitInputs(tokens=tokens, input_row_offsets=input_row_offsets)

# Access tensors
list(inputs) == [tokens, input_row_offsets]  # Output: True

kv_cache_inputs

kv_cache_inputs*: KVCacheInputs | None* = None

ModelOutputs

class max.pipelines.lib.pipeline.ModelOutputs(logits: 'Tensor', next_token_logits: 'Tensor | None' = None, logit_offsets: 'Tensor | None' = None)

logit_offsets

logit_offsets*: Tensor | None* = None

Offsets to access variable length logits for each sequence.

logits

logits*: Tensor*

Logits for a variable number of tokens per sequence.

next_token_logits

next_token_logits*: Tensor | None* = None

Logits for just the next token.

PipelineModel

class max.pipelines.lib.pipeline.PipelineModel(pipeline_config: PipelineConfig, session: InferenceSession, huggingface_config: AutoConfig, encoding: SupportedEncoding, devices: list[Device], kv_cache_config: KVCacheConfig, weights: Weights, adapter: WeightsAdapter | None, return_logits: ReturnLogits)

A pipeline model with setup, input preparation and execution methods.

calculate_max_seq_len()

abstract classmethod calculate_max_seq_len(pipeline_config: PipelineConfig, huggingface_config: AutoConfig) → int

Calculate the optimal max sequence length for the model. Models are expected to implement this method.

The following example shows how to implement this method for a Mistral model:

class MistralModel(PipelineModel):
    @classmethod
    def calculate_max_seq_len(cls, pipeline_config, huggingface_config) -> int:
        try:
            return upper_bounded_default(
                upper_bound=huggingface_config.max_seq_len,
                default=pipeline_config.max_length,
            )
        except ValueError as e:
            msg = (
                "Unable to infer max_length for Mistral, the provided "
                f"max_length ({pipeline_config.max_length}) exceeds the "
                f"model's max_seq_len ({huggingface_config.max_seq_len})."
            )
            raise ValueError(msg) from e

class MistralModel(PipelineModel):
    @classmethod
    def calculate_max_seq_len(cls, pipeline_config, huggingface_config) -> int:
        try:
            return upper_bounded_default(
                upper_bound=huggingface_config.max_seq_len,
                default=pipeline_config.max_length,
            )
        except ValueError as e:
            msg = (
                "Unable to infer max_length for Mistral, the provided "
                f"max_length ({pipeline_config.max_length}) exceeds the "
                f"model's max_seq_len ({huggingface_config.max_seq_len})."
            )
            raise ValueError(msg) from e

• Parameters:

  • pipeline_config – Configuration for the pipeline.
  • huggingface_config – Hugging Face model configuration.

• Returns:

  The maximum sequence length to use.

• Return type:

  int

compute_log_probabilities()

compute_log_probabilities(model_inputs: ModelInputs, model_outputs: ModelOutputs, next_tokens: Tensor, batch_top_n: list[int], batch_echo: list[bool]) → list[max.pipelines.core.interfaces.response.LogProbabilities | None] | None

Optional method that can be overridden to compute log probabilities.

• Parameters:

  • model_inputs – Inputs to the model returned by prepare_*_token_inputs().
  • model_outputs – Outputs returned by execute().
  • next_tokens – Sampled tokens. Should have shape=[batch size]
  • batch_top_n – Number of top log probabilities to return per input in the batch. For any element where top_n == 0, the LogProbabilities is skipped.
  • batch_echo – Whether to include input tokens in the returned log probabilities.

• Returns:

  List of log probabilities.

dtype

property dtype*: DType*

estimate_weights_size()

classmethod estimate_weights_size(pipeline_config: PipelineConfig) → int

Calculates the estimated memory consumption of our model.

execute()

abstract execute(model_inputs: ModelInputs) → ModelOutputs

Executes the graph with the given inputs.

• Parameters:

  model_inputs – The model inputs to execute, containing tensors and any other required data for model execution.

• Returns:

  ModelOutputs containing the pipeline’s output tensors.

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

get_kv_params()

abstract classmethod get_kv_params(huggingface_config: AutoConfig, n_devices: int, kv_cache_config: KVCacheConfig, cache_dtype: DType) → KVCacheParams

Returns the KV cache params for the pipeline model.

get_num_layers()

abstract classmethod get_num_layers(huggingface_config: AutoConfig) → int

Returns the number of layers for the pipeline model.

infer_optimal_batch_size()

classmethod infer_optimal_batch_size(pipeline_config: PipelineConfig, available_cache_memory: int, huggingface_config: AutoConfig, devices: list[Device], kv_cache_config: KVCacheConfig, cache_dtype: DType) → int

Returns the estimated optimal batch size to run the model given current memory constraints.

prepare_initial_token_inputs()

abstract prepare_initial_token_inputs(context_batch: Sequence[T], kv_cache_inputs: KVCacheInputs | None = None, return_n_logits: int = 1) → ModelInputs

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

The inputs and functionality of this method can vary per model. For example, the model inputs could include:

• Encoded tensors
• A unique IDs for each tensor if this model uses a KV Cache manager.
• kv_cache_inputs: The kv cache inputs required for the model. This should be None if the model does not use KV Cache. This function would batch the encoded tensors, claim a slot in the kv cache if the ID hasn’t been seen before, and return the inputs and caches as a list of tensors.

prepare_next_token_inputs()

abstract prepare_next_token_inputs(next_tokens: Tensor, prev_model_inputs: ModelInputs) → ModelInputs

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

While prepare_initial_token_inputs is responsible for managing the initial inputs. This function is responsible for updating the inputs, for each step in a multi-step execution pattern.

TextGenerationPipeline

class max.pipelines.lib.pipeline.TextGenerationPipeline(pipeline_config: PipelineConfig, pipeline_model: type[PipelineModel], eos_token_id: int, weight_adapters: dict[WeightsFormat, WeightsAdapter])

Generalized token generator pipeline.

calculate_num_steps()

calculate_num_steps(num_steps: int, context: T) → int

next_token()

next_token(batch: dict[str, T], num_steps: int) → dict[str, max.pipelines.core.interfaces.response.TextGenerationResponse]

Provided a batch, process batch inputs, execute the graph for num_steps in a multi-step scenario, then decode the tokens holistically and return the list of decoded tokens.

prepare_batch()

prepare_batch(batch: list[T], num_steps: int) → tuple[max.pipelines.lib.pipeline.ModelInputs, int, Optional[torch.Tensor]]

release()

release(context: T) → None

Mark the context as complete, releasing the cache slot from the KV manager.

sample_logits()

sample_logits(logits: Tensor, prev_tokens: Tensor, logit_offsets: Tensor | None, bitmask: Tensor | None) → tuple[max._core.driver.Tensor, max._core.driver.Tensor]

upper_bounded_default()

max.pipelines.lib.pipeline.upper_bounded_default(upper_bound: int, default: int | None) → int

Given an upper bound and an optional default value, returns a final value that cannot exceed the upper bound.

• Parameters:

  • default – The default value to use, or None to use the upper bound.
  • upper_bound – The upper bound to use.

• Raises:

  ValueError – If the provided default value exceeds the upper bound.

• Returns:

  The final value.

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