MOCAP: Wafer-Scale-Chip-Oriented Memory-Orchestrated Chunked Pipelining Framework for Prefill-Only LLM Inference

2026-06-22

AI Summary

This work addresses the memory and latency bottlenecks in long-context prefill inference caused by inter-chunk causal dependencies, which lead to imbalanced KV cache distribution and computational load. Targeting wafer-scale processors, the paper proposes a memory-coordinated chunked pipelining framework that introduces two key innovations: Memory-Balanced KV Reallocation (MBKR) for equitable redistribution of KV caches and Latency-Balanced Chunk Partitioning (LBCP) for optimizing chunking strategies to minimize latency. Evaluated against state-of-the-art baselines, the proposed approach reduces end-to-end latency by 76.4% and achieves a 3.24× throughput improvement over GPipe, while extending the maximum supported sequence length by 1.31× compared to Terapipe, thereby significantly advancing the efficiency and scalability of long-sequence inference.

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.

Problem

• prefill-only LLM inference
• wafer-scale chips
• chunked pipelining
• memory imbalance
• latency imbalance

Innovation

• wafer-scale chip
• chunked pipelining
• KV cache reallocation
• memory balancing
• prefill-only inference