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_dspark_deepseekV41

DeepSeek-V4.1-Flash speculative decoding with its built-in DSpark drafter.

UnifiedDSparkDeepseekV41​

class max.pipelines.architectures.unified_dspark_deepseekV41.UnifiedDSparkDeepseekV41(config, speculative_config, specs, *, model, trailing_input_types=(), enable_structured_output=False)

source

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

Verifies drafts with the target model and drafts blocks with DSpark.

Target and draft are one DeepseekV41WithDSpark; draft names its drafter so the driver’s module tree has both roles, but the weights are loaded and registered through target alone. On more than one device both roles run on the target’s per-device copies (shard()): tensor-parallel, or data-parallel with expert-parallel MoE when the config’s data-parallel degree is the device count.

The model takes graph values at construction (its Engram tables are graph inputs), so the graph’s signature has to exist before it does: built with model=None, this driver only answers input_types() (and decode_inputs()), from the same spec a model-carrying one uses.

trailing_input_types are this model’s own graph inputs, appended after the driver’s standard inputs, in order: the Engram lookback window (when the config has Engram), then the pinned Engram tables, one weight/scale pair per device and Engram layer.

Parameters:

shard()​

shard(devices, engram_host_tables)

source

Runs both roles on per-device copies of target.

Call it once, after the weights are loaded into target, which stays the module the weights registry is built from. Returns the per-device copies, one per device.

Parameters:

Return type:

list[DeepseekV41WithDSpark]

UnifiedDSparkDeepseekV41BatchProcessor​

class max.pipelines.architectures.unified_dspark_deepseekV41.UnifiedDSparkDeepseekV41BatchProcessor(config, runtime)

source

Bases: UnifiedSpecDecodeBatchProcessor[UnifiedDSparkDeepseekV41Inputs, DeepseekV41Config]

Ragged batching with the sampling seed and the Engram lookback window.

Parameters:

bind_ep_buffers()​

bind_ep_buffers(buffers)

source

Binds the expert-parallel buffers created at model load.

Parameters:

buffers (Sequence[Buffer])

Return type:

None

prepare_initial_token_inputs()​

prepare_initial_token_inputs(replica_batches, kv_cache_inputs=None, return_n_logits=1)

source

Prepares ragged token inputs with persistent buffers and a per-step seed.

Parameters:

Return type:

UnifiedDSparkDeepseekV41Inputs

UnifiedDSparkDeepseekV41Inputs​

class max.pipelines.architectures.unified_dspark_deepseekV41.UnifiedDSparkDeepseekV41Inputs(tokens, input_row_offsets, return_n_logits, signal_buffers, host_input_row_offsets=None, data_parallel_splits=None, batch_context_lengths=<factory>, ep_buffers=(), engram_lookback_ids=None, engram_lookback_source=None, engram_lookback_history=None, engram_host_buffers=(), *, 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, draft_slot_ids=None, draft_block_offsets=None, pinned_bitmask=None, wait_payload=None, device_bitmask_scratch=None, structured_output=False)

source

Bases: UnifiedSpecDecodeInputs

Graph inputs for one step of target-plus-drafter decoding.

buffers lists them in graph-input order: the ragged tokens; the per-device image span rows and their token rows (set by the pipeline’s vision encoder cache, zero rows on a step without images); the row offsets (and, under data parallelism, their host copy) and logit count; under data parallelism the replica splits; the signal buffers (only with more than one device); the KV cache buffers; under data parallelism each device’s context length and the expert-parallel buffers; the speculative-decoding inputs; the Engram lookback window, its per-row source, the lookback history and the pinned Engram tables (only when the model has Engram); and last, a placeholder buffer whose length sets MERGED_ROWS_DIM.

Parameters:

batch_context_lengths​

batch_context_lengths: list[Buffer]

source

Per device, its replica’s summed context length, int32 [1] on CPU.

Only under data parallelism.

buffers​

property buffers: tuple[Buffer, ...]

source

Returns positional Buffer inputs for model ABI calls.

data_parallel_splits​

data_parallel_splits: Buffer | None = None

source

Int64 [replicas + 1] request boundaries of each replica, on CPU.

Only under data parallelism, where the batch is replica-major.

engram_host_buffers​

engram_host_buffers: Sequence[Buffer] = ()

source

The pinned host Engram tables, the graph’s last inputs but one.

Held here so buffers is the complete input list, which device graph capture needs.

engram_lookback_history​

engram_lookback_history: Buffer | None = None

source

[rows, max_ngram_size - 1] device-0 windows the previous step recorded for this one, one row per row of the previous batch.

Rows follow the replica-major batch, as the lookback window does, so one buffer serves every data-parallel replica.

engram_lookback_ids​

engram_lookback_ids: Buffer | None = None

source

[batch, max_ngram_size - 1] int64 token ids before each chunk.

Each request’s raw token ids just before its chunk, on device 0.

engram_lookback_source​

engram_lookback_source: Buffer | None = None

source

[batch] int64 row of engram_lookback_history to read.

-1 reads engram_lookback_ids instead.

ep_buffers​

ep_buffers: Sequence[Buffer] = ()

source

The routed experts’ expert-parallel communication buffers; only under data parallelism.

host_input_row_offsets​

host_input_row_offsets: Buffer | None = None

source

CPU copy of input_row_offsets; only under data parallelism.

input_row_offsets​

input_row_offsets: Buffer

source

return_n_logits​

return_n_logits: Buffer

source

signal_buffers​

signal_buffers: list[Buffer]

source

tokens​

tokens: Buffer

source

UnifiedDSparkDeepseekV41Model​

class max.pipelines.architectures.unified_dspark_deepseekV41.UnifiedDSparkDeepseekV41Model(*args, **kwargs)

source

Bases: _UnifiedSpecDecodeModelMixin, DeepseekV41Model

DeepSeek-V4.1 that verifies and drafts tokens in one graph per step.

The drafter’s stages store their context KV in the target model’s sliding-window cache, as layers numbered from num_hidden_layers up, so there is no separate draft cache to size or feed.

Parameters:

batch_processor_cls​

batch_processor_cls

source

alias of UnifiedDSparkDeepseekV41BatchProcessor

execute()​

execute(model_inputs)

source

Runs one step, binding the pinned Engram tables as graph inputs.

The tables follow the batch’s own inputs, in the graph’s input order.

Parameters:

model_inputs (ModelInputs)

Return type:

UnifiedEagleOutputs

Models

Cookbook

Puzzles

Blog

Community