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

DFlash2 draft model for Qwen3.5-family targets.

DFlash v1’s non-causal block drafter plus a two-tap grouped dynamic convolution around each sublayer and a top-k candidate-path selector.

DFlash2CandidateSelector​

class max.pipelines.architectures.dflash2_qwen3_5.DFlash2CandidateSelector(hidden_size, *, vocab_size, rank, top_k, dtype, device)

source

Bases: Module

Low-rank bilinear scorer over adjacent candidate tokens, plus its walk.

S[l, p, c] = unary[l, c] + <pred_book[pred_id[l, p]] * proj(h_l), succ_book[cand_id[l, c]]>, where pred_id at slot l is slot l - 1’s candidate list and, at l == 0, the anchor (the last verified token) broadcast across the predecessor axis.

The predecessor/successor role assignment is read from _score_edges in the authors’ fork, not inferred from the tensor names.

Parameters:

score_edges()​

score_edges(candidate_ids, unary_logits, hidden_states, anchor_token_ids)

source

Scores every adjacent candidate pair in the block.

Parameters:

  • candidate_ids (TensorValue) – [batch, steps, top_k] top-k ids per mask slot.
  • unary_logits (TensorValue) – [batch, steps, top_k] their draft-head logits.
  • hidden_states (TensorValue) – [batch, steps, hidden] drafter output per slot.
  • anchor_token_ids (TensorValue) – [batch] id of the last verified token.

Returns:

[batch, steps, top_k, top_k] indexed [b, l, p, c].

Return type:

TensorValue

select_path()​

select_path(scores, candidate_ids, sampler=None)

source

Walks the score tensor left to right, following the best successor.

A single greedy pass, unrolled at graph-build time over the static slot axis: no beam and no Viterbi. Slot 0’s predecessor row is the anchor broadcast across top_k, so seeding the walk at index 0 selects the same row as any other index.

Parameters:

Returns:

[batch, steps] chosen token ids.

Return type:

TensorValue

DFlash2GroupedConv​

class max.pipelines.architectures.dflash2_qwen3_5.DFlash2GroupedConv(hidden_size, *, taps, group_size, block_size, dtype, device)

source

Bases: Module

One sublayer’s pair of dynamic grouped depthwise convolutions.

Applied along the block axis, once before the sublayer and once after it:

out[t] = sum_tap (base_kernel[side, tap] + delta[t, tap, group(c)])
         * x[t - tap]

base_kernel is per-channel at full resolution; delta is per group of group_size channels and broadcast across the group. t is block-local, so the t - tap taps are zero for the first tap positions of every block and the convolution never reads across a block boundary.

kernel_projection runs once, on the sublayer’s input, and emits the deltas for both sides at once — hence its 2 * taps * num_groups output width. prepare() applies side 0 and hands back side 1’s coefficients for finish() to reuse after the sublayer.

Parameters:

finish()​

finish(hidden_states, coefficients)

source

Applies the post-sublayer convolution with prepare()’s coefficients.

Parameters:

Return type:

TensorValue

prepare()​

prepare(hidden_states)

source

Applies the pre-sublayer convolution.

Parameters:

hidden_states (TensorValue) – [total_tokens, hidden_size], one dense block_size-row block per sequence.

Returns:

The convolved states and the [total_tokens, taps, num_groups] coefficients finish() must reuse.

Return type:

tuple[TensorValue, TensorValue]

DFlash2Qwen3_5​

class max.pipelines.architectures.dflash2_qwen3_5.DFlash2Qwen3_5(config, *, num_context_features, block_size, conv_kernel_size, conv_group_size, selector_rank, selector_top_k, layer_types=None)

source

Bases: Module

DFlash2 draft transformer for a Qwen3.5 target.

Builds the draft stack.

Parameters:

  • config (Llama3Config) – The draft’s own Llama3-shaped config.
  • num_context_features (int) – Number of target hidden-state taps the fc projection consumes.
  • block_size (int) – Query tokens per request: the anchor plus the mask tokens. The convolution’s block-local axis.
  • conv_kernel_size (int) – Convolution taps.
  • conv_group_size (int) – Channels sharing one dynamic coefficient.
  • selector_rank (int) – Codebook rank.
  • selector_top_k (int) – Candidates kept per mask slot.
  • layer_types (Sequence[str] | None) – Per-layer "sliding_attention" / "full_attention" selection from the draft checkpoint. None applies config.sliding_window to every layer.

compute_candidates()​

compute_candidates(logits)

source

Top-k candidate ids and their logits.

Parameters:

logits (TensorValue) – [batch, steps, vocab] float32 logits of the borrowed head at each mask slot.

Returns:

(candidate_ids, unary_logits), both [batch, steps, selector_top_k].

Return type:

tuple[TensorValue, TensorValue]

ctx_at_draft_cache_length​

ctx_at_draft_cache_length = False

source

forward_block()​

forward_block(input_embeds, kv_collections, input_row_offsets, signal_buffers)

source

Runs block_size query rows per sequence through the stack.

input_embeds must be laid out as dense block_size-row blocks, one per sequence: the convolution’s block axis is a reshape of the token axis, not a function of input_row_offsets.

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

Writes per-layer context K/V projected from the target’s states.

No input_layernorm and no convolution on this path. Q is not skipped, though: materialize_kv_from_hidden runs the fused qkv_proj and, because the drafter sets use_qk_norm, normalizes Q alongside K before rope_split_store_ragged keeps only K/V. Storing the context K/V therefore costs a full QKV projection, not a k_proj/v_proj pair.

Parameters:

Return type:

None

supports_zero_draft_rows​

supports_zero_draft_rows = False

source

DFlash2Qwen3_5ArchConfig​

class max.pipelines.architectures.dflash2_qwen3_5.DFlash2Qwen3_5ArchConfig(*, hidden_size, num_attention_heads, num_key_value_heads, num_hidden_layers, rope_theta, rope_scaling_params, max_seq_len, intermediate_size, interleaved_rope_weights, vocab_size, dtype, model_quantization_encoding, quantization_config, kv_params, return_logits=ReturnLogits.LAST_TOKEN, norm_method='rms_norm', norm_dtype=None, attention_bias=False, rms_norm_eps=None, tie_word_embeddings=False, stacked_mlp=False, stacked_qkv=False, attention_multiplier, embedding_multiplier, residual_multiplier, devices, clip_qkv, quant_config=None, longrope_scaling_params=None, logits_scaling=1.0, return_hidden_states=ReturnHiddenStates.NONE, target_layer_ids=None, use_subgraphs=True, data_parallel_degree=1, sliding_window=None, quantization_encoding=None)

source

Bases: Llama3Config

The draft’s Llama 3 config, sizing its own windowed KV leaf and building the draft’s module.

Parameters:

build_module()​

classmethod build_module(context)

source

Builds the DFlash2 draft from the draft checkpoint context names.

The draft computes in its checkpoint’s own precision and loads strictly: a mis-named tensor would load as nothing and show only as collapsed acceptance.

Raises:

ValueError – If the method isn’t dflash2, or the checkpoint’s tap ids repeat a layer.

Parameters:

context (ModelLoadContext)

Return type:

tuple[DFlash2Qwen3_5, DFlash2Qwen3_5Config]

construct_kv_params()​

classmethod construct_kv_params(huggingface_config, pipeline_config, devices, kv_cache_config, cache_dtype, *, allow_kv_head_replication=False)

source

Builds the draft’s bounded KV leaf from its own geometry.

Two things differ from the target’s leaf and neither can be inherited: the drafter’s dtype (its K/V projections are bf16 whatever the target’s cache is) and its window. The window is what makes this leaf a per-slot constant: the drafter’s mask never reaches further back than DRAFT_SLIDING_WINDOW, so a page older than the window is neither read nor written.

Parameters:

Return type:

KVCacheParams

DFlash2TransformerBlock​

class max.pipelines.architectures.dflash2_qwen3_5.DFlash2TransformerBlock(*, attention, mlp, attention_norm, mlp_norm, attention_conv, mlp_conv)

source

Bases: Module

A DFlash draft block with a dynamic convolution around each sublayer.

r = x; h = norm(r); h = conv_a.prepare(h); h = attn(h); h = conv_a.finish(h); r += h and the same again for the MLP — four convolution applications and two kernel_projection evaluations per layer.

Parameters:

Models

Cookbook

Puzzles

Blog

Community