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

Qwen3.5 fused with a DFlash2 block drafter for speculative decoding.

UnifiedDflash2Qwen3_5Config​

class max.pipelines.architectures.unified_dflash2_qwen3_5.UnifiedDflash2Qwen3_5Config(*, target, draft, draft_kv_params, speculative_config, target_layer_ids=<factory>, layer_types=<factory>, num_target_layers=None, mask_token_id=0, block_size=0, conv_kernel_size=0, conv_group_size=0, selector_rank=0, selector_top_k=0)

source

Bases: ArchConfigWithKVCache

Target, drafter and the structural constants that bind them.

Parameters:

DEFAULT_ENCODING​

DEFAULT_ENCODING: ClassVar[max.pipelines.lib.config.SupportedEncoding] = 'bfloat16'

source

SUPPORTED_ENCODINGS​

SUPPORTED_ENCODINGS: ClassVar[set[max.pipelines.lib.config.SupportedEncoding]] = {'bfloat16', 'float32', 'float4_e2m1fnx2', 'float8_e4m3fn'}

source

block_size​

block_size: int = 0

source

calculate_max_seq_len()​

classmethod calculate_max_seq_len(huggingface_config, model_config)

source

Returns the resolved maximum sequence length.

Bounds or defaults the user’s model_config.max_length with the model’s own limits. Construction runs this once and stores the result on model_config.max_length; memory planning may lower it further, but only on the memory plan.

Parameters:

  • huggingface_config (AutoConfig) – The HuggingFace config to read model bounds from.
  • model_config (MAXModelConfig) – The model config whose max_length carries the user’s setting.

Return type:

int

conv_group_size​

conv_group_size: int = 0

source

conv_kernel_size​

conv_kernel_size: int = 0

source

devices​

property devices: list[DeviceRef]

source

draft​

draft: Llama3Config

source

draft_kv_params​

draft_kv_params: KVCacheParams

source

get_kv_params()​

get_kv_params()

source

KV cache parameters to use when running the model.

Return type:

KVCacheParamInterface

get_max_seq_len()​

get_max_seq_len()

source

Returns the effective maximum sequence length for the model.

For configs that store a deployment length, this is the value initialize received; for metadata-only configs it derives from the checkpoint.

Return type:

int

initialize()​

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

source

Initialize the config from a PipelineConfig.

Parameters:

  • pipeline_config (PipelineConfig) – The pipeline configuration.
  • model_config (MAXModelConfig | None) – The model configuration to read from. When None (the default), pipeline_config.model is used. Pass an explicit config (e.g. pipeline_config.draft_model) to initialize the arch config for a different model.
  • max_seq_len (int) – The effective maximum sequence length to store on the config. The value is received, never derived here: the pipeline model passes the memory plan’s VRAM-clamped length, while memory planning (which runs before a plan exists) passes the construction-resolved model_config.max_length. Configs whose sequence length is pure model metadata (e.g. diffusion components) ignore it.

Return type:

Self

layer_types​

layer_types: list[str]

source

mask_token_id​

mask_token_id: int = 0

source

num_speculative_tokens​

property num_speculative_tokens: int

source

the anchor slot never predicts.

Type:

Mask slots per block

num_target_layers​

num_target_layers: int | None = None

source

Target depth the drafter was trained against; None when the checkpoint omits the optional field.

selector_rank​

selector_rank: int = 0

source

selector_top_k​

selector_top_k: int = 0

source

speculative_config​

speculative_config: SpeculativeConfig

source

target​

target: Qwen3_5Config

source

target_layer_ids​

target_layer_ids: list[int]

source

validate_dflash2_fields()​

validate_dflash2_fields()

source

Checks the drafter against the target it was trained to draft for.

Every one of these is a silent-wrong-answer if it is wrong: a bad tap id reads a different layer, a hidden or vocab mismatch misreads the borrowed fc / lm_head contract, and a bad mask id drafts from an embedding row the drafter never saw.

Return type:

None

UnifiedDflash2Qwen3_5Inputs​

class max.pipelines.architectures.unified_dflash2_qwen3_5.UnifiedDflash2Qwen3_5Inputs(tokens, input_row_offsets, host_input_row_offsets, return_n_logits, data_parallel_splits, signal_buffers, batch_context_lengths, live_conv_pools, live_recurrent_pools, live_conv_row_ids, live_recurrent_row_ids, shadow_conv_pools, shadow_recurrent_pools, *, kv_cache_inputs=None, lora_buffers=(), vision_embeddings=<factory>, vision_scatter_indices=<factory>, hidden_states=None, draft_tokens=None, draft_probs_full=None, seed=None, temperature=None, top_k=None, max_k=None, top_p=None, min_top_p=None, in_thinking_phase=None, pinned_bitmask=None, wait_payload=None, device_bitmask_scratch=None, structured_output=False, sampled_draft_proposal=False)

source

Bases: UnifiedSpecDecodeInputs

Inputs for the fused Qwen3.5 DFlash2 graph.

Identical to the Qwen3.5 MTP graph’s packing: the canonical spec-decode prefix and tail, then this target’s state-pool tail. Only the draft KV leaf’s shapes differ between the two graphs.

Parameters:

batch_context_lengths​

batch_context_lengths: list[Buffer]

source

buffers​

property buffers: tuple[Buffer, ...]

source

Returns positional Buffer inputs for model ABI calls.

data_parallel_splits​

data_parallel_splits: Buffer

source

host_input_row_offsets​

host_input_row_offsets: Buffer

source

input_row_offsets​

input_row_offsets: Buffer

source

live_conv_pools​

live_conv_pools: list[Buffer]

source

live_conv_row_ids​

live_conv_row_ids: list[Buffer]

source

live_recurrent_pools​

live_recurrent_pools: list[Buffer]

source

live_recurrent_row_ids​

live_recurrent_row_ids: list[Buffer]

source

return_n_logits​

return_n_logits: Buffer

source

shadow_conv_pools​

shadow_conv_pools: list[Buffer]

source

shadow_recurrent_pools​

shadow_recurrent_pools: list[Buffer]

source

signal_buffers​

signal_buffers: list[Buffer]

source

tokens​

tokens: Buffer

source

UnifiedDflash2Qwen3_5Model​

class max.pipelines.architectures.unified_dflash2_qwen3_5.UnifiedDflash2Qwen3_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: _UnifiedSpecDecodeModelMixin, AlwaysSignalBuffersMixin, GraphPipelineModelWithKVCache[Qwen3VLTextAndVisionContext]

Qwen3.5 with a DFlash2 block drafter, in one compiled graph.

Parameters:

batch_processor_cls​

batch_processor_cls

source

alias of UnifiedDflash2Qwen3_5BatchProcessor

get_kv_params()​

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

source

The target’s full-attention leaf plus the drafter’s windowed one.

Called during memory planning, before _create_model_config; both must agree, so both go through the same two constructors.

Parameters:

Return type:

MultiKVCacheParams

model​

model: Model

source

model_config_cls​

model_config_cls

source

alias of UnifiedDflash2Qwen3_5Config

Models

Cookbook

Puzzles

Blog

Community

Contact