github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-06 11:32:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix the sp

Browse Source
This commit is contained in:
wangbluo
2024-09-13 02:32:03 +00:00
parent a35a078f08
commit fdd84b9087
2 changed files with 30 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/kernel/kernel_loader.py
View File

@@ -118,6 +118,8 @@ class FlashAttentionLoader(KernelLoader):
FlashAttentionSdpaCudaExtension,
]
class FlashAttentionDaoLoader(KernelLoader):
REGISTRY = [FlashAttentionDaoCudaExtension]
class FlashAttentionWithCustomMaskLoader(KernelLoader):
REGISTRY = [FlashAttentionNpuExtension, FlashAttentionSdpaCudaExtension]