github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-12 12:47:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

[builder] builder for scaled_upper_triang_masked_softmax (#2234)

Browse Source
This commit is contained in:
Jiarui Fang
2022-12-30 09:58:00 +08:00
committed by GitHub
parent 31fe84237b
commit db4cbdc7fb
6 changed files with 53 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -23,27 +23,20 @@ class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
@staticmethod
def forward(ctx, inputs, scale):
try:
import colossalai._C.scaled_upper_triang_masked_softmax
except ImportError:
raise RuntimeError('ScaledUpperTriangMaskedSoftmax requires cuda extensions')
from colossalai.kernel import scaled_upper_triang_masked_softmax
scale_t = torch.tensor([scale])
softmax_results = colossalai._C.scaled_upper_triang_masked_softmax.forward(inputs, scale_t[0])
softmax_results = scaled_upper_triang_masked_softmax.forward(inputs, scale_t[0])
ctx.save_for_backward(softmax_results, scale_t)
return softmax_results
@staticmethod
def backward(ctx, output_grads):
try:
import colossalai._C.scaled_upper_triang_masked_softmax
except ImportError:
raise RuntimeError('ScaledUpperTriangMaskedSoftmax requires cuda extensions')
from colossalai.kernel import scaled_upper_triang_masked_softmax
softmax_results, scale_t = ctx.saved_tensors
input_grads = colossalai._C.scaled_upper_triang_masked_softmax.backward(output_grads, softmax_results,
scale_t[0])
input_grads = scaled_upper_triang_masked_softmax.backward(output_grads, softmax_results, scale_t[0])
return input_grads, None