quantized_linear.py

"""Portable Mage-VL FP8 and NVFP4 linear modules.

This file is loaded as Hugging Face remote code.  The checkpoint config
selects one format before the state dictionary is materialized, so the
original BF16 language projection weights are never allocated or requested.
"""

from __future__ import annotations

import hashlib
import os
from functools import lru_cache
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F
from torch import nn


LANGUAGE_PROJECTION_ROLES = {
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "gate_proj",
    "up_proj",
    "down_proj",
}
_SMALLM_SOURCE_ROOT: Path | None = None


def _configure_smallm_source(model_name_or_path: str) -> None:
    """Resolve native sources from a local repo or Hugging Face snapshot."""

    global _SMALLM_SOURCE_ROOT
    candidate = Path(model_name_or_path).expanduser()
    local_root = candidate / "native" / "smallm_gemv"
    required = ("smallm_gemv.cpp", "smallm_gemv.cu", "smallm_gemv.h")
    if all((local_root / name).is_file() for name in required):
        _SMALLM_SOURCE_ROOT = local_root.resolve()
        return
    if not model_name_or_path:
        raise RuntimeError("Mage-VL small-M source repository is unspecified")

    from transformers.utils.hub import cached_file

    resolved = [
        Path(
            cached_file(
                model_name_or_path,
                f"native/smallm_gemv/{name}",
            )
        )
        for name in required
    ]
    parents = {path.parent.resolve() for path in resolved}
    if len(parents) != 1:
        raise RuntimeError(
            "small-M native sources resolved to different directories: "
            f"{sorted(str(value) for value in parents)}"
        )
    _SMALLM_SOURCE_ROOT = parents.pop()


def _smallm_source_root() -> Path:
    if _SMALLM_SOURCE_ROOT is None:
        raise RuntimeError(
            "small-M native sources were not configured during model setup"
        )
    return _SMALLM_SOURCE_ROOT


@lru_cache(maxsize=1)
def _load_smallm_extension() -> Any:
    from torch.utils.cpp_extension import load

    source_root = _smallm_source_root()
    configured_build = os.environ.get("MAGE_VL_SMALLM_BUILD_DIR")
    build_root = (
        Path(configured_build).expanduser().resolve()
        if configured_build
        else Path(__file__).resolve().parent / ".native_build"
    )
    build_root.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
    os.environ.setdefault("MAX_JOBS", "4")
    source_hash = hashlib.sha256(
        b"".join(
            (source_root / name).read_bytes()
            for name in (
                "smallm_gemv.cpp",
                "smallm_gemv.cu",
                "smallm_gemv.h",
            )
        )
    ).hexdigest()[:12]
    return load(
        name=f"mage_vl_smallm_gemv_{source_hash}",
        sources=[
            str(source_root / "smallm_gemv.cpp"),
            str(source_root / "smallm_gemv.cu"),
        ],
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
        extra_include_paths=[str(source_root)],
        build_directory=str(build_root),
        with_cuda=True,
        verbose=False,
        is_python_module=True,
    )


def _smallm_nvfp4_linear(
    value: torch.Tensor,
    *,
    qdata: torch.Tensor,
    weight_block_scale: torch.Tensor,
    weight_scale: torch.Tensor,
    bias: torch.Tensor | None,
) -> torch.Tensor:
    return _load_smallm_extension().linear(
        value.contiguous(),
        qdata,
        weight_block_scale,
        weight_scale,
        bias,
    )


def _resolve_parent(root: nn.Module, module_name: str) -> tuple[nn.Module, str]:
    parent_name, separator, leaf = module_name.rpartition(".")
    if not separator:
        return root, module_name
    return root.get_submodule(parent_name), leaf


def _empty_like_source(
    source: nn.Linear,
    shape: tuple[int, ...],
    dtype: torch.dtype,
) -> torch.Tensor:
    return torch.empty(shape, dtype=dtype, device=source.weight.device)


class MageVLScaledFP8Linear(nn.Module):
    """W8A8 prefill with optional resident-weight W8A16 small-M decode."""

    def __init__(
        self,
        source: nn.Linear,
        *,
        role: str,
        smallm_backend: str,
        smallm_threshold: int,
        smallm_roles: set[str],
    ) -> None:
        super().__init__()
        if smallm_backend not in {"off", "w8a16_gemv"}:
            raise ValueError(f"unsupported FP8 small-M backend: {smallm_backend}")
        if smallm_threshold <= 0:
            raise ValueError("FP8 small-M threshold must be positive")
        self.in_features = int(source.in_features)
        self.out_features = int(source.out_features)
        self.role = role
        self.smallm_backend = (
            smallm_backend if role in smallm_roles else "off"
        )
        self.smallm_threshold = int(smallm_threshold)
        self.register_buffer(
            "qdata",
            _empty_like_source(
                source,
                (self.out_features, self.in_features),
                torch.float8_e4m3fn,
            ),
            persistent=True,
        )
        self.register_buffer(
            "weight_scale",
            _empty_like_source(source, (), torch.float32),
            persistent=True,
        )
        if source.bias is None:
            self.bias_bf16 = None
        else:
            self.register_buffer(
                "bias_bf16",
                _empty_like_source(
                    source,
                    (self.out_features,),
                    torch.bfloat16,
                ),
                persistent=True,
            )

    def _weight_quantized_tensor(self) -> Any:
        from comfy_kitchen.tensor import QuantizedTensor, TensorCoreFP8Layout

        params = TensorCoreFP8Layout.Params(
            scale=self.weight_scale,
            orig_dtype=torch.bfloat16,
            orig_shape=(self.out_features, self.in_features),
        )
        return QuantizedTensor(
            self.qdata,
            "TensorCoreFP8Layout",
            params,
        )

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        from comfy_kitchen.tensor import QuantizedTensor

        input_shape = tuple(value.shape)
        flattened = value.reshape(-1, input_shape[-1]).contiguous()
        if (
            self.smallm_backend == "w8a16_gemv"
            and flattened.shape[0] <= self.smallm_threshold
        ):
            from .fp8_decode_runtime import smallm_fp8_linear

            output = smallm_fp8_linear(
                flattened,
                qdata=self.qdata,
                weight_scale=self.weight_scale,
                bias=self.bias_bf16,
            )
            return output.reshape(*input_shape[:-1], self.out_features)
        quantized_input = QuantizedTensor.from_float(
            flattened,
            "TensorCoreFP8Layout",
        )
        output = F.linear(
            quantized_input,
            self._weight_quantized_tensor(),
            None,
        )
        if self.bias_bf16 is not None:
            output = output + self.bias_bf16
        return output.reshape(*input_shape[:-1], self.out_features)


class MageVLNVFP4Linear(nn.Module):
    """Native W4A4 prefill with optional packed-weight W4A16 small-M decode."""

    def __init__(
        self,
        source: nn.Linear,
        *,
        role: str,
        smallm_backend: str,
        smallm_threshold: int,
        smallm_roles: set[str],
    ) -> None:
        super().__init__()
        self.in_features = int(source.in_features)
        self.out_features = int(source.out_features)
        self.role = role
        self.smallm_backend = (
            smallm_backend if role in smallm_roles else "off"
        )
        self.smallm_threshold = int(smallm_threshold)
        self.register_buffer(
            "qdata",
            _empty_like_source(
                source,
                (self.out_features, self.in_features // 2),
                torch.uint8,
            ),
            persistent=True,
        )
        self.register_buffer(
            "weight_scale",
            _empty_like_source(source, (), torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "weight_block_scale",
            _empty_like_source(
                source,
                (self.out_features, self.in_features // 16),
                torch.float8_e4m3fn,
            ),
            persistent=True,
        )
        if source.bias is None:
            self.bias_bf16 = None
        else:
            self.register_buffer(
                "bias_bf16",
                _empty_like_source(
                    source,
                    (self.out_features,),
                    torch.bfloat16,
                ),
                persistent=True,
            )

    def _weight_quantized_tensor(self) -> Any:
        from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout

        params = TensorCoreNVFP4Layout.Params(
            scale=self.weight_scale,
            orig_dtype=torch.bfloat16,
            orig_shape=(self.out_features, self.in_features),
            block_scale=self.weight_block_scale,
        )
        return QuantizedTensor(
            self.qdata,
            "TensorCoreNVFP4Layout",
            params,
        )

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        from comfy_kitchen.tensor import QuantizedTensor

        input_shape = tuple(value.shape)
        flattened = value.reshape(-1, input_shape[-1]).contiguous()
        if (
            self.smallm_backend == "w4a16_gemv"
            and flattened.shape[0] <= self.smallm_threshold
        ):
            output = _smallm_nvfp4_linear(
                flattened,
                qdata=self.qdata,
                weight_block_scale=self.weight_block_scale,
                weight_scale=self.weight_scale,
                bias=self.bias_bf16,
            )
            return output.reshape(*input_shape[:-1], self.out_features)

        quantized_input = QuantizedTensor.from_float(
            flattened,
            "TensorCoreNVFP4Layout",
        )
        output = F.linear(
            quantized_input,
            self._weight_quantized_tensor(),
            None,
        )
        if self.bias_bf16 is not None:
            output = output + self.bias_bf16
        return output.reshape(*input_shape[:-1], self.out_features)


def _smallm_policy(
    quantization: dict[str, Any],
    *,
    format_name: str,
) -> tuple[str, int, set[str]]:
    backend = os.environ.get(
        "MAGE_VL_SMALLM_BACKEND",
        str(quantization.get("smallm_backend", "off")),
    )
    supported_backend = (
        "w8a16_gemv"
        if format_name == "scaled_fp8_w8a8"
        else "w4a16_gemv"
    )
    if backend not in {"off", supported_backend}:
        raise ValueError(f"unsupported MAGE_VL_SMALLM_BACKEND: {backend}")
    threshold = int(
        os.environ.get(
            "MAGE_VL_SMALLM_THRESHOLD",
            str(quantization.get("smallm_threshold", 1)),
        )
    )
    if threshold <= 0:
        raise ValueError("MAGE_VL_SMALLM_THRESHOLD must be positive")
    configured_roles = quantization.get(
        "smallm_roles",
        sorted(LANGUAGE_PROJECTION_ROLES),
    )
    role_text = os.environ.get(
        "MAGE_VL_SMALLM_ROLES",
        ",".join(str(value) for value in configured_roles),
    )
    roles = {value.strip() for value in role_text.split(",") if value.strip()}
    if not roles <= LANGUAGE_PROJECTION_ROLES:
        raise ValueError(
            f"invalid small-M roles: {sorted(roles - LANGUAGE_PROJECTION_ROLES)}"
        )
    return backend, threshold, roles


def _environment_flag(name: str, default: bool) -> bool:
    value = os.environ.get(name)
    if value is None:
        return bool(default)
    normalized = value.strip().lower()
    if normalized in {"1", "true", "yes", "on"}:
        return True
    if normalized in {"0", "false", "no", "off"}:
        return False
    raise ValueError(f"{name} must be one of 1/0, true/false, yes/no, or on/off")


def apply_mage_vl_quantization(
    model: nn.Module,
    config: Any,
) -> None:
    """Replace all 252 Qwen language projections before checkpoint loading."""

    quantization = getattr(config, "mage_vl_quantization", None)
    if not quantization:
        return
    if not isinstance(quantization, dict):
        raise TypeError("mage_vl_quantization must be a dictionary")
    format_name = quantization.get("format")
    if format_name not in {"scaled_fp8_w8a8", "native_nvfp4_w4a4"}:
        raise ValueError(f"unsupported Mage-VL quantization: {format_name}")

    backend, threshold, smallm_roles = _smallm_policy(
        quantization,
        format_name=format_name,
    )
    fused_gate_up = False
    fused_qkv = False
    fused_gate_up_threshold = 1
    fused_qkv_threshold = 1
    if format_name == "scaled_fp8_w8a8":
        fused_gate_up = _environment_flag(
            "MAGE_VL_FP8_FUSED_GATE_UP",
            bool(quantization.get("fused_gate_up", False)),
        )
        fused_qkv = _environment_flag(
            "MAGE_VL_FP8_FUSED_QKV",
            bool(quantization.get("fused_qkv", False)),
        )
        fused_gate_up_threshold = int(
            os.environ.get(
                "MAGE_VL_FP8_FUSED_GATE_UP_THRESHOLD",
                str(quantization.get("fused_gate_up_threshold", 1)),
            )
        )
        fused_qkv_threshold = int(
            os.environ.get(
                "MAGE_VL_FP8_FUSED_QKV_THRESHOLD",
                str(quantization.get("fused_qkv_threshold", 1)),
            )
        )
        if fused_gate_up_threshold <= 0 or fused_qkv_threshold <= 0:
            raise ValueError("FP8 fusion thresholds must be positive")
        if backend != "off" or fused_gate_up or fused_qkv:
            from .fp8_decode_runtime import configure_fp8_decode_sources

            configure_fp8_decode_sources(
                str(getattr(config, "_name_or_path", ""))
            )
    if format_name == "native_nvfp4_w4a4" and backend != "off":
        _configure_smallm_source(str(getattr(config, "_name_or_path", "")))
    installed = []
    for layer in range(36):
        for branch, roles in (
            ("self_attn", ("q_proj", "k_proj", "v_proj", "o_proj")),
            ("mlp", ("gate_proj", "up_proj", "down_proj")),
        ):
            for role in roles:
                name = f"language_model.layers.{layer}.{branch}.{role}"
                parent, leaf = _resolve_parent(model, name)
                source = getattr(parent, leaf)
                if not isinstance(source, nn.Linear):
                    raise TypeError(
                        f"{name}: expected nn.Linear, got "
                        f"{type(source).__name__}"
                    )
                if format_name == "scaled_fp8_w8a8":
                    replacement = MageVLScaledFP8Linear(
                        source,
                        role=role,
                        smallm_backend=backend,
                        smallm_threshold=threshold,
                        smallm_roles=smallm_roles,
                    )
                else:
                    replacement = MageVLNVFP4Linear(
                        source,
                        role=role,
                        smallm_backend=backend,
                        smallm_threshold=threshold,
                        smallm_roles=smallm_roles,
                    )
                setattr(parent, leaf, replacement)
                installed.append(name)
    if len(installed) != 252:
        raise RuntimeError(
            f"expected 252 quantized language projections, got {len(installed)}"
        )
    runtime_manifest = {
        "format": format_name,
        "smallm_backend": backend,
        "smallm_threshold": threshold,
        "smallm_roles": sorted(smallm_roles),
        "fused_gate_up": fused_gate_up,
        "fused_gate_up_threshold": fused_gate_up_threshold,
        "fused_qkv": fused_qkv,
        "fused_qkv_threshold": fused_qkv_threshold,
    }
    if format_name == "scaled_fp8_w8a8":
        from .fp8_decode_runtime import (
            install_fp8_fused_gate_up,
            install_fp8_fused_qkv,
        )

        if fused_gate_up:
            runtime_manifest["gate_up_install"] = install_fp8_fused_gate_up(
                model,
                threshold=fused_gate_up_threshold,
            )
        if fused_qkv:
            runtime_manifest["qkv_install"] = install_fp8_fused_qkv(
                model,
                threshold=fused_qkv_threshold,
            )
    object.__setattr__(model, "_mage_vl_runtime_manifest", runtime_manifest)


__all__ = [
    "MageVLNVFP4Linear",
    "MageVLScaledFP8Linear",
    "apply_mage_vl_quantization",
]