Rebase experimental_proxy on main

[avrittrohwer]

experimental_proxy only has 2 commits on it since branching:

$ git log --no-merges origin/experimental_proxy ^origin/main
commit 13f519e39e0d904e320c7d8a472161e4bcf03408 (HEAD -> avritt/noocdbt, origin/vivianrwu_experimental_proxy, origin/experimental_proxy, experimental_proxy)
Author: Zhihao Shan <[email protected]>
Date:   Mon Sep 9 15:15:08 2024 -0700

    correct model name for llama-2 70b

commit f8980b6a4d5cb01fa07258bda88b7497bdc01320
Author: Zhihao Shan <[email protected]>
Date:   Thu Sep 5 11:30:47 2024 -0700

    Add llama3 70b inference experimental support on v5e-16

    In addition, fetch the KV cache sharding from model forward in decode
    phase and use different sharding for prefill and decode.

We want https://github.com/AI-Hypercomputer/maxtext/commit/6c749c7e624ebb6d4a24716c2aa1dcc75f006b91 changes in experimental_proxy branch. Rebasing was easier than cherry-picking.

Steps:

git fetch origin
git branch -c avritt/noocdbt
git reset --hard origin/experimental_proxy
git rebase -X theirs --onto origin/main HEAD~2