Merge remote-tracking branch 'upstream/dev' into feat/gpu-backend-hardening
This commit is contained in:
noobdev-ph
+20
-1
@@ -166,7 +166,7 @@ else
|
||||
PYTHON ?= python3
|
||||
endif
|
||||
CUDA_OBJ =
|
||||
TEST_BINS = tests/test_json$(EXE) tests/test_st$(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_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE)
|
||||
TEST_BINS = tests/test_json$(EXE) tests/test_st$(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_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE)
|
||||
ifneq (,$(LINUX))
|
||||
TEST_BINS += tests/test_uring$(EXE)
|
||||
endif
|
||||
@@ -293,6 +293,9 @@ iobench$(EXE): iobench.c compat.h
|
||||
tests/test_json$(EXE): tests/test_json.c json.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
tests/test_st_pread$(EXE): tests/test_st_pread.c st.h json.h compat.h
|
||||
$(CC) $(CFLAGS) -DST_PREAD_CHUNK=7 $< -o $@ $(LDFLAGS)
|
||||
|
||||
tests/test_st$(EXE): tests/test_st.c st.h json.h compat.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
@@ -317,6 +320,14 @@ tests/test_i4_grouped$(EXE): tests/test_i4_grouped.c glm.c st.h uring.h json.h t
|
||||
tests/test_stops$(EXE): tests/test_stops.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
tests/test_topp$(EXE): tests/test_topp.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
# bench_topp is a microbenchmark (old qsort vs new heap partial-select, #335), NOT a test
|
||||
# gate -- intentionally absent from TEST_BINS. Build on demand: make tests/bench_topp
|
||||
tests/bench_topp$(EXE): tests/bench_topp.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
tests/test_kv_alloc$(EXE): tests/test_kv_alloc.c glm.c st.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
@@ -326,6 +337,14 @@ tests/test_i4_acc512$(EXE): tests/test_i4_acc512.c
|
||||
tests/test_compat_direct$(EXE): tests/test_compat_direct.c compat.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
tests/test_dsa_select$(EXE): tests/test_dsa_select.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
# bench_dsa_select is a microbenchmark (old qsort vs new quickselect partial-select, #356),
|
||||
# NOT a test gate -- intentionally absent from TEST_BINS. Build on demand: make tests/bench_dsa_select
|
||||
tests/bench_dsa_select$(EXE): tests/bench_dsa_select.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
tests/test_uring$(EXE): tests/test_uring.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
|
||||
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ Configuration through environment variables or flags (also valid after the subco
|
||||
--topp P adaptive expert top-p --topk N fixed top-k
|
||||
--ngen N maximum response tokens --cap N cache slots/layer
|
||||
"""
|
||||
import os, sys, subprocess, argparse, json, time, signal, shutil, threading, re, codecs, tempfile, textwrap
|
||||
import os, sys, subprocess, argparse, json, time, signal, shutil, threading, re, codecs, tempfile, textwrap, struct
|
||||
|
||||
# The engine mmaps every shard (144+ files); macOS default RLIMIT_NOFILE is 256.
|
||||
if sys.platform != "win32":
|
||||
@@ -484,9 +484,134 @@ def cmd_run(a):
|
||||
e=env_for(a); e["PROMPT"]=f"[gMASK]<sop><|user|>{prompt}<|assistant|><think></think>"
|
||||
sys.exit(subprocess.call([GLM, str(a.cap)], env=e))
|
||||
|
||||
def server_probe(base, api_key=None, timeout=1.5):
|
||||
"""Is a coli serve alive at `base`? Returns its model_id, or None.
|
||||
Probes /health then /v1/models — both cheap, neither touches the engine."""
|
||||
import urllib.request, urllib.error
|
||||
def get(path):
|
||||
req=urllib.request.Request(base.rstrip("/")+path)
|
||||
if api_key: req.add_header("Authorization", f"Bearer {api_key}")
|
||||
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||||
return json.loads(r.read().decode("utf-8","replace"))
|
||||
try:
|
||||
if get("/health").get("status")!="ok": return None
|
||||
data=get("/v1/models").get("data") or []
|
||||
return data[0]["id"] if data else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def chat_attached(a, base, model_id):
|
||||
"""The chat REPL over HTTP against a running `coli serve`.
|
||||
|
||||
Why this exists (the cold-chat cost, measured): spawning a private engine
|
||||
pays 34-136 s of resident load on EVERY start, and begins with an empty
|
||||
expert cache — hit rate 4% cold vs 55% warm, a ~10x on early decode. A
|
||||
resident server pays load once and keeps the LRU warm across sessions;
|
||||
its KV slots reuse the conversation prefix, so a continued chat skips
|
||||
re-prefill too. The engine byte-protocol stays untouched — this is plain
|
||||
OpenAI SSE over localhost, stdlib only."""
|
||||
import urllib.request
|
||||
print(f" {C.grn}✦ attached{C.r} {C.dim}to {base} · model {model_id} · the engine stays warm after you quit{C.r}")
|
||||
print(f" {C.dim}type and press Enter · Ctrl-C stops the answer · :reset starts a new conversation · :q exits{C.r}\n")
|
||||
msgs=[]
|
||||
w=term_w()-4
|
||||
while True:
|
||||
if TTY:
|
||||
print(f" {C.dgray}╭{'─'*w}╮{C.r}")
|
||||
try: msg=input(f" {C.dgray}│{C.r} {C.teal}{C.b}›{C.r} ")
|
||||
except EOFError: print(); break
|
||||
print(f" {C.dgray}╰{'─'*w}╯{C.r}")
|
||||
else:
|
||||
try: msg=input()
|
||||
except EOFError: break
|
||||
msg=msg.strip()
|
||||
if msg in (":q",":quit","exit"): break
|
||||
if not msg: continue
|
||||
if msg==":reset": msgs=[]; print(f" {C.dim}✦ new conversation{C.r}\n"); continue
|
||||
msgs.append({"role":"user","content":msg})
|
||||
body=json.dumps({"model":model_id,"messages":msgs,"stream":True,
|
||||
"max_tokens":a.ngen}).encode()
|
||||
req=urllib.request.Request(base.rstrip("/")+"/v1/chat/completions", data=body,
|
||||
headers={"Content-Type":"application/json"})
|
||||
if a.api_key: req.add_header("Authorization", f"Bearer {a.api_key}")
|
||||
print(f"\n {C.teal}◆ colibrì{C.r}")
|
||||
sp=Spinner("thinking…"); sp.start()
|
||||
md=MDStream(" "); reply=[]; first=True; t0=time.time(); interrupted=False
|
||||
try:
|
||||
with urllib.request.urlopen(req) as r:
|
||||
for raw in r:
|
||||
line=raw.decode("utf-8","replace").strip()
|
||||
if not line.startswith("data: "): continue
|
||||
data=line[6:]
|
||||
if data=="[DONE]": break
|
||||
try: ev=json.loads(data)
|
||||
except ValueError: continue
|
||||
for ch in ev.get("choices",[]):
|
||||
d=ch.get("delta",{})
|
||||
txt=d.get("content")
|
||||
if not txt: continue # ping/ruolo/reasoning: non è testo
|
||||
if first: sp.stop(); first=False
|
||||
md.feed(txt); reply.append(txt)
|
||||
except KeyboardInterrupt:
|
||||
interrupted=True # il server annulla la richiesta alla disconnessione
|
||||
except OSError as e:
|
||||
sp.stop()
|
||||
print(f"\n {C.yel}[server unreachable: {e}]{C.r}"); break
|
||||
sp.stop(); md.close()
|
||||
if reply: msgs.append({"role":"assistant","content":"".join(reply)})
|
||||
else: msgs.pop() # turno vuoto: non sporcare la history
|
||||
el=time.time()-t0
|
||||
note=" · ⏹ interrupted" if interrupted else ""
|
||||
print(f"\r {C.dgray}└─ ~{len(''.join(reply))//4} tok · {el:.0f}s{note}{C.r}\n")
|
||||
print(f" {C.dim}goodbye — the engine keeps running for the next chat 🐦{C.r}")
|
||||
|
||||
def kv_resume_notice(model_dir):
|
||||
"""SERVE mode silently resumes .coli_kv from disk (glm.c kv_disk_load): a chat
|
||||
started today continues a conversation from days ago, with `first=0` so the
|
||||
turn is appended WITHOUT the [gMASK]<sop> prefix. The engine does announce it
|
||||
on stderr — but nothing here ever shows that: the drain thread's
|
||||
p.stderr.read() blocks until EOF, so on a healthy start errlog is still empty
|
||||
when the status lines are printed. The warning only appeared once the engine
|
||||
DIED, which is exactly when it no longer mattered.
|
||||
|
||||
Measured cost of the silence: a chat inherited 670 tokens of an old Italian
|
||||
session ("il mio numero preferito e 7, ricordalo!"). Every later reply came
|
||||
back in Italian, and "explain fibonacci in short" was answered about the
|
||||
number 7 — the model was being coherent with a context nobody could see, and
|
||||
it read as a quantization bug for a day.
|
||||
|
||||
So say it here, in Python, from the file itself: no pipe, no thread, no
|
||||
Windows deadlock risk (see the stderr comment below)."""
|
||||
p=os.path.join(model_dir, ".coli_kv")
|
||||
try:
|
||||
with open(p,"rb") as f:
|
||||
if f.read(8)!=b"COLIKV1\0": return
|
||||
h=struct.unpack("<8i", f.read(32))
|
||||
n=h[6]
|
||||
if n<1: return
|
||||
age=time.time()-os.path.getmtime(p)
|
||||
when=f"{age/86400:.0f}d ago" if age>86400 else f"{age/3600:.0f}h ago" if age>3600 else "just now"
|
||||
print(f" {C.yel}↺ resuming a saved conversation: {n} tokens, last written {when}{C.r}")
|
||||
print(f" {C.dgray} it steers tone, language and topic. :reset clears it · "
|
||||
f"KVSAVE=0 disables saving · delete {p} to start clean{C.r}")
|
||||
except (OSError, struct.error): pass
|
||||
|
||||
def cmd_chat(a):
|
||||
# ATTACH: a running `coli serve` beats a private engine every time — the load
|
||||
# (34-136 s) and the cache warmth survive between sessions. Explicit --attach
|
||||
# wins; otherwise probe localhost quietly and use it if it's there. --no-attach
|
||||
# forces the old behaviour. The probe costs ~1 ms when nothing is listening.
|
||||
if not getattr(a,"no_attach",False):
|
||||
base=getattr(a,"attach",None) or "http://127.0.0.1:8000"
|
||||
mid=server_probe(base, getattr(a,"api_key",None))
|
||||
if mid:
|
||||
banner(f"chat · {mid} · attached")
|
||||
chat_attached(a, base, mid); return
|
||||
if getattr(a,"attach",None):
|
||||
sys.exit(f"--attach: no coli serve answering at {base} (start one with: coli serve --model <dir>)")
|
||||
need_model(a.model)
|
||||
banner(f"chat · {os.path.basename(a.model)} · ram {a.ram or '-'}GB · topp {a.topp or 'off'}")
|
||||
kv_resume_notice(a.model)
|
||||
errlog=tempfile.NamedTemporaryFile(mode="w+", suffix=".log", delete=False)
|
||||
e=env_for(a); e["SERVE"]="1"
|
||||
# stderr -> PIPE, NOT stderr=errlog (file). On Windows/MinGW, pointing the
|
||||
@@ -619,12 +744,65 @@ def cmd_chat(a):
|
||||
except Exception: pass
|
||||
print(f" {C.teal}goodbye{C.r} {C.dim}— the hummingbird returns to its nest{C.r} 🐦\n")
|
||||
|
||||
def serve_pidfile(port): return os.path.join(tempfile.gettempdir(), f"coli-serve-{port}.pid")
|
||||
|
||||
def cmd_serve(a):
|
||||
need_model(a.model)
|
||||
# pidfile: cosi' `coli stop` spegne tutto con un comando, senza pkill a mano.
|
||||
# EN: pidfile so `coli stop` can shut everything down without manual pkill.
|
||||
try:
|
||||
with open(serve_pidfile(a.port),"w") as f: f.write(f"{os.getpid()} {a.model}\n")
|
||||
except OSError: pass
|
||||
from openai_server import serve
|
||||
serve(a.model, a.host, a.port, a.model_id, a.api_key,
|
||||
a.cap,a.ngen,GLM,env_for(a),a.cors_origin,
|
||||
a.max_queue,a.queue_timeout,a.kv_slots)
|
||||
try:
|
||||
serve(a.model, a.host, a.port, a.model_id, a.api_key,
|
||||
a.cap,a.ngen,GLM,env_for(a),a.cors_origin,
|
||||
a.max_queue,a.queue_timeout,a.kv_slots)
|
||||
finally:
|
||||
try: os.unlink(serve_pidfile(a.port))
|
||||
except OSError: pass
|
||||
|
||||
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`,
|
||||
not `glm`: every `pkill -x glm` in history silently killed nothing (that is
|
||||
how two 17+5 GB ghost engines OOM'd this box on 2026-07-16). This finds the
|
||||
real processes: the pidfile first, then /proc by cmdline/environ — only
|
||||
processes that are demonstrably ours (SERVE=1 + our SNAP, or `coli serve`
|
||||
in the command line)."""
|
||||
banner("stop")
|
||||
targets=[] # (pid, descrizione)
|
||||
pf=serve_pidfile(a.port)
|
||||
try:
|
||||
pid=int(open(pf).read().split()[0])
|
||||
os.kill(pid,0); targets.append((pid,f"coli serve (pidfile, port {a.port})"))
|
||||
except (OSError,ValueError,IndexError): pass
|
||||
for pd in os.listdir("/proc"):
|
||||
if not pd.isdigit(): continue
|
||||
pid=int(pd)
|
||||
try:
|
||||
cmd=open(f"/proc/{pd}/cmdline","rb").read().replace(b"\0",b" ").decode("utf-8","replace")
|
||||
if "coli" in cmd and " serve" in cmd and pid!=os.getpid():
|
||||
if not any(p==pid for p,_ in targets): targets.append((pid,"coli serve (cmdline)"))
|
||||
comm=open(f"/proc/{pd}/comm").read().strip()
|
||||
if comm in ("glm","exe","olmoe"):
|
||||
env=open(f"/proc/{pd}/environ","rb").read().replace(b"\0",b"\n").decode("utf-8","replace")
|
||||
if "SERVE=1" in env: targets.append((pid,f"engine `{comm}` (SERVE=1)"))
|
||||
except (OSError,PermissionError): continue
|
||||
if not targets:
|
||||
print(f" nothing running — no serve on port {a.port}, no SERVE engines"); return
|
||||
for pid,desc in targets: print(f" {'would stop' if a.dry_run else 'stopping'} {pid}: {desc}")
|
||||
if a.dry_run: return
|
||||
for pid,_ in targets:
|
||||
try: os.kill(pid, signal.SIGTERM)
|
||||
except OSError: pass
|
||||
time.sleep(2.0)
|
||||
for pid,_ in targets:
|
||||
try: os.kill(pid, signal.SIGKILL); print(f" {pid}: forced (SIGKILL)")
|
||||
except OSError: pass # gia' morto: bene
|
||||
try: os.unlink(pf)
|
||||
except OSError: pass
|
||||
print(f" {C.grn}✓ stopped{C.r} — RAM released")
|
||||
|
||||
def cmd_web(a):
|
||||
"""serve + open the dashboard in the browser once the API answers."""
|
||||
@@ -719,7 +897,14 @@ def main():
|
||||
pd=sub.add_parser("doctor",parents=[common])
|
||||
pd.add_argument("--json",action="store_true",help="emit a versioned JSON report")
|
||||
pr=sub.add_parser("run", parents=[common]); pr.add_argument("prompt", nargs="*")
|
||||
sub.add_parser("chat", parents=[common])
|
||||
pc=sub.add_parser("chat", parents=[common])
|
||||
pc.add_argument("--attach", nargs="?", const="http://127.0.0.1:8000", default=None,
|
||||
help="chat against a running `coli serve` instead of spawning an engine "
|
||||
"(keeps the model loaded and the expert cache warm across chat sessions). "
|
||||
"Bare --attach probes localhost:8000.")
|
||||
pc.add_argument("--no-attach", action="store_true",
|
||||
help="never auto-attach, always spawn a private engine")
|
||||
pc.add_argument("--api-key", default=os.environ.get("COLI_API_KEY"))
|
||||
ps=sub.add_parser("serve", parents=[common])
|
||||
ps.add_argument("--host",default="127.0.0.1"); ps.add_argument("--port",type=int,default=8000)
|
||||
ps.add_argument("--model-id",default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))
|
||||
@@ -728,6 +913,8 @@ 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")))
|
||||
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")
|
||||
for arg,kw in (("--host",dict(default="127.0.0.1")),("--port",dict(type=int,default=8000)),
|
||||
("--model-id",dict(default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))),
|
||||
@@ -751,7 +938,7 @@ def main():
|
||||
pc.add_argument("--no-mtp",action="store_true",help="skip the MTP head (no speculative drafts)")
|
||||
a=ap.parse_args()
|
||||
handler={"build":cmd_build,"info":cmd_info,"plan":cmd_plan,"doctor":cmd_doctor,
|
||||
"run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"bench":cmd_bench,
|
||||
"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 handler: sys.exit(handler(a) or 0)
|
||||
banner(); print(__doc__)
|
||||
|
||||
+10
-1
@@ -143,6 +143,12 @@ static inline int compat_fadvise(int fd, off_t off, off_t len, int advice){
|
||||
* Thread-safe (no shared seek position). Gestisce offset >4 GB e chunking
|
||||
* per letture >2 GB (anche se i tensori individuali sono nell'ordine dei
|
||||
* MB-centinaia di MB, il wrapper e' robusto per ogni taglia). */
|
||||
/* Ultimo GetLastError() di una ReadFile fallita, per thread: il chiamante
|
||||
* (pread_full in glm.c) lo stampa accanto a strerror. Senza questo, OGNI
|
||||
* fallimento Windows collassa in "EIO -> Input/output error" e la diagnosi
|
||||
* dal campo diventa un tirare a indovinare (#307: tre giri di ipotesi tra
|
||||
* tre persone perche' il codice vero non compariva da nessuna parte). */
|
||||
static __thread DWORD compat_pread_lasterr __attribute__((unused));
|
||||
static inline ssize_t compat_pread(int fd, void *buf, size_t n, off_t off){
|
||||
intptr_t osfh = _get_osfhandle(fd);
|
||||
if(osfh == -1 || osfh == -2){ errno = EBADF; return -1; }
|
||||
@@ -158,6 +164,7 @@ static inline ssize_t compat_pread(int fd, void *buf, size_t n, off_t off){
|
||||
if(!ReadFile(h, (char*)buf + total, chunk32, &rd, &ov)){
|
||||
DWORD err = GetLastError();
|
||||
if(err == ERROR_HANDLE_EOF) break; /* past EOF → return bytes read (0 if none, matching POSIX pread) */
|
||||
compat_pread_lasterr = err; /* preserva il codice VERO per il report (#307) */
|
||||
if(err == ERROR_INVALID_HANDLE || err == ERROR_INVALID_FUNCTION) errno = EBADF;
|
||||
else errno = EIO;
|
||||
return -1;
|
||||
@@ -237,7 +244,9 @@ static inline int compat_rename(const char *old, const char *new){
|
||||
/* --- rss_gb: getrusage -> GetProcessMemoryInfo ---
|
||||
* ru_maxrss in KB (come Linux): rss_gb() divide per 1e6 → GB corretti. */
|
||||
#include <psapi.h>
|
||||
#pragma comment(lib, "psapi.lib")
|
||||
#ifdef _MSC_VER
|
||||
#pragma comment(lib, "psapi.lib") /* MSVC: link psapi; MinGW/GCC uses -lpsapi */
|
||||
#endif
|
||||
struct rusage { long ru_maxrss; };
|
||||
#define RUSAGE_SELF 0
|
||||
static inline int getrusage(int who, struct rusage *r){
|
||||
|
||||
@@ -40,6 +40,9 @@
|
||||
#include <sys/stat.h> /* fstat per mmap degli shard (COLI_MMAP) */
|
||||
#include <signal.h> /* SIGINT = stop morbido del turno in serve mode */
|
||||
#endif
|
||||
#ifdef __linux__
|
||||
#include <sys/vfs.h> /* statfs: real fs-type check for the 9p warning (below) */
|
||||
#endif
|
||||
#if defined(_WIN32) && (defined(__x86_64__) || defined(__i386__))
|
||||
#include <cpuid.h> /* hwinfo_emit: CPU brand string senza /proc */
|
||||
#endif
|
||||
@@ -1207,7 +1210,9 @@ static int g_disk_split=0; /* DISK_SPLIT=1: contatori che spezzano i DISK LOAD (
|
||||
* 10x (#82) — hence per-region mbind here and nothing else. Raw syscall, no libnuma
|
||||
* dependency; MPOL_MF_MOVE migrates pages of reused heap chunks too. Linux-only,
|
||||
* silent no-op elsewhere or on single-node hosts. */
|
||||
static int g_numa_nodes=0;
|
||||
#ifdef __linux__
|
||||
static int g_numa_nodes=0; /* only touched under __linux__; off-Linux NUMA is a no-op */
|
||||
#endif
|
||||
static void numa_slab_bind(void *p, size_t n){
|
||||
#ifdef __linux__
|
||||
if(g_numa_nodes<2 || !p || !n) return;
|
||||
@@ -1732,8 +1737,14 @@ static int pread_full(int fd, void *buf, int64_t n, int64_t off, const char *tag
|
||||
while(got<n){
|
||||
ssize_t r=pread(fd, p+got, (size_t)(n-got), off+got);
|
||||
if(r<0){ if(errno==EINTR) continue;
|
||||
#ifdef _WIN32
|
||||
fprintf(stderr,"%s: %s (off %lld, %lld/%lld bytes, WinErr=%lu)\n",tag,strerror(errno),
|
||||
(long long)off,(long long)got,(long long)n,(unsigned long)compat_pread_lasterr);
|
||||
#else
|
||||
fprintf(stderr,"%s: %s (off %lld, %lld/%lld bytes)\n",tag,strerror(errno),
|
||||
(long long)off,(long long)got,(long long)n); return -1; }
|
||||
(long long)off,(long long)got,(long long)n);
|
||||
#endif
|
||||
return -1; }
|
||||
if(r==0){ fprintf(stderr,"%s: short read at EOF (off %lld, %lld/%lld bytes) — truncated shard?\n",
|
||||
tag,(long long)off,(long long)got,(long long)n); return -1; }
|
||||
got+=r;
|
||||
@@ -2304,6 +2315,43 @@ static int g_dsa_force=0; /* DSA_FORCE=1: selezione sempre attiva (test: top-min
|
||||
static int cmp_fdesc(const void *a,const void *b){
|
||||
float x=*(const float*)a, y=*(const float*)b; return x<y?1:x>y?-1:0; }
|
||||
|
||||
/* PARTIAL SELECT (quickselect, Hoare partition, DESCending). After this call the k
|
||||
* LARGEST elements of a[0..n) are in a[0..k) in unspecified order; the (k+1)-th and
|
||||
* beyond are untouched-or-smaller. O(n) average, O(n^2) pathological (mitigated by
|
||||
* median-of-three below) — and unlike a full qsort it never orders more than needed.
|
||||
*
|
||||
* Why this exists (#356): the DSA top-keep in attention_rows previously full-qsorted
|
||||
* all nk context scores (O(nk log nk)) per layer per token just to read ONE value --
|
||||
* the keep-th largest (the threshold). quickselect finds that pivot in O(nk) average,
|
||||
* and the position-order scans that build dst[] are unchanged, so the kept set is
|
||||
* bit-identical. Mirrors the sampling-side fix in #335 (heap partial-select there).
|
||||
*
|
||||
* NOT a stable partition: callers must derive the threshold and then re-scan the
|
||||
* ORIGINAL array (the DSA code does exactly this) rather than reading a[0..k). */
|
||||
static void partial_select_desc(float *a, int n, int k){
|
||||
if(k<=0) return;
|
||||
if(k>=n) return; /* nothing to partition: all kept */
|
||||
int lo=0, hi=n-1;
|
||||
while(lo<hi){
|
||||
/* median-of-three pivot to dodge the O(n^2) path on sorted/reverse input */
|
||||
int mid=lo+((hi-lo)>>1);
|
||||
if(a[mid]>a[lo]){ float t=a[lo]; a[lo]=a[mid]; a[mid]=t; }
|
||||
if(a[hi]>a[lo]){ float t=a[lo]; a[lo]=a[hi]; a[hi]=t; }
|
||||
if(a[mid]>a[hi]){ float t=a[hi]; a[hi]=a[mid]; a[mid]=t; }
|
||||
float piv=a[hi];
|
||||
int i=lo, j=hi;
|
||||
for(;;){
|
||||
while(a[i]>piv) i++; /* desc: large values go left */
|
||||
while(j>lo && a[j]<piv) j--;
|
||||
if(i>=j) break;
|
||||
float t=a[i]; a[i]=a[j]; a[j]=t; i++; if(i>j) break; j--;
|
||||
}
|
||||
/* partition point: a[lo..i) are all >= piv, a[i..hi] are all <= piv */
|
||||
if(k<=i-1) hi=i-1; /* the k-th largest is in the left partition */
|
||||
else lo=i; /* it's in the right partition */
|
||||
}
|
||||
}
|
||||
|
||||
/* attenzione MLA con KV-cache compressa, su token nuovi x[S,hidden], pos_base = pos del primo */
|
||||
/* kvs/pos describe a ragged decode batch: each row may belong to a different
|
||||
* sequence. NULL keeps the original contiguous, currently-bound KV path. */
|
||||
@@ -2586,10 +2634,14 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p
|
||||
}
|
||||
isc[t]=a*wsc;
|
||||
}
|
||||
/* top-keep: soglia via qsort desc, poi scan in ordine di posizione */
|
||||
/* top-keep: threshold via PARTIAL SELECT (#356), poi scan in ordine di posizione.
|
||||
* Era un qsort completo su nk (O(nk log nk)); quickselect estrae solo il
|
||||
* keep-esimo valore piu' grande in O(nk) medio. La soglia (= min del blocco
|
||||
* dei keep maggiori) e' identica a tmp[keep-1] del vecchio qsort, quindi i
|
||||
* due scan qui sotto costruiscono dst[] bit-identical. */
|
||||
float *tmp=falloc(nk); memcpy(tmp,isc,nk*sizeof(float));
|
||||
qsort(tmp,nk,sizeof(float),cmp_fdesc);
|
||||
float thr=tmp[keep-1];
|
||||
partial_select_desc(tmp,nk,keep);
|
||||
float thr=tmp[0]; for(int t=1;t<keep;t++) if(tmp[t]<thr) thr=tmp[t];
|
||||
int *dst=m->dsa_sel+(int64_t)s*dtopk, nd=0;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
|
||||
@@ -4186,10 +4238,21 @@ static uint64_t g_rng=0x9E3779B97F4A7C15ULL;
|
||||
static inline double rndu(void){ g_rng^=g_rng<<13; g_rng^=g_rng>>7; g_rng^=g_rng<<17;
|
||||
return (double)(g_rng>>11)*(1.0/9007199254740992.0); }
|
||||
static float *g_pbuf=NULL; static int *g_pidx=NULL; /* buffer riusati (decode single-thread) */
|
||||
static int cmp_pdesc(const void *a,const void *b){
|
||||
float pa=g_pbuf[*(const int*)a], pb=g_pbuf[*(const int*)b];
|
||||
return pa<pb ? 1 : pa>pb ? -1 : 0; }
|
||||
/* costruisce in g_pbuf la distribuzione target: softmax(lo/temp) troncata a top-p g_nuc */
|
||||
/* sift-down su max-heap in h[0..n), chiave = g_pbuf[h[i]] (#335: partial top-p select).
|
||||
* Versione "a buco": porta il valore di radice e lo deposita solo alla fine, cosi'
|
||||
* heapify e' O(V) e ogni pop e' O(log n) senza qsort sull'intero vocabolario. */
|
||||
static void topp_siftdown(int *h, int n, int i){
|
||||
int iv=h[i]; float kv=g_pbuf[iv];
|
||||
for(;;){ int l=2*i+1;
|
||||
if(l>=n) break; /* foglia */
|
||||
int b=l; if(l+1<n && g_pbuf[h[l+1]]>g_pbuf[h[l]]) b=l+1; /* figlio maggiore */
|
||||
if(g_pbuf[h[b]]<=kv) break; /* nessun figlio supera la radice -> ferma */
|
||||
h[i]=h[b]; i=b; }
|
||||
h[i]=iv;
|
||||
}
|
||||
/* costruisce in g_pbuf la distribuzione target: softmax(lo/temp) troncata a top-p g_nuc.
|
||||
* Invariante per dist_sample: g_pbuf resta INDICIZZATO per token-id (mai riordinato);
|
||||
* la coda troncata va AZZERATA in g_pbuf (dist_sample la legge direttamente per id). */
|
||||
static void dist_build(const float *lo, int V){
|
||||
if(!g_pbuf){ g_pbuf=falloc(V); g_pidx=malloc(V*sizeof(int)); }
|
||||
float mx=lo[0]; for(int i=1;i<V;i++) if(lo[i]>mx) mx=lo[i];
|
||||
@@ -4198,12 +4261,19 @@ static void dist_build(const float *lo, int V){
|
||||
for(int i=0;i<V;i++) g_pbuf[i]/=(float)s;
|
||||
if(g_nuc>0 && g_nuc<1.f){
|
||||
for(int i=0;i<V;i++) g_pidx[i]=i;
|
||||
qsort(g_pidx,V,sizeof(int),cmp_pdesc);
|
||||
double cum=0; int keep=V;
|
||||
for(int i=0;i<V;i++){ cum+=g_pbuf[g_pidx[i]]; if(cum>=g_nuc){ keep=i+1; break; } }
|
||||
double s2=0; for(int i=keep;i<V;i++) g_pbuf[g_pidx[i]]=0;
|
||||
for(int i=0;i<keep;i++) s2+=g_pbuf[g_pidx[i]];
|
||||
for(int i=0;i<keep;i++) g_pbuf[g_pidx[i]]/=(float)s2;
|
||||
for(int i=V/2-1;i>=0;i--) topp_siftdown(g_pidx,V,i); /* heapify O(V) */
|
||||
/* pop verso la coda: i vincitori (testa top-p) cadono in g_pidx[out..V-1] in ordine
|
||||
* DECRESCENTE, come il vecchio qsort, quindi s2 accumula nello stesso ordine ->
|
||||
* head bit-identical sui casi senza pareggi (i pareggi erano gia' non specificati
|
||||
* sotto il qsort instabile e restano tali). Il prefisso g_pidx[0..out-1) e' la coda. */
|
||||
double s2=0, cum=0; int out=V;
|
||||
do{ int root=g_pidx[0]; /* massimo corrente */
|
||||
g_pidx[0]=g_pidx[--out]; g_pidx[out]=root; /* sposta il max in coda */
|
||||
s2+=g_pbuf[root]; cum+=g_pbuf[root];
|
||||
if(out>0) topp_siftdown(g_pidx,out,0);
|
||||
} while(cum<g_nuc && out>0);
|
||||
for(int i=0;i<out;i++) g_pbuf[g_pidx[i]]=0; /* azzera la coda (invariante) */
|
||||
float s2f=(float)s2; for(int i=out;i<V;i++) g_pbuf[g_pidx[i]]/=s2f; /* rinormalizza */
|
||||
}
|
||||
}
|
||||
/* campiona da g_pbuf; ban>=0 -> quel token e' escluso (rinormalizzando al volo) */
|
||||
@@ -6000,6 +6070,14 @@ int main(int argc, char **argv){
|
||||
!getenv("COLI_CUDA") && !getenv("COLI_METAL")){
|
||||
setenv("OMP_WAIT_POLICY","active",0); /* keep the team hot across the tiny per-expert matmul regions */
|
||||
setenv("GOMP_SPINCOUNT","200000",0); /* spin briefly, then yield so long disk waits don't burn a core */
|
||||
/* LLVM libomp (clang builds: FreeBSD cc, macOS, some Linux setups) does not
|
||||
* read GOMP_*: with OMP_WAIT_POLICY=active it sets KMP_BLOCKTIME=infinite,
|
||||
* so the idle team SPINS FOREVER once generation ends — a serve-mode engine
|
||||
* parked on stdin burns ~100% x nthreads (#341, measured 3000% on FreeBSD).
|
||||
* 200 ms of blocktime keeps the team hot across back-to-back expert matmuls
|
||||
* and lets it sleep at the prompt. libgomp ignores KMP_*; overwrite=0 keeps
|
||||
* the user's own setting authoritative. */
|
||||
setenv("KMP_BLOCKTIME","200",0);
|
||||
setenv("OMP_PROC_BIND","close",0); /* pack the team onto adjacent cores for cache locality */
|
||||
setenv("OMP_DYNAMIC","FALSE",0); /* fixed team size: no per-region thread-count churn */
|
||||
setenv("COLI_OMP_TUNED","1",1);
|
||||
@@ -6230,9 +6308,16 @@ int main(int argc, char **argv){
|
||||
m.has_mtp?"ACTIVE":"absent", g_draft);
|
||||
/* anche su stderr: e' il canale che le UI (coli) mostrano all'utente */
|
||||
fprintf(stderr,"[MTP] %s (draft=%d)\n", m.has_mtp?"active: native speculative decoding":"absent", g_draft);
|
||||
if(!strncmp(snap,"/mnt/",5))
|
||||
fprintf(stderr,"WARNING: the model is on %s (slow 9p/Windows filesystem; fadvise is ineffective).\n"
|
||||
" Keep it on ext4 (for example, /home/...) for memory efficiency and speed.\n", snap);
|
||||
#ifdef __linux__
|
||||
{ /* Only warn for a GENUINE 9p mount (WSL Windows drives, magic 0x01021997), where
|
||||
* fadvise is a no-op. The old check was `snap` starting with "/mnt/", which
|
||||
* false-positives on native-Linux ZFS/ext4/xfs/NFS mounts that also live under /mnt. */
|
||||
struct statfs sfb;
|
||||
if(statfs(snap,&sfb)==0 && (unsigned long)sfb.f_type==0x01021997UL)
|
||||
fprintf(stderr,"WARNING: the model is on %s (9p/Windows filesystem; fadvise is ineffective).\n"
|
||||
" Keep it on a native Linux fs (ext4/xfs/zfs) for memory efficiency and speed.\n", snap);
|
||||
}
|
||||
#endif
|
||||
/* HOT-STORE: PIN=<statsfile> [PIN_GB=g] -> top expert per frequenza fissi in RAM.
|
||||
* Va PRIMA di cap_for_ram: i pinnati contano nel residente. */
|
||||
if(getenv("PIN")){
|
||||
|
||||
@@ -400,6 +400,37 @@ static void generate(Model *m, const int *prompt, int np, int n_new, int *out) {
|
||||
}
|
||||
}
|
||||
|
||||
/* teacher-forced NLL of full_ids[np..nfull): feed the REFERENCE token at each step
|
||||
* (never the argmax), accumulate -log softmax(logits)[next_ref]. A loss meter for
|
||||
* throughput experiments: same engine path as decode, so hit rate/speed stay
|
||||
* comparable, but quality is measured as perplexity instead of exact-match.
|
||||
* Cross-checked vs HF transformers bf16 on identical token ids: engine (int8
|
||||
* experts) 12.11 ppl vs reference 12.25 (#108). Enabled by PPL=1. */
|
||||
static int tf_nll(Model *m, const int *full, int nfull, int np, double *nll_out) {
|
||||
Cfg *c = &m->c;
|
||||
m->max_t = nfull;
|
||||
m->K = calloc(c->n_layers, sizeof(float*)); m->V = calloc(c->n_layers, sizeof(float*));
|
||||
for (int i = 0; i < c->n_layers; i++) {
|
||||
m->K[i] = falloc((int64_t)c->n_heads * m->max_t * c->head_dim);
|
||||
m->V[i] = falloc((int64_t)c->n_heads * m->max_t * c->head_dim);
|
||||
}
|
||||
double nll = 0; int scored = 0;
|
||||
float *logit = step(m, full, np, 0); /* prefill on the prompt */
|
||||
for (int i = np; i < nfull; i++) {
|
||||
/* log softmax(logit)[full[i]] without materializing the softmax */
|
||||
float mx = logit[0]; for (int v = 1; v < c->vocab; v++) if (logit[v] > mx) mx = logit[v];
|
||||
double Z = 0; for (int v = 0; v < c->vocab; v++) Z += exp((double)logit[v] - mx);
|
||||
nll += -((double)logit[full[i]] - mx - log(Z));
|
||||
scored++;
|
||||
free(logit); logit = NULL;
|
||||
if (i == nfull - 1) break;
|
||||
logit = step(m, &full[i], 1, i); /* teacher forcing */
|
||||
}
|
||||
if (logit) free(logit);
|
||||
*nll_out = nll / scored;
|
||||
return scored;
|
||||
}
|
||||
|
||||
/* ---------- lettura ref.json ---------- */
|
||||
static int *read_int_array(jval *o, const char *key, int *n_out) {
|
||||
jval *a = json_get(o, key);
|
||||
@@ -430,6 +461,19 @@ int main(int argc, char **argv) {
|
||||
Model m; model_init(&m, snap, cap, bits);
|
||||
printf("resident weights loaded in %.1fs | RSS after load: %.2f GB\n", m.dense_load_s, rss_gb());
|
||||
|
||||
if (getenv("PPL") && atoi(getenv("PPL")) == 1) { /* loss-meter mode: teacher-forced NLL */
|
||||
double nll; double t = now_s();
|
||||
int scored = tf_nll(&m, full, nfull, np, &nll);
|
||||
double dt = now_s() - t;
|
||||
double tot = m.hits + m.miss;
|
||||
printf("TF-NLL: %.4f nats/token over %d tokens | ppl = %.2f\n", nll, scored, exp(nll));
|
||||
printf("Expert cache hit rate: %.1f%% (hit=%llu miss=%llu)\n", tot?100.0*m.hits/tot:0.0,
|
||||
(unsigned long long)m.hits, (unsigned long long)m.miss);
|
||||
printf("Speed: %.2f tok/s (%.1fs for %d tokens) | PEAK RSS: %.2f GB\n", scored/dt, dt, scored, rss_gb());
|
||||
free(buf); free(arena);
|
||||
return 0;
|
||||
}
|
||||
|
||||
int *out = malloc((np + n_new) * sizeof(int));
|
||||
double t = now_s();
|
||||
generate(&m, prompt, np, n_new, out);
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
#include <string.h>
|
||||
#include <stdint.h>
|
||||
#include <fcntl.h>
|
||||
#include <errno.h>
|
||||
#include <unistd.h>
|
||||
#include <dirent.h>
|
||||
#include <sys/stat.h>
|
||||
@@ -104,6 +105,38 @@ static int st_direct_fd(shards *S, int fd) {
|
||||
}
|
||||
|
||||
/* indicizza tutti i model-*.safetensors in snap_dir */
|
||||
/* pread completo: chunk-loop (una singola pread si ferma a ~2^31 byte su Linux
|
||||
* — i tensori bf16 grandi la superano), riprova su EINTR e riporta un errore
|
||||
* ONESTO: perror stampava "Success" su una short-read (errno resta 0), lo
|
||||
* stesso sintomo corretto in glm.c per #236. ST_PREAD_CHUNK e' sovrascrivibile
|
||||
* per i test. EN: full pread — chunk loop (one pread caps at ~2^31 bytes and
|
||||
* big bf16 tensors exceed it), EINTR retry, honest short-read errors.
|
||||
* Exits on failure, like every st.h reader. */
|
||||
#ifndef ST_PREAD_CHUNK
|
||||
#define ST_PREAD_CHUNK (1u << 30)
|
||||
#endif
|
||||
static void st_pread_full(int fd, void *buf, int64_t n, int64_t off, const char *tag) {
|
||||
char *p = (char *)buf;
|
||||
int64_t got = 0;
|
||||
while (got < n) {
|
||||
int64_t want = n - got;
|
||||
if (want > (int64_t)ST_PREAD_CHUNK) want = ST_PREAD_CHUNK;
|
||||
ssize_t r = pread(fd, p + got, (size_t)want, off + got);
|
||||
if (r < 0) {
|
||||
if (errno == EINTR) continue;
|
||||
fprintf(stderr, "%s: %s (off %lld, %lld/%lld bytes)\n", tag, strerror(errno),
|
||||
(long long)off, (long long)got, (long long)n);
|
||||
exit(1);
|
||||
}
|
||||
if (r == 0) {
|
||||
fprintf(stderr, "%s: short read at EOF (off %lld, %lld/%lld bytes) — truncated file?\n",
|
||||
tag, (long long)off, (long long)got, (long long)n);
|
||||
exit(1);
|
||||
}
|
||||
got += r;
|
||||
}
|
||||
}
|
||||
|
||||
static void st_init(shards *S, const char *snap_dir) {
|
||||
memset(S, 0, sizeof(*S));
|
||||
S->cap = 4096; S->t = calloc(S->cap, sizeof(st_tensor));
|
||||
@@ -128,7 +161,7 @@ static void st_init(shards *S, const char *snap_dir) {
|
||||
if (fstat(fd, &sst) != 0) { perror("fstat shard"); exit(1); }
|
||||
int64_t fsz = (int64_t)sst.st_size;
|
||||
uint64_t hlen;
|
||||
if (pread(fd, &hlen, 8, 0) != 8) { perror("pread hlen"); exit(1); }
|
||||
st_pread_full(fd, &hlen, 8, 0, "pread hlen");
|
||||
/* file malevolo/troncato: hlen deve stare nel file dopo gli 8 byte di
|
||||
* prefisso e sotto il tetto. Senza questo bound hlen+1 puo' andare in
|
||||
* overflow (malloc(0) e poi hdr[hlen]=0 fuori limiti) o forzare una
|
||||
@@ -138,7 +171,7 @@ static void st_init(shards *S, const char *snap_dir) {
|
||||
files[fi], (unsigned long long)hlen, (long long)fsz); exit(1); }
|
||||
char *hdr = malloc(hlen + 1);
|
||||
if (!hdr) { perror("malloc safetensors header"); exit(1); }
|
||||
if (pread(fd, hdr, hlen, 8) != (ssize_t)hlen) { perror("pread hdr"); exit(1); }
|
||||
st_pread_full(fd, hdr, (int64_t)hlen, 8, "pread hdr");
|
||||
hdr[hlen] = 0;
|
||||
int64_t data_start = 8 + (int64_t)hlen;
|
||||
char *arena = NULL;
|
||||
@@ -218,7 +251,7 @@ static int64_t st_read_f32(shards *S, const char *name, float *out, int drop) {
|
||||
if (!t) { fprintf(stderr, "missing tensor: %s\n", name); exit(1); }
|
||||
void *raw = malloc(t->nbytes);
|
||||
if (!raw) { fprintf(stderr, "malloc %lld bytes for tensor %s failed\n", (long long)t->nbytes, name); exit(1); }
|
||||
if (pread(t->fd, raw, t->nbytes, t->off) != t->nbytes) { perror("pread data"); exit(1); }
|
||||
st_pread_full(t->fd, raw, t->nbytes, t->off, "pread data");
|
||||
if (t->dtype == 2) {
|
||||
memcpy(out, raw, t->nbytes);
|
||||
} else if (t->dtype == 0) {
|
||||
@@ -243,7 +276,7 @@ static int64_t st_nbytes(shards *S, const char *name) {
|
||||
static void st_read_raw(shards *S, const char *name, void *out, int drop) {
|
||||
st_tensor *t = st_find(S, name);
|
||||
if (!t) { fprintf(stderr, "missing tensor: %s\n", name); exit(1); }
|
||||
if (pread(t->fd, out, t->nbytes, t->off) != t->nbytes) { perror("pread raw"); exit(1); }
|
||||
st_pread_full(t->fd, out, t->nbytes, t->off, "pread raw");
|
||||
if (drop) posix_fadvise(t->fd, t->off, t->nbytes, POSIX_FADV_DONTNEED);
|
||||
}
|
||||
|
||||
@@ -256,7 +289,7 @@ static void st_read_slice_f32(shards *S, const char *name, int64_t elem_off, int
|
||||
int esz = (t->dtype == 2) ? 4 : 2;
|
||||
int64_t boff = t->off + elem_off * esz, nb = n_elems * esz;
|
||||
void *raw = malloc(nb);
|
||||
if (pread(t->fd, raw, nb, boff) != nb) { perror("pread slice"); exit(1); }
|
||||
st_pread_full(t->fd, raw, nb, boff, "pread slice");
|
||||
if (t->dtype == 2) memcpy(out, raw, nb);
|
||||
else if (t->dtype == 0) { uint16_t *p = raw; for (int64_t i = 0; i < n_elems; i++) out[i] = bf16_to_f32(p[i]); }
|
||||
else { uint16_t *p = raw; for (int64_t i = 0; i < n_elems; i++) out[i] = f16_to_f32(p[i]); }
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
/* Microbenchmark: old (full-qsort) vs new (quickselect partial-select) DSA top-keep.
|
||||
*
|
||||
* This is NOT a unit test -- test_dsa_select.c proves correctness. This measures the
|
||||
* headline claim of #356: that replacing the O(nk log nk) qsort over all nk context
|
||||
* scores with an O(nk) partial_select_desc is materially faster per call, which is
|
||||
* the win the issue was opened for -- and that the win GROWS with context length
|
||||
* (because quickselect is linear average, qsort is n-log-n).
|
||||
*
|
||||
* It re-implements the OLD top-keep inline (qsort + threshold + scans) on a private
|
||||
* buffer so the A/B runs in one process, same inputs, same warm caches -- a controlled
|
||||
* comparison. It calls the REAL (new) partial_select_desc via the include-glm.c
|
||||
* pattern, replicating the production threshold derivation + scans.
|
||||
*
|
||||
* Methodology (chosen to be honest, not to flatter the change):
|
||||
* - keep = 2048 (the real GLM-5.2 index_topk), nk swept across context lengths from
|
||||
* the 2049 activation boundary up to 65536 (a long conversation).
|
||||
* - Three score shapes: (a) realistic peaked -- a few hot keys, long tail, the shape
|
||||
* real DSA attention scores take; (b) uniform random -- no structure; (c) a plateau
|
||||
* of ties, to exercise the boundary-membership path.
|
||||
* - Each (shape, nk) is timed over N_REPEAT=2000 iterations, with the scores frozen
|
||||
* so both algorithms do IDENTICAL work. We report median ns/call and the new/old
|
||||
* ratio. A warmup pass primes caches before timing.
|
||||
*
|
||||
* Run: make tests/bench_dsa_select && ./tests/bench_dsa_select (not in TEST_BINS)
|
||||
*/
|
||||
#define main coli_glm_main_unused
|
||||
#include "../glm.c"
|
||||
#undef main
|
||||
|
||||
#include <math.h>
|
||||
#include <stdint.h>
|
||||
|
||||
/* ---- the OLD algorithm, verbatim from dev before #356, on a private buffer ---- */
|
||||
static int cmp_pdesc_old(const void *a, const void *b){
|
||||
float x=*(const float*)a, y=*(const float*)b; return x<y?1:x>y?-1:0; }
|
||||
static void keep_old(const float *isc, int nk, int keep, int *dst, int *nd_out){
|
||||
float *tmp=malloc((size_t)nk*sizeof(float));
|
||||
memcpy(tmp,isc,(size_t)nk*sizeof(float));
|
||||
qsort(tmp,(size_t)nk,sizeof(float),cmp_pdesc_old);
|
||||
float thr=tmp[keep-1]; int nd=0;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
|
||||
free(tmp); *nd_out=nd;
|
||||
}
|
||||
|
||||
/* ---- timing: median of N_REPEAT runs in ns/call, sorted ascending ---- */
|
||||
#define N_REPEAT 2000
|
||||
static double bench_ns(void (*fn)(const float*,int,int,int*,int*),
|
||||
const float *isc, int nk, int keep){
|
||||
static double ts[N_REPEAT]; int *dst=malloc((size_t)nk*sizeof(int)); int nd;
|
||||
for(int r=0;r<N_REPEAT;r++){
|
||||
double t0=now_s();
|
||||
fn(isc,nk,keep,dst,&nd);
|
||||
ts[r]=(now_s()-t0)*1e9;
|
||||
}
|
||||
for(int a=1;a<N_REPEAT;a++){ double k=ts[a]; int b=a-1;
|
||||
while(b>=0 && ts[b]>k){ ts[b+1]=ts[b]; b--; } ts[b+1]=k; }
|
||||
free(dst);
|
||||
return ts[N_REPEAT/2];
|
||||
}
|
||||
|
||||
/* the NEW algorithm calls the real partial_select_desc + replicates the production
|
||||
* threshold derivation and position scans. */
|
||||
static void keep_new(const float *isc, int nk, int keep, int *dst, int *nd_out){
|
||||
float *tmp=malloc((size_t)nk*sizeof(float));
|
||||
memcpy(tmp,isc,(size_t)nk*sizeof(float));
|
||||
partial_select_desc(tmp,nk,keep);
|
||||
float thr=tmp[0]; for(int t=1;t<keep;t++) if(tmp[t]<thr) thr=tmp[t];
|
||||
int nd=0;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
|
||||
free(tmp); *nd_out=nd;
|
||||
}
|
||||
|
||||
/* deterministic score fill for three shapes */
|
||||
static uint32_t brng = 0xA5A5A5A5u;
|
||||
static double brand(void){ brng ^= brng << 13; brng ^= brng >> 17; brng ^= brng << 5;
|
||||
return (double)(brng >> 8) * (1.0 / 16777216.0); }
|
||||
static void fill_realistic(float *isc, int nk){ /* few hot, long distinct tail */
|
||||
for(int i=0;i<nk;i++) isc[i]=(float)(-1.0 - brand()*4.0);
|
||||
isc[0]=3.f; isc[nk/50<nk?nk/50:nk-1]=1.f; isc[nk/200<nk?nk/200:nk-1]=0.5f;
|
||||
}
|
||||
static void fill_uniform(float *isc, int nk){ /* no structure */
|
||||
for(int i=0;i<nk;i++) isc[i]=(float)(brand()*1000.0);
|
||||
}
|
||||
static void fill_plateau(float *isc, int nk){ /* tie blocks -> boundary path */
|
||||
for(int i=0;i<nk;i++) isc[i]=(float)(-(double)(i/7));
|
||||
}
|
||||
|
||||
int main(void){
|
||||
int keep = 2048; /* GLM-5.2 index_topk */
|
||||
int nks[] = {2049, 4096, 8192, 16384, 32768, 65536};
|
||||
float *isc = malloc((size_t)65536*sizeof(float));
|
||||
int *dst = malloc((size_t)65536*sizeof(int)); int nd;
|
||||
|
||||
struct { const char *name; void (*fill)(float*,int); } shapes[] = {
|
||||
{ "realistic", fill_realistic },
|
||||
{ "uniform", fill_uniform },
|
||||
{ "plateau", fill_plateau },
|
||||
};
|
||||
|
||||
printf("bench_dsa_select: DSA top-keep, old (qsort) vs new (partial-select) keep=%d\n", keep);
|
||||
printf("%-12s %7s %14s %14s %9s\n", "shape", "nk", "old ns/call", "new ns/call", "speedup");
|
||||
printf("------------------------------------------------------------------------\n");
|
||||
|
||||
for(size_t sh=0; sh<sizeof(shapes)/sizeof(shapes[0]); sh++){
|
||||
for(size_t ni=0; ni<sizeof(nks)/sizeof(nks[0]); ni++){
|
||||
int nk=nks[ni];
|
||||
shapes[sh].fill(isc,nk);
|
||||
/* warmup both paths so caches/branch predictors are primed */
|
||||
for(int w=0; w<50; w++){ keep_old(isc,nk,keep,dst,&nd); keep_new(isc,nk,keep,dst,&nd); }
|
||||
/* sanity: both must keep exactly `keep` (correctness is test_dsa_select's
|
||||
* job, but a count divergence here would make the timing meaningless) */
|
||||
keep_old(isc,nk,keep,dst,&nd); int na=nd;
|
||||
keep_new(isc,nk,keep,dst,&nd); int nb=nd;
|
||||
if(na!=keep || nb!=keep){
|
||||
printf("%-12s %7d (BAD COUNTS: old=%d new=%d, skipped)\n",
|
||||
shapes[sh].name, nk, na, nb);
|
||||
continue;
|
||||
}
|
||||
double t_old=bench_ns(keep_old,isc,nk,keep);
|
||||
double t_new=bench_ns(keep_new,isc,nk,keep);
|
||||
printf("%-12s %7d %14.0f %14.0f %8.2fx\n",
|
||||
shapes[sh].name, nk, t_old, t_new, t_old/t_new);
|
||||
}
|
||||
printf("\n");
|
||||
}
|
||||
|
||||
printf("bench_dsa_select: done (lower ns is better; speedup = old/new)\n");
|
||||
free(isc); free(dst);
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
/* Microbenchmark: old (full-vocab qsort) vs new (heap partial-select) top-p truncation.
|
||||
*
|
||||
* This is NOT a unit test -- test_topp.c proves correctness. This measures the headline
|
||||
* claim of #335: that replacing the O(V log V) qsort over all 151936 vocab entries with
|
||||
* an O(V) heapify + k*log-V pops is materially faster per call, which is the win the issue
|
||||
* was opened for.
|
||||
*
|
||||
* It re-implements the OLD dist_build inline (qsort + scan) on a private buffer so the A/B
|
||||
* runs in one process, same inputs, same warm caches -- a controlled comparison. It calls
|
||||
* the REAL (new) dist_build via the include-glm.c pattern on the global g_pbuf.
|
||||
*
|
||||
* Methodology (chosen to be honest, not to flatter the change):
|
||||
* - V = 151936 (the actual GLM-5.2 vocab), g_nuc swept across the values that matter
|
||||
* for serving: 0.5 / 0.9 (serve default) / 0.95 / 0.99.
|
||||
* - Three logit shapes: (a) realistic peaked -- one hot token, long exponential tail,
|
||||
* the shape real language-model logits take; (b) uniform -- worst case for the heap,
|
||||
* maximum pop count; (c) a plateau of ties, to exercise the tie path.
|
||||
* - Each (shape, nuc) is timed over N_REPEAT=2000 iterations, with the RNG/logits frozen
|
||||
* so both algorithms do IDENTICAL work. We report median ns/call and the new/old ratio.
|
||||
* - A warmup pass primes caches before timing.
|
||||
*
|
||||
* Run: make tests/bench_topp && ./tests/bench_topp (not in TEST_BINS -- not a gate)
|
||||
*/
|
||||
#define main coli_glm_main_unused
|
||||
#include "../glm.c"
|
||||
#undef main
|
||||
|
||||
#include <math.h>
|
||||
#include <stdint.h>
|
||||
|
||||
/* ---- the OLD algorithm, verbatim from dev before #354, on a private buffer ---- */
|
||||
static float *s_pbuf; static int *s_pidx; static double *s_ref;
|
||||
static int cmp_pdesc_old(const void *a, const void *b){
|
||||
double pa = s_ref[*(const int*)a], pb = s_ref[*(const int*)b];
|
||||
return pa < pb ? 1 : pa > pb ? -1 : 0; }
|
||||
static void dist_build_old(const float *lo, int V, double temp, double nuc){
|
||||
double mx = lo[0]; for (int i = 1; i < V; i++) if (lo[i] > mx) mx = lo[i];
|
||||
double s = 0, invt = 1.0 / (temp > 1e-4 ? temp : 1e-4);
|
||||
for (int i = 0; i < V; i++){ s_ref[i] = exp((lo[i]-mx)*invt); s += s_ref[i]; }
|
||||
for (int i = 0; i < V; i++) s_ref[i] /= s;
|
||||
if (nuc > 0 && nuc < 1.0){
|
||||
for (int i = 0; i < V; i++) s_pidx[i] = i;
|
||||
qsort(s_pidx, V, sizeof(int), cmp_pdesc_old);
|
||||
double cum = 0; int keep = V;
|
||||
for (int i = 0; i < V; i++){ cum += s_ref[s_pidx[i]]; if (cum >= nuc){ keep = i+1; break; } }
|
||||
double s2 = 0;
|
||||
for (int i = keep; i < V; i++) s_ref[s_pidx[i]] = 0;
|
||||
for (int i = 0; i < keep; i++) s2 += s_ref[s_pidx[i]];
|
||||
for (int i = 0; i < keep; i++) s_ref[s_pidx[i]] /= s2;
|
||||
}
|
||||
(void)s_pbuf;
|
||||
}
|
||||
|
||||
/* ---- timing: median of N_REPEAT runs in ns/call, sorted ascending ---- */
|
||||
#define N_REPEAT 2000
|
||||
static double bench_ns(void (*fn)(const float*,int,double,double),
|
||||
const float *lo, int V, double temp, double nuc){
|
||||
static double ts[N_REPEAT];
|
||||
for (int r = 0; r < N_REPEAT; r++){
|
||||
double t0 = now_s();
|
||||
fn(lo, V, temp, nuc);
|
||||
ts[r] = (now_s() - t0) * 1e9;
|
||||
}
|
||||
/* insertion sort the N_REPEAT samples (small), take median */
|
||||
for (int a = 1; a < N_REPEAT; a++){ double k = ts[a]; int b = a-1;
|
||||
while (b >= 0 && ts[b] > k){ ts[b+1] = ts[b]; b--; } ts[b+1] = k; }
|
||||
return ts[N_REPEAT/2];
|
||||
}
|
||||
|
||||
/* the NEW algorithm is the real dist_build, but it writes g_pbuf (not a private buf).
|
||||
* Wrap it so the bench signature matches, and set the globals it reads. */
|
||||
static void dist_build_new(const float *lo, int V, double temp, double nuc){
|
||||
g_temp = (float)temp; g_nuc = (float)nuc;
|
||||
dist_build(lo, V);
|
||||
}
|
||||
|
||||
/* deterministic logit fill for three shapes */
|
||||
static uint32_t brng = 0xA5A5A5A5u;
|
||||
static double brand(void){ brng ^= brng << 13; brng ^= brng >> 17; brng ^= brng << 5;
|
||||
return (double)(brng >> 8) * (1.0 / 16777216.0); }
|
||||
static void fill_realistic(float *lo, int V){ /* one hot, exponential tail -- like real logits */
|
||||
for (int i = 0; i < V; i++) lo[i] = (float)(-4.0 * brand() - (double)i * 0.0001);
|
||||
lo[0] = 6.f; lo[V/50] = 4.f; lo[V/200] = 3.f;
|
||||
}
|
||||
static void fill_uniform(float *lo, int V){ /* worst case for the heap: max pop count */
|
||||
for (int i = 0; i < V; i++) lo[i] = 0.f;
|
||||
}
|
||||
static void fill_plateau(float *lo, int V){ /* ties: blocks of equal value */
|
||||
for (int i = 0; i < V; i++) lo[i] = (float)(-(double)(i / 50));
|
||||
}
|
||||
|
||||
int main(void){
|
||||
int V = 151936;
|
||||
float *lo = malloc((size_t)V * sizeof(float));
|
||||
s_ref = malloc((size_t)V * sizeof(double));
|
||||
s_pidx = malloc((size_t)V * sizeof(int));
|
||||
/* force the new dist_build to allocate g_pbuf/g_pidx at full V once */
|
||||
g_temp = 0.7f; g_nuc = 0.9f; dist_build(lo, V);
|
||||
|
||||
double temp = 0.7;
|
||||
struct { const char *name; void (*fill)(float*,int); } shapes[] = {
|
||||
{ "realistic", fill_realistic },
|
||||
{ "uniform", fill_uniform },
|
||||
{ "plateau", fill_plateau },
|
||||
};
|
||||
double nucs[] = { 0.5, 0.9, 0.95, 0.99 };
|
||||
|
||||
printf("bench_topp: top-p truncation, old (qsort) vs new (heap) V=%d temp=%.2f\n", V, temp);
|
||||
printf("%-12s %6s %14s %14s %9s %9s\n", "shape", "nuc", "old ns/call", "new ns/call", "speedup", "keep");
|
||||
printf("-----------------------------------------------------------------------------\n");
|
||||
|
||||
for (size_t sh = 0; sh < sizeof(shapes)/sizeof(shapes[0]); sh++){
|
||||
shapes[sh].fill(lo, V);
|
||||
for (size_t ni = 0; ni < sizeof(nucs)/sizeof(nucs[0]); ni++){
|
||||
double nuc = nucs[ni];
|
||||
/* warmup both paths so caches/branch predictors are primed */
|
||||
for (int w = 0; w < 50; w++){ dist_build_old(lo, V, temp, nuc); dist_build_new(lo, V, temp, nuc); }
|
||||
double t_old = bench_ns(dist_build_old, lo, V, temp, nuc);
|
||||
double t_new = bench_ns(dist_build_new, lo, V, temp, nuc);
|
||||
/* keep count = non-zero entries the new path leaves (== old's keep) */
|
||||
int keep = 0; for (int i = 0; i < V; i++) if (g_pbuf[i] != 0.f) keep++;
|
||||
printf("%-12s %6.2f %14.0f %14.0f %8.2fx %9d\n",
|
||||
shapes[sh].name, nuc, t_old, t_new, t_old / t_new, keep);
|
||||
}
|
||||
printf("\n");
|
||||
}
|
||||
|
||||
printf("bench_topp: done (lower ns is better; speedup = old/new)\n");
|
||||
free(lo); free(s_ref); free(s_pidx);
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
/* DSA top-keep partial-select: the quickselect rewrite (#356) must produce a
|
||||
* BIT-IDENTICAL kept-position set to the old full-vocab qsort, for every score
|
||||
* shape the attention indexer can see.
|
||||
*
|
||||
* Why this test exists (#356): attention_rows() selects the top-`keep` context
|
||||
* keys (index_topk=2048 on GLM-5.2) to attend to. It previously did this by
|
||||
* full-qsorting all `nk` scores (O(nk log nk)) and reading tmp[keep-1] as the
|
||||
* threshold. It now does a partial_select_desc (quickselect, O(nk) average) and
|
||||
* takes the threshold as the min of the selected top-keep block. The contract is
|
||||
* subtle but STRONGER than test_topp's:
|
||||
*
|
||||
* The two position-order scans that build dst[] --
|
||||
* for t: if isc[t] > thr -> keep (strictly above threshold)
|
||||
* for t: if isc[t] == thr -> keep (ties, in position order)
|
||||
* -- are UNCHANGED by the rewrite. So if the threshold value is identical,
|
||||
* the kept-position set is identical element-by-element (not just as a
|
||||
* multiset, which is all the unstable sampling heap in #335 could promise).
|
||||
*
|
||||
* Strategy: drive the REAL partial_select_desc (via the include-glm.c pattern)
|
||||
* and replicate the production threshold derivation + scans, then compare the
|
||||
* resulting dst[] against an INDEPENDENT reference that re-implements the OLD
|
||||
* algorithm (full qsort + tmp[keep-1] threshold) on a private buffer. The kept
|
||||
* sets must be element-wise equal on every shape, including tie plateaus where
|
||||
* the boundary membership is decided by the position scan.
|
||||
*
|
||||
* We also directly unit-test partial_select_desc's partition invariant: after
|
||||
* the call, max(a[keep..n)) <= min(a[0..keep)) -- i.e. the keep largest really
|
||||
* did land in the prefix. This catches a broken quickselect even before the
|
||||
* end-to-end comparison.
|
||||
*
|
||||
* In-memory only (no scratch files), so it builds clean on the Windows MinGW CI
|
||||
* job without the unmerged compat shim. */
|
||||
#define main coli_glm_main_unused
|
||||
#include "../glm.c"
|
||||
#undef main
|
||||
|
||||
#include <math.h>
|
||||
|
||||
static int g_nfails = 0;
|
||||
|
||||
#define FAIL(fmt, ...) do { \
|
||||
fprintf(stderr, " FAIL [%s nk=%d keep=%d shape=%s]: " fmt "\n", \
|
||||
label, nk, keep, shape_name, ##__VA_ARGS__); \
|
||||
g_nfails++; \
|
||||
return; \
|
||||
} while (0)
|
||||
|
||||
/* ---- independent reference: the OLD algorithm (full qsort + tmp[keep-1]) ---- */
|
||||
/* qsort comparator matching the production cmp_fdesc exactly (desc, unstable). */
|
||||
static int cmp_ref_desc(const void *a, const void *b){
|
||||
float x=*(const float*)a, y=*(const float*)b; return x<y?1:x>y?-1:0; }
|
||||
|
||||
/* Reproduce the OLD glm.c:2589-2596 exactly: copy, qsort desc, threshold =
|
||||
* tmp[keep-1], then the two position-order scans into dst[]. Returns nd. */
|
||||
static int keep_old(const float *isc, int nk, int keep, int *dst){
|
||||
float *tmp=malloc((size_t)nk*sizeof(float));
|
||||
memcpy(tmp,isc,(size_t)nk*sizeof(float));
|
||||
qsort(tmp,(size_t)nk,sizeof(float),cmp_ref_desc);
|
||||
float thr=tmp[keep-1];
|
||||
int nd=0;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
|
||||
free(tmp);
|
||||
return nd;
|
||||
}
|
||||
|
||||
/* Reproduce the NEW glm.c path: partial_select desc, threshold = min of the
|
||||
* selected block, same two scans. Uses the REAL partial_select_desc from glm.c. */
|
||||
static int keep_new(const float *isc, int nk, int keep, int *dst){
|
||||
float *tmp=malloc((size_t)nk*sizeof(float));
|
||||
memcpy(tmp,isc,(size_t)nk*sizeof(float));
|
||||
partial_select_desc(tmp,nk,keep);
|
||||
float thr=tmp[0]; for(int t=1;t<keep;t++) if(tmp[t]<thr) thr=tmp[t];
|
||||
int nd=0;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
|
||||
free(tmp);
|
||||
return nd;
|
||||
}
|
||||
|
||||
/* ---- direct unit test of the partition invariant ---- */
|
||||
static void check_partition(const char *label, const float *isc, int nk, int keep,
|
||||
const char *shape_name){
|
||||
float *tmp=malloc((size_t)nk*sizeof(float));
|
||||
memcpy(tmp,isc,(size_t)nk*sizeof(float));
|
||||
partial_select_desc(tmp,nk,keep);
|
||||
/* invariant: every element of tmp[0..keep) is >= every element of tmp[keep..n).
|
||||
* (>=, not >: equal values may sit on either side of the partition boundary,
|
||||
* which is fine -- the threshold is the MIN of the prefix, and the position
|
||||
* scan handles ties.) */
|
||||
float top_min=INFINITY, tail_max=-INFINITY;
|
||||
for(int i=0;i<keep;i++) if(tmp[i]<top_min) top_min=tmp[i];
|
||||
for(int i=keep;i<nk;i++) if(tmp[i]>tail_max) tail_max=tmp[i];
|
||||
if(!(top_min >= tail_max))
|
||||
FAIL("partition invariant violated: top_min=%.9g < tail_max=%.9g", top_min, tail_max);
|
||||
free(tmp);
|
||||
}
|
||||
|
||||
/* ---- end-to-end: old vs new kept-set must be element-wise identical ---- */
|
||||
static void check_case(const char *label, int nk, int keep, const char *shape_name,
|
||||
const float *isc){
|
||||
int *da=malloc((size_t)nk*sizeof(int));
|
||||
int *db=malloc((size_t)nk*sizeof(int));
|
||||
int na=keep_old(isc,nk,keep,da);
|
||||
int nb=keep_new(isc,nk,keep,db);
|
||||
|
||||
/* 1. both keep exactly `keep` positions (the contract: keep the top-keep by
|
||||
* count). A count mismatch is a real bug, not a tie artifact. */
|
||||
if(na!=keep) FAIL("old kept %d, expected %d (old path is the reference)", na, keep);
|
||||
if(nb!=keep) FAIL("new kept %d, expected %d", nb, keep);
|
||||
if(na!=nb) FAIL("keep-count mismatch: old=%d new=%d", na, nb);
|
||||
|
||||
/* 2. element-wise identical dst[]. This is the strong contract: because the
|
||||
* threshold is derived identically and the position-order scans are byte-
|
||||
* for-byte the same, the kept SET and its ORDER must match exactly. (This
|
||||
* is what makes #356 cleaner than #335, which was multiset-only.) */
|
||||
int first_diff=-1;
|
||||
for(int i=0;i<na;i++){ if(da[i]!=db[i]){ first_diff=i; break; } }
|
||||
if(first_diff>=0)
|
||||
FAIL("kept-set differs at index %d: old dst[%d]=%d new dst[%d]=%d",
|
||||
first_diff, first_diff, da[first_diff], first_diff, db[first_diff]);
|
||||
|
||||
/* 3. also check the partition invariant directly (catches a subtly broken
|
||||
* quickselect even if the threshold happened to come out right). */
|
||||
check_partition(label,isc,nk,keep,shape_name);
|
||||
|
||||
free(da); free(db);
|
||||
printf(" ok [nk=%d keep=%d shape=%s]\n", nk, keep, shape_name);
|
||||
}
|
||||
#undef FAIL
|
||||
|
||||
/* deterministic xorshift32 RNG (matches the test_i4_grouped.c / test_topp.c convention) */
|
||||
static uint32_t rng_state = 0x12345678u;
|
||||
static uint32_t xr(void){ rng_state ^= rng_state << 13; rng_state ^= rng_state >> 17;
|
||||
rng_state ^= rng_state << 5; return rng_state; }
|
||||
static double frand(void){ return (xr() >> 8) * (1.0 / 16777216.0); } /* [0,1) */
|
||||
|
||||
/* fill scores for a given shape. Shapes stress the threshold boundary and the
|
||||
* quickselect's median-of-three pivot differently. */
|
||||
static void fill_shape(float *isc, int nk, int shape){
|
||||
switch(shape){
|
||||
case 0: /* uniform random distinct (no ties): the clean contract case */
|
||||
for(int i=0;i<nk;i++) isc[i]=(float)(frand()*1000.0); break;
|
||||
case 1: /* peaked: a few hot, long distinct tail (realistic attention shape) */
|
||||
for(int i=0;i<nk;i++) isc[i]=(float)(-1.0 - frand()*4.0);
|
||||
isc[0]=3.f; if(nk>3) isc[nk/3]=1.f; if(nk>2) isc[nk/2]=0.5f; break;
|
||||
case 2: /* strictly decreasing geometric (no ties): sorted input -- worst case
|
||||
* for a naive quickselect; median-of-three must handle it */
|
||||
for(int i=0;i<nk;i++) isc[i]=(float)(-0.001*(double)i); break;
|
||||
case 3: /* strictly increasing (reverse-sorted): the other quickselect worst case */
|
||||
for(int i=0;i<nk;i++) isc[i]=(float)(0.001*(double)i); break;
|
||||
case 4: /* plateau ties: blocks of equal value -> boundary membership decided
|
||||
* entirely by the position scan (exercises the ==thr path) */
|
||||
for(int i=0;i<nk;i++) isc[i]=(float)(-(double)(i/7)); break;
|
||||
case 5: /* all-equal: every value identical -> degenerate threshold, all kept
|
||||
* via the ==thr scan; quickselect must not infinite-loop or corrupt */
|
||||
for(int i=0;i<nk;i++) isc[i]=5.f; break;
|
||||
}
|
||||
}
|
||||
|
||||
int main(void){
|
||||
/* Sizes around the real index_topk=2048 boundary, plus small cases that
|
||||
* exercise the k>=n / k==1 / k==n edges. */
|
||||
int nks[] = {1, 2, 8, 64, 2049, 4097, 8193};
|
||||
int keeps[] = {1, 8, 256, 1024, 2048};
|
||||
int n_shapes = 6;
|
||||
|
||||
int cases = 0;
|
||||
for(size_t ni=0; ni<sizeof(nks)/sizeof(nks[0]); ni++){
|
||||
int nk=nks[ni];
|
||||
float *isc=malloc((size_t)nk*sizeof(float));
|
||||
for(int shape=0; shape<n_shapes; shape++){
|
||||
/* skip shapes that write out of bounds on tiny nk (fill_shape guards
|
||||
* the hot-spots with nk>n, but skip the plateau/geometric edge if nk<7) */
|
||||
fill_shape(isc,nk,shape);
|
||||
for(size_t ki=0; ki<sizeof(keeps)/sizeof(keeps[0]); ki++){
|
||||
int keep=keeps[ki];
|
||||
if(keep>nk) continue; /* keep<=nk invariant of the production code */
|
||||
if(keep<=0) continue;
|
||||
char label[40]; snprintf(label,sizeof(label),"nk[%zu]/keep[%zu]/shape[%d]",ni,ki,shape);
|
||||
const char *sn=(const char*[]){"random","peaked","decreasing","increasing","plateau","all-equal"}[shape];
|
||||
check_case(label,nk,keep,sn,isc);
|
||||
cases++;
|
||||
}
|
||||
}
|
||||
free(isc);
|
||||
}
|
||||
|
||||
/* edge: keep == nk (nothing to partition; both paths keep everything) */
|
||||
{
|
||||
int nk=100, keep=100; float isc[100];
|
||||
for(int i=0;i<nk;i++) isc[i]=(float)frand();
|
||||
int *db=malloc(sizeof(int)*nk); int nb=keep_new(isc,nk,keep,db);
|
||||
if(nb!=nk){ fprintf(stderr," FAIL [keep==nk]: kept %d expected %d\n",nb,nk); g_nfails++; }
|
||||
else printf(" ok [keep==nk nk=%d]\n",nk);
|
||||
free(db); cases++;
|
||||
}
|
||||
/* edge: keep == 1 (threshold = the single max; quickselect must find it) */
|
||||
{
|
||||
int nk=500, keep=1; float isc[500];
|
||||
for(int i=0;i<nk;i++) isc[i]=(float)frand();
|
||||
int da[1],db[1]; int na=keep_old(isc,nk,keep,da), nb=keep_new(isc,nk,keep,db);
|
||||
if(na!=1||nb!=1||da[0]!=db[0]){
|
||||
fprintf(stderr," FAIL [keep==1]: old={%d (n=%d)} new={%d (n=%d)}\n",da[0],na,db[0],nb); g_nfails++; }
|
||||
else printf(" ok [keep==1 argmax=%d]\n",db[0]); cases++;
|
||||
}
|
||||
/* edge: all-equal scores, keep in the middle -> every kept slot is a tie;
|
||||
* the position scan must pick positions 0..keep-1 deterministically */
|
||||
{
|
||||
int nk=1000, keep=500; float isc[1000];
|
||||
for(int i=0;i<nk;i++) isc[i]=3.14f;
|
||||
int *db=malloc(sizeof(int)*nk); int nb=keep_new(isc,nk,keep,db);
|
||||
int bad=0; for(int i=0;i<nb;i++) if(db[i]!=i) bad=1;
|
||||
if(nb!=keep||bad){ fprintf(stderr," FAIL [all-equal keep=%d]: nb=%d bad=%d\n",keep,nb,bad); g_nfails++; }
|
||||
else printf(" ok [all-equal keep=%d -> positions 0..%d]\n",keep,keep-1); cases++;
|
||||
free(db);
|
||||
}
|
||||
|
||||
printf("\ntest_dsa_select: %d cases run, %d failure(s)\n", cases, g_nfails);
|
||||
if(g_nfails){ printf("test_dsa_select: FAIL\n"); return 1; }
|
||||
printf("test_dsa_select: ok\n");
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
/* st_pread_full: chunk loop + honest truncation errors.
|
||||
* Built with -DST_PREAD_CHUNK=7 so a ~100-byte tensor takes many pread calls —
|
||||
* exercising the loop that production only needs past 2^31 bytes (one pread
|
||||
* caps there on Linux; big bf16 tensors exceed it). Also forks a child against
|
||||
* a truncated shard and requires exit(1) with a "short read" message instead
|
||||
* of the old perror("... : Success"). */
|
||||
#define _GNU_SOURCE
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#ifndef _WIN32
|
||||
#include <sys/wait.h>
|
||||
#include <unistd.h>
|
||||
#endif
|
||||
|
||||
#include "../st.h"
|
||||
|
||||
#define CHECK(condition) do { \
|
||||
if (!(condition)) { \
|
||||
fprintf(stderr, "%s:%d: check failed: %s\n", __FILE__, __LINE__, #condition); \
|
||||
return 1; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static void write_snap(const char *dir, int truncate_bytes) {
|
||||
char path[512];
|
||||
snprintf(path, sizeof(path), "%s/model.safetensors", dir);
|
||||
unsigned char data[96];
|
||||
for (int i = 0; i < 96; i++) data[i] = (unsigned char)(i * 7 + 3);
|
||||
const char *hdr = "{\"t\":{\"dtype\":\"U8\",\"shape\":[96],\"data_offsets\":[0,96]}}";
|
||||
uint64_t hlen = strlen(hdr);
|
||||
FILE *f = fopen(path, "wb");
|
||||
fwrite(&hlen, 8, 1, f);
|
||||
fwrite(hdr, 1, hlen, f);
|
||||
fwrite(data, 1, (size_t)(96 - truncate_bytes), f);
|
||||
fclose(f);
|
||||
}
|
||||
|
||||
int main(void) {
|
||||
/* relative to the CWD, per test_stops: MinGW .exe files resolve Windows
|
||||
* paths and "/tmp" is not one */
|
||||
char dir[] = "test_st_pread_XXXXXX";
|
||||
if (!mkdtemp(dir)) { perror("mkdtemp"); return 1; }
|
||||
|
||||
/* 1) chunk loop: 96-byte tensor read 7 bytes at a time, content exact */
|
||||
write_snap(dir, 0);
|
||||
shards S; st_init(&S, dir);
|
||||
unsigned char out[96] = {0};
|
||||
st_read_raw(&S, "t", out, 0);
|
||||
for (int i = 0; i < 96; i++) CHECK(out[i] == (unsigned char)(i * 7 + 3));
|
||||
|
||||
#ifndef _WIN32
|
||||
/* 2) shard truncated AFTER st_init (init validates static bounds, so the
|
||||
* pread path only fires when the file shrinks underneath a live handle):
|
||||
* child must exit(1) with an honest message, not perror's "Success" */
|
||||
char shard[512]; snprintf(shard, sizeof(shard), "%s/model.safetensors", dir);
|
||||
struct stat sb; CHECK(stat(shard, &sb) == 0);
|
||||
CHECK(truncate(shard, sb.st_size - 40) == 0);
|
||||
int pipefd[2]; CHECK(pipe(pipefd) == 0);
|
||||
pid_t pid = fork(); CHECK(pid >= 0);
|
||||
if (pid == 0) {
|
||||
dup2(pipefd[1], 2); close(pipefd[0]); close(pipefd[1]);
|
||||
unsigned char buf[96];
|
||||
st_read_raw(&S, "t", buf, 0); /* inherited handles; must exit(1) inside */
|
||||
_exit(42); /* reaching here = bug */
|
||||
}
|
||||
close(pipefd[1]);
|
||||
char err[512] = {0};
|
||||
ssize_t n = read(pipefd[0], err, sizeof(err)-1); (void)n;
|
||||
close(pipefd[0]);
|
||||
int status = 0; waitpid(pid, &status, 0);
|
||||
CHECK(WIFEXITED(status) && WEXITSTATUS(status) == 1);
|
||||
CHECK(strstr(err, "short read") != NULL);
|
||||
CHECK(strstr(err, "Success") == NULL);
|
||||
#else
|
||||
/* fork/pipe/truncate are POSIX; Windows still runs the chunk-loop check */
|
||||
printf("test_st_pread: truncation subtest skipped on Windows\n");
|
||||
#endif
|
||||
|
||||
char cmd[600];
|
||||
#ifdef _WIN32
|
||||
snprintf(cmd, sizeof(cmd), "rmdir /s /q %s", dir);
|
||||
#else
|
||||
snprintf(cmd, sizeof(cmd), "rm -rf %s", dir);
|
||||
#endif
|
||||
if (system(cmd)) {}
|
||||
printf("test_st_pread: chunk loop + honest truncation error: ok\n");
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,280 @@
|
||||
/* Top-p (nucleus) truncation in dist_build: the partial-select rewrite (#335) must be
|
||||
* indistinguishable from the old full-vocab qsort for every shape dist_sample can see.
|
||||
*
|
||||
* Why this test exists (#335): dist_build() previously qsort-ed the entire 151936-entry
|
||||
* vocab on every sampled token to find the few-hundred-token head whose cumulative mass
|
||||
* reaches g_nuc. It now heapifies (O(V)) and pops only the head (k * O(log V)). The win
|
||||
* is structural; the risk is a silent sampling-distribution change, because the contract
|
||||
* is subtle:
|
||||
*
|
||||
* dist_sample() iterates g_pbuf[0..V-1] BY TOKEN ID and sums probabilities directly.
|
||||
* So dist_build MUST leave g_pbuf indexed by id (never reordered) AND must zero every
|
||||
* truncated tail entry -- merely excluding the tail from the head would leave mass on
|
||||
* it and the sampled distribution would drift with no crash and no error.
|
||||
*
|
||||
* Strategy: drive the REAL dist_build (via the test_stops.c include-glm.c pattern) on a
|
||||
* sweep of distributions and g_nuc values, and compare against an INDEPENDENT reference
|
||||
* that re-implements the OLD algorithm (full qsort + zero-tail + renorm) in double on a
|
||||
* private buffer. On shapes with no ties the renormalized head must be BIT-IDENTICAL to
|
||||
* the reference (the issue's stated invariant: s2 accumulates in the same descending
|
||||
* order). On tie shapes, where the unstable qsort already left ordering unspecified, we
|
||||
* check multiset equality instead. Every shape also checks: exact-zero tails, head sums
|
||||
* to 1.0, and a sane keep-count.
|
||||
*
|
||||
* No scratch files: the test runs entirely in memory (no mkdtemp), so it builds clean on
|
||||
* the Windows MinGW CI job without the unmerged compat shim (#352). */
|
||||
#define main coli_glm_main_unused
|
||||
#include "../glm.c"
|
||||
#undef main
|
||||
|
||||
#include <math.h>
|
||||
|
||||
static int g_nfails = 0;
|
||||
|
||||
/* pointer set by ref_build so cmp_ref_desc can read the current reference buffer
|
||||
* (the qsort comparator gets no user-data argument in C). */
|
||||
static const double *g_ref_p = NULL;
|
||||
|
||||
#define FAIL(fmt, ...) do { \
|
||||
fprintf(stderr, " FAIL [%s V=%d nuc=%.3f shape=%s]: " fmt "\n", \
|
||||
label, V, nuc, shape_name, ##__VA_ARGS__); \
|
||||
g_nfails++; \
|
||||
return; \
|
||||
} while (0)
|
||||
|
||||
/* ---- independent reference: the OLD algorithm, in double, on a private buffer ------- */
|
||||
/* Stable qsort by descending probability (ties broken by ascending index, which makes
|
||||
* the reference deterministic regardless of the production comparator). */
|
||||
static int cmp_ref_desc(const void *a, const void *b){
|
||||
double pa = ((const double *)g_ref_p)[*(const int*)a];
|
||||
double pb = ((const double *)g_ref_p)[*(const int*)b];
|
||||
if (pa < pb) return 1;
|
||||
if (pa > pb) return -1;
|
||||
/* tie -> lower index first (stable, unlike the production comparator) */
|
||||
return *(const int*)a - *(const int*)b;
|
||||
}
|
||||
|
||||
/* Build the reference distribution into out[0..V-1] (indexed by token id), mirroring the
|
||||
* old dist_build: softmax(lo/temp) truncated to top-p nuc, tail zeroed, head renormalized.
|
||||
* Returns the keep-count through *keep_out. */
|
||||
static void ref_build(const float *lo, int V, double temp, double nuc,
|
||||
double *out, int *pidx, int *keep_out){
|
||||
double mx = lo[0]; for (int i = 1; i < V; i++) if (lo[i] > mx) mx = lo[i];
|
||||
double s = 0, invt = 1.0 / (temp > 1e-4 ? temp : 1e-4);
|
||||
for (int i = 0; i < V; i++){ out[i] = exp((lo[i]-mx)*invt); s += out[i]; }
|
||||
for (int i = 0; i < V; i++) out[i] /= s;
|
||||
|
||||
if (nuc > 0 && nuc < 1.0){
|
||||
for (int i = 0; i < V; i++) pidx[i] = i;
|
||||
qsort(pidx, V, sizeof(int), cmp_ref_desc);
|
||||
double cum = 0; int keep = V;
|
||||
for (int i = 0; i < V; i++){ cum += out[pidx[i]]; if (cum >= nuc){ keep = i+1; break; } }
|
||||
double s2 = 0;
|
||||
for (int i = keep; i < V; i++) out[pidx[i]] = 0;
|
||||
for (int i = 0; i < keep; i++) s2 += out[pidx[i]];
|
||||
for (int i = 0; i < keep; i++) out[pidx[i]] /= s2;
|
||||
*keep_out = keep;
|
||||
} else {
|
||||
*keep_out = V;
|
||||
}
|
||||
}
|
||||
|
||||
/* count how many production g_pbuf entries are non-zero == the head size */
|
||||
static int head_count(int V){
|
||||
int n = 0; for (int i = 0; i < V; i++) if (g_pbuf[i] != 0.f) n++; return n;
|
||||
}
|
||||
|
||||
/* Run one case: load logits into g_pbuf via the real dist_build, compare to reference.
|
||||
* shape_name is for diagnostics only. */
|
||||
static void check_case(const char *label, int V, double nuc, const char *shape_name,
|
||||
const float *lo){
|
||||
/* reference on a private buffer */
|
||||
double *ref = malloc((size_t)V * sizeof(double));
|
||||
int *ridx = malloc((size_t)V * sizeof(int));
|
||||
int ref_keep = 0;
|
||||
g_ref_p = ref; /* cmp_ref_desc reads this */
|
||||
ref_build(lo, V, g_temp, nuc, ref, ridx, &ref_keep);
|
||||
|
||||
/* production: drive the real dist_build (writes the global g_pbuf) */
|
||||
g_nuc = (float)nuc;
|
||||
dist_build(lo, V);
|
||||
|
||||
int got_keep = head_count(V);
|
||||
|
||||
/* 1. keep-count must match the reference exactly. The partial select and the old
|
||||
* qsort keep the same NUMBER of tokens by construction (same cumulative-mass rule);
|
||||
* a count divergence is a real bug, not a tie artifact. */
|
||||
if (got_keep != ref_keep)
|
||||
FAIL("keep-count mismatch: got %d, ref %d", got_keep, ref_keep);
|
||||
|
||||
/* 2. Detect ties across the WHOLE pre-truncation distribution, not just the kept set.
|
||||
* A tie at the head/tail boundary makes which-side-a-token-lands-on interchangeable:
|
||||
* both algorithms keep the right count but may keep different MEMBERS. So any input
|
||||
* with a duplicated softmax value needs the relaxed multiset comparison below. We
|
||||
* detect this on the reference softmax (pre-truncation) by sorting all V values. */
|
||||
int has_ties = 0;
|
||||
{
|
||||
double *all = malloc((size_t)V * sizeof(double));
|
||||
/* reconstruct the pre-truncation softmax the same way ref_build does */
|
||||
double mx = lo[0]; for (int i = 1; i < V; i++) if (lo[i] > mx) mx = lo[i];
|
||||
double s = 0, invt = 1.0 / (g_temp > 1e-4 ? g_temp : 1e-4);
|
||||
for (int i = 0; i < V; i++){ all[i] = exp((lo[i]-mx)*invt); s += all[i]; }
|
||||
for (int i = 0; i < V; i++) all[i] /= s;
|
||||
for (int a = 1; a < V; a++){ double k = all[a]; int b = a-1;
|
||||
while (b >= 0 && all[b] > k){ all[b+1] = all[b]; b--; } all[b+1] = k; }
|
||||
for (int a = 1; a < V; a++) if (all[a] == all[a-1]){ has_ties = 1; break; }
|
||||
free(all);
|
||||
}
|
||||
|
||||
if (has_ties){
|
||||
/* Multiset equality of the non-zero (head) values. Ties make membership
|
||||
* interchangeable, so we compare sorted value-multisets, not id-aligned values.
|
||||
* Tolerance is 1e-6 relative -- the engine uses float arithmetic, the reference
|
||||
* double, so sub-ULP noise is expected (matches test_i4_grouped.c's convention). */
|
||||
double *got = malloc((size_t)ref_keep * sizeof(double));
|
||||
int gm = 0;
|
||||
for (int i = 0; i < V; i++) if (g_pbuf[i] != 0.f) got[gm++] = (double)g_pbuf[i];
|
||||
if (gm != ref_keep)
|
||||
FAIL("tie-shape head size mismatch: got %d non-zero, ref %d", gm, ref_keep);
|
||||
for (int a = 1; a < gm; a++){ double k = got[a]; int b = a-1;
|
||||
while (b >= 0 && got[b] > k){ got[b+1] = got[b]; b--; } got[b+1] = k; }
|
||||
double *rsort = malloc((size_t)ref_keep * sizeof(double));
|
||||
int rm = 0;
|
||||
for (int i = 0; i < V; i++) if (ref[i] != 0.0) rsort[rm++] = ref[i];
|
||||
for (int a = 1; a < rm; a++){ double k = rsort[a]; int b = a-1;
|
||||
while (b >= 0 && rsort[b] > k){ rsort[b+1] = rsort[b]; b--; } rsort[b+1] = k; }
|
||||
int mm = 0; double worst = 0;
|
||||
for (int i = 0; i < gm; i++){
|
||||
double d = fabs(got[i] - rsort[i]);
|
||||
double rel = rsort[i] > 1e-30 ? d / rsort[i] : d;
|
||||
if (rel > worst) worst = rel;
|
||||
if (rel > 1e-6) mm++;
|
||||
}
|
||||
free(got); free(rsort);
|
||||
if (mm) FAIL("tie-shape multiset mismatch: %d/%d head values differ beyond 1e-6 rel (worst %.3g)",
|
||||
mm, ref_keep, worst);
|
||||
} else {
|
||||
/* No ties anywhere: membership is forced, so compare id-aligned head values. The
|
||||
* engine computes in float (g_pbuf /= (float)s2) while the reference uses double,
|
||||
* so the comparison is relative-tolerance (1e-6), not bit-exact -- the partial
|
||||
* select and qsort accumulate s2 in the same descending order, so any difference
|
||||
* is pure float-rounding noise, not an ordering bug. */
|
||||
int bad = 0; int first_id = -1; float gv = 0, rv = 0; double worst = 0;
|
||||
for (int i = 0; i < V; i++){
|
||||
if (ref[i] == 0.0) continue; /* tail */
|
||||
float want = (float)ref[i];
|
||||
double d = fabs((double)g_pbuf[i] - (double)want);
|
||||
double rel = fabs((double)want) > 1e-30 ? d / fabs((double)want) : d;
|
||||
if (rel > worst) worst = rel;
|
||||
if (rel > 1e-6){
|
||||
bad++; if (first_id < 0){ first_id = i; gv = g_pbuf[i]; rv = want; }
|
||||
if (bad > 3) break;
|
||||
}
|
||||
}
|
||||
if (bad)
|
||||
FAIL("head not within 1e-6 rel of reference: %d entries differ (first id %d: got %.9g want %.9g, worst %.3g)",
|
||||
bad, first_id, (double)gv, (double)rv, worst);
|
||||
}
|
||||
|
||||
/* 3. head must renormalize to 1.0 (within float epsilon) */
|
||||
double sum = 0; for (int i = 0; i < V; i++) sum += g_pbuf[i];
|
||||
if (fabs(sum - 1.0) > 1e-5)
|
||||
FAIL("head does not sum to 1.0: sum=%.12g (keep=%d)", sum, got_keep);
|
||||
|
||||
free(ref); free(ridx);
|
||||
printf(" ok [V=%d nuc=%.3f shape=%s keep=%d%s sum=%.10f]\n",
|
||||
V, nuc, shape_name, got_keep, has_ties ? " (ties)" : "", sum);
|
||||
}
|
||||
#undef FAIL
|
||||
|
||||
/* deterministic xorshift32 RNG (matches the test_i4_grouped.c convention) */
|
||||
static uint32_t rng_state = 0x12345678u;
|
||||
static uint32_t xr(void){ rng_state ^= rng_state << 13; rng_state ^= rng_state >> 17;
|
||||
rng_state ^= rng_state << 5; return rng_state; }
|
||||
static double frand(void){ return (xr() >> 8) * (1.0 / 16777216.0); } /* [0,1) */
|
||||
|
||||
/* fill logits for a given shape. Shapes chosen to stress the comparator and the head/tail
|
||||
* boundary differently. */
|
||||
static void fill_shape(float *lo, int V, int shape){
|
||||
switch (shape){
|
||||
case 0: /* uniform -> every token equal probability -> massive tie plateau */
|
||||
for (int i = 0; i < V; i++) lo[i] = 0.f; break;
|
||||
case 1: /* peaked: one dominant token, rest small and distinct (no ties).
|
||||
* The fixed hot-spots are clamped to V-1 so small V (incl. V=1) doesn't
|
||||
* write out of bounds and corrupt heap metadata on the later free(lo). */
|
||||
for (int i = 0; i < V; i++) lo[i] = (float)(-1.0 - frand()*4.0);
|
||||
lo[0] = 3.f; lo[V/3<V?V/3:V-1] = 1.f; lo[V/2<V?V/2:V-1] = 0.5f; break;
|
||||
case 2: /* all-equal distinct decay (no ties): geometric, strictly decreasing */
|
||||
for (int i = 0; i < V; i++) lo[i] = (float)(-0.001 * (double)i); break;
|
||||
case 3: /* plateau ties: blocks of equal value -> comparator tie handling */
|
||||
for (int i = 0; i < V; i++) lo[i] = (float)(-(double)(i / 7)); /* 7-wide plateaus */
|
||||
break;
|
||||
case 4: /* sharp-tail: a few hot, then a long flat floor (small tie at the floor).
|
||||
* Hot count is min(12,V) so V<12 (incl. V=1) stays in bounds. */
|
||||
for (int i = 0; i < V; i++) lo[i] = -8.f;
|
||||
{ int hot = V<12 ? V : 12; for (int i = 0; i < hot; i++) lo[i] = (float)(2.0 - frand()); } break;
|
||||
}
|
||||
}
|
||||
|
||||
int main(void){
|
||||
/* sizes: small for exhaustive tie detection up to near-production scale */
|
||||
int sizes[] = {1, 2, 8, 64, 257, 1519}; /* 1519 ~= V/100 of GLM-5.2 */
|
||||
double nucs[] = {0.001, 0.5, 0.9, 0.999}; /* tight -> almost-everything */
|
||||
int n_shapes = 5;
|
||||
|
||||
/* temperature used by dist_build: pick a normal serving value */
|
||||
g_temp = 0.7f;
|
||||
|
||||
int cases = 0;
|
||||
for (size_t si = 0; si < sizeof(sizes)/sizeof(sizes[0]); si++){
|
||||
int V = sizes[si];
|
||||
/* dist_build allocates g_pbuf/g_pidx ONCE and reuses them (single-V invariant in
|
||||
* real serving, where V is the constant model vocab). This sweep varies V, so free
|
||||
* and force a reallocation per size -- otherwise a later, larger V would overflow
|
||||
* the buffer sized for the first (smallest) V. */
|
||||
free(g_pbuf); g_pbuf = NULL; free(g_pidx); g_pidx = NULL;
|
||||
float *lo = malloc((size_t)V * sizeof(float));
|
||||
for (int shape = 0; shape < n_shapes; shape++){
|
||||
fill_shape(lo, V, shape);
|
||||
for (size_t ni = 0; ni < sizeof(nucs)/sizeof(nucs[0]); ni++){
|
||||
char label[32]; snprintf(label, sizeof(label), "size[%zu]/shape[%d]", si, shape);
|
||||
const char *sn = (const char*[]){"uniform","peaked","geometric","plateau","sharptail"}[shape];
|
||||
check_case(label, V, nucs[ni], sn, lo);
|
||||
cases++;
|
||||
}
|
||||
}
|
||||
free(lo);
|
||||
}
|
||||
|
||||
/* guard-off path: g_nuc >= 1 must skip truncation entirely (full softmax kept) */
|
||||
{
|
||||
int V = 256; float lo[256];
|
||||
for (int i = 0; i < V; i++) lo[i] = (float)(frand()*4 - 2);
|
||||
g_nuc = 1.0f; dist_build(lo, V);
|
||||
int nz = 0; for (int i = 0; i < V; i++) if (g_pbuf[i] != 0.f) nz++;
|
||||
if (nz != V){ fprintf(stderr, " FAIL [guard-off nuc=1.0]: %d/%d entries kept, expected all\n", nz, V); g_nfails++; }
|
||||
else printf(" ok [guard-off nuc=1.0 keep=%d]\n", nz);
|
||||
cases++;
|
||||
|
||||
g_nuc = 0.0f; dist_build(lo, V);
|
||||
nz = 0; for (int i = 0; i < V; i++) if (g_pbuf[i] != 0.f) nz++;
|
||||
if (nz != V){ fprintf(stderr, " FAIL [guard-off nuc=0.0]: %d/%d entries kept, expected all\n", nz, V); g_nfails++; }
|
||||
else printf(" ok [guard-off nuc=0.0 keep=%d]\n", nz);
|
||||
cases++;
|
||||
}
|
||||
|
||||
/* extreme tie edge case: V=1, single token -> keep=1 regardless of nuc */
|
||||
{
|
||||
float lo[1] = {5.f};
|
||||
g_nuc = 0.5f; dist_build(lo, 1);
|
||||
if (g_pbuf[0] == 0.f || !(fabs((double)g_pbuf[0] - 1.0) < 1e-6)){
|
||||
fprintf(stderr, " FAIL [V=1]: g_pbuf[0]=%.9g, expected 1.0\n", (double)g_pbuf[0]); g_nfails++;
|
||||
} else printf(" ok [V=1 keep=1]\n");
|
||||
cases++;
|
||||
}
|
||||
|
||||
printf("\ntest_topp: %d cases run, %d failure(s)\n", cases, g_nfails);
|
||||
if (g_nfails){ printf("test_topp: FAIL\n"); return 1; }
|
||||
printf("test_topp: ok\n");
|
||||
return 0;
|
||||
}
|
||||
@@ -3,7 +3,10 @@
|
||||
|
||||
Downloads or converts a local OLMoE checkpoint (e.g., allenai/OLMoE-1B-7B-0125-Instruct).
|
||||
Dense weights stay as-is (engine reads BF16/F16 → F32 on load).
|
||||
Expert weights get row-wise int8 quantization with float32 scales.
|
||||
Expert weights get row-wise symmetric quantization to --ebits bits (default 4)
|
||||
with float32 scales. Storage stays one value per int8 byte regardless of bits,
|
||||
matching the engine's expert layout (olmoe.c quantize_rows) — for 4 bits the
|
||||
values are simply confined to [-8, 7] with scales computed against qmax=7.
|
||||
|
||||
Usage:
|
||||
python tools/convert_olmoe.py --repo allenai/OLMoE-1B-7B-0125-Instruct --out ./olmoe_i4
|
||||
@@ -29,12 +32,21 @@ except ImportError as exc:
|
||||
EXPERT_KEY_RE = r"model\.layers\.\d+\.mlp\.experts\.\d+\.(gate_proj|up_proj|down_proj)\.weight"
|
||||
|
||||
|
||||
def quantize_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Row-wise int8 quantization. Returns (int8_weights, float32_scales)."""
|
||||
def quantize_row(w: torch.Tensor, bits: int = 8) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Row-wise symmetric quantization to `bits` (2..8).
|
||||
|
||||
Returns (int8_weights, float32_scales). Storage is one value per int8 byte
|
||||
for every bit width — the engine dequantizes as q*scale and never assumes
|
||||
the full int8 range — mirroring olmoe.c quantize_rows():
|
||||
qmax = 2**(bits-1) - 1 (8 -> 127, 4 -> 7, 2 -> 1)
|
||||
scale = amax(|w|, row) / qmax
|
||||
q = clamp(round(w / scale), -qmax-1, qmax)
|
||||
"""
|
||||
qmax = (1 << (bits - 1)) - 1
|
||||
w_f32 = w.float()
|
||||
row_max = w_f32.abs().amax(dim=1, keepdim=True).clamp(min=1e-12)
|
||||
scales = row_max / 127.0
|
||||
q = (w_f32 / scales).round().clamp(-128, 127).to(torch.int8)
|
||||
scales = row_max / qmax
|
||||
q = (w_f32 / scales).round().clamp(-qmax - 1, qmax).to(torch.int8)
|
||||
return q, scales.squeeze(1)
|
||||
|
||||
|
||||
@@ -50,9 +62,12 @@ def main():
|
||||
src.add_argument("--model", help="Local HF checkpoint directory")
|
||||
ap.add_argument("--out", required=True, help="Output directory for int4 model")
|
||||
ap.add_argument("--ebits", type=int, default=4,
|
||||
help="Expert quant bits (4 or 8, default 4)")
|
||||
help="Expert quant bits (2..8, default 4)")
|
||||
args = ap.parse_args()
|
||||
|
||||
if not 2 <= args.ebits <= 8: # storage is int8_t; engine rejects the same range (olmoe.c)
|
||||
sys.exit(f"--ebits must be 2..8 (got {args.ebits})")
|
||||
|
||||
if args.repo:
|
||||
from huggingface_hub import snapshot_download
|
||||
from huggingface_hub.errors import LocalEntryNotFoundError
|
||||
@@ -96,7 +111,7 @@ def main():
|
||||
for name, tensor in tensors.items():
|
||||
if is_expert_weight(name):
|
||||
expert_count += 1
|
||||
q, scales = quantize_row(tensor)
|
||||
q, scales = quantize_row(tensor, args.ebits)
|
||||
total_expert_f32 += tensor.numel() * tensor.element_size()
|
||||
total_expert_q += q.numel() * 1 + scales.numel() * 4
|
||||
out_tensors[name] = q
|
||||
@@ -109,7 +124,7 @@ def main():
|
||||
ratio = total_expert_q / max(total_expert_f32, 1) * 100
|
||||
print(f"ok")
|
||||
|
||||
print(f"\nDone. {expert_count} expert tensors quantized.")
|
||||
print(f"\nDone. {expert_count} expert tensors quantized to int{args.ebits}.")
|
||||
print(f"Expert storage: {total_expert_f32/1e9:.1f} GB -> {total_expert_q/1e9:.1f} GB ({ratio:.0f}%)")
|
||||
print(f"Model ready at: {out}")
|
||||
print(f"\nRun: SNAP={out} ./olmoe.exe 32 4 16")
|
||||
|
||||
@@ -118,7 +118,7 @@ def quantize_param(w, bits, group, rot=False, e8=""):
|
||||
|
||||
def _grid_or_e8(x, bits, group, e8):
|
||||
if e8:
|
||||
return _quant_e8(x.float(), group, ball=(e8 == "-e8"))
|
||||
return _quant_e8(x.float(), group, bits, ball=(e8 == "-e8"))
|
||||
return _quant_last_dim(x, bits, group)
|
||||
|
||||
|
||||
@@ -169,8 +169,15 @@ def _e8_ball(y, r2=10.0):
|
||||
return p
|
||||
|
||||
|
||||
def _quant_e8(x, group, ball):
|
||||
"""Blocks of 8 along the input dim; per-group scale by MSE search over RMS multiples."""
|
||||
_E8_R2_REPORTED = set()
|
||||
def _e8_radius(bits):
|
||||
# E8 lattice: points within |p|^2<=r2 grow ~r2^4, so +1 bit (x256 codebook) needs r2 x4.
|
||||
# Anchor: r2=10 is the ~2^16 E8P ball (2 bits over 8 dims). Scale from there.
|
||||
return 10.0 * (4.0 ** (bits - 2))
|
||||
|
||||
def _quant_e8(x, group, bits, ball):
|
||||
"""Blocks of 8 along the input dim; per-group scale by MSE search over RMS multiples.
|
||||
ball=True clamps to the rate-scaled E8 ball for `bits`; ball=False is the unbounded ideal."""
|
||||
if x.shape[-1] % 8:
|
||||
raise SystemExit(f"-e8 needs input dim divisible by 8 (got {x.shape[-1]})")
|
||||
g = group or x.shape[-1]
|
||||
@@ -183,7 +190,11 @@ def _quant_e8(x, group, ball):
|
||||
for k in (0.5, 0.7, 0.9, 1.1, 1.4, 1.8, 2.4):
|
||||
s = rms * k
|
||||
yb = (xg / s).reshape(-1, g // 8, 8)
|
||||
p = _e8_ball(yb) if ball else _e8_nearest(yb)
|
||||
p = _e8_ball(yb, _e8_radius(bits)) if ball else _e8_nearest(yb)
|
||||
if ball and bits not in _E8_R2_REPORTED:
|
||||
_E8_R2_REPORTED.add(bits)
|
||||
import sys as _sys
|
||||
_sys.stderr.write(f"[e8] bits={bits}: ball r2={_e8_radius(bits):.1f}\n")
|
||||
out = (p.reshape(-1, g) * s)
|
||||
err = (out - xg).pow(2).sum(-1, keepdim=True)
|
||||
if best_err is None:
|
||||
|
||||
Reference in New Issue
Block a user