Memory and communication approach improves large AI model inference
MCF: A Memory- and Communication-Aware Co-Designed Framework for Long-Context LLM Inference
Hardware Architecture
Summary
Large language models (LLMs) need to remember and process long conversations, but this creates delays and heavy use of memory and communication between computers. The authors present MCF, a system that smartly manages memory and communication by changing how data is stored and shared during different stages of processing. MCF also uses specialized hardware to speed up data routing, reducing wait times and increasing efficiency as conversations get longer. Their tests show MCF performs better than existing methods, especially with longer context sizes and bigger computer clusters.
What this means in practice
- •For ai infrastructure engineers: Design distributed systems that handle long-context large language model inference efficiently by managing memory replication and communication dynamically.
- •For network hardware developers: Develop hardware-accelerated network devices to support metadata-driven data routing and collective operations directly in communication paths for AI inference workloads.
Authors
Huaizhi Zhang, Jiacheng Zhu, Xuqi Zhu, Junkyu Lee, Klaus McDonald-Maier, Xiaojun Zhai
Abstract
This paper presents Memory- and Communication-Aware Framework for Long-Context LLM Inference (MCF) that optimizes long-context inference through communication-aware and stage-adaptive KV-cache layout transitions. At the core of MCF is {Memory-Aware Replication and Scheduling} (MARS), a dynamic orchestration mechanism for chunked prefill. MARS exploits early-stage memory headroom by selectively replicating highly reused KV blocks, improving locality and hiding KV propagation latency behind attention computation. As prefill progresses and memory pressure increases, MARS progressively drains redundant replicas and converges the KV cache into a sharded decode-ready layout, avoiding expensive layout conversion at the prefill--decode boundary. For decoding, MCF adopts a query-centric execution model in which KV shards remain stationary while lightweight query tensors are routed across workers. This changes communication from KV-cache-dependent transfer to query-dependent transfer, reducing communication growth with context length. MCF further integrates a hardware-software co-designed execution layer based on an FPGA compute-NIC datapath, enabling metadata-driven KV routing and collective operations directly in the network path. Experiments show that MCF improves TTFT, throughput, and scalability over state-of-the-art distributed inference systems, with larger benefits as context length and cluster size increase.