[Raiden Weight Sync 3/7] MoE 128-lane layout and padding for TPU GMM kernels - #5168

[Raiden Weight Sync 3/7] MoE 128-lane layout and padding for TPU GMM kernels#5168

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

yixuann-m2-moe-fused-layoutAI-Hypercomputer/maxtext:yixuann-m2-moe-fused-layoutCopy head branch name to clipboard