From 65b8d65813f4c0af7498efa58a5a3ad02a8df65b Mon Sep 17 00:00:00 2001
From: Mingyan Jiang <1829166702@qq.com>
Date: Thu, 14 Sep 2023 17:02:55 +0800
Subject: [PATCH 1/2] [kenel] fix kernel import error due to uninstalling
 triton

---
 requirements/requirements.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 9aa5f2822e40..5f56f2ddcc0c 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -11,3 +11,4 @@ ninja
 torch>=1.12
 safetensors
 einops
+triton==2.0.0.dev20221202

From 92cd7f9c728c0bd9cfe5d49836e879aeb6fdee71 Mon Sep 17 00:00:00 2001
From: Mingyan Jiang <1829166702@qq.com>
Date: Thu, 14 Sep 2023 17:50:37 +0800
Subject: [PATCH 2/2] [kenel] fix kernel import error due to uninstalling
 triton

---
 colossalai/kernel/__init__.py                              | 7 -------
 colossalai/kernel/triton/__init__.py                       | 5 -----
 requirements/requirements.txt                              | 1 -
 .../test_infer_ops/triton/test_llama_context_attention.py  | 2 +-
 4 files changed, 1 insertion(+), 14 deletions(-)

diff --git a/colossalai/kernel/__init__.py b/colossalai/kernel/__init__.py
index a99cb497c3e7..8933fc0a3c2f 100644
--- a/colossalai/kernel/__init__.py
+++ b/colossalai/kernel/__init__.py
@@ -1,14 +1,7 @@
 from .cuda_native import FusedScaleMaskSoftmax, LayerNorm, MultiHeadAttention
-from .triton import llama_context_attn_fwd, bloom_context_attn_fwd
-from .triton import softmax
-from .triton import copy_kv_cache_to_dest
 
 __all__ = [
     "LayerNorm",
     "FusedScaleMaskSoftmax",
     "MultiHeadAttention",
-    "llama_context_attn_fwd",
-    "bloom_context_attn_fwd",
-    "softmax",
-    "copy_kv_cache_to_dest",
 ]
diff --git a/colossalai/kernel/triton/__init__.py b/colossalai/kernel/triton/__init__.py
index eb0335c01ce2..e69de29bb2d1 100644
--- a/colossalai/kernel/triton/__init__.py
+++ b/colossalai/kernel/triton/__init__.py
@@ -1,5 +0,0 @@
-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 .rms_norm import rmsnorm_forward
-from .softmax import softmax
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 5f56f2ddcc0c..9aa5f2822e40 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -11,4 +11,3 @@ ninja
 torch>=1.12
 safetensors
 einops
-triton==2.0.0.dev20221202
diff --git a/tests/test_infer_ops/triton/test_llama_context_attention.py b/tests/test_infer_ops/triton/test_llama_context_attention.py
index 4ea6095d4109..0ff06b5007e4 100644
--- a/tests/test_infer_ops/triton/test_llama_context_attention.py
+++ b/tests/test_infer_ops/triton/test_llama_context_attention.py
@@ -10,7 +10,7 @@
     import triton
     import triton.language as tl
 
-    from colossalai.kernel.triton import llama_context_attn_fwd
+    from colossalai.kernel.triton.context_attention import llama_context_attn_fwd
     from tests.test_infer_ops.triton.kernel_utils import torch_context_attention
     HAS_TRITON = True
 except ImportError: