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

The native DSpark draft for Gemma4 targets.

DSparkGemma4

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

source

Bases: Module

DSpark draft transformer for a Gemma4 target.

Parameters:

• config (DSparkGemma4DraftConfig)
• kv_params (KVCacheParams)
• devices (list[DeviceRef])
• dtype (DType)

forward_block()

forward_block(input_embeds, kv_collection, input_row_offsets)

source

Parameters:

• input_embeds (TensorValue)
• kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue])
• input_row_offsets (TensorValue)

Return type:

TensorValue

materialize_kv()

materialize_kv(ctx_hidden, input_row_offsets, kv_collection)

source

Parameters:

• ctx_hidden (TensorValue)
• input_row_offsets (TensorValue)
• kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue])

Return type:

None

project_target_hidden()

project_target_hidden(target_hs_concat)

source

Parameters:

target_hs_concat (TensorValue)

Return type:

TensorValue

DSparkGemma4DraftConfig

class max.pipelines.architectures.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 (such as 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