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)

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(hidden_states)

source

Top-k candidate ids and their logits from the borrowed head.

Parameters:

hidden_states (TensorValue) – [batch, steps, hidden] mask-slot outputs.

Returns:

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

Return type:

tuple[TensorValue, TensorValue]

forward_block()​

forward_block(input_embeds, kv_collection, input_row_offsets)

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:

TensorValue

materialize_kv()​

materialize_kv(ctx_hidden, input_row_offsets, kv_collection)

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

project_target_hidden()​

project_target_hidden(target_hs_concat)

source

Parameters:

target_hs_concat (TensorValue)

Return type:

TensorValue

DFlash2Qwen3_5Model​

class max.pipelines.architectures.dflash2_qwen3_5.DFlash2Qwen3_5Model(pipeline_config, session, devices, kv_cache_config, weights, *, memory_plan, adapter=None, return_logits=ReturnLogits.LAST_TOKEN, return_hidden_states=ReturnHiddenStates.NONE, max_batch_size=1)

source

Bases: LlamaModelBase

Placeholder pipeline model for the DFlash2 draft architecture.

See module docstring. execute raises because the draft is only ever run via the unified pipeline.

Parameters:

execute()​

execute(model_inputs)

source

Executes the graph with the given inputs.

Parameters:

model_inputs (ModelInputs) – The model inputs to execute, containing tensors and any other required data for model execution.

Returns:

ModelOutputs containing the pipeline’s output tensors.

Return type:

ModelOutputs

This is an abstract method that must be implemented by concrete PipelineModels to define their specific execution logic.

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

Contact