feat(bench): Add features and fix some bugs for pipeline flashattention.

[KuangjuX]

No description provided.

[KuangjuX]

The pseudocode for FlashAttention-2 can be represented in the following form:

Iterate (k, v) in K, V :
  qk = dot(q, k) (1)
  mij = max(max(qk), lsei) (2)
  p = exp(qk − mij ) (3)
  lij = sum(p) (4)
  // renormalize o
  acc_o_scale = exp(mi − mij ) (5)
  acc_o = acc_o_scale ∗ acc_o (6)
  acc_o = acc_o + dot(p, v) (7)
  // update statistics
  mi = mij (8)
  li_new = exp (lsei − mij ) + lij (9)
  lsei = mij + log(li_new) (10)
// o_scale is the denominator of the softmax function
o_scale = exp(mi − lsei) (11)
acc_o = acc_o ∗ o_scale (12)

In the implementation of FractalTensor, some implementation details seem to differ, and I am not sure if this is compliant with the specifications. Below, I will point them out one by one.

Can you help me check these issues? @lcy-seso

[KuangjuX]

In FractalTensor, the update of LSE is performed outside the loop(https://github.com/microsoft/FractalTensor/blob/artifact/artifacts/FractalTensor/benchmarks/multi-head_attention/fractaltensor/kernel.h#L278), whereas in the pseudocode, the update of LSE is done inside the loop.

[KuangjuX]

I moved the update of LSE inside the loop.

[KuangjuX]

In the pseudocode, o_scale = exp(mi − lsei) is used to rescale the final result, while in FractalTensor, the operation is performed using 1 / lse_new(ax0)(https://github.com/microsoft/FractalTensor/blob/artifact/artifacts/FractalTensor/benchmarks/multi-head_attention/fractaltensor/kernel.h#L277).

[haruhi55]

I would like to give quick answers to your two questions and can provide detailed explanations for the reasons. First of all, I have carefully derived it, and the pseudocode is correct. I believe you can follow it. The normalization factor is outside the loop intentionally to reduce computational complexity.