ColossalAI/colossalai/kernel/triton/__init__.py

try:
    import triton

    HAS_TRITON = True
except ImportError:
    HAS_TRITON = False
    print("Triton is not installed. Please install Triton to use Triton kernels.")

# There may exist import error even if we have triton installed.
if HAS_TRITON:
    from .context_attn_unpad import context_attention_unpadded
    from .flash_decoding import flash_decoding_fwd
    from .fused_layernorm import layer_norm
    from .gptq_triton import gptq_fused_linear_triton
    from .kvcache_copy import copy_kv_to_blocked_cache
    from .no_pad_rotary_embedding import rotary_embedding
    from .softmax import softmax

    __all__ = [
        "context_attention_unpadded",
        "flash_decoding_fwd",
        "copy_kv_to_blocked_cache",
        "softmax",
        "layer_norm",
        "gptq_fused_linear_triton",
        "rotary_embedding",
    ]