For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Python module
max.pipelines.architectures.unified_mtp_deepseekV3
DeepSeek-V3 multi-token prediction draft model for speculative decoding with unified graph compilation.
UnifiedMTPDeepseekV3Inputs
class max.pipelines.architectures.unified_mtp_deepseekV3.UnifiedMTPDeepseekV3Inputs(tokens, input_row_offsets, signal_buffers, host_input_row_offsets, batch_context_lengths, draft_tokens=None, draft_kv_blocks=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, *, kv_cache_inputs=None, lora_ids=None, lora_ranks=None, hidden_states=None, return_n_logits, data_parallel_splits, ep_inputs=())
Bases: DeepseekV3Inputs
Inputs for the UnifiedMTPDeepseekV3 model.
-
Parameters:
-
- tokens (Buffer)
- input_row_offsets (Buffer)
- signal_buffers (list[Buffer])
- host_input_row_offsets (Buffer)
- batch_context_lengths (list[Buffer])
- draft_tokens (Buffer | None)
- draft_kv_blocks (list[Buffer] | None)
- seed (Buffer | None)
- temperature (Buffer | None)
- top_k (Buffer | None)
- max_k (Buffer | None)
- top_p (Buffer | None)
- min_top_p (Buffer | None)
- in_thinking_phase (Buffer | None)
- pinned_bitmask (Buffer | None)
- wait_payload (Buffer | None)
- device_bitmask_scratch (Buffer | None)
- kv_cache_inputs (KVCacheInputs[Buffer, Buffer] | None)
- lora_ids (Buffer | None)
- lora_ranks (Buffer | None)
- hidden_states (Buffer | list[Buffer] | None)
- return_n_logits (Buffer)
- data_parallel_splits (Buffer)
- ep_inputs (tuple[Buffer, ...])
buffers
Returns positional Buffer inputs for model ABI calls.
device_bitmask_scratch
Device scratch buffer that receives the in-graph H2D from
pinned_bitmask; the acceptance sampler reads from it. Only set
when structured output is enabled.
draft_kv_blocks
draft_tokens
in_thinking_phase
Per-batch bool flag marking rows currently inside a
<think>...</think> block; consumed by relaxed acceptance.
max_k
min_top_p
pinned_bitmask
Pinned host bitmask for constrained decoding.
Shape [batch_size, num_speculative_tokens + 1, vocab_size].
Position i contains the valid-token mask given the FSM state after
consuming draft[0:i-1]; position num_speculative_tokens is for
the bonus token. None when structured output is disabled.
seed
temperature
top_k
top_p
wait_payload
CPU int64[2] payload = [flag._unsafe_ptr, 1] consumed by
the in-graph mo.wait_host_value_with_dep op. Only set when
structured output is enabled.
UnifiedMTPDeepseekV3Model
class max.pipelines.architectures.unified_mtp_deepseekV3.UnifiedMTPDeepseekV3Model(*args, **kwargs)
Bases: DeepseekV3Model
DeepseekV3 with MTP: merge + target + rejection + shift in one graph.
execute()
execute(model_inputs)
Execute and return all 3 graph outputs for speculative decoding.
-
Parameters:
-
model_inputs (ModelInputs)
-
Return type:
load_model()
load_model(session)
Load the model with the given weights.
-
Parameters:
-
session (InferenceSession)
-
Return type:
prepare_initial_token_inputs()
prepare_initial_token_inputs(replica_batches, kv_cache_inputs=None, return_n_logits=1, draft_tokens=None, draft_kv_cache_buffers=None, **kwargs)
Prepares the initial inputs to be passed to execute().
The inputs and functionality can vary per model. For example, model
inputs could include encoded tensors, unique IDs per tensor when using
a KV cache manager, and kv_cache_inputs (or None if the model does
not use KV cache). This method typically batches encoded tensors,
claims a KV cache slot if needed, and returns the inputs and caches.
-
Parameters:
-
- replica_batches (Sequence[Sequence[TextContext]])
- kv_cache_inputs (KVCacheInputs[Buffer, Buffer] | None)
- return_n_logits (int)
- draft_tokens (Buffer | None)
- draft_kv_cache_buffers (list[Buffer] | None)
-
Return type:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!