[Draft] Support Flash Attention#501

cmathw wants to merge 5 commits intoTransformerLensOrg:devTransformerLensOrg/TransformerLens:devfrom

cmathw:flash-attentioncmathw/TransformerLens:flash-attentionCopy head branch name to clipboard