Implement QK attention head chunking in MLA to reduce HBM footprint#4564

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

zjiahao/DSA3.2-qk-chunkAI-Hypercomputer/maxtext:zjiahao/DSA3.2-qk-chunkCopy head branch name to clipboard