fix import path

hpcaitech · Feb 28, 2024 · 6607316 · 6607316
1 parent b4a2157
commit 6607316
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 216 deletions.
diff --git a/tests/test_infer/test_ops/cuda/kernel_utils.py b/tests/test_infer/test_ops/cuda/kernel_utils.py
diff --git a/tests/test_infer/test_ops/cuda/test_kv_cache_memcpy.py b/tests/test_infer/test_ops/cuda/test_kv_cache_memcpy.py
@@ -1,12 +1,12 @@
 import pytest
 import torch
-from kernel_utils import generate_caches_and_block_tables_v2, mock_alloc_single_token
 from packaging import version
 
 from colossalai.inference.modeling.layers.attention import copy_to_cache
 from colossalai.kernel.kernel_loader import InferenceOpsLoader
 from colossalai.kernel.triton import copy_kv_to_blocked_cache
 from colossalai.utils import get_current_device
+from tests.test_infer.test_ops.triton.kernel_utils import generate_caches_and_block_tables_v2, mock_alloc_single_token
 
 try:
     import triton  # noqa

diff --git a/tests/test_infer/test_ops/triton/test_kvcache_copy.py b/tests/test_infer/test_ops/triton/test_kvcache_copy.py
@@ -1,11 +1,11 @@
 import pytest
 import torch
-from kernel_utils import generate_caches_and_block_tables_v2, mock_alloc_single_token
 from packaging import version
 
 from colossalai.inference.modeling.layers.attention import copy_to_cache
 from colossalai.kernel.triton import copy_kv_to_blocked_cache
 from colossalai.utils import get_current_device
+from tests.test_infer.test_ops.triton.kernel_utils import generate_caches_and_block_tables_v2, mock_alloc_single_token
 
 try:
     import triton  # noqa