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_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 .int8_rotary_embedding_kernel import int8_rotary_embedding_fwd from .smooth_attention import smooth_llama_context_attn_fwd, smooth_token_attention_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", "copy_kv_cache_to_dest", "token_attention_fwd", "gptq_fused_linear_triton", "int8_rotary_embedding_fwd", "smooth_llama_context_attn_fwd", "smooth_token_attention_fwd", ]