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

Code before:
                    kv_seqlen = None
                    if batch_size > 1:
                        query, key, value = self.unpad(torch.stack([query, key, value], dim=2), q_indices).unbind(dim=2)
                else:
                    q_indices = torch.arange(batch_size * tgt_len, dtype=torch.int32, device=query.device)
                    q_seqlen = torch.LongTensor([tgt_len] * batch_size, device=query.device)
                    kv_indices, kv_seqlen = self.get_seq_info_from_mask(attn_mask)
                    if batch_size > 1:
                        query = rearrange(query, "b s ... -> c (b s) ...", c=1)
                        key, value = self.unpad(torch.stack([query, key, value], dim=2), kv_indices).unbind(dim=2)
                if attn_mask_type == AttnMaskType.padding:
                    attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen)
                elif attn_mask_type == AttnMaskType.paddedcausal:
                    attn_bias = BlockDiagonalCausalMask.from_seqlens(q_seqlen, kv_seqlen)
            elif attn_mask_type == AttnMaskType.causal:    # gpt style
                attn_bias = LowerTriangularMask()

            if bias is not None:    # alibi / relative position embedding
                assert allow_alibi, "flash attention with bias is not supported in this system."
                assert attn_mask_type == AttnMaskType.causal, \
                    "attention with bias is only supported for causal attention so far."
                attn_bias = attn_bias.add_bias(bias)

            out = memory_efficient_attention(query, key, value, attn_bias=attn_bias, p=self.dropout, scale=self.scale)

            if attn_mask_type in [AttnMaskType.padding, AttnMaskType.paddedcausal] and batch_size > 1:
                out = self.repad(out, q_indices, batch_size, tgt_len)

            out = rearrange(out, 'b s h d -> b s (h d)')
            return out


##########################################################################
# the flash attention functions below that are copied
# from the OpenAI/triton repository will be deprecated
# You can find the repository in Triton https://github.com/openai/triton
# You can find the source file in https://github.com/openai/triton/blob/main/python/tutorials/06-fused-attention.py
# Reference:
# 1. Dao et al., https://arxiv.org/pdf/2205.14135v2.pdf
# 2. Rabe and Staats https://arxiv.org/pdf/2112.05682v2.pdf


def triton_cuda_check():
    cuda_home = os.getenv("CUDA_HOME", default="/usr/local/cuda")
    cuda_version = subprocess.check_output([os.path.join(cuda_home, "bin/nvcc"), "--version"]).decode().strip()
    cuda_version = cuda_version.split('release ')[1]
    cuda_version = cuda_version.split(',')[0]
    cuda_version = cuda_version.split('.')
    if len(cuda_version) == 2 and \
        (int(cuda_version[0]) == 11 and int(cuda_version[1]) >= 4) or \
        int(cuda_version[0]) > 11:

Review comment (suggestion):
How about this way to check if padded: `attn_mask_type.value % 2 == 1`

Code after:
                    kv_seqlen = None
                    if batch_size > 1:
                        query, key, value = self.unpad(torch.stack([query, key, value], dim=2), q_indices).unbind(dim=2)
                else:
                    q_indices = torch.arange(batch_size * tgt_len, dtype=torch.int32, device=query.device)
                    q_seqlen = torch.LongTensor([tgt_len] * batch_size, device=query.device)
                    kv_indices, kv_seqlen = self.get_seq_info_from_mask(attn_mask)
                    if batch_size > 1:
                        query = rearrange(query, "b s ... -> c (b s) ...", c=1)
                        key, value = self.unpad(torch.stack([query, key, value], dim=2), kv_indices).unbind(dim=2)
                if attn_mask_type == AttnMaskType.padding:
                    attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen)
                elif attn_mask_type == AttnMaskType.paddedcausal:
                    attn_bias = BlockDiagonalCausalMask.from_seqlens(q_seqlen, kv_seqlen)
            elif attn_mask_type == AttnMaskType.causal:    # gpt style
                attn_bias = LowerTriangularMask()

            if bias is not None:    # alibi / relative position embedding
                assert allow_alibi, "flash attention with bias is not supported in this system."
                assert attn_mask_type == AttnMaskType.causal, \
                    "attention with bias is only supported for causal attention so far."
                attn_bias = attn_bias.add_bias(bias)

            out = memory_efficient_attention(query, key, value, attn_bias=attn_bias, p=self.dropout, scale=self.scale)

            if attn_mask_type and attn_mask_type.value % 2 == 1 and batch_size > 1:
                out = self.repad(out, q_indices, batch_size, tgt_len)

            out = rearrange(out, 'b s h d -> b s (h d)')
            return out


##########################################################################
# the flash attention functions below that are copied
# from the OpenAI/triton repository will be deprecated
# You can find the repository in Triton https://github.com/openai/triton
# You can find the source file in https://github.com/openai/triton/blob/main/python/tutorials/06-fused-attention.py
# Reference:
# 1. Dao et al., https://arxiv.org/pdf/2205.14135v2.pdf
# 2. Rabe and Staats https://arxiv.org/pdf/2112.05682v2.pdf


def triton_cuda_check():
    cuda_home = os.getenv("CUDA_HOME", default="/usr/local/cuda")
    cuda_version = subprocess.check_output([os.path.join(cuda_home, "bin/nvcc"), "--version"]).decode().strip()
    cuda_version = cuda_version.split('release ')[1]
    cuda_version = cuda_version.split(',')[0]
    cuda_version = cuda_version.split('.')
    if len(cuda_version) == 2 and \
        (int(cuda_version[0]) == 11 and int(cuda_version[1]) >= 4) or \
        int(cuda_version[0]) > 11: