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.")