vllm.models.muse_spark

Muse Spark HF model entry point.

Modules:

• amd –

  AMD ROCm implementation of Muse Spark.

• common –

  Platform-independent MuseSpark preprocessing.

• nvidia –

  NVIDIA-specific MuseSpark model implementations.

Classes:

• MuseSparkDSparkForCausalLM –

  Muse DSpark head with the target's output-distribution transform.

• MuseSparkForConditionalGeneration –

  Muse Spark Omni: ROCm text stack plus NaViT vision and waveform audio.

MuseSparkDSparkForCausalLM

Bases: Qwen3DSparkForCausalLM

Muse DSpark head with the target's output-distribution transform.

Source code in vllm/models/muse_spark/amd/dspark.py

class MuseSparkDSparkForCausalLM(Qwen3DSparkForCausalLM):
    """Muse DSpark head with the target's output-distribution transform."""

    has_own_lm_head = True

    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_stacked={
            ".self_attn.q_proj.": (".self_attn.qkv_proj.", "q"),
            ".self_attn.k_proj.": (".self_attn.qkv_proj.", "k"),
            ".self_attn.v_proj.": (".self_attn.qkv_proj.", "v"),
            ".self_attn.o_gate_proj.": (".self_attn.qkv_proj.", "o_gate"),
            ".feed_forward.gate_proj.": (".feed_forward.gate_up_proj.", 0),
            ".feed_forward.up_proj.": (".feed_forward.gate_up_proj.", 1),
        }
    )

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        nn.Module.__init__(self)
        assert vllm_config.speculative_config is not None
        self.draft_model_config = vllm_config.speculative_config.draft_model_config
        self.config = self.draft_model_config.hf_config
        target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
        self.model = MuseSparkDSparkModel(
            vllm_config=vllm_config,
            prefix=maybe_prefix(prefix, "model"),
            start_layer_id=target_layer_num,
        )
        self.lm_head = ParallelLMHead(
            self.config.draft_vocab_size,
            self.config.hidden_size,
            prefix=maybe_prefix(prefix, "lm_head"),
        )
        self.logits_processor = LogitsProcessor(self.config.draft_vocab_size)
        self.target_vocab_size = vllm_config.model_config.get_vocab_size()
        self.draft_id_to_target_id = nn.Parameter(
            torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
            requires_grad=False,
        )
        metap = self.config.metap
        if metap.use_metap and metap.router_has_one_inf_dim:
            width_mult = self.config.hidden_size / metap.base_width
            if metap.metap_mode == "mup":
                self.output_multiplier = width_mult**-1
            elif metap.metap_mode == "sp":
                self.output_multiplier = width_mult**-0.5
            else:
                raise ValueError(f"Unknown Metap mode {metap.metap_mode!r}")
        else:
            self.output_multiplier = 1.0
        self.output_soft_cap_temp = self.config.output_soft_cap_temp

    def apply_output_logit_transform(self, logits: torch.Tensor) -> torch.Tensor:
        logits = logits.float() * self.output_multiplier
        if self.output_soft_cap_temp is not None:
            logits = self.output_soft_cap_temp * torch.tanh(
                logits / self.output_soft_cap_temp
            )
        return logits

    def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # Identical full-vocabulary heads may be shared with an FP32 target head.
        return self.logits_processor(
            self.lm_head, hidden_states.to(self.lm_head.weight.dtype)
        )

    def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # Target residuals are FP32; the training collator casts features before FC.
        return super().combine_hidden_states(
            hidden_states.to(self.model.fc.weight.dtype)
        )

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        model_weights = {}
        for name, loaded_weight in weights:
            if name == "t2d":
                continue
            if name == "d2t":
                name = "draft_id_to_target_id"
            elif "lm_head" not in name:
                name = "model." + name
            model_weights[name] = loaded_weight
        mapper = WeightsMapper(
            orig_to_new_substr={
                "model.embed_tokens": None,
                "model.mask_embedding": None,
            },
            orig_to_new_stacked=self.hf_to_vllm_mapper.orig_to_new_stacked,
        )
        return AutoWeightsLoader(self).load_weights(
            model_weights.items(), mapper=mapper
        )

    def process_weights_after_loading(self) -> None:
        super().process_weights_after_loading()
        materialize_residual_path_constants(self.model)

MuseSparkForConditionalGeneration

Bases: MuseSparkOmniMixin, MuseSparkForCausalLM, SupportsMultiModal

Muse Spark Omni: ROCm text stack plus NaViT vision and waveform audio.

Source code in vllm/models/muse_spark/amd/multimodal.py

@register_omni_processor
class MuseSparkForConditionalGeneration(
    MuseSparkOmniMixin, MuseSparkForCausalLM, SupportsMultiModal
):
    """Muse Spark Omni: ROCm text stack plus NaViT vision and waveform audio."""

    hf_to_vllm_mapper = MuseSparkForCausalLM.hf_to_vllm_mapper | OMNI_WEIGHTS_MAPPER