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_nemotron_h

Unified DSpark speculative decoding for Nemotron-H.

DSparkQwen3​

class max.pipelines.architectures.unified_dspark_nemotron_h.DSparkQwen3(config, kv_params, *, num_devices, max_seq_len, device)

source

Bases: Module

The DSpark draft: context projection, layers and Markov head.

Builds the draft.

Parameters:

  • config (DSparkQwen3DraftConfig) – The draft’s hyperparameters.
  • kv_params (KVCacheParams) – The draft’s KV cache leaf.
  • num_devices (int) – The tensor-parallel degree.
  • max_seq_len (int) – The longest sequence, which sizes the rope table.
  • device (Any) – The device the rope table is built on.

forward()​

forward(embeds, kv_collection, input_row_offsets)

source

Defines the computation performed by the module.

Users must override this method in their subclass to define the module’s computation.

Parameters:

  • *args – Positional arguments for the computation.
  • **kwargs – Keyword arguments for the computation.
  • embeds (Tensor)
  • kv_collection (PagedCacheValues)
  • input_row_offsets (Tensor)

Returns:

The result of applying the module to the input.

Raises:

NotImplementedError – If the subclass does not override this method.

Return type:

Tensor

forward_block()​

forward_block(embeds, kv_collection, input_row_offsets)

source

Runs the block, whose K and V land at the cache’s lengths.

Parameters:

  • embeds (Tensor)
  • kv_collection (PagedCacheValues)
  • input_row_offsets (Tensor)

Return type:

Tensor

materialize()​

materialize(ctx, kv_collection, input_row_offsets)

source

Writes every layer’s K and V of the context rows into the cache.

Parameters:

  • ctx (Tensor)
  • kv_collection (PagedCacheValues)
  • input_row_offsets (Tensor)

Return type:

None

project_context()​

project_context(taps)

source

Projects the concatenated target taps to the draft’s context.

Parameters:

taps (Tensor)

Return type:

Tensor

DSparkQwen3DraftConfig​

class max.pipelines.architectures.unified_dspark_nemotron_h.DSparkQwen3DraftConfig(hidden_size, intermediate_size, num_hidden_layers, num_attention_heads, num_key_value_heads, head_dim, rms_norm_eps, vocab_size, rope_theta, sliding_window, max_position_embeddings, block_size, mask_token_id, target_layer_ids, markov_rank, sample_from_anchor, quantized)

source

Bases: object

The DSpark draft’s hyperparameters, from its checkpoint config.

Parameters:

  • hidden_size (int)
  • intermediate_size (int)
  • num_hidden_layers (int)
  • num_attention_heads (int)
  • num_key_value_heads (int)
  • head_dim (int)
  • rms_norm_eps (float)
  • vocab_size (int)
  • rope_theta (float)
  • sliding_window (int)
  • max_position_embeddings (int)
  • block_size (int)
  • mask_token_id (int)
  • target_layer_ids (tuple[int, ...])
  • markov_rank (int)
  • sample_from_anchor (bool)
  • quantized (bool)

block_size​

block_size: int

source

head_dim​

head_dim: int

source

hidden_size​

hidden_size: int

source

intermediate_size​

intermediate_size: int

source

markov_rank​

markov_rank: int

source

mask_token_id​

mask_token_id: int

source

max_position_embeddings​

max_position_embeddings: int

source

num_attention_heads​

num_attention_heads: int

source

num_context_features​

property num_context_features: int

source

How many target taps feed the fc projection.

num_drafts​

property num_drafts: int

source

How many tokens a block proposes.

num_hidden_layers​

num_hidden_layers: int

source

num_key_value_heads​

num_key_value_heads: int

source

quantized​

quantized: bool

source

Whether the checkpoint stores its linears NVFP4 weight-only.

rms_norm_eps​

rms_norm_eps: float

source

rope_theta​

rope_theta: float

source

sample_from_anchor​

sample_from_anchor: bool

source

Whether the anchor slot of the block predicts a draft token.

False is the DFlash layout: the block’s first slot holds the committed token and predicts nothing, so a block of block_size carries block_size - 1 drafts.

sliding_window​

sliding_window: int

source

tap_layer_inputs​

property tap_layer_inputs: tuple[int, ...]

source

The target layers whose inputs the target taps.

The checkpoint names layer outputs and the target taps layer inputs, and a layer’s output is the next layer’s input.

target_layer_ids​

target_layer_ids: tuple[int, ...]

source

The target layers whose outputs the draft reads, as the checkpoint names them. Tapping a layer’s output is tapping the next layer’s input.

vocab_size​

vocab_size: int

source

UnifiedDSparkNemotronH​

class max.pipelines.architectures.unified_dspark_nemotron_h.UnifiedDSparkNemotronH(config, target, draft, *, enable_structured_output=False)

source

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

Merge, verify on scratch Mamba state, accept, commit, then block draft.

The phases live in BlockDriver. This architecture’s own are the two adapters in spec_adapters: the target’s snapshot, verify and replay of the Mamba state, and the draft’s use of the target’s embedding and head.

Parameters:

UnifiedDSparkNemotronHConfig​

class max.pipelines.architectures.unified_dspark_nemotron_h.UnifiedDSparkNemotronHConfig(*, target, draft, draft_kv_params, speculative_config, target_layer_ids=<factory>, mask_token_id=0, block_size=0)

source

Bases: ArchConfigWithKVCache

Nemotron-H and its DSpark draft, side by side.

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', 'float4_e2m1fnx2'}

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

devices​

property devices: list[DeviceRef]

source

The devices the unified model runs on.

draft​

draft: DSparkQwen3DraftConfig

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 (such as 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 (such as diffusion components) ignore it.

Return type:

Self

mask_token_id​

mask_token_id: int = 0

source

speculative_config​

speculative_config: SpeculativeConfig

source

target​

target: NemotronHConfig

source

target_layer_ids​

target_layer_ids: list[int]

source

The target layers whose inputs are tapped for the draft.

validate()​

validate()

source

Checks the draft against the target it is paired with.

Raises:

ValueError – If the pair cannot run together.

Return type:

None

UnifiedDSparkNemotronHInputs​

class max.pipelines.architectures.unified_dspark_nemotron_h.UnifiedDSparkNemotronHInputs(tokens, input_row_offsets, return_n_logits, signal_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

Inputs for the unified DSpark Nemotron-H graph.

The spec-decode fields and trailing buffer packing come from UnifiedSpecDecodeInputs; tokens / input_row_offsets / return_n_logits plus the KV cache tree form the graph’s prefix. The signal buffers of a multi-device mesh trail every declared input, because the graph is traced by Module.compile.

Parameters:

buffers​

property buffers: tuple[Buffer, ...]

source

Returns positional Buffer inputs for model ABI calls.

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

UnifiedDSparkNemotronHModel​

class max.pipelines.architectures.unified_dspark_nemotron_h.UnifiedDSparkNemotronHModel(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, NemotronHModel

Nemotron-H with the DSpark block draft, compiled as one graph.

Parameters:

batch_processor_cls​

batch_processor_cls

source

alias of UnifiedDSparkNemotronHBatchProcessor

get_kv_params()​

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

source

Returns the KV cache params for the pipeline model.

Delegates to model_config_cls.construct_kv_params(...). Subclasses with custom KV behavior should override this method.

Parameters:

Return type:

KVCacheParamInterface

load_model()​

load_model()

source

Build and compile the ModuleV3 callable.

Return type:

Any

model_config_cls​

model_config_cls

source

alias of UnifiedDSparkNemotronHConfig

UnifiedDSparkNemotronHModule​

class max.pipelines.architectures.unified_dspark_nemotron_h.UnifiedDSparkNemotronHModule(config, target, draft, *, kv_params, speculative_config, enable_structured_output)

source

Bases: Module

The model one compiled graph runs: target, draft and the driver.

The weights are target.* and draft.*. Takes the unified spec-decode signature’s inputs, as input_types() declares them, and returns the driver’s outputs.

Parameters:

driver()​

driver()

source

Builds the driver over this module’s target and draft.

Built per trace rather than held, because compiling traces a copy of this module whose parameters are the graph’s weights, and the driver must call that copy’s layers.

Return type:

UnifiedDSparkNemotronH

forward()​

forward(*inputs)

source

Defines the computation performed by the module.

Users must override this method in their subclass to define the module’s computation.

Parameters:

  • *args – Positional arguments for the computation.
  • **kwargs – Keyword arguments for the computation.
  • inputs (Tensor)

Returns:

The result of applying the module to the input.

Raises:

NotImplementedError – If the subclass does not override this method.

Return type:

tuple[Tensor, …]

input_types()​

input_types()

source

The types forward() takes, one per positional input.

Return type:

tuple[TensorType | BufferType, …]

Models

Cookbook

Puzzles

Blog

Community