A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 17:00 UTC
[Performance] Add prefetch for checkpoints to OS page cache (#36012)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>
Artem Perevedentsev committed
Mar 16, 2026 at 13:32 UTC
f5e59ee7a6c3a07aad8f814b261bc0a1db2dcaf1
2 files changed
+76
-1
vllm/config/load.py
+3
@@ -62,6 +62,9 @@ class LoadConfig:
62
This is recommended for models on network filesystems (e.g., Lustre, NFS)
63
as it avoids inefficient random reads, significantly speeding up model
64
initialization. However, it uses more CPU RAM.
65
+ - "prefetch": Checkpoint files are read into the OS page cache before
66
+ workers load them, speeding up the model loading phase. Useful on
67
+ network or high-latency storage.
68
- "torchao": Weights are loaded in upfront and then reconstructed
69
into torchao tensor subclasses. This is used when the checkpoint
70
was quantized using torchao and saved using safetensors.
vllm/model_executor/model_loader/weight_utils.py
+73
-1
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Utilities for downloading and initializing model weights."""
4
5
+import asyncio
6
import concurrent.futures
7
import fnmatch
8
import glob
@@ -9,6 +10,7 @@ import hashlib
10
import json
11
import os
12
import tempfile
13
+import threading
14
import time
15
from collections import defaultdict
16
from collections.abc import Callable, Generator
@@ -720,6 +722,71 @@ def np_cache_weights_iterator(
722
yield name, torch.from_numpy(param)
723
724
725
+def _prefetch_checkpoint(file_path: str) -> None:
726
+ """Prefetch a checkpoint file into the OS page cache.
727
+
728
+ Reads the file in 16MB blocks so the kernel caches its pages before
729
+ workers load the same file.
730
+ """
731
+ block_size = 16 * 1024 * 1024 # 16MB
732
+ with open(file_path, "rb") as f:
733
+ while f.read(block_size):
734
+ pass
735
+
736
+
737
+def _prefetch_all_checkpoints(sorted_files: list[str]) -> None:
738
+ """Start prefetching checkpoint files into page cache in a background thread."""
739
+ if torch.distributed.is_initialized():
740
+ rank = torch.distributed.get_rank()
741
+ world_size = torch.distributed.get_world_size()
742
+ else:
743
+ rank = 0
744
+ world_size = 1
745
+ num_prefetch_threads = 8
746
+ paths_to_prefetch = sorted_files[rank::world_size]
747
+ total_for_rank = len(paths_to_prefetch)
748
+
749
+ async def _prefetch_all() -> None:
750
+ semaphore = asyncio.Semaphore(num_prefetch_threads)
751
+ completed = 0
752
+ next_log_pct = 10
753
+
754
+ async def prefetch_one(path: str) -> None:
755
+ nonlocal completed, next_log_pct
756
+ try:
757
+ async with semaphore:
758
+ await asyncio.to_thread(_prefetch_checkpoint, path)
759
+ completed += 1
760
+ if total_for_rank > 0 and next_log_pct <= 100:
761
+ pct = 100 * completed / total_for_rank
762
+ if pct >= next_log_pct:
763
+ logger.info(
764
+ "Prefetching checkpoint files: %d%% (%d/%d)",
765
+ next_log_pct,
766
+ completed,
767
+ total_for_rank,
768
+ )
769
+ next_log_pct += 10
770
+ except Exception:
771
+ logger.warning(
772
+ "Failed to prefetch checkpoint file %r.", path, exc_info=True
773
+ )
774
+
775
+ await asyncio.gather(*(prefetch_one(p) for p in paths_to_prefetch))
776
+
777
+ def _run_prefetch() -> None:
778
+ start = time.perf_counter()
779
+ asyncio.run(_prefetch_all())
780
+ elapsed = time.perf_counter() - start
781
+ logger.info(
782
+ "Prefetching checkpoint files into page cache finished in %.2fs",
783
+ elapsed,
784
+ )
785
+
786
+ logger.info("Prefetching checkpoint files into page cache started (in background)")
787
+ threading.Thread(target=_run_prefetch, daemon=True).start()
788
+
789
+
790
def safetensors_weights_iterator(
791
hf_weights_files: list[str],
792
use_tqdm_on_load: bool,
@@ -736,9 +803,14 @@ def safetensors_weights_iterator(
803
if safetensors_load_strategy == "eager":
804
loading_desc += " (eager)"
805
806
+ sorted_files = sorted(hf_weights_files, key=_natural_sort_key)
807
+
808
+ if safetensors_load_strategy == "prefetch":
809
+ _prefetch_all_checkpoints(sorted_files)
810
+
811
leftover_state_dict: dict[str, torch.Tensor] = {}
812
for st_file in tqdm(
741
- sorted(hf_weights_files, key=_natural_sort_key),
813
+ sorted_files,
814
desc=loading_desc,
815
disable=not enable_tqdm(use_tqdm_on_load),
816
bar_format=_BAR_FORMAT,