fp8_decode_runtime.py

"""Optional native FP8 decode paths for the standalone Mage-VL package.

The checkpoint remains resident E4M3 FP8.  Large-M work keeps the portable
dynamic W8A8 path, while M=1 decode can use W8A16 kernels over the same stored
weights.  Every patch has an exact feature-off and larger-M fallback.
"""

from __future__ import annotations

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

import torch
from torch import nn


_SOURCE_FILES = {
    "smallm_fp8_gemv": (
        "smallm_fp8_gemv.cpp",
        "smallm_fp8_gemv.cu",
        "smallm_fp8_gemv.h",
    ),
    "fused_fp8_gate_up": (
        "fused_fp8_gate_up.cpp",
        "fused_fp8_gate_up.cu",
        "fused_fp8_gate_up.h",
    ),
    "fused_fp8_qkv": (
        "fused_fp8_qkv.cpp",
        "fused_fp8_qkv.cu",
        "fused_fp8_qkv.h",
    ),
}
_SOURCE_ROOTS: dict[str, Path] = {}


def configure_fp8_decode_sources(model_name_or_path: str) -> None:
    """Resolve every native source set from a local or Hub model snapshot."""

    candidate = Path(model_name_or_path).expanduser()
    for feature, filenames in _SOURCE_FILES.items():
        local_root = candidate / "native" / feature
        if all((local_root / name).is_file() for name in filenames):
            _SOURCE_ROOTS[feature] = local_root.resolve()
            continue
        if not model_name_or_path:
            raise RuntimeError(
                f"Mage-VL native source repository is unspecified for {feature}"
            )
        from transformers.utils.hub import cached_file

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


def _source_root(feature: str) -> Path:
    try:
        return _SOURCE_ROOTS[feature]
    except KeyError as exc:
        raise RuntimeError(
            f"{feature} sources were not configured during model setup"
        ) from exc


def fp8_decode_source_manifest() -> dict[str, str]:
    result = {}
    for feature in sorted(_SOURCE_FILES):
        root = _source_root(feature)
        for filename in _SOURCE_FILES[feature]:
            result[f"native/{feature}/{filename}"] = hashlib.sha256(
                (root / filename).read_bytes()
            ).hexdigest()
    return result


def _build_root(feature: str) -> Path:
    configured = os.environ.get("MAGE_VL_NATIVE_BUILD_DIR")
    if not configured and feature == "smallm_fp8_gemv":
        configured = os.environ.get("MAGE_VL_SMALLM_BUILD_DIR")
    base = (
        Path(configured).expanduser().resolve()
        if configured
        else Path(__file__).resolve().parent / ".native_build"
    )
    root = base / feature
    root.mkdir(parents=True, exist_ok=True)
    return root


@lru_cache(maxsize=None)
def _load_extension(feature: str) -> Any:
    from torch.utils.cpp_extension import load

    root = _source_root(feature)
    filenames = _SOURCE_FILES[feature]
    source_hash = hashlib.sha256(
        b"".join((root / name).read_bytes() for name in filenames)
    ).hexdigest()[:12]
    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
    os.environ.setdefault("MAX_JOBS", "4")
    return load(
        name=f"mage_vl_{feature}_{source_hash}",
        sources=[str(root / name) for name in filenames if name.endswith((".cpp", ".cu"))],
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
        extra_include_paths=[str(root)],
        build_directory=str(_build_root(feature)),
        with_cuda=True,
        verbose=False,
        is_python_module=True,
    )


def smallm_fp8_linear(
    value: torch.Tensor,
    *,
    qdata: torch.Tensor,
    weight_scale: torch.Tensor,
    bias: torch.Tensor | None,
) -> torch.Tensor:
    return _load_extension("smallm_fp8_gemv").linear(
        value.contiguous(), qdata, weight_scale, bias
    )


def fused_fp8_gate_up_silu(
    value: torch.Tensor,
    *,
    gate_qdata: torch.Tensor,
    gate_scale: torch.Tensor,
    up_qdata: torch.Tensor,
    up_scale: torch.Tensor,
) -> torch.Tensor:
    return _load_extension("fused_fp8_gate_up").gate_up_silu(
        value.contiguous(),
        gate_qdata,
        gate_scale,
        up_qdata,
        up_scale,
    )


def fused_fp8_qkv(
    value: torch.Tensor,
    *,
    q_qdata: torch.Tensor,
    q_scale: torch.Tensor,
    k_qdata: torch.Tensor,
    k_scale: torch.Tensor,
    v_qdata: torch.Tensor,
    v_scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    outputs = _load_extension("fused_fp8_qkv").qkv(
        value.contiguous(),
        q_qdata,
        q_scale,
        k_qdata,
        k_scale,
        v_qdata,
        v_scale,
    )
    return outputs[0], outputs[1], outputs[2]


def _is_fp8_linear(module: nn.Module) -> bool:
    return (
        module.__class__.__name__ == "MageVLScaledFP8Linear"
        and hasattr(module, "qdata")
        and hasattr(module, "weight_scale")
        and hasattr(module, "bias_bf16")
    )


def install_fp8_fused_gate_up(
    model: nn.Module,
    *,
    threshold: int = 1,
) -> dict[str, Any]:
    """Patch Qwen MLP forwards with the accepted fused M=1 implementation."""

    if threshold <= 0:
        raise ValueError("fused gate/up threshold must be positive")
    installed = []
    for layer in range(36):
        name = f"language_model.layers.{layer}.mlp"
        mlp = model.get_submodule(name)
        if hasattr(mlp, "_fp8_fused_gate_up_original_forward"):
            raise RuntimeError(f"{name}: fused gate/up is already installed")
        for role in ("gate_proj", "up_proj", "down_proj"):
            module = getattr(mlp, role, None)
            if not _is_fp8_linear(module):
                raise TypeError(
                    f"{name}.{role}: expected MageVLScaledFP8Linear, "
                    f"got {type(module).__name__}"
                )
        if mlp.gate_proj.bias_bf16 is not None or mlp.up_proj.bias_bf16 is not None:
            raise RuntimeError(f"{name}: fused gate/up requires bias-free projections")
        original_forward = mlp.forward
        object.__setattr__(mlp, "_fp8_fused_gate_up_original_forward", original_forward)
        object.__setattr__(mlp, "_fp8_fused_gate_up_threshold", int(threshold))

        def fused_forward(self, value: torch.Tensor):
            input_shape = tuple(value.shape)
            flattened = value.reshape(-1, input_shape[-1]).contiguous()
            if flattened.shape[0] <= self._fp8_fused_gate_up_threshold:
                intermediate = fused_fp8_gate_up_silu(
                    flattened,
                    gate_qdata=self.gate_proj.qdata,
                    gate_scale=self.gate_proj.weight_scale,
                    up_qdata=self.up_proj.qdata,
                    up_scale=self.up_proj.weight_scale,
                ).reshape(*input_shape[:-1], self.gate_proj.out_features)
                return self.down_proj(intermediate)
            return self._fp8_fused_gate_up_original_forward(value)

        object.__setattr__(mlp, "forward", types.MethodType(fused_forward, mlp))
        installed.append(name)
    return {
        "feature": "fused_fp8_w8a16_gate_up_silu",
        "threshold": int(threshold),
        "installed_module_count": len(installed),
        "fallback": "original_qwen_mlp_forward",
        "stored_weight_payload": "unchanged",
    }


def restore_fp8_fused_gate_up(model: nn.Module) -> dict[str, Any]:
    restored = []
    for layer in range(36):
        name = f"language_model.layers.{layer}.mlp"
        mlp = model.get_submodule(name)
        original = getattr(mlp, "_fp8_fused_gate_up_original_forward", None)
        if original is None:
            continue
        object.__setattr__(mlp, "forward", original)
        object.__delattr__(mlp, "_fp8_fused_gate_up_original_forward")
        object.__delattr__(mlp, "_fp8_fused_gate_up_threshold")
        restored.append(name)
    return {"feature": "fused_fp8_w8a16_gate_up_silu", "restored": len(restored)}


class _FP8FusedQKVCoordinator:
    """Preserve the upstream Q->K->V call order while grouping M=1 dispatch."""

    def __init__(self, q_proj: nn.Module, k_proj: nn.Module, v_proj: nn.Module, *, threshold: int) -> None:
        self.projections = {"q": q_proj, "k": k_proj, "v": v_proj}
        self.original_forwards = {
            role: module.forward for role, module in self.projections.items()
        }
        self.threshold = int(threshold)
        self.pending: dict[str, Any] | None = None

    @staticmethod
    def _signature(value: torch.Tensor) -> tuple[Any, ...]:
        return (
            value.data_ptr(),
            tuple(value.shape),
            tuple(value.stride()),
            value.storage_offset(),
            value.device,
            value.dtype,
        )

    def forward(self, role: str, value: torch.Tensor) -> torch.Tensor:
        input_shape = tuple(value.shape)
        flattened = value.reshape(-1, input_shape[-1]).contiguous()
        if flattened.shape[0] > self.threshold:
            self.pending = None
            return self.original_forwards[role](value)
        signature = self._signature(value)
        if role == "q":
            q_proj = self.projections["q"]
            k_proj = self.projections["k"]
            v_proj = self.projections["v"]
            q_output, k_output, v_output = fused_fp8_qkv(
                flattened,
                q_qdata=q_proj.qdata,
                q_scale=q_proj.weight_scale,
                k_qdata=k_proj.qdata,
                k_scale=k_proj.weight_scale,
                v_qdata=v_proj.qdata,
                v_scale=v_proj.weight_scale,
            )
            self.pending = {
                "signature": signature,
                "k": k_output.reshape(*input_shape[:-1], k_proj.out_features),
                "v": v_output.reshape(*input_shape[:-1], v_proj.out_features),
            }
            return q_output.reshape(*input_shape[:-1], q_proj.out_features)
        if self.pending is None or self.pending["signature"] != signature:
            return self.original_forwards[role](value)
        output = self.pending[role]
        if role == "v":
            self.pending = None
        return output


def install_fp8_fused_qkv(
    model: nn.Module,
    *,
    threshold: int = 1,
) -> dict[str, Any]:
    """Patch Q/K/V projection forwards with one grouped M=1 dispatch."""

    if threshold <= 0:
        raise ValueError("fused QKV threshold must be positive")
    installed = []
    for layer in range(36):
        name = f"language_model.layers.{layer}.self_attn"
        attention = model.get_submodule(name)
        if hasattr(attention, "_fp8_fused_qkv_coordinator"):
            raise RuntimeError(f"{name}: fused QKV is already installed")
        projections = {}
        for role in ("q", "k", "v"):
            module = getattr(attention, f"{role}_proj", None)
            if not _is_fp8_linear(module):
                raise TypeError(
                    f"{name}.{role}_proj: expected MageVLScaledFP8Linear, "
                    f"got {type(module).__name__}"
                )
            if module.bias_bf16 is not None:
                raise RuntimeError(f"{name}.{role}_proj: fused QKV requires no bias")
            projections[role] = module
        coordinator = _FP8FusedQKVCoordinator(
            projections["q"], projections["k"], projections["v"], threshold=threshold
        )
        object.__setattr__(attention, "_fp8_fused_qkv_coordinator", coordinator)
        for role, module in projections.items():
            def fused_forward(
                self,
                value: torch.Tensor,
                *,
                _role: str = role,
                _coordinator: _FP8FusedQKVCoordinator = coordinator,
            ):
                return _coordinator.forward(_role, value)

            object.__setattr__(module, "forward", types.MethodType(fused_forward, module))
        installed.append(name)
    return {
        "feature": "fused_fp8_w8a16_qkv",
        "threshold": int(threshold),
        "installed_module_count": len(installed),
        "fallback": "separate_q_k_v_projection_forwards",
        "stored_weight_payload": "unchanged",
    }


def restore_fp8_fused_qkv(model: nn.Module) -> dict[str, Any]:
    restored = []
    for layer in range(36):
        name = f"language_model.layers.{layer}.self_attn"
        attention = model.get_submodule(name)
        coordinator = getattr(attention, "_fp8_fused_qkv_coordinator", None)
        if coordinator is None:
            continue
        for role in ("q", "k", "v"):
            module = getattr(attention, f"{role}_proj")
            object.__setattr__(module, "forward", coordinator.original_forwards[role])
        coordinator.pending = None
        object.__delattr__(attention, "_fp8_fused_qkv_coordinator")
        restored.append(name)
    return {"feature": "fused_fp8_w8a16_qkv", "restored": len(restored)}


__all__ = [
    "configure_fp8_decode_sources",
    "fp8_decode_source_manifest",
    "install_fp8_fused_gate_up",
    "install_fp8_fused_qkv",
    "restore_fp8_fused_gate_up",
    "restore_fp8_fused_qkv",
    "smallm_fp8_linear",
]