vllm.model_executor.kernels.linear.mxfp4.emulation ¶

class EmulationMxfp4LinearKernel(MxFp4LinearKernel):
    """Software emulation fallback for OCP MXFP4/MXFP6 (dequant + F.linear)."""

    def __init__(self, config: MxFp4LinearLayerConfig) -> None:
        super().__init__(config)
        if config.activation_quant_key is None:
            # no input Q/DQ for weight-only
            self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = (
                lambda x: x
            )
        else:
            self.quant_dequant_func = _ACTIVATION_QUANT_DEQUANT_FUNCS[
                config.activation_quant_key
            ]

    @classmethod
    def is_supported(
        cls, compute_capability: int | None = None
    ) -> tuple[bool, str | None]:
        return True, None

    @classmethod
    def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
        if config.activation_quant_key not in (
            None,
            kMxfp4Dynamic,
            kMxfp6E3M2Dynamic,
            kMxfp6E2M3Dynamic,
        ):
            return False, "only supports MXFP4 or MXFP6 or unquantized activations"

        if (
            current_platform.is_rocm()
            and current_platform.supports_mx()
            and config.activation_quant_key != kMxfp4Dynamic
        ):
            logger.warning_once(
                "The current platform supports native MXFP4/MXFP6 computation, "
                f"but kernels for activation_quant_key={config.activation_quant_key} "
                f"are not yet integrated in vLLM. Using EmulationMxfp4LinearKernel, "
                "with simulated weight dequantization and activation "
                "QDQ (quantize and dequantize), with the linear "
                "layers computed in high precision."
            )

        if not current_platform.supports_mx():
            logger.warning_once(
                "The current platform does not support native MXFP4 "
                "computation. Using EmulationMxfp4LinearKernel, with simulated weight "
                "dequantization and activation QDQ (quantize and dequantize), with "
                "the linear layers computed in high precision."
            )

        return True, None

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)

    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype)
        qdq_x = self.quant_dequant_func(x)
        return F.linear(qdq_x, dq_w, bias)