[Loader] Pinned host staging for the Run:ai streamer path

docs/models/extensions/runai_model_streamer.md

@@ -82,6 +82,14 @@ vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
     --model-loader-extra-config '{"memory_limit":5368709120}'
 ```
 
+You can set `pinned_staging` to page-lock the streamer's CPU buffer and have the streamer copy every tensor to the GPU out of it, so the host-to-device copies run as DMA copies from pinned memory instead of pageable copies. This helps on hosts where the weight load is bound by the copy path rather than by storage (for example a large checkpoint on a fast local disk or a parallel filesystem). The buffer is registered once per load, at roughly 0.2 s per GB, so unless `memory_limit` is set it is bounded to 16 GB. If the buffer cannot be pinned, loading falls back to the regular path.
+
+```bash
+vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
+    --load-format runai_streamer \
+    --model-loader-extra-config '{"pinned_staging":true}'
+```
+
 !!! note
     For further instructions about tunable parameters and additional parameters configurable through environment variables, read the [Environment Variables Documentation](https://github.com/run-ai/runai-model-streamer/blob/master/docs/src/env-vars.md).
 

tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py

@@ -96,6 +96,7 @@ def _runai_loader(extra):
         ({"distributed": "yes"}, "distributed must be a bool"),
         ({"concurrency": "16"}, "concurrency must be a positive integer"),
         ({"concurrency": -1}, "concurrency must be a positive integer"),
+        ({"pinned_staging": "yes"}, "pinned_staging must be a bool"),
     ],
 )
 def test_runai_rejects_invalid_extra_config(extra, match):
@@ -116,6 +117,29 @@ def test_runai_accepts_valid_extra_config():
         assert os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] == "1024"
 
 
+def test_runai_pinned_staging_bounds_default_memory_limit():
+    # Page-locking is paid per byte of staging buffer, so the streamer's 40 GB
+    # default is bounded unless the user chose a limit themselves.
+    with patch.dict(os.environ, {}, clear=False):
+        os.environ.pop("RUNAI_STREAMER_MEMORY_LIMIT", None)
+        loader = _runai_loader({"pinned_staging": True})
+        assert loader._pinned_staging is True
+        assert os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] == str(
+            rsl.RUNAI_PINNED_STAGING_DEFAULT_BYTES
+        )
+
+    with patch.dict(os.environ, {}, clear=False):
+        os.environ.pop("RUNAI_STREAMER_MEMORY_LIMIT", None)
+        _runai_loader({"pinned_staging": True, "memory_limit": 1024})
+        assert os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] == "1024"
+
+    with patch.dict(os.environ, {"RUNAI_STREAMER_MEMORY_LIMIT": "2048"}, clear=False):
+        _runai_loader({"pinned_staging": True})
+        assert os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] == "2048"
+
+    assert _runai_loader({"distributed": True})._pinned_staging is False
+
+
 def test_runai_invalid_extra_config_leaves_environ_untouched():
     # A later invalid key must not leave an earlier valid key applied to
     # os.environ (all values are validated before any global mutation).

tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py

@@ -36,6 +36,31 @@ def test_runai_safetensors_weights_iterator_clones_reused_buffers(
         assert torch.equal(actual_tensors[name], expected_tensor)
 
 
+def test_runai_safetensors_weights_iterator_pinned_staging_fallback(
+    tmp_path, monkeypatch
+):
+    # Without a CUDA device the staging buffer cannot be registered; the
+    # iterator must then yield the same (copied-out) tensors as the plain path.
+    monkeypatch.setenv("RUNAI_STREAMER_MEMORY_LIMIT", "0")
+    weights_file = tmp_path / "model.safetensors"
+    expected_tensors = {
+        "first": torch.tensor([1.0, 2.0]),
+        "second": torch.tensor([3.0, 4.0]),
+    }
+    save_file(expected_tensors, weights_file)
+
+    actual_tensors = dict(
+        runai_safetensors_weights_iterator(
+            [str(weights_file)], False, pinned_staging=True
+        )
+    )
+
+    assert actual_tensors.keys() == expected_tensors.keys()
+    assert actual_tensors["first"].data_ptr() != actual_tensors["second"].data_ptr()
+    for name, expected_tensor in expected_tensors.items():
+        assert torch.equal(actual_tensors[name], expected_tensor)
+
+
 def test_runai_model_loader():
     with tempfile.TemporaryDirectory() as tmpdir:
         huggingface_hub.constants.HF_HUB_OFFLINE = False

vllm/model_executor/model_loader/runai_streamer_loader.py

@@ -11,6 +11,7 @@
 from vllm.config.load import LoadConfig
 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
 from vllm.model_executor.model_loader.weight_utils import (
+    RUNAI_PINNED_STAGING_DEFAULT_BYTES,
     download_safetensors_index_file_from_hf,
     download_weights_from_hf,
     runai_safetensors_weights_iterator,
@@ -28,10 +29,16 @@ def __init__(self, load_config: LoadConfig):
         super().__init__(load_config)
 
         self._is_distributed: bool = False
+        self._pinned_staging: bool = False
         if load_config.model_loader_extra_config:
             extra_config = load_config.model_loader_extra_config
 
-            allowed_keys = {"distributed", "concurrency", "memory_limit"}
+            allowed_keys = {
+                "distributed",
+                "concurrency",
+                "memory_limit",
+                "pinned_staging",
+            }
             if unexpected_keys := set(extra_config) - allowed_keys:
                 raise ValueError(
                     "Unexpected extra config keys for runai_streamer: "
@@ -44,6 +51,14 @@ def __init__(self, load_config: LoadConfig):
                     raise ValueError(f"distributed must be a bool, got {distributed!r}")
                 self._is_distributed = distributed
 
+            if "pinned_staging" in extra_config:
+                pinned_staging = extra_config["pinned_staging"]
+                if not isinstance(pinned_staging, bool):
+                    raise ValueError(
+                        f"pinned_staging must be a bool, got {pinned_staging!r}"
+                    )
+                self._pinned_staging = pinned_staging
+
             # Validate every value before mutating os.environ, so a later
             # invalid key cannot leave an earlier one partially applied.
             env_updates: dict[str, str] = {}
@@ -70,6 +85,15 @@ def __init__(self, load_config: LoadConfig):
                         f"memory_limit must be an integer >= -1, got {memory_limit!r}"
                     )
                 env_updates["RUNAI_STREAMER_MEMORY_LIMIT"] = str(memory_limit)
+            elif (
+                self._pinned_staging
+                and os.getenv("RUNAI_STREAMER_MEMORY_LIMIT") is None
+            ):
+                # Page-locking is paid once per byte of staging buffer, so
+                # bound the streamer's default (40 GB) unless the user chose.
+                env_updates["RUNAI_STREAMER_MEMORY_LIMIT"] = str(
+                    RUNAI_PINNED_STAGING_DEFAULT_BYTES
+                )
             os.environ.update(env_updates)
 
             runai_streamer_s3_endpoint = os.getenv("RUNAI_STREAMER_S3_ENDPOINT")
@@ -127,7 +151,10 @@ def _get_weights_iterator(
         """Get an iterator for the model weights based on the load format."""
         hf_weights_files = self._prepare_weights(model_or_path, revision)
         return runai_safetensors_weights_iterator(
-            hf_weights_files, self.load_config.use_tqdm_on_load, self._is_distributed
+            hf_weights_files,
+            self.load_config.use_tqdm_on_load,
+            self._is_distributed,
+            pinned_staging=self._pinned_staging,
         )
 
     def download_model(self, model_config: ModelConfig) -> None:

vllm/model_executor/model_loader/weight_utils.py

@@ -1243,25 +1243,129 @@ def _load_file(st_file: str):
                 yield key, state_dict.pop(key)
 
 
+# Default size of the Run:ai streamer's CPU staging buffer when it is
+# page-locked (``pinned_staging``). Registering the buffer costs roughly
+# 0.2 s/GB once per load; 16 GB keeps that to a few seconds while leaving the
+# streamer enough room to keep every reader busy.
+RUNAI_PINNED_STAGING_DEFAULT_BYTES = 16_000_000_000
+# Threads used to pre-fault the staging buffer before it is page-locked.
+_RUNAI_PINNED_STAGING_TOUCH_THREADS = 16
+
+
+def _pin_runai_streamer_buffer(streamer) -> int | None:
+    """Page-lock the Run:ai streamer's CPU staging buffer.
+
+    The streamer reads every tensor into one reusable host buffer and the
+    host-to-device copies out of it run as pageable copies, i.e. a driver
+    staged memcpy per tensor. Registering the buffer with ``cudaHostRegister``
+    turns those into DMA copies out of pinned memory. The pages are faulted in
+    with ``MADV_POPULATE_WRITE`` from several threads first, so the
+    single-threaded registration does not pay for the page faults itself.
+
+    Returns the registered address, to be passed to ``cudaHostUnregister``
+    once the iterator is exhausted, or ``None`` if the buffer could not be
+    pinned. In that case loading continues unchanged on the pageable path.
+    """
+    try:
+        buf = streamer.file_streamer.file_streamer.requests_iterator.buffer
+    except AttributeError as e:
+        logger.warning(
+            "Run:ai streamer staging buffer not found (%s); "
+            "loading from pageable memory.",
+            e,
+        )
+        return None
+    nbytes = int(buf.nbytes)
+    if nbytes == 0:
+        return None
+    addr = int(buf.ctypes.data)
+    start_time = time.perf_counter()
+    try:
+        import ctypes
+
+        libc = ctypes.CDLL("libc.so.6", use_errno=True)
+        page = 4096
+        madv_hugepage = 14  # MADV_HUGEPAGE
+        madv_populate_write = 23  # MADV_POPULATE_WRITE, Linux >= 5.14
+        begin = addr & ~(page - 1)
+        end = (addr + nbytes + page - 1) & ~(page - 1)
+        libc.madvise(
+            ctypes.c_void_p(begin), ctypes.c_size_t(end - begin), madv_hugepage
+        )
+        step = (end - begin) // _RUNAI_PINNED_STAGING_TOUCH_THREADS
+        step = max((step + page - 1) & ~(page - 1), page)
+
+        def _populate(lo: int, hi: int) -> None:
+            libc.madvise(
+                ctypes.c_void_p(lo), ctypes.c_size_t(hi - lo), madv_populate_write
+            )
+
+        threads = [
+            threading.Thread(target=_populate, args=(lo, min(lo + step, end)))
+            for lo in range(begin, end, step)
+        ]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+        populated_time = time.perf_counter()
+
+        rc = torch.cuda.cudart().cudaHostRegister(addr, nbytes, 0)
+        if int(rc) != 0:
+            logger.warning(
+                "cudaHostRegister on the Run:ai streamer staging buffer returned "
+                "%s; loading from pageable memory.",
+                rc,
+            )
+            return None
+        logger.info(
+            "Page-locked the %.1f GB Run:ai streamer staging buffer "
+            "(populate %.2f s, register %.2f s).",
+            nbytes / 1e9,
+            populated_time - start_time,
+            time.perf_counter() - populated_time,
+        )
+        return addr
+    except Exception as e:
+        logger.warning(
+            "Could not page-lock the Run:ai streamer staging buffer (%s); "
+            "loading from pageable memory.",
+            e,
+        )
+        return None
+
+
 def runai_safetensors_weights_iterator(
     hf_weights_files: list[str],
     use_tqdm_on_load: bool,
     is_distributed: bool = False,
+    pinned_staging: bool = False,
 ) -> Generator[tuple[str, torch.Tensor], None, None]:
-    """Iterate over the weights in the model safetensor files."""
+    """Iterate over the weights in the model safetensor files.
+
+    Args:
+        hf_weights_files: safetensors files to stream.
+        use_tqdm_on_load: show a progress bar.
+        is_distributed: use the streamer's distributed mode.
+        pinned_staging: page-lock the streamer's CPU staging buffer and let
+            the streamer copy each tensor to the current device out of it, so
+            the weight loaders only ever see device tensors. Falls back to the
+            pageable path if the buffer cannot be pinned.
+
+    """
     with SafetensorsStreamer() as streamer:
         is_cuda_alike = current_platform.is_cuda_alike()
+        stream_to_device = is_cuda_alike and (is_distributed or pinned_staging)
         device = (
-            f"cuda:{current_platform.current_device()}"
-            if is_distributed and is_cuda_alike
-            else "cpu"
+            f"cuda:{current_platform.current_device()}" if stream_to_device else "cpu"
         )
 
         streamer.stream_files(
             hf_weights_files,
             device=device,
             is_distributed=is_distributed,
         )
+        pinned_addr = _pin_runai_streamer_buffer(streamer) if pinned_staging else None
         total_tensors = sum(
             len(tensors_meta)
             for tensors_meta in streamer.files_to_tensors_metadata.values()
@@ -1276,8 +1380,23 @@ def runai_safetensors_weights_iterator(
             mininterval=2,
         )
 
-        for name, tensor in tensor_iter:
-            yield name, tensor.clone()
+        try:
+            for name, tensor in tensor_iter:
+                # The streamer reuses its staging buffers, so a tensor must be
+                # copied out before the next chunk lands in the same memory.
+                # On the pinned path (non-distributed) the streamer has already
+                # moved the tensor to the device, which is that copy.
+                if (
+                    pinned_addr is not None
+                    and not is_distributed
+                    and tensor.device.type != "cpu"
+                ):
+                    yield name, tensor
+                else:
+                    yield name, tensor.clone()
+        finally:
+            if pinned_addr is not None:
+                torch.cuda.cudart().cudaHostUnregister(pinned_addr)
 
 
 def fastsafetensors_weights_iterator(