Fix flash attention ops calculation (#758)

Author: azaidy

python/perf-kernels/flash-attention.py

@@ -1979,15 +1979,28 @@ def bench_flash_attention(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, causal
                 if causal:
                     # If seqlen_q != seqlen_k then the causal mask ignores computation
                     # depending on which seqlen is larger. Either the lower triangle, or right triangle
-                    causal_correction = seqlen_k if seqlen_q > seqlen_k else seqlen_q
-                    flops_per_matmul += (seqlen_q * seqlen_k - (causal_correction**2) / 2) * HQ * D_HEAD * 2
+                    # If seqlen_q is greater than seqlen_k, the lower triangle is non zero
+                    # where the last row has seqlen_k valid element, the second last row has
+                    # seqlen_k - 1 valid elements and so on until one element is valid in the
+                    # seqlen_q - seqlen_k row, hence total valid elements are 1+2+...+seqlen_k
+                    # which is seqlen_k*(seqlen_k+1)/2
+                    # If seqlen_q is less than seqlen_k, then we count the zero elements
+                    # the first row has seqlen_q-1 zero elements, the second row has seqlen_q-2
+                    # zero elements and so on until the second last row has 1 zero element
+                    # Total zero elements are 1+2+...+(seqlen_q-1) = seqlen_q*(seqlen_q-1)/2
+                    # Total non zero elements are seqlen_q*seqlen_k - (seqlen_q*(seqlen_q-1)/2)
+                    valid_out_elements = ((seqlen_k**2 + seqlen_k) / 2) if seqlen_q > seqlen_k else \
+                            (seqlen_q * seqlen_k - ((seqlen_q**1 - seqlen_q) / 2))
+                    flops_per_matmul += valid_out_elements * HQ * D_HEAD * 2
                 else:
                     flops_per_matmul += seqlen_q * seqlen_k * HQ * D_HEAD * 2
         else:
             q, k, v, input_metadata = input_helper(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, args.layout)
             if causal:
-                causal_correction = N_CTX_K if N_CTX_Q > N_CTX_K else N_CTX_Q
-                flops_per_matmul = 2.0 * BATCH * HQ * (N_CTX_Q * N_CTX_K - (causal_correction**2) / 2) * D_HEAD
+                # Same calculation as if varlen/if causal above
+                valid_out_elements = ((N_CTX_K**2 + N_CTX_K) / 2) if N_CTX_Q > N_CTX_K else \
+                        (N_CTX_Q * N_CTX_K - ((N_CTX_Q**2 - N_CTX_Q) / 2))
+                flops_per_matmul = 2.0 * BATCH * HQ * valid_out_elements * D_HEAD
             else:
                 flops_per_matmul = 2.0 * BATCH * HQ * N_CTX_Q * N_CTX_K * D_HEAD
         if causal: