maba_sparse/kernels/dispatcher.py

import logging
import os
from typing import Any, Callable, Dict, Optional, Set, Tuple, Union
import warnings

import torch

from maba_sparse.kernels.common import (
    reference_compute_centroids,
    reference_dgda_prefill,
    reference_dgda_step,
    reference_index_topk,
    reference_stream_superposition,
    ref_compute_centroids,
    ref_dgda_prefill,
    ref_dgda_step,
    ref_index_topk,
    ref_stream_superposition,
    ref_centroids,
    ref_topk,
    ref_superposition,
)
from maba_sparse.kernels.cpu_dgda import (
    cpu_dgda_prefill,
    cpu_dgda_step,
)

logger = logging.getLogger(__name__)

_KERNEL_REGISTRY: Dict[str, Dict[str, Callable]] = {
    "reference": {
        "dgda_prefill": reference_dgda_prefill,
        "dgda_step": reference_dgda_step,
        "compute_centroids": reference_compute_centroids,
        "index_topk": reference_index_topk,
        "stream_superposition": reference_stream_superposition,
    },
    "cpu": {
        "dgda_prefill": cpu_dgda_prefill,
        "dgda_step": cpu_dgda_step,
        "compute_centroids": reference_compute_centroids,
        "index_topk": reference_index_topk,
        "stream_superposition": reference_stream_superposition,
    },
    "triton": {},
    "xla": {},
}

_WARNED_FALLBACKS: Set[str] = set()

_BACKEND_CACHE: Dict[Tuple[Optional[str], Optional[int], Optional[str]], str] = {}
_SM75_CACHE: Dict[Optional[int], bool] = {}
_TRITON_CACHE: Optional[bool] = None
_OPENMP_CACHE: Optional[bool] = None


def clear_dispatcher_cache() -> None:
    _BACKEND_CACHE.clear()
    _SM75_CACHE.clear()
    global _TRITON_CACHE, _OPENMP_CACHE
    _TRITON_CACHE = None
    _OPENMP_CACHE = None


def clear_fallback_warnings() -> None:
    _WARNED_FALLBACKS.clear()
    clear_dispatcher_cache()


def register_kernel(backend: str, op_name: str) -> Callable:
    def decorator(fn: Callable) -> Callable:
        if backend not in _KERNEL_REGISTRY:
            _KERNEL_REGISTRY[backend] = {}
        _KERNEL_REGISTRY[backend][op_name] = fn
        return fn
    return decorator


def get_kernel(backend: str, op_name: str) -> Optional[Callable]:
    if backend == "triton" and op_name not in _KERNEL_REGISTRY.get("triton", {}):
        if op_name in ("dgda_prefill", "dgda_step"):
            try:
                import maba_sparse.kernels.triton_dgda  # noqa: F401
            except Exception:
                pass
        elif op_name in ("compute_centroids", "index_topk", "stream_superposition"):
            try:
                import maba_sparse.kernels.triton_indexer  # noqa: F401
            except Exception:
                pass
    elif backend == "xla" and op_name not in _KERNEL_REGISTRY.get("xla", {}):
        if op_name in ("dgda_prefill", "dgda_step"):
            try:
                import maba_sparse.kernels.xla_dgda  # noqa: F401
            except Exception:
                pass
        elif op_name in ("compute_centroids", "index_topk", "stream_superposition"):
            try:
                import maba_sparse.kernels.xla_indexer  # noqa: F401
            except Exception:
                pass
    elif backend == "cpu":
        if op_name in ("compute_centroids", "index_topk", "stream_superposition"):
            reg_fn = _KERNEL_REGISTRY.get("cpu", {}).get(op_name)
            if reg_fn is None or reg_fn.__name__.startswith("ref"):
                try:
                    import maba_sparse.kernels.cpu_indexer  # noqa: F401
                except Exception:
                    pass
    return _KERNEL_REGISTRY.get(backend, {}).get(op_name, None)


def is_cuda_sm75_available(device: Optional[Union[torch.device, str, int]] = None) -> bool:
    if not torch.cuda.is_available():
        return False

    dev_idx = 0
    if device is not None:
        if isinstance(device, int):
            dev_idx = device
        elif isinstance(device, str):
            try:
                d = torch.device(device)
                dev_idx = d.index if d.index is not None else 0
            except Exception:
                return False
        elif isinstance(device, torch.device):
            dev_idx = device.index if device.index is not None else 0

    if dev_idx in _SM75_CACHE:
        return _SM75_CACHE[dev_idx]

    try:
        cap = torch.cuda.get_device_capability(dev_idx)
        res = cap >= (7, 5)
        _SM75_CACHE[dev_idx] = res
        return res
    except Exception:
        return False


def is_triton_available() -> bool:
    global _TRITON_CACHE
    if _TRITON_CACHE is not None:
        return _TRITON_CACHE
    try:
        import triton  # noqa: F401
        _TRITON_CACHE = True
        return True
    except (ImportError, ModuleNotFoundError, Exception):
        _TRITON_CACHE = False
        return False


def is_xla_available(device: Optional[Union[torch.device, str]] = None) -> bool:
    if device is not None:
        dev = torch.device(device) if isinstance(device, str) else device
        if dev.type == "xla":
            return True
    try:
        import torch_xla  # noqa: F401
        return True
    except (ImportError, ModuleNotFoundError, Exception):
        return False


def is_openmp_available() -> bool:
    global _OPENMP_CACHE
    if _OPENMP_CACHE is not None:
        return _OPENMP_CACHE
    try:
        res = torch.backends.openmp.is_available()
        _OPENMP_CACHE = res
        return res
    except Exception:
        _OPENMP_CACHE = False
        return False


def get_backend(device: Union[torch.device, str, None] = None) -> str:
    dev_type: Optional[str] = None
    dev_idx: Optional[int] = None
    if device is not None:
        if isinstance(device, str):
            try:
                d = torch.device(device)
                dev_type, dev_idx = d.type, (d.index if d.index is not None else 0)
            except Exception:
                return "reference"
        elif isinstance(device, torch.device):
            dev_type, dev_idx = device.type, (device.index if device.index is not None else 0)

    env_override = os.environ.get("MABA_BACKEND")
    cache_key = (dev_type, dev_idx, env_override)
    cached = _BACKEND_CACHE.get(cache_key)
    if cached is not None:
        return cached

    if env_override and env_override.lower() != "auto":
        val = env_override.lower()
        if val in ("triton", "cuda", "gpu"):
            res = "triton"
        elif val in ("xla", "tpu"):
            res = "xla"
        elif val in ("cpu", "openmp"):
            res = "cpu"
        elif val in ("reference", "pytorch", "ref"):
            res = "reference"
        else:
            raise ValueError(
                f"Unsupported MABA_BACKEND '{env_override}'. Supported: ['auto', 'triton', 'cpu', 'xla', 'reference', 'ref']"
            )
        _BACKEND_CACHE[cache_key] = res
        return res

    if dev_type is None:
        if torch.cuda.is_available() and is_cuda_sm75_available(0) and is_triton_available():
            res = "triton"
        else:
            res = "cpu"
    elif dev_type == "cuda":
        if torch.cuda.is_available() and is_cuda_sm75_available(dev_idx) and is_triton_available():
            res = "triton"
        else:
            res = "reference"
    elif dev_type == "xla":
        if is_xla_available(device):
            res = "xla"
        else:
            res = "reference"
    elif dev_type == "cpu":
        res = "cpu"
    else:
        res = "reference"

    _BACKEND_CACHE[cache_key] = res
    return res


def _log_fallback_warning(op_name: str, backend: str, exc: Optional[Exception] = None) -> None:
    strict = os.environ.get("MABA_STRICT_BACKEND", "0").lower() in ("1", "true")
    err_detail = f" ({type(exc).__name__}: {exc})" if exc is not None else ""
    msg = (
        f"[Maba Dispatcher Fallback] Native kernel for '{op_name}' on backend '{backend}' "
        f"is unavailable or failed{err_detail}. Falling back to reference."
    )
    if strict:
        raise RuntimeError(msg)

    key = f"{backend}:{op_name}"
    if key not in _WARNED_FALLBACKS:
        _WARNED_FALLBACKS.add(key)
        warnings.warn(msg, RuntimeWarning, stacklevel=3)
        logger.debug(msg)


def dispatch_dgda_prefill(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    alpha: Optional[torch.Tensor] = None,
    b: Optional[torch.Tensor] = None,
    w: Optional[torch.Tensor] = None,
    chunk_size: int = 16,
    initial_state: Optional[torch.Tensor] = None,
    log_alpha: Optional[torch.Tensor] = None,
    **kwargs: Any,
) -> Tuple[torch.Tensor, torch.Tensor]:
    if alpha is None and log_alpha is None:
        raise ValueError("Either alpha or log_alpha must be provided")
    if b is None or w is None:
        raise ValueError("Both b and w gate tensors must be provided")

    q_dev = q.device
    q_dt = q.dtype

    supported_dtypes = (torch.float32, torch.float16, torch.bfloat16, torch.float64)
    if q_dt not in supported_dtypes:
        raise TypeError(f"Unsupported dtype {q_dt}. Supported: {supported_dtypes}")

    tensors = [("k", k), ("v", v), ("b", b), ("w", w)]
    if alpha is not None:
        tensors.append(("alpha", alpha))
    if log_alpha is not None:
        tensors.append(("log_alpha", log_alpha))

    for name, t in tensors:
        if t.device != q_dev:
            raise RuntimeError(
                f"Expected all tensors to be on the same device, but got {q_dev} and {t.device} for {name}"
            )
        if t.dtype != q_dt:
            raise TypeError(
                f"Tensor dtypes must match, but got {q_dt} for q and {t.dtype} for {name}"
            )
    if initial_state is not None:
        if initial_state.device != q_dev:
            raise RuntimeError(
                f"Expected all tensors to be on the same device, but got {q_dev} and {initial_state.device} for initial_state"
            )
        if initial_state.dtype != q_dt:
            raise TypeError(
                f"Tensor dtypes must match, but got {q_dt} for q and {initial_state.dtype} for initial_state"
            )

    target_backend = get_backend(q_dev)
    kernel_fn = get_kernel(target_backend, "dgda_prefill")

    if kernel_fn is not None and target_backend != "reference":
        try:
            return kernel_fn(
                q=q, k=k, v=v, alpha=alpha, b=b, w=w,
                chunk_size=chunk_size, initial_state=initial_state,
                log_alpha=log_alpha, **kwargs
            )
        except Exception as exc:
            _log_fallback_warning("dgda_prefill", target_backend, exc)
    elif target_backend != "reference" and kernel_fn is None:
        _log_fallback_warning("dgda_prefill", target_backend, None)

    ref_fn = get_kernel("reference", "dgda_prefill") or reference_dgda_prefill
    return ref_fn(
        q=q, k=k, v=v, alpha=alpha, b=b, w=w,
        chunk_size=chunk_size, initial_state=initial_state,
        log_alpha=log_alpha, **kwargs
    )


def dispatch_dgda_step(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    alpha: Optional[torch.Tensor] = None,
    b: Optional[torch.Tensor] = None,
    w: Optional[torch.Tensor] = None,
    state: Optional[torch.Tensor] = None,
    log_alpha: Optional[torch.Tensor] = None,
    **kwargs: Any,
) -> Tuple[torch.Tensor, torch.Tensor]:
    if alpha is None and log_alpha is None:
        raise ValueError("Either alpha or log_alpha must be provided")
    if b is None or w is None:
        raise ValueError("Both b and w gate tensors must be provided")

    q_dev = q.device
    q_dt = q.dtype

    supported_dtypes = (torch.float32, torch.float16, torch.bfloat16, torch.float64)
    if q_dt not in supported_dtypes:
        raise TypeError(f"Unsupported dtype {q_dt}. Supported: {supported_dtypes}")

    tensors = [("k", k), ("v", v), ("b", b), ("w", w)]
    if alpha is not None:
        tensors.append(("alpha", alpha))
    if log_alpha is not None:
        tensors.append(("log_alpha", log_alpha))

    for name, t in tensors:
        if t.device != q_dev:
            raise RuntimeError(
                f"Expected all tensors to be on the same device, but got {q_dev} and {t.device} for {name}"
            )
        if t.dtype != q_dt:
            raise TypeError(
                f"Tensor dtypes must match, but got {q_dt} for q and {t.dtype} for {name}"
            )
    if state is not None:
        if state.device != q_dev:
            raise RuntimeError(
                f"Expected all tensors to be on the same device, but got {q_dev} and {state.device} for state"
            )
        if state.dtype != q_dt:
            raise TypeError(
                f"Tensor dtypes must match, but got {q_dt} for q and {state.dtype} for state"
            )

    target_backend = get_backend(q_dev)
    kernel_fn = get_kernel(target_backend, "dgda_step")

    if kernel_fn is not None and target_backend != "reference":
        try:
            return kernel_fn(
                q=q, k=k, v=v, alpha=alpha, b=b, w=w,
                state=state, log_alpha=log_alpha, **kwargs
            )
        except Exception as exc:
            _log_fallback_warning("dgda_step", target_backend, exc)
    elif target_backend != "reference" and kernel_fn is None:
        _log_fallback_warning("dgda_step", target_backend, None)

    ref_fn = get_kernel("reference", "dgda_step") or reference_dgda_step
    return ref_fn(
        q=q, k=k, v=v, alpha=alpha, b=b, w=w,
        state=state, log_alpha=log_alpha, **kwargs
    )


def dispatch_compute_centroids(
    k_idx: torch.Tensor,
    block_size: int = 64,
    **kwargs: Any,
) -> torch.Tensor:
    from maba_sparse.kernels.common import validate_indexer_inputs
    validate_indexer_inputs(k_idx=k_idx)
    if block_size <= 0:
        raise ValueError(f"block_size must be positive, got {block_size}")

    target_backend = get_backend(k_idx.device)
    kernel_fn = get_kernel(target_backend, "compute_centroids")

    if kernel_fn is not None and target_backend != "reference":
        try:
            return kernel_fn(k_idx=k_idx, block_size=block_size, **kwargs)
        except Exception as exc:
            _log_fallback_warning("compute_centroids", target_backend, exc)
    elif target_backend != "reference" and kernel_fn is None:
        _log_fallback_warning("compute_centroids", target_backend, None)

    ref_fn = get_kernel("reference", "compute_centroids") or reference_compute_centroids
    return ref_fn(k_idx=k_idx, block_size=block_size, **kwargs)


def dispatch_index_topk(
    q_idx: torch.Tensor,
    centroids: torch.Tensor,
    lambda_dist: float = 0.5,
    top_k: int = 32,
    block_size: int = 64,
    **kwargs: Any,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
    if q_idx.device != centroids.device:
        raise RuntimeError(
            f"Expected all tensors to be on the same device, but got {q_idx.device} and {centroids.device}"
        )
    if q_idx.dim() != 3 or centroids.dim() != 3:
        raise ValueError(f"q_idx and centroids must be 3D, got {q_idx.dim()}D and {centroids.dim()}D")
    if q_idx.shape[0] != centroids.shape[0] or q_idx.shape[2] != centroids.shape[2]:
        raise ValueError(
            f"Dimension mismatch between q_idx {list(q_idx.shape)} and centroids {list(centroids.shape)}"
        )
    if top_k <= 0:
        raise ValueError(f"top_k must be positive, got {top_k}")
    if block_size <= 0:
        raise ValueError(f"block_size must be positive, got {block_size}")

    target_backend = get_backend(q_idx.device)
    kernel_fn = get_kernel(target_backend, "index_topk")

    if kernel_fn is not None and target_backend != "reference":
        try:
            return kernel_fn(
                q_idx=q_idx, centroids=centroids, lambda_dist=lambda_dist,
                top_k=top_k, block_size=block_size, **kwargs
            )
        except Exception as exc:
            _log_fallback_warning("index_topk", target_backend, exc)
    elif target_backend != "reference" and kernel_fn is None:
        _log_fallback_warning("index_topk", target_backend, None)

    ref_fn = get_kernel("reference", "index_topk") or reference_index_topk
    return ref_fn(
        q_idx=q_idx, centroids=centroids, lambda_dist=lambda_dist,
        top_k=top_k, block_size=block_size, **kwargs
    )


def dispatch_stream_superposition(
    o_local: torch.Tensor,
    o_sparse: torch.Tensor,
    o_hca: torch.Tensor,
    gate_logits: Optional[torch.Tensor] = None,
    **kwargs: Any,
) -> torch.Tensor:
    device = o_local.device
    if o_sparse.device != device or o_hca.device != device:
        raise RuntimeError(
            f"Expected all tensors to be on the same device, but got {device}, {o_sparse.device}, {o_hca.device}"
        )
    gate_weights = kwargs.get("gate_weights")
    gate_t = gate_logits if gate_logits is not None else gate_weights
    if gate_t is not None and gate_t.device != device:
        raise RuntimeError(
            f"Expected all tensors to be on the same device, but got {device} and {gate_t.device}"
        )

    from maba_sparse.kernels.common import validate_superposition_inputs
    validate_superposition_inputs(o_local, o_sparse, o_hca, gate_t)

    target_backend = get_backend(device)
    kernel_fn = get_kernel(target_backend, "stream_superposition")

    if kernel_fn is not None and target_backend != "reference":
        try:
            return kernel_fn(
                o_local=o_local, o_sparse=o_sparse, o_hca=o_hca,
                gate_logits=gate_logits, **kwargs
            )
        except Exception as exc:
            _log_fallback_warning("stream_superposition", target_backend, exc)
    elif target_backend != "reference" and kernel_fn is None:
        _log_fallback_warning("stream_superposition", target_backend, None)

    ref_fn = get_kernel("reference", "stream_superposition") or reference_stream_superposition
    return ref_fn(
        o_local=o_local, o_sparse=o_sparse, o_hca=o_hca,
        gate_logits=gate_logits, **kwargs
    )


dgda_prefill = dispatch_dgda_prefill
dgda_step = dispatch_dgda_step
compute_centroids = dispatch_compute_centroids
index_topk = dispatch_index_topk
stream_superposition = dispatch_stream_superposition


__all__ = [
    "get_backend",
    "register_kernel",
    "get_kernel",
    "clear_fallback_warnings",
    "clear_dispatcher_cache",
    "is_cuda_sm75_available",
    "is_triton_available",
    "is_xla_available",
    "is_openmp_available",
    "dispatch_dgda_prefill",
    "dispatch_dgda_step",
    "dispatch_compute_centroids",
    "dispatch_index_topk",
    "dispatch_stream_superposition",
    "dgda_prefill",
    "dgda_step",
    "compute_centroids",
    "index_topk",
    "stream_superposition",
    "ref_dgda_prefill",
    "ref_dgda_step",
    "ref_centroids",
    "ref_topk",
    "ref_superposition",
]