Skip to content

GitLab

Unverified Commit 32e7f994 authored by Xuanlei Zhao's avatar Xuanlei Zhao Committed by GitHub
Browse files

[kernel] update triton init #4740 (#4740)

parent d151dcab
Hide whitespace changes
Inline Side-by-side
Showing with 19 additions and 11 deletions
colossalai/kernel/triton/__init__.py
View file @ 32e7f994
from .context_attention import bloom_context_attn_fwd, llama_context_attn_fwd try:
from .copy_kv_cache_dest import copy_kv_cache_to_dest import triton
from .fused_layernorm import layer_norm HAS_TRITON = True
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__ = [ from .context_attention import bloom_context_attn_fwd, llama_context_attn_fwd
"llama_context_attn_fwd", "bloom_context_attn_fwd", "softmax", "layer_norm", "rmsnorm_forward", from .copy_kv_cache_dest import copy_kv_cache_to_dest
"copy_kv_cache_to_dest", "rotary_embedding_fwd", "token_attention_fwd" from .fused_layernorm import layer_norm
] 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"
]
except ImportError:
HAS_TRITON = False
print("Triton is not installed. Please install Triton to use Triton kernels.")
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment