ColossalAI/tests/test_infer/test_kernels/cuda/__init__.py