kernel: clamp zero-size host registrations; WDDM lock-pool hint on pin failure

python/freetoken/kernel/csrc/pinned_tensor.cpp

@@ -95,9 +95,14 @@ int64_t host_device_ptr(int64_t host_ptr) {
 }
 
 void host_register(int64_t addr, int64_t nbytes) {
-  const cudaError_t err =
-      cudaHostRegister(reinterpret_cast<void *>(addr), static_cast<size_t>(nbytes),
-                       cudaHostRegisterPortable | cudaHostRegisterMapped);
+  // Zero-size registrations are rejected by the driver with cudaErrorInvalidValue.
+  // Callers pin-after-fill small scale banks whose padded size can round to zero
+  // on dense layers; clamp so they register as a 1-byte mapped region instead.
+  const size_t reg_nbytes =
+      static_cast<size_t>(nbytes > 0 ? nbytes : 1);
+  const cudaError_t err = cudaHostRegister(
+      reinterpret_cast<void *>(addr), reg_nbytes,
+      cudaHostRegisterPortable | cudaHostRegisterMapped);
   TORCH_CHECK(err == cudaSuccess,
               "cudaHostRegister failed: ", cudaGetErrorString(err));
 }

python/freetoken/moe/host_banks.py

@@ -22,6 +22,7 @@
 import mmap
 import os
 import queue
+import sys
 import threading
 from concurrent.futures import ThreadPoolExecutor
 from enum import Enum
@@ -88,8 +89,41 @@ def pin(self) -> None:
         try:
             host_register(self.addr, len(self._buf))
         except RuntimeError as exc:
+            msg = str(exc).lower()
+            lock_words = {"lock", "locked", "locking", "memlock"}
+            if "out of memory" in msg or bool(lock_words & set(msg.split())) or "memlock" in msg:
+                if sys.platform == "win32":
+                    hint = (
+                        "the driver refused to lock more host memory (on Windows/WDDM "
+                        "the pageable-locking pool is roughly half of system RAM); "
+                        "close other processes, add RAM, or serve some layers with "
+                        "--moe-backend cpu"
+                    )
+                else:
+                    try:
+                        import resource
+
+                        soft = resource.getrlimit(resource.RLIMIT_MEMLOCK)[0]
+                        if soft == resource.RLIM_INFINITY:
+                            rlimit = "unlimited"
+                        elif soft >= 1073741824:
+                            rlimit = f"{soft / 1073741824:.0f} GB"
+                        else:
+                            rlimit = f"{soft // 1048576} MB"
+                    except Exception:
+                        rlimit = "?"
+                    hint = (
+                        "the kernel refused to lock more host memory: RLIMIT_MEMLOCK "
+                        f"is {rlimit} (often RAM/8 by default). Check `ulimit -l`, "
+                        "raise it via /etc/security/limits.d/*.conf and re-login -- "
+                        "an existing SSH ControlMaster session keeps serving the old "
+                        "limit until `ssh -O exit <host>`"
+                    )
+            else:
+                hint = "see the driver error above"
             raise RuntimeError(
-                f"cudaHostRegister failed for {len(self._buf) / 2**30:.1f} GiB"
+                f"cudaHostRegister failed for a {len(self._buf) / 2**30:.2f} GiB bank "
+                f"({msg}): {hint}"
             ) from exc
         self._pinned = True