| """Several chats at once on the Ling-mini-2.0 megakernel: a local web app.
|
|
|
| python server.py [--port 8790]
|
|
|
| Then open http://127.0.0.1:8790. Every pane is its own conversation with its own cache. One thread
|
| owns the GPU and works in rounds: each round, every chat that is reading a prompt or writing an
|
| answer gets its own launch of the megakernel on a share of the SMs (82 / chats working, at most 5
|
| at once; more wait their turn), all launches run together, and their words go out to the panes as
|
| the round ends. One chat alone gets the whole GPU (~850 words/s); three together ~460 each.
|
|
|
| Greedy decoding, thinking off, up to 4,096 tokens per conversation. Binds to 127.0.0.1 only.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import argparse
|
| import json
|
| import queue
|
| import sys
|
| import threading
|
| import time
|
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
| from pathlib import Path
|
| from urllib.parse import parse_qs, urlparse
|
|
|
| import torch
|
|
|
| from ling_kernel import HERE, NQ, LingKernel, State
|
| CHUNK = 32
|
| PREFILL = 64
|
| MAX_AT_ONCE = 5
|
| MAX_POS = 4096
|
| MAX_NEW = 1024
|
| PANES = 8
|
|
|
|
|
| class Chat:
|
| """One conversation: its GPU state (KV cache and buffers), its history, its listeners."""
|
|
|
| def __init__(self, mk: LingKernel, cid: int):
|
| self.mk, self.cid = mk, cid
|
| self.state = State(MAX_POS)
|
| self.out = self.state.out
|
| self.stream = torch.cuda.Stream()
|
| self.history: list[dict] = []
|
| self.cached: list[int] = []
|
| self.listeners: list[queue.Queue] = []
|
| self.job: dict | None = None
|
|
|
| def launch(self, ins: list[int], pos0: int, n: int, forced: bool, nblocks: int) -> None:
|
| self.mk.launch_async(self.state, ins, pos0, n, forced, nblocks, self.stream)
|
|
|
| def send(self, event: dict) -> None:
|
| for q in list(self.listeners):
|
| q.put(event)
|
|
|
|
|
| class Server:
|
| def __init__(self) -> None:
|
| from transformers import AutoTokenizer
|
| self.tok = AutoTokenizer.from_pretrained(HERE)
|
| print("loading the Ling-mini-2.0 megakernel (4-bit) ...", flush=True)
|
| self.mk = LingKernel(max_pos=MAX_POS, log=lambda s: print(s, flush=True))
|
| self.sms = self.mk.sms
|
| self.eos = {i for i in (self.tok.convert_tokens_to_ids("<|role_end|>"), self.tok.eos_token_id,
|
| self.mk.config.get("eos_token_id")) if isinstance(i, int) and i >= 0}
|
| self.chats = [Chat(self.mk, i) for i in range(PANES)]
|
| self.scratch = Chat(self.mk, -1)
|
| self.lock = threading.Lock()
|
| self.wake = threading.Event()
|
| self.recent: list[tuple[float, int]] = []
|
| self.last_work = 0.0
|
| threading.Thread(target=self.loop, daemon=True).start()
|
|
|
|
|
| def start(self, cid: int, text: str) -> str | None:
|
| with self.lock:
|
| c = self.chats[cid]
|
| if c.job is not None:
|
| return "this chat is still writing"
|
| c.job = {"text": text, "state": "new"}
|
| c.send({"type": "user", "text": text})
|
| self.wake.set()
|
| return None
|
|
|
| def reset(self, cid: int) -> str | None:
|
| with self.lock:
|
| c = self.chats[cid]
|
| if c.job is not None:
|
| return "this chat is still writing"
|
| c.history, c.cached = [], []
|
| c.send({"type": "cleared"})
|
| return None
|
|
|
| def transcript(self, cid: int) -> dict:
|
| with self.lock:
|
| c = self.chats[cid]
|
| job = c.job
|
| msgs = list(c.history)
|
| if job is not None:
|
| msgs.append({"role": "user", "content": job["text"]})
|
| msgs.append({"role": "assistant", "content": job.get("shown", "")})
|
| return {"messages": msgs, "writing": job is not None}
|
|
|
| def total_rate(self) -> float:
|
| now = time.perf_counter()
|
| with self.lock:
|
| self.recent = [(t, n) for t, n in self.recent if now - t < 1.0]
|
| n = sum(k for _, k in self.recent)
|
| return n / 1.0
|
|
|
|
|
| def begin(self, c: Chat) -> None:
|
| """Turn a new message into prompt tokens; keep the cache's common prefix."""
|
| job = c.job
|
| messages = c.history + [{"role": "user", "content": job["text"]}]
|
| ids = self.tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
|
| if hasattr(ids, "keys"):
|
| ids = ids["input_ids"]
|
| ids = list(ids)
|
| if len(ids) + 2 * CHUNK + 1 > MAX_POS:
|
| c.job = None
|
| c.send({"type": "error", "text": f"this conversation is full ({MAX_POS} tokens): press Clear"})
|
| return
|
| n = 0
|
| while n < min(len(ids) - 1, len(c.cached)) and ids[n] == c.cached[n]:
|
| n += 1
|
| job.update(state="read", messages=messages, ids=ids, pos=n, out=[], shown="", seconds=0.0, steps=0,
|
| max_new=min(MAX_NEW, MAX_POS - len(ids) - CHUNK))
|
| c.send({"type": "start", "prompt_tokens": len(ids) - n})
|
|
|
| def finish(self, c: Chat, note: str = "") -> None:
|
| job = c.job
|
| text = job["shown"].strip()
|
| c.history = job["messages"] + [{"role": "assistant", "content": text}]
|
| c.cached = job["ids"] + job["out"][:-1]
|
| rate = job["steps"] / job["seconds"] if job["seconds"] else 0.0
|
| c.send({"type": "done", "tokens": len(job["out"]), "rate": rate, "note": note})
|
| c.job = None
|
|
|
| def warm(self) -> None:
|
| """After the GPU idled, Windows lowers its clocks: run the scratch chat until launches are fast."""
|
| start = time.perf_counter()
|
| best = float("inf")
|
| while time.perf_counter() - start < 2.0:
|
| t = time.perf_counter()
|
| self.scratch.launch([1000], 0, CHUNK, False, self.sms)
|
| torch.cuda.synchronize()
|
| d = time.perf_counter() - t
|
| if d > 0.97 * best and d < 1.2 * best:
|
| break
|
| best = min(best, d)
|
|
|
| def loop(self) -> None:
|
| while True:
|
| with self.lock:
|
| for c in self.chats:
|
| if c.job is not None and c.job["state"] == "new":
|
| self.begin(c)
|
| working = [c for c in self.chats if c.job is not None][:MAX_AT_ONCE]
|
| if not working:
|
| self.wake.wait(0.5)
|
| self.wake.clear()
|
| continue
|
| if time.perf_counter() - self.last_work > 1.0:
|
| self.warm()
|
| nb = max(NQ, self.sms // len(working))
|
| plan = []
|
| try:
|
| for c in working:
|
| job = c.job
|
| if job["state"] == "read":
|
| k = min(PREFILL, len(job["ids"]) - job["pos"])
|
| c.launch(job["ids"][job["pos"]:job["pos"] + k], job["pos"], k, True, nb)
|
| else:
|
| k = min(CHUNK, job["max_new"] - len(job["out"]))
|
| c.launch([job["out"][-1]], job["pos"], k, False, nb)
|
| plan.append((c, k))
|
| t = time.perf_counter()
|
| torch.cuda.synchronize()
|
| dt = time.perf_counter() - t
|
| except Exception as e:
|
| torch.cuda.synchronize()
|
| with self.lock:
|
| for c, _ in plan or [(c, 0) for c in working]:
|
| if c.job is not None:
|
| c.send({"type": "error", "text": f"GPU error: {e}"})
|
| c.job = None
|
| continue
|
| self.last_work = time.perf_counter()
|
| made = 0
|
| with self.lock:
|
| for c, k in plan:
|
| job = c.job
|
| outs = c.out[:k].tolist()
|
| if job["state"] == "read":
|
| job["pos"] += k
|
| if job["pos"] == len(job["ids"]):
|
| job["state"] = "write"
|
| job["out"] = [outs[k - 1]]
|
| made += 1
|
| self.emit(c)
|
| if outs[k - 1] in self.eos:
|
| self.finish(c)
|
| continue
|
| job["seconds"] += dt
|
| job["steps"] += k
|
| job["pos"] += k
|
| for t_ in outs:
|
| job["out"].append(t_)
|
| made += 1
|
| if t_ in self.eos:
|
| break
|
| self.emit(c)
|
| if job["out"][-1] in self.eos:
|
| self.finish(c)
|
| elif len(job["out"]) >= job["max_new"]:
|
| self.finish(c, "stopped at the length limit")
|
| self.recent.append((time.perf_counter(), made))
|
|
|
| def emit(self, c: Chat) -> None:
|
| job = c.job
|
| words = [t for t in job["out"] if t not in self.eos]
|
| text = self.tok.decode(words, skip_special_tokens=True)
|
| if text.endswith("�"):
|
| text = text.rstrip("�")
|
| if len(text) > len(job["shown"]) and text.startswith(job["shown"]):
|
| c.send({"type": "delta", "text": text[len(job["shown"]):]})
|
| job["shown"] = text
|
| elif text != job["shown"] and not text.startswith(job["shown"]):
|
| c.send({"type": "replace", "text": text})
|
| job["shown"] = text
|
|
|
|
|
| def handler(server: Server):
|
| page = (HERE / "multichat.html").read_bytes()
|
|
|
| class H(BaseHTTPRequestHandler):
|
| def log_message(self, *a):
|
| pass
|
|
|
| def reply(self, code: int, body: bytes, ctype: str = "application/json") -> None:
|
| self.send_response(code)
|
| self.send_header("Content-Type", ctype)
|
| self.send_header("Content-Length", str(len(body)))
|
| self.send_header("Cache-Control", "no-store")
|
| self.end_headers()
|
| self.wfile.write(body)
|
|
|
| def do_GET(self):
|
| u = urlparse(self.path)
|
| if u.path == "/":
|
| return self.reply(200, page, "text/html; charset=utf-8")
|
| if u.path == "/history":
|
| cid = int(parse_qs(u.query).get("chat", ["0"])[0])
|
| if not 0 <= cid < PANES:
|
| return self.reply(404, b"{}")
|
| return self.reply(200, json.dumps(server.transcript(cid)).encode())
|
| if u.path == "/status":
|
| with server.lock:
|
| busy = [c.cid for c in server.chats if c.job is not None]
|
| return self.reply(200, json.dumps({"total_rate": server.total_rate(), "busy": busy,
|
| "sms": server.sms, "panes": PANES}).encode())
|
| if u.path == "/events":
|
| cid = int(parse_qs(u.query).get("chat", ["0"])[0])
|
| if not 0 <= cid < PANES:
|
| return self.reply(404, b"{}")
|
| q: queue.Queue = queue.Queue()
|
| c = server.chats[cid]
|
| with server.lock:
|
| c.listeners.append(q)
|
| self.send_response(200)
|
| self.send_header("Content-Type", "text/event-stream")
|
| self.send_header("Cache-Control", "no-store")
|
| self.end_headers()
|
| try:
|
| self.wfile.write(b": open\n\n")
|
| self.wfile.flush()
|
| while True:
|
| try:
|
| ev = q.get(timeout=15)
|
| self.wfile.write(f"data: {json.dumps(ev)}\n\n".encode())
|
| except queue.Empty:
|
| self.wfile.write(b": keep-alive\n\n")
|
| self.wfile.flush()
|
| except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError, OSError):
|
| pass
|
| finally:
|
| with server.lock:
|
| if q in c.listeners:
|
| c.listeners.remove(q)
|
| return
|
| return self.reply(404, b"{}")
|
|
|
| def do_POST(self):
|
| u = urlparse(self.path)
|
| n = int(self.headers.get("Content-Length", "0") or 0)
|
| try:
|
| body = json.loads(self.rfile.read(n) or b"{}")
|
| cid = int(body.get("chat", -1))
|
| except (ValueError, TypeError):
|
| return self.reply(400, b'{"error": "bad request"}')
|
| if not 0 <= cid < PANES:
|
| return self.reply(404, b'{"error": "no such chat"}')
|
| if u.path == "/send":
|
| text = str(body.get("text", "")).strip()
|
| if not text:
|
| return self.reply(400, b'{"error": "empty message"}')
|
| err = server.start(cid, text)
|
| elif u.path == "/reset":
|
| err = server.reset(cid)
|
| else:
|
| return self.reply(404, b"{}")
|
| if err:
|
| return self.reply(409, json.dumps({"error": err}).encode())
|
| return self.reply(200, b'{"ok": true}')
|
|
|
| return H
|
|
|
|
|
| def main() -> None:
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--port", type=int, default=8790)
|
| args = ap.parse_args()
|
| sys.stdout.reconfigure(encoding="utf-8")
|
| server = Server()
|
| httpd = ThreadingHTTPServer(("127.0.0.1", args.port), handler(server))
|
| httpd.daemon_threads = True
|
| print(f"open http://127.0.0.1:{args.port} ({server.sms} SMs; up to {MAX_AT_ONCE} chats writing at once)", flush=True)
|
| httpd.serve_forever()
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
