ColossalAI/colossalai/_C/scaled_upper_triang_masked_softmax.pyi

from torch import Tensor

def forward(input: Tensor, scale: float) -> Tensor:
    ...


def backward(output_grads: Tensor, softmax_results: Tensor, scale: float) -> Tensor:
    ...