Enable Pipeline Parallelism on torchax path#1055

Chenyaaang merged 2 commits intomainvllm-project/tpu-inference:mainfrom

chenyangli/pp-5vllm-project/tpu-inference:chenyangli/pp-5Copy head branch name to clipboard