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

Code before:
            layer.w2_input_scale = torch.nn.Parameter(
                layer.w2_input_scale.max().to(torch.float32), requires_grad=False
            )

            from triton_kernels.numerics import InFlexData

            lhs_data13 = InFlexData(scale=layer.w13_input_scale)
            lhs_data2 = InFlexData(scale=layer.w2_input_scale)

            self.w13_precision_config = PrecisionConfig(
                weight_scale=w13_scale,
                flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13),
            )
            self.w2_precision_config = PrecisionConfig(
                weight_scale=w2_scale,
                flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2),
            )

    def get_fused_moe_quant_config(
        self, layer: torch.nn.Module
    ) -> FusedMoEQuantConfig | None:
        w1_scale = self.w13_precision_config
        w2_scale = self.w2_precision_config

        # TODO: how to set scale?
        return mxfp4_w4a4_moe_quant_config(
            w1_bias=layer.w13_bias,
            w2_bias=layer.w2_bias,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
        )

    @property
    def is_monolithic(self) -> bool:
        return True

    def apply_monolithic(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        router_logits: torch.Tensor,
        expert_map: torch.Tensor | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        if layer.enable_eplb:
            raise NotImplementedError(
                "EPLB not supported for `QuarkW4MXFp4MoEMethod_OSS` yet."
            )

        from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (  # noqa: E501
            triton_kernel_moe_oss_forward,
        )

Review comment (suggestion):
as commented elsewhere, suggest return `mxfp4_w4a8_moe_quant_config`

Code after:
            )

            from triton_kernels.numerics import InFlexData

            lhs_data13 = InFlexData(scale=layer.w13_input_scale)
            lhs_data2 = InFlexData(scale=layer.w2_input_scale)

            self.w13_precision_config = PrecisionConfig(
                weight_scale=w13_scale,
                flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13),
            )

            self.w2_precision_config = PrecisionConfig(
                weight_scale=w2_scale,
                flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2),
            )

    def get_fused_moe_quant_config(
        self, layer: torch.nn.Module
    ) -> FusedMoEQuantConfig | None:
        return mxfp4_w4a8_moe_quant_config(
            w1_scale=self.w13_precision_config,
            w2_scale=self.w2_precision_config,
            a1_scale=layer.w13_input_scale,
            a2_scale=layer.w2_input_scale,
            w1_bias=layer.w13_bias,
            w2_bias=layer.w2_bias,
            block_shape=None,
        )

    @property
    def is_monolithic(self) -> bool:
        return True

    def apply_monolithic(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        router_logits: torch.Tensor,
        expert_map: torch.Tensor | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        if layer.enable_eplb:
            raise NotImplementedError(
                "EPLB not supported for `QuarkW4MXFp4MoEMethod_OSS` yet."
            )

        from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (  # noqa: E501
            triton_kernel_moe_forward,
        )

        return triton_kernel_moe_forward(