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
[Feature] Add InstantTensor weight loader (#36139)
arlo committed
Mar 15, 2026 at 01:05 UTC
8c29042bb98e79546576ff1a46c9def863046258
13 files changed
+177
-5
docker/Dockerfile.cpu
+2
-2
@@ -31,7 +31,7 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
31
--mount=type=cache,target=/var/lib/apt,sharing=locked \
32
apt-get update -y \
33
&& apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates \
34
- gcc-12 g++-12 libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof xz-utils \
34
+ gcc-12 g++-12 libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof make xz-utils \
35
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \
36
&& curl -LsSf https://astral.sh/uv/install.sh | sh
37
@@ -154,7 +154,7 @@ WORKDIR /vllm-workspace
154
155
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
156
--mount=type=cache,target=/var/lib/apt,sharing=locked \
157
- apt-get install -y --no-install-recommends vim numactl make clangd-14
157
+ apt-get install -y --no-install-recommends vim numactl clangd-14
158
159
RUN ln -s /usr/bin/clangd-14 /usr/bin/clangd
160
docs/models/extensions/instanttensor.md
new
+31
@@ -0,0 +1,31 @@
1
+# Loading Model Weights with InstantTensor
2
+
3
+InstantTensor accelerates loading Safetensors weights on CUDA devices through distributed loading, pipelined prefetching, and direct I/O. InstantTensor also supports GDS (GPUDirect Storage) when available.
4
+For more details, see the [InstantTensor GitHub repository](https://github.com/scitix/InstantTensor).
5
+
6
+## Installation
7
+
8
+```bash
9
+pip install instanttensor
10
+```
11
+
12
+## Use InstantTensor in vLLM
13
+
14
+Add `--load-format instanttensor` as a command-line argument.
15
+
16
+For example:
17
+
18
+```bash
19
+vllm serve Qwen/Qwen2.5-0.5B --load-format instanttensor
20
+```
21
+
22
+## Benchmarks
23
+
24
+| Model | GPU | Backend | Load Time (s) | Throughput (GB/s) | Speedup |
25
+| --- | ---: | --- | ---: | ---: | --- |
26
+| Qwen3-30B-A3B | 1*H200 | Safetensors | 57.4 | 1.1 | 1x |
27
+| Qwen3-30B-A3B | 1*H200 | InstantTensor | 1.77 | 35 | <span style="color: green">**32.4x**</span> |
28
+| DeepSeek-R1 | 8*H200 | Safetensors | 160 | 4.3 | 1x |
29
+| DeepSeek-R1 | 8*H200 | InstantTensor | 15.3 | 45 | <span style="color: green">**10.5x**</span> |
30
+
31
+For the full benchmark results, see <https://github.com/scitix/InstantTensor/blob/main/docs/benchmark.md>.
requirements/nightly_torch_test.txt
+1
@@ -44,4 +44,5 @@ numba == 0.61.2 # Required for N-gram speculative decoding
44
numpy
45
runai-model-streamer[s3,gcs]==0.15.3
46
fastsafetensors>=0.2.2
47
+instanttensor>=0.1.5
48
pydantic>=2.12 # 2.11 leads to error on python 3.13
requirements/test.in
+1
@@ -57,6 +57,7 @@ numba == 0.61.2 # Required for N-gram speculative decoding
57
numpy
58
runai-model-streamer[s3,gcs]==0.15.3
59
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
60
+instanttensor>=0.1.5
61
pydantic>=2.12 # 2.11 leads to error on python 3.13
62
decord==0.6.0
63
terratorch >= 1.2.2 # Required for Prithvi tests
requirements/test.txt
+3
@@ -375,6 +375,8 @@ inflect==5.6.2
375
# via datamodel-code-generator
376
iniconfig==2.0.0
377
# via pytest
378
+instanttensor==0.1.5
379
+ # via -r requirements/test.in
380
isoduration==20.11.0
381
# via jsonschema
382
isort==5.13.2
@@ -1169,6 +1171,7 @@ torch==2.10.0+cu129
1171
# accelerate
1172
# bitsandbytes
1173
# encodec
1174
+ # instanttensor
1175
# kornia
1176
# lightly
1177
# lightning
setup.py
+1
@@ -969,6 +969,7 @@ setup(
969
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
970
"tensorizer": ["tensorizer==2.10.1"],
971
"fastsafetensors": ["fastsafetensors >= 0.2.2"],
972
+ "instanttensor": ["instanttensor >= 0.1.5"],
973
"runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
974
"audio": [
975
"librosa",
tests/model_executor/model_loader/instanttensor_loader/__init__.py
tests/model_executor/model_loader/instanttensor_loader/test_instanttensor_loader.py
new
+28
@@ -0,0 +1,28 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm import SamplingParams
7
+from vllm.platforms import current_platform
8
+
9
+test_model = "openai-community/gpt2"
10
+
11
+prompts = [
12
+ "Hello, my name is",
13
+ "The president of the United States is",
14
+ "The capital of France is",
15
+ "The future of AI is",
16
+]
17
+# Create a sampling params object.
18
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
19
+
20
+
21
+@pytest.mark.skipif(
22
+ not current_platform.is_cuda(),
23
+ reason="InstantTensor requires NVIDIA GPUs",
24
+)
25
+def test_model_loader_download_files(vllm_runner):
26
+ with vllm_runner(test_model, load_format="instanttensor") as llm:
27
+ deserialized_outputs = llm.generate(prompts, sampling_params)
28
+ assert deserialized_outputs
tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py
new
+52
@@ -0,0 +1,52 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import glob
5
+import tempfile
6
+
7
+import huggingface_hub.constants
8
+import pytest
9
+import torch
10
+
11
+from vllm.model_executor.model_loader.weight_utils import (
12
+ download_weights_from_hf,
13
+ instanttensor_weights_iterator,
14
+ safetensors_weights_iterator,
15
+)
16
+from vllm.platforms import current_platform
17
+
18
+
19
+@pytest.mark.skipif(
20
+ not current_platform.is_cuda(),
21
+ reason="InstantTensor requires NVIDIA GPUs",
22
+)
23
+def test_instanttensor_model_loader():
24
+ with tempfile.TemporaryDirectory() as tmpdir:
25
+ huggingface_hub.constants.HF_HUB_OFFLINE = False
26
+ download_weights_from_hf(
27
+ "openai-community/gpt2", allow_patterns=["*.safetensors"], cache_dir=tmpdir
28
+ )
29
+ safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
30
+ assert len(safetensors) > 0
31
+
32
+ instanttensor_tensors = {}
33
+ hf_safetensors_tensors = {}
34
+
35
+ for name, tensor in instanttensor_weights_iterator(safetensors, True):
36
+ # Copy the tensor immediately as it is a reference to the internal
37
+ # buffer of instanttensor.
38
+ instanttensor_tensors[name] = tensor.to("cpu")
39
+
40
+ for name, tensor in safetensors_weights_iterator(safetensors, True):
41
+ hf_safetensors_tensors[name] = tensor
42
+
43
+ assert len(instanttensor_tensors) == len(hf_safetensors_tensors)
44
+
45
+ for name, instanttensor_tensor in instanttensor_tensors.items():
46
+ assert instanttensor_tensor.dtype == hf_safetensors_tensors[name].dtype
47
+ assert instanttensor_tensor.shape == hf_safetensors_tensors[name].shape
48
+ assert torch.all(instanttensor_tensor.eq(hf_safetensors_tensors[name]))
49
+
50
+
51
+if __name__ == "__main__":
52
+ test_instanttensor_model_loader()
vllm/config/load.py
+4
-1
@@ -29,6 +29,9 @@ class LoadConfig:
29
back to the pytorch bin format if safetensors format is not available.\n
30
- "pt" will load the weights in the pytorch bin format.\n
31
- "safetensors" will load the weights in the safetensors format.\n
32
+ - "instanttensor" will load the Safetensors weights on CUDA devices using
33
+ InstantTensor, which enables distributed loading with pipelined prefetching
34
+ and fast direct I/O.\n
35
- "npcache" will load the weights in pytorch format and store a numpy cache
36
to speed up the loading.\n
37
- "dummy" will initialize the weights with random values, which is mainly
@@ -46,7 +49,7 @@ class LoadConfig:
49
- "gguf" will load weights from GGUF format files (details specified in
50
https://github.com/ggml-org/ggml/blob/master/docs/gguf.md).\n
51
- "mistral" will load weights from consolidated safetensors files used by
49
- Mistral models.
52
+ Mistral models.\n
53
- Other custom values can be supported via plugins."""
54
download_dir: str | None = None
55
"""Directory to download and load the weights, default to the default
vllm/model_executor/model_loader/__init__.py
+2
@@ -35,6 +35,7 @@ LoadFormats = Literal[
35
"dummy",
36
"fastsafetensors",
37
"gguf",
38
+ "instanttensor",
39
"mistral",
40
"npcache",
41
"pt",
@@ -51,6 +52,7 @@ _LOAD_FORMAT_TO_MODEL_LOADER: dict[str, type[BaseModelLoader]] = {
52
"dummy": DummyModelLoader,
53
"fastsafetensors": DefaultModelLoader,
54
"gguf": GGUFModelLoader,
55
+ "instanttensor": DefaultModelLoader,
56
"mistral": DefaultModelLoader,
57
"npcache": DefaultModelLoader,
58
"pt": DefaultModelLoader,
vllm/model_executor/model_loader/default_loader.py
+11
-1
@@ -23,6 +23,7 @@ from vllm.model_executor.model_loader.weight_utils import (
23
filter_duplicate_safetensors_files,
24
filter_files_not_needed_for_inference,
25
get_quant_config,
26
+ instanttensor_weights_iterator,
27
maybe_download_from_modelscope,
28
multi_thread_pt_weights_iterator,
29
multi_thread_safetensors_weights_iterator,
@@ -121,7 +122,11 @@ class DefaultModelLoader(BaseModelLoader):
122
# Some quantized models use .pt files for storing the weights.
123
if load_format == "hf":
124
allow_patterns = ["*.safetensors", "*.bin"]
124
- elif load_format == "safetensors" or load_format == "fastsafetensors":
125
+ elif (
126
+ load_format == "safetensors"
127
+ or load_format == "fastsafetensors"
128
+ or load_format == "instanttensor"
129
+ ):
130
use_safetensors = True
131
allow_patterns = ["*.safetensors"]
132
elif load_format == "mistral":
@@ -219,6 +224,11 @@ class DefaultModelLoader(BaseModelLoader):
224
hf_weights_files,
225
self.load_config.use_tqdm_on_load,
226
)
227
+ elif self.load_config.load_format == "instanttensor":
228
+ weights_iterator = instanttensor_weights_iterator(
229
+ hf_weights_files,
230
+ self.load_config.use_tqdm_on_load,
231
+ )
232
else:
233
if extra_config.get("enable_multithread_load"):
234
weights_iterator = multi_thread_safetensors_weights_iterator(
vllm/model_executor/model_loader/weight_utils.py
+41
-1
@@ -29,7 +29,7 @@ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
29
from vllm import envs
30
from vllm.config import ModelConfig
31
from vllm.config.load import LoadConfig
32
-from vllm.distributed import get_tensor_model_parallel_rank
32
+from vllm.distributed import get_tensor_model_parallel_rank, get_world_group
33
from vllm.logger import init_logger
34
from vllm.model_executor.layers.quantization import (
35
QuantizationConfig,
@@ -909,6 +909,46 @@ def fastsafetensors_weights_iterator(
909
loader.close()
910
911
912
+def instanttensor_weights_iterator(
913
+ hf_weights_files: list[str],
914
+ use_tqdm_on_load: bool,
915
+) -> Generator[tuple[str, torch.Tensor], None, None]:
916
+ """Iterate over the weights in the model safetensor files
917
+ using instanttensor library."""
918
+ try:
919
+ import instanttensor
920
+ except ImportError as e:
921
+ raise ImportError(
922
+ "Please install instanttensor via `pip install instanttensor`"
923
+ ) from e
924
+
925
+ if not current_platform.is_cuda():
926
+ raise ValueError("InstantTensor requires NVIDIA GPUs")
927
+
928
+ try:
929
+ world_group = get_world_group()
930
+ except AssertionError:
931
+ # Entering here only in unit tests where the world group is not initialized.
932
+ process_group = None
933
+ else:
934
+ process_group = world_group.device_group if world_group.world_size > 1 else None
935
+
936
+ device = current_platform.current_device()
937
+
938
+ with instanttensor.safe_open(
939
+ hf_weights_files, framework="pt", device=device, process_group=process_group
940
+ ) as f:
941
+ yield from tqdm(
942
+ f.tensors(),
943
+ desc="Loading safetensors using InstantTensor loader",
944
+ disable=not enable_tqdm(use_tqdm_on_load),
945
+ bar_format=_BAR_FORMAT,
946
+ position=tqdm._get_free_pos(),
947
+ total=len(f.keys()),
948
+ mininterval=1.0,
949
+ )
950
+
951
+
952
def pt_weights_iterator(
953
hf_weights_files: list[str],
954
use_tqdm_on_load: bool,