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 3, 22:45 UTC
[Performance] Make safetensors checkpoint prefetch settings configurable (#41499)
Signed-off-by: Mohammad Miadh Angkad <MAngkad.BSDSBA2027@aim.edu>
Mohammad Miadh Angkad committed
May 10, 2026 at 23:55 UTC
21943d4c258983c4b8eb56d50029aca4f18e4629
4 files changed
+87
-15
vllm/config/load.py
+12
@@ -9,6 +9,9 @@ from vllm.config.utils import config
9
from vllm.logger import init_logger
10
from vllm.utils.hashing import safe_hash
11
12
+DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS = 8
13
+DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE = 16 * 1024 * 1024
14
+
15
if TYPE_CHECKING:
16
from vllm.model_executor.model_loader import LoadFormats
17
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
@@ -79,6 +82,15 @@ class LoadConfig:
82
was quantized using torchao and saved using safetensors.
83
Needs `torchao >= 0.14.0`.
84
"""
85
+ safetensors_prefetch_num_threads: int = Field(
86
+ default=DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, ge=1
87
+ )
88
+ """Number of worker threads used to prefetch safetensors checkpoint files
89
+ into the OS page cache when safetensors prefetching is enabled."""
90
+ safetensors_prefetch_block_size: int = Field(
91
+ default=DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE, ge=1
92
+ )
93
+ """Read size in bytes for each safetensors checkpoint file prefetch."""
94
model_loader_extra_config: dict | TensorizerConfig = Field(default_factory=dict)
95
"""Extra config for model loader. This will be passed to the model loader
96
corresponding to the chosen load_format."""
vllm/engine/arg_utils.py
+17
-1
@@ -355,7 +355,11 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]:
355
if name == "max_model_len":
356
kwargs[name]["type"] = human_readable_int_or_auto
357
kwargs[name]["help"] += f"\n\n{human_readable_int_or_auto.__doc__}"
358
- elif name in ("max_num_batched_tokens", "kv_cache_memory_bytes"):
358
+ elif name in (
359
+ "max_num_batched_tokens",
360
+ "kv_cache_memory_bytes",
361
+ "safetensors_prefetch_block_size",
362
+ ):
363
kwargs[name]["type"] = human_readable_int
364
kwargs[name]["help"] += f"\n\n{human_readable_int.__doc__}"
365
else:
@@ -424,6 +428,8 @@ class EngineArgs:
428
allowed_media_domains: list[str] | None = ModelConfig.allowed_media_domains
429
download_dir: str | None = LoadConfig.download_dir
430
safetensors_load_strategy: str | None = LoadConfig.safetensors_load_strategy
431
+ safetensors_prefetch_num_threads: int = LoadConfig.safetensors_prefetch_num_threads
432
+ safetensors_prefetch_block_size: int = LoadConfig.safetensors_prefetch_block_size
433
load_format: str | LoadFormats = LoadConfig.load_format
434
config_format: str = ModelConfig.config_format
435
dtype: ModelDType = ModelConfig.dtype
@@ -844,6 +850,14 @@ class EngineArgs:
850
load_group.add_argument(
851
"--safetensors-load-strategy", **load_kwargs["safetensors_load_strategy"]
852
)
853
+ load_group.add_argument(
854
+ "--safetensors-prefetch-num-threads",
855
+ **load_kwargs["safetensors_prefetch_num_threads"],
856
+ )
857
+ load_group.add_argument(
858
+ "--safetensors-prefetch-block-size",
859
+ **load_kwargs["safetensors_prefetch_block_size"],
860
+ )
861
load_group.add_argument(
862
"--model-loader-extra-config", **load_kwargs["model_loader_extra_config"]
863
)
@@ -1584,6 +1598,8 @@ class EngineArgs:
1598
load_format=self.load_format,
1599
download_dir=self.download_dir,
1600
safetensors_load_strategy=self.safetensors_load_strategy,
1601
+ safetensors_prefetch_num_threads=self.safetensors_prefetch_num_threads,
1602
+ safetensors_prefetch_block_size=self.safetensors_prefetch_block_size,
1603
model_loader_extra_config=self.model_loader_extra_config,
1604
ignore_patterns=self.ignore_patterns,
1605
use_tqdm_on_load=self.use_tqdm_on_load,
vllm/model_executor/model_loader/default_loader.py
+6
@@ -256,6 +256,12 @@ class DefaultModelLoader(BaseModelLoader):
256
self.load_config.use_tqdm_on_load,
257
self.load_config.safetensors_load_strategy,
258
local_expert_ids=self.local_expert_ids,
259
+ safetensors_prefetch_num_threads=(
260
+ self.load_config.safetensors_prefetch_num_threads
261
+ ),
262
+ safetensors_prefetch_block_size=(
263
+ self.load_config.safetensors_prefetch_block_size
264
+ ),
265
)
266
else:
267
if extra_config.get("enable_multithread_load"):
vllm/model_executor/model_loader/weight_utils.py
+52
-14
@@ -30,7 +30,11 @@ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
30
31
from vllm import envs
32
from vllm.config import ModelConfig
33
-from vllm.config.load import LoadConfig
33
+from vllm.config.load import (
34
+ DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
35
+ DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS,
36
+ LoadConfig,
37
+)
38
from vllm.distributed import get_tensor_model_parallel_rank, get_world_group
39
from vllm.logger import init_logger
40
from vllm.model_executor.layers.quantization import (
@@ -810,40 +814,57 @@ def _get_fs_type(files: list[str]) -> str:
814
return ""
815
816
813
-def _prefetch_checkpoint(file_path: str) -> None:
817
+def _prefetch_checkpoint(
818
+ file_path: str,
819
+ block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
820
+) -> None:
821
"""Prefetch a checkpoint file into the OS page cache.
822
816
- Reads the file in 16MB blocks so the kernel caches its pages before
817
- workers load the same file.
823
+ Reads the file in blocks so the kernel caches its pages before workers load
824
+ the same file.
825
"""
819
- block_size = 16 * 1024 * 1024 # 16MB
826
+ if block_size < 1:
827
+ raise ValueError("safetensors prefetch block size must be >= 1")
828
+
829
with open(file_path, "rb") as f:
830
while f.read(block_size):
831
pass
832
833
825
-def _prefetch_all_checkpoints(sorted_files: list[str]) -> None:
834
+def _prefetch_all_checkpoints(
835
+ sorted_files: list[str],
836
+ num_prefetch_threads: int = DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS,
837
+ block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
838
+) -> None:
839
"""Start prefetching checkpoint files into page cache in a background thread."""
840
+ if num_prefetch_threads < 1:
841
+ raise ValueError("safetensors prefetch num threads must be >= 1")
842
+ if block_size < 1:
843
+ raise ValueError("safetensors prefetch block size must be >= 1")
844
+
845
if torch.distributed.is_initialized():
846
rank = torch.distributed.get_rank()
847
world_size = torch.distributed.get_world_size()
848
else:
849
rank = 0
850
world_size = 1
833
- num_prefetch_threads = 8
851
paths_to_prefetch = sorted_files[rank::world_size]
852
total_for_rank = len(paths_to_prefetch)
853
854
async def _prefetch_all() -> None:
838
- semaphore = asyncio.Semaphore(num_prefetch_threads)
855
+ loop = asyncio.get_running_loop()
856
completed = 0
857
next_log_pct = 10
858
842
- async def prefetch_one(path: str) -> None:
859
+ async def prefetch_one(
860
+ path: str,
861
+ executor: concurrent.futures.ThreadPoolExecutor,
862
+ ) -> None:
863
nonlocal completed, next_log_pct
864
try:
845
- async with semaphore:
846
- await asyncio.to_thread(_prefetch_checkpoint, path)
865
+ await loop.run_in_executor(
866
+ executor, _prefetch_checkpoint, path, block_size
867
+ )
868
completed += 1
869
if total_for_rank > 0 and next_log_pct <= 100:
870
pct = 100 * completed / total_for_rank
@@ -860,7 +881,12 @@ def _prefetch_all_checkpoints(sorted_files: list[str]) -> None:
881
"Failed to prefetch checkpoint file %r.", path, exc_info=True
882
)
883
863
- await asyncio.gather(*(prefetch_one(p) for p in paths_to_prefetch))
884
+ with concurrent.futures.ThreadPoolExecutor(
885
+ max_workers=num_prefetch_threads
886
+ ) as executor:
887
+ await asyncio.gather(
888
+ *(prefetch_one(p, executor) for p in paths_to_prefetch)
889
+ )
890
891
def _run_prefetch() -> None:
892
start = time.perf_counter()
@@ -871,7 +897,12 @@ def _prefetch_all_checkpoints(sorted_files: list[str]) -> None:
897
elapsed,
898
)
899
874
- logger.info("Prefetching checkpoint files into page cache started (in background)")
900
+ logger.info(
901
+ "Prefetching checkpoint files into page cache started "
902
+ "(in background, num_threads=%d, block_size=%d bytes)",
903
+ num_prefetch_threads,
904
+ block_size,
905
+ )
906
threading.Thread(target=_run_prefetch, daemon=True).start()
907
908
@@ -880,6 +911,9 @@ def safetensors_weights_iterator(
911
use_tqdm_on_load: bool,
912
safetensors_load_strategy: str | None = None,
913
local_expert_ids: set[int] | None = None,
914
+ *,
915
+ safetensors_prefetch_num_threads: int = DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS,
916
+ safetensors_prefetch_block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
917
) -> Generator[tuple[str, torch.Tensor], None, None]:
918
"""Iterate over the weights in the model safetensor files.
919
@@ -951,7 +985,11 @@ def safetensors_weights_iterator(
985
)
986
987
if should_prefetch:
954
- _prefetch_all_checkpoints(sorted_files)
988
+ _prefetch_all_checkpoints(
989
+ sorted_files,
990
+ num_prefetch_threads=safetensors_prefetch_num_threads,
991
+ block_size=safetensors_prefetch_block_size,
992
+ )
993
994
leftover_state_dict: dict[str, torch.Tensor] = {}
995
for st_file in tqdm(