Clamp kv_stage to avoid SMEM overflow for small head_dims on SM100#2594

Johnsonms merged 2 commits into

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

Johnsonms:johnson/fix-hd-small-smem-overflowJohnsonms/flash-attention:johnson/fix-hd-small-smem-overflowCopy head branch name to clipboard