For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python module

max.pipelines.architectures.diffusion_gemma

DiffusionGemma block-diffusion architecture.

DiffusionGemmaForBlockDiffusionConfig

class max.pipelines.architectures.diffusion_gemma.DiffusionGemmaForBlockDiffusionConfig(*, devices, dtype, unquantized_dtype=bfloat16, kv_params, image_token_index, video_token_index=262144, audio_token_index=None, text_config, vision_config, tie_word_embeddings=False, quantization_encoding=None, layer_ids=None, canvas_length=256, boi_token_id=255999, eoi_token_id=258882)

source

Bases: Gemma4ForConditionalGenerationConfig

Top-level MAX config for DiffusionGemma block diffusion.

Extends the Gemma4 multimodal config with the block-diffusion canvas geometry. KV cache parameters, per-layer-type attention geometry, MoE sizing, and the vision tower config are all inherited from the donor, which reads them through the view installed by initialize_from_config.

Parameters:

• devices (list[DeviceRef])
• dtype (DType)
• unquantized_dtype (DType)
• kv_params (MultiKVCacheParams)
• image_token_index (int)
• video_token_index (int)
• audio_token_index (int | None)
• text_config (Gemma4TextConfig)
• vision_config (Gemma4VisionConfig | None)
• tie_word_embeddings (bool)
• quantization_encoding (SupportedEncoding | None)
• layer_ids (range | None)
• canvas_length (int)
• boi_token_id (int)
• eoi_token_id (int)

DEFAULT_ENCODING

DEFAULT_ENCODING: ClassVar[SupportedEncoding] = 'float4_e2m1fnx2'

source

SUPPORTED_ENCODINGS

SUPPORTED_ENCODINGS: ClassVar[set[SupportedEncoding]] = {'bfloat16', 'float4_e2m1fnx2'}

source

boi_token_id

boi_token_id: int = 255999

source

Begin-of-image token id wrapping image prompts.

canvas_length

canvas_length: int = 256

source

Number of tokens denoised per block-diffusion canvas.

eoi_token_id

eoi_token_id: int = 258882

source

End-of-image token id wrapping image prompts.

finalize()

finalize(huggingface_config, state_dict, return_logits)

source

Finalize with state_dict-dependent fields.

Parses quantization config from the weights and finalizes the text sub-config.

Parameters:

• huggingface_config (AutoConfig) – HuggingFace model configuration.
• state_dict (dict[str, WeightData]) – Model weights dictionary.
• return_logits (ReturnLogits) – Return logits configuration.

Return type:

None

initialize_from_config()

classmethod initialize_from_config(pipeline_config, huggingface_config, *, max_seq_len)

source

Initializes from pipeline and HuggingFace configs.

Fields that depend on the state_dict should be set via finalize().

Parameters:

• pipeline_config (PipelineConfig) – The MAX Engine pipeline configuration.
• huggingface_config (AutoConfig) – Top-level HuggingFace model configuration.
• max_seq_len (int)

Returns:

A config instance ready for finalization.

Return type:

Self

DiffusionGemmaForBlockDiffusionModel

class max.pipelines.architectures.diffusion_gemma.DiffusionGemmaForBlockDiffusionModel(pipeline_config, session, devices, kv_cache_config, weights, *, memory_plan, adapter=None, return_logits=ReturnLogits.LAST_TOKEN, max_batch_size=1)

source

Bases: Gemma3_MultiModalModel

Compiles the vision/encoder/decoder graphs for block diffusion.

Differences from the donor pipeline model:

• load_model uses this port’s weight converters (decoder-canonical checkpoint layout) and compiles an extra decoder graph.
• execute_decoder_step runs one denoise step: canvas K/V are written into the cache slots after each request’s committed length, so calling it repeatedly without advancing context lengths overwrites the same slots (read-only cache semantics from the encoder’s perspective).

Parameters:

• pipeline_config (PipelineConfig)
• session (InferenceSession)
• devices (list[Device])
• kv_cache_config (KVCacheConfig)
• weights (Weights)
• memory_plan (MemoryPlan)
• adapter (WeightsAdapter | None)
• return_logits (ReturnLogits)
• max_batch_size (int)

compute_log_probabilities()

compute_log_probabilities(session, model_inputs, model_outputs, next_tokens, batch_top_n, batch_echo)

source

Optional method that can be overridden to compute log probabilities.

Parameters:

• session (InferenceSession) – Inference session to compute log probabilities within.
• model_inputs (ModelInputs) – Inputs to the model returned by prepare_initial_token_inputs().
• model_outputs (ModelOutputs) – Outputs returned by execute().
• next_tokens (Buffer) – Sampled tokens. Should have shape=[batch size]
• batch_top_n (list[int]) – 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 (list[bool]) – Whether to include input tokens in the returned log probabilities.

Returns:

List of log probabilities.

Return type:

list[LogProbabilities | None]

decoder_model

decoder_model: Model

source

execute_decoder_step()

execute_decoder_step(canvas_tokens, input_row_offsets, sc_logits, sc_enabled, temperature, kv_cache_inputs)

source

Runs one denoise step.

Returns (sc_logits_out, argmax, topk_probs, topk_idx, entropy); sc_logits_out is device-resident bf16 and feeds the next step’s sc_logits; the [N, 64] top-k pair feeds host-side categorical sampling.

Parameters:

• canvas_tokens (Buffer)
• input_row_offsets (Buffer)
• sc_logits (Buffer)
• sc_enabled (Buffer)
• temperature (Buffer)
• kv_cache_inputs (tuple[KVCacheInputsPerDevice[Buffer, Buffer], ...])

Return type:

tuple[Buffer, Buffer, Buffer, Buffer, Buffer]

load_graphs()

load_graphs(session)

source

Loads the graphs, then binds the EP comm initializer.

Parameters:

session (InferenceSession) – The session that compiles and initializes the graphs.

Returns:

The initialized models, keyed by role.

Return type:

dict[str, Model]