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

[DS V4.1][Engram] Support async prefetch for offloaded engram lookups and engram DP sharding (#56512)

Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: ywang96 <> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Codex <noreply@openai.com>

Juntian Liu committed Sep 13, 2026 at 02:44 UTC d2d649e674c75425d2d6975c87eb89fd4d55fff8
12 files changed +1820 -194
.buildkite/test_areas/distributed.yaml
+4
@@ -206,6 +206,9 @@ steps:
206 - tests/v1/distributed
207 - tests/v1/engine/test_engine_core_client.py
208 - tests/distributed/test_utils
209 + - tests/distributed/test_engram_dp_shard.py
210 + - vllm/config/engram.py
211 + - vllm/models/deepseek_v4_1/
212 commands:
213 # https://github.com/NVIDIA/nccl/issues/1838
214 - export NCCL_CUMEM_HOST_ENABLE=0
@@ -217,6 +220,7 @@ steps:
220 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
221 - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
222 - pytest -v -s distributed/test_utils.py
223 + - pytest -v -s distributed/test_engram_dp_shard.py
224
225 - label: ":nvidia: (L4) Distributed Compile + Comm"
226 key: distributed-compile-comm-4-gpus
tests/distributed/test_engram_dp_shard.py new
+776
@@ -0,0 +1,776 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Engram tables shared by co-located DP replicas.
4 +
5 +Four GPUs on one node split the hash heads four ways, whether as DEP
6 +(TP1 x DP4, the single-node MoE layout) or TP2 x DP2. A lookup only returns
7 +the right rows if the ids of every replica are gathered, each rank looks up
8 +its own heads, and each replica trades those tokens back for the heads its
9 +peers own.
10 +
11 +With shared host storage, DP replicas instead map the same TP slice and
12 +prefetch only their own tokens without DP gathers.
13 +"""
14 +
15 +from itertools import product
16 +from pathlib import Path
17 +from tempfile import TemporaryDirectory
18 +from threading import Event
19 +from types import SimpleNamespace
20 +
21 +import pytest
22 +import torch
23 +import torch.multiprocessing as mp
24 +
25 +from vllm.config import (
26 + EngramConfig,
27 + VllmConfig,
28 + get_current_vllm_config_or_none,
29 + set_current_vllm_config,
30 +)
31 +from vllm.config.load import LoadConfig
32 +from vllm.config.parallel import ParallelConfig
33 +from vllm.config.scheduler import SchedulerConfig
34 +from vllm.distributed import cleanup_dist_env_and_memory, parallel_state
35 +from vllm.distributed.parallel_state import (
36 + get_engram_dp_group,
37 + get_engram_dp_size,
38 + init_distributed_environment,
39 + initialize_model_parallel,
40 +)
41 +from vllm.forward_context import set_forward_context
42 +from vllm.models.deepseek_v4_1.common.engram import EngramLayout
43 +from vllm.models.deepseek_v4_1.nvidia import engram as engram_ops
44 +from vllm.models.deepseek_v4_1.nvidia.engram import (
45 + Engram,
46 + ParallelEngramEmbedding,
47 + engram_head_shard_rank,
48 + gather_engram_hashes,
49 +)
50 +from vllm.platforms import current_platform
51 +from vllm.utils.network_utils import get_open_port
52 +from vllm.utils.system_utils import update_environment_variables
53 +
54 +WORLD_SIZE = 4
55 +DIM, BLOCK = 64, 32
56 +# Prime-sized buckets, as the model's layout builds them.
57 +HEAD_SIZES = (97, 101, 103, 107, 109, 113, 127, 131)
58 +
59 +
60 +def _token_counts(dp_size: int) -> tuple[tuple[int, ...], ...]:
61 + """Even loads, then uneven ones with an idle replica to pad around."""
62 + return ((8,) * dp_size, (7, 0, 5, 1)[:dp_size], (7,) * dp_size)
63 +
64 +
65 +def _full_table(rows: int) -> tuple[torch.Tensor, torch.Tensor]:
66 + """The unsharded table, identical on every rank."""
67 + torch.manual_seed(0)
68 + weight = (torch.randn(rows, DIM) * 4).to(torch.float8_e4m3fn)
69 + scale_inv = torch.randint(120, 134, (rows, DIM // BLOCK), dtype=torch.uint8)
70 + return weight, scale_inv
71 +
72 +
73 +def _make_ids(head_sizes: tuple[int, ...], num_tokens: int, seed: int) -> torch.Tensor:
74 + """Ids of each head, drawn from that head's own bucket range."""
75 + torch.manual_seed(seed)
76 + offset = 0
77 + columns = []
78 + for size in head_sizes:
79 + columns.append(torch.randint(offset, offset + size, (num_tokens, 1)))
80 + offset += size
81 + return torch.cat(columns, dim=1).to(torch.int32).cuda()
82 +
83 +
84 +def _reference(weight, scale_inv, ids) -> torch.Tensor:
85 + """Dequantized lookup against the whole table."""
86 + valid = ids >= 0
87 + safe_ids = ids.clamp_min(0).long()
88 + values = torch.nn.functional.embedding(safe_ids, weight).float()
89 + scales = torch.nn.functional.embedding(safe_ids, scale_inv)
90 + scales = (scales.to(torch.int32) << 23).view(torch.float32)
91 + values = values.unflatten(-1, (-1, BLOCK)) * scales.unsqueeze(-1)
92 + return values.flatten(-2).to(torch.bfloat16).masked_fill(~valid.unsqueeze(-1), 0)
93 +
94 +
95 +@pytest.mark.parametrize(
96 + "architecture,options,edp_ranks,etp_ranks",
97 + [
98 + ("DeepseekV41ForCausalLM", {}, [0, 2], [0, 1]),
99 + ("Qwen4ExpForCausalLM", {}, None, [0, 1]),
100 + (
101 + "Qwen4ExpForCausalLM",
102 + {"embedding_across_dp": True},
103 + None,
104 + [0, 1, 2, 3],
105 + ),
106 + ("Qwen4ExpForConditionalGeneration", {}, None, [0, 1]),
107 + ],
108 +)
109 +def test_engram_group_creation_honors_model_and_options(
110 + monkeypatch, architecture, options, edp_ranks, etp_ranks
111 +):
112 + config = SimpleNamespace(
113 + model_config=SimpleNamespace(architecture=architecture, is_moe=False),
114 + parallel_config=ParallelConfig(
115 + tensor_parallel_size=2,
116 + data_parallel_size=2,
117 + distributed_executor_backend="mp",
118 + ),
119 + engram_config=EngramConfig(**options),
120 + )
121 + for name in (
122 + "_TP",
123 + "_ETP",
124 + "_ENGRAM_DP",
125 + "_DCP",
126 + "_PCP",
127 + "_PP",
128 + "_DP",
129 + "_EP",
130 + "_EPLB",
131 + ):
132 + monkeypatch.setattr(parallel_state, name, None)
133 + monkeypatch.setattr(parallel_state, "_WORLD", SimpleNamespace(local_rank=0))
134 + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
135 + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 4)
136 + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0)
137 + monkeypatch.setattr(parallel_state, "get_node_count", lambda: 1)
138 +
139 + def make_group(group_ranks, *args, **kwargs):
140 + ranks = next(ranks for ranks in group_ranks if 0 in ranks)
141 + return SimpleNamespace(ranks=ranks, rank_in_group=0, world_size=len(ranks))
142 +
143 + monkeypatch.setattr(parallel_state, "init_model_parallel_group", make_group)
144 + with set_current_vllm_config(config):
145 + initialize_model_parallel(tensor_model_parallel_size=2, backend="nccl")
146 + group = get_engram_dp_group()
147 + assert (group.ranks if group is not None else None) == edp_ranks
148 + assert parallel_state.get_etp_group().ranks == etp_ranks
149 +
150 +
151 +def _worker(rank: int, tp_size: int, port: int) -> None:
152 + torch.accelerator.set_device_index(rank)
153 + update_environment_variables(
154 + {
155 + "RANK": str(rank),
156 + "LOCAL_RANK": str(rank),
157 + "WORLD_SIZE": str(WORLD_SIZE),
158 + "MASTER_ADDR": "localhost",
159 + "MASTER_PORT": str(port),
160 + }
161 + )
162 + dp_size = WORLD_SIZE // tp_size
163 + dp_rank, tp_rank = divmod(rank, tp_size)
164 + vllm_config = VllmConfig(
165 + parallel_config=ParallelConfig(
166 + tensor_parallel_size=tp_size,
167 + data_parallel_size=dp_size,
168 + data_parallel_rank=dp_rank,
169 + ),
170 + scheduler_config=SchedulerConfig(
171 + max_model_len=8,
172 + is_encoder_decoder=False,
173 + max_num_batched_tokens=8,
174 + max_num_seqs=8,
175 + ),
176 + )
177 + # Supply the resolved Engram settings without constructing a full model.
178 + vllm_config.engram_config = EngramConfig()
179 + vllm_config.model_config = SimpleNamespace(
180 + architecture="DeepseekV41ForCausalLM", is_moe=True
181 + )
182 + try:
183 + init_distributed_environment()
184 + with set_current_vllm_config(vllm_config):
185 + initialize_model_parallel(tensor_model_parallel_size=tp_size)
186 + assert get_engram_dp_size() == dp_size
187 + assert get_engram_dp_group().ranks == [
188 + tp_rank + replica * tp_size for replica in range(dp_size)
189 + ]
190 + assert engram_head_shard_rank() == tp_rank * dp_size + dp_rank
191 +
192 + for failure in ("create", "open", "mmap", "register", "pinned", None):
193 + _check_shared_storage_lifetime_and_failures(failure)
194 +
195 + # Reuse the processes and communication groups across storage modes.
196 + for (cpu_offload, dp_shared_memory), n_heads, sequence_parallel in product(
197 + [(False, False), (True, False), (True, True)],
198 + [8, 7],
199 + [False, True] if tp_size > 1 else [False],
200 + ):
201 + _check_table(
202 + vllm_config,
203 + cpu_offload,
204 + dp_shared_memory,
205 + n_heads,
206 + sequence_parallel,
207 + )
208 + finally:
209 + cleanup_dist_env_and_memory()
210 +
211 +
212 +def _check_shared_storage_lifetime_and_failures(failure):
213 + """One peer's failure must roll back every mapping/registration, without hangs."""
214 + import gc
215 + import weakref
216 +
217 + group = get_engram_dp_group()
218 + runtime = torch.cuda.cudart()
219 + make_file = engram_ops.tempfile.NamedTemporaryFile
220 + map_file = engram_ops.mmap.mmap
221 + paths, mappings, registrations, unregistrations = [], [], [], []
222 + failed_rank = 0 if failure == "create" else 1
223 +
224 + def fail():
225 + raise OSError(f"injected {failure} failure")
226 +
227 + def temporary_file(*args, **kwargs):
228 + if failure == "create":
229 + fail()
230 + file = make_file(*args, **kwargs)
231 + paths.append(Path(file.name))
232 + return file
233 +
234 + def open_file(*args, **kwargs):
235 + if failure == "open" and group.rank_in_group == failed_rank:
236 + fail()
237 + return open(*args, **kwargs)
238 +
239 + def mapping(*args, **kwargs):
240 + if failure == "mmap" and group.rank_in_group == failed_rank:
241 + fail()
242 + result = map_file(*args, **kwargs)
243 + mappings.append(weakref.ref(result))
244 + return result
245 +
246 + def register(pointer, size, flags):
247 + if failure == "register" and group.rank_in_group == failed_rank:
248 + return SimpleNamespace(value=1)
249 + result = runtime.cudaHostRegister(pointer, size, flags)
250 + assert result.value == 0
251 + registrations.append(pointer)
252 + return result
253 +
254 + def unregister(pointer):
255 + unregistrations.append(pointer)
256 + result = runtime.cudaHostUnregister(pointer)
257 + assert result.value == 0
258 + return result
259 +
260 + with pytest.MonkeyPatch.context() as patch:
261 + patch.setattr(
262 + engram_ops,
263 + "tempfile",
264 + SimpleNamespace(NamedTemporaryFile=temporary_file),
265 + )
266 + patch.setattr(engram_ops, "open", open_file, raising=False)
267 + patch.setattr(
268 + engram_ops,
269 + "mmap",
270 + SimpleNamespace(mmap=mapping, MAP_SHARED=engram_ops.mmap.MAP_SHARED),
271 + )
272 + patch.setattr(
273 + torch.cuda,
274 + "cudart",
275 + lambda: SimpleNamespace(
276 + cudaHostRegister=register,
277 + cudaHostUnregister=unregister,
278 + ),
279 + )
280 + if failure == "pinned" and group.rank_in_group == failed_rank:
281 + patch.setattr(torch.Tensor, "is_pinned", lambda self: False)
282 + if failure is not None:
283 + stage = "cudaHostRegister" if failure in ("register", "pinned") else failure
284 + with pytest.raises(RuntimeError, match=f"EDP rank {failed_rank}.*{stage}"):
285 + engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group)
286 + assert unregistrations == registrations
287 + else:
288 + storage = engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group)
289 + assert all(not path.exists() for path in paths)
290 + storage_ref = weakref.ref(storage)
291 + parameter = torch.nn.Parameter(storage.weight, requires_grad=False)
292 + views = storage.get_views(parameter, storage.weight_scale_inv)
293 + del storage
294 + gc.collect()
295 + assert storage_ref() is None
296 + assert not unregistrations
297 + del parameter
298 + gc.collect()
299 + assert not unregistrations # UVA views still retain the CPU allocation.
300 + del views
301 + gc.collect()
302 + assert unregistrations == registrations
303 + for ref in mappings:
304 + mapped = ref()
305 + assert mapped is None or mapped.closed
306 + assert all(not path.exists() for path in paths)
307 +
308 +
309 +def _check_table(
310 + vllm_config, cpu_offload, dp_shared_memory, n_heads, sequence_parallel
311 +):
312 + parallel = vllm_config.parallel_config
313 + dp_size, tp_size = parallel.data_parallel_size, parallel.tensor_parallel_size
314 + dp_rank = parallel.data_parallel_rank
315 + tp_rank = engram_ops.get_tensor_model_parallel_rank()
316 + multithread = (
317 + dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1)
318 + )
319 + load_config = (
320 + LoadConfig(
321 + load_format="safetensors",
322 + use_tqdm_on_load=False,
323 + model_loader_extra_config={
324 + "enable_multithread_load": True,
325 + "num_threads": 2,
326 + },
327 + )
328 + if multithread
329 + else vllm_config.load_config
330 + )
331 + with pytest.MonkeyPatch.context() as m:
332 + head_sizes = HEAD_SIZES[:n_heads]
333 + rows = sum(head_sizes)
334 + config = SimpleNamespace(
335 + hidden_size=DIM,
336 + hc_mult=2,
337 + rms_norm_eps=1e-6,
338 + engram_layer_ids=[0],
339 + engram_num_embeddings=[rows],
340 + engram_max_ngram_size=2,
341 + engram_n_heads=n_heads,
342 + engram_head_dim=DIM,
343 + engram_compressed_vocab_size=32,
344 + engram_pad_token_id=0,
345 + engram_vocab_size=HEAD_SIZES[0],
346 + )
347 + with set_current_vllm_config(vllm_config):
348 + weight, scale_inv = _full_table(rows)
349 + layout = EngramLayout(config)
350 + # Stub configuration inputs, not Engram's initialization or execution.
351 + component_config = SimpleNamespace(
352 + engram_config=EngramConfig(
353 + cpu_offload=cpu_offload,
354 + dp_shared_memory=dp_shared_memory,
355 + ),
356 + scheduler_config=vllm_config.scheduler_config,
357 + load_config=load_config,
358 + )
359 + with m.context() as init_patch, torch.device("cuda"):
360 + init_patch.setattr(
361 + engram_ops, "get_current_vllm_config", lambda: component_config
362 + )
363 + engram = Engram(
364 + config,
365 + quant_config=None,
366 + layout=layout,
367 + layer_hash_index=0,
368 + use_sequence_parallel=sequence_parallel,
369 + prefix="model.layers.0.engram",
370 + )
371 + layer = engram.embed_tokens
372 + if multithread:
373 + _load_shared_multithread(layer, weight, scale_inv, load_config)
374 + else:
375 + # Only the leader has valid checkpoint payload in shared mode.
376 + loader_weight = weight if not dp_shared_memory or dp_rank == 0 else None
377 + loader_scale = (
378 + scale_inv if not dp_shared_memory or dp_rank == 0 else None
379 + )
380 + layer.weight.weight_loader(layer.weight, loader_weight)
381 + layer.weight_scale_inv.weight_loader(
382 + layer.weight_scale_inv, loader_scale
383 + )
384 + if dp_shared_memory:
385 + assert layer.dp_size == 1
386 + assert layer.head_start == tp_rank * layer.part_n_hash_cols
387 + assert layer.weight.is_pinned() and layer.weight_scale_inv.is_pinned()
388 +
389 + def unexpected_collective(*args, **kwargs):
390 + raise AssertionError("shared lookup must not use DP collectives")
391 +
392 + m.setattr(get_engram_dp_group(), "all_gather", unexpected_collective)
393 +
394 + token_counts = _token_counts(dp_size)
395 +
396 + # Production forward runs after the construction-time config has exited.
397 + assert get_current_vllm_config_or_none() is None
398 + for batch_idx, counts in enumerate(token_counts):
399 + num_tokens = counts[dp_rank]
400 + # Ids differ per replica but must agree across its TP ranks.
401 + ids = _make_ids(head_sizes, num_tokens, seed=100 + dp_rank)
402 + if batch_idx == 2 and dp_rank == 1:
403 + ids.fill_(engram_ops.DEAD_ID)
404 + with set_forward_context(
405 + None,
406 + vllm_config,
407 + num_tokens=num_tokens,
408 + num_tokens_across_dp=torch.tensor(counts, dtype=torch.int32),
409 + ):
410 + gathered = gather_engram_hashes(
411 + ids, dp_shared_memory=layer.dp_shared_memory
412 + )
413 + expected_tokens = (
414 + num_tokens if dp_shared_memory else max(counts) * dp_size
415 + )
416 + assert gathered.shape[0] == expected_tokens
417 + engram.prepare_embeddings(gathered)
418 + looked_up = engram.embed(ids)
419 + standalone = layer(ids)
420 + expected = _reference(weight.cuda(), scale_inv.cuda(), ids)
421 + torch.testing.assert_close(standalone, expected, atol=0, rtol=0)
422 + if sequence_parallel:
423 + # SP hands back only this rank's token shard, zero-padded.
424 + chunk = -(-num_tokens // tp_size)
425 + expected = torch.nn.functional.pad(
426 + expected, (0, 0, 0, 0, 0, chunk * tp_size - num_tokens)
427 + )
428 + expected = expected[tp_rank * chunk : (tp_rank + 1) * chunk]
429 + torch.testing.assert_close(looked_up, expected, atol=0, rtol=0)
430 +
431 + if (
432 + cpu_offload
433 + and n_heads == 7
434 + and not sequence_parallel
435 + and (not dp_shared_memory or tp_size == 1)
436 + ):
437 + _check_dummy_hash_model_forward(
438 + vllm_config,
439 + engram,
440 + head_sizes,
441 + weight,
442 + scale_inv,
443 + dp_shared_memory,
444 + )
445 +
446 + replay = n_heads == 7 and (
447 + (dp_shared_memory and sequence_parallel == (tp_size > 1))
448 + or (
449 + cpu_offload
450 + and not dp_shared_memory
451 + and sequence_parallel == (tp_size > 1)
452 + )
453 + )
454 + if replay:
455 + captures = ("full", "breakable")
456 + for capture in captures:
457 + _check_prefetch_replay(
458 + engram,
459 + head_sizes,
460 + weight,
461 + scale_inv,
462 + dp_shared_memory,
463 + capture,
464 + )
465 + if dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1):
466 + with m.context() as storage_patch:
467 + storage_patch.setattr(layer.weight, "data", layer.weight.data.clone())
468 + with pytest.raises(RuntimeError, match="storage must not be replaced"):
469 + layer(ids)
470 +
471 +
472 +def _check_dummy_hash_model_forward(
473 + vllm_config, engram, head_sizes, weight, scales, dp_shared_memory
474 +):
475 + """A metadata-less replica joins DP sharding but skips a shared lookup."""
476 + from vllm.models.deepseek_v4_1.common.engram import NgramHashState
477 + from vllm.models.deepseek_v4_1.nvidia import model as model_ops
478 +
479 + dp_rank = vllm_config.parallel_config.data_parallel_rank
480 + if dp_shared_memory and dp_rank != 1:
481 + return
482 + tokens = 4
483 + ids = _make_ids(head_sizes, tokens, seed=300 + dp_rank)
484 + state = NgramHashState.__new__(NgramHashState)
485 + torch.nn.Module.__init__(state)
486 + state.multipliers = torch.empty(1, 2, dtype=torch.int64, device="cuda")
487 + state.primes = torch.empty(1, 1, len(head_sizes), dtype=torch.int64, device="cuda")
488 + state.ensure_cache = lambda: True
489 + # Hash arithmetic is covered separately; keep the real model's branch and
490 + # NgramHashState.dummy_hashes, plus real preparation/lookup collectives.
491 + state.forward = lambda *args: ids.unsqueeze(1)
492 +
493 + class Decoder(SimpleNamespace):
494 + def __call__(self, hidden, positions, input_ids, *args):
495 + hashes, keep = args[-2:]
496 + if dp_shared_memory and dp_rank == 1:
497 + assert hashes is None and keep is None
498 + output = hidden.new_zeros(
499 + tokens, len(head_sizes), DIM, dtype=torch.bfloat16
500 + )
501 + else:
502 + assert hashes is not None and keep is not None
503 + assert bool(keep.all()) == (dp_rank != 1)
504 + if dp_rank == 1:
505 + assert torch.all(hashes == engram_ops.DEAD_ID)
506 + output = engram.embed(hashes[:, 0])
507 + return output, None, None, None, None
508 +
509 + model = SimpleNamespace(
510 + use_mega_moe=False,
511 + use_sequence_parallel=False,
512 + engram_hash=state,
513 + engram_swa_prefix="swa",
514 + engram_dp_shared_memory=dp_shared_memory,
515 + layers=[Decoder(engram=engram)],
516 + start_layer=0,
517 + end_layer=1,
518 + aux_hidden_state_layers=(),
519 + )
520 + metadata = (
521 + None
522 + if dp_rank == 1
523 + else {
524 + "swa": SimpleNamespace(
525 + query_start_loc=None, slot_mapping=None, block_table=None
526 + )
527 + }
528 + )
529 + with pytest.MonkeyPatch.context() as patch:
530 + patch.setattr(
531 + model_ops,
532 + "get_pp_group",
533 + lambda: SimpleNamespace(
534 + is_first_rank=True,
535 + is_last_rank=False,
536 + ),
537 + )
538 + patch.setattr(model_ops, "mhc_post_tilelang", lambda hidden, *args: hidden)
539 + with set_forward_context(
540 + metadata,
541 + vllm_config,
542 + num_tokens=tokens,
543 + num_tokens_across_dp=torch.full(
544 + (vllm_config.parallel_config.data_parallel_size,),
545 + tokens,
546 + dtype=torch.int32,
547 + ),
548 + ):
549 + output = model_ops.DeepseekV4Model.forward(
550 + model,
551 + torch.arange(tokens, device="cuda"),
552 + torch.arange(tokens, device="cuda"),
553 + intermediate_tensors=None,
554 + inputs_embeds=torch.zeros(tokens, DIM, device="cuda"),
555 + lookback_token_ids=torch.full((1, 1), -1, device="cuda"),
556 + )["hidden_states"]
557 + expected_ids = torch.full_like(ids, engram_ops.DEAD_ID) if dp_rank == 1 else ids
558 + expected = _reference(weight.cuda(), scales.cuda(), expected_ids)
559 + torch.testing.assert_close(output, expected, atol=0, rtol=0)
560 +
561 +
562 +def _load_shared_multithread(layer, weight, scale_inv, load_config):
563 + """Opposite shard completion orders must leave every shared parameter ready."""
564 + from safetensors.torch import save_file
565 +
566 + from vllm.model_executor.model_loader import weight_utils
567 + from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
568 +
569 + tensors = {
570 + "weight": weight,
571 + "weight_scale_inv": scale_inv.view(torch.float8_e8m0fnu),
572 + }
573 + order = list(tensors)
574 + if get_engram_dp_group().rank_in_group % 2:
575 + order.reverse()
576 + first_loaded = Event()
577 + original_load = weight_utils.load_file
578 +
579 + def load_file(path, **kwargs):
580 + if Path(path).stem != order[0]:
581 + assert first_loaded.wait(timeout=30), "First shared weight load stalled"
582 + return original_load(path, **kwargs)
583 +
584 + with TemporaryDirectory() as directory, pytest.MonkeyPatch.context() as m:
585 + for name, tensor in tensors.items():
586 + save_file({name: tensor}, str(Path(directory) / f"{name}.safetensors"))
587 + m.setattr(weight_utils, "load_file", load_file)
588 + loader = DefaultModelLoader(load_config)
589 + source = DefaultModelLoader.Source(directory, revision=None)
590 + loaded = []
591 + for name, tensor in loader._get_weights_iterator(source):
592 + param = getattr(layer, name)
593 + param.weight_loader(param, tensor)
594 + loaded.append(name)
595 + first_loaded.set()
596 + assert loaded == order
597 +
598 + # Check before any further collective could mask incomplete shared writes.
599 + for name, tensor in tensors.items():
600 + param = getattr(layer, name)
601 + expected = tensor.narrow(0, param.engram_vocab_start, param.shape[0])
602 + torch.testing.assert_close(
603 + param.view(torch.uint8), expected.view(torch.uint8), atol=0, rtol=0
604 + )
605 +
606 +
607 +def _check_prefetch_replay(
608 + engram, head_sizes, weight, scale_inv, dp_shared_memory, capture
609 +):
610 + """DP-sharded and shared prefetches must replay with new IDs."""
611 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
612 +
613 + ids = _make_ids(head_sizes, 8, seed=501)
614 + tp_rank = engram_ops.get_tensor_model_parallel_rank()
615 + tokens = 8 // engram.embed_tokens.tp_size if engram.use_sequence_parallel else 8
616 + output = torch.empty(
617 + tokens, len(head_sizes), DIM, device="cuda", dtype=torch.bfloat16
618 + )
619 +
620 + def step():
621 + gathered = gather_engram_hashes(ids.clone(), dp_shared_memory=dp_shared_memory)
622 + engram.prepare_embeddings(gathered)
623 + output.copy_(engram.embed(ids))
624 +
625 + with pytest.MonkeyPatch.context() as patch:
626 + patch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 8)
627 + warmup = torch.cuda.Stream()
628 + warmup.wait_stream(torch.cuda.current_stream())
629 + with torch.cuda.stream(warmup):
630 + step()
631 + if capture == "breakable":
632 + graph = BreakableCUDAGraphCapture()
633 + with torch.cuda.stream(warmup), graph:
634 + step()
635 + assert graph.num_eager_breaks == 2
636 + else:
637 + graph = torch.cuda.CUDAGraph()
638 + with torch.cuda.graph(graph, stream=warmup):
639 + step()
640 + torch.cuda.current_stream().wait_stream(warmup)
641 + for iteration in range(2):
642 + ids.copy_(_make_ids(head_sizes, 8, seed=600 + iteration))
643 + graph.replay()
644 + expected = _reference(weight.cuda(), scale_inv.cuda(), ids)
645 + if engram.use_sequence_parallel:
646 + expected = expected[tp_rank * tokens : (tp_rank + 1) * tokens]
647 + torch.testing.assert_close(output, expected, atol=0, rtol=0)
648 + torch.accelerator.synchronize()
649 +
650 +
651 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
652 +@pytest.mark.parametrize("tp_size", [1, 2])
653 +def test_engram_table_shared_across_dp_replicas(tp_size, monkeypatch):
654 + """Check all storage modes with one process launch per TP/DP topology."""
655 + if torch.accelerator.device_count() < WORLD_SIZE:
656 + pytest.skip(f"Need {WORLD_SIZE} GPUs to run the test.")
657 + monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1")
658 + mp.spawn(_worker, args=(tp_size, get_open_port()), nprocs=WORLD_SIZE)
659 + cleanup_dist_env_and_memory()
660 +
661 +
662 +@pytest.mark.parametrize(
663 + "cpu_offload,dp_size,error",
664 + [
665 + (False, 4, "requires cpu_offload=True"),
666 + (True, 1, "effective Engram DP size is 1"),
667 + ],
668 +)
669 +def test_engram_dp_shared_memory_runtime_requirements(
670 + monkeypatch, cpu_offload, dp_size, error
671 +):
672 + """Report unmet storage/topology requirements before allocating CUDA memory."""
673 + monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: dp_size)
674 + with pytest.raises(ValueError, match=error):
675 + ParallelEngramEmbedding(
676 + sum(HEAD_SIZES),
677 + DIM,
678 + HEAD_SIZES,
679 + cpu_offload=cpu_offload,
680 + dp_shared_memory=True,
681 + )
682 +
683 +
684 +@pytest.mark.parametrize(
685 + "node_ids,dp_size,replica_size,expected",
686 + [
687 + ([0] * 4, 1, 4, 1),
688 + ([0] * 4, 4, 1, 4),
689 + ([0] * 4, 2, 2, 2),
690 + ([0] * 8 + [1] * 8, 8, 2, 4),
691 + ([0] * 6 + [1] * 6, 4, 1, 2),
692 + ([0] * 3 + [1] * 4, 7, 1, 1),
693 + ([0] * 8 + [1] * 8 + [2] * 8, 8, 3, 1),
694 + ([0, 1] * 4, 4, 2, 1),
695 + ([0] * 3 + [1] * 5, 4, 2, 1),
696 + ],
697 +)
698 +def test_engram_dp_shard_size_respects_node_boundaries(
699 + node_ids, dp_size, replica_size, expected, monkeypatch
700 +):
701 + """Sharing must fall back when replicas or node rank ranges do not align."""
702 + monkeypatch.setattr(parallel_state, "get_node_count", lambda: len(set(node_ids)))
703 + monkeypatch.setattr(
704 + parallel_state, "get_world_group", lambda: SimpleNamespace(cpu_group=None)
705 + )
706 + monkeypatch.setattr(
707 + parallel_state,
708 + "in_the_same_node_as",
709 + lambda group, src: [node == node_ids[src] for node in node_ids],
710 + )
711 + assert (
712 + parallel_state._engram_dp_shard_size(len(node_ids), dp_size, replica_size)
713 + == expected
714 + )
715 +
716 +
717 +@pytest.mark.parametrize("num_tokens", [0, 2, 4, 5])
718 +def test_engram_hash_padding_has_no_valid_rows(num_tokens, monkeypatch):
719 + """Idle/padded slots must not query row zero; oversized batches are rejected."""
720 + monkeypatch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 4)
721 + monkeypatch.setattr(
722 + engram_ops,
723 + "get_engram_dp_group",
724 + lambda: SimpleNamespace(all_gather=lambda ids, dim: torch.cat([ids, ids], dim)),
725 + )
726 + ids = torch.full((num_tokens, 2, 3), 17, dtype=torch.int32)
727 + if num_tokens > 4:
728 + with pytest.raises(ValueError, match="exceeds the DP token slot"):
729 + gather_engram_hashes(ids)
730 + return
731 + gathered = gather_engram_hashes(ids).reshape(2, 4, 2, 3)
732 + for replica in gathered:
733 + torch.testing.assert_close(replica[:num_tokens], ids)
734 + assert torch.all(replica[num_tokens:] == engram_ops.DEAD_ID)
735 +
736 +
737 +@pytest.mark.parametrize("tp_size", [1, 2])
738 +@pytest.mark.parametrize("uniform", [False, True])
739 +def test_engram_hash_padding_ignores_other_nodes(tp_size, uniform, monkeypatch):
740 + """Two four-GPU nodes pad only to their own maximum, including TP x DP."""
741 + group_size = 4 // tp_size
742 + counts = (
743 + [4] * (8 // tp_size)
744 + if uniform
745 + else [4096] + [2] * (group_size - 1) + [1, 3] + [0] * (group_size - 2)
746 + )
747 + monkeypatch.setattr(
748 + engram_ops,
749 + "get_forward_context",
750 + lambda: SimpleNamespace(
751 + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=torch.tensor(counts))
752 + ),
753 + )
754 + for dp_rank in range(8 // tp_size):
755 + slot = 4 if uniform else (4096 if dp_rank < group_size else 3)
756 + ids = torch.full((counts[dp_rank], 2, 3), 17, dtype=torch.int32)
757 + monkeypatch.setattr(
758 + engram_ops,
759 + "get_dp_group",
760 + lambda rank=dp_rank: SimpleNamespace(rank_in_group=rank),
761 + )
762 + monkeypatch.setattr(
763 + engram_ops,
764 + "get_engram_dp_group",
765 + lambda rank=dp_rank: SimpleNamespace(
766 + rank_in_group=rank % group_size,
767 + world_size=group_size,
768 + all_gather=lambda tensor, dim: torch.cat([tensor] * group_size, dim),
769 + ),
770 + )
771 + gathered = gather_engram_hashes(ids)
772 + expected = ids.new_full((slot, 2, 3), engram_ops.DEAD_ID)
773 + expected[: ids.shape[0]] = ids
774 + torch.testing.assert_close(
775 + gathered, expected.repeat(group_size, 1, 1), msg=f"DP rank {dp_rank}"
776 + )
tests/engine/test_arg_utils.py
+22 -8
@@ -89,7 +89,10 @@ def test_watermark_config_cli():
89 @pytest.mark.parametrize(
90 "options",
91 [
92 - ["--engram-config", '{"cpu_offload": false, "embedding_across_dp": true}'],
92 + [
93 + "--engram-config",
94 + '{"cpu_offload": false, "embedding_across_dp": true}',
95 + ],
96 [
97 "--engram-config.cpu_offload",
98 "false",
@@ -98,9 +101,8 @@ def test_watermark_config_cli():
101 ],
102 ],
103 )
101 -def test_engram_config_cli(options, monkeypatch):
102 - """CLI settings take precedence over the legacy offload environment."""
103 - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1")
104 +def test_engram_config_cli(options):
105 + """JSON and dotted CLI options independently control Engram settings."""
106 parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
107 args = EngineArgs.from_cli_args(parser.parse_args(options))
108 assert args.engram_config is not None
@@ -109,14 +111,26 @@ def test_engram_config_cli(options, monkeypatch):
111
112
113 @pytest.mark.parametrize(
112 - "options, provided",
113 - [([], False), (["--engram-config", "{}"], True)],
114 + "options,provided,dp_shared_memory",
115 + [
116 + ([], False, False),
117 + (["--engram-config", "{}"], True, False),
118 + (
119 + ["--engram-config", '{"dp_shared_memory": true}'],
120 + True,
121 + True,
122 + ),
123 + (["--engram-config.dp_shared_memory", "true"], True, True),
124 + ],
125 )
115 -def test_engram_config_cli_optional(options, provided):
116 - """An explicit empty config must remain distinct from an omitted config."""
126 +def test_engram_config_cli_optional(options, provided, dp_shared_memory):
127 + """Explicit configs honor defaults and the JSON/dotted DP shared-memory flag."""
128 parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
129 args = EngineArgs.from_cli_args(parser.parse_args(options))
130 assert (args.engram_config is not None) == provided
131 + if provided:
132 + assert args.engram_config.cpu_offload is True
133 + assert args.engram_config.dp_shared_memory is dp_shared_memory
134
135
136 @pytest.mark.parametrize(
tests/kernels/test_engram.py
+226 -14
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import bisect
4 +import inspect
5 from types import SimpleNamespace
6
7 import pytest
@@ -8,10 +9,13 @@ import torch
9
10 from vllm.models.deepseek_v4_1.common import engram as engram_ops
11 from vllm.models.deepseek_v4_1.common.engram import (
11 - Engram,
12 + Engram as CommonEngram,
13 +)
14 +from vllm.models.deepseek_v4_1.common.engram import (
15 NgramHashState,
13 - ParallelEngramEmbedding,
16 )
17 +from vllm.models.deepseek_v4_1.nvidia import engram as nvidia_engram_ops
18 +from vllm.models.deepseek_v4_1.nvidia.engram import Engram, ParallelEngramEmbedding
19 from vllm.platforms import current_platform
20
21
@@ -107,7 +111,7 @@ def test_fused_engram_post_wkv_matches_reference(
111 )
112 monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: tp_rank)
113
110 - module = Engram.__new__(Engram)
114 + module = CommonEngram.__new__(CommonEngram)
115 torch.nn.Module.__init__(module)
116 module.dim = dim
117 module.hc_mult = hc_mult
@@ -116,7 +120,7 @@ def test_fused_engram_post_wkv_matches_reference(
120 module.use_sequence_parallel = use_sequence_parallel
121 module.embed_tokens = torch.nn.Identity()
122 # `forward` reads rows staged by `prepare_embeddings`, so inject kv there.
119 - module.embed_tokens.tp_size = 1
123 + module.embed_tokens.tp_size = module.embed_tokens.n_hash_cols = 1
124 module.staged_rows = kv.unsqueeze(1)
125 if use_sequence_parallel:
126 padded = torch.nn.functional.pad(kv, (0, 0, 0, (-num_kv_tokens) % tp_size))
@@ -151,6 +155,25 @@ def _hash_state(use_slot_cache: bool) -> NgramHashState:
155 return state
156
157
158 +@pytest.mark.parametrize("num_tokens", [0, 7])
159 +def test_engram_dummy_hashes_leave_history_untouched(num_tokens):
160 + """Dummy forwards have no valid table IDs and do not alter cached history."""
161 + state = _hash_state(use_slot_cache=True)
162 + state.multipliers = torch.empty(2, 4, dtype=torch.int64)
163 + state.primes = torch.empty(2, 3, 8, dtype=torch.int64)
164 + state._cache = torch.arange(16, dtype=torch.int32)
165 + history = state._cache.clone()
166 +
167 + hashes, keep = state.dummy_hashes(torch.arange(num_tokens))
168 +
169 + assert hashes.shape == (num_tokens, 2, 24)
170 + assert hashes.dtype == torch.int32
171 + assert torch.all(hashes == engram_ops.DEAD_ID)
172 + assert keep.shape == (num_tokens,) and keep.dtype == torch.bool
173 + assert not keep.any()
174 + torch.testing.assert_close(state._cache, history)
175 +
176 +
177 @pytest.mark.parametrize("num_blocks", [4, 100])
178 def test_engram_cache_follows_kv_cache_binding(num_blocks):
179 """Discard profiling history when the real KV cache replaces its storage."""
@@ -533,9 +556,11 @@ def _make_embedding(cpu_offload, rows=4096, dim=256, block=32):
556 layer = ParallelEngramEmbedding.__new__(ParallelEngramEmbedding)
557 torch.nn.Module.__init__(layer)
558 layer.dim, layer.block_size, layer.tp_size = dim, block, 1
559 + layer.dp_size = 1
560 layer.n_hash_cols = layer.part_n_hash_cols = 24
561 layer.head_start = 0
562 layer.cpu_offload = cpu_offload
563 + layer.dp_shared_memory = False
564 layer.part_num_embeddings = rows
565 # A window strictly inside the table, so unowned rows are exercised too.
566 layer.vocab_start_idx, layer.vocab_end_idx = rows // 4, rows // 4 + rows // 2
@@ -559,12 +584,27 @@ def _make_embedding(cpu_offload, rows=4096, dim=256, block=32):
584 return layer
585
586
587 +@pytest.mark.parametrize(
588 + "tp_size,dp_size,n_heads", [(1, 4, 5), (2, 2, 5), (4, 1, 6), (8, 1, 6)]
589 +)
590 +def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch):
591 + """Reject empty owners before allocating weights or accessing CUDA."""
592 + monkeypatch.setattr(
593 + engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size
594 + )
595 + monkeypatch.setattr(nvidia_engram_ops, "get_engram_dp_size", lambda: dp_size)
596 + monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
597 + monkeypatch.setattr(nvidia_engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
598 + with pytest.raises(AssertionError, match="ranks without hash heads"):
599 + ParallelEngramEmbedding(n_heads * 17, 64, (17,) * n_heads)
600 +
601 +
602 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
603 @pytest.mark.parametrize("cpu_offload", [False, True])
604 @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
605 def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch):
606 """Keep complete buckets and reconstruct head order, including TP padding."""
567 - head_sizes = (17, 19, 23, 29, 31, 37)
607 + head_sizes = (17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73)
608 num_rows, dim = sum(head_sizes), 64
609 torch.manual_seed(0)
610 weight = torch.randn(num_rows + 7, dim).to(torch.float8_e4m3fn)
@@ -623,7 +663,7 @@ def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeyp
663 module.staged_rows = torch.empty_like(shards[0])
664 module.prepare_embeddings(ids)
665 torch.testing.assert_close(module.embed(ids), expected, rtol=0, atol=0)
626 - # Slicing before head reordering must preserve padded heads and empty owners.
666 + # Slicing before head reordering must preserve padded heads and tokens.
667 module.use_sequence_parallel = True
668 chunk = (len(ids) + tp_size - 1) // tp_size
669 padded = torch.nn.functional.pad(expected, (0, 0, 0, 0, 0, (-len(ids)) % tp_size))
@@ -636,6 +676,31 @@ def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeyp
676 )
677
678
679 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
680 +def test_engram_lookup_reuses_jit_across_token_shapes():
681 + """Runtime token counts and launch grids must share one JIT variant."""
682 + layer = _make_embedding(cpu_offload=False)
683 + kernel = engram_ops._engram_lookup_kernel
684 + kernel_cache = kernel.device_caches[torch.accelerator.current_device_index()][0]
685 + kernel_cache.clear()
686 +
687 + cache_sizes = []
688 + for background in (False, True):
689 + for num_tokens in (1, 7, 256):
690 + ids = torch.zeros(num_tokens, 24, dtype=torch.int32, device="cuda")
691 + out = torch.empty(
692 + num_tokens,
693 + 24,
694 + layer.dim,
695 + dtype=torch.bfloat16,
696 + device="cuda",
697 + )
698 + layer.lookup(ids, out, background=background)
699 + cache_sizes.append(len(kernel_cache))
700 +
701 + assert cache_sizes == [1] * len(cache_sizes)
702 +
703 +
704 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
705 @pytest.mark.parametrize("cpu_offload", [False, True])
706 @pytest.mark.parametrize("background", [False, True])
@@ -660,19 +725,147 @@ def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens):
725
726
727 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
663 -@pytest.mark.parametrize("cpu_offload", [False, True])
728 +@pytest.mark.parametrize(
729 + "backend,cpu_offload", [("nvidia", None), ("nvidia", False), ("common", False)]
730 +)
731 +def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload):
732 + """Default offload uses pinned storage and a side stream; False uses HBM."""
733 + from vllm.config import EngramConfig
734 +
735 + config = SimpleNamespace(hidden_size=16, hc_mult=1, rms_norm_eps=1e-6)
736 + layout = SimpleNamespace(
737 + num_embeddings=(72,),
738 + head_dim=32,
739 + max_ngram_size=2,
740 + n_heads=24,
741 + primes=(((3,) * 24,),),
742 + )
743 + vllm_config = SimpleNamespace(
744 + engram_config=EngramConfig()
745 + if cpu_offload is None
746 + else EngramConfig(cpu_offload=cpu_offload),
747 + scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
748 + )
749 + offloaded = backend == "nvidia" and vllm_config.engram_config.cpu_offload
750 + if backend == "common":
751 + vllm_config.engram_config = None
752 + monkeypatch.setattr(engram_ops, "get_current_vllm_config", lambda: vllm_config)
753 + monkeypatch.setattr(
754 + nvidia_engram_ops, "get_current_vllm_config", lambda: vllm_config
755 + )
756 + monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1)
757 + monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
758 + monkeypatch.setattr(
759 + engram_ops, "ReplicatedLinear", lambda *a, **k: torch.nn.Identity()
760 + )
761 + with torch.device("cuda"):
762 + cls = Engram if backend == "nvidia" else CommonEngram
763 + module = cls(config, None, layout, 0, False, "engram")
764 + layer = module.embed_tokens
765 + if offloaded:
766 + assert layer.weight.device.type == "cpu" and layer.weight.is_pinned()
767 + else:
768 + assert layer.weight.device.type == "cuda"
769 + layer.weight.data.copy_(torch.randn(layer.weight.shape).to(torch.float8_e4m3fn))
770 + layer.weight_scale_inv.data.fill_(127)
771 + streams = []
772 + lookup = layer.lookup
773 +
774 + def track_lookup(ids, out, background=False):
775 + streams.append(torch.cuda.current_stream())
776 + lookup(ids, out, background=background)
777 +
778 + layer.lookup = track_lookup
779 + main = torch.cuda.current_stream()
780 + ids = torch.randint(
781 + layer.vocab_start_idx,
782 + layer.vocab_end_idx,
783 + (8, 24),
784 + device="cuda",
785 + dtype=torch.int32,
786 + )
787 + module.prepare_embeddings(ids)
788 + expected = _reference_lookup(
789 + layer.weight.cuda(),
790 + layer.weight_scale_inv.cuda(),
791 + ids,
792 + layer.vocab_start_idx,
793 + layer.vocab_end_idx,
794 + )
795 + torch.testing.assert_close(module.embed(ids), expected, atol=0, rtol=0)
796 + if offloaded:
797 + assert streams == [module._prefetch_stream]
798 + assert streams[0] != main
799 + else:
800 + assert streams == [main]
801 +
802 +
803 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
804 +@pytest.mark.parametrize(
805 + "cpu_offload,delay",
806 + [(False, None), (True, None), (True, "producer"), (True, "lookup")],
807 +)
808 @pytest.mark.parametrize("capture", ["eager", "full", "breakable"])
665 -def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture):
666 - """Consume early lookup results after a break, with fresh IDs each replay."""
667 - _run_engram_prepared_rows(cpu_offload, capture)
809 +def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture, delay):
810 + """Temporary lookup IDs survive allocator reuse and graph replay."""
811 + _run_engram_prepared_rows(cpu_offload, capture, delay=delay)
812 +
813 +
814 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
815 +@pytest.mark.parametrize(
816 + "missing_dependency", ["record_stream", "producer_wait", "consumer_wait"]
817 +)
818 +def test_engram_prefetch_detects_missing_dependency(monkeypatch, missing_dependency):
819 + """Negative controls must expose incorrect rows when a dependency is removed."""
820 + if (
821 + missing_dependency == "record_stream"
822 + and torch.cuda.memory.get_allocator_backend() != "native"
823 + ):
824 + pytest.skip(
825 + "The allocator reuse negative control requires the native allocator"
826 + )
827 +
828 + def start(self, hash_ids, rows, stream):
829 + if missing_dependency != "producer_wait":
830 + stream.wait_stream(torch.cuda.current_stream())
831 + if missing_dependency != "record_stream":
832 + hash_ids.record_stream(stream)
833 + with torch.cuda.stream(stream):
834 + self.embed_tokens.lookup(hash_ids, rows, background=True)
835 +
836 + def finish(self, stream):
837 + pass
838 +
839 + monkeypatch.setattr(Engram, "_start_prefetch", start)
840 + if missing_dependency == "consumer_wait":
841 + monkeypatch.setattr(Engram, "_finish_prefetch", finish)
842 + delay = "producer" if missing_dependency == "producer_wait" else "lookup"
843 + try:
844 + with pytest.raises(
845 + AssertionError, match=r"Tensor-likes are not (?:equal|close)!"
846 + ):
847 + _run_engram_prepared_rows(True, "eager", delay=delay)
848 + finally:
849 + torch.accelerator.synchronize()
850
851
852 def _run_engram_prepared_rows(
671 - cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False
853 + cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False, delay=None
854 ):
673 - from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
855 + from vllm.compilation.breakable_cudagraph import (
856 + BreakableCUDAGraphCapture,
857 + eager_break_during_capture,
858 + )
859
860 layer = _make_embedding(cpu_offload)
861 + if delay == "lookup":
862 + lookup = layer.lookup
863 +
864 + def delayed_lookup(indices, out, background=False):
865 + torch.cuda._sleep(2_000_000)
866 + lookup(indices, out, background=background)
867 +
868 + layer.lookup = delayed_lookup
869 cols, num_tokens = (23, 65) if use_sequence_parallel else (24, 64)
870 layer.n_hash_cols = cols
871 layer.tp_size = tp_size
@@ -681,6 +874,17 @@ def _run_engram_prepared_rows(
874 engram = Engram.__new__(Engram)
875 torch.nn.Module.__init__(engram)
876 engram.embed_tokens = layer
877 + engram._prefetch_stream = torch.cuda.Stream() if cpu_offload else None
878 + # Exercise the production eager boundaries even when the test process
879 + # imported Engram before breakable graphs were enabled.
880 + if capture == "breakable":
881 + import vllm.envs as envs
882 +
883 + with pytest.MonkeyPatch.context() as patch:
884 + patch.setattr(envs, "VLLM_USE_BREAKABLE_CUDAGRAPH", True)
885 + for name in ("_start_prefetch", "_finish_prefetch"):
886 + fn = eager_break_during_capture(inspect.unwrap(getattr(Engram, name)))
887 + setattr(engram, name, fn.__get__(engram, Engram))
888 engram.use_sequence_parallel = use_sequence_parallel
889 engram.staged_rows = torch.empty(
890 num_tokens,
@@ -709,7 +913,12 @@ def _run_engram_prepared_rows(
913 embed = torch.compile(embed, backend="eager", fullgraph=True, dynamic=True)
914
915 def step(cap=None):
712 - engram.prepare_embeddings(src)
916 + if delay == "producer":
917 + torch.cuda._sleep(2_000_000)
918 + # Drop the last reference to a non-contiguous input after launching lookup.
919 + engram.prepare_embeddings(hashes.clone()[:, 1])
920 + # Exercise same-size allocator reuse before consuming the prefetched rows.
921 + torch.empty_like(hashes).fill_(layer.part_num_embeddings - 1)
922 if cap is not None:
923 cap.add_eager(lambda: None)
924 out.copy_(embed(src))
@@ -730,9 +939,12 @@ def _run_engram_prepared_rows(
939 with torch.cuda.stream(warmup), graph:
940 step(graph)
941 torch.cuda.current_stream().wait_stream(warmup)
733 - assert graph.num_graphs == 2
942 + assert graph.num_eager_breaks == (3 if cpu_offload else 1)
943
944 for _ in range(3):
945 + # Eager prefill can overwrite staging rows between decode replays.
946 + hashes.random_(0, layer.part_num_embeddings)
947 + step()
948 hashes.random_(0, layer.part_num_embeddings)
949 if graph is None:
950 step()
tests/test_config.py
+110 -48
@@ -1158,31 +1158,68 @@ def test_engram_tensor_parallel_size(dp_size: int, across_dp: bool, expected: in
1158 assert config.get_parallel_size(parallel) == expected
1159
1160
1161 -def test_engram_rejects_elastic_cross_dp():
1161 +@pytest.mark.parametrize("option", ["embedding_across_dp", "dp_shared_memory"])
1162 +def test_engram_rejects_elastic_cross_dp(option):
1163 parallel = ParallelConfig(
1164 tensor_parallel_size=4,
1165 data_parallel_size=2,
1166 distributed_executor_backend="mp",
1167 )
1168 parallel.enable_elastic_ep = True
1168 - with pytest.raises(ValueError, match="embedding_across_dp.*elastic EP"):
1169 - EngramConfig(embedding_across_dp=True).verify_parallel_config(parallel)
1169 + with pytest.raises(ValueError, match=f"{option}.*elastic EP"):
1170 + EngramConfig(**{option: True}).verify_parallel_config(parallel)
1171
1172
1172 -@pytest.mark.parametrize("legacy", [None, "0", "1"])
1173 -def test_engram_cpu_offload_environment_fallback(monkeypatch, legacy):
1174 - """Explicit settings must override the legacy environment fallback."""
1175 - monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False)
1176 - if legacy is not None:
1177 - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
1178 - assert EngramConfig().cpu_offload == (legacy == "1")
1179 - assert EngramConfig(cpu_offload=False).cpu_offload is False
1180 - assert EngramConfig(cpu_offload=True).cpu_offload is True
1173 +def test_engram_dp_shared_memory_requires_cpu_offload():
1174 + with pytest.raises(ValueError, match="requires cpu_offload"):
1175 + EngramConfig(cpu_offload=False, dp_shared_memory=True)
1176 +
1177 +
1178 +@pytest.mark.parametrize(
1179 + "dp_size,load_format,multithread,error",
1180 + [
1181 + (1, "auto", False, "requires data_parallel_size > 1"),
1182 + (2, "dummy", False, "requires load_format"),
1183 + (2, "sharded_state", False, "requires load_format"),
1184 + (2, "auto", False, None),
1185 + (2, "safetensors", True, None),
1186 + (2, "pt", True, None),
1187 + ],
1188 +)
1189 +def test_engram_dp_shared_memory_config_validation(
1190 + monkeypatch, dp_size, load_format, multithread, error
1191 +):
1192 + """Reject invalid shared configs before distributed init; allow threaded loads."""
1193 + monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
1194 + config = cast(
1195 + VllmConfig,
1196 + SimpleNamespace(
1197 + model_config=SimpleNamespace(
1198 + architecture="DeepseekV41ForCausalLM",
1199 + hf_text_config=SimpleNamespace(engram_layer_ids=[1]),
1200 + ),
1201 + speculative_config=None,
1202 + engram_config=EngramConfig(cpu_offload=True, dp_shared_memory=True),
1203 + parallel_config=ParallelConfig(data_parallel_size=dp_size),
1204 + load_config=LoadConfig(
1205 + load_format=load_format,
1206 + model_loader_extra_config={"enable_multithread_load": multithread},
1207 + ),
1208 + ),
1209 + )
1210 + if error:
1211 + with pytest.raises(ValueError, match=error):
1212 + VllmConfig._resolve_and_verify_engram_config(config)
1213 + else:
1214 + VllmConfig._resolve_and_verify_engram_config(config)
1215
1216
1217 @pytest.mark.parametrize(
1218 "architecture, ple_layers, cuda, supported",
1219 [
1220 + ("DeepseekV41ForCausalLM", [1], True, True),
1221 + ("DeepseekV41ForCausalLM", [], True, False),
1222 + ("DeepseekV41ForCausalLM", [1], False, False),
1223 ("Qwen4ExpForCausalLM", [1], True, True),
1224 ("Qwen4ExpForConditionalGeneration", [1], True, True),
1225 ("Qwen4ExpForCausalLM", [], True, False),
@@ -1201,7 +1238,9 @@ def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, suppo
1238 ModelConfig,
1239 SimpleNamespace(
1240 architecture=architecture,
1204 - hf_text_config=SimpleNamespace(ple_layer_ids=ple_layers),
1241 + hf_text_config=SimpleNamespace(
1242 + ple_layer_ids=ple_layers, engram_layer_ids=ple_layers
1243 + ),
1244 ),
1245 )
1246 if architecture is not None
@@ -1214,55 +1253,73 @@ def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, suppo
1253 with pytest.raises(ValueError, match="requires a model with supported Engram"):
1254 config.verify_model_config(model)
1255
1256 + resolved = cast(
1257 + VllmConfig,
1258 + SimpleNamespace(
1259 + model_config=model,
1260 + speculative_config=None,
1261 + parallel_config=ParallelConfig(),
1262 + load_config=LoadConfig(load_format="dummy"),
1263 + engram_config=None,
1264 + ),
1265 + )
1266 + VllmConfig._resolve_and_verify_engram_config(resolved)
1267 + assert (resolved.engram_config is not None) == supported
1268 + if supported:
1269 + assert resolved.engram_config.cpu_offload is True
1270
1218 -def test_engram_config_defaults_to_none(monkeypatch):
1271 +
1272 +@pytest.mark.parametrize(
1273 + ("value", "expected"), [(None, True), ("0", False), ("1", True)]
1274 +)
1275 +def test_engram_cpu_offload_environment_default(monkeypatch, value, expected):
1276 monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False)
1277 + if value is not None:
1278 + monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", value)
1279 + assert EngramConfig().cpu_offload is expected
1280 + assert EngramConfig(cpu_offload=False).cpu_offload is False
1281 + assert EngramConfig(cpu_offload=True).cpu_offload is True
1282 +
1283 +
1284 +def test_engram_config_defaults_to_none():
1285 config = VllmConfig()
1286 assert config.engram_config is None
1287 assert config.compute_hash()
1288
1289
1225 -@pytest.mark.parametrize("legacy", ["0", "1"])
1226 -def test_engram_none_resolves_legacy_offload(monkeypatch, legacy):
1227 - """Legacy enablement materializes a config before model validation."""
1228 - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
1229 - monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
1290 +@pytest.mark.parametrize("explicit", [False, True])
1291 +@pytest.mark.parametrize("enable_dbo,ubatch_size", [(True, 0), (False, 2)])
1292 +def test_deepseek_engram_rejects_microbatching(explicit, enable_dbo, ubatch_size):
1293 + """Reject overlapping Engram staging even without an explicit EngramConfig."""
1294 config = cast(
1295 VllmConfig,
1296 SimpleNamespace(
1297 model_config=SimpleNamespace(
1234 - architecture="Qwen4ExpForCausalLM",
1235 - hf_text_config=SimpleNamespace(ple_layer_ids=[1]),
1298 + architecture="DeepseekV41ForCausalLM",
1299 + hf_text_config=SimpleNamespace(engram_layer_ids=[1]),
1300 ),
1301 speculative_config=None,
1238 - engram_config=None,
1239 - parallel_config=ParallelConfig(),
1302 + engram_config=EngramConfig(cpu_offload=False) if explicit else None,
1303 + parallel_config=ParallelConfig(
1304 + enable_dbo=enable_dbo, ubatch_size=ubatch_size
1305 + ),
1306 ),
1307 )
1242 - VllmConfig._resolve_and_verify_engram_config(config)
1243 - if legacy == "1":
1244 - assert config.engram_config is not None
1245 - assert config.engram_config.cpu_offload is True
1246 - else:
1247 - assert config.engram_config is None
1308 + with pytest.raises(
1309 + ValueError, match="Engram does not support DBO or microbatching"
1310 + ):
1311 + VllmConfig._resolve_and_verify_engram_config(config)
1312
1313
1250 -@pytest.mark.parametrize("legacy", ["0", "1"])
1251 -def test_engram_explicit_config_requires_supported_model(monkeypatch, legacy):
1314 +def test_engram_explicit_config_requires_supported_model():
1315 """Explicit all-false settings still opt into model validation."""
1253 - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
1316 with pytest.raises(ValueError, match="requires a model with supported Engram"):
1317 VllmConfig(engram_config=EngramConfig(cpu_offload=False))
1318
1319
1258 -def test_engram_legacy_offload_requires_supported_model(monkeypatch):
1259 - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1")
1260 - with pytest.raises(ValueError, match="requires a model with supported Engram"):
1261 - VllmConfig()
1262 -
1263 -
1320 @pytest.mark.parametrize("target_has_ple", [False, True])
1265 -def test_engram_draft_config_validates_target(monkeypatch, target_has_ple):
1321 +@pytest.mark.parametrize("explicit", [False, True])
1322 +def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, explicit):
1323 """MTP may inherit cross-DP sharding without having its own PLE layers."""
1324 monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
1325 target = SimpleNamespace(
@@ -1277,24 +1334,29 @@ def test_engram_draft_config_validates_target(monkeypatch, target_has_ple):
1334 speculative_config=SimpleNamespace(
1335 draft_model_config=draft, target_model_config=target
1336 ),
1280 - engram_config=EngramConfig(embedding_across_dp=True),
1337 + engram_config=EngramConfig(embedding_across_dp=True) if explicit else None,
1338 parallel_config=ParallelConfig(),
1339 + load_config=LoadConfig(load_format="dummy"),
1340 ),
1341 )
1284 - if target_has_ple:
1285 - VllmConfig._resolve_and_verify_engram_config(config)
1286 - else:
1342 + if explicit and not target_has_ple:
1343 with pytest.raises(ValueError, match="requires a model with supported Engram"):
1344 VllmConfig._resolve_and_verify_engram_config(config)
1345 + else:
1346 + VllmConfig._resolve_and_verify_engram_config(config)
1347 + assert (config.engram_config is not None) == target_has_ple
1348 + if target_has_ple:
1349 + assert config.engram_config.cpu_offload is True
1350
1351
1291 -def test_engram_hash_tracks_storage_and_sharding():
1352 +def test_engram_hash_tracks_execution_options():
1353 configs = [
1293 - EngramConfig(cpu_offload=offload, embedding_across_dp=across_dp)
1294 - for offload in (False, True)
1295 - for across_dp in (False, True)
1354 + EngramConfig(cpu_offload=True),
1355 + EngramConfig(cpu_offload=False),
1356 + EngramConfig(cpu_offload=True, embedding_across_dp=True),
1357 + EngramConfig(cpu_offload=True, dp_shared_memory=True),
1358 ]
1297 - assert len({config.compute_hash() for config in configs}) == 4
1359 + assert len({config.compute_hash() for config in configs}) == len(configs)
1360
1361
1362 @pytest.mark.parametrize("port", [1, 29550, 65535])
vllm/config/engram.py
+45 -15
@@ -3,18 +3,17 @@
3
4 from typing import TYPE_CHECKING
5
6 -from pydantic import Field
6 +from pydantic import Field, model_validator
7 +from typing_extensions import Self
8
9 import vllm.envs as envs
10 from vllm.config.utils import config, get_hash_factors, hash_factors
10 -from vllm.logger import init_logger
11
12 if TYPE_CHECKING:
13 + from vllm.config.load import LoadConfig
14 from vllm.config.model import ModelConfig
15 from vllm.config.parallel import ParallelConfig
16
16 -logger = init_logger(__name__)
17 -
17 # Architecture -> the hf_text_config field naming its n-gram layers. A model is
18 # only configurable here if it actually has such layers to store.
19 _NGRAM_LAYER_FIELDS = {
@@ -25,28 +24,44 @@ _NGRAM_LAYER_FIELDS = {
24
25
26 def _default_cpu_offload() -> bool:
28 - """Honor the legacy environment variable only when the field is omitted."""
29 - if envs.is_set("VLLM_PLE_CPU_OFFLOAD"):
30 - logger.warning_once(
31 - "VLLM_PLE_CPU_OFFLOAD is a legacy setting and may be removed in a "
32 - "future release. Use --engram-config.cpu_offload instead."
33 - )
27 return envs.VLLM_PLE_CPU_OFFLOAD
28
29
30 +def model_has_engram_layers(model_config: "ModelConfig | None") -> bool:
31 + """Whether the model carries n-gram embedding layers."""
32 + if model_config is None:
33 + return False
34 + field = _NGRAM_LAYER_FIELDS.get(model_config.architecture)
35 + if field is None:
36 + return False
37 + return bool(getattr(model_config.hf_text_config, field, None))
38 +
39 +
40 @config
41 class EngramConfig:
42 """Configuration for Engram embedding storage and sharding."""
43
44 cpu_offload: bool = Field(default_factory=_default_cpu_offload)
45 """Store embedding weights in pinned CPU memory for UVA lookup.
43 - Defaults to False, or VLLM_PLE_CPU_OFFLOAD when set for compatibility.
44 - An explicit value takes precedence over the legacy environment variable."""
46 + Defaults to VLLM_PLE_CPU_OFFLOAD, which is enabled by default. An explicit
47 + value takes precedence over the environment variable."""
48
49 embedding_across_dp: bool = False
50 """Shard embeddings across TP and all DP ranks when enabled.
51 Otherwise, each DP rank has a separate TP-sharded embedding replica."""
52
53 + dp_shared_memory: bool = False
54 + """Share CPU-offloaded embedding weights between co-located
55 + DP replicas. Each node stores one copy of every TP shard, reducing host
56 + memory without per-step Engram DP collectives. Requires sufficient
57 + /dev/shm capacity and a shared IPC namespace."""
58 +
59 + @model_validator(mode="after")
60 + def _validate_shared_memory(self) -> Self:
61 + if self.dp_shared_memory and not self.cpu_offload:
62 + raise ValueError("dp_shared_memory requires cpu_offload=True")
63 + return self
64 +
65 def verify_model_config(self, model_config: "ModelConfig | None") -> None:
66 """Reject Engram configuration for models without n-gram embeddings."""
67 from vllm.platforms import current_platform
@@ -64,13 +79,16 @@ class EngramConfig:
79 ):
80 raise ValueError(
81 "EngramConfig requires a model with supported Engram "
67 - "embeddings. Currently only the CUDA Qwen4Exp and DeepSeek "
68 - "V4.1 implementations with non-empty n-gram layer ids are "
69 - "supported."
82 + "embeddings, non-empty n-gram layer ids, and CUDA."
83 )
84
85 def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None:
86 """Reject unsupported embedding parallel topologies."""
87 + if self.dp_shared_memory:
88 + if parallel_config.data_parallel_size <= 1:
89 + raise ValueError("dp_shared_memory requires data_parallel_size > 1.")
90 + if parallel_config.enable_elastic_ep:
91 + raise ValueError("dp_shared_memory is not supported with elastic EP.")
92 if (
93 self.embedding_across_dp
94 and parallel_config.data_parallel_size > 1
@@ -80,6 +98,18 @@ class EngramConfig:
98 "Engram embedding_across_dp is not supported with elastic EP yet."
99 )
100
101 + def verify_load_config(self, load_config: "LoadConfig") -> None:
102 + """Shared tables require a loader that invokes parameter weight callbacks."""
103 + if self.dp_shared_memory and load_config.load_format not in (
104 + "auto",
105 + "safetensors",
106 + "pt",
107 + ):
108 + raise ValueError(
109 + "dp_shared_memory requires load_format 'auto', "
110 + f"'safetensors' or 'pt'; got {load_config.load_format!r}."
111 + )
112 +
113 def get_parallel_size(self, parallel_config: "ParallelConfig") -> int:
114 """Derive the embedding group size from the parallel configuration."""
115 size = parallel_config.tensor_parallel_size
vllm/config/vllm.py
+23 -8
@@ -34,7 +34,7 @@ from .device import DeviceConfig
34 from .diffusion import DiffusionConfig
35 from .ec_manager_config import EncoderCacheManagerConfig
36 from .ec_transfer import ECTransferConfig
37 -from .engram import EngramConfig
37 +from .engram import EngramConfig, model_has_engram_layers
38 from .kernel import KernelConfig
39 from .kv_events import KVEventsConfig
40 from .kv_transfer import KVTransferConfig
@@ -381,7 +381,7 @@ class VllmConfig:
381 attention_config: AttentionConfig = Field(default_factory=AttentionConfig)
382 """Attention configuration."""
383 engram_config: EngramConfig | None = None
384 - """Optional Engram configuration, only valid for supported PLE models."""
384 + """N-gram embedding storage and sharding settings."""
385 mamba_config: MambaConfig = Field(default_factory=MambaConfig)
386 """Mamba configuration."""
387 kernel_config: KernelConfig = Field(default_factory=KernelConfig)
@@ -1218,22 +1218,37 @@ class VllmConfig:
1218 )
1219
1220 def _resolve_and_verify_engram_config(self) -> None:
1221 - """Resolve legacy offload settings and validate model and parallel configs."""
1222 - if self.engram_config is None:
1223 - if not envs.VLLM_PLE_CPU_OFFLOAD:
1224 - return
1225 - self.engram_config = EngramConfig()
1221 + """Resolve defaults and validate n-gram embedding settings."""
1222 + from vllm.platforms import current_platform
1223 +
1224 model_config = self.model_config
1225 speculative_config = self.speculative_config
1226 # Draft configs inherit the target's communication groups and settings.
1229 - # Qwen4Exp MTP itself disables PLE, so validate its target instead.
1227 + # Validate the target because the draft may disable n-gram embeddings.
1228 if (
1229 speculative_config is not None
1230 and model_config is speculative_config.draft_model_config
1231 ):
1232 model_config = speculative_config.target_model_config
1233 + if (
1234 + model_config is not None
1235 + and model_config.architecture == "DeepseekV41ForCausalLM"
1236 + and getattr(model_config.hf_text_config, "engram_layer_ids", None)
1237 + and self.parallel_config.use_ubatching
1238 + ):
1239 + raise ValueError(
1240 + "DeepSeek V4.1 Engram does not support DBO or microbatching. "
1241 + "Disable --enable-dbo and set --ubatch-size to 0."
1242 + )
1243 + if self.engram_config is None:
1244 + if not current_platform.is_cuda() or not model_has_engram_layers(
1245 + model_config
1246 + ):
1247 + return
1248 + self.engram_config = EngramConfig()
1249 self.engram_config.verify_model_config(model_config)
1250 self.engram_config.verify_parallel_config(self.parallel_config)
1251 + self.engram_config.verify_load_config(self.load_config)
1252 logger.info_once("Resolved Engram configuration: %s", str(self.engram_config))
1253
1254 def __post_init__(self):
vllm/distributed/parallel_state.py
+80
@@ -32,6 +32,7 @@ from collections.abc import Callable
32 from contextlib import AbstractContextManager, contextmanager, nullcontext
33 from dataclasses import dataclass
34 from datetime import timedelta
35 +from math import gcd
36 from multiprocessing import shared_memory
37 from typing import TYPE_CHECKING, Any, Protocol
38 from unittest.mock import patch
@@ -1558,6 +1559,23 @@ def get_etp_group() -> GroupCoordinator:
1559 return _ETP
1560
1561
1562 +_ENGRAM_DP: GroupCoordinator | None = None
1563 +
1564 +
1565 +def get_engram_dp_group() -> GroupCoordinator | None:
1566 + """Return the DP replicas that share one engram embedding table.
1567 +
1568 + None when every replica holds a full (TP-sharded) copy, which is the case
1569 + for models without engram layers and for replicas that span nodes.
1570 + """
1571 + return _ENGRAM_DP
1572 +
1573 +
1574 +def get_engram_dp_size() -> int:
1575 + """Number of DP replicas one engram embedding table is sharded over."""
1576 + return _ENGRAM_DP.world_size if _ENGRAM_DP is not None else 1
1577 +
1578 +
1579 _DCP: GroupCoordinator | None = None
1580
1581
@@ -1933,6 +1951,29 @@ def _elastic_join_warmup_ctx() -> AbstractContextManager[Any]:
1951 return nullcontext()
1952
1953
1954 +def _engram_dp_shard_size(
1955 + world_size: int,
1956 + data_parallel_size: int,
1957 + replica_size: int,
1958 +) -> int:
1959 + """Size of DP shards aligned with complete replicas and node boundaries."""
1960 + node_count = get_node_count()
1961 + replicas_per_node, remainder = divmod(world_size, node_count * replica_size)
1962 + if remainder:
1963 + return 1
1964 + # Shards must divide both the DP group and each node's replica count.
1965 + shard_size = gcd(data_parallel_size, replicas_per_node)
1966 + if shard_size == 1 or node_count == 1:
1967 + return shard_size
1968 + ranks_per_node = replicas_per_node * replica_size
1969 + for start in range(0, world_size, ranks_per_node):
1970 + same_node = in_the_same_node_as(get_world_group().cpu_group, start)
1971 + node_ranks = [rank for rank, local in enumerate(same_node) if local]
1972 + if node_ranks != list(range(start, start + ranks_per_node)):
1973 + return 1
1974 + return shard_size
1975 +
1976 +
1977 def initialize_model_parallel(
1978 tensor_model_parallel_size: int = 1,
1979 pipeline_model_parallel_size: int = 1,
@@ -2055,6 +2096,40 @@ def initialize_model_parallel(
2096 group_name="etp",
2097 )
2098
2099 + # Build the node-local DP group used to shard or share an Engram table.
2100 + global _ENGRAM_DP
2101 + assert _ENGRAM_DP is None, "engram data parallel group is already initialized"
2102 + engram_dp_size = 1
2103 + engram_config = config.engram_config
2104 + if (
2105 + engram_config is not None
2106 + and config.model_config is not None
2107 + and config.model_config.architecture == "DeepseekV41ForCausalLM"
2108 + and not enable_elastic_ep
2109 + ):
2110 + engram_dp_size = _engram_dp_shard_size(
2111 + world_size,
2112 + data_parallel_size,
2113 + tensor_model_parallel_size
2114 + * pipeline_model_parallel_size
2115 + * prefill_context_model_parallel_size,
2116 + )
2117 + if engram_dp_size > 1:
2118 + group_ranks = (
2119 + all_ranks.permute(0, 2, 3, 4, 1).reshape(-1, engram_dp_size).unbind(0)
2120 + )
2121 + _ENGRAM_DP = init_model_parallel_group(
2122 + [x.tolist() for x in group_ranks],
2123 + get_world_group().local_rank,
2124 + backend,
2125 + group_name="edp",
2126 + )
2127 + logger.info_once(
2128 + "Engram node-local DP group size: %d (TP=%d)",
2129 + engram_dp_size,
2130 + tensor_model_parallel_size,
2131 + )
2132 +
2133 # Build the DCP model-parallel groups.
2134 global _DCP
2135 assert _DCP is None, "decode context model parallel group is already initialized"
@@ -2304,6 +2379,11 @@ def destroy_model_parallel():
2379 _TP.destroy()
2380 _TP = None
2381
2382 + global _ENGRAM_DP
2383 + if _ENGRAM_DP:
2384 + _ENGRAM_DP.destroy()
2385 + _ENGRAM_DP = None
2386 +
2387 global _DCP
2388 if _DCP:
2389 _DCP.destroy()
vllm/envs.py
+3 -6
@@ -154,7 +154,7 @@ if TYPE_CHECKING:
154 VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False
155 VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
156 VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
157 - VLLM_PLE_CPU_OFFLOAD: bool = False
157 + VLLM_PLE_CPU_OFFLOAD: bool = True
158 VLLM_DISABLE_COMPILE_CACHE: bool = False
159 VLLM_REPLICATE_EMBED: bool = False
160 VLLM_USE_LAYERNAME: bool = True
@@ -2059,11 +2059,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
2059 "VLLM_LOG_MODEL_INSPECTION": lambda: bool(
2060 int(os.getenv("VLLM_LOG_MODEL_INSPECTION", "0"))
2061 ),
2062 - # Keep Qwen4Exp PLE embedding tables in pinned CPU memory and gather their
2063 - # rows through UVA on a dedicated CUDA stream.
2064 - # Legacy fallback for EngramConfig.cpu_offload, which takes precedence.
2065 - # This environment variable may be removed in a future release.
2066 - "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "0"))),
2062 + # Store n-gram embedding tables in pinned CPU memory for UVA lookup.
2063 + "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "1"))),
2064 # Debug logging for --enable-mfu-metrics
2065 "VLLM_DEBUG_MFU_METRICS": lambda: bool(
2066 int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0"))
vllm/models/deepseek_v4_1/common/engram.py
+112 -93
@@ -51,8 +51,6 @@ from vllm.model_executor.layers.linear import ReplicatedLinear
51 from vllm.model_executor.layers.quantization import QuantizationConfig
52 from vllm.model_executor.utils import set_weight_attrs
53 from vllm.triton_utils import tl, triton
54 -from vllm.utils.platform_utils import is_uva_available
55 -from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
54
55 logger = init_logger(__name__)
56
@@ -236,6 +234,7 @@ def _write_hash_cache_kernel(
234 tl.store(cache + slot, value, valid)
235
236
237 +# Keep request shapes out of the cache key so one warmed variant covers runtime batches.
238 @triton.jit(
239 do_not_specialize=[
240 "num_tokens",
@@ -462,6 +461,21 @@ class NgramHashState(nn.Module):
461 self._kv_cache_ref = weakref.ref(kv_cache)
462 return True
463
464 + def dummy_hashes(
465 + self, input_ids: torch.Tensor
466 + ) -> tuple[torch.Tensor, torch.Tensor]:
467 + """Participate in DP lookups without valid rows or hash-cache updates."""
468 + num_tokens = input_ids.shape[0]
469 + num_layers, max_ngram = self.multipliers.shape
470 + num_heads = self.primes.shape[-1]
471 + hashes = input_ids.new_full(
472 + (num_tokens, num_layers, (max_ngram - 1) * num_heads),
473 + DEAD_ID,
474 + dtype=torch.int32,
475 + )
476 + keep = torch.zeros(num_tokens, dtype=torch.bool, device=input_ids.device)
477 + return hashes, keep
478 +
479 def forward(
480 self,
481 input_ids: torch.Tensor,
@@ -561,7 +575,16 @@ def _engram_head_shard_weight_loader(
575 param.data.copy_(shard)
576
577
564 -@triton.jit
578 +@triton.jit(
579 + do_not_specialize=[
580 + "vocab_start",
581 + "vocab_end",
582 + "num_rows",
583 + "ids_stride_t",
584 + "ids_stride_h",
585 + "GRID",
586 + ]
587 +)
588 def _engram_lookup_kernel(
589 weight,
590 scales,
@@ -578,7 +601,7 @@ def _engram_lookup_kernel(
601 DIM: tl.constexpr,
602 QUANT_BLOCK: tl.constexpr,
603 BLOCK_R: tl.constexpr,
581 - GRID: tl.constexpr,
604 + GRID,
605 ):
606 """Gather fp8 rows, apply their ue8m0 block scales, write bf16.
607
@@ -620,12 +643,9 @@ def _engram_lookup_kernel(
643
644
645 class ParallelEngramEmbedding(nn.Module):
623 - """The n-gram hash table, sharded by complete hash heads over TP ranks.
624 - Rows stay fp8 and are dequantized with ue8m0 per-32 scales on lookup.
646 + """TP-sharded hash heads with FP8 rows and per-block E8M0 scales."""
647
626 - With `cpu_offload` the shard lives in pinned host memory and is read over
627 - UVA instead of HBM; the TP sharding is unchanged either way.
628 - """
648 + _weight_loader = staticmethod(_engram_head_shard_weight_loader)
649
650 def __init__(
651 self,
@@ -633,89 +653,55 @@ class ParallelEngramEmbedding(nn.Module):
653 dim: int,
654 head_sizes: tuple[int, ...],
655 block_size: int = 32,
636 - cpu_offload: bool = False,
637 - ):
656 + ) -> None:
657 super().__init__()
639 - tp_size = get_tensor_model_parallel_world_size()
640 - tp_rank = get_tensor_model_parallel_rank()
658 assert head_sizes and all(size > 0 for size in head_sizes)
659 assert sum(head_sizes) <= num_embeddings
643 - if cpu_offload and not is_uva_available():
644 - raise RuntimeError("Engram CPU offload requires UVA support")
660 self.num_embeddings = num_embeddings
661 self.dim = dim
662 self.block_size = block_size
663 self.n_hash_cols = len(head_sizes)
649 - self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, tp_size)
650 - self.head_start = tp_rank * self.part_n_hash_cols
664 + self.tp_size = get_tensor_model_parallel_world_size()
665 + num_shards, head_rank = self._get_shard_info()
666 + self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, num_shards)
667 + # TODO: Support row-wise sharding when there are too few hash heads.
668 + assert (num_shards - 1) * self.part_n_hash_cols < self.n_hash_cols, (
669 + f"Engram sharding leaves ranks without hash heads: "
670 + f"{self.n_hash_cols} heads over {num_shards} shards"
671 + )
672 + self.head_start = head_rank * self.part_n_hash_cols
673 head_end = self.head_start + self.part_n_hash_cols
674 self.vocab_start_idx = sum(head_sizes[: self.head_start])
675 self.vocab_end_idx = sum(head_sizes[:head_end])
676 self.part_num_embeddings = self.vocab_end_idx - self.vocab_start_idx
655 - self.tp_size = tp_size
656 - self.cpu_offload = cpu_offload
657 - self._views: tuple[torch.Tensor, torch.Tensor] | None = None
658 - self._view_src: tuple[int, int] | None = None
677 self._num_sms = torch.cuda.get_device_properties(
678 torch.accelerator.current_device_index()
679 ).multi_processor_count
662 -
663 - # Explicit device: model init runs under a `torch.device("cuda")`
664 - # context, which would otherwise put the shard in HBM.
665 - kwargs = {"device": "cpu", "pin_memory": True} if cpu_offload else {}
666 - self.weight = nn.Parameter(
667 - torch.empty(
668 - self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn, **kwargs
669 - ),
670 - requires_grad=False,
671 - )
672 - self.weight_scale_inv = nn.Parameter(
673 - torch.empty(
674 - self.part_num_embeddings,
675 - dim // block_size,
676 - dtype=torch.uint8,
677 - **kwargs,
678 - ),
679 - requires_grad=False,
680 - )
680 + weight, scales = self._allocate_weights()
681 + self.weight = nn.Parameter(weight, requires_grad=False)
682 + self.weight_scale_inv = nn.Parameter(scales, requires_grad=False)
683 for param in (self.weight, self.weight_scale_inv):
684 set_weight_attrs(
685 param,
686 {
685 - "weight_loader": _engram_head_shard_weight_loader,
687 + "weight_loader": self._weight_loader,
688 "engram_vocab_start": self.vocab_start_idx,
689 },
690 )
689 - if cpu_offload:
690 - # Constant dummy values avoid randomizing huge CPU lookup tables.
691 - set_weight_attrs(self.weight, {"dummy_weight_value": 1.0})
692 - # The ue8m0 encoding of scale 1.0 is exponent byte 127.
693 - set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127})
694 - logger.info(
695 - "Engram table offloaded to pinned host memory: %d rows x %d, "
696 - "%.2f GiB per rank",
697 - self.part_num_embeddings,
698 - dim,
699 - self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
700 - )
691
702 - def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
703 - """Parameters when resident, else cached UVA views of the pinned shard.
692 + def _get_shard_info(self) -> tuple[int, int]:
693 + return self.tp_size, get_tensor_model_parallel_rank()
694
705 - Rebuilt if anything swaps `.data`, so a stale device pointer cannot
706 - survive silently.
707 - """
708 - if not self.cpu_offload:
709 - return self.weight.data, self.weight_scale_inv.data
710 - src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr())
711 - if self._view_src != src:
712 - self._views = (
713 - get_accelerator_view_from_cpu_tensor(self.weight.data),
714 - get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data),
715 - )
716 - self._view_src = src
717 - assert self._views is not None
718 - return self._views
695 + def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
696 + return (
697 + torch.empty(self.part_num_embeddings, self.dim, dtype=torch.float8_e4m3fn),
698 + torch.empty(
699 + self.part_num_embeddings, self.dim // self.block_size, dtype=torch.uint8
700 + ),
701 + )
702 +
703 + def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
704 + return self.weight.data, self.weight_scale_inv.data
705
706 def lookup(
707 self, indices: torch.Tensor, out: torch.Tensor, background: bool = False
@@ -753,7 +739,7 @@ class ParallelEngramEmbedding(nn.Module):
739
740 def forward(self, indices: torch.Tensor) -> torch.Tensor:
741 """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim]
756 - bf16, gathered from all TP shards."""
742 + bf16, gathered from all shards for this replica's tokens."""
743 out = torch.empty(
744 (indices.shape[0], self.part_n_hash_cols, self.dim),
745 dtype=torch.bfloat16,
@@ -762,8 +748,7 @@ class ParallelEngramEmbedding(nn.Module):
748 self.lookup(indices, out)
749 if self.tp_size > 1:
750 out = tensor_model_parallel_all_gather(out, dim=1)
765 - out = out[:, : self.n_hash_cols]
766 - return out
751 + return out[:, : self.n_hash_cols]
752
753
754 @triton.jit(do_not_specialize=["num_kv_tokens"])
@@ -859,7 +844,7 @@ def _fused_engram_post_wkv_kernel(
844
845
846 @triton.jit(do_not_specialize=["num_tokens", "token_start", "num_elements"])
862 -def _engram_sp_rows_kernel(
847 +def _engram_select_rows_kernel(
848 gathered,
849 output,
850 num_tokens,
@@ -880,6 +865,32 @@ def _engram_sp_rows_kernel(
865 tl.store(output + offsets, values, offsets < num_elements)
866
867
868 +def _engram_select_rows(
869 + gathered: torch.Tensor,
870 + output: torch.Tensor,
871 + source_tokens: int,
872 + token_start: int,
873 + local_width: int,
874 +) -> None:
875 + """Copy one token window out of a rank-major gathered buffer.
876 +
877 + Both gathers land rank-major ([rank][token][local width]); this walks the
878 + window the rank keeps and lays its ranks out side by side as width.
879 + """
880 + if output.numel() == 0:
881 + return
882 + _engram_select_rows_kernel[(triton.cdiv(output.numel(), 1024),)](
883 + gathered,
884 + output,
885 + source_tokens,
886 + token_start,
887 + output.numel(),
888 + local_width,
889 + output.shape[1] * output.shape[2],
890 + BLOCK_SIZE=1024,
891 + )
892 +
893 +
894 class Engram(nn.Module):
895 """Writes an n-gram lookup into the residual stream, gated by how well it
896 matches that stream.
@@ -910,13 +921,7 @@ class Engram(nn.Module):
921 # Named ``embed_tokens`` so the checkpoint's ``engram.embed.weight``
922 # survives the mapper's ``embed.weight`` -> ``embed_tokens.weight``
923 # suffix rule.
913 - engram_config = get_current_vllm_config().engram_config
914 - self.embed_tokens = ParallelEngramEmbedding(
915 - layout.num_embeddings[layer_hash_index],
916 - layout.head_dim,
917 - tuple(size for order in layout.primes[layer_hash_index] for size in order),
918 - cpu_offload=engram_config.cpu_offload if engram_config else True,
919 - )
924 + self.embed_tokens = self._create_embedding(layout, layer_hash_index)
925 n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads
926 self.wkv = ReplicatedLinear(
927 n_hash_cols * layout.head_dim,
@@ -936,40 +941,54 @@ class Engram(nn.Module):
941 )
942
943 max_tokens = get_current_vllm_config().scheduler_config.max_num_batched_tokens
939 - # Keep lookup results alive across breakable graph segments.
944 + self._init_staging(max_tokens, layout.head_dim)
945 +
946 + def _create_embedding(
947 + self, layout: EngramLayout, layer_hash_index: int
948 + ) -> ParallelEngramEmbedding:
949 + return ParallelEngramEmbedding(
950 + layout.num_embeddings[layer_hash_index],
951 + layout.head_dim,
952 + tuple(size for order in layout.primes[layer_hash_index] for size in order),
953 + )
954 +
955 + def _init_staging(self, max_tokens: int, head_dim: int) -> None:
956 + # Persistent storage keeps lookup addresses stable across graph replays.
957 self.staged_rows = torch.empty(
958 max_tokens,
959 self.embed_tokens.part_n_hash_cols,
943 - layout.head_dim,
960 + head_dim,
961 dtype=torch.bfloat16,
962 )
963
964 def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
948 - """Gather this layer's rows on the main stream before decoder layers."""
949 - self.embed_tokens.lookup(hash_ids, self.staged_rows[: hash_ids.shape[0]])
965 + """Look up rows before the decoder layers consume them."""
966 + rows = self.staged_rows[: hash_ids.shape[0]]
967 + assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small"
968 + self.embed_tokens.lookup(hash_ids, rows)
969 +
970 + def _ready_rows(self, num_tokens: int) -> torch.Tensor:
971 + return self.staged_rows[:num_tokens]
972
973 def embed(self, hash_ids: torch.Tensor) -> torch.Tensor:
974 """Gather heads, returning only local tokens when SP is enabled."""
953 - rows = self.staged_rows[: hash_ids.shape[0]]
975 + rows = self._ready_rows(hash_ids.shape[0])
976 if self.embed_tokens.tp_size == 1:
955 - return rows
977 + return rows[:, : self.embed_tokens.n_hash_cols]
978 if self.use_sequence_parallel:
979 tp_size = self.embed_tokens.tp_size
980 num_tokens, local_heads, dim = rows.shape
981 gathered = tensor_model_parallel_all_gather(rows, dim=0)
982 chunk = (num_tokens + tp_size - 1) // tp_size
961 - rows = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim))
962 - _engram_sp_rows_kernel[(triton.cdiv(rows.numel(), 1024),)](
983 + out = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim))
984 + _engram_select_rows(
985 gathered,
964 - rows,
986 + out,
987 num_tokens,
988 get_tensor_model_parallel_rank() * chunk,
967 - rows.numel(),
989 local_heads * dim,
969 - self.embed_tokens.n_hash_cols * dim,
970 - BLOCK_SIZE=1024,
990 )
972 - return rows
991 + return out
992 rows = tensor_model_parallel_all_gather(rows, dim=1)
993 return rows[:, : self.embed_tokens.n_hash_cols]
994
vllm/models/deepseek_v4_1/nvidia/engram.py new
+403
@@ -0,0 +1,403 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch."""
4 +
5 +import mmap
6 +import tempfile
7 +import weakref
8 +from contextlib import ExitStack
9 +
10 +import numpy as np
11 +import torch
12 +
13 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
14 +from vllm.config import get_current_vllm_config
15 +from vllm.distributed import (
16 + get_dp_group,
17 + get_engram_dp_group,
18 + get_engram_dp_size,
19 + get_tensor_model_parallel_rank,
20 + tensor_model_parallel_all_gather,
21 +)
22 +from vllm.distributed.parallel_state import GroupCoordinator
23 +from vllm.forward_context import get_forward_context
24 +from vllm.logger import init_logger
25 +from vllm.model_executor.utils import set_weight_attrs
26 +from vllm.models.deepseek_v4_1.common.engram import (
27 + DEAD_ID,
28 + EngramLayout,
29 + _engram_head_shard_weight_loader,
30 + _engram_select_rows,
31 +)
32 +from vllm.models.deepseek_v4_1.common.engram import (
33 + Engram as BaseEngram,
34 +)
35 +from vllm.models.deepseek_v4_1.common.engram import (
36 + ParallelEngramEmbedding as BaseParallelEngramEmbedding,
37 +)
38 +from vllm.utils.platform_utils import is_uva_available
39 +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
40 +
41 +logger = init_logger(__name__)
42 +
43 +
44 +def engram_head_shard_rank() -> int:
45 + """This rank's slot among the hash-head shards of one engram table.
46 +
47 + TP-major, so the shards a DP gather brings in are contiguous heads and
48 + the following TP gather completes the head order.
49 + """
50 + dp_group = get_engram_dp_group()
51 + dp_size = dp_group.world_size if dp_group is not None else 1
52 + dp_rank = dp_group.rank_in_group if dp_group is not None else 0
53 + return get_tensor_model_parallel_rank() * dp_size + dp_rank
54 +
55 +
56 +def engram_gathered_num_tokens() -> int:
57 + """Per-replica token slot for the node-local Engram DP group."""
58 + dp_metadata = get_forward_context().dp_metadata
59 + if dp_metadata is None:
60 + raise RuntimeError("a DP-shared engram table needs DP token metadata")
61 + group = get_engram_dp_group()
62 + assert group is not None
63 + # Engram groups are contiguous slices of the full DP group.
64 + start = get_dp_group().rank_in_group - group.rank_in_group
65 + return int(
66 + dp_metadata.num_tokens_across_dp_cpu[start : start + group.world_size].max()
67 + )
68 +
69 +
70 +def gather_engram_hashes(
71 + hash_ids: torch.Tensor, *, dp_shared_memory: bool = False
72 +) -> torch.Tensor:
73 + """Collect the n-gram ids of every DP replica sharing one table.
74 +
75 + Replicas are padded to a common token slot, so the gathered shape is
76 + static under CUDA graph capture (where DP already pads alike).
77 + """
78 + dp_group = get_engram_dp_group()
79 + if dp_group is None or dp_shared_memory:
80 + return hash_ids
81 + slot = engram_gathered_num_tokens()
82 + if hash_ids.shape[0] > slot:
83 + raise ValueError("Engram token count exceeds the DP token slot")
84 + if hash_ids.shape[0] < slot:
85 + pad = hash_ids.new_full(
86 + (slot - hash_ids.shape[0], *hash_ids.shape[1:]), DEAD_ID
87 + )
88 + hash_ids = torch.cat((hash_ids, pad))
89 + return dp_group.all_gather(hash_ids, dim=0)
90 +
91 +
92 +class DPSharedEngramStorage:
93 + """Registered host weights shared by a node-local DP group with one writer."""
94 +
95 + def __init__(
96 + self, num_rows: int, dim: int, block_size: int, group: GroupCoordinator
97 + ) -> None:
98 + self.group = group
99 + weight_bytes = num_rows * dim
100 + storage = self._allocate(weight_bytes + weight_bytes // block_size)
101 + self.weight = storage[:weight_bytes].view(torch.float8_e4m3fn).view(-1, dim)
102 + self.weight_scale_inv = storage[weight_bytes:].view(-1, dim // block_size)
103 + self._views: tuple[torch.Tensor, torch.Tensor] | None = None
104 +
105 + def _allocate(self, num_bytes: int) -> torch.Tensor:
106 + """Map and register one physical allocation across a node-local DP group."""
107 + from vllm.distributed.device_communicators.shm_broadcast import (
108 + check_shm_free_space,
109 + )
110 +
111 + group = self.group
112 + with ExitStack() as stack:
113 + path, error = None, None
114 + if group.rank_in_group == 0:
115 + try:
116 + check_shm_free_space(num_bytes)
117 + backing_file = stack.enter_context(
118 + tempfile.NamedTemporaryFile(
119 + prefix="vllm_engram_", dir="/dev/shm"
120 + )
121 + )
122 + backing_file.truncate(num_bytes)
123 + path = backing_file.name
124 + except Exception as exc:
125 + error = f"{type(exc).__name__}: {exc}"
126 + path, error = group.broadcast_object((path, error))
127 + if error is not None:
128 + raise RuntimeError(
129 + "Engram shared-memory creation failed on EDP rank 0: " + error
130 + )
131 +
132 + mapping = owner = tensor = finalizer = None
133 + stage = "open"
134 + try:
135 + try:
136 + with open(path, "r+b") as file:
137 + stage = "mmap"
138 + mapping = mmap.mmap(
139 + file.fileno(), num_bytes, flags=mmap.MAP_SHARED
140 + )
141 + stage = "cudaHostRegister"
142 + owner = np.frombuffer(mapping, dtype=np.uint8)
143 + pointer = owner.ctypes.data
144 + tensor = torch.from_numpy(owner)
145 + result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
146 + if result.value != 0:
147 + raise RuntimeError(f"cudaHostRegister failed: {result}")
148 + finalizer = weakref.finalize(
149 + owner, self._unregister, mapping, pointer
150 + )
151 + finalizer.atexit = False # type: ignore[misc]
152 + # The UVA helper otherwise allocates a private pinned copy.
153 + if not tensor.is_pinned():
154 + raise RuntimeError(
155 + "CUDA did not recognize the shared Engram registration"
156 + )
157 + except Exception as exc:
158 + error = f"{stage}: {type(exc).__name__}: {exc}"
159 +
160 + errors: list[str | None] = [None] * group.world_size
161 + # Also fences peer mappings before the leader unlinks the file.
162 + torch.distributed.all_gather_object(
163 + errors, error, group=group.cpu_group
164 + )
165 + failures = "; ".join(
166 + f"EDP rank {rank}: {error}"
167 + for rank, error in enumerate(errors)
168 + if error is not None
169 + )
170 + if failures:
171 + raise RuntimeError(
172 + "Engram shared-memory initialization failed: " + failures
173 + )
174 + assert tensor is not None
175 + return tensor
176 + except Exception:
177 + if finalizer is not None:
178 + finalizer()
179 + tensor = owner = None
180 + if mapping is not None:
181 + mapping.close()
182 + raise
183 +
184 + @staticmethod
185 + def _unregister(mapping: mmap.mmap, pointer: int) -> None:
186 + # Torch storage retains the numpy owner, including through cached UVA views.
187 + # Keep its mmap alive until CUDA has released the registration.
188 + result = torch.cuda.cudart().cudaHostUnregister(pointer)
189 + if result.value != 0:
190 + logger.warning("Engram cudaHostUnregister failed: %s", result)
191 +
192 + def load_weight(
193 + self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
194 + ) -> None:
195 + if self.group.rank_in_group == 0:
196 + _engram_head_shard_weight_loader(param, loaded_weight)
197 + # Read order may differ across ranks. Equal load counts ensure all shared
198 + # weights are ready after the last weight-loader call returns.
199 + torch.distributed.barrier(group=self.group.cpu_group)
200 +
201 + def get_views(
202 + self, weight: torch.Tensor, scales: torch.Tensor
203 + ) -> tuple[torch.Tensor, torch.Tensor]:
204 + if (weight.data_ptr(), scales.data_ptr()) != (
205 + self.weight.data_ptr(),
206 + self.weight_scale_inv.data_ptr(),
207 + ):
208 + raise RuntimeError("Shared Engram parameter storage must not be replaced")
209 + if self._views is None:
210 + self._views = (
211 + get_accelerator_view_from_cpu_tensor(self.weight),
212 + get_accelerator_view_from_cpu_tensor(self.weight_scale_inv),
213 + )
214 + return self._views
215 +
216 +
217 +class ParallelEngramEmbedding(BaseParallelEngramEmbedding):
218 + """Extend TP lookup with DP head sharding or shared, CPU-offloaded TP slices."""
219 +
220 + _shared_memory: DPSharedEngramStorage | None = None
221 +
222 + def __init__(
223 + self,
224 + num_embeddings: int,
225 + dim: int,
226 + head_sizes: tuple[int, ...],
227 + block_size: int = 32,
228 + cpu_offload: bool = False,
229 + dp_shared_memory: bool = False,
230 + ) -> None:
231 + self.cpu_offload = cpu_offload
232 + self.dp_shared_memory = dp_shared_memory
233 + self.dp_size = get_engram_dp_size()
234 + if dp_shared_memory:
235 + if not cpu_offload:
236 + raise ValueError("dp_shared_memory requires cpu_offload=True")
237 + if self.dp_size <= 1:
238 + raise ValueError(
239 + "dp_shared_memory requires a node-local Engram DP "
240 + f"group with size > 1; effective Engram DP size is {self.dp_size}. "
241 + "Check that the node layout and rank placement allow complete "
242 + "DP replicas to be co-located."
243 + )
244 + self.dp_size = 1
245 + if cpu_offload and not is_uva_available():
246 + raise RuntimeError("Engram CPU offload requires UVA support")
247 + self._views: tuple[torch.Tensor, torch.Tensor] | None = None
248 + self._view_src: tuple[int, int] | None = None
249 + super().__init__(num_embeddings, dim, head_sizes, block_size)
250 + if cpu_offload:
251 + # Constant dummy values avoid randomizing huge CPU lookup tables.
252 + set_weight_attrs(self.weight, {"dummy_weight_value": 1.0})
253 + # The ue8m0 encoding of scale 1.0 is exponent byte 127.
254 + set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127})
255 + logger.info(
256 + "Engram table offloaded to pinned host memory: %d rows x %d, "
257 + "%.2f GiB %s",
258 + self.part_num_embeddings,
259 + dim,
260 + self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
261 + "shared across DP replicas" if dp_shared_memory else "per rank",
262 + )
263 +
264 + def _get_shard_info(self) -> tuple[int, int]:
265 + if self.dp_size == 1:
266 + return super()._get_shard_info()
267 + return self.tp_size * self.dp_size, engram_head_shard_rank()
268 +
269 + def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
270 + if self.dp_shared_memory:
271 + group = get_engram_dp_group()
272 + assert group is not None
273 + storage = DPSharedEngramStorage(
274 + self.part_num_embeddings, self.dim, self.block_size, group
275 + )
276 + self._shared_memory = storage
277 + self._weight_loader = storage.load_weight
278 + return storage.weight, storage.weight_scale_inv
279 + if not self.cpu_offload:
280 + return super()._allocate_weights()
281 + # Model initialization may be inside a CUDA device context.
282 + return (
283 + torch.empty(
284 + self.part_num_embeddings,
285 + self.dim,
286 + dtype=torch.float8_e4m3fn,
287 + device="cpu",
288 + pin_memory=True,
289 + ),
290 + torch.empty(
291 + self.part_num_embeddings,
292 + self.dim // self.block_size,
293 + dtype=torch.uint8,
294 + device="cpu",
295 + pin_memory=True,
296 + ),
297 + )
298 +
299 + def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
300 + if self._shared_memory is not None:
301 + return self._shared_memory.get_views(self.weight, self.weight_scale_inv)
302 + if not self.cpu_offload:
303 + return super()._storage()
304 + src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr())
305 + if self._view_src != src:
306 + self._views = (
307 + get_accelerator_view_from_cpu_tensor(self.weight.data),
308 + get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data),
309 + )
310 + self._view_src = src
311 + assert self._views is not None
312 + return self._views
313 +
314 + def forward(self, indices: torch.Tensor) -> torch.Tensor:
315 + if self.dp_size == 1:
316 + return super().forward(indices)
317 + num_tokens = indices.shape[0]
318 + indices = gather_engram_hashes(indices)
319 + out = torch.empty(
320 + (indices.shape[0], self.part_n_hash_cols, self.dim),
321 + dtype=torch.bfloat16,
322 + device=indices.device,
323 + )
324 + self.lookup(indices, out)
325 + out = _gather_engram_rows(out, num_tokens)
326 + if self.tp_size > 1:
327 + out = tensor_model_parallel_all_gather(out, dim=1)
328 + return out[:, : self.n_hash_cols]
329 +
330 +
331 +def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor:
332 + """Exchange DP tokens for heads, retaining only this replica's tokens."""
333 + dp_group = get_engram_dp_group()
334 + assert dp_group is not None
335 + slot, remainder = divmod(staged.shape[0], dp_group.world_size)
336 + assert remainder == 0 and 0 <= num_tokens <= slot
337 + gathered = dp_group.all_gather(staged, dim=0)
338 + local_heads, dim = staged.shape[1:]
339 + rows = staged.new_empty((num_tokens, dp_group.world_size * local_heads, dim))
340 + _engram_select_rows(
341 + gathered,
342 + rows,
343 + staged.shape[0],
344 + dp_group.rank_in_group * slot,
345 + local_heads * dim,
346 + )
347 + return rows
348 +
349 +
350 +class Engram(BaseEngram):
351 + """NVIDIA Engram with asynchronous offload and node-local DP lookup."""
352 +
353 + _prefetch_stream: torch.cuda.Stream | None = None
354 +
355 + def _create_embedding(
356 + self, layout: EngramLayout, layer_hash_index: int
357 + ) -> ParallelEngramEmbedding:
358 + engram_config = get_current_vllm_config().engram_config
359 + assert engram_config is not None
360 + return ParallelEngramEmbedding(
361 + layout.num_embeddings[layer_hash_index],
362 + layout.head_dim,
363 + tuple(size for order in layout.primes[layer_hash_index] for size in order),
364 + cpu_offload=engram_config.cpu_offload,
365 + dp_shared_memory=engram_config.dp_shared_memory,
366 + )
367 +
368 + def _init_staging(self, max_tokens: int, head_dim: int) -> None:
369 + super()._init_staging(max_tokens * self.embed_tokens.dp_size, head_dim)
370 + if self.embed_tokens.cpu_offload:
371 + self._prefetch_stream = torch.cuda.Stream(device=self.staged_rows.device)
372 +
373 + def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
374 + """Prefetch local shared rows or the DP group's gathered hash IDs."""
375 + if self._prefetch_stream is None:
376 + return super().prepare_embeddings(hash_ids)
377 + rows = self.staged_rows[: hash_ids.shape[0]]
378 + assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small"
379 + self._start_prefetch(hash_ids, rows, self._prefetch_stream)
380 +
381 + @eager_break_during_capture
382 + def _start_prefetch(
383 + self, hash_ids: torch.Tensor, rows: torch.Tensor, stream: torch.cuda.Stream
384 + ) -> None:
385 + # Eager boundaries let the lookup span piecewise graph segments.
386 + stream.wait_stream(torch.cuda.current_stream())
387 + # Keep temporary hash storage alive until lookup finishes reading it.
388 + hash_ids.record_stream(stream)
389 + with torch.cuda.stream(stream):
390 + self.embed_tokens.lookup(hash_ids, rows, background=True)
391 +
392 + @eager_break_during_capture
393 + def _finish_prefetch(self, stream: torch.cuda.Stream) -> None:
394 + torch.cuda.current_stream().wait_stream(stream)
395 +
396 + def _ready_rows(self, num_tokens: int) -> torch.Tensor:
397 + if self._prefetch_stream is not None:
398 + self._finish_prefetch(self._prefetch_stream)
399 + if self.embed_tokens.dp_size > 1:
400 + slot = engram_gathered_num_tokens()
401 + staged = self.staged_rows[: slot * self.embed_tokens.dp_size]
402 + return _gather_engram_rows(staged, num_tokens)
403 + return super()._ready_rows(num_tokens)
vllm/models/deepseek_v4_1/nvidia/model.py
+16 -2
@@ -12,6 +12,7 @@ import vllm.envs as envs
12 from vllm.config import VllmConfig
13 from vllm.config.kernel import MEGA_MOE_BACKENDS
14 from vllm.distributed import (
15 + get_engram_dp_size,
16 get_pp_group,
17 get_tensor_model_parallel_rank,
18 get_tensor_model_parallel_world_size,
@@ -74,8 +75,9 @@ from vllm.utils.math_utils import cdiv
75 from vllm.v1.attention.backends.registry import AttentionBackendEnum
76 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
77
77 -from ..common.engram import Engram, EngramLayout, NgramHashState
78 +from ..common.engram import EngramLayout, NgramHashState
79 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
80 +from .engram import Engram, gather_engram_hashes
81
82 if typing.TYPE_CHECKING:
83 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -466,6 +468,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
468 # layer's sliding-window KV cache. Only PP ranks owning an engram
469 # layer need it.
470 self.engram_hash: NgramHashState | None = None
471 + self.engram_dp_shared_memory = bool(
472 + vllm_config.engram_config and vllm_config.engram_config.dp_shared_memory
473 + )
474 self.engram_swa_prefix: str | None = None
475 if self.engram_layout is not None:
476 local_engram = any(
@@ -594,12 +599,21 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
599 swa_metadata.slot_mapping,
600 swa_metadata.block_table,
601 )
602 + elif not self.engram_dp_shared_memory and get_engram_dp_size() > 1:
603 + # DP-sharded lookups are collective, so a replica skipping the
604 + # hash still has to reach them.
605 + engram_hashes, engram_mask = self.engram_hash.dummy_hashes(input_ids)
606 + if engram_hashes is not None:
607 # Gather all Engram rows before entering the decoder layers.
608 + # One gather feeds every layer sharing the DP-split table.
609 + gathered_hashes = gather_engram_hashes(
610 + engram_hashes, dp_shared_memory=self.engram_dp_shared_memory
611 + )
612 for layer in islice(self.layers, self.start_layer, self.end_layer):
613 engram = getattr(layer, "engram", None)
614 if engram is not None:
615 engram.prepare_embeddings(
602 - engram_hashes[:, engram.layer_hash_index]
616 + gathered_hashes[:, engram.layer_hash_index]
617 )
618
619 full_num_tokens = positions.shape[0]