accelerators.py

"""Optional CUDA/Triton kernels with a correctness-first PyTorch fallback.

The accelerator is deliberately parameterless.  It only changes how existing
source-derived tensors are evaluated; it never registers weights or persistent
model buffers.  Imports are lazy so CPU use and installations without FLA keep
working without importing Triton.
"""

from __future__ import annotations

import os
import warnings
from dataclasses import asdict, dataclass
from typing import Any

import torch


@dataclass(frozen=True, slots=True)
class DendroAcceleratorStatus:
    requested: str
    active: str
    available: bool
    reason: str | None
    kernel_cache: str | None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


_FLA_KERNELS: tuple[Any, Any, Any, Any] | None = None
_FLA_FAILURE: str | None = None
_WARNED_FAILURE = False


def _requested_backend(configured: str = "auto") -> str:
    requested = os.environ.get("DENDRO_ACCELERATOR", configured).strip().lower()
    if requested not in {"auto", "fla", "torch"}:
        warnings.warn(
            f"Unknown DENDRO_ACCELERATOR={requested!r}; using the PyTorch fallback",
            RuntimeWarning,
            stacklevel=3,
        )
        return "torch"
    return requested


def _prepare_kernel_cache() -> str | None:
    root = os.environ.get("DENDRO_KERNEL_CACHE")
    if not root:
        return os.environ.get("TRITON_CACHE_DIR") or os.environ.get("TRITON_HOME")
    os.environ.setdefault("TRITON_HOME", root)
    os.environ.setdefault("TRITON_CACHE_DIR", os.path.join(root, "cache"))
    return os.environ["TRITON_CACHE_DIR"]


def _load_fla(*, warn: bool = False) -> tuple[Any, Any, Any, Any] | None:
    global _FLA_KERNELS, _FLA_FAILURE, _WARNED_FAILURE
    if _FLA_KERNELS is not None:
        return _FLA_KERNELS
    if _FLA_FAILURE is not None:
        return None
    _prepare_kernel_cache()
    try:
        from fla.modules.convolution import causal_conv1d, causal_conv1d_update
        from fla.ops.gated_delta_rule import (
            chunk_gated_delta_rule,
            fused_recurrent_gated_delta_rule,
        )

        _FLA_KERNELS = (
            chunk_gated_delta_rule,
            fused_recurrent_gated_delta_rule,
            causal_conv1d,
            causal_conv1d_update,
        )
        return _FLA_KERNELS
    except Exception as error:  # optional dependency: every failure must fall back
        _FLA_FAILURE = f"{type(error).__name__}: {error}"
        if warn and not _WARNED_FAILURE:
            warnings.warn(
                f"FLA kernels are unavailable ({_FLA_FAILURE}); using PyTorch kernels",
                RuntimeWarning,
                stacklevel=3,
            )
            _WARNED_FAILURE = True
        return None


def _can_accelerate(tensor: torch.Tensor, configured: str) -> bool:
    requested = _requested_backend(configured)
    return (
        requested != "torch"
        and tensor.device.type == "cuda"
        and tensor.dtype in {torch.float16, torch.bfloat16}
        and _load_fla(warn=requested == "fla") is not None
    )


def accelerator_status(configured: str = "auto", *, probe: bool = False) -> DendroAcceleratorStatus:
    requested = _requested_backend(configured)
    if requested == "torch":
        return DendroAcceleratorStatus(requested, "torch", True, None, _prepare_kernel_cache())
    kernels = _load_fla(warn=requested == "fla") if probe else _FLA_KERNELS
    available = kernels is not None
    return DendroAcceleratorStatus(
        requested=requested,
        active="fla" if available else "torch",
        available=available,
        reason=None if available else (_FLA_FAILURE or "not probed"),
        kernel_cache=_prepare_kernel_cache(),
    )


def fla_chunk_gated_delta_rule(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    *,
    return_state: bool,
    configured: str = "auto",
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None:
    # Triton launch/import overhead dominates short prefills on the release RTX
    # 3060.  Explicit ``fla`` still permits benchmarking or overriding the guard.
    if _requested_backend(configured) == "auto" and query.shape[1] < 256:
        return None
    if not _can_accelerate(query, configured):
        return None
    assert _FLA_KERNELS is not None
    try:
        output, state = _FLA_KERNELS[0](
            query,
            key,
            value,
            g=g,
            beta=beta,
            output_final_state=return_state,
            use_qk_l2norm_in_kernel=True,
        )
        return (output, state) if return_state else output
    except Exception as error:
        _disable_after_runtime_failure("gated-delta chunk", error)
        return None


def fla_recurrent_gated_delta_rule(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    state: torch.Tensor,
    *,
    configured: str = "auto",
) -> tuple[torch.Tensor, torch.Tensor] | None:
    if not _can_accelerate(query, configured):
        return None
    assert _FLA_KERNELS is not None
    try:
        return _FLA_KERNELS[1](
            query,
            key,
            value,
            g=g,
            beta=beta,
            initial_state=state,
            output_final_state=True,
            use_qk_l2norm_in_kernel=True,
        )
    except Exception as error:
        _disable_after_runtime_failure("gated-delta recurrent", error)
        return None


def fla_causal_conv1d(
    sequence: torch.Tensor,
    weight: torch.Tensor,
    *,
    activation: str | None = "silu",
    configured: str = "auto",
) -> torch.Tensor | None:
    """Evaluate ``[batch, time, channels]`` with FLA's Triton convolution."""

    if _requested_backend(configured) == "auto" and sequence.shape[1] < 256:
        return None
    if not _can_accelerate(sequence, configured):
        return None
    assert _FLA_KERNELS is not None
    try:
        output, _ = _FLA_KERNELS[2](
            sequence,
            weight=weight,
            bias=None,
            activation=activation,
            backend="triton",
        )
        return output
    except Exception as error:
        _disable_after_runtime_failure("causal convolution", error)
        return None


def fla_causal_conv1d_update(
    token: torch.Tensor,
    state: torch.Tensor,
    weight: torch.Tensor,
    *,
    activation: str | None = "silu",
    configured: str = "auto",
) -> tuple[torch.Tensor, torch.Tensor] | None:
    """Advance one convolution token with FLA's in-place Triton state kernel.

    ``token`` is ``[batch, 1, channels]`` and ``state`` is
    ``[batch, channels, kernel]``.  The function is parameterless and mutates
    only the activation cache supplied by the caller.
    """

    if token.ndim != 3 or token.shape[1] != 1:
        return None
    if state.ndim != 3 or state.shape[0] != token.shape[0]:
        return None
    if state.shape[1] != token.shape[2] or state.shape[2] != weight.shape[1]:
        return None
    if not _can_accelerate(token, configured):
        return None
    assert _FLA_KERNELS is not None
    try:
        output, updated = _FLA_KERNELS[3](
            token,
            state,
            weight=weight,
            bias=None,
            activation=activation,
        )
        return output, updated
    except Exception as error:
        _disable_after_runtime_failure("causal convolution update", error)
        return None


def _disable_after_runtime_failure(operation: str, error: Exception) -> None:
    global _FLA_KERNELS, _FLA_FAILURE, _WARNED_FAILURE
    _FLA_KERNELS = None
    _FLA_FAILURE = f"{operation}: {type(error).__name__}: {error}"
    if not _WARNED_FAILURE:
        warnings.warn(
            f"FLA {_FLA_FAILURE}; disabling it and continuing with PyTorch kernels",
            RuntimeWarning,
            stacklevel=3,
        )
        _WARNED_FAILURE = True