Pooling Engram Conditional Memory in Large Language Models using CXL

This paper proposes a scalable and cost-efficient solution for Large Language Models by integrating Compute Express Link (CXL) memory pools into SGLang to store Engram conditional memory, achieving near-DRAM end-to-end performance while overcoming the latency limitations of traditional RDMA approaches.

The Gist

Imagine you are running a massive, super-smart library (a Large Language Model, or LLM) that can write stories, answer questions, and solve problems. To be this smart, the library needs to remember a huge amount of facts and vocabulary.

The Problem: The "Too-Heavy" Backpack

Traditionally, this library tries to carry all its knowledge in its immediate backpack (the computer's fast RAM). But as the library grows smarter, the backpack gets so heavy and expensive that it becomes impossible to carry. You either have to buy a super-expensive, giant backpack for every single librarian, or the library slows down because it's struggling to carry the weight.

Recently, a new idea called Engram was introduced. Think of Engram as a "smart lookup system." Instead of memorizing every single fact inside the main brain, the model keeps a massive, static dictionary of word combinations (like "how to," "once upon a," "the cat sat") in a separate, huge storage room. When the model needs to write a sentence, it quickly glances at this dictionary to grab the right phrase.

The Catch: This "dictionary" is huge (hundreds of gigabytes), but the librarian only needs to grab tiny, scattered pages from it very quickly. If you try to move this dictionary to a cheaper, slower storage room (like a hard drive), the librarian gets stuck waiting for the pages, and the whole library slows down.

The Solution: The "CXL Super-Hallway"

This paper proposes a brilliant new way to store that dictionary using a technology called CXL (Compute Express Link).

Here is the analogy:

1. The Old Way (RDMA): Imagine the librarian has to run to a different building to get a book. They have to fill out a form, wait in line at the front desk, get the book, and run back. Even if the book is just a single page, the process takes too long. This is like using standard network cables (RDMA) to fetch data; it's great for moving big boxes, but terrible for grabbing tiny, scattered pages quickly.
2. The New Way (CXL): Now, imagine building a magic hallway (CXL) that connects the librarian's desk directly to the storage room. This hallway is so fast and direct that it feels like the book is right on the desk, even though it's actually in the other room.
   • Fine-Grained Access: The librarian can grab a single page without waiting in line.
   • Low Latency: The travel time is almost zero.
   • Shared Resource: Instead of every librarian having their own giant, expensive storage room, they all share one massive, centralized storage room connected by this magic hallway.

What the Researchers Did

The team built a prototype system to test this idea:

• The Setup: They took a modern AI system (SGLang) and connected it to a shared pool of memory using this "CXL hallway."
• The Test: They tried to run the AI with the massive dictionary stored in this shared pool instead of the expensive local memory.
• The Result: It worked almost perfectly! The speed of the AI was nearly identical to if the dictionary were stored right on the main computer. The "magic hallway" was fast enough to keep up with the librarian's need for instant information.

Why This Matters (The "Aha!" Moment)

• Cost Savings: In the old world, if you wanted to run a super-smart AI, you had to buy expensive memory for every single computer. With CXL, you can buy one big, shared memory bank and let dozens of computers share it. It's like sharing a single, massive library building among many small reading rooms instead of building a library in every house.
• Scalability: As AI models get bigger and need more "knowledge," this system can easily expand by just adding more shelves to the shared room, without needing to upgrade every single computer.

In a Nutshell

This paper shows that by using a new type of "super-fast hallway" (CXL), we can store the massive, heavy knowledge bases of future AI models in a cheap, shared location without slowing them down. It's the difference between carrying a heavy backpack everywhere versus having a teleportation device that instantly fetches exactly what you need, right when you need it. This makes building super-smart AI much cheaper and more efficient for the future.

Technical Summary

Here is a detailed technical summary of the paper "Pooling Engram Conditional Memory in Large Language Models using CXL".

1. Problem Statement

Large Language Models (LLMs) increasingly rely on Engram, a conditional memory architecture that decouples static knowledge lookup (N-Gram embeddings) from dynamic computation. While Engram offers superior scaling laws and efficiency compared to traditional Mixture-of-Experts (MoE) models, it introduces significant memory overhead.

• Memory Footprint: Future Engram tables could scale to hundreds of GBs, creating a "memory wall" for GPU-accelerated inference.
• Access Pattern Mismatch: Engram exhibits a unique access profile:
  • Sparse & Discrete: Retrieval involves fetching small, scattered data segments (e.g., 5 KB per token per layer) rather than large contiguous blocks.
  • Read-Only: The data is static during inference, allowing for offloading.
  • Latency-Sensitive: Retrieval must occur within a narrow prefetch window (approx. 56µs for early layers) to overlap with computation.
• Limitations of Existing Solutions:
  • Local DRAM: Too expensive to scale for massive Engram tables.
  • RDMA-based Pooling: While used for KV Cache offloading, RDMA suffers from high latency and throughput degradation for small, discrete packets (often <25% of peak bandwidth for 64-byte messages) due to protocol stack overhead.

2. Methodology

The authors propose a system that offloads Engram parameters to a Compute Express Link (CXL) memory pool, leveraging CXL's hardware-level load/store semantics to overcome RDMA limitations.

System Architecture

• Hardware Setup: A cluster of compute nodes connected via a CXL Switch (e.g., XConn XC50256) to a centralized pool of CXL memory devices (e.g., 256 GB Samsung cards).
• Memory Exposure: The CXL memory is exposed as a DAX (Direct Access) device in the OS. It is mapped to user-space virtual addresses via mmap, allowing applications to treat remote memory as local.
• Access Routines:
  • CXL → CPU: Utilizes multi-threaded parallel memcpy (OpenMP) to saturate PCIe bandwidth, treating CXL memory like local DRAM.
  • CXL → GPU: Implements a custom CUDA kernel for direct Peer-to-Peer (P2P) transfer. The CXL region is registered as CUDA host memory (cudaHostRegister), allowing the GPU DMA engine to fetch data directly over the PCIe bus without CPU intervention. A "wide-grid" kernel fuses thousands of discrete requests to minimize launch overhead.
• Integration: The system is integrated into SGLang, a state-of-the-art LLM inference framework. The Engram retrieval is triggered asynchronously at specific Transformer layers (e.g., layers 2 and 15) to overlap with preceding computation.

Optimization Strategy

• Prefetching: Since Engram retrieval depends on token IDs (not hidden states), the system initiates memory fetches immediately upon token generation, hiding latency behind the execution of other Transformer blocks.
• Parallelism: The system maximizes concurrency by distributing retrieval requests across multiple threads and GPU blocks to fully utilize the CXL switch's aggregate bandwidth (up to 512 GB/s).

3. Key Contributions

1. First CXL-Engram System: Proposes and implements the first system to offload Engram parameters to a CXL-based memory pool.
2. Comparative Analysis: Provides a rigorous analysis demonstrating that CXL is superior to RDMA for Engram workloads. While RDMA struggles with small-packet latency, CXL provides near-local DRAM latency due to its native load/store interface.
3. High-Performance Implementation: Develops specialized CXL access routines (CPU and GPU) that achieve near-DRAM end-to-end performance within the SGLang framework.
4. Cost-Benefit Validation: Demonstrates that CXL pooling offers a scalable, cost-effective solution for future memory-augmented LLMs without compromising inference throughput.

4. Experimental Results

The system was evaluated on a testbed with Intel Xeon CPUs, NVIDIA L20 GPUs, and a CXL switch/memory setup.

• Latency Performance:
  • CXL-to-CPU: Latency is comparable to local DRAM.
  • CXL-to-GPU: While slightly higher than DRAM, the latency remains within the strict prefetch window requirements (e.g., <56µs for early layers).
  • Comparison: RDMA pooling showed orders of magnitude higher latency for small batch sizes compared to CXL.
• End-to-End Throughput:
  • Tested with Qwen3-4B and Qwen3-8B models.
  • Result: CXL-based Engram pooling achieved throughput nearly identical to DRAM-based Engram.
    • Example (Qwen3-4B): Baseline (6183 tok/s) vs. Engram-DRAM (5683 tok/s) vs. Engram-CXL (5614 tok/s). The overhead of CXL was negligible (~1.2% drop compared to DRAM).
• Scalability:
  • Scaling to 2 nodes and higher Data Parallelism (DP=2) showed negligible performance degradation, proving the CXL pool can handle concurrent multi-host access effectively.
• Cost Analysis:
  • For small clusters (2 nodes), CXL has higher fixed costs due to the switch and adapters.
  • Economies of Scale: As the cluster size and model complexity grow (e.g., 16 nodes, 400B Engram), CXL pooling becomes significantly cheaper (saving >$160k in the 16-node scenario) by eliminating the need to replicate massive Engram tables on every node's local DRAM.

5. Significance

This paper addresses a critical bottleneck in the evolution of next-generation LLMs: the trade-off between massive knowledge storage and inference cost/latency.

• Infrastructure Shift: It validates CXL as the preferred interconnect for sparse, fine-grained memory access patterns in AI, distinguishing it from RDMA which is better suited for bulk data transfers (like KV Cache).
• Feasibility of Engram: By solving the memory and latency constraints, this work makes the deployment of massive Engram tables feasible, enabling LLMs to utilize "conditional memory" for better reasoning and knowledge retention without prohibitive hardware costs.
• Future-Proofing: The proposed architecture provides a blueprint for scalable, memory-augmented LLM infrastructures that can grow exponentially in capacity while maintaining linear cost efficiency.