@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 145 lines 6.96 KB
Rendered Raw
1 # Engram: conditional memory via n-gram lookups
2
3 Engram is a *conditional memory* module attached to specific layers of a
4 transformer backbone. At each position, the preceding few tokens are hashed
5 into large static embedding tables, and the retrieved rows are fused into the
6 hidden state, gated by how well they match it. Where Mixture-of-Experts scales
7 capacity through conditional computation, Engram scales it through conditional
8 memory: parameters that are never computed over, only looked up, so capacity
9 grows at O(1) per-token cost. The mechanism and its design rationale are
10 described in [Conditional Memory via Scalable Lookup: A New Axis of Sparsity
11 for Large Language Models](https://arxiv.org/abs/2601.07372).
12
13 vLLM enables Engram automatically whenever the checkpoint declares Engram
14 layers, so basic serving needs no extra flags. This page explains what Engram
15 does and how to control its memory footprint and multi-GPU topology.
16
17 ## Supported models
18
19 Engram layers are declared by the checkpoint, not chosen at serve time:
20
21 | Architecture | HF config field | Notes |
22 | --- | --- | --- |
23 | `DeepseekV41ForCausalLM` | `engram_layer_ids` | e.g. `deepseek-ai/DeepSeek-V4.1-Flash` |
24 | `Qwen4ExpForCausalLM` / `Qwen4ExpForConditionalGeneration` | `ple_layer_ids` | same mechanism, called PLE |
25
26 Passing `--engram-config` for a model without n-gram embedding layers fails at
27 startup, as does running on a non-CUDA-alike platform.
28
29 ## How it works
30
31 Engram processes each position in two phases: *retrieval* and *fusion*.
32
33 **Retrieval** maps the local token context to static table rows:
34
35 1. **Tokenizer compression.** At startup, vLLM builds a many-to-one map from
36 the tokenizer that collapses tokens which normalize alike (NFKC, accent
37 stripping, lowercasing, whitespace collapsing), so `" The"`, `"the"` and
38 `"THE"` hash identically. This concentrates semantic density into a much
39 smaller id space; its size must match `engram_compressed_vocab_size` from
40 the checkpoint, or startup fails — every hash multiplier derives from it.
41 2. **Multi-head hashing.** For each n-gram order `n` (up to
42 `engram_max_ngram_size - 1`), several hash heads each apply an independent
43 multiplicative-XOR hash over the compressed suffix n-gram. Every (n-gram
44 order, head) pair indexes its own prime-sized bucket range in the layer's
45 table — disjoint across pairs, so collisions cannot collide twice — and
46 the retrieved rows are concatenated into one memory vector per position.
47
48 **Fusion** turns the retrieved rows into a context-aware contribution to the
49 hidden state. The rows are context-independent priors and may be noisy (hash
50 collisions, polysemy), so they are gated: a `wkv` projection produces one key
51 per hyper-connection branch plus a shared value, and each branch computes a
52 sigmoid gate over the normalized dot product between its hidden state —
53 which has already aggregated global context through preceding attention —
54 and its key. If the retrieved memory contradicts the current context, the
55 gate tends toward zero and suppresses it. The gated value is added
56 residually into the branch's stream (`hidden + gate * value`); positions that
57 take no part in an n-gram (e.g. image spans) are masked so they pass through
58 untouched.
59
60 Engram sits only on the layers listed in the checkpoint. Placement is a
61 deliberate hardware-algorithm trade-off: early enough that the backbone is
62 relieved from reconstructing static patterns in its first layers, deep enough
63 that vLLM can hide table lookups behind the compute of the preceding layers
64 (see below).
65
66 In vLLM, hashing, lookup, and gating run as fused Triton kernels over FP8
67 rows with ue8m0 block scales. The layer layout — table sizes, n-gram size,
68 head count — is fixed by the checkpoint and is not user-configurable. Only the
69 storage and sharding described below are configurable, via
70 `--engram-config`, which takes a JSON object and can also be set field by
71 field as `--engram-config.<field>`
72 ([CLI reference](../cli/serve.md#-engram-config)).
73
74 ## CPU offload
75
76 Engram tables dwarf typical embedding weights: [DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash)
77 carries two Engram layers of roughly 384M table rows each, with 256-byte FP8
78 entries plus their ue8m0 block scales — about 200 GB of table weights alone.
79 Because lookups are deterministic — the indices are known from the token ids
80 alone, before the forward pass — the tables need not sit in GPU memory at
81 all. By default vLLM stores them in **pinned host memory** and serves lookups
82 directly from the Triton kernel through unified virtual addressing (UVA),
83 prefetching on a side CUDA stream so transfers overlap with the compute of
84 preceding layers. The freed GPU memory goes to the KV cache.
85
86 `cpu_offload` defaults to `true` (on)
87
88 To keep the tables resident on GPU instead (e.g. when host memory is
89 scarce):
90
91 ```bash
92 vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
93 --engram-config.cpu_offload false
94 ```
95
96 CPU offload requires a GPU with UVA support;
97 vLLM fails fast when it is unavailable.
98
99 ## Data-parallel topologies (DeepSeek V4.1)
100
101 By default (`embedding_across_dp: false`), each DP replica keeps its own
102 TP-sharded copy of the tables — no cross-replica communication, but one host
103 copy per replica.
104
105 With `embedding_across_dp: true`, the hash heads are sharded across all
106 TP × DP ranks into a single table copy. Each step gathers the hash ids of
107 every co-located DP replica, each rank looks up the heads it owns, and the
108 rows are exchanged back. This trades per-step DP collectives for a much
109 smaller table footprint:
110
111 ```bash
112 vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
113 --tensor-parallel-size 2 --data-parallel-size 4 \
114 --engram-config '{"embedding_across_dp": true}'
115 ```
116
117 `embedding_across_dp` is not supported with elastic expert parallelism yet.
118 Qwen4Exp PLE tables are ETP-sharded instead and honor `cpu_offload` only.
119
120 ## Sharing host tables across DP replicas
121
122 When the tables are CPU-offloaded and DP replicas are co-located on one node,
123 `dp_shared_memory` stores **one copy of each TP shard in `/dev/shm`**, shared
124 by all co-located replicas through `cudaHostRegister`. Each replica then
125 prefetches only its own tokens — no per-step Engram DP collectives — and host
126 memory drops from one table copy per replica to one per node.
127
128 `dp_shared_memory` defaults to enabled whenever the other settings allow it,
129 and falls back to per-replica (or DP-sharded) tables with a warning when:
130
131 - the DP replicas are not co-located on a single node,
132 - `/dev/shm` cannot hold the full tables — in containers, raise the limit
133 with `--shm-size` or `--ipc=host`,
134 - expert parallelism is elastic, which is unsupported.
135
136 Requirements: `cpu_offload` enabled, `--data-parallel-size > 1`, and a shared
137 IPC namespace across the co-located replicas.
138
139 ## Limitations
140
141 - CUDA-alike platforms only.
142 - DeepSeek V4.1 Engram does not support DBO or microbatching; disable
143 `--enable-dbo` and set `--ubatch-size 0`.
144 - `dp_shared_memory` requires `cpu_offload`, `data_parallel_size > 1`, and is
145 unsupported with elastic expert parallelism.