feat: distributed expert workers — rebased, fmt=6 rotation fixed, token-exact parity gate

.github/workflows/ci.yml

@@ -402,6 +402,42 @@ jobs:
             tests.test_inefficiency.TinyEfficiencyTest.test_profile_phases_present_and_nonneg \
             tests.test_inefficiency.TinyEfficiencyTest.test_disk_wait_not_dominant
 
+  # The token-exact parity gate from #7 (c/tests/test_cluster_sharding.py): the
+  # engine's local CPU run must already reproduce the transformers oracle
+  # token-exactly, and delegating the routed-expert shard to a cluster worker
+  # must not perturb a single position. It exercises the fmt=6 (E8/IQ3,
+  # rotation-bearing) path and its fmt=4 (grouped int4, no-rotation) control.
+  #
+  # Not in `make check`: the two fixtures are gitignored (regenerated, not
+  # committed -- see .gitignore) and need the pinned torch/transformers to
+  # build, the same rationale as inkling-oracle and efficiency above. The test
+  # skips cleanly when the fixtures are absent, so leaving fixture generation
+  # out of CI would turn the gate into a silent no-op on a clean host -- the
+  # worker path would never actually run. This job regenerates both fixtures
+  # first so the gate is exercised, not dormant.
+  #
+  # Socket safety: the test spawns a cluster worker subprocess and a
+  # coordinator, each binding a 127.0.0.1 port obtained from _free_port()
+  # (bind to port 0), so it is safe on a shared runner.
+  cluster-parity:
+    name: Cluster parity gate (fmt6/fmt4, token-exact)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+          cache: pip
+          cache-dependency-path: c/tools/oracle-requirements.txt
+      - name: Install torch (CPU) + transformers
+        run: pip install -r c/tools/oracle-requirements.txt
+      - name: Build colibri
+        run: make -C c colibri
+      - name: Generate the fmt6/fmt4 quantized fixtures
+        run: cd c && python3 tools/make_glm_oracle.py --fmt6 && python3 tools/make_glm_oracle.py --fmt4
+      - name: Cluster parity gate (token-exact)
+        run: cd c && python3 -m unittest -v tests.test_cluster_sharding
+
   engine-hip-syntax:
     name: HIP syntax check
     runs-on: ubuntu-latest

.gitignore

@@ -82,6 +82,8 @@ c/glm_tiny/
 c/glm_tiny_i2/
 c/glm_tiny_i4/
 c/glm_tiny_mix/
+c/glm_tiny_fmt6/
+c/glm_tiny_fmt4/
 c/glm_bench_medium/
 c/bench/
 

README.md

@@ -162,8 +162,40 @@ the right experts get. It works because routing has measurable structure (see
 the [expert atlas](https://github.com/JustVugg/colibri/issues/175)) — and
 structure is cacheable.
 
-The engine is a single C file (`c/glm.c`) plus small headers. No BLAS, no Python
-at runtime, no GPU required.
+The engine is a single C file (`c/colibri.c`) plus small headers. No BLAS, no
+Python at runtime, no GPU required.
+
+### Local cluster mode
+
+The coordinator keeps token generation, routing, and KV state local while
+disk-backed expert workers execute routed FFNs on other Macs. A layer's routed
+batch-union is sent as one persistent TCP request, so a token does not incur one
+round trip per expert.
+
+Start the optional registration service:
+
+```bash
+./coli cluster coordinator --host 0.0.0.0 --port 8765
+```
+
+On each worker, with the same converted model available locally:
+
+```bash
+./coli cluster worker --model /nvme/glm52_i4 --port 9100 \
+  --coordinator http://COORDINATOR:8765 --advertise-host WORKER_IP
+```
+
+Run the coordinator with discovery, or provide `--cluster-workers
+HOST:PORT,...` for a static setup:
+
+```bash
+./coli serve --model /nvme/glm52_i4 \
+  --cluster-coordinator http://127.0.0.1:8765
+```
+
+The transport is disabled unless workers are configured, so the existing
+single-machine path remains unchanged. Dense-layer sharding and browser/WebGPU
+workers are separate follow-up seams.
 
 ## How it works
 

c/Makefile

@@ -894,6 +894,9 @@ iobench$(EXE): iobench.c compat.h
 tests/test_serve_sentinel$(EXE): tests/test_serve_sentinel.c compat.h
 	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
 
+tests/test_cluster_protocol$(EXE): tests/test_cluster_protocol.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h route_trace.h
+	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
+
 tests/test_ue8m0$(EXE): tests/test_ue8m0.c st.h json.h compat.h
 	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
 
@@ -1244,7 +1247,7 @@ install: colibri$(EXE) inkling$(EXE) kimi_k3$(EXE) olmoe$(EXE) \
 	  $(INSTALL) -m 755 deepseek_v4$(EXE) $(DESTDIR)$(LIBEXECDIR)/deepseek_v4$(EXE); \
 	fi
 	$(INSTALL) -m 644 family_registry.py resource_plan.py doctor.py autotune.py \
-		openai_server.py v4_dsml.py version.py $(DESTDIR)$(LIBEXECDIR)/
+		openai_server.py cluster.py v4_dsml.py version.py $(DESTDIR)$(LIBEXECDIR)/
 	$(INSTALL) -m 644 tools/*.py $(DESTDIR)$(LIBEXECDIR)/tools/
 	@# The dashboard is an optional build artifact (cd web && npm run build), so install
 	@# it only when it exists. It goes NEXT TO openai_server.py, which probes ./web/dist.

c/cluster.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python3
+"""Registration and discovery control plane for local expert workers."""
+
+import argparse
+import json
+import os
+import sys
+import threading
+import time
+from http.server import ThreadingHTTPServer
+from urllib.request import Request, urlopen
+
+import openai_server
+
+
+PROTOCOL_VERSION = 1
+
+
+class ClusterRegistry:
+    def __init__(self, stale_after=30.0):
+        self.stale_after = float(stale_after)
+        self._nodes = {}
+        self._lock = threading.Lock()
+
+    def register(self, node):
+        required = {"node_id", "host", "port", "role"}
+        missing = sorted(required - set(node))
+        if missing:
+            raise ValueError("missing node fields: " + ", ".join(missing))
+        port = int(node["port"])
+        if not 1 <= port <= 65535:
+            raise ValueError("port must be between 1 and 65535")
+        role = str(node["role"])
+        if role not in ("expert", "dense", "coordinator"):
+            raise ValueError("role must be expert, dense, or coordinator")
+        record = dict(node)
+        record.update(protocol_version=PROTOCOL_VERSION, port=port, last_seen=time.time())
+        with self._lock:
+            self._nodes[str(node["node_id"])] = record
+        return record
+
+    def heartbeat(self, node_id):
+        with self._lock:
+            node = self._nodes.get(str(node_id))
+            if node is None:
+                raise KeyError(node_id)
+            node["last_seen"] = time.time()
+            return dict(node)
+
+    def snapshot(self):
+        now = time.time()
+        with self._lock:
+            nodes = [dict(node) for node in self._nodes.values()
+                     if now - node["last_seen"] <= self.stale_after]
+        nodes.sort(key=lambda node: (node["role"], node["node_id"]))
+        return {"protocol_version": PROTOCOL_VERSION, "nodes": nodes}
+
+    def expert_endpoints(self):
+        return [f"{node['host']}:{node['port']}"
+                for node in self.snapshot()["nodes"] if node["role"] == "expert"]
+
+
+class _Handler(openai_server.APIHandler):
+    # The original registry served HTTP/1.0: one request per connection, no
+    # keep-alive. Keep that wire shape; the shared APIHandler machinery (Host
+    # guard, body cap, deadline reader, connection bookkeeping) still applies.
+    protocol_version = "HTTP/1.0"
+    server_version = "colibri-cluster/1"
+
+    def log_message(self, *_args):
+        return
+
+    def _fail(self, error):
+        self.send_json(error.status, {"error": error.message})
+
+    def do_GET(self):  # noqa: N802 - stdlib handler API
+        try:
+            self._check_host()
+        except openai_server.APIError as error:
+            self._fail(error)
+            return
+        if self.path in ("/health", "/v1/cluster/topology"):
+            self.send_json(200, self.server.registry.snapshot())
+            return
+        self.send_json(404, {"error": "not found"})
+
+    def do_POST(self):  # noqa: N802 - stdlib handler API
+        try:
+            self._check_host()
+            body = self.read_json()
+            if self.path == "/v1/cluster/register":
+                self.send_json(200, self.server.registry.register(body))
+            elif self.path == "/v1/cluster/heartbeat":
+                self.send_json(200, self.server.registry.heartbeat(body["node_id"]))
+            else:
+                self.send_json(404, {"error": "not found"})
+        except openai_server.APIError as error:
+            self._fail(error)
+        except (KeyError, ValueError, TypeError) as error:
+            self.send_json(400, {"error": str(error)})
+
+
+class ClusterServer(openai_server.APIServer):
+    def __init__(self, address, registry, allowed_hosts=()):
+        ThreadingHTTPServer.__init__(self, address, _Handler)
+        self.registry = registry
+        # Shared APIHandler plumbing reads these off the server. The registry has
+        # no API key or CORS surface, and its Host guard stays loopback + bind
+        # address (the same default openai_server enforces). A cross-host worker
+        # registers from a LAN IP, so the operator opts those hosts in via
+        # --allowed-host (the same #597 escape hatch coli serve exposes); the
+        # default stays loopback + bind address.
+        self.cors_origins = ()
+        self.allowed_hosts = tuple(
+            h.strip().lower() for h in allowed_hosts if h and h.strip())
+        # Connection-cap bookkeeping for the inherited APIServer.process_request /
+        # _release / close_request. There is no engine or model to carry, so
+        # initialise the tracking state directly rather than APIServer.__init__.
+        self._conn_lock = threading.Lock()
+        self._conn_live = 0
+        self._conn_by_ip = {}
+        self._conn_owner = {}
+
+
+def _endpoint(coordinator, path):
+    return coordinator.rstrip("/") + "/v1/cluster/" + path
+
+
+def serve(host="127.0.0.1", port=8765, stale_after=30.0, allowed_hosts=()):
+    if allowed_hosts and "*" in allowed_hosts:
+        print("WARNING: --allowed-host '*' accepts ANY Host header "
+              "(DNS-rebinding guard disabled)", file=sys.stderr)
+    server = ClusterServer((host, port), ClusterRegistry(stale_after), allowed_hosts)
+    print(f"colibri cluster coordinator listening on http://{host}:{port}", flush=True)
+    try:
+        server.serve_forever()
+    finally:
+        server.server_close()
+
+
+def register(coordinator, node):
+    request = Request(_endpoint(coordinator, "register"),
+                      data=json.dumps(node).encode(),
+                      headers={"Content-Type": "application/json"}, method="POST")
+    with urlopen(request, timeout=5) as response:
+        return json.load(response)
+
+
+def heartbeat(coordinator, node_id):
+    request = Request(_endpoint(coordinator, "heartbeat"),
+                      data=json.dumps({"node_id": node_id}).encode(),
+                      headers={"Content-Type": "application/json"}, method="POST")
+    with urlopen(request, timeout=5) as response:
+        return json.load(response)
+
+
+def discover_workers(coordinator):
+    with urlopen(_endpoint(coordinator, "topology"), timeout=5) as response:
+        topology = json.load(response)
+    return [f"{node['host']}:{int(node['port'])}"
+            for node in topology.get("nodes", []) if node.get("role") == "expert"]
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8765)
+    parser.add_argument("--stale-after", type=float, default=30.0)
+    parser.add_argument("--allowed-host", action="append",
+                        default=[h.strip() for h in os.environ.get("COLI_ALLOWED_HOSTS", "").split(",") if h.strip()],
+                        help="additional Host header accepted by the DNS-rebinding guard; repeat as needed")
+    args = parser.parse_args()
+    serve(args.host, args.port, args.stale_after, args.allowed_host)
+
+
+if __name__ == "__main__":
+    main()