From 9c17dc6d292095cf56ca2c4cc92712e65a0e1e3a Mon Sep 17 00:00:00 2001 From: ZacharyZcR Date: Sat, 18 Jul 2026 18:38:40 +0800 Subject: [PATCH] plan NUMA interleave on multi-socket Linux --- c/resource_plan.py | 24 +++++++++++++++++++++++- c/tests/test_resource_plan.py | 19 +++++++++++++++++-- docs/ENVIRONMENT.md | 2 +- docs/benchmarks.md | 8 ++++++-- 4 files changed, 47 insertions(+), 6 deletions(-) diff --git a/c/resource_plan.py b/c/resource_plan.py index 6b3f2fe..5f17a6a 100644 --- a/c/resource_plan.py +++ b/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"] diff --git a/c/tests/test_resource_plan.py b/c/tests/test_resource_plan.py index 023cb93..b184c80 100644 --- a/c/tests/test_resource_plan.py +++ b/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, diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index 535288a..9f0d749 100644 --- a/docs/ENVIRONMENT.md +++ b/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. | diff --git a/docs/benchmarks.md b/docs/benchmarks.md index f9c83f2..09d3752 100644 --- a/docs/benchmarks.md +++ b/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 -a GPU host (measured 10× regression via the DMA staging pages). +([#82](https://github.com/JustVugg/colibri/issues/82)). On a 2-socket Xeon Silver +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