vllm-omni/vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Format-specific live-weight validation and reference A16 materialization."""

from __future__ import annotations

import torch
import torch.nn.functional as F
from vllm.config.quantization import QuantSpec
from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod
from vllm.model_executor.layers.quantization.utils import nvfp4_emulation_utils
from vllm.model_executor.layers.quantization.utils.quant_utils import (
    kFp8StaticTensorSym,
    kNvfp4Dynamic,
    kNvfp4Static,
)

_NVFP4_BLOCK_SIZE = 16

# vLLM #49381 replaced the per-format ModelOpt linear methods
# (ModelOptFp8LinearMethod / ModelOptNvFp4LinearMethod, reached through the
# removed ``LinearMethodCls`` attributes) with one generic
# ``ModelOptLinearMethod`` built from a ``(weight, activation)`` QuantSpec pair
# produced by ``resolve()``. Identify the formats the same way upstream does:
# by that exact spec, so the set of accepted methods stays as narrow as it was
# with the class checks (PcPt / PbWo / W4A16 are still not accepted here).
_FP8_W8A8_SPEC = QuantSpec(weight=kFp8StaticTensorSym, activation=kFp8StaticTensorSym)
_NVFP4_W4A4_SPEC = QuantSpec(weight=kNvfp4Static, activation=kNvfp4Dynamic)


def _has_spec(method: object | None, spec: QuantSpec) -> bool:
    return isinstance(method, ModelOptLinearMethod) and getattr(method, "spec", None) == spec


class Cosmos3PrecisionStrategy:
    """Quantization-format boundary behind the shared precision schedule."""

    def accepts(self, method: object | None) -> bool:
        """Return whether this strategy owns a native linear method."""
        raise NotImplementedError

    def validate_before_processing(
        self,
        method: object,
        layer: torch.nn.Module,
        module_name: str,
    ) -> None:
        """Reject native transforms that cannot support live dequantization."""

    def validate_after_processing(
        self,
        layer: torch.nn.Module,
        module_name: str,
    ) -> None:
        """Validate the live native weight representation."""
        raise NotImplementedError

    def materialize(self, layer: torch.nn.Module) -> torch.Tensor:
        """Materialize the live quantized weight as a dense BF16 matrix."""
        raise NotImplementedError

    def apply_high(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None,
    ) -> torch.Tensor:
        """Apply the reference dense A16 path."""
        input_size = int(layer.input_size_per_partition)
        output_size = int(layer.output_size_per_partition)
        weight = self.materialize(layer)
        expected_shape = (output_size, input_size)
        if tuple(weight.shape) != expected_shape:
            raise RuntimeError(f"Dense weight shape {tuple(weight.shape)} does not match {expected_shape}")
        x_2d = x.reshape(-1, x.shape[-1])
        if x_2d.shape[1] != input_size:
            raise ValueError(f"Expected activation width {input_size}, got {x_2d.shape[1]}")
        weight = weight.to(dtype=x.dtype)
        if bias is not None and bias.dtype != x.dtype:
            bias = bias.to(dtype=x.dtype)
        output = F.linear(x_2d, weight, bias)
        return output.view(*x.shape[:-1], output_size)


class Fp8W8A8W8A16Strategy(Cosmos3PrecisionStrategy):
    """Use native ModelOpt W8A8 or reference dense W8A16."""

    def accepts(self, method: object | None) -> bool:
        return _has_spec(method, _FP8_W8A8_SPEC)

    def validate_before_processing(
        self,
        method: object,
        layer: torch.nn.Module,
        module_name: str,
    ) -> None:
        if hasattr(layer, "pre_quant_scale"):
            raise ValueError(f"{module_name} uses unsupported SmoothQuant pre_quant_scale")
        weight_scale = getattr(layer, "weight_scale", None)
        if not isinstance(weight_scale, torch.Tensor) or weight_scale.numel() != 1:
            raise ValueError(f"{module_name} requires serialized tensorwise FP8 weights")
        fp8_kernel = getattr(method, "fp8_linear", None)
        from vllm.model_executor.kernels.linear.scaled_mm.marlin import (
            MarlinFP8ScaledMMLinearKernel,
        )

        if isinstance(fp8_kernel, MarlinFP8ScaledMMLinearKernel):
            raise ValueError(
                f"{module_name} selected Marlin FP8, which repacks the only weight copy; "
                "use a backend that retains canonical FP8 weights"
            )

    def validate_after_processing(
        self,
        layer: torch.nn.Module,
        module_name: str,
    ) -> None:
        weight = getattr(layer, "weight", None)
        scale = getattr(layer, "weight_scale", None)
        if not isinstance(weight, torch.Tensor) or weight.dtype != torch.float8_e4m3fn:
            dtype = getattr(weight, "dtype", None)
            raise TypeError(f"{module_name} requires a canonical FP8 weight, got {dtype}")
        if weight.ndim != 2:
            raise ValueError(f"{module_name} requires a rank-2 FP8 weight, got {tuple(weight.shape)}")
        if not isinstance(scale, torch.Tensor):
            raise ValueError(f"{module_name} has an FP8 weight but no weight_scale")
        if scale.numel() != 1:
            raise ValueError(f"{module_name} requires one live FP8 weight scale")
        input_size = int(layer.input_size_per_partition)
        output_size = int(layer.output_size_per_partition)
        if weight.shape[0] < input_size or weight.shape[1] < output_size:
            raise ValueError(
                f"{module_name} live FP8 weight shape {tuple(weight.shape)} does not cover "
                f"(K, N)=({input_size}, {output_size})"
            )
        _validate_positive_finite_scale(scale, module_name)

    def materialize(self, layer: torch.nn.Module) -> torch.Tensor:
        input_size = int(layer.input_size_per_partition)
        output_size = int(layer.output_size_per_partition)
        weight = layer.weight[:input_size, :output_size]
        scale = layer.weight_scale.reshape(1).to(device=weight.device, dtype=torch.float32)
        return (weight.to(torch.float32) * scale).t().to(torch.bfloat16)


class Nvfp4W4A4W4A16Strategy(Cosmos3PrecisionStrategy):
    """Use native ModelOpt W4A4 or reference dense W4A16."""

    def accepts(self, method: object | None) -> bool:
        return _has_spec(method, _NVFP4_W4A4_SPEC)

    def validate_before_processing(
        self,
        method: object,
        layer: torch.nn.Module,
        module_name: str,
    ) -> None:
        kernel = getattr(method, "kernel", None)
        from vllm.model_executor.kernels.linear.nvfp4.cutlass import (
            CutlassNvFp4LinearKernel,
        )
        from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
            FlashInferCuteDslNvFp4LinearKernel,
            FlashInferCutlassNvFp4LinearKernel,
        )

        compatible_kernels = (
            CutlassNvFp4LinearKernel,
            FlashInferCuteDslNvFp4LinearKernel,
            FlashInferCutlassNvFp4LinearKernel,
        )
        if not isinstance(kernel, compatible_kernels):
            kernel_name = type(kernel).__name__ if kernel is not None else "none"
            raise ValueError(
                f"{module_name} selected unsupported NVFP4 backend {kernel_name}; "
                "live W4A16 dequantization requires a CUTLASS-compatible layout"
            )
        global_scale = _nvfp4_global_scale(layer)
        if global_scale is None or global_scale.numel() != 1:
            raise ValueError(f"{module_name} is missing ModelOpt NVFP4 scales")

    def validate_after_processing(
        self,
        layer: torch.nn.Module,
        module_name: str,
    ) -> None:
        packed = getattr(layer, "weight", None)
        scale = getattr(layer, "weight_scale", None)
        global_scale = getattr(layer, "weight_global_scale", None)
        if not isinstance(packed, torch.Tensor) or packed.dtype != torch.uint8:
            dtype = getattr(packed, "dtype", None)
            raise TypeError(f"{module_name} requires a live packed NVFP4 weight, got {dtype}")
        if not isinstance(scale, torch.Tensor) or not isinstance(global_scale, torch.Tensor):
            raise ValueError(f"{module_name} is missing live NVFP4 scales")
        output_size = int(layer.output_size_per_partition)
        input_size = int(layer.input_size_per_partition)
        if input_size % _NVFP4_BLOCK_SIZE != 0:
            raise ValueError(f"{module_name} input size {input_size} is not divisible by {_NVFP4_BLOCK_SIZE}")
        expected_weight = (output_size, input_size // 2)
        if packed.shape[0] < expected_weight[0] or packed.shape[1] < expected_weight[1]:
            raise ValueError(
                f"{module_name} live NVFP4 weight shape {tuple(packed.shape)} does not cover {expected_weight}"
            )
        _validate_positive_finite_scale(global_scale, module_name)
        lut_handle = nvfp4_emulation_utils.kE2M1ToFloat_handle
        if lut_handle.val.device != packed.device:
            lut_handle.val = lut_handle.val.to(packed.device)

    def materialize(self, layer: torch.nn.Module) -> torch.Tensor:
        output_size = int(layer.output_size_per_partition)
        input_size = int(layer.input_size_per_partition)
        packed = layer.weight[:output_size, : input_size // 2]
        return nvfp4_emulation_utils.dequantize_to_dtype(
            packed,
            layer.weight_scale,
            layer.weight_global_scale,
            torch.bfloat16,
            _NVFP4_BLOCK_SIZE,
            True,
        )


def _validate_positive_finite_scale(scale: torch.Tensor, module_name: str) -> None:
    values = scale.detach().float()
    if not torch.isfinite(values).all() or not (values > 0).all():
        raise ValueError(f"{module_name} has a non-finite or non-positive weight scale")


def _nvfp4_global_scale(layer: torch.nn.Module) -> torch.Tensor | None:
    for name in ("weight_scale_2", "weight_global_scale"):
        value = getattr(layer, name, None)
        if isinstance(value, torch.Tensor):
            return value
    return None