Optimize MLA generate_mask with scatter update#4261

JHCuc3m wants to merge 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

zjiahao/DSA3.2-scatter-mask-genAI-Hypercomputer/maxtext:zjiahao/DSA3.2-scatter-mask-genCopy head branch name to clipboard