From 13a9782969e7b4a56172bb8bb0c58f66ee5bda8d Mon Sep 17 00:00:00 2001 From: kitty <2165990891@qq.com> Date: Thu, 27 Aug 2026 11:45:49 +0800 Subject: [PATCH] fix(engine): probe the real WSL pin budget instead of guessing 40% RAM _pin_budget_bytes() returned 40% of physical RAM (25.1 GiB on the test box, ~25x above the real cumulative CUDA pin ceiling), so split residency never engaged and loads died in cudaHostRegister. Now on WSL2 it measures the ceiling once (lru_cache): cudaHostAlloc in 256 MiB chunks until the driver refuses, faults the pages to match pin-after-fill banks, returns 80% of the measured total for headroom, and frees the probe buffers. Plain Linux (no "microsoft" kernel tag) stays uncapped (None); FREETOKEN_PIN_BUDGET_GB still overrides. The refused alloc leaves a *sticky* CUDA error every later call replays as "out of memory"; cudaGetLastError() clears it before return. Experiment env: WSL2 kernel 6.18.33.2-microsoft-standard-WSL2, CUDA 13.0 (torch, libcudart.so.13), RTX 3090 24 GiB, 62.8 GiB RAM. Probed wall 0.75-1.0 GiB cumulative (256 MiB chunks); host_ptr_identity (UVA) False. --- python/freetoken/engine/engine.py | 42 ++++++++++++++++++++++++++++--- 1 file changed, 38 insertions(+), 4 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cd6505d2d..7e33d5b3d 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1,5 +1,6 @@ from __future__ import annotations +import functools import gc import math import os @@ -1158,15 +1159,48 @@ def _cpu_moe_executor_viable(model_config) -> bool: return fmt == "mxfp4" or fmt in _WFMT_IDS -def _pin_budget_bytes() -> int | None: - """Bytes this process can safely cudaHostRegister, or None when the platform does not cap pinning (plain Linux). +@functools.lru_cache(maxsize=1) +def _probe_wsl_pin_budget() -> int | None: + """Measured cumulative CUDA pin ceiling; None if unavailable.""" + try: + import ctypes + if not torch.cuda.is_available(): + return None + torch.zeros(1, device="cuda") # warm the CUDA context + rt = ctypes.CDLL(f"libcudart.so.{(torch.version.cuda or '0').split('.')[0]}") + for fn in ("cudaHostAlloc", "cudaFreeHost", "cudaGetLastError"): + getattr(rt, fn).restype = ctypes.c_int + rt.cudaHostAlloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t, ctypes.c_uint] + rt.cudaFreeHost.argtypes = [ctypes.c_void_p] + except Exception: + return None + chunk = 256 << 20 # small chunks get the most usable budget + max_probe = 8 << 30 # bound probe time/RAM on uncapped hosts + held: list[ctypes.c_void_p] = [] + total = 0 + try: + while total + chunk <= max_probe: + ptr = ctypes.c_void_p() + if rt.cudaHostAlloc(ctypes.byref(ptr), chunk, 0) != 0 or not ptr.value: + break # hit the pin wall + ctypes.memset(ptr, 0, 1 << 20) # fault pages, matching pin-after-fill banks + held.append(ptr) + total += chunk + finally: + for ptr in held: # free the probe buffers so the real banks get the budget + rt.cudaFreeHost(ptr) + rt.cudaGetLastError() # clear sticky error from the refused alloc or torch OOMs + return int(total * 0.8) if total else None + - WSL's WDDM-backed CUDA caps pinning near half of RAM, shared across processes -- budget 40%. FREETOKEN_PIN_BUDGET_GB overrides anywhere.""" +def _pin_budget_bytes() -> int | None: + """Bytes safe to cudaHostRegister, or None when the platform does not cap pinning (plain Linux). + WSL/WDDM caps near ~1 GiB; FREETOKEN_PIN_BUDGET_GB overrides anywhere.""" if env := os.environ.get("FREETOKEN_PIN_BUDGET_GB"): return int(float(env) * 2**30) if not hasattr(os, "uname") or "microsoft" not in os.uname().release.lower(): # WSL kernel tag return None - return int(os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") * 0.4) + return _probe_wsl_pin_budget() def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int) -> frozenset[int]: