Skip to content

Commit e76acbb

Browse files
authored
[inference] moved ops tests to test_infer (#5354)
1 parent db1a763 commit e76acbb

File tree

9 files changed

+7
-4
lines changed

9 files changed

+7
-4
lines changed

tests/test_infer/test_config_and_struct.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,9 @@ def check_config_and_inference():
6363
batch.add_seqs([sequence])
6464
batch.add_seqs([sequence2, sequence3])
6565

66+
# add duplicated sequence to test that it will not be counted twice
67+
batch.add_seqs([sequence])
68+
6669
assert batch.is_empty == False
6770
assert batch.get_batch_size() == 3
6871
batch.update_batch_tokens([1, 2, 3])

tests/test_infer_ops/triton/test_context_attn_unpad.py renamed to tests/test_infer/test_ops/triton/test_context_attn_unpad.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from colossalai.inference.modeling.layers.attention import PagedAttention
77
from colossalai.kernel.triton import context_attention_unpadded
88
from colossalai.utils import get_current_device
9-
from tests.test_infer_ops.triton.kernel_utils import generate_caches_and_block_tables_v2, torch_attn_ref
9+
from tests.test_infer.test_ops.triton.kernel_utils import generate_caches_and_block_tables_v2, torch_attn_ref
1010

1111
try:
1212
import triton # noqa

tests/test_infer_ops/triton/test_decoding_attn.py renamed to tests/test_infer/test_ops/triton/test_decoding_attn.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
from colossalai.kernel.triton import flash_decoding_attention
66
from colossalai.utils import get_current_device
7-
from tests.test_infer_ops.triton.kernel_utils import (
7+
from tests.test_infer.test_ops.triton.kernel_utils import (
88
convert_kv_unpad_to_padded,
99
generate_caches_and_block_tables_v2,
1010
prepare_padding_mask,

tests/test_infer_ops/triton/test_kvcache_copy.py renamed to tests/test_infer/test_ops/triton/test_kvcache_copy.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from colossalai.inference.modeling.layers.attention import copy_to_cache
66
from colossalai.kernel.triton import copy_kv_to_blocked_cache
77
from colossalai.utils import get_current_device
8-
from tests.test_infer_ops.triton.kernel_utils import generate_caches_and_block_tables_v2, mock_alloc_single_token
8+
from tests.test_infer.test_ops.triton.kernel_utils import generate_caches_and_block_tables_v2, mock_alloc_single_token
99

1010
try:
1111
import triton # noqa

tests/test_infer_ops/triton/test_rotary_embdding_unpad.py renamed to tests/test_infer/test_ops/triton/test_rotary_embdding_unpad.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb
55

66
from colossalai.kernel.triton import rotary_embedding
7-
from tests.test_infer_ops.triton.kernel_utils import mock_alloc_block_table_and_kvcache_v2
7+
from tests.test_infer.test_ops.triton.kernel_utils import mock_alloc_block_table_and_kvcache_v2
88

99
try:
1010
import triton # noqa

0 commit comments

Comments
 (0)