vllm.model_executor.layers.fused_moe.expert_substitution

Constant expert substitution for MoNE-pruned checkpoints.

MoNE (https://arxiv.org/abs/2507.00390) replaces redundant experts with "novices" that return a constant vector. Checkpoints list the replaced experts per decoder layer either in config.approximate_experts, with constants named <experts prefix>.<expert id>.approx_value, or in compression_config.transform_config.expert_substitution, with explicitly named constants. Retained experts keep their logical checkpoint names and are loaded into compact physical rows.

Classes:

• ConstantExpertSubstitution –

  Constant outputs for the substituted experts of one MoE layer.

• SubstitutedRoutedExperts –

  Routed experts that hold weights only for retained experts.

Functions:

• decoder_layer_index –

  Return the only integer component of a module path, if unique.

• get_expert_substitution_spec –

  Parse the substituted experts declared in a model's HF config.

ConstantExpertSubstitution

Bases: Module

Constant outputs for the substituted experts of one MoE layer.

Router outputs keep logical expert IDs. transform_routes maps retained experts to compact physical rows and turns substituted routes into zero-weight routes to physical expert 0, so any decomposed MoE backend can run them unchanged.

Methods:

• physical_expert_id –

  Physical row of a retained expert, or -1 for a substituted one.

• transform_routes –

  Remap routes in place and return the constant experts' output.

Source code in vllm/model_executor/layers/fused_moe/expert_substitution.py

class ConstantExpertSubstitution(nn.Module):
    """Constant outputs for the substituted experts of one MoE layer.

    Router outputs keep logical expert IDs. ``transform_routes`` maps retained
    experts to compact physical rows and turns substituted routes into
    zero-weight routes to physical expert 0, so any decomposed MoE backend can
    run them unchanged.
    """

    def __init__(
        self,
        layer_idx: int,
        num_logical_experts: int,
        substituted_expert_ids: tuple[int, ...],
        hidden_size: int,
        params_dtype: torch.dtype,
    ) -> None:
        super().__init__()
        substituted = set(substituted_expert_ids)
        if not substituted <= set(range(num_logical_experts)):
            raise ValueError(
                f"approximate_experts for layer {layer_idx} must be within "
                f"[0, {num_logical_experts}), got {sorted(substituted)}"
            )
        if len(substituted) == num_logical_experts:
            raise ValueError(
                f"approximate_experts for layer {layer_idx} must retain at least "
                "one expert"
            )
        self.layer_idx = layer_idx
        self.substituted_expert_ids = tuple(sorted(substituted))
        retained = [i for i in range(num_logical_experts) if i not in substituted]
        self.num_compute_experts = len(retained)

        # Retained experts map to their physical row r >= 0, substituted
        # experts to -1 - (their row in `values`).
        routes = [0] * num_logical_experts
        for row, expert_id in enumerate(retained):
            routes[expert_id] = row
        for row, expert_id in enumerate(self.substituted_expert_ids):
            routes[expert_id] = -1 - row
        self._routes = tuple(routes)
        self.register_buffer(
            "expert_substitution_routes",
            torch.tensor(routes, dtype=torch.int32),
            persistent=False,
        )

        # NaN marks values that the checkpoint has not provided yet.
        self.values = nn.Parameter(
            torch.full((len(substituted), hidden_size), torch.nan, dtype=params_dtype),
            requires_grad=False,
        )
        self.values.weight_loader = self.weight_loader

    @property
    def num_logical_experts(self) -> int:
        return len(self._routes)

    def physical_expert_id(self, expert_id: int) -> int:
        """Physical row of a retained expert, or -1 for a substituted one."""
        return max(self._routes[expert_id], -1)

    def weight_loader(
        self, param: nn.Parameter, loaded_weight: torch.Tensor, expert_id: int
    ) -> None:
        if not 0 <= expert_id < self.num_logical_experts or (
            self._routes[expert_id] >= 0
        ):
            raise ValueError(
                f"expert {expert_id} of layer {self.layer_idx} is not listed in "
                "approximate_experts"
            )
        row = param.data[-1 - self._routes[expert_id]]
        if loaded_weight.shape != row.shape:
            raise ValueError(
                f"approx_value for expert {expert_id} of layer {self.layer_idx} has "
                f"shape {tuple(loaded_weight.shape)}, expected {tuple(row.shape)}"
            )
        row.copy_(loaded_weight)

    def validate_loaded_values(self, prefix: str) -> None:
        missing_rows = torch.isnan(self.values).any(dim=1).nonzero().flatten().tolist()
        if missing_rows:
            missing = [self.substituted_expert_ids[row] for row in missing_rows]
            raise ValueError(
                f"{prefix} is missing approx_value for {len(missing)} "
                f"expert(s), first missing expert IDs: {missing[:8]}"
            )

    def transform_routes(
        self,
        hidden_states: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
    ) -> torch.Tensor:
        """Remap routes in place and return the constant experts' output."""
        valid = (topk_ids >= 0) & (topk_ids < self.num_logical_experts)
        routes = self.expert_substitution_routes[
            torch.where(valid, topk_ids, torch.zeros_like(topk_ids)).long()
        ].long()
        substituted = valid & (routes < 0)

        # Sum each token's router weights per substituted expert in FP32; other
        # routes land in a trailing column that is dropped.
        num_values, value_size = self.values.shape
        scales = topk_weights.new_zeros(
            (topk_ids.shape[0], num_values + 1), dtype=torch.float32
        )
        scales.scatter_add_(
            1,
            torch.where(substituted, -1 - routes, num_values),
            topk_weights.float(),
        )
        output = torch.zeros_like(hidden_states)
        output[..., :value_size] = scales[:, :num_values] @ self.values.float()

        topk_weights.masked_fill_(~valid | substituted, 0.0)
        topk_ids.copy_(routes.clamp_min(0).to(topk_ids.dtype))
        return output

physical_expert_id(expert_id)

Physical row of a retained expert, or -1 for a substituted one.

Source code in vllm/model_executor/layers/fused_moe/expert_substitution.py

def physical_expert_id(self, expert_id: int) -> int:
    """Physical row of a retained expert, or -1 for a substituted one."""
    return max(self._routes[expert_id], -1)

transform_routes(hidden_states, topk_weights, topk_ids)

Remap routes in place and return the constant experts' output.

Source code in vllm/model_executor/layers/fused_moe/expert_substitution.py

def transform_routes(
    self,
    hidden_states: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
) -> torch.Tensor:
    """Remap routes in place and return the constant experts' output."""
    valid = (topk_ids >= 0) & (topk_ids < self.num_logical_experts)
    routes = self.expert_substitution_routes[
        torch.where(valid, topk_ids, torch.zeros_like(topk_ids)).long()
    ].long()
    substituted = valid & (routes < 0)

    # Sum each token's router weights per substituted expert in FP32; other
    # routes land in a trailing column that is dropped.
    num_values, value_size = self.values.shape
    scales = topk_weights.new_zeros(
        (topk_ids.shape[0], num_values + 1), dtype=torch.float32
    )
    scales.scatter_add_(
        1,
        torch.where(substituted, -1 - routes, num_values),
        topk_weights.float(),
    )
    output = torch.zeros_like(hidden_states)
    output[..., :value_size] = scales[:, :num_values] @ self.values.float()

    topk_weights.masked_fill_(~valid | substituted, 0.0)
    topk_ids.copy_(routes.clamp_min(0).to(topk_ids.dtype))
    return output

SubstitutedRoutedExperts

Bases: RoutedExperts

Routed experts that hold weights only for retained experts.

Source code in vllm/model_executor/layers/fused_moe/expert_substitution.py

class SubstitutedRoutedExperts(RoutedExperts):
    """Routed experts that hold weights only for retained experts."""

    def __init__(
        self, *args: Any, expert_substitution: ConstantExpertSubstitution, **kwargs: Any
    ):
        super().__init__(*args, **kwargs)
        self.expert_substitution = expert_substitution
        moe_config = self.moe_config
        parallel_config = moe_config.moe_parallel_config
        unsupported = [
            name
            for name, enabled in (
                (
                    "quantized MoE weights",
                    not isinstance(self.quant_method, UnquantizedFusedMoEMethod),
                ),
                ("monolithic MoE backends", self.quant_method.is_monolithic),
                ("MoE LoRA", moe_config.is_lora_enabled),
                (
                    "EPLB",
                    parallel_config.enable_eplb
                    or moe_config.num_experts
                    != expert_substitution.num_compute_experts,
                ),
                (
                    "fused shared experts",
                    self.expert_map_manager.num_fused_shared_experts > 0,
                ),
                ("expert parallelism", parallel_config.use_ep),
                ("data parallelism", parallel_config.dp_size > 1),
                ("prefill context parallelism", parallel_config.pcp_size > 1),
                ("sequence parallelism", parallel_config.is_sequence_parallel),
                ("router weights on expert inputs", self.apply_router_weight_on_input),
            )
            if enabled
        ]
        if unsupported:
            raise NotImplementedError(
                "expert substitution does not support: " + ", ".join(unsupported)
            )

    def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int:
        # Checkpoints name retained experts by their logical ID.
        physical_id = self.expert_substitution.physical_expert_id(expert_id)
        if physical_id < 0:
            return -1
        return super()._map_global_expert_id_to_local_expert_id(physical_id)

    def get_expert_mapping(
        self,
        ckpt_gate_proj_name: str | None = None,
        ckpt_down_proj_name: str | None = None,
        ckpt_up_proj_name: str | None = None,
        include_fused: bool = False,
    ) -> list[tuple[str, str, int, str]]:
        # Fused checkpoint tensors would cover every logical expert at once.
        return self.build_expert_params_mapping(
            ckpt_gate_proj_name or self.ckpt_gate_proj_name,
            ckpt_down_proj_name or self.ckpt_down_proj_name,
            ckpt_up_proj_name or self.ckpt_up_proj_name,
            num_experts=self.expert_substitution.num_logical_experts,
            routed_experts_prefix="",
            lora_base_layer_prefix=self.lora_base_layer_prefix,
        )

    def forward_modular(
        self,
        x: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        shared_experts: "SharedExperts | None" = None,
        shared_experts_input: torch.Tensor | None = None,
    ) -> torch.Tensor:
        constant_output = self.expert_substitution.transform_routes(
            x, topk_weights, topk_ids
        )
        output = super().forward_modular(
            x, topk_weights, topk_ids, shared_experts, shared_experts_input
        )
        if isinstance(output, UnfinalizedMoEOutput):
            raise NotImplementedError(
                "expert substitution does not support deferred MoE finalize"
            )
        # The constant output is replicated across TP ranks, whose partial
        # outputs are all-reduced later, so add it on one rank only.
        if self.moe_config.tp_rank == 0:
            output.add_(constant_output[..., : output.shape[-1]])
        return output

decoder_layer_index(path)

Return the only integer component of a module path, if unique.

Source code in vllm/model_executor/layers/fused_moe/expert_substitution.py

def decoder_layer_index(path: str) -> int | None:
    """Return the only integer component of a module path, if unique."""
    indices = [int(part) for part in path.split(".") if part.isdigit()]
    return indices[0] if len(indices) == 1 else None

get_expert_substitution_spec(model_config)

Parse the substituted experts declared in a model's HF config.

Source code in vllm/model_executor/layers/fused_moe/expert_substitution.py

def get_expert_substitution_spec(model_config: Any) -> ExpertSubstitutionSpec | None:
    """Parse the substituted experts declared in a model's HF config."""
    hf_configs = [
        getattr(model_config, name, None) for name in ("hf_text_config", "hf_config")
    ]
    approximate = None
    for config in hf_configs:
        approximate = getattr(config, "approximate_experts", None)
        if approximate is not None:
            break
    compression_config = getattr(hf_configs[1], "compression_config", None)
    transform_config = (
        compression_config.get("transform_config")
        if isinstance(compression_config, Mapping)
        else None
    )
    declared = (
        transform_config.get("expert_substitution")
        if isinstance(transform_config, Mapping)
        else None
    )
    if approximate is not None and declared is not None:
        raise ValueError(
            "config declares both approximate_experts and "
            "compression_config.transform_config.expert_substitution"
        )
    if approximate is not None:
        spec = _parse_approximate_experts(approximate)
    elif declared is not None:
        spec = _parse_expert_substitution(declared)
    else:
        return None
    return spec if spec.experts else None