github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-02 16:28:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
flybird11111 2025-04-10 10:19:38 +08:00
parent 99298c6a6d
commit eaef783ec3
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tests/test_infer/test_kernels/cuda/test_flash_decoding_attention.py
View File

@ -18,6 +18,7 @@ from tests.test_infer.test_kernels.triton.kernel_utils import (
generate_caches_and_block_tables_vllm, generate_caches_and_block_tables_vllm,
torch_attn_ref, torch_attn_ref,
) )
from colossalai.testing import clear_cache_before_run
q_len = 1 q_len = 1
PARTITION_SIZE = 512 PARTITION_SIZE = 512
@ -55,7 +56,7 @@ def numpy_allclose(x, y, rtol, atol):
np.testing.assert_allclose(x_numpy, y_numpy, rtol=rtol, atol=atol) np.testing.assert_allclose(x_numpy, y_numpy, rtol=rtol, atol=atol)
@clear_cache_before_run()
@pytest.mark.parametrize("BATCH_SIZE", [1, 4, 7, 32]) @pytest.mark.parametrize("BATCH_SIZE", [1, 4, 7, 32])
@pytest.mark.parametrize("BLOCK_SIZE", [8, 16, 32]) @pytest.mark.parametrize("BLOCK_SIZE", [8, 16, 32])
@pytest.mark.parametrize("MAX_NUM_BLOCKS_PER_SEQ", [1, 8, 32, 256, 512]) @pytest.mark.parametrize("MAX_NUM_BLOCKS_PER_SEQ", [1, 8, 32, 256, 512])
@ -196,7 +197,7 @@ except ImportError:
HAS_VLLM = False HAS_VLLM = False
print("The subsequent test requires vllm. Please refer to https://github.com/vllm-project/vllm") print("The subsequent test requires vllm. Please refer to https://github.com/vllm-project/vllm")
@clear_cache_before_run()
@pytest.mark.skipif(not HAS_VLLM, reason="requires vllm") @pytest.mark.skipif(not HAS_VLLM, reason="requires vllm")
@pytest.mark.parametrize("BATCH_SIZE", [1, 7, 32]) @pytest.mark.parametrize("BATCH_SIZE", [1, 7, 32])
@pytest.mark.parametrize("BLOCK_SIZE", [6, 32]) @pytest.mark.parametrize("BLOCK_SIZE", [6, 32])