ColossalAI/tests/test_infer/test_kernels/__init__.py