Keep HF cuFile/hipFile weights on device when memory permits

tests/model_executor/model_loader/test_direct_io.py

@@ -108,6 +108,7 @@ def test_large_cufile_tensor_bounds_compat_memory_and_reads_exact_tail(
     payload_offset = 8 + int.from_bytes(payload[:8], "little")
     reads, allocated, closed = [], [], []
     limit = 17
+    monkeypatch.setattr(direct_io, "_try_device_buffer", lambda _: None)
 
     class File:
         def __init__(self, filename, flags):
@@ -339,3 +340,55 @@ def test_destination_requires_current_accelerator(monkeypatch):
         )
         assert direct_io._can_load_into(tensor) is expected
     assert not direct_io._can_load_into(torch.empty(1))
+
+
+@pytest.mark.parametrize("budget", ["available", "low", "allocation_race"])
+def test_large_hf_weight_uses_device_buffer_or_bounded_fallback(
+    monkeypatch, tmp_path, budget
+):
+    value = torch.arange(37, dtype=torch.float32)
+    path = tmp_path / "hf.safetensors"
+    save_file({"weight": value}, path)
+    payload = path.read_bytes()
+    allocated, attempts = [], []
+    size = value.numel() * value.element_size()
+
+    def device_empty(length):
+        attempts.append(length)
+        if budget == "allocation_race" and length == size:
+            raise torch.OutOfMemoryError("competing allocation")
+        tensor = torch.empty(length, dtype=torch.uint8)
+        allocated.append(tensor)
+        return tensor
+
+    @contextmanager
+    def open_reader(_):
+        def read(raw, offset, length):
+            raw.copy_(
+                torch.frombuffer(
+                    bytearray(payload[offset : offset + length]), dtype=torch.uint8
+                )
+            )
+
+        yield read
+
+    monkeypatch.setattr(direct_io, "_empty", device_empty)
+    free = 0 if budget == "low" else 1024 * 1024 * 1024
+    monkeypatch.setattr(torch.accelerator, "get_memory_info", lambda: (free, free))
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
+    loaded = dict(
+        direct_io.direct_weights_iterator(
+            [str(path)],
+            backend="test",
+            mode="compat",
+            open_reader=open_reader,
+            coalesce_bytes=17,
+        )
+    )["weight"]
+    torch.testing.assert_close(loaded, value)
+    if budget == "available":
+        assert attempts == [size]
+        assert loaded.data_ptr() == allocated[0].data_ptr()
+    else:
+        assert attempts == ([size, 17] if budget == "allocation_race" else [17])
+        assert loaded.data_ptr() != allocated[0].data_ptr()

vllm/model_executor/model_loader/direct_io.py

@@ -120,6 +120,19 @@ def _iter_contiguous_groups(
         yield group
 
 
+def _try_device_buffer(length: int) -> torch.Tensor | None:
+    free, _ = torch.accelerator.get_memory_info()
+    # Leave room for the model loader's own temporary tensors. The free-memory
+    # snapshot is advisory: another allocation may win the race below.
+    reserve = max(256 * 1024 * 1024, length // 4)
+    if length + reserve > free:
+        return None
+    try:
+        return _empty(length)
+    except torch.OutOfMemoryError:
+        return None
+
+
 def _can_load_into(tensor: torch.Tensor) -> bool:
     return (
         tensor.device.type == "cuda"
@@ -161,8 +174,9 @@ def direct_weights_iterator(
     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.
+    Oversized compatibility-mode tensors stay on device when free memory
+    permits. Otherwise bounded device reads assemble them on CPU, so TP slicing
+    can still proceed without a full tensor's temporary VRAM.
     """
     if coalesce_bytes is not None and coalesce_bytes < 0:
         raise ValueError("coalesce_bytes must be nonnegative")
@@ -204,6 +218,22 @@ def direct_weights_iterator(
                     yield group[0][0], destination
                     continue
                 if mode == "compat" and length > compat_buffer_bytes:
+                    device_raw = _try_device_buffer(length)
+                    if device_raw is not None:
+                        _read_into(read, device_raw, first_offset, compat_buffer_bytes)
+                        for name, dtype, shape, offset, tensor_length in group:
+                            tensor = (
+                                device_raw.narrow(
+                                    0, offset - first_offset, tensor_length
+                                )
+                                .view(dtype)
+                                .reshape(shape)
+                            )
+                            total += tensor_length
+                            yield name, tensor
+                            del tensor
+                        del device_raw
+                        continue
                     logger.info(
                         "%s compat bounded staging tensor=%s bytes=%d "
                         "device_buffer_bytes=%d",