Improve unified memory usage on the RTX Spark

toolkit/util/unified_memory.py

@@ -1,13 +1,19 @@
-"""Unified memory (integrated GPU) support, e.g. NVIDIA DGX Spark (GB10).
+"""Unified memory (integrated GPU) support, e.g. NVIDIA DGX Spark (GB10) and
+the Windows RTX Spark (N1X).
 
 On these devices the CPU and GPU share the same physical memory, so moving a
 module to the CPU frees nothing. Worse, device -> host copies are extremely
-slow (~100MB/s on GB10), so every "offload to cpu" of a large transformer or
-text encoder stalls for minutes. When a unified memory device is detected we
-make module level moves to the CPU a no-op (dtype casts still apply). Tensor
-level .cpu()/.to('cpu') is left alone since saving, numpy, etc. need it.
+slow (~100MB/s on GB10 under Linux, ~0.5-2GB/s pageable on Windows/WDDM), so
+every "offload to cpu" of a large transformer or text encoder stalls. When a
+unified memory device is detected we make module level moves to the CPU a
+no-op (dtype casts still apply). Tensor level .cpu()/.to('cpu') is left alone
+since saving, numpy, etc. need it.
 
-Set AITK_DISABLE_UNIFIED_MEMORY=1 to turn this off.
+Note for Windows: WDDM gives CUDA a fixed budget (~80% of system RAM, e.g.
+~45GB of 54GB) that host allocations do not count against. nvidia-smi / Task
+Manager only show the small dedicated segment (8GB); the real pool is what
+torch.cuda.mem_get_info() reports. If a job needs to keep modules on the CPU
+to stay under that budget, set AITK_DISABLE_UNIFIED_MEMORY=1 to turn this off.
 """
 import os
 import sys
@@ -20,8 +26,8 @@ _patched = False
 def is_unified_memory() -> bool:
     if os.environ.get("AITK_DISABLE_UNIFIED_MEMORY", "0") == "1":
         return False
-    # linux only (DGX Spark). The Windows RTX Spark stack does not behave the same way
-    if not sys.platform.startswith("linux"):
+    # macOS has no CUDA; MPS unified memory is handled elsewhere
+    if sys.platform == "darwin":
         return False
     # NVIDIA only. ROCm APUs report integrated too, but may have a dedicated VRAM carveout
     if torch.version.hip is not None:

ui/src/server/monitor.ts

@@ -36,11 +36,20 @@ const NV_BATCH_FLUSH_MS = 100;
 const CPU_TEMP_REFRESH_MS = 5000;
 
 // Power limits for GPUs whose driver reports none (nvidia-smi gives [N/A]).
-// GB10 (DGX Spark) is a unified SoC with no settable GPU limit; 140 W is
-// NVIDIA's rated TDP for the GB10 superchip.
-const KNOWN_POWER_LIMITS_W: Record<string, number> = {
-  'NVIDIA GB10': 140,
-};
+// These unified SoCs have no settable GPU limit, so show the rated TDP of the
+// whole chip. GB10 (DGX Spark) is 140 W. RTX Spark N1X: the 5120-core part is
+// laptop-only, rated 45-80 W (OEM configurable); the 6144-core part is 140 W
+// on desktop and 45-80 W in laptops but shares its name, so it falls through
+// to the desktop figure.
+const KNOWN_POWER_LIMITS_W: [RegExp, number][] = [
+  [/\bGB10\b/, 140],
+  [/RTX Spark.*5120-core/, 80],
+  [/RTX Spark/, 140],
+];
+
+function knownPowerLimit(name: string): number {
+  return KNOWN_POWER_LIMITS_W.find(([re]) => re.test(name))?.[1] ?? NaN;
+}
 
 function parseGpuLine(line: string): GpuInfo | null {
   const [
@@ -76,7 +85,7 @@ function parseGpuLine(line: string): GpuInfo | null {
     },
     power: {
       draw: parseFloat(powerDraw),
-      limit: isNaN(parseFloat(powerLimit)) ? (KNOWN_POWER_LIMITS_W[name] ?? NaN) : parseFloat(powerLimit),
+      limit: isNaN(parseFloat(powerLimit)) ? knownPowerLimit(name) : parseFloat(powerLimit),
     },
     clocks: {
       graphics: parseInt(clockGraphics),
@@ -88,19 +97,28 @@ function parseGpuLine(line: string): GpuInfo | null {
   };
 }
 
+// Integrated NVIDIA GPUs that share system RAM with the CPU.
+const UNIFIED_MEMORY_GPU = /\b(GB10|RTX Spark)\b/i;
+
 /**
- * Unified-memory GPUs (GB10 / DGX Spark) have no dedicated VRAM, so
- * nvidia-smi reports memory as [N/A]. Report the system RAM pool they share
- * instead, like the Mac path does for Apple silicon.
+ * Unified-memory GPUs (GB10 / DGX Spark, RTX Spark) have no dedicated VRAM.
+ * On Linux nvidia-smi reports their memory as [N/A]; on Windows (WDDM) it
+ * reports only the small dedicated segment (8 GB) even though CUDA can use
+ * most of system RAM. Report the system RAM pool they share instead, like
+ * the Mac path does for Apple silicon.
  */
+function isSharedMemoryGpu(g: GpuInfo): boolean {
+  return isNaN(g.memory.total) || UNIFIED_MEMORY_GPU.test(g.name);
+}
+
 function withSharedMemory(gpu: GPUApiResponse, cpu: CpuInfo | null): GPUApiResponse {
-  if (!cpu || !gpu.gpus.some(g => isNaN(g.memory.total))) return gpu;
+  if (!cpu || !gpu.gpus.some(isSharedMemoryGpu)) return gpu;
   const total = Math.round(cpu.totalMemory);
   const free = Math.round(cpu.availableMemory);
   return {
     ...gpu,
     gpus: gpu.gpus.map(g =>
-      isNaN(g.memory.total) ? { ...g, memory: { total, free, used: total - free, shared: true } } : g,
+      isSharedMemoryGpu(g) ? { ...g, memory: { total, free, used: total - free, shared: true } } : g,
     ),
   };
 }