fix(llm): size VRAM headroom from reclaimable memory on integrated GPUs

src/nemo_safe_synthesizer/llm/utils.py

@@ -436,12 +436,35 @@ class _VRAMAllocation:
     memory_bytes: int
 
 
+def _reclaimable_available_bytes() -> int | None:
+    """Return kernel-reclaimable system memory (``MemAvailable``) in bytes.
+
+    ``MemAvailable`` in ``/proc/meminfo`` accounts for reclaimable page cache, so
+    it reflects what the kernel can actually hand out. Returns ``None`` when it
+    cannot be read (for example on non-Linux hosts, where ``/proc/meminfo`` is
+    absent).
+    """
+    try:
+        with open("/proc/meminfo", encoding="utf-8") as meminfo:
+            for line in meminfo:
+                if line.startswith("MemAvailable:"):
+                    return int(line.split()[1]) * 1024  # reported in kibibytes
+    except (OSError, ValueError, IndexError):
+        return None
+    return None
+
+
 def _get_vram_allocations(max_vram_fraction: float | None = None) -> dict[int, _VRAMAllocation]:
     """Calculate maximum memory allocation for each available GPU.
 
     Reserves a 2 GiB safety buffer on each device, then applies
     ``max_vram_fraction`` to the remaining free memory.
 
+    On systems with integrated GPUs (e.g. DGX Spark) GPU memory *is* system memory,
+    so ``torch.cuda.mem_get_info`` reports only unallocated pages and ignores the
+    reclaimable page cache -- badly under-reporting real headroom. There we use
+    the kernel's ``MemAvailable`` (capped at device total) instead.
+
     Args:
         max_vram_fraction: Fraction of total GPU memory to allocate.
             Defaults to ``0.8`` (80 %).
@@ -457,8 +480,12 @@ def _get_vram_allocations(max_vram_fraction: float | None = None) -> dict[int, _
 
     if torch.cuda.is_available():
         num_gpus = torch.cuda.device_count()
+        # System-wide value; read once rather than per device.
+        reclaimable = _reclaimable_available_bytes()
         for i in range(num_gpus):
             free, total = torch.cuda.mem_get_info(device=i)
+            if getattr(torch.cuda.get_device_properties(i), "is_integrated", False) and reclaimable is not None:
+                free = min(max(free, reclaimable), total)
             safe_free = max(free - (2 * 1024**3), 0)
             gpu_memory_utilization = min(max_vram_fraction, safe_free / total) if total > 0 else 0.0
             memory_bytes = int(gpu_memory_utilization * total)

tests/llm/test_utils.py

@@ -4,12 +4,13 @@
 """Unit tests for llm.utils helpers."""
 
 from pathlib import Path
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, mock_open, patch
 
 import pytest
 
 from nemo_safe_synthesizer.llm.utils import (
     ModelRef,
+    _reclaimable_available_bytes,
     get_max_memory_map,
     get_max_vram,
     get_quantization_config,
@@ -73,15 +74,74 @@ def test_trust_remote_code_for_model_requires_configured_cache_root(tmp_path: Pa
 
 def test_vram_helpers_return_fraction_and_hf_memory_map() -> None:
     gib = 1024**3
+    discrete_gpu = MagicMock(is_integrated=False)
     with (
         patch("torch.cuda.is_available", return_value=True),
         patch("torch.cuda.device_count", return_value=1),
         patch("torch.cuda.mem_get_info", return_value=(10 * gib, 16 * gib)),
+        patch("torch.cuda.get_device_properties", return_value=discrete_gpu),
     ):
         assert get_max_vram(max_vram_fraction=0.8) == {0: 0.5}
         assert get_max_memory_map(max_vram_fraction=0.8) == {0: 8 * gib}
 
 
+def test_vram_helpers_use_reclaimable_memory_on_integrated_gpus() -> None:
+    gib = 1024**3
+    integrated_gpu = MagicMock(is_integrated=True)
+    with (
+        patch("torch.cuda.is_available", return_value=True),
+        patch("torch.cuda.device_count", return_value=1),
+        # Unified memory: mem_get_info under-reports free (2 GiB) ...
+        patch("torch.cuda.mem_get_info", return_value=(2 * gib, 120 * gib)),
+        patch("torch.cuda.get_device_properties", return_value=integrated_gpu),
+        # ... but 60 GiB is reclaimable/available per the kernel.
+        patch("nemo_safe_synthesizer.llm.utils._reclaimable_available_bytes", return_value=60 * gib),
+    ):
+        # free -> 60 GiB, safe_free -> 58 GiB, utilization -> 58/120.
+        assert get_max_vram(max_vram_fraction=0.8) == {0: pytest.approx(58 / 120)}
+        assert get_max_memory_map(max_vram_fraction=0.8) == {0: 58 * gib}
+
+
+def test_vram_helpers_cap_reclaimable_memory_at_device_total() -> None:
+    gib = 1024**3
+    integrated_gpu = MagicMock(is_integrated=True)
+    with (
+        patch("torch.cuda.is_available", return_value=True),
+        patch("torch.cuda.device_count", return_value=1),
+        patch("torch.cuda.mem_get_info", return_value=(2 * gib, 120 * gib)),
+        patch("torch.cuda.get_device_properties", return_value=integrated_gpu),
+        # MemAvailable exceeds device total; must be capped so the 2 GiB buffer survives.
+        patch("nemo_safe_synthesizer.llm.utils._reclaimable_available_bytes", return_value=200 * gib),
+    ):
+        # free -> 120 GiB (capped), safe_free -> 118 GiB, utilization -> 118/120.
+        assert get_max_vram(max_vram_fraction=1.0) == {0: pytest.approx(118 / 120)}
+        assert get_max_memory_map(max_vram_fraction=1.0) == {0: 118 * gib}
+
+
+def test_vram_helpers_fall_back_to_cuda_free_when_meminfo_unreadable() -> None:
+    gib = 1024**3
+    integrated_gpu = MagicMock(is_integrated=True)
+    with (
+        patch("torch.cuda.is_available", return_value=True),
+        patch("torch.cuda.device_count", return_value=1),
+        patch("torch.cuda.mem_get_info", return_value=(10 * gib, 16 * gib)),
+        patch("torch.cuda.get_device_properties", return_value=integrated_gpu),
+        patch("nemo_safe_synthesizer.llm.utils._reclaimable_available_bytes", return_value=None),
+    ):
+        assert get_max_vram(max_vram_fraction=0.8) == {0: 0.5}
+
+
+def test_reclaimable_available_bytes_parses_meminfo() -> None:
+    meminfo = "MemTotal:       127603160 kB\nMemFree:             204 kB\nMemAvailable:   68157440 kB\n"
+    with patch("builtins.open", mock_open(read_data=meminfo)):
+        assert _reclaimable_available_bytes() == 68157440 * 1024
+
+
+def test_reclaimable_available_bytes_returns_none_when_absent() -> None:
+    with patch("builtins.open", side_effect=FileNotFoundError):
+        assert _reclaimable_available_bytes() is None
+
+
 def test_model_ref_trusts_snapshot_under_configured_cache_root(tmp_path: Path, hf_cached_snapshot_factory) -> None:
     """HF cache path recognition intentionally tracks Hub snapshot metadata."""
     cache_root = tmp_path / "custom-cache"