fix(engine): size memory pools against the VRAM other processes leave free on Windows

python/freetoken/engine/graph.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import gc
+import sys
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, List
 
 import torch
 from freetoken.core import Batch, Req, get_global_ctx
 from freetoken.distributed import get_tp_info
+from freetoken.gpu_select import gpu_uuid, nvml_free_bytes
 from freetoken.utils import init_logger, mem_GB
 from freetoken.utils.progress import emit_progress
 from tqdm import tqdm
@@ -99,7 +101,17 @@ def _determine_cuda_graph_bs(
 
 
 def get_free_memory(device: torch.device) -> int:
-    return torch.cuda.mem_get_info(device)[0]
+    """Device memory a pool can take without spilling.
+
+    Under WDDM (Windows) cudaMemGetInfo leaves out what other processes hold, so a card the desktop
+    uses reads as free and a pool sized to it pages into shared memory; NVML counts every process.
+    """
+    free = torch.cuda.mem_get_info(device)[0]
+    if sys.platform != "win32":
+        return free
+    index = torch.cuda.current_device() if device.index is None else device.index
+    machine_free = nvml_free_bytes(gpu_uuid(index))
+    return free if machine_free is None else min(free, machine_free)
 
 
 class GraphRunner:

python/freetoken/gpu_select.py

@@ -9,6 +9,7 @@
 Each worker publishes its own entry (set_assigned_gpu / assign_gpu) and binds it when CUDA comes up (bind_assigned_gpu) by matching the UUID against CUDA's visible devices.
 One process runs on one GPU. Binding is unconditional: a process that publishes nothing binds a default ordinal and records it, so assigned_visible_gpu() names that card in every case.
 No process mutates CUDA_VISIBLE_DEVICES, and the UUID match holds under any CUDA_DEVICE_ORDER.
+NVML also reports a GPU's free memory across every process (nvml_free_bytes), which WDDM's cudaMemGetInfo does not.
 
 Stdlib only (torch is imported lazily); not under freetoken.utils, which imports transformers.
 """
@@ -68,12 +69,11 @@ def single_gpu_arg(value: str) -> str:
     return entries[0]
 
 
-def _nvml_uuids() -> "list[str] | None":
-    """Full GPU UUIDs in physical (nvidia-smi) order, or None when NVML is unavailable.
+def _load_nvml():
+    """The NVML library, or None when no candidate loads.
 
     Own ctypes loader instead of torch's _raw_device_uuid_nvml: that helper only knows the Linux library name, raises (not None) when the library is missing, and is private API.
     NVML exports are cdecl on every platform, so CDLL is right on Windows too (same as nvidia-ml-py).
-    None on any failure -- no library, a stub library without the _v2 symbols, WSL, a dead device -- and callers fall back.
     """
     import ctypes
 
@@ -85,14 +85,24 @@ def _nvml_uuids() -> "list[str] | None":
         ]
     else:
         candidates = ["libnvidia-ml.so.1"]
+    for name in candidates:
+        try:
+            return ctypes.CDLL(name)
+        except OSError:
+            continue
+    return None
+
+
+def _nvml_uuids() -> "list[str] | None":
+    """Full GPU UUIDs in physical (nvidia-smi) order, or None when NVML is unavailable.
+
+    None on any failure -- no library, a stub library without the _v2 symbols, WSL, a dead device -- and callers fall back.
+    """
+    import ctypes
+
     try:
-        for name in candidates:
-            try:
-                lib = ctypes.CDLL(name)
-                break
-            except OSError:
-                continue
-        else:
+        lib = _load_nvml()
+        if lib is None:
             return None
         if lib.nvmlInit() != 0:
             return None
@@ -116,6 +126,36 @@ def _nvml_uuids() -> "list[str] | None":
         return None
 
 
+def nvml_free_bytes(uuid: "str | None") -> "int | None":
+    """Free memory of GPU ``uuid`` across every process, as nvidia-smi reports it; None when NVML cannot say.
+
+    Under WDDM, cudaMemGetInfo leaves out what other processes hold; this does not.
+    """
+    import ctypes
+
+    class Memory(ctypes.Structure):
+        _fields_ = [("total", ctypes.c_ulonglong), ("free", ctypes.c_ulonglong), ("used", ctypes.c_ulonglong)]
+
+    if uuid is None:
+        return None
+    try:
+        lib = _load_nvml()
+        if lib is None or lib.nvmlInit() != 0:
+            return None
+        try:
+            handle = ctypes.c_void_p()
+            if lib.nvmlDeviceGetHandleByUUID(uuid.encode("ascii"), ctypes.byref(handle)) != 0:
+                return None
+            memory = Memory()
+            if lib.nvmlDeviceGetMemoryInfo(handle, ctypes.byref(memory)) != 0:
+                return None
+            return int(memory.free)
+        finally:
+            lib.nvmlShutdown()
+    except (OSError, AttributeError):
+        return None
+
+
 def _match_uuid(spec: str, uuids: "list[str]", where: str) -> str:
     """The unique full UUID that ``spec`` prefixes, else ValueError."""
     hits = [u for u in uuids if u.upper().startswith(spec.upper())]
@@ -223,7 +263,7 @@ def _visible_of_physical(uuid: str) -> int:
     seen: list[str] = []
     hits: list[int] = []
     for v in range(torch.cuda.device_count()):
-        u = format_gpu_uuid(getattr(torch.cuda.get_device_properties(v), "uuid", None))
+        u = gpu_uuid(v)
         seen.append(u or "?")
         if u is not None and u.upper().startswith(uuid.upper()):
             hits.append(v)
@@ -274,6 +314,13 @@ def format_gpu_uuid(raw) -> str | None:
     return None if raw is None else f"{UUID_PREFIX}{raw}"
 
 
+def gpu_uuid(index: int) -> str | None:
+    """nvidia-smi form GPU-<uuid> of visible device ``index``; None where torch reports no uuid."""
+    import torch
+
+    return format_gpu_uuid(getattr(torch.cuda.get_device_properties(index), "uuid", None))
+
+
 def gpu_identity(index: int) -> dict:
     """{index, name, uuid, total_bytes} of visible device ``index``."""
     import torch
@@ -282,6 +329,6 @@ def gpu_identity(index: int) -> dict:
     return {
         "index": index,
         "name": props.name,
-        "uuid": format_gpu_uuid(getattr(props, "uuid", None)),
+        "uuid": gpu_uuid(index),
         "total_bytes": int(props.total_memory),
     }

tests/engine/test_free_memory.py

@@ -0,0 +1,74 @@
+"""The free device memory the engine sizes its pools against.
+
+Under WDDM (Windows) cudaMemGetInfo leaves out what other processes hold, so a card the desktop uses
+reads as free and a pool sized to it pages into shared memory. NVML counts every process.
+"""
+
+from __future__ import annotations
+
+import sys
+from types import SimpleNamespace
+
+import torch
+
+from freetoken import gpu_select
+from freetoken.engine import graph
+
+GIB = 1 << 30
+DEVICE_UUID = "80472e39-ad2b-441d-fcff-0f3cc6ab9db2"
+
+
+class FakeNvml:
+    """NVML with one device and ``free`` bytes free on it; records the UUIDs it is asked about."""
+
+    def __init__(self, free: int):
+        self.free = free
+        self.asked: list[bytes] = []
+
+    def nvmlInit(self) -> int:
+        return 0
+
+    def nvmlShutdown(self) -> int:
+        return 0
+
+    def nvmlDeviceGetHandleByUUID(self, uuid: bytes, handle) -> int:
+        self.asked.append(uuid)
+        return 0
+
+    def nvmlDeviceGetMemoryInfo(self, handle, memory) -> int:
+        memory._obj.free = self.free
+        return 0
+
+
+def _machine(monkeypatch, *, cuda_free: int, nvml: FakeNvml | None, platform: str) -> None:
+    """Fakes CUDA's reading, the device's UUID, the NVML library the engine would load and the platform."""
+    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device=None: (cuda_free, 24 * GIB))
+    monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: SimpleNamespace(uuid=DEVICE_UUID))
+    monkeypatch.setattr(gpu_select, "_load_nvml", lambda: nvml)
+    monkeypatch.setattr(sys, "platform", platform)
+
+
+def test_windows_sizes_against_what_other_processes_leave(monkeypatch):
+    # the reading that spilled: CUDA says 23 GiB free while the desktop holds more than 2 GiB of 24
+    nvml = FakeNvml(free=21 * GIB)
+    _machine(monkeypatch, cuda_free=23 * GIB, nvml=nvml, platform="win32")
+    assert graph.get_free_memory(torch.device("cuda", 0)) == 21 * GIB
+    assert nvml.asked == [f"GPU-{DEVICE_UUID}".encode("ascii")]
+
+
+def test_windows_keeps_cudas_figure_when_it_is_the_smaller(monkeypatch):
+    _machine(monkeypatch, cuda_free=20 * GIB, nvml=FakeNvml(free=21 * GIB), platform="win32")
+    assert graph.get_free_memory(torch.device("cuda", 0)) == 20 * GIB
+
+
+def test_windows_without_nvml_keeps_cudas_figure(monkeypatch):
+    _machine(monkeypatch, cuda_free=23 * GIB, nvml=None, platform="win32")
+    assert graph.get_free_memory(torch.device("cuda", 0)) == 23 * GIB
+
+
+def test_linux_reads_cuda_alone(monkeypatch):
+    # cudaMemGetInfo already counts every process there, so NVML is not asked
+    nvml = FakeNvml(free=21 * GIB)
+    _machine(monkeypatch, cuda_free=23 * GIB, nvml=nvml, platform="linux")
+    assert graph.get_free_memory(torch.device("cuda", 0)) == 23 * GIB
+    assert nvml.asked == []