From 3d29ab6acaa31dcad47ba850e4356b6e32f84b2a Mon Sep 17 00:00:00 2001 From: Yuanyuan Chen Date: Thu, 22 Oct 2026 04:16:16 +0000 Subject: [PATCH 1/2] Add shared direct-storage weight destinations Co-authored-by: OpenAI Codex Signed-off-by: Yuanyuan Chen --- .../model_loader/test_direct_io.py | 79 +++++++++++++++++++ vllm/model_executor/model_loader/cufile.py | 3 + vllm/model_executor/model_loader/direct_io.py | 35 +++++++- vllm/model_executor/model_loader/hipfile.py | 3 + .../model_loader/sharded_state_loader.py | 22 +++++- 5 files changed, 137 insertions(+), 5 deletions(-) diff --git a/tests/model_executor/model_loader/test_direct_io.py b/tests/model_executor/model_loader/test_direct_io.py index 855012ea45e3..6bade6f50e3b 100644 --- a/tests/model_executor/model_loader/test_direct_io.py +++ b/tests/model_executor/model_loader/test_direct_io.py @@ -231,3 +231,82 @@ 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] diff --git a/vllm/model_executor/model_loader/cufile.py b/vllm/model_executor/model_loader/cufile.py index a0cd81600174..bcdfe8e34a9c 100644 --- a/vllm/model_executor/model_loader/cufile.py +++ b/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, ) diff --git a/vllm/model_executor/model_loader/direct_io.py b/vllm/model_executor/model_loader/direct_io.py index b474e5fb93ae..15e865fa3b31 100644 --- a/vllm/model_executor/model_loader/direct_io.py +++ b/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,20 @@ def _iter_contiguous_groups( yield group +def _can_load_into(tensor: torch.Tensor) -> bool: + return tensor.device.type == "cuda" 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 +142,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. @@ -151,7 +167,7 @@ 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) @@ -159,6 +175,23 @@ def direct_weights_iterator( 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 " diff --git a/vllm/model_executor/model_loader/hipfile.py b/vllm/model_executor/model_loader/hipfile.py index 4083ad6d6962..a4c775caddcd 100644 --- a/vllm/model_executor/model_loader/hipfile.py +++ b/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, ) diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py index 922bb25e816d..3547ac72a2ef 100644 --- a/vllm/model_executor/model_loader/sharded_state_loader.py +++ b/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,7 +257,9 @@ 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: @@ -253,11 +267,11 @@ def iterate_over_files(self, paths) -> Generator[tuple[str, torch.Tensor]]: 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: From 5defc10ecfba4206abaa8d61d43f06f5332f0363 Mon Sep 17 00:00:00 2001 From: Yuanyuan Chen Date: Thu, 22 Oct 2026 04:19:00 +0000 Subject: [PATCH 2/2] Validate direct destinations and chunk sizes Co-authored-by: OpenAI Codex Signed-off-by: Yuanyuan Chen --- .../model_loader/test_direct_io.py | 29 +++++++++++++++++++ vllm/model_executor/model_loader/direct_io.py | 13 ++++++++- 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/model_loader/test_direct_io.py b/tests/model_executor/model_loader/test_direct_io.py index 6bade6f50e3b..f0bdd85714e7 100644 --- a/tests/model_executor/model_loader/test_direct_io.py +++ b/tests/model_executor/model_loader/test_direct_io.py @@ -310,3 +310,32 @@ def read(raw, offset, length): 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)) diff --git a/vllm/model_executor/model_loader/direct_io.py b/vllm/model_executor/model_loader/direct_io.py index 15e865fa3b31..66f40b9325e6 100644 --- a/vllm/model_executor/model_loader/direct_io.py +++ b/vllm/model_executor/model_loader/direct_io.py @@ -121,7 +121,11 @@ def _iter_contiguous_groups( def _can_load_into(tensor: torch.Tensor) -> bool: - return tensor.device.type == "cuda" and tensor.is_contiguous() + 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: @@ -152,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