support grouped query attention(MQA & GQA) for flash_attn (#22)

* support grouped query attention(GQA) for flash_attn(fwd, bwd, split_kv, total_attention)
* add mqa/gqa into feature list; update documentations and testings for flash attention.

README.md

@@ -233,6 +233,7 @@ The performance of piecewise_attention has improved compared to that in v0.1. In
 - the sequence length of k/v can be different from that of q;
 - support computation of total attention of each `k` gets from all `q`'s;
 - supports returning accumulative attention of each keys.
+- supports [MQA](https://arxiv.org/abs/1911.02150) and [GQA](https://arxiv.org/pdf/2305.13245).
 
 #### Limitations
 

src/flag_attn/flash.py

@@ -8,22 +8,40 @@
 
 __all__ = ["attention"]
 
+
+def maybe_contiguous(x):
+    # only when the inner most dimension is contiguous can LDGSTS be used
+    # so inner-dimension contiguity is enforced.
+    return x.contiguous() if x.stride(-1) != 1 else x
+
+def rounded_multiple(a, b):
+    return (a + b - 1) // b * b
+
 # --------------------------- public API ---------------------------
 class FlashAttention(torch.autograd.Function):
     @staticmethod
     def forward(ctx, q, k, v, causal, sm_scale, return_log_normalizer, return_total_attention):
+        # size, stride, dtype checking
         Dq, Dk, Dv = q.shape[-1], k.shape[-1], v.shape[-1]
-        assert Dq == Dk == Dv
+        assert Dq == Dk == Dv, "feature size of q, k, v should be equal"
         assert Dk in {16, 32, 64, 128}
 
         B, H, M, D = q.shape
         N = k.shape[2]
+        Hk, Hv = k.shape[1], v.shape[1]
+        assert Hk == Hv, "num of heads in k and v should be equal"
+        assert H % Hk == 0, "number of heads in q must be a multiple of that in k & v"
+        num_groups = H // Hk
+
         P_SEQ = N - M
         larger_m = M > N
 
         if sm_scale is None:
             sm_scale = 1. / math.sqrt(D)
 
+        # contiguity
+        q, k, v = maybe_contiguous(q), maybe_contiguous(k), maybe_contiguous(v)
+
         # to work around https://github.com/openai/triton/issues/2441
         device = torch.cuda.device_of(q)
         num_sms = torch.cuda.get_device_properties(device).multi_processor_count
@@ -32,6 +50,7 @@ def forward(ctx, q, k, v, causal, sm_scale, return_log_normalizer, return_total_
             config_for_split_kv = get_fwd_config_kv_split(B, H, M, N, D, causal)
             S = num_splits_herustic(B, H, M, N, config_for_split_kv[0], config_for_split_kv[1], num_sms, 128)
             split_kv: bool = S > 1
+            # print(f"flag_attn choose {S} splits")
 
             if not split_kv:
                 config = get_fwd_config(B, H, M, N, D, causal)
@@ -50,7 +69,7 @@ def forward(ctx, q, k, v, causal, sm_scale, return_log_normalizer, return_total_
                     k.stride(0), k.stride(1), k.stride(2), k.stride(3),
                     v.stride(0), v.stride(1), v.stride(2), v.stride(3),
                     o.stride(0), o.stride(1), o.stride(2), o.stride(3),
-                    B, H, M, N, P_SEQ,
+                    B, H, M, N, P_SEQ, num_groups,
                     BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_DMODEL=D,
                     IS_CAUSAL=causal, LARGER_M=larger_m,
                     DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,
@@ -61,7 +80,6 @@ def forward(ctx, q, k, v, causal, sm_scale, return_log_normalizer, return_total_
 
                 divisible_m = M % BLOCK_M == 0
                 divisible_n = N % BLOCK_N == 0
-
                 # consider using 3d grid to avoid div & rem
                 multiple_l = torch.empty((B, H, S, M), dtype=torch.float32, device="cuda")
                 multiple_o = torch.empty((B, H, S, M, D), dtype=torch.float16, device="cuda")
@@ -74,7 +92,7 @@ def forward(ctx, q, k, v, causal, sm_scale, return_log_normalizer, return_total_
                     k.stride(0), k.stride(1), k.stride(2), k.stride(3),
                     v.stride(0), v.stride(1), v.stride(2), v.stride(3),
                     multiple_o.stride(0), multiple_o.stride(1), multiple_o.stride(2), multiple_o.stride(3), multiple_o.stride(4),
-                    B, H, M, N, P_SEQ, N_SPLIT_SIZE, S,
+                    B, H, M, N, P_SEQ, N_SPLIT_SIZE, S, num_groups,
                     BLOCK_M=BLOCK_M, BLOCK_DMODEL=D, BLOCK_N=BLOCK_N,
                     IS_CAUSAL=causal, LARGER_M=larger_m,
                     DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,
@@ -103,7 +121,7 @@ def forward(ctx, q, k, v, causal, sm_scale, return_log_normalizer, return_total_
                     q, k, L, tot_attn, sm_scale,
                     q.stride(0), q.stride(1), q.stride(2), q.stride(3),
                     k.stride(0), k.stride(1), k.stride(2), k.stride(3),
-                    B, H, M, N, P_SEQ,
+                    B, H, M, N, P_SEQ, num_groups,
                     BLOCK_M=BLOCK_M, BLOCK_DMODEL=D, BLOCK_N=BLOCK_N,
                     CAUSAL=causal,
                     DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,
@@ -133,6 +151,8 @@ def backward(ctx, do, *ignored):
 
         B, H, M, D = q.shape
         N = k.shape[2]
+        Hk = k.shape[1]
+        num_groups = H // Hk
         P_SEQ = N - M
         larger_m = M > N
 
@@ -161,8 +181,9 @@ def backward(ctx, do, *ignored):
                 DIVISIBLE_M=divisible_m,
             )
 
-            dk = torch.empty_like(k)
-            dv = torch.empty_like(v)
+            # NOTE that dk & dv always have the same number of heads as q, instead of q.
+            dk = torch.empty((B, H, N, D), dtype=k.dtype, device=q.device)
+            dv = torch.empty((B, H, N, D), dtype=v.dtype, device=q.device)
             grid = (triton.cdiv(N, BLOCK_N), H, B)
             _bwd_kv_kernel[grid](
                 q, k, v, sm_scale, do,
@@ -175,6 +196,7 @@ def backward(ctx, do, *ignored):
                 dk.stride(0), dk.stride(1), dk.stride(2), dk.stride(3),
                 dv.stride(0), dv.stride(1), dv.stride(2), dv.stride(3),
                 B, H, M, N, P_SEQ,
+                num_groups,
                 BLOCK_M=BLOCK_M, BLOCK_DMODEL=D, BLOCK_N=BLOCK_N, CAUSAL=causal,
                 DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,
                 num_stages=num_stages, num_warps=num_warps,
@@ -192,12 +214,14 @@ def backward(ctx, do, *ignored):
                 do.stride(0), do.stride(1), do.stride(2), do.stride(3),
                 dq.stride(0), dq.stride(1), dq.stride(2), dq.stride(3),
                 B, H, M, N, P_SEQ,
+                num_groups,
                 BLOCK_M=BLOCK_M, BLOCK_DMODEL=D, BLOCK_N=BLOCK_N,
                 CAUSAL=causal, LARGER_M=larger_m,
                 DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,
                 num_stages=num_stages, num_warps = num_warps,
             )
-
+            dk = dk.reshape((B, Hk, num_groups, N, D)).sum(2)
+            dv = dv.reshape((B, Hk, num_groups, N, D)).sum(2)
         return dq, dk, dv, None, None, None, None
 
 
@@ -208,21 +232,24 @@ def attention(q, k, v, causal=False, sm_scale=None,
     An implementation of FlashAttention v2(https://arxiv.org/abs/2307.08691).
 
     Arguments:
-        q(torch.Tensor): The first queries. The shape is (batch_size, nheads, seqlen_q, headdim).
-        k(torch.Tensor): The first keys. The shape is (batch_size, nheads, seqlen_k, headdim).
-        v(torch.Tensor): The values. The shape is (batch_size, nheads, seqlen_k, headdim).
+        q(torch.Tensor): The first queries. The shape is (batch_size, num_heads_q, seqlen_q, headdim).
+        k(torch.Tensor): The first keys. The shape is (batch_size, num_heads_k, seqlen_k, headdim).
+        v(torch.Tensor): The values. The shape is (batch_size, num_heads_k, seqlen_k, headdim).
         causal(bool): Whether causal masking is applied to attention scores before applying softmax.
         sm_scale(float): The scaling of attention scores before applying softmax.
         return_log_normalizer(bool): Whether to return the log normalizer of softmax inside attention.
         return_total_attention(bool): Whether to return the sum of attention along q's sequence dimendion.
 
     Returns:
-        out(torch.Tensor): The output. The shape is (batch_size, nheads, seqlen_q, headdim).
+        out(torch.Tensor): The output. The shape is (batch_size, num_heads_q, seqlen_q, headdim).
 
         If `return_log_normalizer` or `return_total_attention`, return the following results in addition.
 
-        log_normalizer(torch.Tensor): The log normalizer. The shape is (batch_size, nheads, seqlen_q).
-        total_attention(torch.Tensor): The total attention. The shape is (batch_size, nheads, seqlen_k).
+        log_normalizer(torch.Tensor): The log normalizer. The shape is (batch_size, num_heads_q, seqlen_q).
+        total_attention(torch.Tensor): The total attention. The shape is (batch_size, num_heads_q, seqlen_k).
+
+    Notes:
+        `num_heads_q` must be a multiple of `num_heads_k`.
     """
     return FlashAttention.apply(q, k, v, causal, sm_scale, return_log_normalizer, return_total_attention)
 
@@ -272,6 +299,7 @@ def _fwd_kernel(
     stride_vz, stride_vh, stride_vn, stride_vk,
     stride_oz, stride_oh, stride_om, stride_ok,
     Z, H, M, N, P_SEQ,
+    num_groups,
     BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
     IS_CAUSAL: tl.constexpr, LARGER_M: tl.constexpr,
     DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr,
@@ -289,9 +317,10 @@ def _fwd_kernel(
     qk_scale = sm_scale * log2e
 
     # offset pointers for (batch, head)
+    off_hk = off_h // num_groups
     Q += off_z * stride_qz + off_h * stride_qh
-    K += off_z * stride_kz + off_h * stride_kh
-    V += off_z * stride_vz + off_h * stride_vh
+    K += off_z * stride_kz + off_hk * stride_kh
+    V += off_z * stride_vz + off_hk * stride_vh
     O += off_z * stride_oz + off_h * stride_oh
     L += (off_z * H + off_h) * M # l's shape is (B, H, M)
 
@@ -491,6 +520,7 @@ def _bwd_kv_kernel(
     stride_dkz, stride_dkh, stride_dkn, stride_dkk,
     stride_dvz, stride_dvh, stride_dvn, stride_dvk,
     Z, H, M, N, P_SEQ,
+    num_groups,
     BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
     CAUSAL: tl.constexpr,
     DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr,
@@ -504,9 +534,10 @@ def _bwd_kv_kernel(
     qk_scale = sm_scale * log2e
 
     # offset pointers for (batch, head)
+    off_hk = off_h // num_groups
     Q += off_z * stride_qz + off_h * stride_qh
-    K += off_z * stride_kz + off_h * stride_kh
-    V += off_z * stride_vz + off_h * stride_vh
+    K += off_z * stride_kz + off_hk * stride_kh
+    V += off_z * stride_vz + off_hk * stride_vh
     DO += off_z * stride_doz + off_h * stride_doh
 
     # offset pointers for batch/head
@@ -637,6 +668,7 @@ def _bwd_q_kernel(
     stride_doz, stride_doh, stride_dom, stride_dok,
     stride_dqz, stride_dqh, stride_dqm, stride_dqk,
     Z, H, M, N, P_SEQ,
+    num_groups,
     BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
     CAUSAL: tl.constexpr, LARGER_M: tl.constexpr,
     DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr,
@@ -654,9 +686,10 @@ def _bwd_q_kernel(
     qk_scale = sm_scale * log2e
 
     # offset pointers for (batch, head)
+    off_hk = off_h // num_groups
     Q += off_z * stride_qz + off_h * stride_qh
-    K += off_z * stride_kz + off_h * stride_kh
-    V += off_z * stride_vz + off_h * stride_vh
+    K += off_z * stride_kz + off_hk * stride_kh
+    V += off_z * stride_vz + off_hk * stride_vh
     DO += off_z * stride_doz + off_h * stride_doh
     D += (off_z * H + off_h) * M
     L += (off_z * H + off_h) * M

src/flag_attn/split_kv.py

@@ -19,7 +19,7 @@ def _fwd_split_kv_kernel(
     stride_kz, stride_kh, stride_kn, stride_kk,
     stride_vz, stride_vh, stride_vn, stride_vk,
     stride_oz, stride_oh, stride_os, stride_om, stride_ok,
-    Z, H, M, N, P_SEQ, N_SPLIT_SIZE, S,
+    Z, H, M, N, P_SEQ, N_SPLIT_SIZE, S, num_groups,
     BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
     IS_CAUSAL: tl.constexpr, LARGER_M: tl.constexpr,
     DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr,
@@ -31,6 +31,7 @@ def _fwd_split_kv_kernel(
     off_zh = tl.program_id(2)
     off_h = off_zh % H
     off_z = off_zh // H
+    off_hk = off_h // num_groups
 
     # scale sm_scale by log_2(e) and use
     # 2^x instead of exp in the loop because CSE and LICM
@@ -40,8 +41,8 @@ def _fwd_split_kv_kernel(
 
     # offset pointers for (batch & head)
     Q += off_z * stride_qz + off_h * stride_qh
-    K += off_z * stride_kz + off_h * stride_kh
-    V += off_z * stride_vz + off_h * stride_vh
+    K += off_z * stride_kz + off_hk * stride_kh
+    V += off_z * stride_vz + off_hk * stride_vh
 
     # offset pointers for (batch & head, split)
     O += off_z * stride_oz + off_h * stride_oh + n_split_id * stride_os # o's shape is (B, H, S, M, D)
@@ -269,6 +270,10 @@ def attention(q, k, v, causal=False, sm_scale=None):
 
     B, H, M, D = q.shape
     N = k.shape[2]
+    Hk, Hv = k.shape[1], v.shape[1]
+    assert Hk == Hv, "num of heads in k and v should be equal"
+    assert H % Hk == 0, "number of heads in q must be a multiple of that in k & v"
+    num_groups = H // Hk
     P_SEQ = N - M
     larger_m = M > N
 
@@ -299,7 +304,7 @@ def attention(q, k, v, causal=False, sm_scale=None):
             k.stride(0), k.stride(1), k.stride(2), k.stride(3),
             v.stride(0), v.stride(1), v.stride(2), v.stride(3),
             multiple_o.stride(0), multiple_o.stride(1), multiple_o.stride(2), multiple_o.stride(3), multiple_o.stride(4),
-            B, H, M, N, P_SEQ, N_SPLIT_SIZE, S,
+            B, H, M, N, P_SEQ, N_SPLIT_SIZE, S, num_groups,
             BLOCK_M=BLOCK_M, BLOCK_DMODEL=D, BLOCK_N=BLOCK_N,
             IS_CAUSAL=causal, LARGER_M=larger_m,
             DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,

src/flag_attn/testing/flash.py

@@ -16,6 +16,15 @@ def attention(q,
     D = q.shape[-1]
     if sm_scale is None:
         sm_scale = 1. / math.sqrt(D)
+
+    num_heads_q = q.shape[1]
+    num_heads_k = k.shape[1]
+    assert num_heads_q % num_heads_k == 0
+    num_groups = num_heads_q // num_heads_k
+
+    if num_groups > 1:
+        k = torch.repeat_interleave(k, repeats=num_groups, dim=1)
+        v = torch.repeat_interleave(v, repeats=num_groups, dim=1)
     kv_seq_len = k.shape[-2]
     q_seq_len = q.shape[-2]
     p_seq = kv_seq_len - q_seq_len

src/flag_attn/total.py

@@ -14,6 +14,10 @@ def total_attention(q, k, l, causal=False, sm_scale=None):
 
     B, H, M, D = q.shape
     N = k.shape[2]
+    Hk = k.shape[1]
+    assert H % Hk == 0, "number of heads in q must be a multiple of that in k"
+    num_groups = H // Hk
+
     P_SEQ = N - M
 
     if sm_scale is None:
@@ -34,7 +38,7 @@ def total_attention(q, k, l, causal=False, sm_scale=None):
             q, k, l, tot_attn, sm_scale,
             q.stride(0), q.stride(1), q.stride(2), q.stride(3),
             k.stride(0), k.stride(1), k.stride(2), k.stride(3),
-            B, H, M, N, P_SEQ,
+            B, H, M, N, P_SEQ, num_groups,
             BLOCK_M=BLOCK_M, BLOCK_DMODEL=D, BLOCK_N=BLOCK_N,
             CAUSAL=causal,
             DIVISIBLE_M=divisible_m, DIVISIBLE_N=divisible_n,
@@ -48,7 +52,7 @@ def _total_attention_kernel(
     Q, K, L, TA, sm_scale,
     stride_qz, stride_qh, stride_qm, stride_qk,
     stride_kz, stride_kh, stride_kn, stride_kk,
-    Z, H, M, N, P_SEQ,
+    Z, H, M, N, P_SEQ, num_groups,
     BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
     CAUSAL: tl.constexpr,
     DIVISIBLE_M: tl.constexpr, DIVISIBLE_N: tl.constexpr,
@@ -61,8 +65,9 @@ def _total_attention_kernel(
     qk_scale = sm_scale * log2e
 
     # offset pointers for (batch, head)
+    off_hk = off_h // num_groups
     Q += off_z * stride_qz + off_h * stride_qh
-    K += off_z * stride_kz + off_h * stride_kh
+    K += off_z * stride_kz + off_hk * stride_kh
     L += (off_z * H + off_h) * M
     TA += (off_z * H + off_h) * N # (b, h, n)