Repository: vllm-project/vllm
File: vllm/model_executor/layers/quantization/mxfp4.py (Python)

Code before:
        w2_weight_scale = torch.nn.Parameter(
            torch.zeros(
                num_experts,
                hidden_size,
                intermediate_size_per_partition_after_pad // mxfp4_block,
                dtype=scale_dtype,
            ),
            requires_grad=False,
        )
        layer.register_parameter("w2_weight_scale", w2_weight_scale)
        set_weight_attrs(w2_weight_scale, extra_weight_attrs)

        w2_bias = torch.nn.Parameter(
            torch.zeros(
                num_experts,
                hidden_size,
                dtype=torch.bfloat16,
            ),
            requires_grad=False,
        )
        layer.register_parameter("w2_bias", w2_bias)
        set_weight_attrs(w2_bias, extra_weight_attrs)

    def process_weights_after_loading(self, layer):
        if self.mxfp4_backend == Mxfp4Backend.MARLIN:
            if self.moe.use_ep:
                raise NotImplementedError("Mxfp4 Marlin doesn't support EP yet")
            prepare_moe_fp4_layer_for_marlin(layer, input_dtype=self.marlin_input_dtype)

            from vllm.model_executor.layers.fused_moe.prepare_finalize import (
                MoEPrepareAndFinalizeNoEP,
            )

            self.moe_quant_config = self.get_fused_moe_quant_config(layer)
            assert self.moe_quant_config is not None
            self.kernel = mk.FusedMoEModularKernel(
                MoEPrepareAndFinalizeNoEP(),
                MarlinExperts(
                    self.moe,
                    self.moe_quant_config,
                    input_dtype=self.marlin_input_dtype,
                ),
                inplace=not self.moe.disable_inplace,
                shared_experts=None,
            )
        elif (
            self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
            or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
        ):
            from flashinfer.fp4_quantization import nvfp4_block_scale_interleave
            from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache

Review comment (suggestion):
this does support EP, we dont need this assert

`MoEPrepareAndFinalizeNoEP` should be named `MoEPrepareAndFinalizeNoDpEp` - I plan to update this soon

Code after:
        set_weight_attrs(w2_weight, extra_weight_attrs)

        w2_weight_scale = torch.nn.Parameter(
            torch.zeros(
                num_experts,
                hidden_size,
                intermediate_size_per_partition_after_pad // mxfp4_block,
                dtype=scale_dtype,
            ),
            requires_grad=False,
        )
        layer.register_parameter("w2_weight_scale", w2_weight_scale)
        set_weight_attrs(w2_weight_scale, extra_weight_attrs)

        w2_bias = torch.nn.Parameter(
            torch.zeros(
                num_experts,
                hidden_size,
                dtype=torch.bfloat16,
            ),
            requires_grad=False,
        )
        layer.register_parameter("w2_bias", w2_bias)
        set_weight_attrs(w2_bias, extra_weight_attrs)

    def process_weights_after_loading(self, layer):
        if self.mxfp4_backend == Mxfp4Backend.MARLIN:
            prepare_moe_fp4_layer_for_marlin(
                layer, input_dtype=get_marlin_input_dtype()
            )

            self.moe_quant_config = self.get_fused_moe_quant_config(layer)
            assert self.moe_quant_config is not None

            prepare_finalize = maybe_make_prepare_finalize(
                moe=self.moe,
                quant_config=self.moe_quant_config,
                routing_tables=layer._maybe_init_expert_routing_tables(),
                allow_new_interface=True,
            )
            assert prepare_finalize is not None

            self.moe_mk = mk.FusedMoEModularKernel(
                prepare_finalize,
                MarlinExperts(
                    self.moe,
                    self.moe_quant_config,
                ),
                inplace=not self.moe.disable_inplace,
                shared_experts=None,
            )