Add shared cuFile/hipFile weight destinations

tests/model_executor/model_loader/test_direct_io.py

@@ -231,3 +231,111 @@ def read(raw, offset, length):
         )
     )
     assert reads == []
+
+
+@pytest.mark.parametrize("invalid", [False, True])
+def test_destination_reads_preserve_storage_offset_and_reject_shape(
+    monkeypatch, tmp_path, invalid
+):
+    value = torch.arange(37, dtype=torch.float32)
+    path = tmp_path / "destination.safetensors"
+    save_file({"weight": value}, path)
+    payload = path.read_bytes()
+    storage = torch.full((45,), -1.0)
+    destination = storage[4:41] if not invalid else storage[4:40]
+    reads = []
+
+    @contextmanager
+    def open_reader(_):
+        def read(raw, offset, length):
+            reads.append(length)
+            raw.copy_(
+                torch.frombuffer(
+                    bytearray(payload[offset : offset + length]), dtype=torch.uint8
+                )
+            )
+
+        yield read
+
+    monkeypatch.setattr(direct_io, "_can_load_into", lambda _: True)
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
+    iterator = direct_io.direct_weights_iterator(
+        [str(path)],
+        backend="test",
+        mode="compat",
+        open_reader=open_reader,
+        coalesce_bytes=17,
+        get_destination=lambda *_: destination,
+    )
+    if invalid:
+        with pytest.raises(ValueError, match="Invalid direct-load destination"):
+            list(iterator)
+        assert reads == []
+        assert torch.all(storage == -1)
+    else:
+        loaded = dict(iterator)
+        assert loaded["weight"] is destination
+        torch.testing.assert_close(destination, value)
+        assert torch.all(storage[:4] == -1)
+        assert torch.all(storage[41:] == -1)
+        assert reads == [17] * 8 + [12]
+
+
+def test_destination_read_failure_does_not_yield_partial_weight(monkeypatch, tmp_path):
+    path = tmp_path / "failed.safetensors"
+    save_file({"weight": torch.arange(8, dtype=torch.uint8)}, path)
+    destination = torch.empty(8, dtype=torch.uint8)
+    calls = []
+
+    @contextmanager
+    def open_reader(_):
+        def read(raw, offset, length):
+            calls.append(length)
+            if len(calls) == 2:
+                raise EOFError("short read")
+            raw.fill_(1)
+
+        yield read
+
+    monkeypatch.setattr(direct_io, "_can_load_into", lambda _: True)
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
+    iterator = direct_io.direct_weights_iterator(
+        [str(path)],
+        backend="test",
+        mode="compat",
+        open_reader=open_reader,
+        coalesce_bytes=4,
+        get_destination=lambda *_: destination,
+    )
+    with pytest.raises(EOFError, match="short read"):
+        next(iterator)
+    assert calls == [4, 4]
+
+
+def test_destination_rejects_negative_chunk_size_before_io():
+    def unexpected(*args):
+        pytest.fail("invalid chunk size must be rejected before I/O")
+
+    with pytest.raises(ValueError, match="coalesce_bytes must be nonnegative"):
+        list(
+            direct_io.direct_weights_iterator(
+                ["unused"],
+                backend="test",
+                mode="compat",
+                open_reader=unexpected,
+                coalesce_bytes=-1,
+                get_destination=unexpected,
+            )
+        )
+
+
+def test_destination_requires_current_accelerator(monkeypatch):
+    from types import SimpleNamespace
+
+    monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 1)
+    for index, expected in [(0, False), (1, True)]:
+        tensor = SimpleNamespace(
+            device=torch.device("cuda", index), is_contiguous=lambda: True
+        )
+        assert direct_io._can_load_into(tensor) is expected
+    assert not direct_io._can_load_into(torch.empty(1))

vllm/model_executor/model_loader/cufile.py

@@ -10,6 +10,7 @@
 
 from vllm import envs
 from vllm.model_executor.model_loader.direct_io import (
+    WeightDestination,
     check_compat_env,
     direct_weights_iterator,
 )
@@ -20,6 +21,7 @@ def cufile_weights_iterator(
     local_expert_ids: set[int] | None = None,
     *,
     is_unused_weight: Callable[[str], bool] | None = None,
+    get_destination: WeightDestination | None = None,
 ) -> Generator[tuple[str, torch.Tensor]]:
     """Read safetensors payloads into the current CUDA device through cuFile.
 
@@ -55,4 +57,5 @@ def read(raw: torch.Tensor, offset: int, length: int) -> None:
         mode=mode,
         open_reader=open_reader,
         is_unused_weight=is_unused_weight,
+        get_destination=get_destination,
     )

vllm/model_executor/model_loader/direct_io.py

@@ -38,6 +38,7 @@
 # Reads `length` bytes at `file_offset` into the uint8 device tensor `raw`.
 Reader = Callable[[torch.Tensor, int, int], None]
 TensorEntry = tuple[str, torch.dtype, list[int], int, int]
+WeightDestination = Callable[[str, torch.dtype, list[int]], torch.Tensor | None]
 # Compatibility modes accept arbitrary offsets; P2P keeps per-tensor I/O for
 # direct-I/O alignment requirements.
 COMPAT_COALESCE_BYTES = 64 * 1024 * 1024
@@ -119,6 +120,24 @@ def _iter_contiguous_groups(
         yield group
 
 
+def _can_load_into(tensor: torch.Tensor) -> bool:
+    return (
+        tensor.device.type == "cuda"
+        and tensor.device.index == torch.accelerator.current_device_index()
+        and tensor.is_contiguous()
+    )
+
+
+def _read_into(read: Reader, tensor: torch.Tensor, offset: int, limit: int) -> None:
+    raw = tensor.detach().view(-1).view(torch.uint8)
+    # Complete prior initialization/copies before external storage writes.
+    torch.accelerator.synchronize()
+    for begin in range(0, raw.numel(), limit):
+        count = min(limit, raw.numel() - begin)
+        read(raw.narrow(0, begin, count), offset + begin, count)
+    torch.accelerator.synchronize()
+
+
 def direct_weights_iterator(
     files: Iterable[str],
     *,
@@ -127,6 +146,7 @@ def direct_weights_iterator(
     open_reader: Callable[[str], Any],
     coalesce_bytes: int | None = None,
     is_unused_weight: Callable[[str], bool] | None = None,
+    get_destination: WeightDestination | None = None,
 ) -> Generator[tuple[str, torch.Tensor]]:
     """Read tensor payloads with bounded device staging in compatibility mode.
 
@@ -136,9 +156,16 @@ def direct_weights_iterator(
     default loader supplies the complete checkpoint, ShardedStateLoader only
     this TP rank's saved shards.
 
+    In compatibility mode, get_destination may return a matching contiguous
+    tensor on the current accelerator. It is written in place; a read failure
+    leaves its contents unspecified and must abort loading. Other destinations
+    fall back to staging. P2P does not use destination callbacks.
+
     Oversized compatibility-mode tensors are assembled on CPU using bounded
     device reads, so TP slicing does not require a full tensor's temporary VRAM.
     """
+    if coalesce_bytes is not None and coalesce_bytes < 0:
+        raise ValueError("coalesce_bytes must be nonnegative")
     if coalesce_bytes is None:
         coalesce_bytes = COMPAT_COALESCE_BYTES if mode == "compat" else 0
     compat_buffer_bytes = coalesce_bytes or COMPAT_COALESCE_BYTES
@@ -151,14 +178,31 @@ def direct_weights_iterator(
             entries = iter_safetensors_entries(path)
             if is_unused_weight is not None:
                 entries = (entry for entry in entries if not is_unused_weight(entry[0]))
-            if coalesce_bytes:
+            if coalesce_bytes and get_destination is None:
                 groups = _iter_contiguous_groups(entries, coalesce_bytes)
             else:
                 groups = ([entry] for entry in entries)
             for group in groups:
                 first_offset = group[0][3]
                 _, _, _, last_offset, last_length = group[-1]
                 length = last_offset + last_length - first_offset
+                destination = None
+                if mode == "compat" and get_destination is not None:
+                    name, dtype, shape, _, _ = group[0]
+                    destination = get_destination(name, dtype, shape)
+                    if destination is not None:
+                        if (
+                            destination.dtype != dtype
+                            or list(destination.shape) != shape
+                        ):
+                            raise ValueError(f"Invalid direct-load destination: {name}")
+                        if not _can_load_into(destination):
+                            destination = None
+                if destination is not None:
+                    _read_into(read, destination, first_offset, compat_buffer_bytes)
+                    total += length
+                    yield group[0][0], destination
+                    continue
                 if mode == "compat" and length > compat_buffer_bytes:
                     logger.info(
                         "%s compat bounded staging tensor=%s bytes=%d "

vllm/model_executor/model_loader/hipfile.py

@@ -11,6 +11,7 @@
 
 from vllm import envs
 from vllm.model_executor.model_loader.direct_io import (
+    WeightDestination,
     check_compat_env,
     direct_weights_iterator,
 )
@@ -21,6 +22,7 @@ def hipfile_weights_iterator(
     local_expert_ids: set[int] | None = None,
     *,
     is_unused_weight: Callable[[str], bool] | None = None,
+    get_destination: WeightDestination | None = None,
 ) -> Generator[tuple[str, torch.Tensor]]:
     """Read safetensors payloads into the current ROCm device through hipFile."""
     if local_expert_ids is not None:
@@ -59,4 +61,5 @@ def read(raw: torch.Tensor, offset: int, length: int) -> None:
             mode=mode,
             open_reader=open_reader,
             is_unused_weight=is_unused_weight,
+            get_destination=get_destination,
         )

vllm/model_executor/model_loader/sharded_state_loader.py

@@ -18,6 +18,7 @@
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
+from vllm.model_executor.model_loader.direct_io import WeightDestination
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
     runai_safetensors_weights_iterator,
@@ -210,7 +211,18 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
             )
         state_dict = self._filter_subtensors(model.state_dict())
         counter_before_loading_weights = time.perf_counter()
-        for key, tensor in self.iterate_over_files(filepaths):
+
+        def get_destination(name, dtype, shape):
+            param = state_dict[name]
+            if param.dtype != dtype or param.ndim != len(shape):
+                return None
+            if any(size > dim for size, dim in zip(shape, param.shape)):
+                return None
+            for dim, size in enumerate(shape):
+                param = param.narrow(dim, 0, size)
+            return param
+
+        for key, tensor in self.iterate_over_files(filepaths, get_destination):
             # If loading with LoRA enabled, additional padding may
             # be added to certain parameters. We only load into a
             # narrowed view of the parameter data.
@@ -245,19 +257,21 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
         if state_dict:
             raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!")
 
-    def iterate_over_files(self, paths) -> Generator[tuple[str, torch.Tensor]]:
+    def iterate_over_files(
+        self, paths, get_destination: WeightDestination | None = None
+    ) -> Generator[tuple[str, torch.Tensor]]:
         if envs.VLLM_CUFILE_LOAD and envs.VLLM_HIPFILE_LOAD:
             raise ValueError("Select either cuFile or hipFile loading")
         if envs.VLLM_HIPFILE_LOAD:
             from vllm.model_executor.model_loader.hipfile import (
                 hipfile_weights_iterator,
             )
 
-            yield from hipfile_weights_iterator(paths)
+            yield from hipfile_weights_iterator(paths, get_destination=get_destination)
         elif envs.VLLM_CUFILE_LOAD:
             from vllm.model_executor.model_loader.cufile import cufile_weights_iterator
 
-            yield from cufile_weights_iterator(paths)
+            yield from cufile_weights_iterator(paths, get_destination=get_destination)
         elif self.load_config.load_format == "runai_streamer_sharded":
             yield from runai_safetensors_weights_iterator(paths, True)
         else: