File size: 14,758 Bytes
e7654c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 | """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 # answer tokens per launch
PREFILL = 64 # prompt tokens per launch (keeps rounds short while others are writing)
MAX_AT_ONCE = 5 # chats working in the same round (16 SMs each at 5: one attention block per head)
MAX_POS = 4096
MAX_NEW = 1024
PANES = 8 # conversations the page can open
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] = [] # tokens that have cache entries, in order
self.listeners: list[queue.Queue] = []
self.job: dict | None = None # the answer being written
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) # warms the GPU's clocks without touching a conversation
self.lock = threading.Lock()
self.wake = threading.Event()
self.recent: list[tuple[float, int]] = [] # (time, tokens) of finished rounds, for the total rate
self.last_work = 0.0
threading.Thread(target=self.loop, daemon=True).start()
# ---- requests (HTTP threads) ----
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}) # every open page shows it
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
# ---- the GPU loop (one thread) ----
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: # report it and drop the round's jobs
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"]): # the prompt is read: its last step's word is the first answer word
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("�"): # half of a character: wait for the rest
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): # quiet
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()
|