Minor Question about attn_mask

[ZhikangNiu]

Thank you for this very meaningful work.
When I try flash-attn in Text to Speech DiT based model, I have a question about the attn mask

Here is my attn mask example (seq_len=10, The mask is text padding mask and mel spec padding mask, so it has different location padding)

False  False False True True True True True False False
False  False True True True True True True True False
True  True  True True True True True True True False
False True  True True True True True True True True

When I use bert_embedding.unpad_input, we hope the max_seq_len is 10 for pad_input after attention.

But I found that it will calculate the max seq length is 9, because

flash-attention/flash_attn/bert_padding.py

Line 113 in 0e79d71

used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)

So why not use attn_mask.size(-1) to get max_seq_len in a batch?

Here is my code, I change this line
I'm not very familiar with FlashAttention, so please point out any mistakes if there are any.

try:
    from flash_attn import flash_attn_varlen_func
    from flash_attn.bert_padding import unpad_input, pad_input, index_first_axis
except ImportError:
    flash_attn_varlen_func, pad_input, unpad_input = None, None, None
    
def unpad_input(hidden_states, attention_mask, unused_mask=None):
    """
    Arguments:
        hidden_states: (batch, seqlen, ...)
        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
        unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused.
    Return:
        hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask.
        indices: (total_nnz), the indices of masked tokens from the flattened input sequence.
        cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
        max_seqlen_in_batch: int
        seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask.
    """
    all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask
    seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32)
    used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
    indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten()
    # max_seqlen_in_batch = seqlens_in_batch.max().item()
    max_seqlen_in_batch = attention_mask.size(-1)
    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
    # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
    # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
    # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
    # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
    # so we write custom forward and backward to make it a bit faster.
    return (
        index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices),
        indices,
        cu_seqlens,
        max_seqlen_in_batch,
        used_seqlens_in_batch, 
    )


class FlashAttnProcessor:
    def __init__(
        self,
        pe_attn_head: int | None = None,  # number of attention head to apply rope, None for all
    ):
        self.pe_attn_head = pe_attn_head
        if flash_attn_varlen_func is None:
            raise ImportError("FlashAttnProcessor requires flash-attn, please install it with `pip install flash-attn`")

    def __call__(
        self,
        attn: Attention,
        x: float["b n d"],  # noised input x  # noqa: F722
        mask: bool["b n"] | None = None,  # noqa: F722
        rope=None,  # rotary position embedding
    ) -> torch.FloatTensor:
        batch_size = x.shape[0]

        # `sample` projections
        query = attn.to_q(x)
        key = attn.to_k(x)
        value = attn.to_v(x) # torch.Size([64, 169, 768])

        # attention
        inner_dim = key.shape[-1]
        head_dim = inner_dim // attn.heads
        query = query.view(batch_size, -1, attn.heads, head_dim) # b, t, n, d
        key = key.view(batch_size, -1, attn.heads, head_dim)
        value = value.view(batch_size, -1, attn.heads, head_dim)

        # qk norm
        if attn.q_norm is not None:
            query = attn.q_norm(query)
        if attn.k_norm is not None:
            key = attn.k_norm(key)

        # apply rotary position embedding
        if rope is not None:
            freqs, xpos_scale = rope
            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)

            if self.pe_attn_head is not None:
                pn = self.pe_attn_head
                query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale)
                key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale)
            else:
                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)

        # mask. e.g. inference got a batch with different target durations, mask out the padding
        if mask is not None:
            query, indices, q_cu_seqlens, q_max_seqlen_in_batch, _ = unpad_input(
                query, mask.bool()
            )
            key, _, k_cu_seqlens, k_max_seqlen_in_batch, _ = unpad_input(
                key, mask.bool()
            )
            value, _, _, _, _ = unpad_input(
                value, mask.bool()
            )
        x = flash_attn_varlen_func(
            query,
            key,
            value,
            q_cu_seqlens,
            k_cu_seqlens,
            q_max_seqlen_in_batch,
            k_max_seqlen_in_batch,
        )
        x = pad_input(x, indices, batch_size, q_max_seqlen_in_batch)
        x = x.view(batch_size, -1, attn.heads * head_dim)
        x = x.to(query.dtype)

        # linear proj
        x = attn.to_out[0](x)
        # dropout
        x = attn.to_out[1](x)
 
        if mask is not None:
            mask = mask.unsqueeze(-1)
            x = x.masked_fill(~mask, 0.0)

        return x

[ZhikangNiu]

ah, maybe