diff --git a/README.md b/README.md index 86aad79c..c609896a 100644 --- a/README.md +++ b/README.md @@ -89,8 +89,33 @@ 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. + +### Browser WebGPU expert workers + +The optional bridge lets a browser with WebGPU execute exported expert FFNs. +Start the bridge and serve `web/public` from the same origin: + +```bash +./coli webgpu --host 0.0.0.0 --control-port 8765 --data-port 8766 +``` + +Open `web/public/webgpu-worker.html`, choose the bridge WebSocket URL, and load +a manifest exported from a source checkpoint: + +```bash +uv run python c/tools/export_webgpu_expert.py \ + --source /models/glm52_fp8 --output web/public \ + --layer 0-2 --expert 0-7 +``` + +Point the native coordinator at the bridge with `--webgpu-workers BRIDGE_IP:8766`. +The bridge uses the same `COLIEX01` request/response framing as native expert +workers, including the layer, hidden, intermediate, expert IDs, and raw +little-endian f32 activations. This first format is intentionally f32 for +cross-worker validation; quantized WebGPU buffers can be added without +changing the activation protocol. ## How it works diff --git a/c/Makefile b/c/Makefile index c7175681..d407f597 100644 --- a/c/Makefile +++ b/c/Makefile @@ -199,7 +199,7 @@ else PYTHON ?= python3 endif CUDA_OBJ = -TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_mirror$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_sample_nan$(EXE) tests/test_temp_env$(EXE) tests/test_tok_o200k$(EXE) tests/test_kv_alloc$(EXE) tests/test_int3$(EXE) tests/test_int3_load$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE) tests/test_logit_nan$(EXE) tests/test_pipe_block$(EXE) tests/test_e8_kernel$(EXE) +TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_mirror$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_sample_nan$(EXE) tests/test_temp_env$(EXE) tests/test_tok_o200k$(EXE) tests/test_kv_alloc$(EXE) tests/test_int3$(EXE) tests/test_int3_load$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE) tests/test_logit_nan$(EXE) tests/test_pipe_block$(EXE) tests/test_e8_kernel$(EXE) tests/test_webgpu_protocol$(EXE) ifneq (,$(LINUX)) TEST_BINS += tests/test_uring$(EXE) endif @@ -282,7 +282,7 @@ $(file >.build-config,$(BUILD_CONFIG)) endif .build-config: ; -colibri$(EXE): colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h quant.h sample.h kv_persist.h telemetry.h $(CUDA_OBJ) $(METAL_OBJ) .build-config +colibri$(EXE): colibri.c webgpu.h st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h quant.h sample.h kv_persist.h telemetry.h $(CUDA_OBJ) $(METAL_OBJ) .build-config $(CC) $(CFLAGS) colibri.c $(CUDA_OBJ) $(METAL_OBJ) -o colibri$(EXE) $(LDFLAGS) # Windows runtime loader object: resolves coli_cuda_* from coli_cuda.dll. @@ -465,6 +465,9 @@ tests/test_uring$(EXE): tests/test_uring.c colibri.c st.h uring.h json.h tok.h t tests/test_pipe_block$(EXE): tests/test_pipe_block.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) +tests/test_webgpu_protocol$(EXE): tests/test_webgpu_protocol.c colibri.c webgpu.h 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 + $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) + test-c: $(TEST_BINS) $(PYTHON) tools/run_tests.py $(TEST_BINS) diff --git a/c/coli b/c/coli index 5864e689..899c6d48 100755 --- a/c/coli +++ b/c/coli @@ -211,6 +211,7 @@ def env_for(a): e.setdefault("PIPE", "1") e.setdefault("PILOT_REAL", "1") e["COLI_POLICY"]=a.policy + if getattr(a, "webgpu_workers", None): e["WEBGPU_WORKERS"] = a.webgpu_workers if a.ram: e["RAM_GB"]=str(a.ram) if a.ngen: e["NGEN"]=str(a.ngen) if a.topp: e["TOPP"]=str(a.topp) @@ -836,6 +837,11 @@ def cmd_serve(a): try: os.unlink(serve_pidfile(a.port)) except OSError: pass +def cmd_webgpu_coordinator(a): + return subprocess.call([sys.executable, os.path.join(HERE, "webgpu.py"), + "--host", a.host, "--control-port", str(a.control_port), + "--data-port", str(a.data_port)]) + def cmd_stop(a): """Shut down a running `coli serve` AND its engine — one command, no pkill. The engine re-execs itself for OMP tuning, so its process is named `exe`, @@ -963,6 +969,8 @@ def main(): common.add_argument("--cap", type=int, default=8); common.add_argument("--ngen", type=int, default=1024) # rete di sicurezza: la fine vera la decidono gli stop token common.add_argument("--topp", type=float, default=0); common.add_argument("--topk", type=int, default=0) common.add_argument("--temp", type=float, default=None) # temperatura token (0=greedy, default 1.0+nucleus .95) + common.add_argument("--webgpu-workers", default=os.environ.get("WEBGPU_WORKERS"), + help="WebGPU proxy endpoint, host:port") ap=argparse.ArgumentParser(prog="coli", parents=[common], description="colibri — run GLM-5.2 locally") ap.add_argument("--version", action="version", version=f"colibri {_version}") sub=ap.add_subparsers(dest="cmd") @@ -988,6 +996,10 @@ def main(): ps.add_argument("--max-queue",type=int,default=int(os.environ.get("COLI_MAX_QUEUE","8"))) ps.add_argument("--queue-timeout",type=float,default=float(os.environ.get("COLI_QUEUE_TIMEOUT","300"))) ps.add_argument("--kv-slots",type=int,default=int(os.environ.get("COLI_KV_SLOTS","1"))) + pcluster=sub.add_parser("webgpu", help="serve the browser WebGPU coordinator bridge") + pcluster.add_argument("--host",default="127.0.0.1") + pcluster.add_argument("--control-port",type=int,default=8765) + pcluster.add_argument("--data-port",type=int,default=8766) pst=sub.add_parser("stop", parents=[common], help="shut down a running coli serve and its engine") pst.add_argument("--port",type=int,default=8000); pst.add_argument("--dry-run",action="store_true") pw=sub.add_parser("web", parents=[common], help="serve + open the dashboard in a browser") @@ -1015,6 +1027,7 @@ def main(): handler={"build":cmd_build,"info":cmd_info,"plan":cmd_plan,"doctor":cmd_doctor, "run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"stop":cmd_stop,"bench":cmd_bench, "convert":cmd_convert,"web":cmd_web}.get(a.cmd) + if a.cmd=="webgpu": handler=cmd_webgpu_coordinator if handler: sys.exit(handler(a) or 0) banner(); print(__doc__) diff --git a/c/colibri.c b/c/colibri.c index bb24c831..75333bc3 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -30,6 +30,10 @@ #include #if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) #include /* select() serve-loop polling (#68); not on native MinGW */ +#include +#include +#include +#include #endif #if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) #include @@ -75,6 +79,11 @@ static const float *g_pre_sh; /* routing precalcolata dalla GPU (Metal layer CB o device router CUDA, #431): * moe() la usa e salta la FASE A. NULL = router su CPU. */ static const int *g_pre_idx; static const float *g_pre_w; static const int *g_pre_keff; +#if !defined(_WIN32) +typedef struct { int fd; char host[128]; int port; } WebGPUWorker; +static WebGPUWorker g_webgpu_worker; +static int g_webgpu_enabled; +#endif #ifdef __APPLE__ #include /* host_statistics64: MemAvailable di macOS */ #endif @@ -1670,6 +1679,8 @@ static int expert_load(Model *m, int layer, int eid, ESlot *s, int fatal, int de return rc; } +#include "webgpu.h" + #ifdef __linux__ /* io_uring expert batches. One owner prepares all reads for a block, submits * them in one syscall, and reaps CQEs on demand. The kernel, rather than a set @@ -3027,6 +3038,12 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int int shared_on_gpu=0; (void)shared_on_gpu; /* set by the Metal path when Phase E was fused */ for(int base=0;basepin[layer]; @@ -6533,6 +6550,12 @@ int main(int argc, char **argv){ #endif printf("== GLM C engine (glm_moe_dsa), cache=%d experts/layer | experts@%d-bit dense@%d-bit | idot: " IDOT_KERNEL " ==\n", cap, ebits, dbits); g_mem_avail_boot = mem_available_gb(); +#if !defined(_WIN32) + if(getenv("WEBGPU_WORKERS") && *getenv("WEBGPU_WORKERS")){ + webgpu_init(); + atexit(webgpu_close); + } +#endif Model m; double t0=now_s(); model_init(&m,snap,cap,ebits,dbits); if(!g_direct_heat_explicit){ /* COLI_DISKCLASS_WINDOW default, needs m.c (topk/n_layers) */ /* CURRENT-STATE CALIBRATION: the "8" multiplier (recency window ~= the last 8 diff --git a/c/tests/test_webgpu.py b/c/tests/test_webgpu.py new file mode 100644 index 00000000..dd2a6017 --- /dev/null +++ b/c/tests/test_webgpu.py @@ -0,0 +1,94 @@ +import socket +import struct +import sys +import threading +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from webgpu import (WebGPURegistry, _ProxyHandler, _ProxyServer, _recv_exact, + _ws_frame, _ws_read_frame, pack_batch, pack_response, + parse_batch, parse_response) + + +class WebGPURuntimeTests(unittest.TestCase): + def test_batch_and_response_preserve_raw_activation_bytes(self): + batch = {"layer": 4, "hidden": 2, "intermediate": 3, + "items": [{"expert_id": 7, "rows": 1, "activations": b"abcdefgh"}]} + decoded = parse_batch(pack_batch(batch)) + self.assertEqual(decoded["intermediate"], 3) + self.assertEqual(decoded["items"][0]["activations"], b"abcdefgh") + version, status, count, _ = parse_response(pack_response(decoded, [b"12345678"])) + self.assertEqual((version, status, count), (1, 0, 1)) + + def test_native_and_webgpu_workers_share_the_same_request_layout(self): + batch = {"layer": 4, "hidden": 2, "intermediate": 3, + "items": [{"expert_id": 7, "rows": 1, "activations": b"abcdefgh"}]} + native_wire = (b"COLIEX01" + struct.pack("!5I", 1, 4, 2, 3, 1) + + struct.pack("!2I", 7, 1) + b"abcdefgh") + self.assertEqual(pack_batch(batch), native_wire) + self.assertEqual(parse_batch(native_wire), batch | {"version": 1}) + + def test_dispatches_batch_to_a_browser_connection(self): + coordinator, browser = socket.socketpair() + registry = WebGPURegistry() + connection = registry.register(coordinator, {"node_id": "browser-a", "expert_ids": ["4:7"]}) + + def browser_worker(): + opcode, payload = _ws_read_frame(browser) + self.assertEqual(opcode, 2) + incoming = parse_batch(payload) + browser.sendall(_ws_frame(2, pack_response(incoming, [b"12345678"]))) + + worker = threading.Thread(target=browser_worker) + worker.start() + batch = {"layer": 4, "hidden": 2, "intermediate": 3, + "items": [{"expert_id": 7, "rows": 1, "activations": b"abcdefgh"}]} + try: + version, status, count, _ = parse_response(registry.dispatch(pack_batch(batch))) + self.assertEqual((version, status, count), (1, 0, 1)) + worker.join(timeout=1) + self.assertFalse(worker.is_alive()) + finally: + registry.unregister(connection) + browser.close() + + def test_native_tcp_client_can_use_webgpu_proxy_wire_contract(self): + registry = WebGPURegistry() + coordinator, browser = socket.socketpair() + connection = registry.register(coordinator, {"node_id": "browser-a", "expert_ids": ["4:7"]}) + proxy = _ProxyServer(("127.0.0.1", 0), _ProxyHandler) + proxy.registry = registry + thread = threading.Thread(target=proxy.serve_forever, daemon=True) + thread.start() + + def browser_worker(): + opcode, payload = _ws_read_frame(browser) + self.assertEqual(opcode, 2) + incoming = parse_batch(payload) + browser.sendall(_ws_frame(2, pack_response(incoming, [b"12345678"]))) + + worker = threading.Thread(target=browser_worker) + worker.start() + native = socket.create_connection(proxy.server_address) + batch = {"layer": 4, "hidden": 2, "intermediate": 3, + "items": [{"expert_id": 7, "rows": 1, "activations": b"abcdefgh"}]} + try: + native.sendall(pack_batch(batch)) + header = _recv_exact(native, 20) + _, status, count = parse_response(header)[0:3] + payload = _recv_exact(native, 8 + 8) + self.assertEqual((status, count), (0, 1)) + self.assertEqual(parse_response(header + payload)[3][-8:], b"12345678") + worker.join(timeout=1) + self.assertFalse(worker.is_alive()) + finally: + native.close() + proxy.shutdown(); proxy.server_close() + registry.unregister(connection) + browser.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/c/tests/test_webgpu_protocol.c b/c/tests/test_webgpu_protocol.c new file mode 100644 index 00000000..b0655316 --- /dev/null +++ b/c/tests/test_webgpu_protocol.c @@ -0,0 +1,52 @@ +/* Compatibility gate for native expert workers and WebGPU proxies. + * Both consume COLIEX01: network-order u32 headers plus raw f32 bytes. */ +#define main coli_engine_main_unused +#include "../colibri.c" +#undef main + +#include + +#if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) +static void test_shared_wire_header(void) +{ + int sockets[2]; + assert(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0); + uint32_t value; + float input[2] = {1.0f, -2.0f}; + assert(webgpu_io(sockets[0], (void *)COLI_WEBGPU_MAGIC, 8, 1) == 0); + value = 1; assert(webgpu_u32(sockets[0], &value, 1) == 0); + value = 4; assert(webgpu_u32(sockets[0], &value, 1) == 0); + value = 2; assert(webgpu_u32(sockets[0], &value, 1) == 0); + value = 3; assert(webgpu_u32(sockets[0], &value, 1) == 0); + value = 1; assert(webgpu_u32(sockets[0], &value, 1) == 0); + value = 7; assert(webgpu_u32(sockets[0], &value, 1) == 0); + value = 1; assert(webgpu_u32(sockets[0], &value, 1) == 0); + assert(webgpu_io(sockets[0], input, sizeof(input), 1) == 0); + + char magic[8]; + assert(webgpu_io(sockets[1], magic, 8, 0) == 0); + assert(memcmp(magic, COLI_WEBGPU_MAGIC, 8) == 0); + value = 0; assert(webgpu_u32(sockets[1], &value, 0) == 0 && value == 1); + value = 0; assert(webgpu_u32(sockets[1], &value, 0) == 0 && value == 4); + value = 0; assert(webgpu_u32(sockets[1], &value, 0) == 0 && value == 2); + value = 0; assert(webgpu_u32(sockets[1], &value, 0) == 0 && value == 3); + value = 0; assert(webgpu_u32(sockets[1], &value, 0) == 0 && value == 1); + value = 0; assert(webgpu_u32(sockets[1], &value, 0) == 0 && value == 7); + value = 0; assert(webgpu_u32(sockets[1], &value, 0) == 0 && value == 1); + float received[2] = {0}; + assert(webgpu_io(sockets[1], received, sizeof(received), 0) == 0); + assert(received[0] == 1.0f && received[1] == -2.0f); + close(sockets[0]); close(sockets[1]); +} +#endif + +int main(void) +{ +#if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) + test_shared_wire_header(); + puts("webgpu protocol compatibility: ok"); +#else + puts("webgpu protocol compatibility: skipped on Windows"); +#endif + return 0; +} diff --git a/c/tools/export_webgpu_expert.py b/c/tools/export_webgpu_expert.py new file mode 100644 index 00000000..bd606092 --- /dev/null +++ b/c/tools/export_webgpu_expert.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +"""Export selected dense MoE experts as browser-uploadable little-endian f32 buffers.""" + +import argparse +import json +from pathlib import Path + + +def parse_indices(values): + result = [] + for value in values: + for part in value.split(","): + if "-" in part: + start, end = (int(item) for item in part.split("-", 1)) + result.extend(range(start, end + 1)) + elif part.strip(): + result.append(int(part)) + return sorted(set(result)) + + +def export(source, output, layers, experts): + try: + import numpy as np + from safetensors import safe_open + except ImportError as error: + raise SystemExit("install exporter dependencies with: uv pip install numpy safetensors") from error + source, output = Path(source), Path(output) + shards = sorted(source.glob("*.safetensors")) + if not shards: + raise SystemExit(f"no safetensors files found in {source}") + index = {} + for shard in shards: + with safe_open(str(shard), framework="pt", device="cpu") as handle: + for name in handle.keys(): + index.setdefault(name, shard) + manifest = {"schema_version": 1, "dtype": "f32", "experts": {}} + for layer in layers: + for expert in experts: + prefix = f"model.layers.{layer}.mlp.experts.{expert}" + names = {key: f"{prefix}.{key}_proj.weight" for key in ("gate", "up", "down")} + missing = [name for name in names.values() if name not in index] + if missing: + raise SystemExit(f"missing tensors for {layer}:{expert}: {', '.join(missing)}") + arrays = {} + for key, name in names.items(): + with safe_open(str(index[name]), framework="pt", device="cpu") as handle: + arrays[key] = np.asarray(handle.get_tensor(name).detach().cpu().numpy(), dtype="65535) return -1; + char portbuf[16]; snprintf(portbuf,sizeof(portbuf),"%d",port); + struct addrinfo hint={0},*ai=NULL; hint.ai_socktype=SOCK_STREAM; + if(getaddrinfo(copy,portbuf,&hint,&ai)!=0) return -1; + int fd=-1; + for(struct addrinfo *p=ai;p;p=p->ai_next){ + fd=socket(p->ai_family,p->ai_socktype,p->ai_protocol); + if(fd<0) continue; + if(connect(fd,p->ai_addr,p->ai_addrlen)==0) break; + close(fd); fd=-1; + } + freeaddrinfo(ai); if(fd<0) return -1; + out->fd=fd; out->port=port; strncpy(out->host,copy,sizeof(out->host)-1); out->host[sizeof(out->host)-1]=0; + return 0; +} +static void webgpu_close(void){ + if(g_webgpu_enabled && g_webgpu_worker.fd>=0) close(g_webgpu_worker.fd); + g_webgpu_enabled=0; +} +static void webgpu_init(void){ + const char *list=getenv("WEBGPU_WORKERS"); if(!list||!*list) return; + if(strchr(list,',')){ fprintf(stderr,"[WEBGPU] one proxy endpoint is supported in this slice\n"); exit(1); } + if(webgpu_connect(list,&g_webgpu_worker)){ fprintf(stderr,"[WEBGPU] proxy %s is unreachable\n",list); exit(1); } + g_webgpu_enabled=1; + fprintf(stderr,"[WEBGPU] coordinator connected to proxy %s\n",list); +} + +typedef struct { int eid,nr; int *rows; float *weights,*inputs; } WebGPUItem; +static void webgpu_item_free(WebGPUItem *item){ free(item->rows); free(item->weights); free(item->inputs); memset(item,0,sizeof(*item)); } +static int webgpu_item(const int *idxs,const float *ws,const int *keff,int K,int S,int eid, + WebGPUItem *item,int D,const float *x){ + item->eid=eid; item->nr=0; + for(int s=0;snr++; break; } + if(!item->nr) return 0; + item->rows=malloc((size_t)item->nr*sizeof(int)); item->weights=malloc((size_t)item->nr*sizeof(float)); + item->inputs=falloc((int64_t)item->nr*D); int r=0; + for(int s=0;srows[r]=s; item->weights[r]=ws[(int64_t)s*K+k]; + memcpy(item->inputs+(int64_t)r*D,x+(int64_t)s*D,(size_t)D*sizeof(float)); r++; break; + } + return 1; +} +static void webgpu_moe_batch(Model *m,int layer,float *x,int S,float *out, + const int *idxs,const float *ws,const int *keff,int K, + const int *uniq,int base,int nb){ + WebGPUWorker *w=&g_webgpu_worker; WebGPUItem items[64]; memset(items,0,sizeof(items)); int n=0,D=m->c.hidden; + for(int j=0;jfd,(void*)COLI_WEBGPU_MAGIC,8,1)) goto fail; + v=1; if(webgpu_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)layer; if(webgpu_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)D; if(webgpu_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)m->c.moe_inter; if(webgpu_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)n; if(webgpu_u32(w->fd,&v,1)) goto fail; + for(int j=0;jfd,&v,1)) goto fail; + v=(uint32_t)items[j].nr; if(webgpu_u32(w->fd,&v,1)) goto fail; + if(webgpu_io(w->fd,items[j].inputs,(size_t)items[j].nr*D*sizeof(float),1)) goto fail; + } + if(webgpu_io(w->fd,magic,8,0)||memcmp(magic,COLI_WEBGPU_MAGIC,8)) goto fail; + if(webgpu_u32(w->fd,&v,0)||v!=1) goto fail; + if(webgpu_u32(w->fd,&v,0)||v!=0) goto fail; + if(webgpu_u32(w->fd,&v,0)||v!=(uint32_t)n) goto fail; + for(int j=0;jfd,&eid,0)||webgpu_u32(w->fd,&nr,0)||eid!=(uint32_t)items[j].eid||nr!=(uint32_t)items[j].nr) goto fail; + float *y=falloc((int64_t)nr*D); + if(webgpu_io(w->fd,y,(size_t)nr*D*sizeof(float),0)){ free(y); goto fail; } + for(uint32_t r=0;rhost,w->port,layer); exit(1); +} +#endif diff --git a/c/webgpu.py b/c/webgpu.py new file mode 100644 index 00000000..8c6721fa --- /dev/null +++ b/c/webgpu.py @@ -0,0 +1,295 @@ +#!/usr/bin/env python3 +"""Bridge browser WebGPU workers to colibri's binary expert data plane.""" + +import argparse +import base64 +import hashlib +import json +import socketserver +import struct +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + + +MAGIC = b"COLIEX01" +VERSION = 1 +WEBSOCKET_PATH = "/v1/webgpu" + + +def _recv_exact(sock, size): + chunks = [] + while size: + chunk = sock.recv(size) + if not chunk: + raise ConnectionError("socket closed") + chunks.append(chunk) + size -= len(chunk) + return b"".join(chunks) + + +def _ws_read_frame(sock): + first, second = _recv_exact(sock, 2) + opcode = first & 0x0F + masked = second & 0x80 + length = second & 0x7F + if length == 126: + length = struct.unpack("!H", _recv_exact(sock, 2))[0] + elif length == 127: + length = struct.unpack("!Q", _recv_exact(sock, 8))[0] + if length > 64 << 20: + raise ValueError("WebSocket frame too large") + mask = _recv_exact(sock, 4) if masked else b"" + data = bytearray(_recv_exact(sock, length)) + for index, value in enumerate(mask): + for offset in range(index, length, 4): + data[offset] ^= value + return opcode, bytes(data) + + +def _ws_frame(opcode, payload): + length = len(payload) + if length < 126: + return bytes((0x80 | opcode, length)) + payload + if length <= 0xFFFF: + return bytes((0x80 | opcode, 126)) + struct.pack("!H", length) + payload + return bytes((0x80 | opcode, 127)) + struct.pack("!Q", length) + payload + + +def parse_batch(payload): + if len(payload) < 28 or payload[:8] != MAGIC: + raise ValueError("invalid expert batch magic") + version, layer, hidden, intermediate, count = struct.unpack_from("!5I", payload, 8) + if (version != VERSION or count > 64 or not 1 <= hidden <= 65536 + or not 1 <= intermediate <= 1 << 20): + raise ValueError("invalid expert batch header") + offset, items = 28, [] + for _ in range(count): + if offset + 8 > len(payload): + raise ValueError("truncated expert item") + expert_id, rows = struct.unpack_from("!2I", payload, offset) + offset += 8 + if not 1 <= rows <= 65536: + raise ValueError("invalid expert rows") + size = rows * hidden * 4 + if offset + size > len(payload): + raise ValueError("truncated activation payload") + items.append({"expert_id": expert_id, "rows": rows, + "activations": payload[offset:offset + size]}) + offset += size + if offset != len(payload): + raise ValueError("trailing expert bytes") + return {"version": version, "layer": layer, "hidden": hidden, + "intermediate": intermediate, "items": items} + + +def pack_batch(batch): + parts = [MAGIC, struct.pack("!5I", VERSION, batch["layer"], batch["hidden"], + batch["intermediate"], len(batch["items"]))] + for item in batch["items"]: + expected = item["rows"] * batch["hidden"] * 4 + if len(item["activations"]) != expected: + raise ValueError("activation byte count does not match shape") + parts += [struct.pack("!2I", item["expert_id"], item["rows"]), item["activations"]] + return b"".join(parts) + + +def pack_response(batch, outputs, status=0): + parts = [MAGIC, struct.pack("!3I", VERSION, status, len(outputs))] + for item, output in zip(batch["items"], outputs): + expected = item["rows"] * batch["hidden"] * 4 + if len(output) != expected: + raise ValueError("output byte count does not match shape") + parts += [struct.pack("!2I", item["expert_id"], item["rows"]), output] + return b"".join(parts) + + +def parse_response(payload): + if len(payload) < 20 or payload[:8] != MAGIC: + raise ValueError("invalid expert response") + version, status, count = struct.unpack_from("!3I", payload, 8) + if version != VERSION or count > 64: + raise ValueError("invalid expert response header") + return version, status, count, payload[20:] + + +class WebGPUConnection: + def __init__(self, sock, node): + self.sock, self.node = sock, node + self.lock = threading.Lock() + + def request(self, payload): + with self.lock: + self.node["load"] = self.node.get("load", 0) + 1 + self.sock.sendall(_ws_frame(2, payload)) + try: + while True: + opcode, data = _ws_read_frame(self.sock) + if opcode == 2: + return data + if opcode == 9: + self.sock.sendall(_ws_frame(10, data)) + elif opcode == 8: + raise ConnectionError("WebGPU worker closed") + finally: + self.node["load"] = max(0, self.node.get("load", 1) - 1) + + def close(self): + try: + self.sock.sendall(_ws_frame(8, b"")) + except OSError: + pass + self.sock.close() + + +class WebGPURegistry: + def __init__(self): + self._workers = {} + self._lock = threading.Lock() + + def register(self, sock, hello): + if hello.get("role", "webgpu") != "webgpu": + raise ValueError("worker role must be webgpu") + node_id = str(hello.get("node_id") or f"webgpu-{id(sock)}") + node = dict(hello, node_id=node_id, role="webgpu", load=0) + connection = WebGPUConnection(sock, node) + with self._lock: + self._workers[node_id] = connection + return connection + + def unregister(self, connection): + with self._lock: + if self._workers.get(connection.node["node_id"]) is connection: + del self._workers[connection.node["node_id"]] + connection.close() + + def _owner(self, layer, expert_id): + key = f"{layer}:{expert_id}" + with self._lock: + candidates = [worker for worker in self._workers.values() + if "*" in worker.node.get("expert_ids", ["*"]) + or key in worker.node.get("expert_ids", []) + or str(expert_id) in worker.node.get("expert_ids", [])] + return min(candidates, key=lambda worker: worker.node.get("load", 0)) if candidates else None + + def dispatch(self, payload): + batch = parse_batch(payload) + grouped = {} + for item in batch["items"]: + owner = self._owner(batch["layer"], item["expert_id"]) + if owner is None: + return pack_response(batch, [], status=1) + grouped.setdefault(owner, []).append(item) + results = {} + for owner, items in grouped.items(): + sub_batch = dict(batch, items=items) + version, status, count, rest = parse_response(owner.request(pack_batch(sub_batch))) + if status or count != len(items): + return pack_response(batch, [], status=1) + offset = 0 + for item in items: + if offset + 8 > len(rest): + raise ValueError("truncated WebGPU output") + expert_id, rows = struct.unpack_from("!2I", rest, offset) + offset += 8 + size = rows * batch["hidden"] * 4 + if expert_id != item["expert_id"] or rows != item["rows"] or offset + size > len(rest): + raise ValueError("invalid WebGPU output shape") + results[expert_id] = rest[offset:offset + size] + offset += size + if offset != len(rest): + raise ValueError("trailing WebGPU output") + return pack_response(batch, [results[item["expert_id"]] for item in batch["items"]]) + + +class _ControlHandler(BaseHTTPRequestHandler): + def log_message(self, *_args): + return + + def do_GET(self): # noqa: N802 - stdlib handler API + if self.path != WEBSOCKET_PATH or self.headers.get("Upgrade", "").lower() != "websocket": + self.send_error(404) + return + key = self.headers.get("Sec-WebSocket-Key") + if not key: + self.send_error(400, "missing Sec-WebSocket-Key") + return + accept = base64.b64encode(hashlib.sha1( + (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode()).digest()).decode() + self.send_response(101, "Switching Protocols") + self.send_header("Upgrade", "websocket") + self.send_header("Connection", "Upgrade") + self.send_header("Sec-WebSocket-Accept", accept) + self.end_headers() + connection = None + try: + opcode, payload = _ws_read_frame(self.connection) + if opcode != 1: + raise ValueError("WebGPU worker must begin with JSON hello") + connection = self.server.registry.register(self.connection, json.loads(payload)) + while True: + # The data proxy owns reads after registration; reading here + # would race the binary response path on the same socket. + time.sleep(1) + except (ConnectionError, OSError, ValueError, json.JSONDecodeError): + pass + finally: + if connection: + self.server.registry.unregister(connection) + + +class _ControlServer(ThreadingHTTPServer): + daemon_threads = True + + +class _ProxyHandler(socketserver.BaseRequestHandler): + def handle(self): + try: + while True: + header = _recv_exact(self.request, 28) + if header[:8] != MAGIC: + raise ValueError("invalid expert batch magic") + version, layer, hidden, intermediate, count = struct.unpack("!5I", header[8:]) + if (version != VERSION or count > 64 or not 1 <= hidden <= 65536 + or not 1 <= intermediate <= 1 << 20): + raise ValueError("invalid expert batch header") + batch_parts = [header] + for _ in range(count): + item_header = _recv_exact(self.request, 8) + _, rows = struct.unpack("!2I", item_header) + if not 1 <= rows <= 65536: + raise ValueError("invalid expert rows") + batch_parts += [item_header, _recv_exact(self.request, rows * hidden * 4)] + self.request.sendall(self.server.registry.dispatch(b"".join(batch_parts))) + except (ConnectionError, OSError, ValueError): + return + + +class _ProxyServer(socketserver.ThreadingMixIn, socketserver.TCPServer): + allow_reuse_address = True + daemon_threads = True + + +def serve(host="127.0.0.1", control_port=8765, data_port=8766): + registry = WebGPURegistry() + control = _ControlServer((host, control_port), _ControlHandler) + control.registry = registry + proxy = _ProxyServer((host, data_port), _ProxyHandler) + proxy.registry = registry + threading.Thread(target=proxy.serve_forever, daemon=True).start() + print(f"WebGPU coordinator: ws://{host}:{control_port}{WEBSOCKET_PATH} · data {data_port}", flush=True) + try: + control.serve_forever() + finally: + control.server_close() + proxy.shutdown() + proxy.server_close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--control-port", type=int, default=8765) + parser.add_argument("--data-port", type=int, default=8766) + args = parser.parse_args() + serve(args.host, args.control_port, args.data_port) diff --git a/web/public/webgpu-worker.html b/web/public/webgpu-worker.html new file mode 100644 index 00000000..185622de --- /dev/null +++ b/web/public/webgpu-worker.html @@ -0,0 +1,23 @@ + + +colibri WebGPU expert worker + +

colibri WebGPU expert worker

+

This browser tab contributes expert FFN compute to a trusted local coordinator.

+ + + + +
idle
+ diff --git a/web/public/webgpu-worker.js b/web/public/webgpu-worker.js new file mode 100644 index 00000000..9c0d0db0 --- /dev/null +++ b/web/public/webgpu-worker.js @@ -0,0 +1,49 @@ +/* Browser WebGPU expert worker. Weights are little-endian f32 in this first + * format; the manifest keeps conversion explicit and auditable. */ +const MAGIC="COLIEX01", VERSION=1; +const textMagic=bytes=>new TextDecoder().decode(bytes); +const putU32=(view,offset,value)=>view.setUint32(offset,value>>>0,false); +const getU32=(view,offset)=>view.getUint32(offset,false); +function concat(parts){const size=parts.reduce((n,p)=>n+p.byteLength,0),out=new Uint8Array(size);let at=0; + for(const part of parts){out.set(new Uint8Array(part.buffer||part,part.byteOffset||0,part.byteLength),at);at+=part.byteLength}return out} + +export class WebGPUExpertWorker { + constructor({manifest,workerId,experts=["*"],deviceType="browser"}){this.manifest=manifest;this.workerId=workerId;this.experts=experts;this.deviceType=deviceType;this.cache=new Map()} + async connect(url){ + if(!navigator.gpu)throw new Error("WebGPU is unavailable in this browser"); + const adapter=await navigator.gpu.requestAdapter();if(!adapter)throw new Error("no WebGPU adapter found"); + this.device=await adapter.requestDevice();this._createPipelines();this.socket=new WebSocket(url);this.socket.binaryType="arraybuffer"; + await new Promise((resolve,reject)=>{this.socket.addEventListener("open",resolve,{once:true});this.socket.addEventListener("error",()=>reject(new Error("WebGPU worker socket failed")),{once:true})}); + this.socket.send(JSON.stringify({role:"webgpu",node_id:this.workerId,host:"browser",port:0,device_type:this.deviceType,precision:"f32",expert_ids:this.experts})); + this.socket.addEventListener("message",event=>this._onMessage(event.data));return this; + } + _createPipelines(){ + const common=`struct Dims{rows:u32,hidden:u32,intermediate:u32,pad:u32};`; + const gate=`${common}@group(0)@binding(0)varinput:array;@group(0)@binding(1)vargate:array;@group(0)@binding(2)varup:array;@group(0)@binding(3)varhidden:array;@group(0)@binding(4)vardims:Dims;@compute@workgroup_size(64)fn main(@builtin(global_invocation_id)id:vec3){let i=id.x;let total=dims.rows*dims.intermediate;if(i>=total){return}let row=i/dims.intermediate;let unit=i%dims.intermediate;var g=0.0;var u=0.0;for(var d=0u;dhidden:array;@group(0)@binding(1)vardown:array;@group(0)@binding(2)varoutput:array;@group(0)@binding(3)vardims:Dims;@compute@workgroup_size(64)fn main(@builtin(global_invocation_id)id:vec3){let i=id.x;let total=dims.rows*dims.hidden;if(i>=total){return}let row=i/dims.hidden;let dim=i%dims.hidden;var sum=0.0;for(var unit=0u;unitthis.device.createComputePipeline({layout:"auto",compute:{module:this.device.createShaderModule({code}),entryPoint:"main"}}); + this.gatePipeline=pipeline(gate);this.downPipeline=pipeline(down); + } + async _loadExpert(layer,expertId){const key=`${layer}:${expertId}`;if(this.cache.has(key))return this.cache.get(key);const spec=this.manifest.experts[key];if(!spec)throw new Error(`manifest has no expert ${key}`); + const values=await Promise.all([spec.gate,spec.up,spec.down].map(path=>fetch(new URL(path,this.manifest.base_url||location.href)).then(response=>{if(!response.ok)throw new Error(`cannot load ${path}`);return response.arrayBuffer()}).then(buffer=>new Float32Array(buffer)))); + const expert={gate:this._weightBuffer(values[0]),up:this._weightBuffer(values[1]),down:this._weightBuffer(values[2])};this.cache.set(key,expert);return expert; + } + _weightBuffer(values){const buffer=this.device.createBuffer({size:values.byteLength,usage:GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_DST});this.device.queue.writeBuffer(buffer,0,values);return buffer} + _buffer(bytes,usage){const buffer=this.device.createBuffer({size:Math.max(4,bytes.byteLength),usage:usage|GPUBufferUsage.COPY_SRC});this.device.queue.writeBuffer(buffer,0,bytes);return buffer} + async _forward(layer,expertId,rows,hidden,inputBytes){const intermediate=this.manifest.intermediate,expert=await this._loadExpert(layer,expertId); + const input=this._buffer(inputBytes,GPUBufferUsage.STORAGE),hiddenBuffer=this.device.createBuffer({size:rows*intermediate*4,usage:GPUBufferUsage.STORAGE}); + const output=this.device.createBuffer({size:rows*hidden*4,usage:GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_SRC}),readback=this.device.createBuffer({size:rows*hidden*4,usage:GPUBufferUsage.MAP_READ|GPUBufferUsage.COPY_DST}),uniform=this.device.createBuffer({size:16,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST}); + this.device.queue.writeBuffer(uniform,0,new Uint32Array([rows,hidden,intermediate,0])); + const gg=this.device.createBindGroup({layout:this.gatePipeline.getBindGroupLayout(0),entries:[{binding:0,resource:{buffer:input}},{binding:1,resource:{buffer:expert.gate}},{binding:2,resource:{buffer:expert.up}},{binding:3,resource:{buffer:hiddenBuffer}},{binding:4,resource:{buffer:uniform}}]}); + const dg=this.device.createBindGroup({layout:this.downPipeline.getBindGroupLayout(0),entries:[{binding:0,resource:{buffer:hiddenBuffer}},{binding:1,resource:{buffer:expert.down}},{binding:2,resource:{buffer:output}},{binding:3,resource:{buffer:uniform}}]}); + const encoder=this.device.createCommandEncoder(),first=encoder.beginComputePass();first.setPipeline(this.gatePipeline);first.setBindGroup(0,gg);first.dispatchWorkgroups(Math.ceil(rows*intermediate/64));first.end(); + const second=encoder.beginComputePass();second.setPipeline(this.downPipeline);second.setBindGroup(0,dg);second.dispatchWorkgroups(Math.ceil(rows*hidden/64));second.end();encoder.copyBufferToBuffer(output,0,readback,0,rows*hidden*4);this.device.queue.submit([encoder.finish()]); + await readback.mapAsync(GPUMapMode.READ);const result=readback.getMappedRange().slice(0);readback.unmap();[input,hiddenBuffer,output,readback,uniform].forEach(buffer=>buffer.destroy());return result; + } + async _onMessage(data){try{const bytes=new Uint8Array(data),view=new DataView(bytes.buffer,bytes.byteOffset,bytes.byteLength);if(textMagic(bytes.subarray(0,8))!==MAGIC||getU32(view,8)!==VERSION)throw new Error("bad request"); + const layer=getU32(view,12),hidden=getU32(view,16),intermediate=getU32(view,20),count=getU32(view,24);if(hidden!==this.manifest.hidden||intermediate!==this.manifest.intermediate||count>64)throw new Error("unsupported request shape");let offset=28,outputs=[]; + for(let i=0;i65536||offset+size>bytes.byteLength)throw new Error("invalid activation shape");const input=bytes.slice(offset,offset+size);offset+=size;outputs.push({expertId,rows,bytes:await this._forward(layer,expertId,rows,hidden,input)})} + const head=new Uint8Array(20),response=new DataView(head.buffer);for(let i=0;i<8;i++)head[i]=MAGIC.charCodeAt(i);putU32(response,8,VERSION);putU32(response,12,0);putU32(response,16,outputs.length);const parts=[head]; + for(const output of outputs){const item=new Uint8Array(8),view2=new DataView(item.buffer);putU32(view2,0,output.expertId);putU32(view2,4,output.rows);parts.push(item,new Uint8Array(output.bytes))}this.socket.send(concat(parts)); + }catch(error){console.error("WebGPU expert request failed",error);const response=new Uint8Array(20),view=new DataView(response.buffer);for(let i=0;i<8;i++)response[i]=MAGIC.charCodeAt(i);putU32(view,8,VERSION);putU32(view,12,1);putU32(view,16,0);this.socket.send(response)}} +}