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

Merge pull request #393 from ZacharyZcR/feat/auto-tune

Browse Source 
plan: auto-tune heuristics — MTP/PIPE/NUMA decisions from bottleneck classification
This commit is contained in:
Vincenzo Fornaro
2026-07-19 22:40:53 +02:00
committed by GitHub
parent a7ef058fc0 8a9a0fca4d
commit 17fbdfa8f8
2 changed files with 163 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/resource_plan.py
+88 -7
View File
@@ -219,6 +219,54 @@ def cpu_socket_count():
return 1
def _auto_tune(bottleneck_class, projected_hit, gpus, cpu_sockets, plan_has_metal):
"""Derive tuning knobs from the bottleneck classification."""
tune = {}
has_gpu = bool(gpus)
n_gpu = len(gpus)
# MTP: costs more than it saves when compute-bound (#389 measured 42% loss)
if bottleneck_class == "compute":
tune["DRAFT"] = {"value": "0",
"reason": "compute-bound: MTP batch overhead exceeds yield"}
elif bottleneck_class == "disk" and projected_hit < 0.90:
tune["DRAFT"] = {"value": "0",
"reason": "low hit rate: MTP widens expert union, adds disk reads"}
# otherwise leave DRAFT unset (engine default: auto)
# PIPE: resident pipeline mode depends on GPU count
if has_gpu and n_gpu == 1:
tune["COLI_CUDA_PIPE"] = {"value": "1",
"reason": "single GPU: S=1 pipeline gate"}
elif has_gpu and n_gpu > 1:
tune["COLI_CUDA_PIPE"] = {"value": "2",
"reason": "multi-GPU: residual stays on-device across layers"}
elif not has_gpu and bottleneck_class == "disk":
tune["PIPE"] = {"value": "1",
"reason": "overlap disk reads with resident expert compute"}
# NUMA: selective interleave for GPU hosts, blanket hint for CPU-only
if cpu_sockets > 1 and has_gpu:
tune["COLI_NUMA"] = {"value": "1",
"reason": "multi-socket + GPU: interleave expert slabs, protect DMA buffers"}
elif cpu_sockets > 1 and not has_gpu:
tune["COLI_NUMA"] = {"value": "1",
"reason": "multi-socket CPU-only: interleave expert slabs across nodes"}
tune["_numa_hint"] = "numactl --interleave=all may perform better on CPU-only hosts"
# OMP: kill hot-thread spin when GPU/Metal owns the power budget
if plan_has_metal:
tune["COLI_NO_OMP_TUNE"] = {"value": "1",
"reason": "Metal: OMP spin-wait steals GPU power budget"}
# PIN: fully resident if RAM allows and no GPU tier competes
if projected_hit >= 0.99 and not has_gpu:
tune["PIN_GB"] = {"value": "all",
"reason": "enough RAM for full expert residency"}
return tune
POLICIES = {
"quality": {"preserve_quantization": True, "preserve_router": True},
"balanced": {"preserve_quantization": True, "preserve_router": True},
@@ -290,12 +338,28 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
if cold_bytes:
warnings.append("cold expert misses may reach disk; normal decode speed depends on hit rate")
total_expert = info["expert_bytes"]
resident_expert = hot_bytes + warm_bytes
projected_hit = resident_expert / total_expert if total_expert else 1.0
if cold_bytes:
bottleneck = "disk expert misses"
elif warm_bytes:
bottleneck = "CPU expert compute and RAM bandwidth"
bottleneck_class = "disk"
elif warm_bytes and gpus:
bottleneck = "CPU expert tail and GPU compute"
bottleneck_class = "mixed"
elif projected_hit >= 0.99:
if gpus:
bottleneck = "GPU compute and interconnect"
else:
bottleneck = "CPU expert compute (fully resident)"
bottleneck_class = "compute"
else:
bottleneck = "GPU compute and interconnect"
bottleneck = "CPU expert compute and RAM bandwidth"
bottleneck_class = "memory"
tune = _auto_tune(bottleneck_class, projected_hit, gpus, cpu_sockets,
plan_has_metal=False)
return {
"version": 2,
@@ -317,6 +381,9 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
"expert_capacity": vram_experts, "requires_host_backing": False},
},
"expected_bottleneck": bottleneck,
"bottleneck_class": bottleneck_class,
"projected_hit_rate": round(projected_hit, 4),
"tune": tune,
"decisions": [
{"target": "VRAM", "reason": "profile-ranked hot experts"},
{"target": "RAM", "reason": "warm experts execute on CPU without quality loss"},
@@ -336,10 +403,11 @@ 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")
tune = plan.get("tune", {})
for key, entry in tune.items():
if key.startswith("_"):
continue
result.setdefault(key, entry["value"])
if plan["policy"]["name"] == "balanced":
result.setdefault("REPIN", "64")
ram = plan["tiers"]["ram"]
@@ -386,5 +454,18 @@ def format_plan(plan):
else:
lines.append("VRAM no NVIDIA device detected · CPU path")
lines.append(f"limit {plan['expected_bottleneck']}")
hit = plan.get("projected_hit_rate", 0)
lines.append(f"hit {hit:.0%} projected expert residency")
tune = plan.get("tune", {})
if tune:
lines.append("")
lines.append("auto-tune:")
for key, entry in tune.items():
if key.startswith("_"):
continue
lines.append(f" {key}={entry['value']:12s} {entry['reason']}")
hint = tune.get("_numa_hint")
if hint:
lines.append(f" hint: {hint}")
lines.extend(f"warn {warning}" for warning in plan["warnings"])
return "\n".join(lines)
c/tests/test_resource_plan.py
+75
View File
@@ -140,6 +140,81 @@ class ResourcePlanTest(unittest.TestCase):
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_auto_tune_mtp_off_when_compute_bound(self):
# Tiny model with 64 GB RAM and no GPU: all experts fit in RAM with no
# warm tier, so the plan classifies as compute-bound.
plan = build_plan(self.model, ram_gb=64, available_memory=64 * GB,
available_disk=100 * GB, gpus=[], physical_cpus=24,
cpu_sockets=2)
# With such a small model fully in RAM and no GPU, bottleneck is compute
self.assertEqual(plan["bottleneck_class"], "compute")
self.assertIn("DRAFT", plan["tune"])
self.assertEqual(plan["tune"]["DRAFT"]["value"], "0")
env = environment_for_plan(plan)
self.assertEqual(env["DRAFT"], "0")
explicit = environment_for_plan(plan, {"DRAFT": "3"})
self.assertEqual(explicit["DRAFT"], "3")
def test_auto_tune_mtp_off_when_disk_low_hit(self):
# Use a model large enough that 8 GB RAM can't hold all experts.
big = tempfile.TemporaryDirectory()
bigmodel = Path(big.name)
(bigmodel / "config.json").write_text(json.dumps({
"num_hidden_layers": 2, "n_routed_experts": 4,
"kv_lora_rank": 4, "qk_rope_head_dim": 2,
"qk_nope_head_dim": 3, "v_head_dim": 5, "num_attention_heads": 2,
}))
expert_size = 3 * GB # each expert 3 GB → 12 GB total, won't fit in 8 GB budget
write_shard(bigmodel / "out-00000.safetensors", [
("model.embed_tokens.weight", 100),
("model.layers.0.self_attn.q_a_proj.weight", 200),
])
for i in range(4):
write_shard(bigmodel / f"out-{i+1:05d}.safetensors", [
(f"model.layers.1.mlp.experts.{i}.gate_proj.weight", expert_size),
])
plan = build_plan(bigmodel, ram_gb=0, available_memory=4 * GB,
available_disk=100 * GB, gpus=[], physical_cpus=8,
cpu_sockets=1)
big.cleanup()
self.assertEqual(plan["bottleneck_class"], "disk")
self.assertLess(plan["projected_hit_rate"], 0.90)
self.assertEqual(plan["tune"]["DRAFT"]["value"], "0")
def test_auto_tune_pipe_multi_gpu(self):
gpus = [
{"index": 0, "name": "a", "total_bytes": 32 * GB, "free_bytes": 30 * GB},
{"index": 1, "name": "b", "total_bytes": 32 * GB, "free_bytes": 30 * GB},
]
plan = build_plan(self.model, ram_gb=16, available_memory=32 * GB,
available_disk=1, gpus=gpus, cpu_sockets=2)
self.assertEqual(plan["tune"]["COLI_CUDA_PIPE"]["value"], "2")
env = environment_for_plan(plan)
self.assertEqual(env["COLI_CUDA_PIPE"], "2")
def test_auto_tune_pipe_single_gpu(self):
gpus = [{"index": 0, "name": "a", "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, cpu_sockets=1)
self.assertEqual(plan["tune"]["COLI_CUDA_PIPE"]["value"], "1")
def test_auto_tune_numa_hint_for_cpu_only(self):
plan = build_plan(self.model, ram_gb=64, available_memory=64 * GB,
available_disk=1, gpus=[], physical_cpus=64, cpu_sockets=2)
self.assertIn("_numa_hint", plan["tune"])
self.assertIn("numactl", plan["tune"]["_numa_hint"])
self.assertIn("auto-tune", format_plan(plan))
def test_format_plan_shows_tune_and_hit_rate(self):
plan = build_plan(self.model, ram_gb=64, available_memory=64 * GB,
available_disk=100 * GB, gpus=[], physical_cpus=24,
cpu_sockets=1)
text = format_plan(plan)
self.assertIn("hit", text)
self.assertIn("auto-tune", text)
self.assertIn("DRAFT", text)
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,