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

The DSpark draft of Muse Spark.

MuseSparkDSpark​

class max.pipelines.architectures.dspark_muse_spark.MuseSparkDSpark(config, kv_params, *, max_seq_len, device)

source

Bases: Module

The drafter: context projection, layers and heads.

Parameters:

confidence()​

confidence(hidden, markov_embed)

source

Each slot’s predicted acceptance probability.

Parameters:

Return type:

Tensor

forward()​

forward(embeds, kv_collection, input_row_offsets)

source

Runs the embedded block; returns its final-normed rows.

Parameters:

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

Return type:

Tensor

has_own_lm_head​

has_own_lm_head = True

source

markov_embed()​

markov_embed(previous)

source

The rank-r embedding of each row’s previous token.

Parameters:

previous (Tensor)

Return type:

Tensor

materialize()​

materialize(taps, kv_collection, input_row_offsets)

source

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

Parameters:

  • taps (Tensor) – The target’s tapped hidden states, [rows, taps * hidden].
  • kv_collection (PagedCacheValues) – The draft’s cache at the context’s position.
  • input_row_offsets (Tensor) – The context rows’ ragged offsets.

Return type:

None

prepare_freq_cis()​

prepare_freq_cis(mesh)

source

Places the RoPE table on mesh in bfloat16.

Parameters:

mesh (DeviceMesh)

Return type:

None

slot_logits()​

slot_logits(base, markov_embed)

source

One slot’s fp32 logits: the base, the Markov bias, the transform.

Parameters:

Return type:

Tensor

MuseSparkDSparkDraftArchConfig​

class max.pipelines.architectures.dspark_muse_spark.MuseSparkDSparkDraftArchConfig(max_position_embeddings=262144)

source

Bases: DSparkSpeculatorsDraftArchConfig

The Muse Spark DSpark draft’s config, which builds its module.

The draft is only served as –draft-model beside a Muse Spark target.

Parameters:

max_position_embeddings (int)

build_module()​

classmethod build_module(context)

source

Builds the drafter for the draft checkpoint context names.

The module holds no weights: it compiles with those checkpoint_weights() returns.

Raises:

ValueError – If the method isn’t dspark, or the drafter can’t run beside the target.

Parameters:

context (ModelLoadContext)

Return type:

tuple[MuseSparkDSpark, MuseSparkDraftConfig]

checkpoint_weights()​

classmethod checkpoint_weights(context, config)

source

Returns the weights the built drafter compiles with.

Parameters:

Return type:

dict[str, Any]

construct_kv_params()​

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

source

Returns the drafter’s own KV leaf, BF16 and full-length.

The drafter’s sliding window is its attention mask’s, not its cache’s.

Parameters:

Return type:

KVCacheParams

MuseSparkDraftConfig​

class max.pipelines.architectures.dspark_muse_spark.MuseSparkDraftConfig(spec, post_norm_eps, qk_scale_factor, residual_gate_temperature, output_multiplier, final_logit_softcapping, confidence_with_markov, target_num_hidden_layers)

source

Bases: object

The drafter’s hyperparameters, from its checkpoint config.

Parameters:

  • spec (DSparkSpeculatorsDraftConfig)
  • post_norm_eps (float)
  • qk_scale_factor (float)
  • residual_gate_temperature (float)
  • output_multiplier (float)
  • final_logit_softcapping (float)
  • confidence_with_markov (bool)
  • target_num_hidden_layers (int)

block_size​

property block_size: int

source

The drafter’s block width.

confidence_with_markov​

confidence_with_markov: bool

source

final_logit_softcapping​

final_logit_softcapping: float

source

from_huggingface_config()​

classmethod from_huggingface_config(hf)

source

Parses the muse_spark_dspark config.

Raises:

ValueError – If the config describes a draft these layers do not build.

Parameters:

hf (Any)

Return type:

MuseSparkDraftConfig

hidden_size​

property hidden_size: int

source

The drafter’s hidden width.

mask_token_id​

property mask_token_id: int

source

The id that fills the block’s draft slots.

num_target_layers​

property num_target_layers: int

source

The target model’s layer count.

output_multiplier​

output_multiplier: float

source

post_norm_eps​

post_norm_eps: float

source

qk_scale_factor​

qk_scale_factor: float

source

residual_gate_temperature​

residual_gate_temperature: float

source

spec​

spec: DSparkSpeculatorsDraftConfig

source

The fields every speculators-format DSpark draft has.

target_layer_ids​

property target_layer_ids: tuple[int, ...]

source

The target layers whose outputs the drafter reads.

target_num_hidden_layers​

target_num_hidden_layers: int

source

vocab_size​

property vocab_size: int

source

The drafter’s vocabulary size.

check_pairing()​

max.pipelines.architectures.dspark_muse_spark.check_pairing(target_text, draft)

source

Checks the drafter against the target it runs beside.

Parameters:

  • target_text (Any) – The target’s Hugging Face text_config.
  • draft (MuseSparkDraftConfig) – The parsed drafter config.

Raises:

ValueError – If the pair cannot run together.

Return type:

None

Models

Cookbook

Puzzles

Blog

Community