Pooling Engram Conditional Memory in Large Language Models using CXL

本文提出利用 Compute Express Link (CXL) 内存池存储大语言模型的 Engram 条件记忆，通过 SGLang 集成实现了接近 DRAM 的端到端性能，为未来 Engram 集成的 LLM 提供了可扩展且高性价比的存储解决方案。

要点

这篇论文讲述了一个关于如何让大型人工智能（LLM）变得更聪明、更便宜、更快速的创新方案。为了让你轻松理解，我们可以把大模型想象成一个超级图书馆，把其中的技术细节用生活中的例子来类比。

1. 核心问题：图书馆太挤了，但书很少被翻动

想象一下，未来的大型语言模型（LLM）就像一座巨大的图书馆。

• 传统的做法：为了回答你的问题，图书馆必须把所有的书籍（知识）都放在手边的桌子上（显存/DRAM），以便随时取用。但这导致桌子越来越小，放不下那么多书，而且搬运这些书非常昂贵。
• Engram（记忆痕迹）的出现：研究人员发现，其实图书馆里只有极少部分的书是真正需要被频繁翻阅的（比如特定的词汇组合）。大部分时间，模型只是在处理逻辑，不需要翻书。
• 新挑战：Engram 技术就是把这些“特定词汇表”单独拿出来，放在一个巨大的外部仓库里。虽然仓库很大（几百 GB），但每次只需要从仓库里拿几页纸（几 KB 的数据）出来用一下。

痛点：如果仓库离得太远，或者拿书的方式太笨重（比如每次都要填单子、等快递），那模型就会因为“等书”而变慢，甚至卡死。

2. 解决方案：CXL 内存池 —— 就像“共享的高速传送带”

为了解决“等书”的问题，这篇论文提出使用一种叫 CXL 的新科技。

• 旧方案（RDMA）：就像是用快递车去仓库取书。

  • 虽然快递车能跑很远，但它有固定的流程：装车、发车、卸货。如果你只需要拿一张纸，让一辆大卡车跑一趟，效率极低，而且路上还要花时间。
  • 在论文中，RDMA 就像这种“大卡车”，对于 Engram 这种“每次只拿几页纸”的需求来说，太慢、太笨重了。

• 新方案（CXL 内存池）：就像是在图书馆和仓库之间架起了一条超高速的透明传送带。

  • 直接存取：CXL 允许计算机直接像拿自己桌子上的书一样，去拿仓库里的书。不需要填单子，不需要等快递。
  • 细粒度：它不仅能拿整箱书，还能精准地只拿一页纸，而且速度几乎和拿自己桌子上的书一样快。
  • 共享：多个图书馆（服务器）可以共用同一个大仓库，不用每个图书馆都自己建一个巨大的仓库，省下了巨额的建设费。

3. 他们做了什么？

研究团队（来自北京大学、阿里云等）做了一个实验：

1. 搭建系统：他们把 Engram 的“词汇表”放到了 CXL 共享内存池里。
2. 优化取书流程：他们写了一套专门的“取书程序”（代码），让电脑能同时派出一百个人，在传送带上并行地、飞快地把需要的几页纸取回来。
3. 接入大脑：把这个系统接入了目前最先进的 AI 推理框架（SGLang）。

4. 结果如何？

• 速度惊人：虽然书是从“外部仓库”（CXL 内存）拿的，但因为传送带（CXL）太快了，AI 感觉不到任何延迟。它的速度几乎和书放在自己桌子上（本地显存）一样快！
• 省钱巨大：
  • 以前：如果你有 10 个服务器，每个都要买 200GB 的昂贵内存来存这些书，总共要花很多钱。
  • 现在：大家共用一个 200GB 的大仓库，只需要买一个 CXL 交换机和适配器。
  • 比喻：就像 10 个人合租一个大冰箱，而不是每个人都要买一个独立的大冰箱。当规模越大（服务器越多、模型越大），省下的钱就越多。

5. 总结：这对我们意味着什么？

这篇论文的核心思想就是：用“共享的高速传送带”（CXL）来替代“昂贵的独立书桌”（本地显存），专门用来存放 AI 需要频繁查阅但又不占地方的小知识（Engram）。

• 对 AI 来说：它变得更聪明了（可以容纳更多知识），而且反应依然很快。
• 对公司来说：建 AI 服务器的成本大幅降低，让未来的大模型更普及。
• 对技术界：证明了 CXL 这种新技术非常适合处理这种“零碎但频繁”的数据读取任务，是未来 AI 基础设施的重要方向。

简单来说，他们发明了一种**“既快又省”的 AI 记忆外挂**，让大模型不再因为“记性太好（内存不够）”而发愁。

技术摘要

这篇论文提出了一种基于 CXL (Compute Express Link) 内存池化技术来存储和管理大语言模型（LLM）中 Engram 条件记忆（Conditional Memory） 参数的系统方案。该研究旨在解决 Engram 带来的巨大内存开销问题，同时保持推理性能，为下一代内存增强型 LLM 提供可扩展且低成本的存储基础设施。

以下是该论文的详细技术总结：

1. 研究背景与问题 (Problem)

• Engram 的引入与挑战： Engram 是一种由 DeepSeek 提出的架构，旨在通过引入“条件记忆”来补充 MoE（混合专家）架构中的“条件计算”。它将静态的 N-Gram 知识从动态计算中解耦，通过 $O(1)$ 的查找机制实现高效的知识检索。然而，随着模型规模扩大，Engram 的嵌入表（Embedding Table）可能达到数百 GB，给 GPU 显存带来巨大压力。
• 现有方案的局限性：
  • 内存墙： 将如此巨大的参数全部加载到 GPU 显存（DRAM）中成本高昂且不可扩展。
  • RDMA 的不足： 现有的内存池化方案（如基于 Mooncake 的 RDMA 方案）通常针对大块数据传输优化。但 Engram 的访问模式具有**稀疏性（Sparse）和细粒度（Fine-grained）**特征（每次推理仅检索少量离散的小块 Embedding，约几 KB）。RDMA 协议栈在处理这种小包传输时延迟极高，无法满足 Engram 严格的预取延迟要求。
• 核心问题： 如何构建一个既能提供低成本、大容量存储，又能满足 Engram 稀疏、低延迟访问需求的内存池化系统？

2. 方法论 (Methodology)

论文提出利用 CXL (Compute Express Link) 技术构建 Engram 内存池，主要包含以下技术细节：

2.1 系统架构设计

• CXL 内存池： 利用 CXL Switch 连接多个计算节点和 CXL 内存设备，构建一个解耦的、主机独立的共享内存池。
• 硬件拓扑： 每个服务器节点通过 PCIe 5.0 x16 适配器连接到 CXL Switch，Switch 连接至集中的 CXL 内存卡（如 256GB 或更大）。
• 数据流： Engram 参数存储在 CXL 内存池中，计算节点（CPU/GPU）通过 CXL 协议直接进行 Load/Store 操作，无需经过复杂的网络协议栈。

2.2 访问模式分析与优化

• 访问特征： Engram 的访问是只读的，且具有稀疏性（每个 Token 每层仅检索约 5KB 数据，分散在巨大的表中）和低延迟容忍度（检索必须在 Transformer 层计算完成前完成，通常窗口仅为几十微秒）。
• CXL 优势： CXL 提供硬件级的 Load/Store 语义，支持 Cache-line 粒度的访问，延迟接近本地 DRAM，完美契合 Engram 的稀疏访问模式。
• 专用访问例程 (Access Routines)：
  • CXL → CPU： 利用 DAX (Direct Access) 模式将 CXL 内存映射到用户空间，使用 OpenMP 多线程并行 memcpy 进行高效读取。
  • CXL → GPU： 注册 CXL 内存为 CUDA Host Memory，利用自定义 CUDA Kernel 实现 P2P (Peer-to-Peer) 直接传输。通过融合数千个离散请求为一个宽网格（Wide-grid）Kernel，最大化 PCIe 带宽利用率，避免大量 cudaMemcpy 的启动开销。

2.3 框架集成

• 将 CXL 内存池集成到先进的推理框架 SGLang 中。
• 初始化： 仅由一个 Rank 将 Engram 参数加载到共享 CXL 池中。
• 预取 (Prefetching)： 在 Forward 阶段异步触发，直接从 CXL 池拉取 Embedding 到 GPU 显存。
• 计算： 各 Rank 从池中获取所需 Embedding，与隐藏状态融合。

3. 关键贡献 (Key Contributions)

1. 首个 CXL 系统： 提出了第一个将 Engram 参数卸载到基于 CXL 的内存池中的系统。
2. RDMA 与 CXL 的对比分析： 深入分析了 Engram 的访问模式，证明了 CXL 在细粒度、稀疏访问场景下显著优于 RDMA（RDMA 在小包传输时延迟过高）。
3. 高性能实现与集成： 在 SGLang 中实现了 CXL 优化的 Engram 池，通过自定义 CPU/GPU 访问例程，实现了接近本地 DRAM 的端到端性能。
4. 成本效益验证： 证明了该方案在大规模部署下能显著降低硬件成本，同时不牺牲推理吞吐量。

4. 实验结果 (Results)

实验基于双节点服务器（配备 Intel Xeon 和 NVIDIA L20 GPU）和 CXL Switch 构建原型系统。

• 延迟性能：
  • CXL vs. 本地 DRAM： CXL 到 CPU 的读取延迟与本地 DRAM 相当。
  • CXL vs. RDMA： 在 Engram-27B 配置下，RDMA 池的延迟比 CXL 高出几个数量级，无法满足 Engram 的预取窗口要求（约 56μs）。CXL 延迟完全在可接受范围内。
  • CXL → GPU： 虽然略高于 CPU 路径，但通过 P2P 优化，延迟仍在系统容忍范围内，且随参数规模扩展保持稳定。
• 端到端吞吐量 (Throughput)：
  • 在 SGLang 中集成 Engram 模块（使用 Qwen3-4B/8B 模拟）。
  • 对比： 使用 CXL 池的 Engram 方案（5614.4 tokens/s for 4B）与使用本地 DRAM 的方案（5683.7 tokens/s）相比，性能损失极小（<2%），远优于 RDMA 方案（未直接列出但推断性能会大幅下降）。
  • 扩展性： 增加数据并行度（DP）和节点数（nnode）时，吞吐量提升显著，且 CXL 池的访问并未成为瓶颈，显示出良好的可扩展性。
• 成本分析：
  • 对于小规模配置（如 2 节点），CXL 方案因固定基础设施成本（Switch/Adapter）略高于纯 DRAM。
  • 随着规模扩大（如 16 节点、400B Engram 表），CXL 共享池方案展现出巨大的成本优势。例如，在 16 节点 400B 场景下，CXL 方案比全 DRAM 方案节省约 16.6 万美元。

5. 意义与展望 (Significance)

• 打破内存墙： 该方案为未来超大规模 LLM 中引入 Engram 等内存密集型组件提供了可行的存储解决方案，解决了显存容量不足的问题。
• 性能与成本的平衡： 证明了 CXL 技术能够以接近本地内存的性能提供远程存储，使得“内存增强型 LLM"在经济上变得可行。
• 架构创新： 为 LLM 推理基础设施从“计算为中心”向“计算与存储解耦”的演进提供了新的思路，特别是针对稀疏、细粒度访问负载的优化。
• 未来方向： 论文讨论了 RDMA 通过定制优化（如 P2P 接口、热点缓存）在 Engram 场景下的潜力，以及 Engram 与 KV Cache 在共享内存池中的共存管理挑战。

总结： 这篇论文通过引入 CXL 内存池化技术，成功解决了 Engram 架构在大规模部署中的存储瓶颈问题。它不仅在理论上分析了 CXL 优于 RDMA 的原因，更通过系统实现和实验验证，证明了该方案在保持高推理性能的同时，能显著降低大规模 LLM 基础设施的部署成本。