github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-04-11 22:54:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

[kernel] added kernel loader to softmax autograd function (#3093)

Browse Source 
* [kernel] added kernel loader to softmax autograd function

* [release] v0.2.6
This commit is contained in:
Frank Lee
2023-03-10 14:27:09 +08:00
committed by GitHub
parent fff98f06ed
commit 95a36eae63
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/kernel/cuda_native/scaled_softmax.py
View File

@@ -180,4 +180,9 @@ class FusedScaleMaskSoftmax(nn.Module):
return probs
def get_batch_per_block(self, sq, sk, b, np):
# build and load kernel if not pre-built
global scaled_masked_softmax
if scaled_masked_softmax is None:
scaled_masked_softmax = ScaledMaskedSoftmaxBuilder().load()
return scaled_masked_softmax.get_batch_per_block(sq, sk, b, np)