url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

plan NUMA interleave on multi-socket Linux

Browse Source
This commit is contained in:
ZacharyZcR
2026-07-18 18:38:40 +08:00
parent f8e4dc95b5
commit 9c17dc6d29
4 changed files with 47 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/resource_plan.py
+23 -1
View File
@@ -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
+17 -2
View File
@@ -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
+1 -1
View File
@@ -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
+6 -2
View File
@@ -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