github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-04 19:16:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Xuanlei Zhao 2023-09-18 09:44:27 +08:00 committed by GitHub
parent d151dcab74
commit 32e7f99416
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 19 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
colossalai/kernel/triton/__init__.py
View File

@ -1,12 +1,20 @@
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.")