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

fix

pull/6061/head
wangbluo 2024-09-13 07:06:14 +00:00
parent dc032172c3
commit 0b14a5512e
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/shardformer/layer/attn.py
View File

@ -118,7 +118,7 @@ class ColoAttention:
)
if size >= MEMORY_BOUND:
FlashAttentionDaoLoader().load()
flash_kernel = FlashAttentionDaoLoader().load()
# lazy load
if isinstance(ColoAttention._kernel_dispatch_map[dtype][mask_type], KernelLoader):
ColoAttention._kernel_dispatch_map[dtype][mask_type] = ColoAttention._kernel_dispatch_map[dtype][
@ -126,7 +126,7 @@ class ColoAttention:
].load()
if size >= MEMORY_BOUND and mask_type in (AttnMaskType.PADDED_CAUSAL, AttnMaskType.CAUSAL):
return FlashAttentionDaoLoader()
return flash_kernel
else:
return ColoAttention._kernel_dispatch_map[dtype][mask_type]