@tduncan / Nvidia-Optimization-Stack-VLLM / commits / fba33e39cb

[Docs] Add an Engram feature page explaining Engram usage in vLLM (#57910)

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

Nicolò Lucchesi committed Sep 21, 2026 at 19:15 UTC fba33e39cb74c7c31cc648ae530a85e0c6636f92
1 file changed +148
docs/features/engram.md new
+148
@@ -0,0 +1,148 @@
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): it follows the
87 +`VLLM_PLE_CPU_OFFLOAD` environment variable, which also defaults to `true`.
88 +An explicit `--engram-config` value takes precedence over the environment
89 +variable.
90 +
91 +To keep the tables resident on GPU instead (e.g. when host memory is
92 +scarce):
93 +
94 +```bash
95 +vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
96 + --engram-config.cpu_offload false
97 +```
98 +
99 +or equivalently `VLLM_PLE_CPU_OFFLOAD=0`. CPU offload requires a GPU with UVA
100 +support; vLLM fails fast when it is unavailable.
101 +
102 +## Data-parallel topologies (DeepSeek V4.1)
103 +
104 +By default (`embedding_across_dp: false`), each DP replica keeps its own
105 +TP-sharded copy of the tables — no cross-replica communication, but one host
106 +copy per replica.
107 +
108 +With `embedding_across_dp: true`, the hash heads are sharded across all
109 +TP × DP ranks into a single table copy. Each step gathers the hash ids of
110 +every co-located DP replica, each rank looks up the heads it owns, and the
111 +rows are exchanged back. This trades per-step DP collectives for a much
112 +smaller table footprint:
113 +
114 +```bash
115 +vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
116 + --tensor-parallel-size 2 --data-parallel-size 4 \
117 + --engram-config '{"embedding_across_dp": true}'
118 +```
119 +
120 +`embedding_across_dp` is not supported with elastic expert parallelism yet.
121 +Qwen4Exp PLE tables are ETP-sharded instead and honor `cpu_offload` only.
122 +
123 +## Sharing host tables across DP replicas
124 +
125 +When the tables are CPU-offloaded and DP replicas are co-located on one node,
126 +`dp_shared_memory` stores **one copy of each TP shard in `/dev/shm`**, shared
127 +by all co-located replicas through `cudaHostRegister`. Each replica then
128 +prefetches only its own tokens — no per-step Engram DP collectives — and host
129 +memory drops from one table copy per replica to one per node.
130 +
131 +`dp_shared_memory` defaults to enabled whenever the other settings allow it,
132 +and falls back to per-replica (or DP-sharded) tables with a warning when:
133 +
134 +- the DP replicas are not co-located on a single node,
135 +- `/dev/shm` cannot hold the full tables — in containers, raise the limit
136 + with `--shm-size` or `--ipc=host`,
137 +- expert parallelism is elastic, which is unsupported.
138 +
139 +Requirements: `cpu_offload` enabled, `--data-parallel-size > 1`, and a shared
140 +IPC namespace across the co-located replicas.
141 +
142 +## Limitations
143 +
144 +- CUDA-alike platforms only.
145 +- DeepSeek V4.1 Engram does not support DBO or microbatching; disable
146 + `--enable-dbo` and set `--ubatch-size 0`.
147 +- `dp_shared_memory` requires `cpu_offload`, `data_parallel_size > 1`, and is
148 + unsupported with elastic expert parallelism.