[WIP] Enable SPMD#32

Draft

wonjoo-wj wants to merge 5 commits intollama2-google-next-inferencefrom

wonjoo/dynamo-spmd

Commits on Jul 28, 2023

Rebase and add example_2d_sharding.py
wonjoo-wj
committed

Commits on Jul 31, 2023

Update prompt generation and hard-code sharding dims
wonjoo-wj
committed

Commits on Aug 1, 2023

Remove example_2d_sharding.py and enable SPMD
wonjoo-wj
committed

Commits on Aug 2, 2023

Manually shard the KV cache
wonjoo-wj
committed

Commits on Aug 7, 2023

Update to latest
wonjoo-wj
committed