Fix KV head shape mismatch when TP size exceeds num_kv_heads#3426

copybara-service[bot] merged 1 commit intomainAI-Hypercomputer/maxtext:mainfrom

mohit/ep_replicateAI-Hypercomputer/maxtext:mohit/ep_replicateCopy head branch name to clipboard