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

Abstract: 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.

• The paper introduces MOCAP, which harnesses memory-balanced KV reallocation and latency-balanced chunk partitioning to cut prefill latency on wafer-scale chips.
• It achieves up to 76.4% lower end-to-end latency and 3.24× higher throughput compared to GPipe while extending the maximum sequence length.
• MOCAP effectively mitigates stage-wise memory and computation imbalances, offering scalable LLM inference with improved service density and reliability.

MOCAP: Memory-Orchestrated Chunked Pipeline Framework for Prefill-Only LLM Inference on Wafer-Scale Chips

Problem Statement and Motivation

Prefill-only LLM inference–where a model processes long input contexts and emits a single output token–has emerged as a common workload, particularly for applications requiring low service latency per request. The prefill phase in such scenarios dominates overall latency. While conventional GPU systems are increasingly challenged by the substantial communication overhead arising from long-context prefill, recent hardware progress has produced wafer-scale chips (WSCs) capable of higher compute densities, larger memory pools, and intra-wafer communication bandwidths several orders above typical GPU clusters.

Figure 1: The prefill-only LLM inference process and empirical latency comparison between GPU systems and wafer-scale chips with equivalent compute/memory capacity.

On GPUs, the communication volume for prefill grows with input sequence length, rapidly becoming the system bottleneck. WSCs, leveraging dense in-package interconnects and expansive SRAM/HBM, notably reduce this bottleneck by 46.8% in empirical latency comparisons. However, existing system designs have not addressed prefill-only workloads specifically or exploited WSC properties for this target.

Chunked Pipeline Processing: System and Architecture

Pipeline parallelism has been widely adopted for LLM training and inference, typically via microbatch slicing (e.g., GPipe) or token/sequence chunking (e.g., Terapipe). Terapipe’s chunked pipeline, partitioning sequences for forward passes, appears promising for prefill acceleration, but introduces distinctive stage-wise memory imbalances due to causal dependencies—each chunk must attend to KV cache from prior chunks, which accumulate unevenly across pipeline stages.

Figure 2: Comparison of microbatch-based pipeline (GPipe) and token-based chunked pipeline (Terapipe) highlighting sequence partitioning and pipeline granularity.

The WSC hardware further accentuates possible improvements due to its unique die array and memory hierarchy.

Figure 3: Wafer-scale chip architecture with dense neighbor mesh, compute die arrays, and high-bandwidth memory pools.

Memory and Latency Bottlenecks in Chunked Prefill

Memory bottlenecks emerge because earliest pipeline stages hold live KV cache for the longest durations, saturating local memory and truncating feasible sequence length. Later pipeline stages, in contrast, underutilize their memory pools. This “debtor-creditor” imbalance restricts preprocessing of long contexts, regardless of total system memory, and leads to wasted capacity.

Latency bottlenecks originate from chunk-wise execution cost heterogeneity: later chunks incur increased attention computation (as prefix length grows), further introducing severe pipeline bubbles, especially upon remote KV access or reallocation among stages.

Figure 4: Stage-wise memory and execution latency imbalance in chunked pipeline, showing disproportionate KV accumulation and pipeline bubbles.

MOCAP Framework Design

MOCAP systematically unifies memory and pipeline orchestration for prefill-only LLM inference on WSCs, introducing two principal mechanisms:

Memory-Balanced KV Reallocation (MBKR)

MOCAP’s MBKR algorithm leverages dynamic “debtor-creditor” stage pairing to reallocate excess KV cache from overloaded to underutilized stages, extending feasible sequence length without additional hardware provisioning. The pairing is optimized to minimize link contention and communication overhead, preferring locality (adjacent stage assignment) and fixed cross-half pairing.

Figure 5: Overview of MBKR policy mapping optimal stage pairing and runtime reallocation states.

Runtime policies adaptively distinguish whether stages are acting as KV cache debtors or creditors, ensuring that memory is effectively pooled across the whole wafer and communication paths are minimized.

Latency-Balanced Chunk Partitioning (LBCP)

To address compute utilization loss from chunk-level latency heterogeneity, MOCAP employs a two-stage LBCP algorithm. First, dynamic programming yields a compute-balanced partition as an initialization. This partition is then refined via simulated annealing, re-optimizing chunk boundaries to jointly minimize attention and remote KV reallocation overhead. The resulting chunk size distribution counteracts prefix-induced attention latency and communication delays, reducing pipeline bubbles and maximizing throughput.

Experimental Results

Comprehensive evaluation was conducted on a $4\times4$ WSC configuration using Llama3-70B, Mistral-123B, Qwen3-235B, and Llama3-405B. Compared against GPipe and Terapipe baselines, MOCAP achieves:

• 76.4% lower end-to-end latency and 3.24$\times$ higher throughput on average compared to GPipe.
• Up to 1.31$\times$ extension in maximum sequence length relative to Terapipe.
• Robust performance across models, including those with MoE (Qwen3-235B) and high GQA ratios (Llama3-405B), demonstrating architecture-agnostic benefits.

  Figure 6: (a) Throughput and latency scaling of GPipe, Terapipe, and MOCAP on WSCs; failing prefill jobs shown in red. (b) Maximum sequence length support comparison for Terapipe and MOCAP.

MOCAP maintains scalability in the long-context regime, effectively utilizes trans-die memory, and mitigates compute and communication bottlenecks even as sequence lengths grow. The tradeoff between pipeline efficiency and sequence headroom is exposed: fewer chunks permit longer sequences but degrade pipeline fill efficiency.

Implications and Future Directions

MOCAP's architecture-aware orchestration for prefill-only inference exemplifies system-accelerator co-design for wafer-scale environments. It unlocks both higher throughput and extended context headroom, crucial for LLM deployments where input sequences are increasingly long and service latency is stringent.

Practical implications include improved service density and reliability at data center scale, more efficient memory use, and reduced failure rates in high-throughput serving. Theoretical implications involve a new class of scheduling and orchestration algorithms, potentially generalizable to other Transformer workloads with innate memory/computation imbalances.

Future AI acceleration may benefit from further exploiting wafer-scale locality, dynamic partitioning, and hybrid granularity scheduling, as well as algorithm-hardware co-design for multi-die attention and memory pooling. Research into hardware-aware chunk sizing and real-time orchestration for autoregressive and prefill-only workloads will offer significant gains in both performance and resource utilization.

Conclusion

MOCAP provides a robust, memory-orchestrated chunked pipeline for prefill-only LLM inference on wafer-scale chips. By combining MBKR-based KV cache reallocation and LBCP-enabled adaptive chunk partitioning, MOCAP addresses both memory and latency imbalances inherent in chunked pipeline execution. Empirical results demonstrate substantial reductions in end-to-end latency, marked throughput increases, and extension of feasible sequence length, furnishing a scalable paradigm for long-context LLM inference in wafer-scale computing environments (2606.22968).