ColossalAI/colossalai/kernel/triton/__init__.py

try:
    import triton

    HAS_TRITON = True

    from .context_attention import bloom_context_attn_fwd, llama_context_attn_fwd
    from .copy_kv_cache_dest import copy_kv_cache_to_dest
    from .fused_layernorm import layer_norm
    from .gptq_triton import gptq_fused_linear_triton
    from .rms_norm import rmsnorm_forward
    from .rotary_embedding_kernel import rotary_embedding_fwd
    from .softmax import softmax
    from .token_attention_kernel import token_attention_fwd

    __all__ = [
        "llama_context_attn_fwd",
        "bloom_context_attn_fwd",
        "softmax",
        "layer_norm",
        "rmsnorm_forward",
        "copy_kv_cache_to_dest",
        "rotary_embedding_fwd",
        "token_attention_fwd",
        "gptq_fused_linear_triton",
    ]

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