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_dflash_mimo_v2

MiMo-V2.6-Flash with its DFlash drafter: the fused speculative graph.

DrafterExport

class max.pipelines.architectures.unified_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

UnifiedDflashMiMoV2

class max.pipelines.architectures.unified_dflash_mimo_v2.UnifiedDflashMiMoV2(spec, mask_embedding, enable_structured_output=False)

source

Bases: BlockDriver[list[list[TensorValue]], list[TensorValue]]

Verify [anchor, d1..dK] on the target, write the drafter’s context for the verified rows, then draft the next block.

The block is the drafter’s trained 8 wide whatever K is, so the next step’s block always overwrites every rejected row’s context.

Builds the fused graph’s module.

Parameters:

• spec (UnifiedDflashMiMoV2Spec) – The target, drafter and verify width.
• mask_embedding (WeightData) – The drafter’s trained mask slot embedding, baked into the graph as a constant.
• enable_structured_output (bool) – Whether acceptance takes the grammar bitmask.

draft

draft: DFlashMiMoV2

source

signature_kv_params

property signature_kv_params: KVCacheParamInterface

source

The KV params this model owns, for graphs whose caller has none.

The single-device EAGLE and DFlash graphs build their own tree.

target

target: MiMoV2

source

UnifiedDflashMiMoV2Model

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

source

Bases: _UnifiedSpecDecodeModelMixin, MiMoV2Model

Target and drafter in one compiled graph per speculative step.

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)

batch_processor_cls

batch_processor_cls

source

alias of UnifiedDflashMiMoV2BatchProcessor

drafter_export

drafter_export: DrafterExport

source

load_graphs()

load_graphs(session)

source

Loads the fused graph as the decoder.

The base registers MiMo’s own graph through graph specs; this graph still builds through load_model().

Parameters:

session (InferenceSession)

Return type:

dict[str, Model]

model_config_cls

model_config_cls

source

alias of UnifiedDflashMiMoV2Config

UnifiedDflashMiMoV2Spec

class max.pipelines.architectures.unified_dflash_mimo_v2.UnifiedDflashMiMoV2Spec(*, target, draft, speculative_config, num_speculative_tokens, sampleable_vocab_size)

source

Bases: object

What the fused graph is built from.

Parameters:

• target (MiMoV2Config)
• draft (DFlashMiMoV2Config)
• speculative_config (SpeculativeConfig)
• num_speculative_tokens (int)
• sampleable_vocab_size (int)

draft

draft: DFlashMiMoV2Config

source

num_speculative_tokens

num_speculative_tokens: int

source

How many of the block’s proposals a step verifies.

sampleable_vocab_size

sampleable_vocab_size: int

source

The tokenizer’s size. lm_head rows from here up are padding with live weights, so the graph masks them out of every verify.

speculative_config

speculative_config: SpeculativeConfig

source

target

target: MiMoV2Config

source

fused_graph()

max.pipelines.architectures.unified_dflash_mimo_v2.fused_graph(nn_model, kv_params)

source

Builds the fused graph over the {target, draft} KV tree.

Parameters:

• nn_model (UnifiedDflashMiMoV2)
• kv_params (MultiKVCacheParams)

Return type:

Graph