plan NUMA interleave on multi-socket Linux

c/resource_plan.py

@@ -203,6 +203,22 @@ def physical_cpu_count():
     return os.cpu_count() or 1
 
 
+def cpu_socket_count():
+    """Return the number of physical CPU sockets visible to this process."""
+    if not sys.platform.startswith("linux"):
+        return 1
+    try:
+        result = subprocess.run(["lscpu", "-p=socket"], text=True,
+                                capture_output=True, check=True, timeout=5)
+        sockets = {int(line) for line in result.stdout.splitlines()
+                   if line and not line.startswith("#")}
+        if sockets:
+            return len(sockets)
+    except (OSError, ValueError, subprocess.SubprocessError):
+        pass
+    return 1
+
+
 POLICIES = {
     "quality": {"preserve_quantization": True, "preserve_router": True},
     "balanced": {"preserve_quantization": True, "preserve_router": True},
@@ -212,11 +228,12 @@ POLICIES = {
 
 def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
                available_memory=None, available_disk=None, gpus=None,
-               policy="quality", physical_cpus=None):
+               policy="quality", physical_cpus=None, cpu_sockets=None):
     if policy not in POLICIES:
         raise ValueError(f"unknown policy: {policy}")
     info = analyze_model(model)
     physical_cpus = physical_cpu_count() if physical_cpus is None else physical_cpus
+    cpu_sockets = cpu_socket_count() if cpu_sockets is None else cpu_sockets
     cfg = info["config"]
     available_memory = memory_available() if available_memory is None else available_memory
     if available_disk is None:
@@ -286,6 +303,7 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
                    "quality_preserving": policy != "experimental-fast"},
         "model": {key: value for key, value in info.items() if key != "config"},
         "cpu": {"physical_cores": max(1, int(physical_cpus)),
+                "sockets": max(1, int(cpu_sockets)),
                 "thread_policy": "physical-cores"},
         "tiers": {
             "disk": {"role": "cold-backing", "model_bytes": info["model_bytes"],
@@ -318,6 +336,10 @@ def environment_for_plan(plan, env=None, cuda_enabled=True):
         # ("Affinity not supported on this configuration"): non impostarle li'.
         result.setdefault("OMP_PROC_BIND", "spread")
         result.setdefault("OMP_PLACES", "cores")
+    if sys.platform.startswith("linux") and plan["cpu"].get("sockets", 1) > 1:
+        # Selectively interleave large expert/dense slabs across memory controllers.
+        # Unlike blanket numactl interleave, this leaves CUDA staging buffers local.
+        result.setdefault("COLI_NUMA", "1")
     if plan["policy"]["name"] == "balanced":
         result.setdefault("REPIN", "64")
     ram = plan["tiers"]["ram"]

c/tests/test_resource_plan.py

@@ -10,6 +10,7 @@ from resource_plan import (
     GB,
     analyze_model,
     build_plan,
+    cpu_socket_count,
     environment_for_plan,
     format_plan,
     memory_available,
@@ -66,15 +67,19 @@ class ResourcePlanTest(unittest.TestCase):
         # 0 slots/layer. The value must be a sane positive number of bytes.
         self.assertGreater(memory_available(), 0)
 
+    def test_cpu_socket_count_is_positive(self):
+        self.assertGreaterEqual(cpu_socket_count(), 1)
+
     def test_builds_bounded_three_tier_plan(self):
         gpus = [{"index": 0, "name": "test-gpu", "total_bytes": 12 * GB,
                  "free_bytes": 10 * GB}]
         plan = build_plan(self.model, ram_gb=16, context=32, vram_gb=20,
                           available_memory=32 * GB, available_disk=100 * GB, gpus=gpus,
-                          physical_cpus=24)
+                          physical_cpus=24, cpu_sockets=2)
         self.assertEqual(plan["version"], 2)
         self.assertEqual(plan["policy"]["name"], "quality")
         self.assertEqual(plan["cpu"]["physical_cores"], 24)
+        self.assertEqual(plan["cpu"]["sockets"], 2)
         self.assertTrue(plan["policy"]["preserve_quantization"])
         self.assertFalse(plan["tiers"]["vram"]["requires_host_backing"])
         self.assertEqual(plan["tiers"]["ram"]["budget_bytes"], 16 * GB)
@@ -105,7 +110,7 @@ class ResourcePlanTest(unittest.TestCase):
             {"index": 1, "name": "b", "total_bytes": 12 * GB, "free_bytes": 10 * GB},
         ]
         plan = build_plan(self.model, ram_gb=16, available_memory=32 * GB,
-                          available_disk=1, gpus=gpus)
+                          available_disk=1, gpus=gpus, cpu_sockets=2)
         env = environment_for_plan(plan, {"RAM_GB": "12", "PIN": "stats.txt",
                                                "COLI_GPUS": "1"})
         self.assertEqual(env["RAM_GB"], "12")
@@ -125,6 +130,16 @@ class ResourcePlanTest(unittest.TestCase):
                                                         "OMP_PROC_BIND": "close"})
         self.assertEqual(explicit_threads["OMP_NUM_THREADS"], "7")
         self.assertEqual(explicit_threads["OMP_PROC_BIND"], "close")
+
+        if sys.platform.startswith("linux"):
+            self.assertEqual(env["COLI_NUMA"], "1")
+            explicit_numa = environment_for_plan(plan, {"COLI_NUMA": "0"})
+            self.assertEqual(explicit_numa["COLI_NUMA"], "0")
+
+    def test_single_socket_plan_does_not_enable_numa(self):
+        plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
+                          gpus=[], physical_cpus=8, cpu_sockets=1)
+        self.assertNotIn("COLI_NUMA", environment_for_plan(plan))
     def test_cpu_binary_does_not_apply_gpu_tier(self):
         plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
                           gpus=[{"index": 0, "name": "a", "total_bytes": 8 * GB,

docs/ENVIRONMENT.md

@@ -43,7 +43,7 @@ Format: `VAR` — default — effect.
 | `URING` | `0` (off) | Linux-only queued expert I/O. `URING=1` implies `PIPE=1`, forces cold reads through io-wq (`IOSQE_ASYNC`), replaces blocking loader pthreads and spin waits with batched SQEs/CQEs, and batches `PILOT_REAL` loads on a separate ring. Use `DIRECT=1` for cold NVMe to avoid page-cache copy/readahead limits. Fails clearly if the kernel denies io_uring; incompatible with `COLI_MMAP=1`. |
 | `DIRECT` | `0` (off) | Use `O_DIRECT`/unbuffered reads for expert slabs. Helps sustained NVMe; keeps the zero-copy GPU path. |
 | `COLI_NO_OMP_TUNE` | off | **Kill-switch** for the OpenMP hot-thread tuning (`OMP_WAIT_POLICY=active` spin + proc-bind). Set `=1` when the CPU is mostly waiting on the GPU (Metal) so spin doesn't steal the shared power budget. |
-| `COLI_NUMA` | off (Linux only) | `COLI_NUMA=1` interleaves expert slabs across NUMA nodes via `mbind` (raw syscall, no libnuma). Helps multi-socket hosts (+7–40% expert matmul); silent no-op on single-node or non-Linux. |
+| `COLI_NUMA` | auto in generated plans on multi-socket Linux; otherwise off | `COLI_NUMA=1` selectively interleaves large expert and dense slabs across NUMA nodes via `mbind` (raw syscall, no libnuma). Helps multi-socket hosts (+7–40% expert matmul); silent no-op on single-node or non-Linux. Explicit `COLI_NUMA=0` overrides the generated plan. |
 | `MLOCK` | `-1` (auto: on for macOS) | Wire the streamed expert cache into physical RAM (`mlock`) to dodge the memory compressor. `0` off, `1` force. |
 | `CAP_RAISE` | `1` (on) | Let the engine raise the expert-cache cap above `topk` when RAM allows (bigger batches). `0` fixes the cap. |
 | `PREFETCH` | `0` | Prefetch depth for streamed experts. |

docs/benchmarks.md

@@ -116,8 +116,12 @@ can match an RTX 5090 on expert matmul ([#101](https://github.com/JustVugg/colib
 so **the GPU tier earns its VRAM only when the CPU is the weak link**. On
 multi-socket hosts, NUMA placement is a further lever: interleaving the resident
 weights across nodes measured **+13% (2-socket) and +40% (4-socket CPU-only)**
-([#82](https://github.com/JustVugg/colibri/issues/82)) — but never blanket-interleave
+([#82](https://github.com/JustVugg/colibri/issues/82)). On a 2-socket Xeon Silver
-a GPU host (measured 10× regression via the DMA staging pages).
+4510 host with 6× RTX 5090, selective `COLI_NUMA=1` raised effective CPU-expert
+bandwidth from **42.42 to 58.26/65.89 GB/s** and greedy decode from **7.66 to
+9.02/9.17 tok/s** (64 tokens, `TEMP=0 DRAFT=0`, byte-identical output). Do not
+blanket-interleave a GPU host: it also spreads DMA staging pages and has measured
+up to a 10× regression; generated plans enable only the selective slab policy.
 
 ## Quality benchmark