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_mtp_gemma4

Gemma4 with MTP draft model for speculative decoding with unified graph compilation.

UnifiedMTPGemma4Inputs

class max.pipelines.architectures.unified_mtp_gemma4.UnifiedMTPGemma4Inputs(tokens, input_row_offsets, host_input_row_offsets, return_n_logits, data_parallel_splits, signal_buffers, batch_context_lengths, *, kv_cache_inputs=None, lora=None, vision_embeddings=<factory>, vision_scatter_indices=<factory>, 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 UnifiedMTPGemma4 model.

The spec-decode fields and trailing buffer packing come from UnifiedSpecDecodeInputs; the fields below plus the KV cache form this distributed MTP graph’s prefix. The graph binds the per-row in_thinking_phase flag and, when structured output is enabled, the constrained-decoding bitmask triple.

Parameters:

• tokens (Buffer)
• input_row_offsets (Buffer)
• host_input_row_offsets (Buffer)
• return_n_logits (Buffer)
• data_parallel_splits (Buffer)
• signal_buffers (list[Buffer])
• batch_context_lengths (list[Buffer])
• kv_cache_inputs (KVCacheInputsInterface[Buffer, Buffer] | None)
• lora (LoRAInputs | None)
• vision_embeddings (list[Buffer])
• vision_scatter_indices (list[Buffer])
• 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)

batch_context_lengths

batch_context_lengths: list[Buffer]

source

buffers

property buffers: tuple[Buffer, ...]

source

Returns positional Buffer inputs for model ABI calls.

data_parallel_splits

data_parallel_splits: Buffer

source

host_input_row_offsets

host_input_row_offsets: Buffer

source

input_row_offsets

input_row_offsets: Buffer

source

return_n_logits

return_n_logits: Buffer

source

signal_buffers

signal_buffers: list[Buffer]

source

tokens

tokens: Buffer

source

UnifiedMTPGemma4Model

class max.pipelines.architectures.unified_mtp_gemma4.UnifiedMTPGemma4Model(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, AlwaysSignalBuffersMixin, MultiGraphPipelineModelWithKVCache[Gemma4Context]

Gemma4 with MTP: merge + target + rejection + shift in one 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 UnifiedMTPGemma4BatchProcessor

calculate_max_seq_len()

classmethod calculate_max_seq_len(pipeline_config, huggingface_config)

source

Calculates the optimal max sequence length for the model.

Default implementation delegates to model_config_cls. Override when pipeline-model semantics differ from the config (for example, bounding max_length where the config is permissive).

Parameters:

• pipeline_config (PipelineConfig) – Configuration for the pipeline.
• huggingface_config (AutoConfig) – Hugging Face model configuration.

Returns:

The maximum sequence length to use.

Return type:

int

empty_vision_embeddings()

empty_vision_embeddings(devices)

source

Per-device zero-row image embeddings for cached / text-only batches.

Cached: hit on every text-only / decode step, so it must not allocate per call.

Parameters:

devices (list[Device])

Return type:

list[Buffer]

execute()

execute(model_inputs)

source

Execute and return all 3 graph outputs for speculative decoding.

Images appear only during prefill (draft_tokens is [batch, 0]); decode steps carry the empty vision-merge inputs, so the graph’s scatter is a no-op there.

Parameters:

model_inputs (ModelInputs)

Return type:

UnifiedEagleOutputs

model

model: Model

source

The compiled unified MTP graph (target + draft + rejection). Re-executed each step; device graph capture is disabled for this arch (see arch.py: the eager prefill vision encoder produces variable-shape embeddings).

model_config_cls

model_config_cls

source

alias of Gemma4ForConditionalGenerationConfig

pack_vision_inputs()

pack_vision_inputs(selection, devices)

source

Pack the pipeline-selected uncached image pixels to device.

Runs in the pipeline’s prep-ahead window (pinned host-to-device copy) so it overlaps the prior batch, delegating to the shared pack_uncached_images().

Parameters:

• selection (Sequence[tuple[Gemma4Context, Sequence[ImageMetadata]]])
• devices (list[Device])

Return type:

VisionRawInputs | None

vision_execute()

vision_execute(selection, devices, packed)

source

Run the vision encoder on the pixels pack_vision_inputs packed.

Returns embeddings only; the cache derives per-image counts from its selection. When there were no packable patches (packed is None), returns an empty result so the cache assembles from resident entries.

Parameters:

• selection (Sequence[tuple[Gemma4Context, Sequence[ImageMetadata]]])
• devices (list[Device])
• packed (VisionRawInputs | None)

Return type:

VisionEncodeResult

vision_model

vision_model: Model | None

source

The compiled vision encoder graph, or None for text-only checkpoints. Runs eagerly during prefill (outside the captured graph).

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