[Perf][Engram] Serialize offloaded lookups and pack host tables into huge pages (#56926)
Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Kimi Code <noreply@moonshot.cn> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Juntian Liu committed
Sep 24, 2026 at 19:52 UTC
6936e77ba451725e6482eee229839e10ab2bf48c
6 files changed
+257
-12
tests/kernels/test_engram.py
+75
-5
@@ -728,12 +728,79 @@ def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens):
728
assert torch.equal(out, expected)
729
730
731
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
732
+@pytest.mark.parametrize("register_fails", [False, True])
733
+def test_engram_use_thp_lookup_and_fallback(monkeypatch, register_fails):
734
+ """Huge-page tables and their pinned fallback both give exact lookup rows."""
735
+ monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1)
736
+ monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
737
+ if register_fails:
738
+ monkeypatch.setattr(
739
+ torch.cuda.cudart(), "cudaHostRegister", lambda *_: SimpleNamespace(value=1)
740
+ )
741
+ rows, dim = 32769, 64
742
+ with torch.device("cuda"):
743
+ layer = ParallelEngramEmbedding(
744
+ rows, dim, (rows,), cpu_offload=True, use_thp=True
745
+ )
746
+ assert layer.weight.is_pinned() and layer.weight_scale_inv.is_pinned()
747
+ same_storage = layer.weight.untyped_storage().data_ptr() == (
748
+ layer.weight_scale_inv.untyped_storage().data_ptr()
749
+ )
750
+ assert same_storage != register_fails
751
+ layer.weight.data.fill_(2)
752
+ layer.weight_scale_inv.data.fill_(127)
753
+ ids = torch.tensor(
754
+ [[0], [rows - 1], [-1], [rows]], device="cuda", dtype=torch.int32
755
+ )
756
+ out = torch.empty(4, 1, dim, device="cuda", dtype=torch.bfloat16)
757
+ layer.lookup(ids, out, background=True)
758
+ expected = torch.zeros_like(out)
759
+ expected[:2] = 2
760
+ torch.testing.assert_close(out, expected, rtol=0, atol=0)
761
+
762
+
763
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
764
+def test_engram_shared_prefetch_stream_waits_per_layer():
765
+ """Consuming a layer must not wait for later layers on the shared stream."""
766
+ stream = torch.cuda.Stream()
767
+ modules = []
768
+ for _ in range(2):
769
+ module = Engram.__new__(Engram)
770
+ torch.nn.Module.__init__(module)
771
+ module._prefetch_stream = stream
772
+ module.embed_tokens = _make_embedding(cpu_offload=True)
773
+ module.use_sequence_parallel = False
774
+ with torch.device("cuda"):
775
+ module._init_staging(8, module.embed_tokens.dim)
776
+ assert module._prefetch_stream is stream
777
+ modules.append(module)
778
+ hashes = torch.randint(1024, 3072, (8, 2, 24), device="cuda", dtype=torch.int32)
779
+ modules[0].prepare_embeddings(hashes[:, 0])
780
+ with torch.cuda.stream(stream):
781
+ torch.cuda._sleep(100_000_000) # ~50 ms ahead of the second lookup
782
+ modules[1].prepare_embeddings(hashes[:, 1])
783
+ rows = modules[0].embed(hashes[:, 0])
784
+ torch.cuda.current_stream().synchronize()
785
+ assert not modules[1]._prefetch_done.query(), "waited for the second lookup"
786
+ layer = modules[0].embed_tokens
787
+ expected = _reference_lookup(
788
+ layer.weight.cuda(),
789
+ layer.weight_scale_inv.cuda(),
790
+ hashes[:, 0],
791
+ layer.vocab_start_idx,
792
+ layer.vocab_end_idx,
793
+ )
794
+ torch.testing.assert_close(rows, expected, rtol=0, atol=0)
795
+ torch.accelerator.synchronize()
796
+
797
+
798
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
799
@pytest.mark.parametrize(
800
"backend,cpu_offload", [("nvidia", None), ("nvidia", False), ("common", False)]
801
)
802
def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload):
736
- """Default offload uses pinned storage and a side stream; False uses HBM."""
803
+ """Default offload uses pinned storage and the given stream; False uses HBM."""
804
from vllm.config import EngramConfig
805
806
config = SimpleNamespace(hidden_size=16, hc_mult=1, rms_norm_eps=1e-6)
@@ -762,9 +829,11 @@ def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload):
829
monkeypatch.setattr(
830
engram_ops, "ReplicatedLinear", lambda *a, **k: torch.nn.Identity()
831
)
832
+ stream = torch.cuda.Stream()
833
+ kwargs = {"prefetch_stream": stream} if backend == "nvidia" else {}
834
with torch.device("cuda"):
835
cls = Engram if backend == "nvidia" else CommonEngram
767
- module = cls(config, None, layout, 0, False, "engram")
836
+ module = cls(config, None, layout, 0, False, "engram", **kwargs)
837
layer = module.embed_tokens
838
if offloaded:
839
assert layer.weight.device.type == "cpu" and layer.weight.is_pinned()
@@ -798,8 +867,7 @@ def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload):
867
)
868
torch.testing.assert_close(module.embed(ids), expected, atol=0, rtol=0)
869
if offloaded:
801
- assert streams == [module._prefetch_stream]
802
- assert streams[0] != main
870
+ assert streams == [stream]
871
else:
872
assert streams == [main]
873
@@ -836,8 +904,9 @@ def test_engram_prefetch_detects_missing_dependency(monkeypatch, missing_depende
904
hash_ids.record_stream(stream)
905
with torch.cuda.stream(stream):
906
self.embed_tokens.lookup(hash_ids, rows, background=True)
907
+ self._prefetch_done.record(stream)
908
840
- def finish(self, stream):
909
+ def finish(self, event):
910
pass
911
912
monkeypatch.setattr(Engram, "_start_prefetch", start)
@@ -879,6 +948,7 @@ def _run_engram_prepared_rows(
948
torch.nn.Module.__init__(engram)
949
engram.embed_tokens = layer
950
engram._prefetch_stream = torch.cuda.Stream() if cpu_offload else None
951
+ engram._prefetch_done = torch.cuda.Event() if cpu_offload else None
952
# Exercise the production eager boundaries even when the test process
953
# imported Engram before breakable graphs were enabled.
954
if capture == "breakable":
tests/model_executor/test_weight_utils.py
+15
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import tempfile
5
+from unittest.mock import Mock
6
7
import huggingface_hub.constants
8
import pytest
@@ -9,6 +10,7 @@ from huggingface_hub.utils import LocalEntryNotFoundError
10
11
from vllm.model_executor.model_loader.weight_utils import (
12
download_weights_from_hf,
13
+ drop_checkpoint_cache,
14
maybe_remap_kv_scale_name,
15
)
16
@@ -292,5 +294,18 @@ class TestKvCacheScaleMapper:
294
)
295
296
297
+@pytest.mark.parametrize("fails", [False, True])
298
+def test_drop_checkpoint_cache_releases_local_weight_files(
299
+ tmp_path, monkeypatch, fails
300
+):
301
+ """Evict only weights; cache eviction errors must not prevent loading."""
302
+ for name in ("weights.safetensors", "config.json"):
303
+ (tmp_path / name).touch()
304
+ fadvise = Mock(side_effect=OSError("cache eviction failed") if fails else None)
305
+ monkeypatch.setattr("os.posix_fadvise", fadvise)
306
+ drop_checkpoint_cache(str(tmp_path))
307
+ fadvise.assert_called_once()
308
+
309
+
310
if __name__ == "__main__":
311
test_download_weights_from_hf()
vllm/config/engram.py
+9
@@ -50,10 +50,19 @@ class EngramConfig:
50
the other settings allow it, falling back to per-replica tables when DP
51
replicas are not co-located on one node or /dev/shm cannot hold them."""
52
53
+ use_thp: bool = False
54
+ """Back private CPU-offloaded tables with transparent huge pages (best
55
+ effort, falls back to ordinary pinned pages). Prefaulting the tables at
56
+ startup takes longer. Requires cpu_offload without dp_shared_memory."""
57
+
58
@model_validator(mode="after")
59
def _validate_shared_memory(self) -> Self:
60
if self.dp_shared_memory and not self.cpu_offload:
61
raise ValueError("dp_shared_memory requires cpu_offload=True")
62
+ if self.use_thp and (not self.cpu_offload or self.dp_shared_memory):
63
+ raise ValueError(
64
+ "use_thp requires cpu_offload=True and dp_shared_memory=False"
65
+ )
66
return self
67
68
def verify_model_config(self, model_config: "ModelConfig | None") -> None:
vllm/model_executor/model_loader/weight_utils.py
+35
-1
@@ -13,7 +13,7 @@ import tempfile
13
import threading
14
import time
15
from collections.abc import Callable, Generator, Iterable
16
-from contextlib import contextmanager
16
+from contextlib import contextmanager, suppress
17
from pathlib import Path
18
from typing import IO, TYPE_CHECKING, Any
19
@@ -844,6 +844,40 @@ def np_cache_weights_iterator(
844
yield name, torch.from_numpy(param)
845
846
847
+def drop_checkpoint_cache(
848
+ model_name_or_path: str,
849
+ revision: str | None = None,
850
+ cache_dir: str | None = None,
851
+) -> None:
852
+ checkpoint_dir: Path | None = None
853
+ if os.path.isdir(model_name_or_path):
854
+ checkpoint_dir = Path(model_name_or_path)
855
+ else:
856
+ with suppress(OSError, ValueError):
857
+ checkpoint_dir = Path(
858
+ hf_api().snapshot_download(
859
+ model_name_or_path,
860
+ revision=revision,
861
+ cache_dir=cache_dir,
862
+ local_files_only=True,
863
+ )
864
+ )
865
+ if checkpoint_dir is None:
866
+ logger.info_once(
867
+ "No local checkpoint found for %s; skipping page-cache eviction",
868
+ model_name_or_path,
869
+ )
870
+ return
871
+ for path in checkpoint_dir.glob("*.safetensors"):
872
+ try:
873
+ with path.open("rb") as file:
874
+ os.posix_fadvise(file.fileno(), 0, 0, os.POSIX_FADV_DONTNEED)
875
+ except OSError as exc:
876
+ logger.warning(
877
+ "Could not release checkpoint page cache for %s: %s", path, exc
878
+ )
879
+
880
+
881
def _get_checkpoints_size_bytes(files: list[str]) -> int:
882
"""Return the total size of the checkpoint files in bytes."""
883
if not files:
vllm/models/deepseek_v41/nvidia/engram.py
+85
-5
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch."""
4
5
+import ctypes
6
import mmap
7
import os
8
import shutil
@@ -47,6 +48,41 @@ from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
48
logger = init_logger(__name__)
49
50
51
+def _allocate_huge_page_storage(num_bytes: int) -> torch.Tensor | None:
52
+ """Register prefaulted huge pages, or return None for pinned-memory fallback."""
53
+ try:
54
+ mapping = mmap.mmap(-1, num_bytes, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS)
55
+ mapping.madvise(mmap.MADV_HUGEPAGE)
56
+ except OSError as exc:
57
+ logger.warning(
58
+ "Engram huge-page allocation failed (%s); using pinned memory.", exc
59
+ )
60
+ return None
61
+
62
+ owner = np.frombuffer(mapping, dtype=np.uint8)
63
+ # Fault the pages in before CUDA pins them, so they can be huge.
64
+ owner[:: mmap.PAGESIZE] = 0
65
+ tensor = torch.from_numpy(owner)
66
+ pointer = tensor.data_ptr()
67
+ result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
68
+ if result.value != 0:
69
+ logger.warning(
70
+ "Engram cudaHostRegister failed (%s); using pinned memory.", result
71
+ )
72
+ return None
73
+ # Tensor views retain owner; its finalizer retains mapping until unregister.
74
+ finalizer = weakref.finalize(
75
+ owner, DPSharedEngramStorage._unregister, mapping, pointer
76
+ )
77
+ finalizer.atexit = False # type: ignore[misc]
78
+ if not tensor.is_pinned():
79
+ logger.warning(
80
+ "CUDA did not recognize Engram registration; using pinned memory."
81
+ )
82
+ return None
83
+ return tensor
84
+
85
+
86
def engram_head_shard_rank() -> int:
87
"""This rank's slot among the hash-head shards of one engram table.
88
@@ -241,9 +277,12 @@ class ParallelEngramEmbedding(BaseParallelEngramEmbedding):
277
block_size: int = 32,
278
cpu_offload: bool = False,
279
dp_shared_memory: bool = False,
280
+ use_thp: bool = False,
281
) -> None:
282
self.cpu_offload = cpu_offload
283
self.dp_shared_memory = dp_shared_memory
284
+ self.use_thp = use_thp
285
+ self._packed: torch.Tensor | None = None
286
self.dp_size = get_engram_dp_size()
287
if dp_shared_memory:
288
if not cpu_offload:
@@ -292,6 +331,17 @@ class ParallelEngramEmbedding(BaseParallelEngramEmbedding):
331
return storage.weight, storage.weight_scale_inv
332
if not self.cpu_offload:
333
return super()._allocate_weights()
334
+ if self.use_thp:
335
+ weight_bytes = self.part_num_embeddings * self.dim
336
+ packed = _allocate_huge_page_storage(
337
+ weight_bytes + weight_bytes // self.block_size
338
+ )
339
+ if packed is not None:
340
+ self._packed = packed
341
+ return (
342
+ packed[:weight_bytes].view(torch.float8_e4m3fn).view(-1, self.dim),
343
+ packed[weight_bytes:].view(-1, self.dim // self.block_size),
344
+ )
345
# Model initialization may be inside a CUDA device context.
346
return (
347
torch.empty(
@@ -310,6 +360,19 @@ class ParallelEngramEmbedding(BaseParallelEngramEmbedding):
360
),
361
)
362
363
+ def collapse_huge_pages(self) -> None:
364
+ """Best-effort MADV_COLLAPSE (Linux >= 6.1) of pages that faulted small."""
365
+ if self._packed is None:
366
+ return
367
+ addr, num_bytes = self._packed.data_ptr(), self._packed.nbytes
368
+ libc = ctypes.CDLL(None, use_errno=True)
369
+ libc.madvise.argtypes = (ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int)
370
+ if libc.madvise(addr, num_bytes, 25) != 0:
371
+ logger.warning(
372
+ "Engram MADV_COLLAPSE failed; keeping existing pages: %s",
373
+ os.strerror(ctypes.get_errno()),
374
+ )
375
+
376
def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
377
if self._shared_memory is not None:
378
return self._shared_memory.get_views(self.weight, self.weight_scale_inv)
@@ -365,6 +428,15 @@ class Engram(BaseEngram):
428
"""NVIDIA Engram with asynchronous offload and node-local DP lookup."""
429
430
_prefetch_stream: torch.cuda.Stream | None = None
431
+ _prefetch_done: torch.cuda.Event | None = None
432
+
433
+ def __init__(
434
+ self, *args, prefetch_stream: torch.cuda.Stream | None = None, **kwargs
435
+ ) -> None:
436
+ # Layers sharing one stream serialize their offloaded lookups, so they
437
+ # take turns instead of jointly starving decoder compute of SMs.
438
+ self._prefetch_stream = prefetch_stream
439
+ super().__init__(*args, **kwargs)
440
441
def _create_embedding(
442
self, layout: EngramLayout, layer_hash_index: int
@@ -377,12 +449,17 @@ class Engram(BaseEngram):
449
tuple(size for order in layout.primes[layer_hash_index] for size in order),
450
cpu_offload=engram_config.cpu_offload,
451
dp_shared_memory=bool(engram_config.dp_shared_memory),
452
+ use_thp=engram_config.use_thp,
453
)
454
455
def _init_staging(self, max_tokens: int, head_dim: int) -> None:
456
super()._init_staging(max_tokens * self.embed_tokens.dp_size, head_dim)
384
- if self.embed_tokens.cpu_offload:
457
+ if not self.embed_tokens.cpu_offload:
458
+ self._prefetch_stream = None
459
+ return
460
+ if self._prefetch_stream is None:
461
self._prefetch_stream = torch.cuda.Stream(device=self.staged_rows.device)
462
+ self._prefetch_done = torch.cuda.Event()
463
464
def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
465
"""Prefetch local shared rows or the DP group's gathered hash IDs."""
@@ -402,14 +479,17 @@ class Engram(BaseEngram):
479
hash_ids.record_stream(stream)
480
with torch.cuda.stream(stream):
481
self.embed_tokens.lookup(hash_ids, rows, background=True)
482
+ assert self._prefetch_done is not None
483
+ self._prefetch_done.record(stream)
484
485
@eager_break_during_capture
407
- def _finish_prefetch(self, stream: torch.cuda.Stream) -> None:
408
- torch.cuda.current_stream().wait_stream(stream)
486
+ def _finish_prefetch(self, event: torch.cuda.Event) -> None:
487
+ # Wait for this layer's lookup only, not for later layers on the stream.
488
+ torch.cuda.current_stream().wait_event(event)
489
490
def _ready_rows(self, num_tokens: int) -> torch.Tensor:
411
- if self._prefetch_stream is not None:
412
- self._finish_prefetch(self._prefetch_stream)
491
+ if self._prefetch_done is not None:
492
+ self._finish_prefetch(self._prefetch_done)
493
if self.embed_tokens.dp_size > 1:
494
slot = engram_gathered_num_tokens()
495
staged = self.staged_rows[: slot * self.embed_tokens.dp_size]
vllm/models/deepseek_v41/nvidia/model.py
+38
-1
@@ -40,7 +40,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
40
ParallelLMHead,
41
VocabParallelEmbedding,
42
)
43
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43
+from vllm.model_executor.model_loader.weight_utils import (
44
+ default_weight_loader,
45
+ drop_checkpoint_cache,
46
+)
47
from vllm.model_executor.models.interfaces import (
48
EagleModelMixin,
49
MixtureOfExperts,
@@ -226,6 +229,7 @@ class DeepseekV4DecoderLayer(nn.Module):
229
aux_stream_list: list[torch.cuda.Stream] | None = None,
230
candidate_block_buffer: torch.Tensor | None = None,
231
engram_layout: EngramLayout | None = None,
232
+ engram_prefetch_stream: torch.cuda.Stream | None = None,
233
run_gemm_rs: bool = False,
234
mhc_stream: torch.cuda.Stream | None = None,
235
fuse_mhc_all_reduce: bool = False,
@@ -249,6 +253,7 @@ class DeepseekV4DecoderLayer(nn.Module):
253
engram_layout.layer_ids.index(layer_id),
254
use_sequence_parallel=self.use_sequence_parallel,
255
prefix=f"{prefix}.engram",
256
+ prefetch_stream=engram_prefetch_stream,
257
)
258
259
self.rms_norm_eps = config.rms_norm_eps
@@ -621,6 +626,14 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
626
627
self.engram_layout = EngramLayout.from_config(config)
628
engram_config = vllm_config.engram_config
629
+ self.vllm_config = vllm_config
630
+ engram_prefetch_stream = (
631
+ torch.cuda.Stream()
632
+ if self.engram_layout is not None
633
+ and engram_config is not None
634
+ and engram_config.cpu_offload
635
+ else None
636
+ )
637
if (
638
self.engram_layout is not None
639
and engram_config is not None
@@ -628,6 +641,15 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
641
):
642
engram_config.dp_shared_memory = can_share_engram_tables(self.engram_layout)
643
644
+ if self.engram_layout is not None and engram_config and engram_config.use_thp:
645
+ # Release old checkpoint cache before allocating the Engram host tables.
646
+ model_config = vllm_config.model_config
647
+ drop_checkpoint_cache(
648
+ model_config.model_weights or model_config.model,
649
+ revision=model_config.revision,
650
+ cache_dir=vllm_config.load_config.download_dir,
651
+ )
652
+
653
# GEMM-RS uses NCCL symmetric-memory multicast, which requires all TP
654
# ranks to belong to one NVLink domain. Collective: run before layers.
655
self.run_gemm_rs = maybe_init_gemm_rs(vllm_config, self.use_sequence_parallel)
@@ -641,6 +663,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
663
aux_stream_list=aux_stream_list,
664
candidate_block_buffer=self.candidate_block_buffer,
665
engram_layout=self.engram_layout,
666
+ engram_prefetch_stream=engram_prefetch_stream,
667
run_gemm_rs=self.run_gemm_rs,
668
mhc_stream=mhc_stream,
669
fuse_mhc_all_reduce=self.fuse_mhc_all_reduce,
@@ -1329,6 +1352,20 @@ class DeepseekV41LLMForCausalLM(
1352
loader = AutoWeightsLoader(self)
1353
loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1354
self.process_weights_after_loading()
1355
+ config = self.model.vllm_config
1356
+ if config.engram_config and config.engram_config.use_thp:
1357
+ # Loading weights refills the file cache; release it before MADV_COLLAPSE.
1358
+ model_config = config.model_config
1359
+ drop_checkpoint_cache(
1360
+ model_config.model_weights or model_config.model,
1361
+ revision=model_config.revision,
1362
+ cache_dir=config.load_config.download_dir,
1363
+ )
1364
+ for layer in islice(
1365
+ self.model.layers, self.model.start_layer, self.model.end_layer
1366
+ ):
1367
+ if layer.engram is not None:
1368
+ layer.engram.embed_tokens.collapse_huge_pages()
1369
return loaded_params
1370
1371
def process_weights_after_loading(self) -> None: