smallm_gemv_runtime.py

"""JIT loader for the optional Mage-VL packed-NVFP4 small-M GEMV."""

from __future__ import annotations

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

import torch


_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 _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):
        _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(
            f"small-M native sources resolved to different directories: "
            f"{sorted(str(value) for value in parents)}"
        )
    _SOURCE_ROOT = parents.pop()


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


def smallm_source_manifest() -> dict[str, str]:
    source_root = _source_root()
    return {
        path.name: hashlib.sha256(path.read_bytes()).hexdigest()
        for path in sorted(source_root.iterdir())
        if path.is_file()
    }


@lru_cache(maxsize=1)
def load_smallm_gemv_extension(*, verbose: bool = False) -> Any:
    from torch.utils.cpp_extension import load

    package_root = Path(__file__).resolve().parent
    source_root = _source_root()
    configured_build = os.environ.get("MAGE_VL_SMALLM_BUILD_DIR")
    build_root = (
        Path(configured_build).expanduser().resolve()
        if configured_build
        else package_root / ".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")
    return load(
        name="mage_vl_smallm_gemv_v1",
        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=verbose,
        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:
    extension = load_smallm_gemv_extension()
    return extension.linear(
        value.contiguous(),
        qdata,
        weight_block_scale,
        weight_scale,
        bias,
    )


__all__ = [
    "configure_smallm_source",
    "load_smallm_gemv_extension",
    "smallm_nvfp4_linear",
    "smallm_source_manifest",
]