Repository: hpcaitech/ColossalAI
File: colossalai/kernel/cuda_native/scaled_softmax.py (Python)

Code before:
    def is_kernel_available(self, mask, b, np, sq, sk):
        attn_batches = b * np

        if (self.scaled_masked_softmax_fusion    # user want to fuse
                and self.input_in_float16    # input must be fp16
                and mask is not None    # mask tensor must not be None
                and 16 < sk <= 2048    # sk must be 16 ~ 2048
                and sq % 4 == 0    # sq must be divisor of 4
                and attn_batches % 4 == 0    # np * b must be divisor of 4
           ):
            if 0 <= sk <= 2048:
                batch_per_block = self.get_batch_per_block(sq, sk, b, np)

                if self.attn_mask_type in [AttnMaskType.causal, AttnMaskType.paddedcausal] :
                    if attn_batches % batch_per_block == 0:
                        return True
                else:
                    if sq % batch_per_block == 0:
                        return True
        return False

    def forward_fused_softmax(self, input, mask):
        b, np, sq, sk = input.size()
        scale = self.scale if self.scale is not None else 1.0

        if self.attn_mask_type in [AttnMaskType.causal, AttnMaskType.paddedcausal]:
            assert sq == sk, "causal mask is only for self attention"

            # input is 3D tensor (attn_batches, sq, sk)
            input = input.view(-1, sq, sk)
            probs = ScaledUpperTriangMaskedSoftmax.apply(input, scale)
            return probs.view(b, np, sq, sk)
        else:
            # input is 4D tensor (b, np, sq, sk)
            return ScaledMaskedSoftmax.apply(input, mask, scale)

    def forward_torch_softmax(self, input, mask):
        if self.input_in_float16 and self.softmax_in_fp32:
            input = input.float()

        if self.scale is not None:
            input = input * self.scale
        mask_output = self.mask_func(input, mask) if mask is not None else input
        probs = torch.nn.Softmax(dim=-1)(mask_output)

        if self.input_in_float16 and self.softmax_in_fp32:
            if self.input_in_fp16:
                probs = probs.half()
            else:
                probs = probs.bfloat16()


Review comment (suggestion):
Similarly, check if causal by `attn_mask_type.value > 1`

Code after:
    def is_kernel_available(self, mask, b, np, sq, sk):
        attn_batches = b * np

        if (self.scaled_masked_softmax_fusion    # user want to fuse
                and self.input_in_float16    # input must be fp16
                and mask is not None    # mask tensor must not be None
                and 16 < sk <= 2048    # sk must be 16 ~ 2048
                and sq % 4 == 0    # sq must be divisor of 4
                and attn_batches % 4 == 0    # np * b must be divisor of 4
           ):
            if 0 <= sk <= 2048:
                batch_per_block = self.get_batch_per_block(sq, sk, b, np)

                if self.attn_mask_type.value > 1:
                    if attn_batches % batch_per_block == 0:
                        return True
                else:
                    if sq % batch_per_block == 0:
                        return True
        return False

    def forward_fused_softmax(self, input, mask):
        b, np, sq, sk = input.size()
        scale = self.scale if self.scale is not None else 1.0

        if self.attn_mask_type.value > 1:
            assert sq == sk, "causal mask is only for self attention"

            # input is 3D tensor (attn_batches, sq, sk)
            input = input.view(-1, sq, sk)
            probs = ScaledUpperTriangMaskedSoftmax.apply(input, scale)
            return probs.view(b, np, sq, sk)
        else:
            # input is 4D tensor (b, np, sq, sk)
            return ScaledMaskedSoftmax.apply(input, mask, scale)

    def forward_torch_softmax(self, input, mask):
        if self.input_in_float16 and self.softmax_in_fp32:
            input = input.float()

        if self.scale is not None:
            input = input * self.scale
        mask_output = self.mask_func(input, mask) if mask is not None else input
        probs = torch.nn.Softmax(dim=-1)(mask_output)

        if self.input_in_float16 and self.softmax_in_fp32:
            if self.input_in_fp16:
                probs = probs.half()
            else:
                probs = probs.bfloat16()