diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index 45cebcfe57..4003c0516a 100644 --- a/tests/kernels/test_engram.py +++ b/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": diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index 3dea587075..512f518a05 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/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() diff --git a/vllm/config/engram.py b/vllm/config/engram.py index 0f5ec31b33..2cd67c0973 100644 --- a/vllm/config/engram.py +++ b/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: diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 9414abb991..12d557fa1c 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/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: diff --git a/vllm/models/deepseek_v41/nvidia/engram.py b/vllm/models/deepseek_v41/nvidia/engram.py index ddb57d3bd2..7d1e0f8a4b 100644 --- a/vllm/models/deepseek_v41/nvidia/engram.py +++ b/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] diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 22674bf705..8dff34daf9 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/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: