Reproduce “Shift Parallelism: Low-Latency, High-Throughput LLM Inference for Dynamic Workloads” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Mert Hidayetoglu, Aurick Qiao, Michael Wyatt, Jeff Rasley, Yuxiong He, Samyam Rajbhandari
Shift Parallelism adapts GPU parallelism between sequence parallelism for larger batches and tensor parallelism for small batches while keeping the KV cache in a compatible layout. In a Llama-70B burst test on 8×H200 GPUs, it reached median TTFT of 148 ms and peak throughput of 69,147 tokens per second. Its peak throughput exceeded TP’s but remained below DP’s.
Key insights
- Sequence Parallelism (SP) maintains KV cache invariance with Tensor Parallelism (TP), allowing for dynamic switching between these modes without costly memory re-organization.
- Adapting Sequence Parallelism for inference necessitates specific engineering solutions for Grouped Query Attention, KV cache scaling beyond head count, and handling small batch sizes efficiently.
- Using pre-compiled CUDA graphs for both SP and TP configurations, along with distinct but KV-cache-compatible model weight loadings, enables rapid and efficient transitions between parallelism modes.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.