vllm.models.muse_spark.amd.muse_moe_runner

ROCm MoE implementation for MuseSpark1x checkpoints.

Classes:

• MuseMoERunner –

  MoE runner whose custom reducer owns the TP reduction.

• MuseSpark1xRouter –

  MuseSpark1x top-k router.

• MuseSparkAiterExperts –

  AITER experts that return one unweighted output per selected route.

• MuseSparkQuarkOCPMXFusedMoE –

  Quark W4A4 method using route-preserving stock AITER experts.

• MuseSparkQuarkW8A8Fp8FusedMoE –

  Quark W8A8 FP8 method using route-preserving Triton experts.

• MuseSparkRoutedExperts –

  MuseSpark routed experts using ROCm BF16, FP8, or MXFP4 methods.

• MuseSparkTopKWeightAndReduce –

  Apply MuseSpark's per-route norm before router weighting and reduction.

• MuseSparkTritonExperts –

  Triton experts that expose unweighted GEMM2 output per route.

Functions:

• validate_bf16_deployment –

  Reject deployments whose BF16 finalization is not implemented.

• validate_fp8_deployment –

  Reject deployments whose FP8 finalization is not implemented.

• validate_mxfp4_deployment –

  Reject deployments whose MXFP4 finalization is not implemented.

MuseMoERunner

Bases: MoERunner

MoE runner whose custom reducer owns the TP reduction.

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

class MuseMoERunner(MoERunner):
    """MoE runner whose custom reducer owns the TP reduction."""

    def __init__(
        self,
        *args,
        post_expert_norm: torch.nn.Module,
        **kwargs,
    ) -> None:
        super().__init__(*args, **kwargs)
        if not isinstance(self.routed_experts, MuseSparkRoutedExperts):
            raise TypeError(
                "MuseMoERunner requires routed_experts_cls=MuseSparkRoutedExperts."
            )
        self.routed_experts.post_expert_norm_fn = post_expert_norm.forward

    @property
    def _fused_output_is_reduced(self) -> bool:
        return needs_pre_norm_all_reduce(self.moe_config)

MuseSpark1xRouter

Bases: BaseRouter

MuseSpark1x top-k router.

Optional metap output multiplier -> sigmoid/softmax -> add expert bias for SELECTION only -> top-k -> gather UNBIASED scores -> divide by (sum + division_eps) -> scale by route_scale.

Attributes:

• routing_method_type (RoutingMethodType) –

  Classify this router in vLLM's routing vocabulary.

Source code in vllm/models/muse_spark/common/moe.py

class MuseSpark1xRouter(BaseRouter):
    """MuseSpark1x top-k router.

    Optional metap output multiplier -> sigmoid/softmax -> add expert bias for
    SELECTION only -> top-k -> gather UNBIASED scores -> divide by
    (sum + division_eps) -> scale by route_scale.
    """

    def __init__(
        self,
        top_k: int,
        num_experts: int,
        expert_bias: torch.Tensor,
        route_scale: float,
        router_output_multiplier: float | None = None,
        router_score_normalization: str = "sigmoid",
        router_score_division_eps: float | None = None,
    ):
        super().__init__(top_k=top_k, global_num_experts=num_experts)
        if router_score_normalization not in ("sigmoid", "softmax"):
            raise ValueError(
                f"Invalid router_score_normalization: {router_score_normalization}"
            )
        self.expert_bias = expert_bias
        self.route_scale = route_scale
        self.router_output_multiplier = router_output_multiplier
        self.router_score_normalization = router_score_normalization
        self.router_score_division_eps = router_score_division_eps

    @property
    def routing_method_type(self) -> RoutingMethodType:
        """Classify this router in vLLM's routing vocabulary."""
        return get_routing_method_type(
            scoring_func=self.router_score_normalization,
            top_k=self.top_k,
            renormalize=True,
            num_expert_group=None,
            has_e_score_bias=True,
        )

    def _compute_routing(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
        indices_type: torch.dtype | None,
        *,
        input_ids: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if not router_logits.is_cuda:
            return self._compute_routing_native(router_logits, indices_type)

        topk_weights, topk_ids = fused_musespark1x_routing(
            router_logits,
            self.expert_bias,
            self.top_k,
            self.route_scale,
            router_output_multiplier=self.router_output_multiplier,
            router_score_normalization=self.router_score_normalization,
            router_score_division_eps=self.router_score_division_eps,
        )
        if indices_type is not None and indices_type != torch.int32:
            topk_ids = topk_ids.to(indices_type)
        return topk_weights, topk_ids

    def _compute_routing_native(
        self,
        router_logits: torch.Tensor,
        indices_type: torch.dtype | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """The recipe in torch: the kernel's reference, and the CPU path."""
        gating_output = router_logits
        if self.router_output_multiplier is not None:
            gating_output = gating_output * self.router_output_multiplier

        if self.router_score_normalization == "sigmoid":
            gating_output = torch.sigmoid(gating_output)
        else:
            gating_output = torch.softmax(gating_output, dim=-1)

        scores_for_topk = gating_output + self.expert_bias
        _, topk_ids = torch.topk(scores_for_topk, self.top_k, dim=-1)

        topk_weights = torch.gather(gating_output, dim=1, index=topk_ids)
        denom = topk_weights.sum(dim=-1, keepdim=True)
        if self.router_score_division_eps is not None:
            denom = denom + self.router_score_division_eps
        topk_weights = topk_weights / denom * self.route_scale

        return topk_weights.to(torch.float32), topk_ids.to(
            torch.int32 if indices_type is None else indices_type
        )

routing_method_type property

Classify this router in vLLM's routing vocabulary.

_compute_routing_native(router_logits, indices_type)

The recipe in torch: the kernel's reference, and the CPU path.

Source code in vllm/models/muse_spark/common/moe.py

def _compute_routing_native(
    self,
    router_logits: torch.Tensor,
    indices_type: torch.dtype | None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """The recipe in torch: the kernel's reference, and the CPU path."""
    gating_output = router_logits
    if self.router_output_multiplier is not None:
        gating_output = gating_output * self.router_output_multiplier

    if self.router_score_normalization == "sigmoid":
        gating_output = torch.sigmoid(gating_output)
    else:
        gating_output = torch.softmax(gating_output, dim=-1)

    scores_for_topk = gating_output + self.expert_bias
    _, topk_ids = torch.topk(scores_for_topk, self.top_k, dim=-1)

    topk_weights = torch.gather(gating_output, dim=1, index=topk_ids)
    denom = topk_weights.sum(dim=-1, keepdim=True)
    if self.router_score_division_eps is not None:
        denom = denom + self.router_score_division_eps
    topk_weights = topk_weights / denom * self.route_scale

    return topk_weights.to(torch.float32), topk_ids.to(
        torch.int32 if indices_type is None else indices_type
    )

MuseSparkAiterExperts

Bases: _RouteOutputHandoff, AiterMxfp4Experts

AITER experts that return one unweighted output per selected route.

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

class MuseSparkAiterExperts(_RouteOutputHandoff, AiterMxfp4Experts):
    """AITER experts that return one unweighted output per selected route."""

    def __init__(
        self,
        moe_config: FusedMoEConfig,
        quant_config: FusedMoEQuantConfig,
    ) -> None:
        super().__init__(moe_config, quant_config)
        if not isinstance(quant_config, MuseSparkQuarkOCPMXMoEQuantConfig):
            raise TypeError(
                "MuseSparkAiterExperts requires MuseSparkQuarkOCPMXMoEQuantConfig."
            )
        if quant_config.post_expert_norm_fn is None:
            raise ValueError("MuseSpark1x post-expert norm is not configured.")
        self._unit_weights_bufs: dict[int, torch.Tensor] = {}

    def _unit_weights(
        self, num_routes: int, width: int, device: torch.device
    ) -> torch.Tensor:
        """A cached ``[num_routes, width]`` fp32 ones buffer (grown on demand)."""
        buf = self._unit_weights_bufs.get(width)
        if buf is None or buf.shape[0] < num_routes or buf.device != device:
            buf = torch.ones(
                (max(num_routes, 1), width), dtype=torch.float32, device=device
            )
            self._unit_weights_bufs[width] = buf
        return buf[:num_routes]

    def workspace_shapes(
        self,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: mk.ExpertTokensMetadata | None,
        activation: MoEActivation,
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
        workspace13, workspace2, _ = super().workspace_shapes(
            M,
            N,
            K,
            topk,
            global_num_experts,
            local_num_experts,
            expert_tokens_meta,
            activation,
        )
        return workspace13, workspace2, (M, topk, K)

    def apply(
        self,
        output: torch.Tensor,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        a1q_scale: torch.Tensor | None,
        a2_scale: torch.Tensor | None,
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: mk.ExpertTokensMetadata | None,
        apply_router_weight_on_input: bool,
    ) -> None:
        if apply_router_weight_on_input:
            raise NotImplementedError(
                "MuseSpark1x requires router weights after the post-expert norm."
            )

        num_tokens, top_k = topk_ids.shape
        hidden_dim = hidden_states.shape[-1]
        expected_output_shape = (num_tokens, top_k, hidden_dim)
        if output.shape != expected_output_shape:
            raise ValueError(
                "MuseSpark1x AITER output has shape "
                f"{tuple(output.shape)}, expected {expected_output_shape}."
            )

        expanded_hidden_states = take_expanded(hidden_states, top_k)
        if expanded_hidden_states is None:
            expanded_hidden_states = (
                hidden_states[:, None, :]
                .expand(num_tokens, top_k, hidden_dim)
                .reshape(num_tokens * top_k, hidden_dim)
                .contiguous()
            )
        # Only the shape of topk_ids is read after the experts.
        num_routes = num_tokens * top_k
        flat_topk_ids = topk_ids.view(num_routes, 1)
        if expert_map is not None:
            # AITER's EP convention: one extra route per row on the always-masked
            # last expert_mask slot; its tuned-config lookup subtracts it.
            masked_slot = flat_topk_ids.new_full(
                (num_routes, 1), expert_map.shape[0] - 1
            )
            flat_topk_ids = torch.cat((flat_topk_ids, masked_slot), dim=1)
        unit_topk_weights = self._unit_weights(
            num_routes, flat_topk_ids.shape[1], topk_ids.device
        )

        if expert_tokens_meta is not None:
            num_local_tokens = expert_tokens_meta.expert_num_tokens
        else:
            num_local_tokens = None

        route_output = rocm_aiter_fused_experts(
            hidden_states=expanded_hidden_states,
            w1=w1,
            w2=w2,
            topk_weights=unit_topk_weights,
            topk_ids=flat_topk_ids,
            activation=activation,
            expert_mask=expert_map,
            quant_config=self.quant_config,
            moe_config=self.moe_config,
            a1q_scale=a1q_scale,
            num_local_tokens=num_local_tokens,
            output_dtype=output.dtype,
            moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(),
            apply_router_weight_on_input=False,
        )
        self._park_route_output(route_output.view_as(output))

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        quant_config = self.quant_config
        assert isinstance(quant_config, MuseSparkQuarkOCPMXMoEQuantConfig)
        assert quant_config.post_expert_norm_fn is not None
        return MuseSparkTopKWeightAndReduce(
            quant_config.post_expert_norm_fn,
            pre_norm_all_reduce=self.moe_config.moe_parallel_config.tp_size > 1,
            take_route_output=self._take_route_output,
        )

_unit_weights(num_routes, width, device)

A cached [num_routes, width] fp32 ones buffer (grown on demand).

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

def _unit_weights(
    self, num_routes: int, width: int, device: torch.device
) -> torch.Tensor:
    """A cached ``[num_routes, width]`` fp32 ones buffer (grown on demand)."""
    buf = self._unit_weights_bufs.get(width)
    if buf is None or buf.shape[0] < num_routes or buf.device != device:
        buf = torch.ones(
            (max(num_routes, 1), width), dtype=torch.float32, device=device
        )
        self._unit_weights_bufs[width] = buf
    return buf[:num_routes]

MuseSparkOnlineMXFP4FusedMoE

Bases: MuseSparkQuarkOCPMXFusedMoE

Unquantized routed experts rounded to OCP MXFP4 (W4A4) after loading.

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

class MuseSparkOnlineMXFP4FusedMoE(MuseSparkQuarkOCPMXFusedMoE):
    """Unquantized routed experts rounded to OCP MXFP4 (W4A4) after loading."""

    def __init__(self, moe: FusedMoEConfig) -> None:
        super().__init__(
            moe, _WEIGHT_QUANT_KEY_MAP["mxfp4"], _ACTIVATION_QUANT_KEY_MAP["mxfp4"]
        )

    def create_weights(
        self,
        layer: RoutedExperts,
        num_experts: int,
        hidden_size: int,
        intermediate_size_per_partition: int,
        params_dtype: torch.dtype,
        **extra_weight_attrs,
    ) -> None:
        for name, shape in (
            (
                "w13_weight",
                (
                    num_experts,
                    self.moe.w13_num_shards * intermediate_size_per_partition,
                    hidden_size,
                ),
            ),
            ("w2_weight", (num_experts, hidden_size, intermediate_size_per_partition)),
        ):
            param = torch.nn.Parameter(
                torch.empty(shape, dtype=params_dtype), requires_grad=False
            )
            layer.register_parameter(name, param)
            set_weight_attrs(param, extra_weight_attrs)
        layer.w13_bias, layer.w2_bias = None, None
        layer.w13_input_scale, layer.w2_input_scale = None, None

    def process_weights_after_loading(self, layer: RoutedExperts) -> None:
        for name in ("w13_weight", "w2_weight"):
            packed, scales = _quantize_experts_to_mxfp4(getattr(layer, name).data)
            delattr(layer, name)
            layer.register_parameter(
                name, torch.nn.Parameter(packed, requires_grad=False)
            )
            layer.register_parameter(
                f"{name}_scale", torch.nn.Parameter(scales, requires_grad=False)
            )
        torch.accelerator.empty_cache()
        super().process_weights_after_loading(layer)

MuseSparkQuarkOCPMXFusedMoE

Bases: QuarkOCP_MX_MoEMethod

Quark W4A4 method using route-preserving stock AITER experts.

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

class MuseSparkQuarkOCPMXFusedMoE(QuarkOCP_MX_MoEMethod):
    """Quark W4A4 method using route-preserving stock AITER experts."""

    def __init__(
        self,
        moe: FusedMoEConfig,
        weight_quant_key: QuantKey,
        activation_quant_key: QuantKey | None,
    ) -> None:
        validate_mxfp4_deployment(moe)
        super().__init__(moe, weight_quant_key, activation_quant_key)
        if self.ocp_mx_scheme != "w_mxfp4_a_mxfp4":
            raise NotImplementedError(
                "MuseSpark1x MXFP4 on ROCm requires Quark OCP W4A4 with "
                f"dynamic MXFP4 activations; got {self.ocp_mx_scheme!r}."
            )
        if self.mxfp4_backend is not Mxfp4MoeBackend.AITER_MXFP4_MXFP4:
            raise NotImplementedError(
                "MuseSpark1x MXFP4 on ROCm requires the "
                "AITER_MXFP4_MXFP4 backend; got "
                f"{self.mxfp4_backend.value}."
            )
        self.experts_cls = MuseSparkAiterExperts

    def get_fused_moe_quant_config(
        self, layer: RoutedExperts
    ) -> MuseSparkQuarkOCPMXMoEQuantConfig:
        base_config = super().get_fused_moe_quant_config(layer)
        if base_config is None:
            raise RuntimeError("MuseSpark1x MXFP4 quant config is not initialized.")
        post_expert_norm_fn = getattr(layer, "post_expert_norm_fn", None)
        if post_expert_norm_fn is None:
            raise RuntimeError("MuseSpark1x post-expert norm is not configured.")
        return MuseSparkQuarkOCPMXMoEQuantConfig(
            **{
                f.name: getattr(base_config, f.name)
                for f in fields(FusedMoEQuantConfig)
                if f.init
            },
            post_expert_norm_fn=post_expert_norm_fn,
        )

    def process_weights_after_loading(self, layer: RoutedExperts) -> None:
        loaded_slots = getattr(layer, "_quark_expert_loaded_slots", None)
        if loaded_slots is not None:
            _validate_quark_expert_weights(
                loaded_slots, layer.moe_config.num_logical_experts
            )
            del layer._quark_expert_loaded_slots
        super().process_weights_after_loading(layer)

MuseSparkQuarkW8A8Fp8FusedMoE

Bases: QuarkW8A8Fp8MoEMethod

Quark W8A8 FP8 method using route-preserving Triton experts.

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

class MuseSparkQuarkW8A8Fp8FusedMoE(QuarkW8A8Fp8MoEMethod):
    """Quark W8A8 FP8 method using route-preserving Triton experts."""

    def __init__(
        self,
        moe: FusedMoEConfig,
        weight_quant_key: QuantKey,
        activation_quant_key: QuantKey | None,
    ) -> None:
        validate_fp8_deployment(moe)
        super().__init__(moe, weight_quant_key, activation_quant_key)
        if self.fp8_backend is not Fp8MoeBackend.TRITON:
            raise NotImplementedError(
                "MuseSpark1x FP8 on ROCm requires the Triton FP8 MoE backend; "
                f"got {self.fp8_backend.value}."
            )
        self.experts_cls = MuseSparkTritonExperts

    def get_fused_moe_quant_config(
        self, layer: RoutedExperts
    ) -> MuseSparkQuarkFp8MoEQuantConfig:
        base_config = super().get_fused_moe_quant_config(layer)
        post_expert_norm_fn = getattr(layer, "post_expert_norm_fn", None)
        if post_expert_norm_fn is None:
            raise RuntimeError("MuseSpark1x post-expert norm is not configured.")
        return MuseSparkQuarkFp8MoEQuantConfig(
            **{
                f.name: getattr(base_config, f.name)
                for f in fields(FusedMoEQuantConfig)
                if f.init
            },
            post_expert_norm_fn=post_expert_norm_fn,
        )

    def process_weights_after_loading(self, layer: RoutedExperts) -> None:
        loaded_slots = getattr(layer, "_quark_expert_loaded_slots", None)
        if loaded_slots is not None:
            _validate_quark_expert_weights(
                loaded_slots, layer.moe_config.num_logical_experts
            )
            del layer._quark_expert_loaded_slots
        super().process_weights_after_loading(layer)

MuseSparkRoutedExperts

Bases: MuseSparkConsolidatedRoutedExperts

MuseSpark routed experts using ROCm BF16, FP8, or MXFP4 methods.

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

class MuseSparkRoutedExperts(MuseSparkConsolidatedRoutedExperts):
    """MuseSpark routed experts using ROCm BF16, FP8, or MXFP4 methods."""

    def _get_quant_method(
        self,
        prefix: str,
        quant_config: QuantizationConfig | None,
        moe_config: FusedMoEConfig,
    ) -> FusedMoEMethodBase:
        if quant_config is None:
            if _online_mxfp4_applies(prefix):
                return MuseSparkOnlineMXFP4FusedMoE(moe_config)
            return MuseSparkUnquantizedFusedMoE(moe_config)
        if not isinstance(quant_config, QuarkConfig):
            raise NotImplementedError(
                "MuseSpark1x on ROCm supports BF16, Quark FP8, or Quark OCP-MXFP4 "
                f"checkpoints; got {quant_config.get_name()!r}."
            )

        quant_method = quant_config.get_quant_method(self, prefix)
        if isinstance(quant_method, UnquantizedFusedMoEMethod):
            return MuseSparkUnquantizedFusedMoE(moe_config)
        if isinstance(quant_method, QuarkW8A8Fp8MoEMethod):
            return MuseSparkQuarkW8A8Fp8FusedMoE(
                moe_config,
                quant_method.weight_quant_key,
                quant_method.activation_quant_key,
            )
        if not isinstance(quant_method, QuarkOCP_MX_MoEMethod):
            raise NotImplementedError(
                "MuseSpark1x on ROCm requires Quark W8A8 FP8 or OCP-MXFP4 for routed "
                f"experts; got {type(quant_method).__name__}."
            )
        return MuseSparkQuarkOCPMXFusedMoE(
            moe_config,
            quant_method.weight_quant_key,
            quant_method.activation_quant_key,
        )

    def load_weights(
        self, weights: Iterable[tuple[str, torch.Tensor]]
    ) -> Iterable[str]:
        if isinstance(
            self.quant_method, MuseSparkOnlineMXFP4FusedMoE
        ) or not isinstance(
            self.quant_method,
            (MuseSparkQuarkW8A8Fp8FusedMoE, MuseSparkQuarkOCPMXFusedMoE),
        ):
            # Quark's unquantized layers (0 and 61) store per-expert tensors.
            consolidated: list[tuple[str, torch.Tensor]] = []
            per_expert: list[tuple[str, torch.Tensor]] = []
            for name, tensor in weights:
                if _quark_expert_weight_slot(name) is None:
                    consolidated.append((name, tensor))
                else:
                    per_expert.append((name, tensor))
            if consolidated:
                yield from super().load_weights(consolidated)
            if per_expert:
                yield from RoutedExperts.load_weights(self, per_expert)
            return

        loaded_slots = getattr(self, "_quark_expert_loaded_slots", None)
        if loaded_slots is None:
            loaded_slots = set()
            self._quark_expert_loaded_slots = loaded_slots
        tracked_weights = _track_quark_expert_weights(weights, loaded_slots)
        yield from RoutedExperts.load_weights(self, tracked_weights)

MuseSparkTopKWeightAndReduce

Bases: TopKWeightAndReduce

Apply MuseSpark's per-route norm before router weighting and reduction.

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

class MuseSparkTopKWeightAndReduce(mk.TopKWeightAndReduce):
    """Apply MuseSpark's per-route norm before router weighting and reduction."""

    def __init__(
        self,
        post_expert_norm_fn: Callable[..., torch.Tensor],
        *,
        pre_norm_all_reduce: bool,
        take_route_output: Callable[[], torch.Tensor | None] | None = None,
        take_route_expert_map: Callable[[], torch.Tensor | None] | None = None,
    ) -> None:
        self.post_expert_norm_fn = post_expert_norm_fn
        self.pre_norm_all_reduce = pre_norm_all_reduce
        self.take_route_output = take_route_output
        self.take_route_expert_map = take_route_expert_map
        self.reduce_dtype = (
            torch.bfloat16
            if envs.VLLM_ROCM_MUSE_MOE_REDUCE_DTYPE == "bf16"
            else torch.float32
        )
        self.rs_min_tokens = _TOKEN_SHARDED_MIN_TOKENS
        norm = getattr(post_expert_norm_fn, "__self__", None)
        self.fused_norm = (
            norm
            if isinstance(norm, MuseSpark1xRMSNorm)
            and norm.variance_size_override is None
            else None
        )

    @staticmethod
    def _chunk_size(top_k: int, hidden_dim: int) -> int:
        elems_per_token = max(top_k * hidden_dim, 1)
        return max(1, _FINALIZE_CHUNK_ELEMS // elems_per_token)

    def apply(
        self,
        output: torch.Tensor | None,
        fused_expert_output: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
    ) -> torch.Tensor:
        if apply_router_weight_on_input:
            raise NotImplementedError(
                "MuseSpark1x requires router weights after the post-expert norm."
            )

        if self.take_route_output is not None:
            route_output = self.take_route_output()
            if route_output is not None:
                fused_expert_output = route_output
        expert_map = (
            self.take_route_expert_map()
            if self.take_route_expert_map is not None
            else None
        )
        num_tokens, top_k = topk_ids.shape
        hidden_dim = fused_expert_output.shape[-1]
        routed = fused_expert_output.view(num_tokens, top_k, hidden_dim)

        if output is None:
            output = torch.empty(
                (num_tokens, hidden_dim),
                dtype=routed.dtype,
                device=routed.device,
            )
        if output.shape != (num_tokens, hidden_dim):
            raise ValueError(
                "MuseSpark1x MoE output has shape "
                f"{tuple(output.shape)}, expected {(num_tokens, hidden_dim)}."
            )
        if expert_map is not None and (
            self.pre_norm_all_reduce or self.fused_norm is None
        ):
            # Only the fused finalize masks unwritten non-local routes itself.
            routed.masked_fill_((expert_map[topk_ids.long()] < 0).unsqueeze(-1), 0)
            expert_map = None

        chunk_size = self._chunk_size(top_k, hidden_dim)
        for start in range(0, num_tokens, chunk_size):
            end = min(start + chunk_size, num_tokens)
            rows = routed[start:end]
            if (
                self.pre_norm_all_reduce
                and self.fused_norm is not None
                and rows.shape[0] >= self.rs_min_tokens
                and (tp_size := get_tensor_model_parallel_world_size()) > 1
            ):
                self._token_sharded_finalize(
                    rows.to(self.reduce_dtype),
                    topk_weights[start:end],
                    output[start:end],
                    tp_size,
                )
                continue
            if self.pre_norm_all_reduce:
                rows = tensor_model_parallel_all_reduce(rows.to(self.reduce_dtype))

            if self.fused_norm is not None:
                fused_post_expert_finalize(
                    rows,
                    self.fused_norm.weight if self.fused_norm.has_weight else None,
                    topk_weights[start:end],
                    self.fused_norm.variance_epsilon,
                    self.fused_norm.fused_gain_center,
                    self.fused_norm.applies_weight_before_norm,
                    output[start:end],
                    topk_ids=None if expert_map is None else topk_ids[start:end],
                    expert_map=expert_map,
                )
                continue
            normalized = self.post_expert_norm_fn(rows, out_dtype=torch.float32)
            normalized.mul_(topk_weights[start:end].float().unsqueeze(-1))
            output[start:end].copy_(normalized.sum(dim=1).to(output.dtype))

        return output

    def _token_sharded_finalize(
        self,
        rows: torch.Tensor,
        topk_weights: torch.Tensor,
        out: torch.Tensor,
        tp_size: int,
    ) -> None:
        num_tokens, top_k, hidden_dim = rows.shape
        pad = (-num_tokens) % tp_size
        if pad:
            rows = torch.cat([rows, rows.new_zeros((pad, top_k, hidden_dim))])
            topk_weights = torch.cat(
                [topk_weights, topk_weights.new_zeros((pad, top_k))]
            )
        local = tensor_model_parallel_reduce_scatter(rows, dim=0)
        per_rank = local.shape[0]
        rank = get_tensor_model_parallel_rank()
        tw_local = topk_weights[rank * per_rank : (rank + 1) * per_rank]
        out_local = torch.empty(
            (per_rank, hidden_dim), dtype=out.dtype, device=out.device
        )
        norm = self.fused_norm
        assert norm is not None
        fused_post_expert_finalize(
            local,
            norm.weight if norm.has_weight else None,
            tw_local,
            norm.variance_epsilon,
            norm.fused_gain_center,
            norm.applies_weight_before_norm,
            out_local,
        )
        gathered = tensor_model_parallel_all_gather(out_local, dim=0)
        out.copy_(gathered[:num_tokens])

MuseSparkTritonExperts

Bases: _RouteOutputHandoff, TritonExperts

Triton experts that expose unweighted GEMM2 output per route.

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

class MuseSparkTritonExperts(_RouteOutputHandoff, TritonExperts):
    """Triton experts that expose unweighted GEMM2 output per route."""

    def __init__(
        self,
        moe_config: FusedMoEConfig,
        quant_config: FusedMoEQuantConfig,
    ) -> None:
        super().__init__(moe_config, quant_config)
        if not isinstance(
            quant_config,
            (MuseSparkUnquantizedMoEQuantConfig, MuseSparkQuarkFp8MoEQuantConfig),
        ):
            raise TypeError(
                "MuseSparkTritonExperts requires a MuseSpark Triton quant config."
            )
        if quant_config.post_expert_norm_fn is None:
            raise ValueError("MuseSpark1x post-expert norm is not configured.")

    def workspace_shapes(
        self,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: mk.ExpertTokensMetadata | None,
        activation: MoEActivation,
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
        workspace13, workspace2, _ = super().workspace_shapes(
            M,
            N,
            K,
            topk,
            global_num_experts,
            local_num_experts,
            expert_tokens_meta,
            activation,
        )
        return workspace13, workspace2, (M, topk, K)

    def _apply_router_weight_in_gemm2(self, apply_router_weight_on_input: bool) -> bool:
        if apply_router_weight_on_input:
            raise NotImplementedError(
                "MuseSpark1x requires router weights after the post-expert norm."
            )
        return False

    def _finalize_expert_output(
        self,
        expert_output: torch.Tensor,
        output: torch.Tensor,
        topk_ids: torch.Tensor | None = None,
        expert_map: torch.Tensor | None = None,
    ) -> None:
        # Under EP the GEMM skips non-local routes, leaving their rows unwritten.
        self._park_route_output(expert_output, expert_map)

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        quant_config = self.quant_config
        assert isinstance(
            quant_config,
            (MuseSparkUnquantizedMoEQuantConfig, MuseSparkQuarkFp8MoEQuantConfig),
        )
        assert quant_config.post_expert_norm_fn is not None
        return MuseSparkTopKWeightAndReduce(
            quant_config.post_expert_norm_fn,
            pre_norm_all_reduce=self.moe_config.moe_parallel_config.tp_size > 1,
            take_route_output=self._take_route_output,
            take_route_expert_map=self._take_route_expert_map,
        )

MuseSparkUnquantizedFusedMoE

Bases: UnquantizedFusedMoEMethod

BF16 MoE method pinned to the MuseSpark-aware Triton experts.

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

class MuseSparkUnquantizedFusedMoE(UnquantizedFusedMoEMethod):
    """BF16 MoE method pinned to the MuseSpark-aware Triton experts."""

    def __init__(self, moe: FusedMoEConfig) -> None:
        validate_bf16_deployment(moe)
        # Triton-only: skip the generic constructor's backend oracle.
        FusedMoEMethodBase.__init__(self, moe)
        self.unquantized_backend = UnquantizedMoeBackend.TRITON
        self.experts_cls = MuseSparkTritonExperts

    def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
        base_config = super().get_fused_moe_quant_config(layer)
        post_expert_norm_fn = getattr(layer, "post_expert_norm_fn", None)
        if post_expert_norm_fn is None:
            raise RuntimeError("MuseSpark1x post-expert norm is not configured.")
        return MuseSparkUnquantizedMoEQuantConfig(
            **{
                f.name: getattr(base_config, f.name)
                for f in fields(FusedMoEQuantConfig)
                if f.init
            },
            post_expert_norm_fn=post_expert_norm_fn,
        )

_RouteOutputHandoff

Hand the experts' per-route buffer to the finalize instead of copying it.

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

class _RouteOutputHandoff:
    """Hand the experts' per-route buffer to the finalize instead of copying it."""

    _route_output: torch.Tensor | None = None
    _route_expert_map: torch.Tensor | None = None

    def _park_route_output(
        self, route_output: torch.Tensor, expert_map: torch.Tensor | None = None
    ) -> None:
        self._route_output = route_output
        self._route_expert_map = expert_map

    def _take_route_output(self) -> torch.Tensor | None:
        route_output, self._route_output = self._route_output, None
        return route_output

    def _take_route_expert_map(self) -> torch.Tensor | None:
        expert_map, self._route_expert_map = self._route_expert_map, None
        return expert_map

_online_mxfp4_applies(prefix)

Whether this layer's unquantized experts are rounded to MXFP4 on load.

The first and last decoder layers stay BF16, as in the Quark MXFP4 export.

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

def _online_mxfp4_applies(prefix: str) -> bool:
    """Whether this layer's unquantized experts are rounded to MXFP4 on load.

    The first and last decoder layers stay BF16, as in the Quark MXFP4 export.
    """
    if envs.VLLM_ROCM_MUSE_MOE_ONLINE_QUANT != "mxfp4":
        return False
    num_layers = get_current_vllm_config().model_config.hf_text_config.num_hidden_layers
    return extract_layer_index(prefix) not in (0, num_layers - 1)

_quantize_experts_to_mxfp4(weight, experts_per_chunk=32)

Round [experts, rows, cols] weights to packed OCP MXFP4.

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

def _quantize_experts_to_mxfp4(
    weight: torch.Tensor, experts_per_chunk: int = 32
) -> tuple[torch.Tensor, torch.Tensor]:
    """Round ``[experts, rows, cols]`` weights to packed OCP MXFP4."""
    from aiter.ops.triton.quant import dynamic_mxfp4_quant

    num_experts, rows, cols = weight.shape
    packed = torch.empty(
        (num_experts, rows, cols // 2), dtype=torch.uint8, device=weight.device
    )
    scales = torch.empty(
        (num_experts, rows, cols // OCP_MX_BLOCK_SIZE),
        dtype=torch.uint8,
        device=weight.device,
    )
    for start in range(0, num_experts, experts_per_chunk):
        chunk = weight[start : start + experts_per_chunk]
        chunk_packed, chunk_scales = dynamic_mxfp4_quant(chunk.reshape(-1, cols))
        packed[start : start + chunk.shape[0]] = chunk_packed.view(
            chunk.shape[0], rows, -1
        )
        scales[start : start + chunk.shape[0]] = chunk_scales.view(torch.uint8).view(
            chunk.shape[0], rows, -1
        )[..., : cols // OCP_MX_BLOCK_SIZE]
    return packed, scales

_track_quark_expert_weights(weights, loaded)

Track expert tensors across the loader's independently grouped calls.

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

def _track_quark_expert_weights(
    weights: Iterable[tuple[str, torch.Tensor]],
    loaded: set[tuple[str, str, int]],
) -> Iterable[tuple[str, torch.Tensor]]:
    """Track expert tensors across the loader's independently grouped calls."""
    for name, tensor in weights:
        slot = _quark_expert_weight_slot(name)
        if slot is not None:
            if slot in loaded:
                raise ValueError(f"Duplicate MuseSpark1x expert tensor {name!r}.")
            loaded.add(slot)
        yield name, tensor

validate_bf16_deployment(moe_config)

Reject deployments whose BF16 finalization is not implemented.

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

def validate_bf16_deployment(moe_config: FusedMoEConfig) -> None:
    """Reject deployments whose BF16 finalization is not implemented."""
    _validate_tp_post_norm_deployment(moe_config, "BF16")
    if moe_config.moe_backend not in ("auto", "triton"):
        raise NotImplementedError(
            "MuseSpark1x BF16 on ROCm requires the Triton MoE backend; got "
            f"{moe_config.moe_backend!r}."
        )

validate_fp8_deployment(moe_config)

Reject deployments whose FP8 finalization is not implemented.

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

def validate_fp8_deployment(moe_config: FusedMoEConfig) -> None:
    """Reject deployments whose FP8 finalization is not implemented."""
    _validate_tp_post_norm_deployment(moe_config, "FP8")
    if moe_config.moe_backend != "triton":
        raise NotImplementedError(
            "MuseSpark1x FP8 on ROCm requires the Triton MoE backend; got "
            f"{moe_config.moe_backend!r}."
        )

validate_mxfp4_deployment(moe_config)

Reject deployments whose MXFP4 finalization is not implemented.

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

def validate_mxfp4_deployment(moe_config: FusedMoEConfig) -> None:
    """Reject deployments whose MXFP4 finalization is not implemented."""
    _validate_tp_post_norm_deployment(moe_config, "MXFP4")