[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>

tests/kernels/test_engram.py

@@ -728,12 +728,79 @@ def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens):
     assert torch.equal(out, expected)
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("register_fails", [False, True])
+def test_engram_use_thp_lookup_and_fallback(monkeypatch, register_fails):
+    """Huge-page tables and their pinned fallback both give exact lookup rows."""
+    monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
+    if register_fails:
+        monkeypatch.setattr(
+            torch.cuda.cudart(), "cudaHostRegister", lambda *_: SimpleNamespace(value=1)
+        )
+    rows, dim = 32769, 64
+    with torch.device("cuda"):
+        layer = ParallelEngramEmbedding(
+            rows, dim, (rows,), cpu_offload=True, use_thp=True
+        )
+    assert layer.weight.is_pinned() and layer.weight_scale_inv.is_pinned()
+    same_storage = layer.weight.untyped_storage().data_ptr() == (
+        layer.weight_scale_inv.untyped_storage().data_ptr()
+    )
+    assert same_storage != register_fails
+    layer.weight.data.fill_(2)
+    layer.weight_scale_inv.data.fill_(127)
+    ids = torch.tensor(
+        [[0], [rows - 1], [-1], [rows]], device="cuda", dtype=torch.int32
+    )
+    out = torch.empty(4, 1, dim, device="cuda", dtype=torch.bfloat16)
+    layer.lookup(ids, out, background=True)
+    expected = torch.zeros_like(out)
+    expected[:2] = 2
+    torch.testing.assert_close(out, expected, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+def test_engram_shared_prefetch_stream_waits_per_layer():
+    """Consuming a layer must not wait for later layers on the shared stream."""
+    stream = torch.cuda.Stream()
+    modules = []
+    for _ in range(2):
+        module = Engram.__new__(Engram)
+        torch.nn.Module.__init__(module)
+        module._prefetch_stream = stream
+        module.embed_tokens = _make_embedding(cpu_offload=True)
+        module.use_sequence_parallel = False
+        with torch.device("cuda"):
+            module._init_staging(8, module.embed_tokens.dim)
+        assert module._prefetch_stream is stream
+        modules.append(module)
+    hashes = torch.randint(1024, 3072, (8, 2, 24), device="cuda", dtype=torch.int32)
+    modules[0].prepare_embeddings(hashes[:, 0])
+    with torch.cuda.stream(stream):
+        torch.cuda._sleep(100_000_000)  # ~50 ms ahead of the second lookup
+    modules[1].prepare_embeddings(hashes[:, 1])
+    rows = modules[0].embed(hashes[:, 0])
+    torch.cuda.current_stream().synchronize()
+    assert not modules[1]._prefetch_done.query(), "waited for the second lookup"
+    layer = modules[0].embed_tokens
+    expected = _reference_lookup(
+        layer.weight.cuda(),
+        layer.weight_scale_inv.cuda(),
+        hashes[:, 0],
+        layer.vocab_start_idx,
+        layer.vocab_end_idx,
+    )
+    torch.testing.assert_close(rows, expected, rtol=0, atol=0)
+    torch.accelerator.synchronize()
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
 @pytest.mark.parametrize(
     "backend,cpu_offload", [("nvidia", None), ("nvidia", False), ("common", False)]
 )
 def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload):
-    """Default offload uses pinned storage and a side stream; False uses HBM."""
+    """Default offload uses pinned storage and the given stream; False uses HBM."""
     from vllm.config import EngramConfig
 
     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):
     monkeypatch.setattr(
         engram_ops, "ReplicatedLinear", lambda *a, **k: torch.nn.Identity()
     )
+    stream = torch.cuda.Stream()
+    kwargs = {"prefetch_stream": stream} if backend == "nvidia" else {}
     with torch.device("cuda"):
         cls = Engram if backend == "nvidia" else CommonEngram
-        module = cls(config, None, layout, 0, False, "engram")
+        module = cls(config, None, layout, 0, False, "engram", **kwargs)
     layer = module.embed_tokens
     if offloaded:
         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):
     )
     torch.testing.assert_close(module.embed(ids), expected, atol=0, rtol=0)
     if offloaded:
-        assert streams == [module._prefetch_stream]
-        assert streams[0] != main
+        assert streams == [stream]
     else:
         assert streams == [main]
 
@@ -836,8 +904,9 @@ def test_engram_prefetch_detects_missing_dependency(monkeypatch, missing_depende
             hash_ids.record_stream(stream)
         with torch.cuda.stream(stream):
             self.embed_tokens.lookup(hash_ids, rows, background=True)
+            self._prefetch_done.record(stream)
 
-    def finish(self, stream):
+    def finish(self, event):
         pass
 
     monkeypatch.setattr(Engram, "_start_prefetch", start)
@@ -879,6 +948,7 @@ def _run_engram_prepared_rows(
     torch.nn.Module.__init__(engram)
     engram.embed_tokens = layer
     engram._prefetch_stream = torch.cuda.Stream() if cpu_offload else None
+    engram._prefetch_done = torch.cuda.Event() if cpu_offload else None
     # Exercise the production eager boundaries even when the test process
     # imported Engram before breakable graphs were enabled.
     if capture == "breakable":

tests/model_executor/test_weight_utils.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import tempfile
+from unittest.mock import Mock
 
 import huggingface_hub.constants
 import pytest
@@ -9,6 +10,7 @@ from huggingface_hub.utils import LocalEntryNotFoundError
 
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
+    drop_checkpoint_cache,
     maybe_remap_kv_scale_name,
 )
 
@@ -292,5 +294,18 @@ class TestKvCacheScaleMapper:
         )
 
 
+@pytest.mark.parametrize("fails", [False, True])
+def test_drop_checkpoint_cache_releases_local_weight_files(
+    tmp_path, monkeypatch, fails
+):
+    """Evict only weights; cache eviction errors must not prevent loading."""
+    for name in ("weights.safetensors", "config.json"):
+        (tmp_path / name).touch()
+    fadvise = Mock(side_effect=OSError("cache eviction failed") if fails else None)
+    monkeypatch.setattr("os.posix_fadvise", fadvise)
+    drop_checkpoint_cache(str(tmp_path))
+    fadvise.assert_called_once()
+
+
 if __name__ == "__main__":
     test_download_weights_from_hf()

vllm/config/engram.py

@@ -50,10 +50,19 @@ class EngramConfig:
     the other settings allow it, falling back to per-replica tables when DP
     replicas are not co-located on one node or /dev/shm cannot hold them."""
 
+    use_thp: bool = False
+    """Back private CPU-offloaded tables with transparent huge pages (best
+    effort, falls back to ordinary pinned pages). Prefaulting the tables at
+    startup takes longer. Requires cpu_offload without dp_shared_memory."""
+
     @model_validator(mode="after")
     def _validate_shared_memory(self) -> Self:
         if self.dp_shared_memory and not self.cpu_offload:
             raise ValueError("dp_shared_memory requires cpu_offload=True")
+        if self.use_thp and (not self.cpu_offload or self.dp_shared_memory):
+            raise ValueError(
+                "use_thp requires cpu_offload=True and dp_shared_memory=False"
+            )
         return self
 
     def verify_model_config(self, model_config: "ModelConfig | None") -> None:

vllm/model_executor/model_loader/weight_utils.py

@@ -13,7 +13,7 @@ import tempfile
 import threading
 import time
 from collections.abc import Callable, Generator, Iterable
-from contextlib import contextmanager
+from contextlib import contextmanager, suppress
 from pathlib import Path
 from typing import IO, TYPE_CHECKING, Any
 
@@ -844,6 +844,40 @@ def np_cache_weights_iterator(
         yield name, torch.from_numpy(param)
 
 
+def drop_checkpoint_cache(
+    model_name_or_path: str,
+    revision: str | None = None,
+    cache_dir: str | None = None,
+) -> None:
+    checkpoint_dir: Path | None = None
+    if os.path.isdir(model_name_or_path):
+        checkpoint_dir = Path(model_name_or_path)
+    else:
+        with suppress(OSError, ValueError):
+            checkpoint_dir = Path(
+                hf_api().snapshot_download(
+                    model_name_or_path,
+                    revision=revision,
+                    cache_dir=cache_dir,
+                    local_files_only=True,
+                )
+            )
+    if checkpoint_dir is None:
+        logger.info_once(
+            "No local checkpoint found for %s; skipping page-cache eviction",
+            model_name_or_path,
+        )
+        return
+    for path in checkpoint_dir.glob("*.safetensors"):
+        try:
+            with path.open("rb") as file:
+                os.posix_fadvise(file.fileno(), 0, 0, os.POSIX_FADV_DONTNEED)
+        except OSError as exc:
+            logger.warning(
+                "Could not release checkpoint page cache for %s: %s", path, exc
+            )
+
+
 def _get_checkpoints_size_bytes(files: list[str]) -> int:
     """Return the total size of the checkpoint files in bytes."""
     if not files:

vllm/models/deepseek_v41/nvidia/engram.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch."""
 
+import ctypes
 import mmap
 import os
 import shutil
@@ -47,6 +48,41 @@ from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
 logger = init_logger(__name__)
 
 
+def _allocate_huge_page_storage(num_bytes: int) -> torch.Tensor | None:
+    """Register prefaulted huge pages, or return None for pinned-memory fallback."""
+    try:
+        mapping = mmap.mmap(-1, num_bytes, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS)
+        mapping.madvise(mmap.MADV_HUGEPAGE)
+    except OSError as exc:
+        logger.warning(
+            "Engram huge-page allocation failed (%s); using pinned memory.", exc
+        )
+        return None
+
+    owner = np.frombuffer(mapping, dtype=np.uint8)
+    # Fault the pages in before CUDA pins them, so they can be huge.
+    owner[:: mmap.PAGESIZE] = 0
+    tensor = torch.from_numpy(owner)
+    pointer = tensor.data_ptr()
+    result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
+    if result.value != 0:
+        logger.warning(
+            "Engram cudaHostRegister failed (%s); using pinned memory.", result
+        )
+        return None
+    # Tensor views retain owner; its finalizer retains mapping until unregister.
+    finalizer = weakref.finalize(
+        owner, DPSharedEngramStorage._unregister, mapping, pointer
+    )
+    finalizer.atexit = False  # type: ignore[misc]
+    if not tensor.is_pinned():
+        logger.warning(
+            "CUDA did not recognize Engram registration; using pinned memory."
+        )
+        return None
+    return tensor
+
+
 def engram_head_shard_rank() -> int:
     """This rank's slot among the hash-head shards of one engram table.
 
@@ -241,9 +277,12 @@ class ParallelEngramEmbedding(BaseParallelEngramEmbedding):
         block_size: int = 32,
         cpu_offload: bool = False,
         dp_shared_memory: bool = False,
+        use_thp: bool = False,
     ) -> None:
         self.cpu_offload = cpu_offload
         self.dp_shared_memory = dp_shared_memory
+        self.use_thp = use_thp
+        self._packed: torch.Tensor | None = None
         self.dp_size = get_engram_dp_size()
         if dp_shared_memory:
             if not cpu_offload:
@@ -292,6 +331,17 @@ class ParallelEngramEmbedding(BaseParallelEngramEmbedding):
             return storage.weight, storage.weight_scale_inv
         if not self.cpu_offload:
             return super()._allocate_weights()
+        if self.use_thp:
+            weight_bytes = self.part_num_embeddings * self.dim
+            packed = _allocate_huge_page_storage(
+                weight_bytes + weight_bytes // self.block_size
+            )
+            if packed is not None:
+                self._packed = packed
+                return (
+                    packed[:weight_bytes].view(torch.float8_e4m3fn).view(-1, self.dim),
+                    packed[weight_bytes:].view(-1, self.dim // self.block_size),
+                )
         # Model initialization may be inside a CUDA device context.
         return (
             torch.empty(
@@ -310,6 +360,19 @@ class ParallelEngramEmbedding(BaseParallelEngramEmbedding):
             ),
         )
 
+    def collapse_huge_pages(self) -> None:
+        """Best-effort MADV_COLLAPSE (Linux >= 6.1) of pages that faulted small."""
+        if self._packed is None:
+            return
+        addr, num_bytes = self._packed.data_ptr(), self._packed.nbytes
+        libc = ctypes.CDLL(None, use_errno=True)
+        libc.madvise.argtypes = (ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int)
+        if libc.madvise(addr, num_bytes, 25) != 0:
+            logger.warning(
+                "Engram MADV_COLLAPSE failed; keeping existing pages: %s",
+                os.strerror(ctypes.get_errno()),
+            )
+
     def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
         if self._shared_memory is not None:
             return self._shared_memory.get_views(self.weight, self.weight_scale_inv)
@@ -365,6 +428,15 @@ class Engram(BaseEngram):
     """NVIDIA Engram with asynchronous offload and node-local DP lookup."""
 
     _prefetch_stream: torch.cuda.Stream | None = None
+    _prefetch_done: torch.cuda.Event | None = None
+
+    def __init__(
+        self, *args, prefetch_stream: torch.cuda.Stream | None = None, **kwargs
+    ) -> None:
+        # Layers sharing one stream serialize their offloaded lookups, so they
+        # take turns instead of jointly starving decoder compute of SMs.
+        self._prefetch_stream = prefetch_stream
+        super().__init__(*args, **kwargs)
 
     def _create_embedding(
         self, layout: EngramLayout, layer_hash_index: int
@@ -377,12 +449,17 @@ class Engram(BaseEngram):
             tuple(size for order in layout.primes[layer_hash_index] for size in order),
             cpu_offload=engram_config.cpu_offload,
             dp_shared_memory=bool(engram_config.dp_shared_memory),
+            use_thp=engram_config.use_thp,
         )
 
     def _init_staging(self, max_tokens: int, head_dim: int) -> None:
         super()._init_staging(max_tokens * self.embed_tokens.dp_size, head_dim)
-        if self.embed_tokens.cpu_offload:
+        if not self.embed_tokens.cpu_offload:
+            self._prefetch_stream = None
+            return
+        if self._prefetch_stream is None:
             self._prefetch_stream = torch.cuda.Stream(device=self.staged_rows.device)
+        self._prefetch_done = torch.cuda.Event()
 
     def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
         """Prefetch local shared rows or the DP group's gathered hash IDs."""
@@ -402,14 +479,17 @@ class Engram(BaseEngram):
         hash_ids.record_stream(stream)
         with torch.cuda.stream(stream):
             self.embed_tokens.lookup(hash_ids, rows, background=True)
+            assert self._prefetch_done is not None
+            self._prefetch_done.record(stream)
 
     @eager_break_during_capture
-    def _finish_prefetch(self, stream: torch.cuda.Stream) -> None:
-        torch.cuda.current_stream().wait_stream(stream)
+    def _finish_prefetch(self, event: torch.cuda.Event) -> None:
+        # Wait for this layer's lookup only, not for later layers on the stream.
+        torch.cuda.current_stream().wait_event(event)
 
     def _ready_rows(self, num_tokens: int) -> torch.Tensor:
-        if self._prefetch_stream is not None:
-            self._finish_prefetch(self._prefetch_stream)
+        if self._prefetch_done is not None:
+            self._finish_prefetch(self._prefetch_done)
         if self.embed_tokens.dp_size > 1:
             slot = engram_gathered_num_tokens()
             staged = self.staged_rows[: slot * self.embed_tokens.dp_size]

vllm/models/deepseek_v41/nvidia/model.py

@@ -40,7 +40,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    drop_checkpoint_cache,
+)
 from vllm.model_executor.models.interfaces import (
     EagleModelMixin,
     MixtureOfExperts,
@@ -226,6 +229,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         aux_stream_list: list[torch.cuda.Stream] | None = None,
         candidate_block_buffer: torch.Tensor | None = None,
         engram_layout: EngramLayout | None = None,
+        engram_prefetch_stream: torch.cuda.Stream | None = None,
         run_gemm_rs: bool = False,
         mhc_stream: torch.cuda.Stream | None = None,
         fuse_mhc_all_reduce: bool = False,
@@ -249,6 +253,7 @@ class DeepseekV4DecoderLayer(nn.Module):
                     engram_layout.layer_ids.index(layer_id),
                     use_sequence_parallel=self.use_sequence_parallel,
                     prefix=f"{prefix}.engram",
+                    prefetch_stream=engram_prefetch_stream,
                 )
 
         self.rms_norm_eps = config.rms_norm_eps
@@ -621,6 +626,14 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
 
         self.engram_layout = EngramLayout.from_config(config)
         engram_config = vllm_config.engram_config
+        self.vllm_config = vllm_config
+        engram_prefetch_stream = (
+            torch.cuda.Stream()
+            if self.engram_layout is not None
+            and engram_config is not None
+            and engram_config.cpu_offload
+            else None
+        )
         if (
             self.engram_layout is not None
             and engram_config is not None
@@ -628,6 +641,15 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
         ):
             engram_config.dp_shared_memory = can_share_engram_tables(self.engram_layout)
 
+        if self.engram_layout is not None and engram_config and engram_config.use_thp:
+            # Release old checkpoint cache before allocating the Engram host tables.
+            model_config = vllm_config.model_config
+            drop_checkpoint_cache(
+                model_config.model_weights or model_config.model,
+                revision=model_config.revision,
+                cache_dir=vllm_config.load_config.download_dir,
+            )
+
         # GEMM-RS uses NCCL symmetric-memory multicast, which requires all TP
         # ranks to belong to one NVLink domain. Collective: run before layers.
         self.run_gemm_rs = maybe_init_gemm_rs(vllm_config, self.use_sequence_parallel)
@@ -641,6 +663,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
                 aux_stream_list=aux_stream_list,
                 candidate_block_buffer=self.candidate_block_buffer,
                 engram_layout=self.engram_layout,
+                engram_prefetch_stream=engram_prefetch_stream,
                 run_gemm_rs=self.run_gemm_rs,
                 mhc_stream=mhc_stream,
                 fuse_mhc_all_reduce=self.fuse_mhc_all_reduce,
@@ -1329,6 +1352,20 @@ class DeepseekV41LLMForCausalLM(
         loader = AutoWeightsLoader(self)
         loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
         self.process_weights_after_loading()
+        config = self.model.vllm_config
+        if config.engram_config and config.engram_config.use_thp:
+            # Loading weights refills the file cache; release it before MADV_COLLAPSE.
+            model_config = config.model_config
+            drop_checkpoint_cache(
+                model_config.model_weights or model_config.model,
+                revision=model_config.revision,
+                cache_dir=config.load_config.download_dir,
+            )
+            for layer in islice(
+                self.model.layers, self.model.start_layer, self.model.end_layer
+            ):
+                if layer.engram is not None:
+                    layer.engram.embed_tokens.collapse_huge_pages()
         return loaded_params
 
     def process_weights_after_loading(self) -> None: