From 6f6b9c6e42a927f087276536f225147eb8507714 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Tue, 3 Feb 2026 15:21:13 +0800 Subject: [PATCH] [Perf] Use safetensors `load_file` in multithread loader (#18124) --- python/sglang/srt/model_loader/weight_utils.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index c7904c47d..e4371a3b6 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -840,12 +840,8 @@ def multi_thread_safetensors_weights_iterator( def _load_file(st_file: str): if disable_mmap: with open(st_file, "rb") as f: - result = safetensors.torch.load(f.read()) - else: - with safetensors.safe_open(st_file, framework="pt", device="cpu") as f: - result = {k: f.get_tensor(k) for k in f.keys()} - - return result + return safetensors.torch.load(f.read()) + return safetensors.torch.load_file(st_file, device="cpu") with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [executor.submit(_load_file, st_file) for st_file in hf_weights_files]