h3-blackwell-runtime/src/h3_blackwell_runtime/fc2_lt.py

"""Guarded production dispatch for the canonical GB10 FC2 cuBLASLt schedule."""

from __future__ import annotations

import os
from functools import lru_cache
from pathlib import Path

import torch


_TRUTHY = {"1", "true", "yes", "on"}
_CANONICAL_LOGICAL_SHAPE = (37_810, 14_336, 5_376)
_CANONICAL_PACKED_ROWS = 37_824
_CONFIG = {
    "algorithm_id": 70,
    "tile_id": 20,
    "stages_id": 37,
    "split_k": 1,
    "reduction_scheme": 0,
    "custom_option": 0,
    "cta_swizzle": 0,
}
_VALIDATED: set[tuple[int, tuple[int, ...], tuple[int, ...]]] = set()
_STATS = {"attempts": 0, "successes": 0, "fallbacks": 0}


class _Fc2LtStrictError(RuntimeError):
    pass


def fc2_lt_enabled() -> bool:
    return os.getenv("H3_NVFP4_FC2_LT_SPLITK1", "").lower() in _TRUTHY


def _fc2_lt_strict() -> bool:
    return os.getenv("H3_NVFP4_FC2_LT_STRICT", "").lower() in _TRUTHY


def fc2_lt_status() -> dict:
    return {"enabled": fc2_lt_enabled(), **_STATS}


@lru_cache(maxsize=1)
def _fc2_lt_extension():
    from torch.utils.cpp_extension import CUDA_HOME, load

    if CUDA_HOME is None:
        raise RuntimeError("CUDA_HOME is unavailable")
    source = Path(__file__).resolve().parent / "csrc" / "fc2_nvfp4_lt.cpp"
    return load(
        name="h3_fc2_nvfp4_lt",
        sources=[str(source)],
        extra_include_paths=[str(Path(CUDA_HOME) / "include")],
        extra_cflags=["-O2", "-std=c++17"],
        extra_ldflags=[
            "-L" + str(Path(CUDA_HOME) / "lib64"),
            "-lcublasLt",
            "-lcublas",
            "-lcudart",
        ],
        verbose=os.getenv("H3_NVFP4_FC2_LT_VERBOSE", "").lower() in _TRUTHY,
    )


@lru_cache(maxsize=1)
def _fc2_lt_extension_result():
    try:
        return _fc2_lt_extension(), None
    except Exception as error:
        return None, error


@lru_cache(maxsize=1)
def _fc2_lt_compatible_result():
    extension, load_error = _fc2_lt_extension_result()
    if extension is None:
        return None, load_error
    try:
        info = dict(extension.build_info())
        compatible = (
            info.get("cuda_version") == 13_000
            and info.get("cublas_version") == 130_100
            and info.get("cuda_runtime_version") == 13_000
            and info.get("cublaslt_runtime_version") == 130_000
        )
        if not compatible:
            return None, RuntimeError(f"unsupported FC2 cuBLASLt environment: {info}")
        extension.prepare()
        return extension, None
    except Exception as error:
        return None, error


def _canonical_fc2_supported(linear, gate_up: torch.Tensor, qdata: torch.Tensor) -> bool:
    logical_rows, in_features, out_features = _CANONICAL_LOGICAL_SHAPE
    return (
        linear.role == "h3_mlp_fc2"
        and linear.in_features == in_features
        and linear.out_features == out_features
        and linear.output_dtype == torch.bfloat16
        and linear.bias is None
        and tuple(gate_up.shape) == (logical_rows, in_features * 2)
        and tuple(qdata.shape) == (_CANONICAL_PACKED_ROWS, in_features // 2)
        and torch.cuda.get_device_capability(gate_up.device) == (12, 1)
    )


def _fallback_or_raise(message: str, error: Exception | None = None):
    _STATS["fallbacks"] += 1
    if _fc2_lt_strict():
        if error is None:
            raise _Fc2LtStrictError(message)
        raise _Fc2LtStrictError(message) from error
    return None


def prepare_fc2_lt() -> bool:
    """Load and validate the extension before latency-sensitive inference."""
    if not fc2_lt_enabled():
        return False
    extension, error = _fc2_lt_compatible_result()
    if extension is None:
        _fallback_or_raise("FC2 cuBLASLt extension is incompatible", error)
        return False
    return True


def fc2_lt_linear(
    linear,
    gate_up: torch.Tensor,
    tensor_scale: torch.Tensor,
    qdata: torch.Tensor,
    block_scale: torch.Tensor,
) -> torch.Tensor | None:
    """Run the validated canonical schedule, or return ``None`` for fallback."""
    if not fc2_lt_enabled():
        return None
    try:
        supported = _canonical_fc2_supported(linear, gate_up, qdata)
    except Exception as error:
        return _fallback_or_raise("FC2 cuBLASLt eligibility check failed", error)
    if not supported:
        _STATS["fallbacks"] += 1
        return None
    _STATS["attempts"] += 1

    extension, load_error = _fc2_lt_compatible_result()
    if extension is None:
        return _fallback_or_raise(
            "FC2 cuBLASLt extension failed to load or is incompatible", load_error
        )

    try:
        device = gate_up.device.index if gate_up.device.index is not None else torch.cuda.current_device()
        validation_key = (device, tuple(qdata.shape), tuple(linear.weight.shape))
        if validation_key not in _VALIDATED:
            checked = dict(
                extension.check(
                    qdata,
                    block_scale,
                    linear.weight,
                    linear.weight_scale,
                    _CONFIG,
                )
            )
            if not checked.get("valid") or checked.get("required_workspace_bytes") != 0:
                return _fallback_or_raise(
                    f"FC2 cuBLASLt algorithm check rejected the candidate: {checked}"
                )
            _VALIDATED.add(validation_key)

        alpha = (tensor_scale.float() * linear.weight_scale_2.float()).reshape(1).contiguous()
        beta = torch.zeros_like(alpha)
        output = torch.empty(
            (qdata.shape[0], linear.out_features),
            device=qdata.device,
            dtype=torch.bfloat16,
        )
        workspace = torch.empty(0, device=qdata.device, dtype=torch.uint8)
        extension.run(
            qdata,
            block_scale,
            linear.weight,
            linear.weight_scale,
            alpha,
            beta,
            output,
            workspace,
            _CONFIG,
        )
        _STATS["successes"] += 1
        return output[: gate_up.shape[0], : linear.out_features]
    except _Fc2LtStrictError:
        raise
    except Exception as error:
        return _fallback_or_raise("FC2 cuBLASLt candidate execution failed", error)