megatron.model.fused_softmax

Description

Classes

FusedScaleMaskSoftmax(input_in_fp16, ...)	fused operation: scaling + mask + softmax
ScaledMaskedSoftmax(*args, **kwargs)	Fused operation which performs following three operations in sequence 1.
ScaledSoftmax(*args, **kwargs)	Fused operation which performs following two operations in sequence 1.
ScaledUpperTriangMaskedSoftmax(*args, **kwargs)	Fused operation which performs following three operations in sequence 1.