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

DFlash draft model for Kimi K2.5 (MLA target with MHA/GQA draft).

Same DFlash architecture as dflash_llama3 (multi-layer non-causal block transformer with external-KV materialization) but attached to a Kimi K2.5 (DeepseekV3 MLA) target. The draft itself uses standard MHA/GQA attention with Deepseek-yarn RoPE so its KV cache geometry is decoupled from the target’s MLA layout. Supports tensor-parallel and data-parallel device topologies.

DFlashKimiK25​

class max.pipelines.architectures.dflash_kimi_k25.DFlashKimiK25(config)

source

Bases: Module

DFlash draft transformer for a Kimi K2.5 (MLA) target.

Parameters:

config (DFlashKimiK25DraftConfig)

ctx_at_draft_cache_length​

ctx_at_draft_cache_length = True

source

forward_block()​

forward_block(input_embeds, kv_collections, input_row_offsets, signal_buffers)

source

Parameters:

Return type:

list[TensorValue]

mask_embeddings()​

mask_embeddings(target_mask)

source

Returns the target’s embedding of the mask token, as trained.

Parameters:

target_mask (Sequence[TensorValue])

Return type:

list[TensorValue]

materialize()​

materialize(target_hidden, input_row_offsets, kv_collections, signal_buffers)

source

Parameters:

Return type:

None

supports_zero_draft_rows​

supports_zero_draft_rows = False

source

DFlashKimiK25DraftArchConfig​

class max.pipelines.architectures.dflash_kimi_k25.DFlashKimiK25DraftArchConfig(max_seq_len)

source

Bases: ArchConfigWithBoundedMaxSeqLen

The draft architecture’s config: its sequence bound and KV leaf.

Parameters:

max_seq_len (int)

build_module()​

classmethod build_module(context)

source

Builds the draft from the draft checkpoint context names.

It embeds and projects through its target’s embedding and head, so spec decoding shares those onto it.

Raises:

  • ValueError – If the method isn’t dflash.
  • NotImplementedError – If the checkpoint names the target’s first layer, whose input is the raw embeddings.

Parameters:

context (ModelLoadContext)

Return type:

tuple[DFlashKimiK25, DFlashKimiK25DraftConfig]

construct_kv_params()​

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

source

Returns the draft’s cache, in bfloat16 whatever the target caches in, and as wide as the block it writes.

Parameters:

Return type:

KVCacheParams

initialize()​

classmethod initialize(pipeline_config, model_config=None, *, max_seq_len)

source

Records the sequence length the pipeline resolved.

Parameters:

Return type:

DFlashKimiK25DraftArchConfig

max_seq_len​

max_seq_len: int

source

DFlashKimiK25DraftConfig​

class max.pipelines.architectures.dflash_kimi_k25.DFlashKimiK25DraftConfig(*, hidden_size, num_attention_heads, num_key_value_heads, head_dim, intermediate_size, num_hidden_layers, vocab_size, rms_norm_eps, rope_theta, max_position_embeddings, devices, data_parallel_degree, dtype, norm_dtype, kv_params, rope_scaling, target_layer_ids, mask_token_id=0, block_size=0, residual_multiplier=1.0, attention_multiplier=None, attention_bias=False, stacked_qkv=False, clip_qkv=None, sliding_window=None, use_qk_norm=True)

source

Bases: object

Minimal config for a DFlash draft over a Kimi K2.5 (MLA) target.

Held separate from DeepseekV3Config because the DFlash draft is MHA/GQA with its own KV geometry; the target’s MLA-specific fields (kv_lora_rank, v_head_dim and so on) don’t apply.

Parameters:

attention_bias​

attention_bias: bool = False

source

attention_multiplier​

attention_multiplier: float | None = None

source

block_size​

block_size: int = 0

source

The block the draft was trained on: its anchor and its drafts.

clip_qkv​

clip_qkv: float | None = None

source

data_parallel_degree​

data_parallel_degree: int

source

devices​

devices: list[DeviceRef]

source

dtype​

dtype: DType

source

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 = 0

source

The token the block’s drafted slots stand for.

max_position_embeddings​

max_position_embeddings: int

source

norm_dtype​

norm_dtype: DType

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

residual_multiplier​

residual_multiplier: float = 1.0

source

rms_norm_eps​

rms_norm_eps: float

source

rope_scaling​

rope_scaling: dict[str, Any]

source

rope_theta​

rope_theta: float

source

sliding_window​

sliding_window: int | None = None

source

stacked_qkv​

stacked_qkv: bool = False

source

target_layer_ids​

target_layer_ids: list[int]

source

use_qk_norm​

use_qk_norm: bool = True

source

vocab_size​

vocab_size: int

source

Models

Cookbook

Puzzles

Blog

Community