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

The MiMo-V2.6-Flash DFlash drafter (dflash/ in the checkpoint).

DFlashContextWriter

class max.pipelines.architectures.dflash_mimo_v2.DFlashContextWriter(config, layers=None)

source

Bases: Module

Writes drafter context K/V for committed target rows.

Built alone, it declares only the weights the writer reads (fc, hidden_norm and each layer’s k_proj, v_proj and k_norm), under the drafter checkpoint’s names.

fc is row parallel: each device multiplies its slice of the tap features, and one allreduce sums the partial products ahead of hidden_norm.

Builds the writer.

Parameters:

• config (DFlashMiMoV2Config) – The drafter configuration.
• layers (Sequence[DFlashContextLayer] | None) – The drafter’s layers, when the writer is part of the drafter. None builds K/V-only layers.

combine_taps()

combine_taps(taps, signal_buffers=None)

source

Returns fc(cat(taps)) on every device, before hidden_norm.

Parameters:

• taps (Sequence[Sequence[TensorValue]]) – Per device, the [rows, hidden] residual stream after each of target_layer_ids, in that order, replicated.
• signal_buffers (Sequence[BufferValue] | None) – The allreduce buffers; needed on more than one device.

Returns:

The context projection, [rows, hidden] per device.

Return type:

list[TensorValue]

project_features()

project_features(features, signal_buffers=None)

source

Returns fc(features) on every device, before hidden_norm.

Parameters:

• features (Sequence[TensorValue]) – Per device, the taps concatenated in target_layer_ids order, replicated.
• signal_buffers (Sequence[BufferValue] | None) – The allreduce buffers; needed on more than one device.

Returns:

The context projection, [rows, hidden] per device.

Return type:

list[TensorValue]

write()

write(ctx, input_row_offsets, kv_collections)

source

Writes every layer’s context K/V from combine_taps()’s output.

Parameters:

• ctx (Sequence[TensorValue]) – Per device, the context projection.
• input_row_offsets (Sequence[TensorValue]) – Per device, ragged offsets of the rows.
• kv_collections (Sequence[KVCacheInputsPerDevice[TensorValue, BufferValue]]) – Per device, the drafter cache, whose lengths are the rows’ first positions.

Return type:

None

DFlashMiMoV2

class max.pipelines.architectures.dflash_mimo_v2.DFlashMiMoV2(config, mask_embedding=None)

source

Bases: DFlashContextWriter

The drafter: the context writer plus the block forward.

The target’s embed_tokens embeds the anchor and its lm_head reads the drafts; the mask slots take the checkpoint’s trained mask_embedding, not the target’s embedding of mask_token_id, which is an untrained padding row.

Builds the drafter.

Parameters:

• config (DFlashMiMoV2Config) – The drafter configuration.
• mask_embedding (WeightData | None) – The trained mask embedding, to bake into the graph as a constant, so an exported graph carries it rather than naming a weight the engine must supply. None declares the mask_embedding weight instead.

ctx_at_draft_cache_length

ctx_at_draft_cache_length = True

source

forward_block()

forward_block(block_embeds, kv_collections, input_row_offsets, signal_buffers=None)

source

Runs the block, whose K/V lands at the caches’ lengths.

Parameters:

• block_embeds (Sequence[TensorValue]) – Per device, the block’s rows: each anchor’s target embedding, then mask_embeddings()’s row in every mask slot.
• kv_collections (Sequence[KVCacheInputsPerDevice[TensorValue, BufferValue]]) – Per device, the drafter cache, whose lengths are the anchors’ positions and whose dispatch metadata is sized for the block.
• input_row_offsets (Sequence[TensorValue]) – Per device, the block’s ragged offsets.
• signal_buffers (Sequence[BufferValue] | None) – The allreduce buffers; needed on more than one device.

Returns:

Per device, the final-normed [batch * block_size, hidden] hidden states.

Return type:

list[TensorValue]

mask_embeddings()

mask_embeddings(target_mask)

source

Returns the trained mask embedding, [1, hidden] per device.

Parameters:

target_mask (Sequence[TensorValue]) – Per device, the target’s embedding of mask_token_id, whose dtype the rows take.

Return type:

list[TensorValue]

materialize()

materialize(target_hidden, input_row_offsets, kv_collections, signal_buffers)

source

Writes the context K/V of the target’s rows into the drafter cache.

Parameters:

• target_hidden (Sequence[TensorValue]) – Per device, the target’s taps concatenated in target_layer_ids order.
• input_row_offsets (Sequence[TensorValue]) – See write().
• kv_collections (Sequence[KVCacheInputsPerDevice[TensorValue, BufferValue]]) – See write().
• signal_buffers (Sequence[BufferValue]) – See project_features().

Return type:

None

supports_zero_draft_rows

supports_zero_draft_rows = False

source

DFlashMiMoV2Config

class max.pipelines.architectures.dflash_mimo_v2.DFlashMiMoV2Config(*, hidden_size, num_attention_heads, num_key_value_heads, head_dim, intermediate_size, num_hidden_layers, rms_norm_eps, rope_theta, partial_rotary_factor, sliding_window, block_size, target_layer_ids, mask_token_id, attention_value_scale, max_seq_len, devices, kv_params, dtype=bfloat16)

source

Bases: object

The drafter’s geometry and the conventions its checkpoint was trained with.

Parameters:

• hidden_size (int)
• num_attention_heads (int)
• num_key_value_heads (int)
• head_dim (int)
• intermediate_size (int)
• num_hidden_layers (int)
• rms_norm_eps (float)
• rope_theta (float)
• partial_rotary_factor (float)
• sliding_window (int)
• block_size (int)
• target_layer_ids (list[int])
• mask_token_id (int)
• attention_value_scale (float)
• max_seq_len (int)
• devices (list[DeviceRef])
• kv_params (KVCacheParams)
• dtype (DType)

attention_value_scale

attention_value_scale: float

source

Scales context and block V alike, as vLLM does.

block_size

block_size: int

source

devices

devices: list[DeviceRef]

source

dtype

dtype: DType = 80

source

from_dflash_config()

classmethod from_dflash_config(config, *, devices, kv_params, max_seq_len)

source

Reads the drafter config, refusing any convention this module lacks.

Parameters:

• config (Mapping[str, object]) – The parsed dflash/config.json.
• devices (Sequence[DeviceRef]) – The devices the drafter runs on, one shard each.
• kv_params (KVCacheParams) – The drafter’s KV cache parameters.
• max_seq_len (int) – The longest position the RoPE table must cover.

Returns:

The drafter configuration.

Raises:

ValueError – If the config describes a drafter this module would compute differently.

Return type:

DFlashMiMoV2Config

head_dim

head_dim: int

source

hidden_size

hidden_size: int

source

intermediate_size

intermediate_size: int

source

kv_params

kv_params: KVCacheParams

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_hidden_layers

num_hidden_layers: int

source

num_key_value_heads

num_key_value_heads: int

source

partial_rotary_factor

partial_rotary_factor: float

source

rms_norm_eps

rms_norm_eps: float

source

rope_theta

rope_theta: float

source

rotary_dim

property rotary_dim: int

source

Leading dims of each head that RoPE rotates, NeoX style.

sliding_window

sliding_window: int

source

Per query, a key is visible iff key + sliding_window > query.

target_layer_ids

target_layer_ids: list[int]

source

Target decoder layers whose outputs feed fc, as numbered in config.json: id k is the residual stream after layer k.

DFlashMiMoV2DraftConfig

class max.pipelines.architectures.dflash_mimo_v2.DFlashMiMoV2DraftConfig(module, sampleable_vocab_size, drafter_export)

source

Bases: object

The drafter module’s config, and what the pairing reads beside it.

Parameters:

• module (DFlashMiMoV2Config)
• sampleable_vocab_size (int)
• drafter_export (DrafterExport)

block_size

property block_size: int

source

Returns the block width the drafter was trained at.

drafter_export

drafter_export: DrafterExport

source

What the graph bakes in of the drafter, for an export’s sidecar.

hidden_size

property hidden_size: int

source

Returns the drafter’s hidden size.

mask_token_id

property mask_token_id: int

source

Returns the token the block’s drafted slots stand for.

module

module: DFlashMiMoV2Config

source

sampleable_vocab_size

sampleable_vocab_size: int

source

The tokenizer’s size. The target’s head rows from here up are padding with live weights, so neither verify nor the draft samples them.

target_layer_ids

property target_layer_ids: list[int]

source

Returns the target layers the drafter reads.

DrafterExport

class max.pipelines.architectures.dflash_mimo_v2.DrafterExport(directory, target_layer_ids, num_speculative_tokens, spec_block_size)

source

Bases: object

What an export built into the graph from the drafter.

gen_mef records it in the sidecar with the hashes of the files in directory, which Mach re-hashes at boot.

Parameters:

• directory (Path)
• target_layer_ids (list[int])
• num_speculative_tokens (int)
• spec_block_size (int)

directory

directory: Path

source

The drafter directory whose files the graph was built from.

num_speculative_tokens

num_speculative_tokens: int

source

The verify width of the speculative graph.

spec_block_size

spec_block_size: int

source

The drafter’s block width.

target_layer_ids

target_layer_ids: list[int]

source

convert_safetensor_state_dict()

max.pipelines.architectures.dflash_mimo_v2.convert_safetensor_state_dict(state_dict, config, mask_embedding)

source

Checks the drafter’s tensors and adds the mask embedding.

Parameters:

• state_dict (Mapping[str, Weights]) – The tensors of dflash_draft_model.safetensors.
• config (DFlashMiMoV2Config) – The drafter configuration.
• mask_embedding (WeightData) – From load_mask_embedding().

Returns:

Every drafter weight, keyed by its name in DFlashMiMoV2.

Raises:

ValueError – If a tensor is missing, unexpected, or has the wrong dtype or shape.

Return type:

dict[str, WeightData]

load_mask_embedding()

max.pipelines.architectures.dflash_mimo_v2.load_mask_embedding(path, config)

source

Reads mask_embedding.pt, the embedding of every block mask slot.

Parameters:

• path (Path) – The mask_embedding.pt file.
• config (DFlashMiMoV2Config) – The drafter configuration, whose mask_token_id the file must name.

Returns:

The [hidden_size] embedding.

Raises:

ValueError – If the file is not a {mask_token_id, embedding} dict for this drafter.

Return type:

WeightData