MOCAP: Wafer-Scale-Chip-Oriented Memory-Orchestrated Chunked Pipelining Framework for Prefill-Only LLM Inference
2026-06-22 • Hardware Architecture
Hardware Architecture
AI summaryⓘ
The authors focus on speeding up large language models when processing very long sequences, which is slow due to communication and memory challenges in traditional GPU setups. They propose MOCAP, a method that splits input sequences into chunks and manages memory more evenly across devices to handle longer sequences efficiently. Their approach also balances computing work between chunks to avoid delays. Tests show MOCAP is much faster and supports longer sequences than previous methods. Overall, the authors improve how large models run on advanced chip systems by smarter memory and workload management.
Large Language ModelsPrefill PhaseWafer-Scale ChipsKV CachePipeline ParallelismMemory ImbalanceAttention MechanismLatencyGPipeTerapipe
Authors
Zichuan Wang, Huizheng Wang, Yuheng Xiao, Haonan Zuo, Taiquan Wei, Jinyi Deng, Chao Li, Yang Hu, Shouyi Yin
Abstract
Large language models (LLMs) are increasingly used in prefill-only workloads, where end-to-end latency is dominated by the prefill phase. For long-context prefill, communication overhead grows with sequence length and quickly becomes a bottleneck on conventional GPU systems, making wafer-scale chips (WSCs) a promising substrate due to their high communication bandwidth and large aggregate compute and memory capacity. A natural way to accelerate prefill is to partition a long input sequence into multiple chunks and execute them in a finer-grained pipeline across devices. However, directly applying this idea to long-context prefill on WSCs remains challenging. First, causal dependency across chunks causes KV cache to accumulate unevenly across pipeline stages, creating severe memory imbalance and limiting the feasible sequence length. Second, later chunks require more attention computation because each chunk depends on preceding chunks, leading to chunk-level latency imbalance. To address these challenges, we present MOCAP, a memory-orchestrated chunked pipelining framework for prefill-only LLM inference on WSCs. MOCAP introduces Memory-Balanced KV Reallocation (MBKR) to alleviate memory imbalance by redistributing KV cache across pipeline stages, thereby extending the feasible sequence length. It further incorporates Latency-Balanced Chunk Partitioning (LBCP) to balance chunk execution cost under both attention-cost growth and KV reallocation overhead, improving pipeline efficiency. Experimental results show that, compared with GPipe, MOCAP achieves 76.4\% lower end-to-end latency and 3.24$\times$ higher throughput on average. MOCAP also extends the maximum supported sequence length by up to 1.31$\times$ compared with Terapipe.