Pooling Engram Conditional Memory in Large Language Models using CXL

이 논문은 Engram 의 희소 접근 패턴에 최적화된 CXL 메모리 풀을 SGLang 에 통합하여, 대규모 언어 모델의 추론 성능을 저해하지 않으면서도 확장 가능하고 비용 효율적인 메모리 솔루션을 제시합니다.

핵심

1. 문제 상황: "지식 있는 AI"의 고충

지금까지의 AI(거대 언어 모델) 는 두 가지 방식 중 하나를 선택해야 했습니다.

1. 머리 속에 모든 것을 외우기: AI 가 모든 지식을 자신의 뇌 (메모리) 에 다 넣으면, 뇌가 너무 커져서 비싸고 무거워집니다.
2. 계산으로 추측하기: 지식을 외우지 않고 매번 계산해서 답을 찾으려 하면, 시간이 너무 오래 걸려서 답답합니다.

최근 **'엔그램 (Engram)'**이라는 새로운 기술이 등장했습니다. 이는 AI 가 "외우기"와 "계산하기" 사이에서 중간 길을 찾은 것입니다. AI 가 특정 지식 (예: "사과"라는 단어가 나올 때 연관된 정보) 을 별도의 작은 창고에서 찾아와서 사용하는 방식입니다.

하지만 여기서 문제가 생깁니다.
이 '지식 창고'의 크기가 너무 큽니다 (수백 GB). 그런데 AI 가 이 창고를 이용할 때는 매우 드물고, 아주 작은 조각만 가져옵니다. 마치 도서관에서 책 전체를 빌리는 게 아니라, 특정 페이지의 한 줄만 복사해 오는 것과 비슷합니다.

기존 방식 (RDMA 라는 네트워크 기술) 으로 이 작은 조각들을 가져오려니, 도서관 문이 너무 느려서 AI 가 답을 기다리는 시간이 길어집니다.

2. 해결책: "CXL"이라는 초고속 엘리베이터

이 논문은 **CXL(컴퓨트 익스프레스 링크)**이라는 새로운 기술을 이용해 이 문제를 해결했습니다.

비유: "공유 아파트의 공용 냉장고"

• 기존 방식 (RDMA): 각 아파트 (서버) 가 따로따로 냉장고를 가지고 있고, 물건을 가져오려면 엘리베이터를 타고 다른 아파트로 이동해야 합니다. 엘리베이터가 느리고, 작은 물건 (한 줄의 지식) 을 가져오려면 시간이 너무 걸립니다.
• 새로운 방식 (CXL): 모든 아파트가 하나의 거대한 공용 냉장고를 공유합니다. 그리고 이 냉장고는 각 아파트의 부엌 (GPU) 과 **초고속 엘리베이터 (CXL)**로 직접 연결되어 있습니다.

CXL 의 장점:

• 직접 연결: 엘리베이터가 부엌 바로 옆에 있어서, 작은 물건 (지식 조각) 을 가져오는 데 걸리는 시간이 거의 없습니다.
• 공유: 한 개의 거대한 냉장고만 있으면 되므로, 아파트 (서버) 를 여러 개 늘려도 냉장고를 새로 살 필요가 없습니다.

3. 실험 결과: "느린 도서관"이 "초고속 도서관"이 되다

연구진은 이 기술을 실제 AI 시스템 (SGLang) 에 적용해 보았습니다.

• 속도: CXL 을 사용하면, AI 가 지식 창고에서 정보를 가져오는 속도가 로컬 메모리 (AI 뇌 바로 옆에 있는 메모리) 와 거의 똑같아졌습니다. 즉, 창고가 멀리 있어도 AI 는 전혀 느려지지 않습니다.
• 비용: 여러 대의 서버가 하나의 거대한 지식 창고를 공유하므로, 하드웨어 비용을 획기적으로 줄일 수 있습니다. 특히 AI 모델이 커질수록 (지식 창고가 커질수록) 비용 절감 효과가 더 큽니다.

4. 요약: 왜 이것이 중요한가요?

이 논문은 **"AI 가 더 똑똑해지려면 더 많은 지식이 필요하지만, 그 지식을 저장하는 비용과 속도가 문제였다"**는 점을 지적하고, CXL 이라는 기술을 통해 "공유된 초고속 지식 창고"를 만들어 해결했다는 것을 보여줍니다.

한 줄 요약:

"AI 가 필요한 지식을 가져오는 속도를 떨어뜨리지 않으면서, 거대한 지식 창고를 여러 AI 가 함께 쓸 수 있게 만들어 비용은 줄이고 성능은 유지하는 방법을 개발했습니다."

이 기술이 상용화되면, 앞으로 더 똑똑하고 저렴한 AI 서비스를 우리가 더 쉽게 이용할 수 있게 될 것입니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• Engram 의 등장과 메모리 부담: 대규모 언어 모델 (LLM) 은 기존 Mixture-of-Experts (MoE) 아키텍처의 비효율성을 해결하기 위해 'Engram'이라는 조건부 메모리 (Conditional Memory) 를 도입했습니다. Engram 은 정적 지식 (N-Gram 임베딩) 을 동적 계산에서 분리하여 $O(1)$ 검색을 가능하게 하지만, 미래의 대규모 LLM 에서는 수백 GB 에 달하는 방대한 임베딩 테이블을 필요로 합니다.
• 기존 메모리 풀링의 한계 (RDMA):
  • 기존 LLM 서비스 (KV Cache 등) 에서는 RDMA 를 이용한 원격 메모리 풀링이 주로 사용되었습니다.
  • 그러나 Engram 은 **희소성 (Sparse)**과 **미세한 단위 (Fine-grained)**의 접근 패턴을 가집니다. 즉, 매 토큰당 매우 작은 데이터 (약 5KB) 를 불연속적으로 자주 가져와야 합니다.
  • RDMA 는 패킷 오버헤드와 메시지 기반 네트워킹 semantics 로 인해 소규모 데이터 전송 시 성능이 급격히 저하됩니다 (예: 64 바이트 메시지 시 최대 대역폭의 25% 미만). 이는 Engram 의 낮은 지연 시간 요구사항을 충족하지 못합니다.
• 핵심 문제: Engram 의 메모리 접근 패턴 (희소, 미세, 낮은 지연 시간 요구) 을 효율적으로 지원하면서도, GPU 메모리 벽 (Memory Wall) 을 극복하고 비용을 절감할 수 있는 저장 솔루션이 필요합니다.

2. 제안된 방법론 (Methodology)

이 논문은 Compute Express Link (CXL) 기반의 메모리 풀링을 Engram 저장을 위해 최초로 제안하고 구현했습니다.

• CXL 기반 메모리 풀링 아키텍처:
  • 하드웨어: CXL 스위치 (XConn XC50256) 를 통해 여러 컴퓨팅 노드가 중앙 집중식 CXL 메모리 풀 (예: 256GB~4TB) 을 공유합니다.
  • 접근 방식: RDMA 와 달리 CXL 은 하드웨어 레벨의 로드/스토어 (Load/Store) 원시 연산을 지원하며, 캐시 라인 (Cache-line) 단위의 세밀한 접근이 가능합니다. 이는 Engram 의 불연속적이고 미세한 데이터 요청에 최적화되어 있습니다.
• SGLang 프레임워크 통합 및 최적화:
  • 초기화: Engram 파라미터를 공유 CXL 메모리에 로드하며, SGLang 의 ModelRunner 가 이를 관리합니다.
  • 비동기 프리페칭 (Prefetching): 디코딩 단계 시작 시 해시 함수를 통해 필요한 임베딩을 비동기적으로 CXL 풀에서 가져옵니다.
  • 고성능 접근 루틴 (Access Routines):
    • CXL → CPU: DAX (Direct Access) 모드를 활용하여 사용자 공간 가상 주소로 메모리를 매핑하고, OpenMP 기반의 멀티스레드 memcpy 를 사용하여 병렬 전송을 최적화합니다.
    • CXL → GPU: 커스텀 CUDA 커널을 사용하여 CPU 를 우회한 P2P (Peer-to-Peer) 데이터 전송을 구현합니다. 수천 개의 작은 요청을 하나의 와이드 그리드 (Wide-grid) 커널로 융합하여 PCIe 대역폭을 최대한 활용하고 전송 오버헤드를 최소화합니다.

3. 주요 기여 (Key Contributions)

1. 최초의 CXL 기반 Engram 시스템: Engram 파라미터를 CXL 기반 공유 메모리 풀로 오프로딩하는 첫 번째 시스템을 제안했습니다.
2. RDMA vs CXL 비교 분석: Engram 의 희소하고 미세한 메모리 접근 패턴에 대해 RDMA 가 가지는 지연 시간 한계를 분석하고, CXL 이 이를 해결하는 데 우월함을 입증했습니다.
3. 실제 프레임워크 구현 및 성능 검증: SGLang 추론 프레임워크에 CXL 기반 Engram 풀을 통합하여, DRAM 로딩과 유사한 엔드 - 투 - 엔드 성능을 달성함을 증명했습니다.

4. 실험 결과 (Results)

• 지연 시간 (Latency):
  • RDMA: 로컬 DRAM 대비 수 배~수십 배 높은 지연 시간을 보이며 Engram 의 프리페칭 창 (Prefetch window, 약 56μs) 을 충족하지 못했습니다.
  • CXL: 로컬 DRAM 과 유사한 지연 시간 (수십 μs 수준) 을 달성하여 Engram 의 엄격한 지연 시간 요구사항을 충족했습니다. 특히 CXL→GPU 직접 전송 오버헤드도 허용 가능한 범위 내에 있었습니다.
• 엔드 - 투 - 엔드 처리량 (Throughput):
  • Qwen3-4B/8B 모델을 SGLang 에서 테스트한 결과, CXL 기반 Engram 풀을 사용했을 때 DRAM 기반 Engram 과 비교해 처리량 감소가 미미했습니다 (예: Qwen3-4B 기준 DRAM 5683.7 tokens/s vs CXL 5614.4 tokens/s).
  • 이는 CXL 풀링이 추론 성능을 저하시키지 않으면서 확장 가능한 솔루션임을 의미합니다.
• 확장성 (Scalability):
  • 데이터 병렬화 (DP) 와 노드 수를 늘려도 성능 저하가 거의 없었으며, CXL 풀의 병렬 접근 능력이 우수함을 확인했습니다.
• 비용 분석 (Cost Analysis):
  • 소규모 구성 (예: 2 노드, 100B 파라미터) 에서는 CXL 인프라 비용이 높게 나오지만, 노드 수와 모델 규모가 커질수록 (예: 16 노드, 400B 파라미터) DRAM 을 각 노드에 모두 탑재하는 것보다 압도적인 비용 절감 효과 (약 86% 절감) 를 보였습니다.

5. 의의 및 결론 (Significance)

• 차세대 LLM 인프라의 패러다임 전환: Engram 과 같은 메모리 증강형 LLM 을 경제적으로 배포할 수 있는 기반을 마련했습니다.
• 성능과 비용의 동시 달성: CXL 의 저지연, 세밀한 접근 특성을 활용하여, 대규모 정적 지식 테이블을 저비용의 공유 메모리로 관리하면서도 DRAM 수준의 추론 속도를 유지할 수 있음을 증명했습니다.
• 미래 지향성: 이 연구는 LLM 서비스의 메모리 병목 현상을 해결하고, 더 크고 복잡한 모델을 비용 효율적으로 운영할 수 있는 확장 가능한 아키텍처를 제시합니다.

요약하자면, 이 논문은 Engram 의 고유한 메모리 접근 패턴에 RDMA 가 부적합함을 지적하고, CXL 의 기술적 우위를 활용하여 저비용·고성능의 Engram 메모리 풀링 시스템을 구축하고 검증한 선구적인 연구입니다.