IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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.unified_dspark_gemma4

Unified DSpark speculative decoding for Gemma4.

DSparkGemma4​

class max.pipelines.architectures.unified_dspark_gemma4.DSparkGemma4(config, *, kv_params, devices, dtype)

source

Bases: Module

DSpark draft transformer for a Gemma4 target.

Parameters:

forward_block()​

forward_block(input_embeds, kv_collection, input_row_offsets)

source

Parameters:

Return type:

TensorValue

materialize_kv()​

materialize_kv(ctx_hidden, input_row_offsets, kv_collection)

source

Parameters:

Return type:

None

project_target_hidden()​

project_target_hidden(target_hs_concat)

source

Parameters:

target_hs_concat (TensorValue)

Return type:

TensorValue

DSparkGemma4Attention​

class max.pipelines.architectures.unified_dspark_gemma4.DSparkGemma4Attention(*, rope_global, rope_local, num_attention_heads, num_key_value_heads, num_global_key_value_heads, attention_k_eq_v, hidden_size, kv_params, global_head_dim, layer_idx, layer_idx_in_cache, is_sliding, dtype=float32, devices, linear_cls=<class 'max.nn.linear.Linear'>, has_bias=False, qk_norm_eps=1e-06, local_window_size=1024, quant_config=None, fused_qkv=False)

source

Bases: Gemma4Attention

Non-causal DSpark draft attention over [materialized ctx ; block].

Initializes the attention layer.

Parameters:

  • rope_global (RotaryEmbedding) – Rotary embedding used for global (non-sliding window) attention layers.
  • rope_local (RotaryEmbedding) – Rotary embedding used for sliding window attention layers.
  • num_attention_heads (int) – The number of attention heads.
  • num_key_value_heads (int) – The number of key/value heads.
  • hidden_size (int) – The dimension of the hidden states.
  • kv_params (KVCacheParams) – KV Cache Params, including the number of kv heads, the head dim, and data type.
  • layer_idx (int) – The layer number associated with this Attention block.
  • layer_idx_in_cache (int) – The 0-based index of this layer within its sub-cache (sliding or global).
  • is_sliding (bool) – Whether this layer uses sliding window attention.
  • dtype (DType) – DType of the attention inputs and weights.
  • devices (list[DeviceRef]) – Device to place the weights and run the computation. If multiple are provided, the first device is used. Use TensorParallelAttentionWithRope to use all devices during attention computation.
  • linear_cls (Callable[..., Linear]) – Linear class to use for the outputs dense layer.
  • has_bias (bool) – Whether to use an attention bias. Defaults to False.
  • qk_norm_eps (float) – Value to use for numerical stability. Defaults to 1e-6.
  • quant_config (QuantConfig | None) – Scaled quantization configuration. Defaults to None.
  • fused_qkv (bool) – When True, the qkv/qk projection uses a single stacked weight (StackedLinear(stacked=True)) loaded pre-fused from the checkpoint instead of concatenating per-projection weights in-graph (DISTINF-194). Defaults to False.
  • num_global_key_value_heads (int)
  • attention_k_eq_v (bool)
  • global_head_dim (int)
  • local_window_size (int)

mask_variant​

mask_variant: MHAMaskVariant | None = '2'

source

materialize_kv_from_hidden()​

materialize_kv_from_hidden(hidden, kv_collection, input_row_offsets)

source

Projects ctx hidden states to K/V and writes the paged KV cache.

Matches the reference ctx path: K gets k_norm + RoPE at the ctx positions (cache_length + row); V is the scale-free v_norm of the pre-norm K projection and is not roped.

Parameters:

Return type:

None

DSparkGemma4DecoderLayer​

class max.pipelines.architectures.unified_dspark_gemma4.DSparkGemma4DecoderLayer(attention, mlp, hidden_size, rms_norm_eps, dtype)

source

Bases: Module

Single-device Gemma4 decoder sandwich with a per-layer scalar.

Mirrors the reference Gemma4DSparkDecoderLayer: input_ln -> attn -> post_attn_ln -> +residual; pre_ffw_ln -> mlp -> post_ffw_ln -> +residual; then multiply by layer_scalar.

Parameters:

DSparkGemma4DraftConfig​

class max.pipelines.architectures.unified_dspark_gemma4.DSparkGemma4DraftConfig(hidden_size, intermediate_size, num_hidden_layers, num_attention_heads, num_key_value_heads, head_dim, rms_norm_eps, vocab_size, hidden_activation, final_logit_softcapping, rope_theta, partial_rotary_factor, max_seq_len, block_size, mask_token_id, target_layer_ids, markov_rank, markov_head_type)

source

Bases: object

DSpark drafter hyperparameters from the draft HF checkpoint config.

The checkpoint config (e.g. deepseek-ai/dspark_gemma4_12b_block7) is a gemma4_text-style config extended with the DSpark fields (block_size, mask_token_id, target_layer_ids, markov_rank, …). Only full-attention (k_eq_v) draft layers are supported, matching the reference drafter.

Parameters:

  • hidden_size (int)
  • intermediate_size (int)
  • num_hidden_layers (int)
  • num_attention_heads (int)
  • num_key_value_heads (int)
  • head_dim (int)
  • rms_norm_eps (float)
  • vocab_size (int)
  • hidden_activation (str)
  • final_logit_softcapping (float | None)
  • rope_theta (float)
  • partial_rotary_factor (float)
  • max_seq_len (int)
  • block_size (int)
  • mask_token_id (int)
  • target_layer_ids (tuple[int, ...])
  • markov_rank (int)
  • markov_head_type (str)

block_size​

block_size: int

source

final_logit_softcapping​

final_logit_softcapping: float | None

source

from_huggingface_config()​

classmethod from_huggingface_config(huggingface_config, *, max_seq_len=None)

source

Parses the DSpark draft checkpoint config (object or dict).

Parameters:

  • huggingface_config (Any)
  • max_seq_len (int | None)

Return type:

DSparkGemma4DraftConfig

head_dim​

head_dim: int

source

Global head dim (global_head_dim in the HF config).

hidden_activation​

hidden_activation: str

source

hidden_size​

hidden_size: int

source

intermediate_size​

intermediate_size: int

source

markov_head_type​

markov_head_type: str

source

markov_rank​

markov_rank: int

source

mask_token_id​

mask_token_id: int

source

max_seq_len​

max_seq_len: int

source

num_attention_heads​

num_attention_heads: int

source

num_context_features​

property num_context_features: int

source

num_hidden_layers​

num_hidden_layers: int

source

num_key_value_heads​

num_key_value_heads: int

source

Global (k_eq_v) key/value head count; every draft layer is global.

partial_rotary_factor​

partial_rotary_factor: float

source

rms_norm_eps​

rms_norm_eps: float

source

rope_theta​

rope_theta: float

source

target_layer_ids​

target_layer_ids: tuple[int, ...]

source

vocab_size​

vocab_size: int

source

DSparkMarkovHead​

class max.pipelines.architectures.unified_dspark_gemma4.DSparkMarkovHead(vocab_size, markov_rank, dtype, device)

source

Bases: Module

Vanilla markov head: bias = markov_w2(markov_w1[prev_token]).

Mirrors the reference VanillaMarkov (markov_head.py): a low-rank per-token logit bias conditioned only on the previous draft token. The hidden states are ignored by the vanilla head.

Parameters:

compute_step_bias()​

compute_step_bias(prev_tokens)

source

Returns the [batch, vocab] bias for [batch] prev tokens.

Parameters:

prev_tokens (TensorValue)

Return type:

TensorValue

UnifiedDSparkGemma4​

class max.pipelines.architectures.unified_dspark_gemma4.UnifiedDSparkGemma4(config)

source

Bases: Module

Fused module: merge → target → reject → materialize → draft block.

Parameters:

config (UnifiedDSparkGemma4Config)

input_types()​

input_types()

source

Single-device DSpark graph. See build_spec_decode_input_types() for the canonical ordering. Signal buffers are declared even though the graph is single-device: Gemma4’s embedding/lm_head layers use collectives unconditionally.

Return type:

tuple[TensorType | BufferType, …]

UnifiedDSparkGemma4Config​

class max.pipelines.architectures.unified_dspark_gemma4.UnifiedDSparkGemma4Config(*, target: 'Gemma4ForConditionalGenerationConfig', draft: 'DSparkGemma4DraftConfig', draft_kv_params: 'KVCacheParams', speculative_config: 'SpeculativeConfig', target_layer_ids: 'list[int]' = <factory>, mask_token_id: 'int' = 0, block_size: 'int' = 0)

source

Bases: ArchConfigWithKVCache

Parameters:

DEFAULT_ENCODING​

DEFAULT_ENCODING: ClassVar[max.pipelines.lib.config.SupportedEncoding] = 'bfloat16'

source

SUPPORTED_ENCODINGS​

SUPPORTED_ENCODINGS: ClassVar[set[max.pipelines.lib.config.SupportedEncoding]] = {'bfloat16'}

source

block_size​

block_size: int = 0

source

devices​

property devices: list[DeviceRef]

source

Devices the unified model runs on (the memory-planner protocol).

draft​

draft: DSparkGemma4DraftConfig

source

draft_kv_params​

draft_kv_params: KVCacheParams

source

get_kv_params()​

get_kv_params()

source

KV cache parameters to use when running the model.

Return type:

KVCacheParamInterface

get_max_seq_len()​

get_max_seq_len()

source

Returns the default maximum sequence length for the model.

Subclasses should determine whether this value can be overridden by setting the --max-length (pipeline_config.model.max_length) flag.

Return type:

int

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

mask_token_id​

mask_token_id: int = 0

source

resolve_block_size()​

resolve_block_size(*, default=None)

source

Parameters:

default (int | None)

Return type:

int

speculative_config​

speculative_config: SpeculativeConfig

source

target​

target: Gemma4ForConditionalGenerationConfig

source

target_layer_ids​

target_layer_ids: list[int]

source

validate_dspark_fields()​

validate_dspark_fields()

source

Strict validation of the DSpark fields parsed from the draft HF config. Unlike DFlash, ALL block_size positions produce drafts, so num_speculative_tokens must equal block_size (no - 1).

Return type:

None

UnifiedDSparkGemma4Inputs​

class max.pipelines.architectures.unified_dspark_gemma4.UnifiedDSparkGemma4Inputs(tokens, input_row_offsets, return_n_logits, signal_buffers, *, 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 unified DSpark Gemma4 graph.

The spec-decode fields and trailing buffer packing come from UnifiedSpecDecodeInputs; tokens / input_row_offsets / return_n_logits / signal_buffers plus the KV cache tree form this single-device graph’s prefix. Signal buffers are bound because Gemma4’s embedding and lm_head layers use collectives even on one device. The DSpark graph does not bind in_thinking_phase.

Parameters:

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

signal_buffers​

signal_buffers: list[Buffer]

source

tokens​

tokens: Buffer

source

UnifiedDSparkGemma4Model​

class max.pipelines.architectures.unified_dspark_gemma4.UnifiedDSparkGemma4Model(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, GraphPipelineModelWithKVCache[TextContext]

Unified DSpark Gemma4: target + draft in one compiled graph.

Parameters:

batch_processor_cls​

batch_processor_cls

source

alias of UnifiedDSparkGemma4BatchProcessor

get_kv_params()​

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

source

Returns the KV cache params for the pipeline model.

Delegates to model_config_cls.construct_kv_params(...). Subclasses with custom KV behavior should override this method.

Parameters:

Return type:

MultiKVCacheParams

model​

model: Model

source

model_config_cls​

model_config_cls

source

alias of UnifiedDSparkGemma4Config

Models

Cookbook

Puzzles

Blog

Community

Contact