github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Fix test import 

Signed-off-by: char-1ee <xingjianli59@gmail.com>
pull/5771/head
char-1ee 2024-06-10 02:03:30 +00:00
parent f5981e808e
commit b303976a27
3 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -4,7 +4,7 @@ import numpy as np
import pytest
import torch
from colossalai.inference.modeling.models.nopadding_baichuan import get_alibi_slopes
from colossalai.inference.utils import get_alibi_slopes
from colossalai.kernel.kernel_loader import InferenceOpsLoader
from colossalai.utils import get_current_device
from tests.test_infer.test_kernels.triton.test_context_attn_unpad import generate_alibi_mask

2
tests/test_infer/test_kernels/triton/test_context_attn_unpad.py
View File

@ -2,7 +2,7 @@ import pytest
import torch
from packaging import version
from colossalai.inference.modeling.models.nopadding_baichuan import get_alibi_slopes
from colossalai.inference.utils import get_alibi_slopes
from colossalai.kernel.triton import context_attention_unpadded
from colossalai.utils import get_current_device
from tests.test_infer.test_kernels.triton.kernel_utils import (

2
tests/test_infer/test_kernels/triton/test_decoding_attn.py
View File

@ -3,7 +3,7 @@ import pytest
import torch
from packaging import version
from colossalai.inference.modeling.models.nopadding_baichuan import get_alibi_slopes
from colossalai.inference.utils import get_alibi_slopes
from colossalai.kernel.triton import flash_decoding_attention
from colossalai.utils import get_current_device
from tests.test_infer.test_kernels.triton.kernel_utils import (