Engram: n-gram 조회를 통한 조건부 메모리

Engram: n-gram 조회를 통한 조건부 메모리 (Engram: conditional memory via n-gram lookups)

Engram은 트랜스포머 백본의 특정 레이어에 붙는 조건부 메모리(conditional memory) 모듈입니다. 각 위치에서 앞선 몇 개의 토큰을 큰 정적 임베딩 테이블에 해시하고, 가져온 행을 은닉 상태와 융합하되 그 일치 정도로 게이팅합니다. MoE가 조건부 계산으로 용량을 키운다면, Engram은 조건부 메모리로 용량을 키웁니다. 즉 연산 대신 조회만 하므로 용량이 토큰당 O(1) 비용으로 늘어납니다. 메커니즘과 설계 근거는 Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models에 설명되어 있습니다.

출처: 문서

체크포인트가 Engram 레이어를 선언하면 vLLM이 자동으로 Engram을 활성화하므로, 기본 서빙에는 별도 플래그가 필요 없습니다. 이 페이지는 Engram이 무엇을 하는지, 그리고 메모리 사용량과 멀티 GPU 토폴로지를 제어하는 방법을 설명합니다.

본문

지원 모델 (Supported models)

Engram 레이어는 서빙 시점에 선택하는 것이 아니라 체크포인트가 선언합니다.

| Architecture | HF config field | Notes | | DeepseekV41ForCausalLM | engram_layer_ids | e.g. deepseek-ai/DeepSeek-V4.1-Flash | | Qwen4ExpForCausalLM / Qwen4ExpForConditionalGeneration | ple_layer_ids | same mechanism, called PLE |

n-그램 임베딩 레이어가 없는 모델에 --engram-config를 넘기면 시작 시 실패하며, CUDA 유사 플랫폼이 아닌 곳에서 실행해도 실패합니다.

동작 방식 (How it works)

Engram은 각 위치를 *조회(retrieval)*와 융합(fusion) 두 단계로 처리합니다.

**조회(Retrieval)**는 로컬 토큰 컨텍스트를 정적 테이블 행에 매핑합니다.

  • 토크나이저 압축. 시작 시 vLLM은 정규화가 비슷한 토큰(NFKC, 악센트 제거, 소문자화, 공백 축소)을 하나로 합치는 토크나이저의 다대일 맵을 만듭니다. 그래서 " The", "the", "THE"는 동일하게 해시됩니다. 이렇게 하면 의미 밀도가 훨씬 작은 ID 공간에 집중됩니다. 그 크기는 체크포인트의 engram_compressed_vocab_size와 일치해야 하며, 그렇지 않으면 시작이 실패합니다. 모든 해시 배수(hash multiplier)가 이 값에서 파생되기 때문입니다.
  • 멀티헤드 해싱. 각 n-그램 차수 n(최대 engram_max_ngram_size - 1)에 대해 여러 해시 헤드가 압축된 접미사 n-그램에 독립적인 곱셈-XOR 해시를 적용합니다. 각 (n-그램 차수, 헤드) 쌍은 레이어의 테이블에서 고유한 소수 크기 버킷 범위를 인덱스합니다. 쌍끼리 서로소이므로 충돌이 두 번 일어날 수 없고, 가져온 행들은 위치당 하나의 메모리 벡터로 연결됩니다.

**융합(Fusion)**은 가져온 행을 은닉 상태에 기여하는 컨텍스트 인지(context-aware) 성분으로 바꿉니다. 행들은 컨텍스트와 무관한 사전(prior)이고 노이즈가 있을 수 있어(해시 충돌, 다의어) 게이팅됩니다. wkv 프로젝션은 하이퍼커넥션 브랜치마다 키 하나와 공유 값 하나를 만들고, 각 브랜치는 앞선 어텐션으로 이미 전역 컨텍스트를 모은 자신의 은닉 상태와 키 사이의 정규화된 내적에 시그모이드 게이트를 계산합니다. 가져온 메모리가 현재 컨텍스트와 모순되면 게이트는 0으로 수렴해 그것을 억제합니다. 게이트된 값은 브랜치의 스트림에 잔차 방식으로 더해집니다(hidden + gate * value). n-그램에 참여하지 않는 위치(예: 이미지 구간)는 마스킹되어 그대로 통과합니다.

Engram은 체크포인트에 나열된 레이어에만 위치합니다. 배치는 의도적인 하드웨어-알고리즘 트레이드오프입니다. 너무 이르면 백본이 첫 레이어에서 정적 패턴 재구성 부담을 덜고, 충분히 깊으면 vLLM이 테이블 조회를 앞선 레이어들의 연산 뒤에 숨길 수 있습니다(아래 참조).

vLLM에서 해싱·조회·게이팅은 FP8 행과 ue8m0 블록 스케일 위에서 동작하는 융합된 Triton 커널로 실행됩니다. 레이어 레이아웃(테이블 크기, n-그램 크기, 헤드 수)은 체크포인트로 고정되며 사용자가 설정할 수 없습니다. 아래에 설명한 저장 방식과 샤딩만 --engram-config로 설정할 수 있는데, JSON 객체를 받으며 --engram-config.<field> 형식으로 필드별 지정도 가능합니다(CLI reference).

CPU 오프로드 (CPU offload)

Engram 테이블은 일반 임베딩 가중치보다 훨씬 큽니다. DeepSeek-V4.1-Flash는 각각 약 3.84억 개 행을 가진 Engram 레이어 2개를 갖고 있으며, 256바이트 FP8 엔트리와 ue8m0 블록 스케일을 합치면 테이블 가중치만 약 200GB입니다. 조회는 결정적이라 인덱스가 포워드 패스 전에 토큰 ID만으로 알려지므로, 테이블이 GPU 메모리에 있을 필요가 전혀 없습니다. 기본적으로 vLLM은 이를 pinned host memory에 저장하고, UVA(unified virtual addressing)를 통해 Triton 커널이 직접 조회를 서빙하며, 측면 CUDA 스트림에서 프리페치해 앞선 레이어들의 연산과 전송이 겹치도록 합니다. 해제된 GPU 메모리는 KV 캐시로 갑니다.

cpu_offload 기본값은 true(켜짐)입니다. VLLM_PLE_CPU_OFFLOAD 환경 변수를 따르며, 이 변수도 기본값이 true입니다. 명시적 --engram-config 값이 환경 변수보다 우선합니다.

대신 테이블을 GPU에 상주시키려면(예: 호스트 메모리가 부족할 때):

vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
  --engram-config.cpu_offload false

또는 동등하게 VLLM_PLE_CPU_OFFLOAD=0. CPU 오프로드는 UVA를 지원하는 GPU가 필요하며, 불가능하면 vLLM이 빠르게 실패합니다.

데이터 병렬 토폴로지 (Data-parallel topologies, DeepSeek V4.1)

기본적으로(embedding_across_dp: false) 각 DP 복제본은 TP로 샤딩된 자체 테이블 사본을 보유합니다. 복제본 간 통신은 없지만, 복제본당 호스트 사본이 하나씩 생깁니다.

embedding_across_dp: true로 설정하면 해시 헤드가 모든 TP × DP 랭크에 걸쳐 단일 테이블 사본으로 샤딩됩니다. 각 스텝에서 공동 배치된 모든 DP 복제본의 해시 ID를 모으고, 각 랭크가 자기가 소유한 헤드를 조회한 뒤 행을 다시 교환합니다. 이는 스텝당 DP 콜렉티브를 훨씬 작은 테이블 풋프린트로 맞바꾼 것입니다.

vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
  --tensor-parallel-size 2 --data-parallel-size 4 \
  --engram-config '{"embedding_across_dp": true}'

embedding_across_dp는 아직 elastic expert parallelism과 함께 지원되지 않습니다. Qwen4Exp PLE 테이블은 대신 ETP로 샤딩되며 cpu_offload만 따릅니다.

DP 복제본 간 호스트 테이블 공유 (Sharing host tables across DP replicas)

테이블이 CPU로 오프로드되고 DP 복제본이 한 노드에 공동 배치된 경우, dp_shared_memory는 각 TP 샤드의 사본 하나를 /dev/shm에 저장해 cudaHostRegister로 모든 공동 배치 복제본이 공유합니다. 그러면 각 복제본은 자신의 토큰만 프리페치하고, 스텝당 Engram DP 콜렉티브가 없으며, 호스트 메모리는 복제본당 테이블 사본 하나에서 노드당 하나로 줄어듭니다.

dp_shared_memory는 다른 설정이 허용할 때마다 기본으로 활성화되며, 다음 경우에는 경고와 함께 복제본별(또는 DP 샤딩) 테이블로 폴백합니다.

  • DP 복제본이 단일 노드에 공동 배치되지 않았거나,
  • /dev/shm이 전체 테이블을 담을 수 없는 경우(컨테이너에서는 --shm-size 또는 --ipc=host로 한도를 올리세요),
  • expert parallelism이 elastic인 경우(지원 안 됨).

요구 사항: cpu_offload 활성화, --data-parallel-size > 1, 공동 배치 복제본 간 공유 IPC 네임스페이스.

제한 사항 (Limitations)

  • CUDA 유사 플랫폼 전용.
  • DeepSeek V4.1 Engram은 DBO나 microbatching을 지원하지 않습니다. --enable-dbo를 끄고 --ubatch-size 0을 설정하세요.
  • dp_shared_memory는 cpu_offload와 data_parallel_size > 1이 필요하며, elastic expert parallelism에서는 지원되지 않습니다.

더 알아보기 (Learn more)