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

Code before:
            ctx.save_for_backward(indices)
            # [ntokens, ...]
            tensor = tensor.squeeze(0)
            out = torch.zeros((batch_size * seq_len, *tensor.shape[1:]), dtype=tensor.dtype, device=tensor.device)
            # [b*s, ...]
            out[indices] = tensor
            # [b, s, ...]
            out = rearrange(out, '(b s) ... -> b s ...', b=batch_size)
            return out

        @staticmethod
        def backward(ctx, grad_output):
            indices, = ctx.saved_tensors
            # [b*s, ...]
            grad_output = rearrange(grad_output, 'b s ... -> (b s) ...')
            grad = grad_output[indices]
            # [1, ntokens, ...]
            return grad.unsqueeze(0), None, None, None

    class ColoAttention(torch.nn.Module):

        def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0, scale=1.0):
            super().__init__()
            assert embed_dim % num_heads == 0, \
                f"the embed dim ({embed_dim}) is not divisible by the number of attention heads ({num_heads})."
            self.scale = scale
            self.dropout = dropout

        @staticmethod
        def get_seq_info_from_mask(attn_mask: torch.Tensor):
            indices = torch.nonzero(attn_mask.flatten(), as_tuple=False).flatten()
            seqlens = attn_mask.sum(dim=-1, dtype=torch.int32).flatten().tolist()
            return indices, seqlens

        @staticmethod
        def unpad(tensor: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
            return Unpad.apply(tensor, indices)

        @staticmethod
        def repad(tensor: torch.Tensor, indices: torch.Tensor, batch_size: int, seq_len: int) -> torch.Tensor:
            return Repad.apply(tensor, indices, batch_size, seq_len)

        def forward(self,
                    query: torch.Tensor,
                    key: torch.Tensor,
                    value: torch.Tensor,
                    attn_mask: Optional[torch.Tensor] = None,
                    attn_mask_type: Optional[AttnMaskType] = None,
                    bias: Optional[torch.Tensor] = None):
            batch_size, tgt_len, src_len = query.shape[0], query.shape[1], key.shape[1]
            attn_bias = None

Review comment (suggestion):
How about check if the input scale is None, we use `1 / sqrt(embed_dim / num_heads)`, otherwise we set its value

Code after:
        https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py
        """

        @staticmethod
        def forward(ctx, tensor: torch.Tensor, indices: torch.Tensor, batch_size: int, seq_len: int):
            ctx.save_for_backward(indices)
            # [ntokens, ...]
            tensor = tensor.squeeze(0)
            out = torch.zeros((batch_size * seq_len, *tensor.shape[1:]), dtype=tensor.dtype, device=tensor.device)
            # [b*s, ...]
            out[indices] = tensor
            # [b, s, ...]
            out = rearrange(out, '(b s) ... -> b s ...', b=batch_size)
            return out

        @staticmethod
        def backward(ctx, grad_output):
            indices, = ctx.saved_tensors
            # [b*s, ...]
            grad_output = rearrange(grad_output, 'b s ... -> (b s) ...')
            grad = grad_output[indices]
            # [1, ntokens, ...]
            return grad.unsqueeze(0), None, None, None

    class ColoAttention(torch.nn.Module):

        def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0, scale=None):
            super().__init__()
            assert embed_dim % num_heads == 0, \
                f"the embed dim ({embed_dim}) is not divisible by the number of attention heads ({num_heads})."
            if scale is not None:
                self.scale = scale
            else:
                self.scale = 1 / math.sqrt(embed_dim // num_heads)
            self.dropout = dropout

        @staticmethod
        def get_seq_info_from_mask(attn_mask: torch.Tensor):
            indices = torch.nonzero(attn_mask.flatten(), as_tuple=False).flatten()
            seqlens = attn_mask.sum(dim=-1, dtype=torch.int32).flatten().tolist()
            return indices, seqlens

        @staticmethod
        def unpad(tensor: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
            return Unpad.apply(tensor, indices)

        @staticmethod
        def repad(tensor: torch.Tensor, indices: torch.Tensor, batch_size: int, seq_len: int) -> torch.Tensor:
            return Repad.apply(tensor, indices, batch_size, seq_len)

        def forward(self,