replicate/flash-attn2

1import math2 3import torch4from einops import rearrange, repeat5from flash_attn.bert_padding import pad_input, unpad_input6 7 8def generate_random_padding_mask(max_seqlen, batch_size, device, mode="random", zero_lengths=False):9    assert mode in ["full", "random", "third"]10    if mode == "full":11        lengths = torch.full((batch_size, 1), max_seqlen, device=device, dtype=torch.int32)12    elif mode == "random":13        lengths = torch.randint(14            max(0 if zero_lengths else 1, max_seqlen - 20), max_seqlen + 1, (batch_size, 1), device=device15        )16    elif mode == "third":17        lengths = torch.randint(max_seqlen // 3, max_seqlen + 1, (batch_size, 1), device=device)18 19    if zero_lengths:20        # Generate zero-lengths every 5 batches and the last batch.21        for i in range(batch_size):22            if i % 5 == 0:23                lengths[i] = 024        lengths[-1] = 025    padding_mask = (26        repeat(torch.arange(max_seqlen, device=device), "s -> b s", b=batch_size) < lengths27    )28    return padding_mask29 30 31def generate_qkv(32    q, k, v, query_padding_mask=None, key_padding_mask=None, 33    kvpacked=False, qkvpacked=False, add_unused_qkv=False,34    query_unused_mask=None, key_unused_mask=None,35):36    """37    Arguments:38        q: (batch_size, seqlen_q, nheads, d)39        k: (batch_size, seqlen_k, nheads_k, d)40        v: (batch_size, seqlen_k, nheads_k, d)41        query_padding_mask: (batch_size, seqlen), bool42        key_padding_mask: (batch_size, seqlen), bool43    """44    assert not (kvpacked and qkvpacked)45    batch_size, seqlen_q, nheads, d = q.shape46    _, seqlen_k, nheads_k, _ = k.shape47    assert k.shape == (batch_size, seqlen_k, nheads_k, d)48    assert v.shape == (batch_size, seqlen_k, nheads_k, d)49    if query_unused_mask is not None or key_unused_mask is not None:50        assert not kvpacked51        assert not qkvpacked52 53    if query_padding_mask is not None:54        q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, seqused_q = unpad_input(55            q, query_padding_mask, query_unused_mask,56        )57        output_pad_fn = lambda output_unpad: pad_input(58            output_unpad, indices_q, batch_size, seqlen_q59        )60    else:61        q_unpad = rearrange(q, "b s h d -> (b s) h d")62        cu_seqlens_q = torch.arange(63            0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, device=q_unpad.device64        )65        seqused_q = None66        max_seqlen_q = seqlen_q67        output_pad_fn = lambda output_unpad: rearrange(68            output_unpad, "(b s) h d -> b s h d", b=batch_size69        )70 71    if key_padding_mask is not None:72        k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, seqused_k = unpad_input(k, key_padding_mask, key_unused_mask)73        v_unpad, _, _, _, _ = unpad_input(v, key_padding_mask, key_unused_mask)74    else:75        k_unpad = rearrange(k, "b s h d -> (b s) h d")76        v_unpad = rearrange(v, "b s h d -> (b s) h d")77        cu_seqlens_k = torch.arange(78            0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, device=k_unpad.device79        )80        seqused_k = None81        max_seqlen_k = seqlen_k82 83    if qkvpacked:84        assert (query_padding_mask == key_padding_mask).all()85        assert nheads == nheads_k86        qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1)87        qkv = torch.stack([q, k, v], dim=2)88        if query_padding_mask is not None:89            dqkv_pad_fn = lambda dqkv_unpad: pad_input(dqkv_unpad, indices_q, batch_size, seqlen_q)90        else:91            dqkv_pad_fn = lambda dqkv_unpad: rearrange(92                dqkv_unpad, "(b s) t h d -> b s t h d", b=batch_size93            )94        return (95            qkv_unpad.detach().requires_grad_(),96            cu_seqlens_q,97            max_seqlen_q,98            qkv.detach().requires_grad_(),99            output_pad_fn,100            dqkv_pad_fn,101        )102    elif kvpacked:103        kv_unpad = torch.stack([k_unpad, v_unpad], dim=1)104        kv = torch.stack([k, v], dim=2)105        dq_pad_fn = output_pad_fn106        if key_padding_mask is not None:107            dkv_pad_fn = lambda dkv_unpad: pad_input(dkv_unpad, indices_k, batch_size, seqlen_k)108        else:109            dkv_pad_fn = lambda dkv_unpad: rearrange(110                dkv_unpad, "(b s) t h d -> b s t h d", b=batch_size111            )112        return (113            q_unpad.detach().requires_grad_(),114            kv_unpad.detach().requires_grad_(),115            cu_seqlens_q,116            cu_seqlens_k,117            max_seqlen_q,118            max_seqlen_k,119            q.detach().requires_grad_(),120            kv.detach().requires_grad_(),121            output_pad_fn,122            dq_pad_fn,123            dkv_pad_fn,124        )125    else:126        dq_pad_fn = output_pad_fn127        if key_padding_mask is not None:128            dk_pad_fn = lambda dk_unpad: pad_input(dk_unpad, indices_k, batch_size, seqlen_k)129        else:130            dk_pad_fn = lambda dk_unpad: rearrange(dk_unpad, "(b s) h d -> b s h d", b=batch_size)131        return (132            q_unpad.detach().requires_grad_(),133            k_unpad.detach().requires_grad_(),134            v_unpad.detach().requires_grad_(),135            cu_seqlens_q,136            cu_seqlens_k,137            seqused_q,138            seqused_k,139            max_seqlen_q,140            max_seqlen_k,141            q.detach().requires_grad_(),142            k.detach().requires_grad_(),143            v.detach().requires_grad_(),144            output_pad_fn,145            dq_pad_fn,146            dk_pad_fn,147        )148 149 150def construct_local_mask(151    seqlen_q,152    seqlen_k,153    window_size=(-1, -1),  # -1 means infinite window size154    query_padding_mask=None,155    key_padding_mask=None,156    device=None,157    key_leftpad=None,158):159    row_idx = rearrange(torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1")160    col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long)161    if key_leftpad is not None:162        key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1")163        col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0])164        col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32)165    sk = (166        seqlen_k167        if key_padding_mask is None168        else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1")169    )170    sq = (171        seqlen_q172        if query_padding_mask is None173        else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1")174    )175    if window_size[0] < 0:176        return col_idx > row_idx + sk - sq + window_size[1]177    else:178        sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk179        return torch.logical_or(180            col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk),181            col_idx < row_idx + sk - sq - window_size[0],182        )183 184 185def attention_ref(186    q,187    k,188    v,189    query_padding_mask=None,190    key_padding_mask=None,191    attn_bias=None,192    dropout_p=0.0,193    dropout_mask=None,194    causal=False,195    window_size=(-1, -1),  # -1 means infinite window size196    softcap=0.0,197    upcast=True,198    reorder_ops=False,199    key_leftpad=None,200):201    """202    Arguments:203        q: (batch_size, seqlen_q, nheads, head_dim)204        k: (batch_size, seqlen_k, nheads_k, head_dim)205        v: (batch_size, seqlen_k, nheads_k, head_dim)206        query_padding_mask: (batch_size, seqlen_q)207        key_padding_mask: (batch_size, seqlen_k)208        attn_bias: broadcastable to (batch_size, nheads, seqlen_q, seqlen_k)209        dropout_p: float210        dropout_mask: (batch_size, nheads, seqlen_q, seqlen_k)211        causal: whether to apply causal masking212        window_size: (int, int), left and right window size213        upcast: whether to cast all inputs to fp32, do all computation in fp32, then cast214            output back to fp16/bf16.215        reorder_ops: whether to change the order of operations (scaling k instead of scaling q, etc.)216            without changing the math. This is to estimate the numerical error from operation217            reordering.218    Output:219        output: (batch_size, seqlen_q, nheads, head_dim)220        attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout221    """222    if causal:223        window_size = (window_size[0], 0)224    dtype_og = q.dtype225    if upcast:226        q, k, v = q.float(), k.float(), v.float()227    seqlen_q, seqlen_k = q.shape[1], k.shape[1]228    k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2])229    v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2])230    d = q.shape[-1]231    if not reorder_ops:232        scores = torch.einsum("bthd,bshd->bhts", q / math.sqrt(d), k)233    else:234        scores = torch.einsum("bthd,bshd->bhts", q, k / math.sqrt(d))235    if softcap > 0:236        scores /= softcap237        scores = scores.tanh()238        scores *= softcap239    if key_padding_mask is not None:240        scores.masked_fill_(rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf"))241    if window_size[0] >= 0 or window_size[1] >= 0:242        local_mask = construct_local_mask(243            seqlen_q,244            seqlen_k,245            window_size,246            query_padding_mask,247            key_padding_mask,248            q.device,249            key_leftpad=key_leftpad,250        )251        scores.masked_fill_(local_mask, float("-inf"))252    if attn_bias is not None:253        scores = scores + attn_bias254    attention = torch.softmax(scores, dim=-1).to(v.dtype)255    # Some rows might be completely masked out so we fill them with zero instead of NaN256    if window_size[0] >= 0 or window_size[1] >= 0:257        attention = attention.masked_fill(torch.all(local_mask, dim=-1, keepdim=True), 0.0)258    # We want to mask here so that the attention matrix doesn't have any NaNs259    # Otherwise we'll get NaN in dV260    if query_padding_mask is not None:261        attention = attention.masked_fill(rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0)262    dropout_scaling = 1.0 / (1 - dropout_p)263    # attention_drop = attention.masked_fill(~dropout_mask, 0.0) * dropout_scaling264    # output = torch.einsum('bhts,bshd->bthd', attention_drop , v)265    if dropout_mask is not None:266        attention_drop = attention.masked_fill(~dropout_mask, 0.0)267    else:268        attention_drop = attention269    output = torch.einsum("bhts,bshd->bthd", attention_drop, v * dropout_scaling)270    if query_padding_mask is not None:271        output.masked_fill_(rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0)272    if key_padding_mask is not None:273        output.masked_fill_(rearrange(torch.logical_not(torch.any(key_padding_mask, 1)), "b -> b 1 1 1"), 0.0)274    return output.to(dtype=dtype_og), attention.to(dtype=dtype_og)275