"""Kernel interfaces and implementations.

This module provides:
- Kernel: Base class for all kernels
- KernelRegistry: Registry for kernel management
- Basic kernels: LinearKernel, EmbeddingKernel, RMSNormKernel, SiLUKernel
- Fused kernels: FusedSiLUMulKernel, FusedRMSNormKernel, etc.
- CUDA Graph: CUDAGraphRunner, CUDAGraphConfig

Kernels abstract hardware-specific compute operations, allowing
core to be hardware-agnostic.

Example:
    # Basic kernel usage
    linear_kernel = backend.get_kernel("linear")
    output = linear_kernel(input, weight)

    # Fused kernel usage
    fused_silu_mul = backend.get_kernel("fused_silu_mul")
    output = fused_silu_mul(gate, up)

    # CUDA Graph for decode optimization
    from sagellm_backend.kernels import CUDAGraphRunner, CUDAGraphConfig

    config = CUDAGraphConfig(max_batch_size=128)
    runner = CUDAGraphRunner(model.forward, config)
    runner.capture_all(sample_input_fn)
    output = runner(input_batch)
"""

from __future__ import annotations

from sagellm_backend.kernels.activation import SiLUKernel
from sagellm_backend.kernels.base import Kernel, KernelRegistry
from sagellm_backend.kernels.embedding import EmbeddingKernel
from sagellm_backend.kernels.linear import LinearKernel
from sagellm_backend.kernels.normalization import RMSNormKernel

__all__ = [
    # Base classes
    "Kernel",
    "KernelRegistry",
    # Basic kernels
    "LinearKernel",
    "EmbeddingKernel",
    "RMSNormKernel",
    "SiLUKernel",
]

# Conditionally export fused kernels
try:
    from sagellm_backend.kernels.fused_ops import (
        FusedAddRMSNormKernel,
        FusedQKVProjectionKernel,
        FusedRMSNormKernel,
        FusedRotaryEmbeddingKernel,
        FusedSiLUMulKernel,
        get_fused_kernel,
        is_triton_available,
        list_fused_kernels,
    )

    __all__.extend(
        [
            "FusedSiLUMulKernel",
            "FusedRMSNormKernel",
            "FusedAddRMSNormKernel",
            "FusedQKVProjectionKernel",
            "FusedRotaryEmbeddingKernel",
            "get_fused_kernel",
            "list_fused_kernels",
            "is_triton_available",
        ]
    )
except ImportError:
    pass

# Conditionally export CUDA Graph utilities
try:
    from sagellm_backend.kernels.cuda_graph import (
        CapturedGraph,
        CUDAGraphConfig,
        CUDAGraphRunner,
        ModelRunnerWithCUDAGraph,
    )

    __all__.extend(
        [
            "CUDAGraphRunner",
            "CUDAGraphConfig",
            "CapturedGraph",
            "ModelRunnerWithCUDAGraph",
        ]
    )
except ImportError:
    pass
