From 95a36eae637563e75484b73ccb6dc3feb56ce05b Mon Sep 17 00:00:00 2001 From: Frank Lee Date: Fri, 10 Mar 2023 14:27:09 +0800 Subject: [PATCH] [kernel] added kernel loader to softmax autograd function (#3093) * [kernel] added kernel loader to softmax autograd function * [release] v0.2.6 --- colossalai/kernel/cuda_native/scaled_softmax.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/colossalai/kernel/cuda_native/scaled_softmax.py b/colossalai/kernel/cuda_native/scaled_softmax.py index 05c6ee35b..24e458bb3 100644 --- a/colossalai/kernel/cuda_native/scaled_softmax.py +++ b/colossalai/kernel/cuda_native/scaled_softmax.py @@ -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)