Add two-section causal backward for shared K/V#2722

foraxe wants to merge 4 commits into

Dao-AILab:mainDao-AILab/flash-attention:mainfrom

foraxe:agent/two-section-causal-backwardforaxe/flash-attention:agent/two-section-causal-backwardCopy head branch name to clipboard