Skip to content

GitLab

Unverified Commit 38b792aa authored by flybird1111's avatar flybird1111 Committed by GitHub
Browse files

[coloattention] fix import error (#4380) 

fixed an import error
parent 25c57b9f
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 1 deletion
colossalai/kernel/cuda_native/mha/__init__.py 0 → 100644
View file @ 38b792aa
from .mha import ColoAttention
__all__ = ['ColoAttention']
tests/test_utils/test_flash_attention.py
View file @ 38b792aa
...@@ -9,7 +9,7 @@ from colossalai.kernel.cuda_native.mha.mem_eff_attn import HAS_MEM_EFF_ATTN ...@@ -9,7 +9,7 @@ from colossalai.kernel.cuda_native.mha.mem_eff_attn import HAS_MEM_EFF_ATTN
from colossalai.testing import clear_cache_before_run, parameterize from colossalai.testing import clear_cache_before_run, parameterize
if HAS_MEM_EFF_ATTN or HAS_FLASH_ATTN: if HAS_MEM_EFF_ATTN or HAS_FLASH_ATTN:
from colossalai.kernel.cuda_native.mha.mha import ColoAttention from colossalai.kernel.cuda_native import ColoAttention
from colossalai.kernel.cuda_native.scaled_softmax import AttnMaskType from colossalai.kernel.cuda_native.scaled_softmax import AttnMaskType
DTYPE = [torch.float16, torch.bfloat16, torch.float32] DTYPE = [torch.float16, torch.bfloat16, torch.float32]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment