AstrAI/astrai/extension/backend/linear.py

"""Inference-only dispatch for AstrAI linear layers.

The CUDA GEMV path is narrow by construction rather than by a measured
shape table: the kernel streams each weight exactly once, so automatic
selection is keyed on the decode batch size alone (M in [2, 4], where it
sits at the HBM bandwidth floor and beat the cuBLAS small-M path on every
measured family). Every training, prefill-sized, out-of-band, or
unsupported call falls back to PyTorch.

The family stays registered with the shared operator dispatcher, so
``op_backend(linear=...)``, ``ASTR_OPS=linear=...``, and ``resolve`` /
``explain`` keep working like for attention and rotary. The per-layer
hot path only consults the dispatcher when one of those selections is
active, keeping it free of axes dictionaries and record sorting.
"""

from typing import Any, Dict, List, Optional

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

from astrai.extension.dispatch import (
    ImplRecord,
    Spec,
    axis,
    env_mode,
    env_selection,
    get_override,
    register_family,
    resolve,
    tensor_axes,
)
from astrai.extension.loader import is_available
from astrai.extension.ops.gemv import bf16_gemv

# M=1 keeps cuBLAS (its GEMV path is already at the bandwidth floor; only
# OPT 1.3B shapes ever passed the full gate). M >= 5 approaches the cuBLAS
# tensor-core crossover (M=8 regressed at wrapper level on every measured
# family, and cuBLAS clearly wins from M ~ 12).
_AUTO_GEMV_M = frozenset({2, 3, 4})


def _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
    return F.linear(x, weight, bias)


def _inference_bf16_gemv(
    x: Tensor, weight: Tensor, bias: Optional[Tensor] = None
) -> Tensor:
    # Model parameters retain requires_grad=True after eval().  Dispatch is
    # already restricted to no-grad, so detached views preserve storage and
    # layout while satisfying the primitive's explicit autograd guard.
    return bf16_gemv(
        x.detach(),
        weight.detach(),
        bias.detach() if bias is not None else None,
    )


def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
    if (
        torch.is_grad_enabled()
        or not x.is_cuda
        or x.dtype != torch.bfloat16
        or weight.dtype != torch.bfloat16
        or weight.ndim != 2
        or x.ndim not in (1, 2)
        or (x.ndim == 2 and not 1 <= x.shape[0] <= 8)
        or x.shape[-1] != weight.shape[1]
        or x.device != weight.device
        or not x.is_contiguous()
        or not weight.is_contiguous()
        or torch.cuda.get_device_capability(x.get_device()) < (8, 0)
        or not is_available("bf16_gemv")
    ):
        return False
    return bias is None or (
        bias.device == x.device
        and bias.dtype == torch.bfloat16
        and bias.ndim == 1
        and bias.shape[0] == weight.shape[0]
        and bias.is_contiguous()
    )


def _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str, Any]:
    weight_shape = tuple(weight.shape)
    m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None)
    supported_m = m is not None and 1 <= m <= 8
    shape_matches = (
        weight.ndim == 2
        and x.ndim in (1, 2)
        and bool(x.shape)
        and x.shape[-1] == weight_shape[-1]
    )
    same_device = x.device == weight.device and (
        bias is None or bias.device == x.device
    )
    bias_supported = bias is None or (
        bias.ndim == 1
        and weight.ndim == 2
        and bias.shape[0] == weight_shape[0]
        and bias.dtype == torch.bfloat16
        and bias.is_contiguous()
    )
    capability = torch.cuda.get_device_capability(x.device) if x.is_cuda else None
    return tensor_axes(
        x,
        mode=env_mode("ASTRAI_GEMV"),
        m=m,
        supported_m=supported_m,
        auto_m=m in _AUTO_GEMV_M,
        shape_matches=shape_matches,
        same_device=same_device,
        weight_dtype=weight.dtype,
        x_contiguous=x.is_contiguous(),
        weight_contiguous=weight.is_contiguous(),
        bias_supported=bias_supported,
        capability=capability,
    )


_SPEC_CAPABLE = (
    axis("device_cuda").truthy()
    & axis("dtype").in_(torch.bfloat16)
    & axis("weight_dtype").in_(torch.bfloat16)
    & axis("grad_enabled").eq(False)
    & axis("supported_m").truthy()
    & axis("shape_matches").truthy()
    & axis("same_device").truthy()
    & axis("x_contiguous").truthy()
    & axis("weight_contiguous").truthy()
    & axis("bias_supported").truthy()
    & Spec.of(
        lambda ax: ax.get("capability") is not None and ax.get("capability") >= (8, 0),
        "capability>=sm_80",
    )
)

_SPEC_AUTO = _SPEC_CAPABLE & axis("auto_m").truthy()


def _linear_records() -> List[ImplRecord]:
    mode = env_mode("ASTRAI_GEMV")
    gemv_priority = 0 if mode == "1" else 100
    auto_priority = 0 if mode == "auto" else 90
    torch_priority = 0 if mode == "0" else 50
    return [
        ImplRecord(
            family="linear",
            name="gemv",
            obj=_inference_bf16_gemv,
            spec=_SPEC_CAPABLE,
            available=lambda: is_available("bf16_gemv"),
            priority=gemv_priority,
        ),
        ImplRecord(
            family="linear",
            name="auto_gemv",
            obj=_inference_bf16_gemv,
            spec=_SPEC_AUTO,
            available=lambda: is_available("bf16_gemv"),
            priority=auto_priority,
        ),
        ImplRecord(
            family="linear",
            name="torch",
            obj=_torch_linear,
            spec=Spec.always(),
            priority=torch_priority,
        ),
    ]


def _fallback_record() -> ImplRecord:
    return ImplRecord(
        family="linear",
        name="torch",
        obj=_torch_linear,
        spec=Spec.always(),
        priority=999,
    )


register_family("linear", _axes, _linear_records, _fallback_record)


def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
    """Apply a linear projection with safe inference-only GEMV dispatch.

    ``ASTRAI_GEMV=0`` always uses PyTorch, ``1`` forces GEMV whenever the
    primitive can safely handle any M in ``{1, ..., 8}``, and ``auto`` (the
    default) uses GEMV for decode batches with M in ``{2, 3, 4}``.
    """
    # Route through the shared dispatcher whenever a selection is active so
    # explicit/context/env overrides stay honored; otherwise keep the hot
    # path free of axes dictionaries and record sorting.
    if get_override("linear") is not None or env_selection("linear") is not None:
        return resolve("linear", x, weight, bias).record.obj(x, weight, bias)

    mode = env_mode("ASTRAI_GEMV")
    if mode != "0" and _gemv_capable(x, weight, bias):
        m = 1 if x.ndim == 1 else x.shape[0]
        if mode == "1" or m in _AUTO_GEMV_M:
            return _inference_bf16_gemv(x, weight, bias)
    return _torch_linear(x, weight, bias)


__all__ = ["linear"]