Merge branch 'main' into fix/windows-cuda-dll-build

This commit is contained in:
mohamedmastouri2000-boop
2026-07-16 19:05:43 +03:00
committed by GitHub
54 changed files with 5608 additions and 305 deletions
+8
View File
@@ -31,6 +31,10 @@ c/tests/test_tier
c/tests/test_tier.exe
c/tests/test_grammar
c/tests/test_grammar.exe
c/tests/test_schema_gbnf
c/tests/test_schema_gbnf.exe
c/tests/test_compat_direct
c/tests/test_compat_direct.exe
# oracoli tiny generati (make_glm_oracle.py) e dati benchmark scaricati
c/glm_tiny/
@@ -60,3 +64,7 @@ c/backend_metal_test
c/tests/test_tier
c/mio_env/
c/bench/
c/tests/test_decode_batch
c/tests/test_i4_acc512
c/tests/test_idot
c/tests/test_uring
+75 -8
View File
@@ -17,14 +17,40 @@ $ ./coli chat
◆ Ciao! 😊 Come posso aiutarti oggi?
```
## See it running
<p align="center">
<img src="docs/media/colibri-dashboard.png" width="900" alt="colibrì web dashboard — live metrics, hardware panel, expert tiers">
</p>
<p align="center"><em>The web dashboard (<code>./coli web</code>): a 744B model answering at 4+ tok/s end-to-end on 6× RTX 5090 —
with live token metrics, the hardware panel, and the VRAM/RAM/disk expert tiers.</em></p>
<p align="center">
<img src="docs/media/colibri-brain.png" width="900" alt="the Brain page — 19,456 experts as a live cortex">
</p>
<p align="center"><em>The <strong>Brain</strong> page: all 19,456 experts as a living cortex — colour is the storage tier,
brightness is routing heat, and every expert routed in a turn flashes white. Hovering shows the expert's
<a href="https://github.com/JustVugg/colibri/issues/175">measured topic affinity</a>.</em></p>
## Contents
- [The idea](#the-idea)
- [See it running](#see-it-running)
- [What's implemented](#whats-implemented)
- [Honest numbers](#honest-numbers-wsl2-12-cores-25-gb-ram-nvme-via-vhdx)
- [Download the model](#download-the-model)
- [Web dashboard](#web-dashboard)
- [Got a better machine?](#got-a-better-machine-try-it--heres-what-to-expect)
## The idea
A 744B Mixture-of-Experts model activates only ~40B parameters per token — and only ~11 GB of those change from token to token (the routed experts). So:
- the **dense part** (attention, shared experts, embeddings — ~17B params) stays **resident in RAM at int4** (~9.9 GB);
- the **21,504 routed experts** (75 MoE layers × 256 experts + the MTP head, ~19 MB each at int4) live **on disk** (~370 GB) and are **streamed on demand**, with a per-layer LRU cache, an optional pinned hot-store, and the OS page cache as a free L2.
- the **19,456 routed experts** (75 MoE layers × 256 experts + the MTP head, ~19 MB each at int4) live **on disk** (~370 GB) and are **streamed on demand**, with a per-layer LRU cache, an optional pinned hot-store, and the OS page cache as a free L2.
The engine is a single C file (`c/glm.c`, ~2,400 lines) plus small headers. No BLAS, no Python at runtime, no GPU required (an opt-in CUDA tier for pinned experts exists — see below).
The engine is a single C file (`c/glm.c`) plus small headers. No BLAS, no Python at runtime, no GPU required (an opt-in CUDA tier for pinned experts exists — see below).
## What's implemented
@@ -354,6 +380,10 @@ PIN=stats.txt PIN_GB=160 SNAP=/nvme/glm52_i4 ./glm 64 4 4
COLI_CUDA=1 COLI_GPUS=0,1,2,3,4,5 CUDA_EXPERT_GB=150 \
CUDA_DENSE=1 PIN=stats.txt PIN_GB=300 RAM_GB=226 \
SNAP=/nvme/glm52_i4 ./glm 64 4 4
# large-RAM host: fill safe VRAM, then keep every remaining expert in RAM
COLI_CUDA=1 COLI_GPUS=0,1,2,3,4,5 CUDA_EXPERT_GB=auto \
CUDA_DENSE=1 COLI_CUDA_ATTN=1 PIN=stats.txt PIN_GB=all RAM_GB=auto \
SNAP=/nvme/glm52_i4 ./glm 64 4 4
```
Selected experts are uploaded during startup, so capacity failures occur before
@@ -375,6 +405,25 @@ replay from 1.87 to 2.16 tok/s (+15.7%), reduced expert disk wait from 5.144s to
3.948s, and kept the projected RAM peak below `RAM_GB=226`. The cache cap adjusts
down automatically (54 to 40 in that run) so the larger pinned tier does not exceed
the process budget. Start lower on hosts with less available RAM.
`CUDA_EXPERT_GB=auto` fills each selected device only up to its measured free
memory minus projected dense tensors and 2 GB of runtime headroom. `PIN_GB=all`
then loads every remaining routed expert into RAM, eliminating decode-time disk
misses when the host budget permits it. The regular `RAM_GB` guard still clamps
the per-layer working cache and rejects unsafe projections; this mode is intended
for dedicated high-memory inference hosts, not desktops running other workloads.
On a dedicated 251 GiB host with six RTX 5090s, this mode selected a 176.7 GB
VRAM expert tier and a 191.3 GB RAM tier (all 19,456 experts resident). The
mode also adapts the VRAM tier every 16 emitted tokens by swapping hot RAM
experts into existing GPU slots. A real 64-token greedy GLM-5.2 generation
measured **6.00 tok/s decode**, up from
2.20 tok/s end-to-end with the earlier 150 GB tier; expert hit rate was 100%
and disk wait was zero. Prompt prefill is reported separately. This is a
host-specific capacity result, not a portable default.
Text-mode timing reports prefill separately from decode. The decode rate starts
after the prompt KV is built, so it is comparable to `REPLAY` throughput without
hiding time-to-first-token.
MTP speculation defaults off on CUDA because cold draft routes increase expert
traffic; an explicit `DRAFT=n` still overrides the default.
@@ -406,8 +455,8 @@ CUDA, CPU hot-store, and CUDA hot-expert execution with identical replay tokens.
### Web interface
`web/` contains a community-contributed browser UI (React + TypeScript, ~390
lines of source, a pure API client — it never touches the engine directly):
`web/` contains a community-contributed browser UI (React + TypeScript, a pure
API client — it never touches the engine directly):
```bash
cd web
@@ -419,7 +468,7 @@ works against the colibrì OpenAI-compatible server (in review, #21) or any othe
compatible endpoint. Nothing leaves the endpoint you configure. The terminal
`coli chat` remains the first-class interface.
Useful knobs (env or flags): `--temp T` token sampling temperature (default 0.7 + nucleus 0.90 — tuned for int4; 0 = greedy), `--topp 0.7` adaptive expert top-p (3040% less disk), `--ngen N` max tokens per answer (`:more` in chat continues a truncated one), `--repin N` adapt RAM/VRAM hot experts every N emitted tokens, `AUTOPIN=0` disable the learning cache's auto-pin, `THINK=1` enable GLM-5.2's reasoning block, `DRAFT=n` MTP draft depth, `GRAMMAR=g.gbnf` grammar-forced drafts for constrained JSON/NDJSON output (`GRAMMAR_DRAFT=n` caps the forced span), `TF=1` teacher-forcing validation, `PILOT=1` router-lookahead disk prefetch (experimental — see below), `CAP_RAISE=0` don't auto-grow the expert cache.
Useful knobs (env or flags): `--temp T` token sampling temperature (default 0.7 + nucleus 0.90 — tuned for int4; 0 = greedy), `--topp 0.7` adaptive expert top-p (3040% less disk), `--ngen N` max tokens per answer (`:more` in chat continues a truncated one), `--repin N` adapt RAM/VRAM hot experts every N emitted tokens, `AUTOPIN=0` disable the learning cache's auto-pin, `THINK=1` enable GLM-5.2's reasoning block, `DRAFT=n` MTP draft depth, `GRAMMAR=g.gbnf` grammar-forced drafts for constrained JSON/NDJSON output (`GRAMMAR_DRAFT=n` caps the forced span), `TF=1` teacher-forcing validation, `PILOT=1` router-lookahead disk prefetch (experimental — see below), `URING=1` Linux-only batched expert I/O (implies `PIPE=1`; also batches `PILOT_REAL`), `PIPE=0` disable the async expert-load pool (**default ON on Windows** — overlaps expert `pread` with the matmul so the CPU isn't idle waiting on the SSD; measured 18% disk service time), `RAM_GB=<n>` claim more RAM for the expert cache than the conservative auto-detect (e.g. `RAM_GB=31` on a 32 GB host raises the cache cap and hit rate measurably), `CAP_RAISE=0` don't auto-grow the expert cache.
### Resource policy
@@ -444,10 +493,10 @@ Disk is an immutable recovery source, not a normal decode target. If the plan
leaves cold expert bytes on disk, speed depends on cache hit rate; output
quality does not.
Cold expert reads use a deferred pipeline: resident RAM/VRAM experts execute
Cold expert reads can use a deferred pipeline: resident RAM/VRAM experts execute
while missing experts are loaded in a bounded background I/O pool, then the
cold results join before the layer completes. `IO_THREADS=n` overrides the
default eight loader threads when foreground work exists. Profiling reports
cold results join before the layer completes. The pool engages only under
`PIPE=1`; `PIPE_WORKERS=n` sets its worker count (default 8). Profiling reports
both disk service time and the smaller foreground-visible wait time so overlap
is explicit rather than credited as unexplained speedup.
@@ -478,6 +527,23 @@ thrashing. Persistent `.coli_usage` remains the long-term signal and is not deca
**Conversations reopen warm** (`.coli_kv`, since 2026-07-10): `coli chat` persists the compressed MLA KV-cache to disk after every turn (~182 KB/token, appended incrementally, crash-safe). Close the chat, reopen it tomorrow — the model still remembers the whole conversation and **zero re-prefill happens**: validated byte-identical to an uninterrupted session. `:reset` clears it, `KVSAVE=0` disables it.
## Web dashboard
One command serves the OpenAI-compatible API **and** the web console on the same port, then opens your browser when the engine is ready:
```bash
cd web && npm install && npm run build # once
./coli web --model <model-dir>
```
What you get:
- **Chat** with live metrics: a flashing token counter while generating, then tok/s, time-to-first-token, prompt→completion counts and queue wait;
- **Runtime panel**: your hardware (CPU, GPUs + VRAM, RAM, cores), the scheduler, and the live expert-tier bar — how many of the 19,456 experts sit in VRAM / RAM / disk right now;
- **Brain**: the whole model as a 76×256 cortex, one cell per expert. Colour = tier, brightness = routing heat, and the experts routed in each turn flash white and decay — you watch the model think. Hover any cell for its tier, heat and [measured topic affinity](https://github.com/JustVugg/colibri/issues/175) (specialists for code, Chinese, math, law… live in layers 1122).
The dashboard talks to the engine over two tiny protocol lines (`TIERS`, `EMAP`/`HITS`) and plain JSON endpoints — nothing heavier than the engine itself.
## Got a better machine? Try it — here's what to expect
colibrì was built on deliberately humble hardware (12 cores, 25 GB RAM, an older DRAM-less NVMe behind a WSL2 VHDX that measured ~1 GB/s random on *this* drive — note WSL2 VHDX is not inherently slow: a community 5090 box measured 10.5 GB/s O_DIRECT through one, [#101](https://github.com/JustVugg/colibri/issues/101)). **Every one of those constraints is a knob your machine can turn up.** The engine needs: Linux (or WSL2), macOS, or **Windows 11 natively (MinGW-w64)**; gcc with OpenMP, AVX2, ≥16 GB RAM, and the ~370 GB int4 model on a local NVMe (ext4/NTFS — never a network/9p mount).
@@ -591,6 +657,7 @@ c/
├── scripts/ long-running conversion helpers
└── tests/ dependency-free C and Python tests
web/ browser UI (pure OpenAI-API client, community-maintained)
desktop/ Tauri v2 desktop shell wrapping the web UI
```
The runtime path intentionally stays flat and readable: `glm.c` plus its small
+97 -25
View File
@@ -14,7 +14,7 @@ TRIPLET := $(shell $(DETECT_CC) -dumpmachine 2>/dev/null)
MINGW := $(findstring mingw,$(TRIPLET))
CYGWIN := $(findstring cygwin,$(TRIPLET))
DARWIN := $(findstring darwin,$(TRIPLET))
PPC64 := $(findstring powerpc64,$(TRIPLET))
LINUX := $(findstring linux,$(TRIPLET))
IS_WIN := $(MINGW)$(CYGWIN)
# Fallbacks for the rare toolchain that does not answer -dumpmachine: keep the
@@ -26,10 +26,17 @@ ifeq ($(IS_WIN),)
UNAME_S := $(shell uname -s)
UNAME_M := $(shell uname -m)
DARWIN := $(findstring Darwin,$(UNAME_S))
PPC64 := $(findstring ppc64,$(UNAME_M))
endif
endif
TARGET_CPU := $(firstword $(subst -, ,$(TRIPLET)))
ifeq ($(TARGET_CPU),)
TARGET_CPU := $(UNAME_M)
endif
X86_64 := $(filter x86_64 amd64,$(TARGET_CPU))
AARCH64 := $(filter aarch64 arm64,$(TARGET_CPU))
PPC64 := $(filter powerpc64% ppc64%,$(TARGET_CPU))
ifneq (,$(DARWIN))
# --- macOS / Apple Silicon ---
# Apple clang non include il runtime OpenMP: se c'e' libomp di Homebrew lo usa
@@ -37,7 +44,9 @@ ifneq (,$(DARWIN))
# Niente -march: su arm64 NEON e' baseline (i kernel __ARM_NEON si attivano da soli).
CC = clang
OMPDIR := $(shell brew --prefix libomp 2>/dev/null)
ifneq ($(OMPDIR),)
# `brew --prefix libomp` can print the formula's prospective path even when it
# is not installed, so verify both artifacts before adding unusable flags.
ifneq ($(and $(OMPDIR),$(wildcard $(OMPDIR)/include/omp.h),$(wildcard $(OMPDIR)/lib/libomp.*)),)
OMPC = -Xclang -fopenmp -I$(OMPDIR)/include
OMPL = -L$(OMPDIR)/lib -lomp
else
@@ -46,6 +55,12 @@ OMPC =
OMPL =
endif
CFLAGS = -O3 $(OMPC) -Wall -Wextra -Wno-unused-parameter -Wno-misleading-indentation -Wno-unused-function
# Opt-in: ARCH=native appends -mcpu=native (arm64 clang uses -mcpu, not -march),
# which unlocks the i8mm SMMLA int8/int4 dot kernels in glm.c. ARCH unset ->
# no -mcpu, default build byte-identical. Apple clang knows apple-m4 / native.
ifneq ($(ARCH),)
CFLAGS += -mcpu=$(ARCH)
endif
LDFLAGS = -lm $(OMPL)
EXE =
else ifneq ($(IS_WIN),)
@@ -60,7 +75,11 @@ else ifneq ($(IS_WIN),)
CC = gcc
ARCH ?= x86-64-v3
CFLAGS = -D_FILE_OFFSET_BITS=64 -O3 -march=$(ARCH) -fopenmp -Wall -Wextra -Wno-unused-parameter -Wno-misleading-indentation -Wno-unused-function
LDFLAGS = -lm -fopenmp -static
# -lpsapi: compat.h calls GetProcessMemoryInfo (rss_gb). It's linked via
# #pragma comment(lib,"psapi.lib") for MSVC, but MinGW gcc ignores that pragma
# (warns), so link psapi explicitly or the build fails undefined-reference on
# modern GCC (e.g. 16.x under UCRT). Harmless where the pragma also resolves it.
LDFLAGS = -lm -fopenmp -static -lpsapi
EXE = .exe
else
ifneq (,$(PPC64))
@@ -70,18 +89,19 @@ ifneq (,$(PPC64))
# (validated token-exact vs the transformers oracle on a POWER8 S824).
CC = gcc
ARCH ?= native
CFLAGS = -O3 -mcpu=$(ARCH) -fopenmp -Wall -Wextra -Wno-unused-parameter -Wno-misleading-indentation -Wno-unused-function
LDFLAGS = -lm -fopenmp
CFLAGS = -O3 -mcpu=$(ARCH) -fopenmp -pthread -Wall -Wextra -Wno-unused-parameter -Wno-misleading-indentation -Wno-unused-function
LDFLAGS = -lm -fopenmp -pthread
EXE =
else
# --- Linux x86-64 (percorso originale, invariato) ---
# --- Linux / *BSD x86-64 (percorso originale, invariato) ---
CC = gcc
# ARCH=native -> ottimizzato per QUESTA macchina (default, piu' veloce).
# ARCH=x86-64-v3 -> binario PORTABILE su qualsiasi x86-64 moderno con AVX2 (per distribuire).
# ARCH=x86-64 -> massima compatibilita' (niente AVX2: usa il path scalare di fallback).
# -pthread: Linux lo tira dentro via -fopenmp, i *BSD no e le pthread_* non risolvono (#219).
ARCH ?= native
CFLAGS = -O3 -march=$(ARCH) -fopenmp -Wall -Wextra -Wno-unused-parameter -Wno-misleading-indentation -Wno-unused-function
LDFLAGS = -lm -fopenmp
CFLAGS = -O3 -march=$(ARCH) -fopenmp -pthread -Wall -Wextra -Wno-unused-parameter -Wno-misleading-indentation -Wno-unused-function
LDFLAGS = -lm -fopenmp -pthread
EXE =
endif
endif
@@ -105,22 +125,51 @@ INSTALL ?= install
# clean. See backend_loader.c and README "cuda-dll" below.
CUDA ?= 0
CUDA_DLL ?= 0
ifneq ($(IS_WIN),)
# the CUDA installer sets CUDA_PATH system-wide (e.g. C:\Program Files\NVIDIA
# GPU Computing Toolkit\CUDA\v13.2); fall back to it before the POSIX default.
# NVCC defaults to plain `nvcc` from PATH: CUDA_PATH contains spaces, which the
# unquoted recipe checks cannot survive — and the cuda-dll recipe already
# requires an MSVC environment (vcvars64) on PATH, so requiring the CUDA bin
# directory too is symmetric. The installer adds it by default.
CUDA_HOME ?= $(subst \,/,$(CUDA_PATH))
NVCC ?= nvcc
# Host compiler override for nvcc (e.g. Fedora ships g++16 but CUDA needs 15):
# make glm CUDA=1 NVCC_CCBIN=g++-15 (opt-in: unset = nvcc's own default)
ifneq ($(NVCC_CCBIN),)
NVCCFLAGS += -ccbin $(NVCC_CCBIN)
endif
else
CUDA_HOME ?= /usr/local/cuda
NVCC ?= $(CUDA_HOME)/bin/nvcc
endif
CUDA_ARCH ?= native
# -Xcompiler passes flags to nvcc's HOST compiler. On Windows that host is MSVC
# cl.exe, which rejects the gcc/clang-style -Wall,-Wextra (D8021). Gate them
# behind a non-MSVC host so the Linux/macOS (gcc/clang) build keeps the warnings
# and the Windows cuda-dll build doesn't fail. (#306)
NVCC_HOST_WARN ?= $(if $(IS_WIN),,-Xcompiler=-Wall,-Wextra)
NVCCFLAGS ?= -O3 -std=c++17 -arch=$(CUDA_ARCH) $(NVCC_HOST_WARN)
ifeq ($(IS_WIN),)
PYTHON ?= python3
ifneq ($(IS_WIN),)
# nvcc's host compiler on Windows is MSVC cl.exe: the GCC-style
# -Xcompiler=-Wall,-Wextra is rejected ("D8021 invalid numeric argument
# '/Wextra'"), which made `make cuda-dll` unbuildable as shipped. -W3 is
# the MSVC warning level (dash form, NOT /W3: MSYS make would mangle the
# slash-form into a filesystem path).
NVCCFLAGS ?= -O3 -std=c++17 -arch=$(CUDA_ARCH) -Xcompiler=-W3
else
NVCCFLAGS ?= -O3 -std=c++17 -arch=$(CUDA_ARCH) -Xcompiler=-Wall,-Wextra
endif
# PYTHON is a HOST tool (clean/test-c/test-python run it on the build machine),
# so key it off the host, NOT $(IS_WIN) — which is derived from the *target*
# triple ($(CC) -dumpmachine). Otherwise a cross build (make CC=x86_64-w64-
# mingw32-gcc ... on Linux) sets IS_WIN and picks `python`, breaking clean/test
# on hosts where only `python3` exists. $(OS)=Windows_NT in every Windows shell
# (the #129 signal) and is empty on Linux/macOS.
ifeq ($(OS),Windows_NT)
PYTHON ?= python
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_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE)
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_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE)
ifneq (,$(LINUX))
TEST_BINS += tests/test_uring$(EXE)
endif
# Windows CUDA DLL path: host links the loader, NOT cudart.
ifneq ($(IS_WIN),)
@@ -148,7 +197,7 @@ endif
# runtime, so no Xcode / offline metal compiler is required. Default build unchanged.
METAL ?= 0
METAL_OBJ =
METALXX = clang++ -x objective-c++ -fobjc-arc -O3
METALXX = clang++ -x objective-c++ -std=gnu++17 -fobjc-arc -O3
ifeq ($(METAL),1)
ifeq (,$(DARWIN))
$(error METAL=1 is supported only on macOS)
@@ -177,7 +226,7 @@ $(shell printf '%s' '$(BUILD_CONFIG)' > .build-config)
endif
.build-config: ;
glm$(EXE): glm.c st.h json.h tok.h tok_unicode.h compat.h grammar.h $(CUDA_OBJ) $(METAL_OBJ) .build-config
glm$(EXE): glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h $(CUDA_OBJ) $(METAL_OBJ) .build-config
$(CC) $(CFLAGS) glm.c $(CUDA_OBJ) $(METAL_OBJ) -o glm$(EXE) $(LDFLAGS)
# Windows runtime loader object: resolves coli_cuda_* from coli_cuda.dll.
@@ -188,7 +237,7 @@ backend_loader.o: backend_loader.c backend_cuda.h compat.h .build-config
# compiler, required by nvcc on Windows) into coli_cuda.dll. Run this from a
# shell that has the MSVC environment set (e.g. after vcvars64.bat, or from a
# "x64 Native Tools Command Prompt"). COLI_CUDA_BUILDING_DLL enables
# __declspec(dllexport) so the 11 API symbols are exported.
# __declspec(dllexport) so the 15 API symbols are exported.
cuda-dll: backend_cuda.cu backend_cuda.h
@command -v "$(NVCC)" >/dev/null 2>&1 || { echo "nvcc not found: set CUDA_HOME or NVCC" >&2; exit 1; }
@command -v cl >/dev/null 2>&1 || { echo "cl.exe (MSVC) not in PATH — run vcvars64.bat first" >&2; exit 1; }
@@ -220,9 +269,23 @@ cuda-bench: backend_cuda.cu backend_cuda.h tests/bench_tensor_core.cu
olmoe$(EXE): olmoe.c st.h json.h compat.h
$(CC) $(CFLAGS) olmoe.c -o olmoe$(EXE) $(LDFLAGS)
# binario portabile da distribuire su altre macchine x86-64
# Use a baseline that matches the compiler target. macOS already targets a
# portable baseline when ARCH is empty; forcing the x86 value there breaks
# Apple Silicon. Unknown targets use native rather than an invalid x86 flag.
ifneq (,$(DARWIN))
PORTABLE_ARCH =
else ifneq (,$(AARCH64))
PORTABLE_ARCH = armv8-a
else ifneq (,$(PPC64))
PORTABLE_ARCH = power8
else ifneq (,$(X86_64))
PORTABLE_ARCH = x86-64-v3
else
PORTABLE_ARCH = native
endif
portable:
$(MAKE) glm$(EXE) ARCH=x86-64-v3
$(MAKE) glm$(EXE) ARCH=$(PORTABLE_ARCH)
iobench$(EXE): iobench.c compat.h
$(CC) $(CFLAGS) iobench.c -o iobench$(EXE) $(LDFLAGS)
@@ -239,10 +302,16 @@ tests/test_tier$(EXE): tests/test_tier.c tier.h
tests/test_grammar$(EXE): tests/test_grammar.c grammar.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_schema_gbnf$(EXE): tests/test_schema_gbnf.c schema_gbnf.h grammar.h json.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_decode_batch$(EXE): tests/test_decode_batch.c decode_batch.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_idot$(EXE): tests/test_idot.c glm.c st.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
tests/test_idot$(EXE): tests/test_idot.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)
tests/test_i4_acc512$(EXE): tests/test_i4_acc512.c
@@ -251,6 +320,9 @@ 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_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)
test-c: $(TEST_BINS)
$(PYTHON) tools/run_tests.py $(TEST_BINS)
@@ -284,4 +356,4 @@ clean:
bench: iobench$(EXE)
@if [ -n "$(ARGS)" ]; then ./iobench$(EXE) $(ARGS); else echo "built iobench$(EXE) — run: ./iobench$(EXE) <file> <MB> <iters> <threads> <direct 0|1>"; fi
.PHONY: all glm cuda-test cuda-bench cuda-dll portable test-c test-python test check clean install uninstall bench
.PHONY: all glm cuda-test cuda-bench cuda-dll portable test-c test-python test check clean install uninstall bench
+449
View File
@@ -18,12 +18,14 @@ struct ColiCudaTensor {
typedef struct {
int device;
int compute_major,compute_minor;
float *x, *y, *gate, *up;
size_t x_cap, y_cap, gate_cap, up_cap;
uint8_t *qx; float *qscale;
size_t qx_cap, qscale_cap;
float *host_x,*host_y; size_t host_x_cap,host_y_cap;
float *aq,*al,*ar,*ac; size_t aq_cap,al_cap,ar_cap,ac_cap;
float *pipe_buf[24]; size_t pipe_cap[24]; /* scratch persistenti del resident pipeline */
cudaStream_t stream;
void *group_desc; size_t group_desc_cap;
size_t tensor_count, tensor_bytes;
@@ -120,6 +122,68 @@ __global__ static void silu_mul(float *gate, const float *up, size_t n) {
}
}
/* Four warps share one A tile and compute 16x64 outputs. This matters for
* prefill: the first prototype reloaded/converter A once per 16 output cols. */
__global__ static void w4a16_matmul(float *y,const float *x,const uint8_t *w,
const float *scale,int M,int K,int N){
#if __CUDA_ARCH__ >= 700
using namespace nvcuda;int warp=threadIdx.x>>5,lane=threadIdx.x&31;
int m0=blockIdx.y*16,n0=blockIdx.x*64+warp*16;
__shared__ __half ah[256],bh[4][256];
wmma::fragment<wmma::accumulator,16,16,16,float> acc;wmma::fill_fragment(acc,0.f);
size_t rb=(size_t)(K+1)/2;
for(int k0=0;k0<K;k0+=16){
for(int z=threadIdx.x;z<256;z+=blockDim.x){
int m=z/16,k=z%16,gm=m0+m,gk=k0+k;
ah[z]=(gm<M&&gk<K)?__float2half(x[(size_t)gm*K+gk]):__float2half(0.f);
}
for(int z=lane;z<256;z+=32){
int n=z/16,gk=k0+(z%16),gn=n0+n;float v=0.f;
if(gn<N&&gk<K){uint8_t q=w[(size_t)gn*rb+(gk>>1)];int a=(gk&1)?q>>4:q&15;
v=(float)(a&8?a-16:a)*scale[gn];}
bh[warp][z]=__float2half(v); /* [Ntile,Ktile] == B col-major */
}
__syncthreads();
wmma::fragment<wmma::matrix_a,16,16,16,__half,wmma::row_major> af;
wmma::fragment<wmma::matrix_b,16,16,16,__half,wmma::col_major> bf;
wmma::load_matrix_sync(af,ah,16);wmma::load_matrix_sync(bf,bh[warp],16);
wmma::mma_sync(acc,af,bf,acc);__syncthreads();
}
__shared__ float out[4][256];wmma::store_matrix_sync(out[warp],acc,16,wmma::mem_row_major);__syncwarp();
for(int z=lane;z<256;z+=32){int m=z/16,n=z%16;
if(m0+m<M&&n0+n<N)y[(size_t)(m0+m)*N+n0+n]=out[warp][z];}
#endif
}
/* Gate and up use the same input. Eight warps compute both 16x64 projections
* while sharing the FP32->FP16 conversion of A. */
__global__ static void w4a16_gate_up(float *gate,float *up,const float *x,
const uint8_t *gw,const uint8_t *uw,const float *gs,const float *us,
int M,int K,int N){
#if __CUDA_ARCH__ >= 700
using namespace nvcuda;int warp=threadIdx.x>>5,lane=threadIdx.x&31,which=warp&1,tile=warp>>1;
int m0=blockIdx.y*16,n0=blockIdx.x*64+tile*16;const uint8_t *w=which?uw:gw;
const float *scale=which?us:gs;float *y=which?up:gate;size_t rb=(size_t)(K+1)/2;
__shared__ __half ah[256],bh[8][256];
wmma::fragment<wmma::accumulator,16,16,16,float> acc;wmma::fill_fragment(acc,0.f);
for(int k0=0;k0<K;k0+=16){
for(int z=threadIdx.x;z<256;z+=blockDim.x){int m=z/16,k=z%16,gm=m0+m,gk=k0+k;
ah[z]=(gm<M&&gk<K)?__float2half(x[(size_t)gm*K+gk]):__float2half(0.f);}
for(int z=lane;z<256;z+=32){int n=z/16,gk=k0+(z%16),gn=n0+n;float v=0.f;
if(gn<N&&gk<K){uint8_t q=w[(size_t)gn*rb+(gk>>1)];int a=(gk&1)?q>>4:q&15;
v=(float)(a&8?a-16:a)*scale[gn];}bh[warp][z]=__float2half(v);}
__syncthreads();
wmma::fragment<wmma::matrix_a,16,16,16,__half,wmma::row_major> af;
wmma::fragment<wmma::matrix_b,16,16,16,__half,wmma::col_major> bf;
wmma::load_matrix_sync(af,ah,16);wmma::load_matrix_sync(bf,bh[warp],16);
wmma::mma_sync(acc,af,bf,acc);__syncthreads();
}
__shared__ float out[8][256];wmma::store_matrix_sync(out[warp],acc,16,wmma::mem_row_major);__syncwarp();
for(int z=lane;z<256;z+=32){int m=z/16,n=z%16;
if(m0+m<M&&n0+n<N)y[(size_t)(m0+m)*N+n0+n]=out[warp][z];}
#endif
}
__global__ static void quantize_s4_rows(uint8_t *q,float *scale,const float *x,int S,int K){
int s=blockIdx.x; if(s>=S)return; const float *xs=x+(size_t)s*K;
float v=0; for(int i=threadIdx.x;i<K;i+=blockDim.x)v=fmaxf(v,fabsf(xs[i]));
@@ -243,6 +307,37 @@ __global__ static void attention_absorb_kernel(float *ctx,const float *q,const f
for(int k=0;k<K;k++)a+=cl[k]*weight_at(weights,fmt,(size_t)row*rb,k);ctx[(size_t)h*V+v]=a*(fmt?wscale[row]:1.f);}
}
__global__ static void attention_absorb_batch_kernel(float *ctx,const float *q,
const float *latent,const float *rope,const void *weights,const float *wscale,
int fmt,int S,int H,int Q,int R,int V,int K,int T,float scale){
int s=blockIdx.y,h=blockIdx.x,tid=threadIdx.x,nt=T-S+s+1,rbase=h*(Q+V);
if(s>=S||nt<1)return;
extern __shared__ float sm[];float *qa=sm,*cl=qa+K,*scores=cl+K,*red=scores+T;
const float *qs=q+((size_t)s*H+h)*(Q+R);
for(int k=tid;k<K;k+=blockDim.x){float a=0;for(int d=0;d<Q;d++)
a+=qs[d]*weight_at(weights,fmt,(size_t)(rbase+d)*row_bytes(fmt,K),k)*
(fmt?wscale[rbase+d]:1.f);qa[k]=a;}
__syncthreads();
for(int t=tid;t<nt;t+=blockDim.x){float a=0;const float *lt=latent+(size_t)t*K;
const float *rt=rope+(size_t)t*R;for(int k=0;k<K;k++)a+=qa[k]*lt[k];
for(int d=0;d<R;d++)a+=qs[Q+d]*rt[d];scores[t]=a*scale;}
__syncthreads();
float local=-3.402823466e+38F;for(int t=tid;t<nt;t+=blockDim.x)local=fmaxf(local,scores[t]);
red[tid]=local;__syncthreads();
for(int n=blockDim.x>>1;n;n>>=1){if(tid<n)red[tid]=fmaxf(red[tid],red[tid+n]);__syncthreads();}
float mx=red[0];local=0;for(int t=tid;t<nt;t+=blockDim.x){float e=expf(scores[t]-mx);scores[t]=e;local+=e;}
red[tid]=local;__syncthreads();
for(int n=blockDim.x>>1;n;n>>=1){if(tid<n)red[tid]+=red[tid+n];__syncthreads();}
float inv=1.f/red[0];for(int t=tid;t<nt;t+=blockDim.x)scores[t]*=inv;
__syncthreads();
for(int k=tid;k<K;k+=blockDim.x){float a=0;for(int t=0;t<nt;t++)
a+=scores[t]*latent[(size_t)t*K+k];cl[k]=a;}
__syncthreads();
for(int v=tid;v<V;v+=blockDim.x){int row=rbase+Q+v;float a=0;size_t rb=row_bytes(fmt,K);
for(int k=0;k<K;k++)a+=cl[k]*weight_at(weights,fmt,(size_t)row*rb,k);
ctx[((size_t)s*H+h)*V+v]=a*(fmt?wscale[row]:1.f);}
}
static int reserve(float **ptr, size_t *cap, size_t bytes) {
if (*cap >= bytes) return 1;
if (*ptr) cudaFree(*ptr);
@@ -286,6 +381,7 @@ extern "C" int coli_cuda_init(const int *devices, int count) {
if (!select_ctx(ctx)) { g_nctx = 0; return 0; }
cudaDeviceProp prop{};
if (!cuda_ok(cudaGetDeviceProperties(&prop, device), "device properties")) { g_nctx = 0; return 0; }
ctx->compute_major=prop.major;ctx->compute_minor=prop.minor;
if(!cuda_ok(cudaStreamCreateWithFlags(&ctx->stream,cudaStreamNonBlocking),"stream creation")){
g_nctx=0;return 0;
}
@@ -307,6 +403,7 @@ extern "C" void coli_cuda_shutdown(void) {
if (ctx->qx) cudaFree(ctx->qx);
if (ctx->qscale) cudaFree(ctx->qscale);
if(ctx->aq)cudaFree(ctx->aq);if(ctx->al)cudaFree(ctx->al);if(ctx->ar)cudaFree(ctx->ar);if(ctx->ac)cudaFree(ctx->ac);
for(int b=0;b<24;b++) if(ctx->pipe_buf[b]) cudaFree(ctx->pipe_buf[b]);
if (ctx->host_x) cudaFreeHost(ctx->host_x);
if (ctx->host_y) cudaFreeHost(ctx->host_y);
if (ctx->stream) cudaStreamDestroy(ctx->stream);
@@ -388,6 +485,23 @@ extern "C" int coli_cuda_tensor_upload(ColiCudaTensor **tensor,
return 1;
}
extern "C" int coli_cuda_tensor_update(ColiCudaTensor *tensor,
const void *weights,
const float *scales) {
if (!tensor || !weights || (tensor->fmt && !scales)) return 0;
DeviceContext *ctx=find_ctx(tensor->device);
if (!select_ctx(ctx)) return 0;
if (!cuda_ok(cudaMemcpy(tensor->weights,weights,tensor->weight_bytes,
cudaMemcpyHostToDevice),"tensor refresh")) return 0;
if(tensor->fmt==2){
offset_to_signed_s4<<<(unsigned)((tensor->weight_bytes+255)/256),256>>>(
(uint8_t*)tensor->weights,tensor->weight_bytes);
if(!cuda_ok(cudaGetLastError(),"int4 weight refresh")) return 0;
}
return !tensor->fmt || cuda_ok(cudaMemcpy(tensor->scales,scales,
(size_t)tensor->O*sizeof(float),cudaMemcpyHostToDevice),"scale refresh");
}
extern "C" int coli_cuda_matmul(ColiCudaTensor **tensor,
float *y, const float *x,
const void *weights, const float *scales,
@@ -436,6 +550,34 @@ extern "C" int coli_cuda_expert_mlp(ColiCudaTensor *gate, ColiCudaTensor *up,
return 1;
}
extern "C" int coli_cuda_shared_mlp_w4a16(ColiCudaTensor *gate,ColiCudaTensor *up,
ColiCudaTensor *down,float *y,const float *x,int S){
if(!gate||!up||!down||!x||!y||S<1||gate->fmt!=2||up->fmt!=2||down->fmt!=2||
gate->device!=up->device||gate->device!=down->device||gate->I!=up->I||
gate->O!=up->O||down->I!=gate->O||down->O!=gate->I)return 0;
DeviceContext *ctx=find_ctx(gate->device);if(!select_ctx(ctx)||ctx->compute_major<7)return 0;
int D=gate->I,I=gate->O;size_t xb=(size_t)S*D*sizeof(float),ib=(size_t)S*I*sizeof(float);
if(!reserve(&ctx->x,&ctx->x_cap,xb)||!reserve(&ctx->gate,&ctx->gate_cap,ib)||
!reserve(&ctx->up,&ctx->up_cap,ib)||!reserve(&ctx->y,&ctx->y_cap,xb)||
!reserve_pinned(&ctx->host_x,&ctx->host_x_cap,xb)||
!reserve_pinned(&ctx->host_y,&ctx->host_y_cap,xb))return 0;
std::memcpy(ctx->host_x,x,xb);
if(!cuda_ok(cudaMemcpyAsync(ctx->x,ctx->host_x,xb,cudaMemcpyHostToDevice,ctx->stream),
"shared w4a16 input upload"))return 0;
dim3 hidden((unsigned)((I+63)/64),(unsigned)((S+15)/16));
dim3 output((unsigned)((D+63)/64),(unsigned)((S+15)/16));
w4a16_gate_up<<<hidden,256,0,ctx->stream>>>(ctx->gate,ctx->up,ctx->x,
(const uint8_t*)gate->weights,(const uint8_t*)up->weights,gate->scales,up->scales,S,D,I);
silu_mul<<<(unsigned)(((size_t)S*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)S*I);
w4a16_matmul<<<output,128,0,ctx->stream>>>(ctx->y,ctx->gate,(const uint8_t*)down->weights,down->scales,S,I,D);
if(!cuda_ok(cudaGetLastError(),"shared w4a16 launch")||
!cuda_ok(cudaMemcpyAsync(ctx->host_y,ctx->y,xb,cudaMemcpyDeviceToHost,ctx->stream),
"shared w4a16 output download")||
!cuda_ok(cudaStreamSynchronize(ctx->stream),"shared w4a16 synchronize"))return 0;
std::memcpy(y,ctx->host_y,xb);
return 1;
}
extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates,
ColiCudaTensor *const *ups,
ColiCudaTensor *const *downs,
@@ -493,6 +635,38 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates,
silu_mul<<<(unsigned)(((size_t)total*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)total*I);
quantize_s4_rows<<<total,256,0,ctx->stream>>>(ctx->qx,ctx->qscale,ctx->gate,total,I);
grouped_s4_wmma<<<dim3((unsigned)((D+63)/64),(unsigned)count),256,0,ctx->stream>>>(ctx->y,ctx->qx,ctx->qscale,dev,I,D,2);
}else if(all_s4&&ctx->compute_major>=7&&getenv("COLI_CUDA_TC_W4A16")&&
atoi(getenv("COLI_CUDA_TC_W4A16"))){
/* W4A16 Tensor Core per gruppo: attivazioni fp16 per tile (lossless al
* contrario del path W4A4), un lancio per expert dentro lo stream —
* l'overhead di lancio e' trascurabile rispetto ai GEMM. */
int tc16_min=getenv("COLI_CUDA_TC_W4A16_MIN")?atoi(getenv("COLI_CUDA_TC_W4A16_MIN")):16;
int off16=0;
for(int c=0;c<count;c++){
int r=rows[c];
float *g16=ctx->gate+(size_t)off16*I,*u16=ctx->up+(size_t)off16*I;
float *x16=ctx->x+(size_t)off16*D,*y16=ctx->y+(size_t)off16*D;
if(r>=tc16_min){
dim3 hg16((unsigned)((I+63)/64),(unsigned)((r+15)/16));
dim3 og16((unsigned)((D+63)/64),(unsigned)((r+15)/16));
w4a16_gate_up<<<hg16,256,0,ctx->stream>>>(g16,u16,x16,
(const uint8_t*)host[c].g,(const uint8_t*)host[c].u,host[c].gs,host[c].us,r,D,I);
silu_mul<<<(unsigned)(((size_t)r*I+255)/256),256,0,ctx->stream>>>(g16,u16,(size_t)r*I);
w4a16_matmul<<<og16,128,0,ctx->stream>>>(y16,g16,
(const uint8_t*)host[c].d,host[c].ds,r,I,D);
}else{
/* piccoli batch: tile TC quasi vuoti + overhead di lancio — il
* kernel naive per-elemento resta piu' veloce (misurato in decode) */
quant_matmul<<<dim3((unsigned)I,(unsigned)r),256,0,ctx->stream>>>(g16,x16,
host[c].g,host[c].gs,host[c].gf,r,D,I,row_bytes(host[c].gf,D));
quant_matmul<<<dim3((unsigned)I,(unsigned)r),256,0,ctx->stream>>>(u16,x16,
host[c].u,host[c].us,host[c].uf,r,D,I,row_bytes(host[c].uf,D));
silu_mul<<<(unsigned)(((size_t)r*I+255)/256),256,0,ctx->stream>>>(g16,u16,(size_t)r*I);
quant_matmul<<<dim3((unsigned)D,(unsigned)r),256,0,ctx->stream>>>(y16,g16,
host[c].d,host[c].ds,host[c].df,r,I,D,row_bytes(host[c].df,I));
}
off16+=r;
}
}else if(all_s4&&(!getenv("COLI_CUDA_W4_PACKED")||atoi(getenv("COLI_CUDA_W4_PACKED")))){
dim3 hg((unsigned)I,(unsigned)max_rows,(unsigned)count),og((unsigned)D,(unsigned)max_rows,(unsigned)count);
int dual=!getenv("COLI_CUDA_DUAL_PROJ")||atoi(getenv("COLI_CUDA_DUAL_PROJ"));
@@ -530,6 +704,7 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates,
return 1;
}
extern "C" int coli_cuda_attention_absorb(ColiCudaTensor *w,float *ctx,const float *q,
const float *latent,const float *rope,int H,int Q,
int R,int V,int K,int T,float scale){
@@ -552,6 +727,49 @@ extern "C" int coli_cuda_attention_absorb(ColiCudaTensor *w,float *ctx,const flo
return 1;
}
static int attention_absorb_batch_run(ColiCudaTensor *w,ColiCudaTensor *proj,float *out,
const float *q,const float *latent,const float *rope,int S,int H,int Q,int R,int V,
int K,int T,float scale){
if(!w||!out||!q||!latent||!rope||S<1||H<1||Q<1||R<1||V<1||K<1||K>512||
T<S||T>8192||w->I!=K||w->O!=H*(Q+V))return 0;
if(proj&&(proj->device!=w->device||proj->I!=H*V))return 0;
DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0;
size_t qb=(size_t)S*H*(Q+R)*sizeof(float),lb=(size_t)T*K*sizeof(float);
size_t rb=(size_t)T*R*sizeof(float),cb=(size_t)S*H*V*sizeof(float);
if(!reserve(&dc->aq,&dc->aq_cap,qb)||!reserve(&dc->al,&dc->al_cap,lb)||
!reserve(&dc->ar,&dc->ar_cap,rb)||!reserve(&dc->ac,&dc->ac_cap,cb))return 0;
if(!cuda_ok(cudaMemcpyAsync(dc->aq,q,qb,cudaMemcpyHostToDevice,dc->stream),"attention batch q upload")||
!cuda_ok(cudaMemcpyAsync(dc->al,latent,lb,cudaMemcpyHostToDevice,dc->stream),"attention batch latent upload")||
!cuda_ok(cudaMemcpyAsync(dc->ar,rope,rb,cudaMemcpyHostToDevice,dc->stream),"attention batch rope upload"))return 0;
size_t shared=(size_t)(2*K+T+256)*sizeof(float);
attention_absorb_batch_kernel<<<dim3(H,S),256,shared,dc->stream>>>(dc->ac,dc->aq,dc->al,
dc->ar,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale);
if(!cuda_ok(cudaGetLastError(),"attention batch launch"))return 0;
const float *src=dc->ac;size_t ob=cb;
if(proj){
ob=(size_t)S*proj->O*sizeof(float);if(!reserve(&dc->y,&dc->y_cap,ob))return 0;
quant_matmul<<<dim3(proj->O,S),256,0,dc->stream>>>(dc->y,dc->ac,proj->weights,
proj->scales,proj->fmt,S,proj->I,proj->O,row_bytes(proj->fmt,proj->I));
if(!cuda_ok(cudaGetLastError(),"attention o_proj launch"))return 0;src=dc->y;
}
if(!cuda_ok(cudaMemcpyAsync(out,src,ob,cudaMemcpyDeviceToHost,dc->stream),
proj?"attention projected output download":"attention batch context download")||
!cuda_ok(cudaStreamSynchronize(dc->stream),"attention batch synchronize"))return 0;
return 1;
}
extern "C" int coli_cuda_attention_absorb_batch(ColiCudaTensor *w,float *ctx,const float *q,
const float *latent,const float *rope,int S,int H,int Q,int R,int V,int K,int T,
float scale){
return attention_absorb_batch_run(w,nullptr,ctx,q,latent,rope,S,H,Q,R,V,K,T,scale);
}
extern "C" int coli_cuda_attention_project_batch(ColiCudaTensor *w,ColiCudaTensor *proj,
float *out,const float *q,const float *latent,const float *rope,int S,int H,int Q,
int R,int V,int K,int T,float scale){
return attention_absorb_batch_run(w,proj,out,q,latent,rope,S,H,Q,R,V,K,T,scale);
}
extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) {
if (!tensor) return;
DeviceContext *ctx = find_ctx(tensor->device);
@@ -573,3 +791,234 @@ extern "C" size_t coli_cuda_tensor_bytes(const ColiCudaTensor *tensor) {
extern "C" int coli_cuda_tensor_device(const ColiCudaTensor *tensor) {
return tensor ? tensor->device : -1;
}
/* ==== resident-pipeline primitives (Inc.0, 2026-07-13) ====
* Device-side building blocks so the residual stream can stay on the layer's
* home device across a whole layer. Control flow stays on CPU; only the data
* plane lives here. All entry points take DEVICE pointers (no transfers) —
* the caller owns staging via the pipe buffer API below. */
__global__ static void pipe_rmsnorm_rows(float *y,const float *x,const float *w,
int D,float eps,int xstride,int ystride){
const float *xr=x+(size_t)blockIdx.x*xstride; float *yr=y+(size_t)blockIdx.x*ystride;
__shared__ double sh[256];
double a=0; for(int i=threadIdx.x;i<D;i+=blockDim.x){ double v=xr[i]; a+=v*v; }
sh[threadIdx.x]=a; __syncthreads();
for(int s=blockDim.x/2;s>0;s>>=1){ if(threadIdx.x<s) sh[threadIdx.x]+=sh[threadIdx.x+s]; __syncthreads(); }
float r=rsqrtf((float)(sh[0]/D)+eps);
for(int i=threadIdx.x;i<D;i+=blockDim.x) yr[i]=xr[i]*r*w[i];
}
/* RoPE interleaved, identical math to glm.c rope_interleave. One block per row;
* row layout: v + row*stride + offset holds R floats. pos index = row/heads
* (heads=1 for k_rot rows, heads=H for [S,H,qh] query rows). */
__global__ static void pipe_rope_rows(float *v,const int *pos,int pos_base,int stride,
int offset,int R,int heads,float theta){
float *p=v+(size_t)blockIdx.x*stride+offset;
int half=R/2, ps=pos?pos[blockIdx.x/heads]:pos_base+(int)(blockIdx.x/heads);
__shared__ float in[256];
for(int j=threadIdx.x;j<R;j+=blockDim.x) in[j]=p[j];
__syncthreads();
for(int j=threadIdx.x;j<half;j+=blockDim.x){
float inv=__powf(theta,-2.0f*j/R);
float ang=ps*inv, cs=__cosf(ang), sn=__sinf(ang);
float a=in[2*j], b=in[2*j+1];
p[j]=a*cs-b*sn; p[half+j]=b*cs+a*sn;
}
}
__global__ static void pipe_add_n(float *x,const float *t,size_t n){
size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x;
if(i<n) x[i]+=t[i];
}
/* Fixed-order partial merge: block b adds partial row b into x row rows[b].
* Target rows are unique by construction (CPU pre-sums per token), so no
* atomics — the 9.20.7 lesson. */
__global__ static void pipe_rows_add(float *x,const float *partial,const int *rows,
int D){
float *xr=x+(size_t)rows[blockIdx.x]*D;
const float *pr=partial+(size_t)blockIdx.x*D;
for(int i=threadIdx.x;i<D;i+=blockDim.x) xr[i]+=pr[i];
}
/* scratch persistente per (device,slot): cresce e resta — niente cudaMalloc/Free
* per layer (78 x ~10 alloc/richiesta erano puro churn). */
extern "C" float *coli_cuda_pipe_scratch(int device,int slot,size_t bytes){
DeviceContext *ctx=find_ctx(device);
if(slot<0||slot>=24||!select_ctx(ctx)) return NULL;
if(!reserve(&ctx->pipe_buf[slot],&ctx->pipe_cap[slot],bytes)) return NULL;
return ctx->pipe_buf[slot];
}
extern "C" void *coli_cuda_pipe_alloc(int device,size_t bytes){
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return NULL;
void *p=NULL;
if(!cuda_ok(cudaMalloc(&p,bytes),"pipe alloc")) return NULL;
return p;
}
extern "C" void coli_cuda_pipe_free(int device,void *p){
DeviceContext *ctx=find_ctx(device); if(!p||!select_ctx(ctx)) return;
cudaFree(p);
}
extern "C" int coli_cuda_pipe_upload(int device,void *dst,const void *src,size_t bytes){
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0;
return cuda_ok(cudaMemcpy(dst,src,bytes,cudaMemcpyHostToDevice),"pipe upload");
}
extern "C" int coli_cuda_pipe_download(int device,const void *src,void *dst,size_t bytes){
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0;
return cuda_ok(cudaMemcpy(dst,src,bytes,cudaMemcpyDeviceToHost),"pipe download");
}
extern "C" int coli_cuda_pipe_rmsnorm(int device,float *y_dev,const float *x_dev,
const float *w_dev,int S,int D,float eps){
DeviceContext *ctx=find_ctx(device);
if(S<1||D<1||!select_ctx(ctx)) return 0;
pipe_rmsnorm_rows<<<S,256>>>(y_dev,x_dev,w_dev,D,eps,D,D);
return cuda_ok(cudaGetLastError(),"pipe rmsnorm");
}
extern "C" int coli_cuda_pipe_rmsnorm_s(int device,float *y_dev,const float *x_dev,
const float *w_dev,int S,int D,float eps,
int xstride,int ystride){
DeviceContext *ctx=find_ctx(device);
if(S<1||D<1||xstride<D||ystride<D||!select_ctx(ctx)) return 0;
pipe_rmsnorm_rows<<<S,256>>>(y_dev,x_dev,w_dev,D,eps,xstride,ystride);
return cuda_ok(cudaGetLastError(),"pipe rmsnorm strided");
}
extern "C" int coli_cuda_pipe_rope(int device,float *v_dev,const int *pos_dev,
int rows,int stride,int offset,int R,int heads,
float theta){
DeviceContext *ctx=find_ctx(device);
if(rows<1||R<2||R>256||heads<1||!select_ctx(ctx)) return 0;
pipe_rope_rows<<<rows,128>>>(v_dev,pos_dev,0,stride,offset,R,heads,theta);
return cuda_ok(cudaGetLastError(),"pipe rope");
}
extern "C" int coli_cuda_pipe_rope_base(int device,float *v_dev,int pos_base,int rows,
int stride,int offset,int R,int heads,float theta){
DeviceContext *ctx=find_ctx(device);
if(rows<1||R<2||R>256||heads<1||!select_ctx(ctx)) return 0;
pipe_rope_rows<<<rows,128>>>(v_dev,NULL,pos_base,stride,offset,R,heads,theta);
return cuda_ok(cudaGetLastError(),"pipe rope base");
}
extern "C" int coli_cuda_pipe_copy2d(int device,float *dst,int dpitch,const float *src,
int spitch,int width,int height){
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0;
return cuda_ok(cudaMemcpy2D(dst,(size_t)dpitch*4,src,(size_t)spitch*4,
(size_t)width*4,height,cudaMemcpyDeviceToDevice),"pipe copy2d");
}
/* attention batch + fused o_proj with DEVICE-resident q/latent/rope: the whole
* upstream projection chain stayed on this device, so nothing is uploaded here.
* Only the final [S,O] projection is downloaded to host. */
extern "C" int coli_cuda_attention_project_batch_dev(ColiCudaTensor *w,ColiCudaTensor *proj,
float *out,const float *q_dev,const float *latent_dev,const float *rope_dev,
int S,int H,int Q,int R,int V,int K,int T,float scale){
if(!w||!proj||!out||!q_dev||!latent_dev||!rope_dev||S<1||H<1||Q<1||R<1||V<1||
K<1||K>512||T<S||T>8192||w->I!=K||w->O!=H*(Q+V)||
proj->device!=w->device||proj->I!=H*V)return 0;
DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0;
size_t cb=(size_t)S*H*V*sizeof(float);
if(!reserve(&dc->ac,&dc->ac_cap,cb))return 0;
size_t shared=(size_t)(2*K+T+256)*sizeof(float);
attention_absorb_batch_kernel<<<dim3(H,S),256,shared,dc->stream>>>(dc->ac,q_dev,latent_dev,
rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale);
if(!cuda_ok(cudaGetLastError(),"pipe attention launch"))return 0;
size_t ob=(size_t)S*proj->O*sizeof(float);
if(!reserve(&dc->y,&dc->y_cap,ob))return 0;
quant_matmul<<<dim3(proj->O,S),256,0,dc->stream>>>(dc->y,dc->ac,proj->weights,
proj->scales,proj->fmt,S,proj->I,proj->O,row_bytes(proj->fmt,proj->I));
if(!cuda_ok(cudaGetLastError(),"pipe o_proj launch"))return 0;
if(!cuda_ok(cudaMemcpyAsync(out,dc->y,ob,cudaMemcpyDeviceToHost,dc->stream),"pipe attention download")||
!cuda_ok(cudaStreamSynchronize(dc->stream),"pipe attention sync"))return 0;
return 1;
}
extern "C" int coli_cuda_pipe_silu_mul(int device,float *gate_dev,const float *up_dev,
size_t n){
DeviceContext *ctx=find_ctx(device); if(!n||!select_ctx(ctx)) return 0;
silu_mul<<<(unsigned)((n+255)/256),256>>>(gate_dev,up_dev,n);
return cuda_ok(cudaGetLastError(),"pipe silu mul");
}
extern "C" int coli_cuda_pipe_add(int device,float *x_dev,const float *t_dev,size_t n){
DeviceContext *ctx=find_ctx(device); if(!n||!select_ctx(ctx)) return 0;
pipe_add_n<<<(unsigned)((n+255)/256),256>>>(x_dev,t_dev,n);
return cuda_ok(cudaGetLastError(),"pipe add");
}
extern "C" int coli_cuda_pipe_rows_add(int device,float *x_dev,const float *partial_dev,
const int *rows_dev,int nrows,int D){
DeviceContext *ctx=find_ctx(device); if(nrows<1||D<1||!select_ctx(ctx)) return 0;
pipe_rows_add<<<nrows,256>>>(x_dev,partial_dev,rows_dev,D);
return cuda_ok(cudaGetLastError(),"pipe rows add");
}
/* GEMM with device-resident activations: same quant_matmul kernel as
* coli_cuda_matmul, zero host transfers. */
extern "C" int coli_cuda_pipe_gemm(ColiCudaTensor *t,float *y_dev,const float *x_dev,
int S){
if(!t||S<1) return 0;
DeviceContext *ctx=find_ctx(t->device); if(!select_ctx(ctx)) return 0;
dim3 grid((unsigned)t->O,(unsigned)S);
quant_matmul<<<grid,256>>>(y_dev,x_dev,t->weights,t->scales,t->fmt,S,t->I,t->O,
row_bytes(t->fmt,t->I));
return cuda_ok(cudaGetLastError(),"pipe gemm");
}
/* copia diretta scheda->scheda (P2P se disponibile, altrimenti staging driver) */
extern "C" int coli_cuda_pipe_peer_copy(int dst_dev,float *dst,int src_dev,
const float *src,size_t bytes){
if(!dst||!src) return 0;
if(dst_dev==src_dev){ DeviceContext *c=find_ctx(dst_dev); if(!select_ctx(c)) return 0;
return cuda_ok(cudaMemcpy(dst,src,bytes,cudaMemcpyDeviceToDevice),"pipe intra copy"); }
return cuda_ok(cudaMemcpyPeer(dst,dst_dev,src,src_dev,bytes),"pipe peer copy");
}
/* come attention_project_batch_dev ma l'uscita di o_proj RESTA sul device (out_dev). */
extern "C" int coli_cuda_attention_project_batch_dev_out(ColiCudaTensor *w,ColiCudaTensor *proj,
float *out_dev,const float *q_dev,const float *latent_dev,const float *rope_dev,
int S,int H,int Q,int R,int V,int K,int T,float scale){
if(!w||!proj||!out_dev||!q_dev||!latent_dev||!rope_dev||S<1||H<1||Q<1||R<1||V<1||
K<1||K>512||T<S||T>8192||w->I!=K||w->O!=H*(Q+V)||
proj->device!=w->device||proj->I!=H*V)return 0;
DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0;
size_t cb=(size_t)S*H*V*sizeof(float);
if(!reserve(&dc->ac,&dc->ac_cap,cb))return 0;
size_t shared=(size_t)(2*K+T+256)*sizeof(float);
attention_absorb_batch_kernel<<<dim3(H,S),256,shared,dc->stream>>>(dc->ac,q_dev,latent_dev,
rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale);
if(!cuda_ok(cudaGetLastError(),"pipe attention launch (dev out)"))return 0;
quant_matmul<<<dim3(proj->O,S),256,0,dc->stream>>>(out_dev,dc->ac,proj->weights,
proj->scales,proj->fmt,S,proj->I,proj->O,row_bytes(proj->fmt,proj->I));
if(!cuda_ok(cudaGetLastError(),"pipe o_proj launch (dev out)"))return 0;
return cuda_ok(cudaStreamSynchronize(dc->stream),"pipe attention sync (dev out)");
}
/* absorb batch con TUTTO su device (q/latent/rope gia' residenti sulla scheda
* dello shard, ctx resta sul device): il cuore della attention head-shardata
* dentro il pipeline. Nessun trasferimento host. */
extern "C" int coli_cuda_attention_absorb_batch_dev(ColiCudaTensor *w,float *ctx_dev,
const float *q_dev,const float *latent_dev,const float *rope_dev,
int S,int H,int Q,int R,int V,int K,int T,float scale){
if(!w||!ctx_dev||!q_dev||!latent_dev||!rope_dev||S<1||H<1||Q<1||R<1||V<1||
K<1||K>512||T<S||T>8192||w->I!=K||w->O!=H*(Q+V))return 0;
DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0;
size_t shared=(size_t)(2*K+T+256)*sizeof(float);
attention_absorb_batch_kernel<<<dim3(H,S),256,shared,dc->stream>>>(ctx_dev,q_dev,latent_dev,
rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale);
if(!cuda_ok(cudaGetLastError(),"pipe shard attention launch"))return 0;
return cuda_ok(cudaStreamSynchronize(dc->stream),"pipe shard attention sync");
}
/* absorb per il DECODE con KV gia' residente: carica solo q (poche KB),
* latent/rope arrivano dall'ombra device. ctx torna a host (S piccolo). */
extern "C" int coli_cuda_attention_absorb_kvdev(ColiCudaTensor *w,float *ctx,const float *q,
const float *latent_dev,const float *rope_dev,int H,int Q,int R,int V,int K,int T,
float scale){
if(!w||!ctx||!q||!latent_dev||!rope_dev||H<1||Q<1||R<1||V<1||K<1||K>512||T<1||T>8192||
w->I!=K||w->O!=H*(Q+V))return 0;
DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0;
size_t qb=(size_t)H*(Q+R)*sizeof(float),cb=(size_t)H*V*sizeof(float);
if(!reserve(&dc->aq,&dc->aq_cap,qb)||!reserve(&dc->ac,&dc->ac_cap,cb))return 0;
if(!cuda_ok(cudaMemcpyAsync(dc->aq,q,qb,cudaMemcpyHostToDevice,dc->stream),"kvdev q upload"))return 0;
size_t shared=(size_t)(2*K+T+256)*sizeof(float);
attention_absorb_batch_kernel<<<dim3(H,1),256,shared,dc->stream>>>(dc->ac,dc->aq,latent_dev,
rope_dev,w->weights,w->scales,w->fmt,1,H,Q,R,V,K,T,scale);
if(!cuda_ok(cudaGetLastError(),"kvdev absorb launch")||
!cuda_ok(cudaMemcpyAsync(ctx,dc->ac,cb,cudaMemcpyDeviceToHost,dc->stream),"kvdev ctx download")||
!cuda_ok(cudaStreamSynchronize(dc->stream),"kvdev absorb sync"))return 0;
return 1;
}
extern "C" int coli_cuda_pipe_sync(int device){
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0;
return cuda_ok(cudaDeviceSynchronize(),"pipe sync");
}
+66
View File
@@ -56,6 +56,14 @@ COLI_CUDA_DLLEXPORT int coli_cuda_matmul(ColiCudaTensor **tensor,
COLI_CUDA_DLLEXPORT int coli_cuda_expert_mlp(ColiCudaTensor *gate, ColiCudaTensor *up,
ColiCudaTensor *down, float *y, const float *x, int S);
/* Prefill-oriented shared expert path. INT4 weights stay packed in global
* memory, activations are converted to FP16 per tile, and Tensor Cores
* accumulate into FP32. Unlike COLI_CUDA_TC_INT4 this does not quantize the
* activation to INT4. */
COLI_CUDA_DLLEXPORT int coli_cuda_shared_mlp_w4a16(ColiCudaTensor *gate, ColiCudaTensor *up,
ColiCudaTensor *down, float *y,
const float *x, int S);
/* Packed group of same-shaped experts. Inputs and outputs contain sum(rows)
* consecutive [D] rows in call order. */
COLI_CUDA_DLLEXPORT int coli_cuda_expert_group(ColiCudaTensor *const *gates,
@@ -69,12 +77,70 @@ COLI_CUDA_DLLEXPORT int coli_cuda_attention_absorb(ColiCudaTensor *kv_b,float *c
const float *latent,const float *rope,int H,int Q,
int R,int V,int K,int T,float attention_scale);
/* Causal MLA absorption for S contiguous rows from one sequence. The KV
* arrays contain T rows ending at the final query; query s attends T-S+s+1
* rows. One transfer and one launch replace S host round-trips. */
COLI_CUDA_DLLEXPORT int coli_cuda_attention_absorb_batch(ColiCudaTensor *kv_b,float *ctx,const float *q,
const float *latent,const float *rope,int S,
int H,int Q,int R,int V,int K,int T,
float attention_scale);
/* Same attention batch followed immediately by resident o_proj on the same
* device. Only the final [S,D] tensor crosses back to the host. */
COLI_CUDA_DLLEXPORT int coli_cuda_attention_project_batch(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj,
float *out,const float *q,const float *latent,
const float *rope,int S,int H,int Q,int R,
int V,int K,int T,float attention_scale);
COLI_CUDA_DLLEXPORT void coli_cuda_tensor_free(ColiCudaTensor *tensor);
COLI_CUDA_DLLEXPORT size_t coli_cuda_tensor_bytes(const ColiCudaTensor *tensor);
COLI_CUDA_DLLEXPORT int coli_cuda_tensor_device(const ColiCudaTensor *tensor);
/* Replace a resident tensor's contents without reallocating its device slot. */
COLI_CUDA_DLLEXPORT int coli_cuda_tensor_update(ColiCudaTensor *tensor,
const void *weights, const float *scales);
/* ---- resident-pipeline primitives (Inc.0): device-pointer entry points ---- */
COLI_CUDA_DLLEXPORT float *coli_cuda_pipe_scratch(int device,int slot,size_t bytes);
COLI_CUDA_DLLEXPORT void *coli_cuda_pipe_alloc(int device,size_t bytes);
COLI_CUDA_DLLEXPORT void coli_cuda_pipe_free(int device,void *p);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_upload(int device,void *dst,const void *src,size_t bytes);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_download(int device,const void *src,void *dst,size_t bytes);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_rmsnorm(int device,float *y_dev,const float *x_dev,
const float *w_dev,int S,int D,float eps);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_rope(int device,float *v_dev,const int *pos_dev,int rows,
int stride,int offset,int R,int heads,float theta);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_silu_mul(int device,float *gate_dev,const float *up_dev,size_t n);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_add(int device,float *x_dev,const float *t_dev,size_t n);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_rows_add(int device,float *x_dev,const float *partial_dev,
const int *rows_dev,int nrows,int D);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_gemm(ColiCudaTensor *t,float *y_dev,const float *x_dev,int S);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_rmsnorm_s(int device,float *y_dev,const float *x_dev,
const float *w_dev,int S,int D,float eps,
int xstride,int ystride);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_rope_base(int device,float *v_dev,int pos_base,int rows,
int stride,int offset,int R,int heads,float theta);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_copy2d(int device,float *dst,int dpitch,const float *src,
int spitch,int width,int height);
COLI_CUDA_DLLEXPORT int coli_cuda_attention_project_batch_dev(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj,
float *out,const float *q_dev,const float *latent_dev,const float *rope_dev,
int S,int H,int Q,int R,int V,int K,int T,float scale);
COLI_CUDA_DLLEXPORT int coli_cuda_attention_absorb_batch_dev(ColiCudaTensor *kv_b_shard,float *ctx_dev,
const float *q_dev,const float *latent_dev,const float *rope_dev,
int S,int H,int Q,int R,int V,int K,int T,float scale);
COLI_CUDA_DLLEXPORT int coli_cuda_attention_absorb_kvdev(ColiCudaTensor *kv_b,float *ctx,const float *q,
const float *latent_dev,const float *rope_dev,int H,int Q,int R,int V,int K,int T,
float scale);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_peer_copy(int dst_dev,float *dst,int src_dev,
const float *src,size_t bytes);
COLI_CUDA_DLLEXPORT int coli_cuda_attention_project_batch_dev_out(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj,
float *out_dev,const float *q_dev,const float *latent_dev,const float *rope_dev,
int S,int H,int Q,int R,int V,int K,int T,float scale);
COLI_CUDA_DLLEXPORT int coli_cuda_pipe_sync(int device);
#ifdef __cplusplus
}
#endif
#endif
+229 -3
View File
@@ -22,6 +22,7 @@
#include <stdio.h>
#include <stddef.h>
#include <string.h>
#include <windows.h>
#include "backend_cuda.h"
@@ -52,6 +53,35 @@ typedef void (*fn_tensor_free)(ColiCudaTensor *tensor);
typedef size_t (*fn_tensor_bytes)(const ColiCudaTensor *tensor);
typedef int (*fn_tensor_device)(const ColiCudaTensor *tensor);
/* --- #111 GPU resident pipeline additions (matched to backend_cuda.h) --- */
/* --- #111 GPU resident pipeline additions (matched to backend_cuda.h) --- */
typedef int (*fn_attention_absorb_batch)(ColiCudaTensor *kv_b,float *ctx,const float *q, const float *latent,const float *rope,int S, int H,int Q,int R,int V,int K,int T, float attention_scale);
typedef int (*fn_attention_absorb_batch_dev)(ColiCudaTensor *kv_b_shard,float *ctx_dev, const float *q_dev,const float *latent_dev,const float *rope_dev, int S,int H,int Q,int R,int V,int K,int T,float scale);
typedef int (*fn_attention_absorb_kvdev)(ColiCudaTensor *kv_b,float *ctx,const float *q, const float *latent_dev,const float *rope_dev,int H,int Q,int R,int V,int K,int T, float scale);
typedef int (*fn_attention_project_batch)(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, float *out,const float *q,const float *latent, const float *rope,int S,int H,int Q,int R, int V,int K,int T,float attention_scale);
typedef int (*fn_attention_project_batch_dev)(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, float *out,const float *q_dev,const float *latent_dev,const float *rope_dev, int S,int H,int Q,int R,int V,int K,int T,float scale);
typedef int (*fn_attention_project_batch_dev_out)(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, float *out_dev,const float *q_dev,const float *latent_dev,const float *rope_dev, int S,int H,int Q,int R,int V,int K,int T,float scale);
typedef int (*fn_pipe_add)(int device,float *x_dev,const float *t_dev,size_t n);
typedef void * (*fn_pipe_alloc)(int device,size_t bytes);
typedef int (*fn_pipe_copy2d)(int device,float *dst,int dpitch,const float *src, int spitch,int width,int height);
typedef int (*fn_pipe_download)(int device,const void *src,void *dst,size_t bytes);
typedef void (*fn_pipe_free)(int device,void *p);
typedef int (*fn_pipe_gemm)(ColiCudaTensor *t,float *y_dev,const float *x_dev,int S);
typedef int (*fn_pipe_peer_copy)(int dst_dev,float *dst,int src_dev, const float *src,size_t bytes);
typedef int (*fn_pipe_rmsnorm)(int device,float *y_dev,const float *x_dev, const float *w_dev,int S,int D,float eps);
typedef int (*fn_pipe_rmsnorm_s)(int device,float *y_dev,const float *x_dev, const float *w_dev,int S,int D,float eps, int xstride,int ystride);
typedef int (*fn_pipe_rope)(int device,float *v_dev,const int *pos_dev,int rows, int stride,int offset,int R,int heads,float theta);
typedef int (*fn_pipe_rope_base)(int device,float *v_dev,int pos_base,int rows, int stride,int offset,int R,int heads,float theta);
typedef int (*fn_pipe_rows_add)(int device,float *x_dev,const float *partial_dev, const int *rows_dev,int nrows,int D);
typedef float * (*fn_pipe_scratch)(int device,int slot,size_t bytes);
typedef int (*fn_pipe_silu_mul)(int device,float *gate_dev,const float *up_dev,size_t n);
typedef int (*fn_pipe_sync)(int device);
typedef int (*fn_pipe_upload)(int device,void *dst,const void *src,size_t bytes);
typedef int (*fn_shared_mlp_w4a16)(ColiCudaTensor *gate, ColiCudaTensor *up, ColiCudaTensor *down, float *y, const float *x, int S);
typedef int (*fn_tensor_update)(ColiCudaTensor *tensor, const void *weights, const float *scales);
/* Resolved pointers, plus a flag so we attempt the load at most once. */
static struct {
int loaded; /* 1 = load attempted (success or fail), 0 = not yet */
@@ -72,6 +102,31 @@ static struct {
fn_tensor_free tensor_free;
fn_tensor_bytes tensor_bytes;
fn_tensor_device tensor_device;
fn_attention_absorb_batch attention_absorb_batch;
fn_attention_absorb_batch_dev attention_absorb_batch_dev;
fn_attention_absorb_kvdev attention_absorb_kvdev;
fn_attention_project_batch attention_project_batch;
fn_attention_project_batch_dev attention_project_batch_dev;
fn_attention_project_batch_dev_out attention_project_batch_dev_out;
fn_pipe_add pipe_add;
fn_pipe_alloc pipe_alloc;
fn_pipe_copy2d pipe_copy2d;
fn_pipe_download pipe_download;
fn_pipe_free pipe_free;
fn_pipe_gemm pipe_gemm;
fn_pipe_peer_copy pipe_peer_copy;
fn_pipe_rmsnorm pipe_rmsnorm;
fn_pipe_rmsnorm_s pipe_rmsnorm_s;
fn_pipe_rope pipe_rope;
fn_pipe_rope_base pipe_rope_base;
fn_pipe_rows_add pipe_rows_add;
fn_pipe_scratch pipe_scratch;
fn_pipe_silu_mul pipe_silu_mul;
fn_pipe_sync pipe_sync;
fn_pipe_upload pipe_upload;
fn_shared_mlp_w4a16 shared_mlp_w4a16;
fn_tensor_update tensor_update;
} g_cuda;
/* Resolve the DLL and all 11 symbols. Returns 1 on success, 0 otherwise.
@@ -82,9 +137,30 @@ static int coli_cuda_load(void){
if(g_cuda.loaded) return g_cuda.available;
g_cuda.loaded = 1;
/* Search the model directory first (so a DLL shipped next to the model
* wins), then the engine directory, then the default DLL search path. */
g_cuda.dll = LoadLibraryA("coli_cuda.dll");
/* Load coli_cuda.dll from the engine's OWN directory, by absolute path —
* never a bare name. LoadLibraryA("coli_cuda.dll") searches the current
* working directory (and, without SafeDllSearchMode, other writable dirs):
* an attacker who drops a coli_cuda.dll where the user launches glm.exe (or
* inside a downloaded model directory the user cd's into) would get their
* DllMain run at load — classic DLL hijacking -> arbitrary code execution.
* Resolving the path next to glm.exe and loading THAT specific file with
* LOAD_WITH_ALTERED_SEARCH_PATH anchors both the DLL and its dependency
* search to the trusted install directory instead of the CWD. */
char dllpath[MAX_PATH];
DWORD mn = GetModuleFileNameA(NULL, dllpath, (DWORD)sizeof(dllpath));
if(mn > 0 && mn < sizeof(dllpath)){
char *slash = strrchr(dllpath, '\\');
if(slash && (size_t)(slash + 1 - dllpath) + sizeof("coli_cuda.dll") <= sizeof(dllpath)){
strcpy(slash + 1, "coli_cuda.dll");
g_cuda.dll = LoadLibraryExA(dllpath, NULL, LOAD_WITH_ALTERED_SEARCH_PATH);
}
}
if(!g_cuda.dll){
/* fallback (GetModuleFileNameA praticamente non fallisce): cerca solo
* nella dir dell'applicazione e in System32, MAI la CWD. */
g_cuda.dll = LoadLibraryExA("coli_cuda.dll", NULL,
LOAD_LIBRARY_SEARCH_APPLICATION_DIR | LOAD_LIBRARY_SEARCH_SYSTEM32);
}
if(!g_cuda.dll){
fprintf(stderr, "[CUDA] coli_cuda.dll not found; GPU tier disabled "
"(CPU path remains active).\n");
@@ -119,6 +195,31 @@ static int coli_cuda_load(void){
RESOLVE(tensor_free, fn_tensor_free)
RESOLVE(tensor_bytes, fn_tensor_bytes)
RESOLVE(tensor_device, fn_tensor_device)
RESOLVE(attention_absorb_batch, fn_attention_absorb_batch)
RESOLVE(attention_absorb_batch_dev, fn_attention_absorb_batch_dev)
RESOLVE(attention_absorb_kvdev, fn_attention_absorb_kvdev)
RESOLVE(attention_project_batch, fn_attention_project_batch)
RESOLVE(attention_project_batch_dev, fn_attention_project_batch_dev)
RESOLVE(attention_project_batch_dev_out, fn_attention_project_batch_dev_out)
RESOLVE(pipe_add, fn_pipe_add)
RESOLVE(pipe_alloc, fn_pipe_alloc)
RESOLVE(pipe_copy2d, fn_pipe_copy2d)
RESOLVE(pipe_download, fn_pipe_download)
RESOLVE(pipe_free, fn_pipe_free)
RESOLVE(pipe_gemm, fn_pipe_gemm)
RESOLVE(pipe_peer_copy, fn_pipe_peer_copy)
RESOLVE(pipe_rmsnorm, fn_pipe_rmsnorm)
RESOLVE(pipe_rmsnorm_s, fn_pipe_rmsnorm_s)
RESOLVE(pipe_rope, fn_pipe_rope)
RESOLVE(pipe_rope_base, fn_pipe_rope_base)
RESOLVE(pipe_rows_add, fn_pipe_rows_add)
RESOLVE(pipe_scratch, fn_pipe_scratch)
RESOLVE(pipe_silu_mul, fn_pipe_silu_mul)
RESOLVE(pipe_sync, fn_pipe_sync)
RESOLVE(pipe_upload, fn_pipe_upload)
RESOLVE(shared_mlp_w4a16, fn_shared_mlp_w4a16)
RESOLVE(tensor_update, fn_tensor_update)
#undef RESOLVE
g_cuda.available = 1;
@@ -216,4 +317,129 @@ int coli_cuda_tensor_device(const ColiCudaTensor *tensor){
return g_cuda.tensor_device(tensor);
}
/* ---- #111 pipeline wrappers ---- */
/* ---- #111 pipeline wrappers (see header for semantics) ---- */
int coli_cuda_attention_absorb_batch(ColiCudaTensor *kv_b,float *ctx,const float *q, const float *latent,const float *rope,int S, int H,int Q,int R,int V,int K,int T, float attention_scale){
if(!g_cuda.available){ return 0; }
return g_cuda.attention_absorb_batch(kv_b, ctx, q, latent, rope, S, H, Q, R, V, K, T, attention_scale);
}
int coli_cuda_attention_absorb_batch_dev(ColiCudaTensor *kv_b_shard,float *ctx_dev, const float *q_dev,const float *latent_dev,const float *rope_dev, int S,int H,int Q,int R,int V,int K,int T,float scale){
if(!g_cuda.available){ return 0; }
return g_cuda.attention_absorb_batch_dev(kv_b_shard, ctx_dev, q_dev, latent_dev, rope_dev, S, H, Q, R, V, K, T, scale);
}
int coli_cuda_attention_absorb_kvdev(ColiCudaTensor *kv_b,float *ctx,const float *q, const float *latent_dev,const float *rope_dev,int H,int Q,int R,int V,int K,int T, float scale){
if(!g_cuda.available){ return 0; }
return g_cuda.attention_absorb_kvdev(kv_b, ctx, q, latent_dev, rope_dev, H, Q, R, V, K, T, scale);
}
int coli_cuda_attention_project_batch(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, float *out,const float *q,const float *latent, const float *rope,int S,int H,int Q,int R, int V,int K,int T,float attention_scale){
if(!g_cuda.available){ return 0; }
return g_cuda.attention_project_batch(kv_b, o_proj, out, q, latent, rope, S, H, Q, R, V, K, T, attention_scale);
}
int coli_cuda_attention_project_batch_dev(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, float *out,const float *q_dev,const float *latent_dev,const float *rope_dev, int S,int H,int Q,int R,int V,int K,int T,float scale){
if(!g_cuda.available){ return 0; }
return g_cuda.attention_project_batch_dev(kv_b, o_proj, out, q_dev, latent_dev, rope_dev, S, H, Q, R, V, K, T, scale);
}
int coli_cuda_attention_project_batch_dev_out(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, float *out_dev,const float *q_dev,const float *latent_dev,const float *rope_dev, int S,int H,int Q,int R,int V,int K,int T,float scale){
if(!g_cuda.available){ return 0; }
return g_cuda.attention_project_batch_dev_out(kv_b, o_proj, out_dev, q_dev, latent_dev, rope_dev, S, H, Q, R, V, K, T, scale);
}
int coli_cuda_pipe_add(int device,float *x_dev,const float *t_dev,size_t n){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_add(device, x_dev, t_dev, n);
}
void * coli_cuda_pipe_alloc(int device,size_t bytes){
if(!g_cuda.available){ return NULL; }
return g_cuda.pipe_alloc(device, bytes);
}
int coli_cuda_pipe_copy2d(int device,float *dst,int dpitch,const float *src, int spitch,int width,int height){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_copy2d(device, dst, dpitch, src, spitch, width, height);
}
int coli_cuda_pipe_download(int device,const void *src,void *dst,size_t bytes){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_download(device, src, dst, bytes);
}
void coli_cuda_pipe_free(int device,void *p){
if(!g_cuda.available){ return; }
g_cuda.pipe_free(device, p);
}
int coli_cuda_pipe_gemm(ColiCudaTensor *t,float *y_dev,const float *x_dev,int S){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_gemm(t, y_dev, x_dev, S);
}
int coli_cuda_pipe_peer_copy(int dst_dev,float *dst,int src_dev, const float *src,size_t bytes){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_peer_copy(dst_dev, dst, src_dev, src, bytes);
}
int coli_cuda_pipe_rmsnorm(int device,float *y_dev,const float *x_dev, const float *w_dev,int S,int D,float eps){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_rmsnorm(device, y_dev, x_dev, w_dev, S, D, eps);
}
int coli_cuda_pipe_rmsnorm_s(int device,float *y_dev,const float *x_dev, const float *w_dev,int S,int D,float eps, int xstride,int ystride){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_rmsnorm_s(device, y_dev, x_dev, w_dev, S, D, eps, xstride, ystride);
}
int coli_cuda_pipe_rope(int device,float *v_dev,const int *pos_dev,int rows, int stride,int offset,int R,int heads,float theta){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_rope(device, v_dev, pos_dev, rows, stride, offset, R, heads, theta);
}
int coli_cuda_pipe_rope_base(int device,float *v_dev,int pos_base,int rows, int stride,int offset,int R,int heads,float theta){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_rope_base(device, v_dev, pos_base, rows, stride, offset, R, heads, theta);
}
int coli_cuda_pipe_rows_add(int device,float *x_dev,const float *partial_dev, const int *rows_dev,int nrows,int D){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_rows_add(device, x_dev, partial_dev, rows_dev, nrows, D);
}
float * coli_cuda_pipe_scratch(int device,int slot,size_t bytes){
if(!g_cuda.available){ return NULL; }
return g_cuda.pipe_scratch(device, slot, bytes);
}
int coli_cuda_pipe_silu_mul(int device,float *gate_dev,const float *up_dev,size_t n){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_silu_mul(device, gate_dev, up_dev, n);
}
int coli_cuda_pipe_sync(int device){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_sync(device);
}
int coli_cuda_pipe_upload(int device,void *dst,const void *src,size_t bytes){
if(!g_cuda.available){ return 0; }
return g_cuda.pipe_upload(device, dst, src, bytes);
}
int coli_cuda_shared_mlp_w4a16(ColiCudaTensor *gate, ColiCudaTensor *up, ColiCudaTensor *down, float *y, const float *x, int S){
if(!g_cuda.available){ return 0; }
return g_cuda.shared_mlp_w4a16(gate, up, down, y, x, S);
}
int coli_cuda_tensor_update(ColiCudaTensor *tensor, const void *weights, const float *scales){
if(!g_cuda.available){ return 0; }
return g_cuda.tensor_update(tensor, weights, scales);
}
#endif /* _WIN32 */
+5
View File
@@ -0,0 +1,5 @@
@echo off
call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat"
cd /d C:\Users\Mark\Desktop\Projects\colibri\c
"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\nvcc" -O3 -std=c++17 -arch=sm_120 -Xcompiler=-W3 -shared -DCOLI_CUDA_BUILDING_DLL -L"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\lib/x64" -lcudart backend_cuda.cu -o coli_cuda.dll
echo EXITCODE=%ERRORLEVEL%
+102 -10
View File
@@ -161,6 +161,34 @@ def resource_request(a, env):
def env_for(a):
e = dict(os.environ, SNAP=a.model)
if sys.platform == "win32":
# COLI_NO_OMP_TUNE spegne SOLO il blocco OMP (stesso perimetro del
# self-exec di glm.c; presence-based come nel motore: impostarla a
# qualsiasi valore, anche 0, disattiva). I default I/O piu' sotto
# restano attivi: kill-switch dedicati = le var stesse (DIRECT=0 ecc.)
if not e.get("COLI_NO_OMP_TUNE"):
# parita' col tuning OMP self-exec di glm.c (solo Linux/FreeBSD, e
# comunque saltato sotto CUDA/Metal): libgomp legge queste variabili
# prima di main, quindi su Windows vanno nell'ambiente del figlio.
# niente OMP_PROC_BIND/OMP_PLACES: la libgomp di MinGW non supporta
# l'affinity su Windows ("Affinity not supported on this configuration")
from resource_plan import physical_cpu_count
for k, v in (("OMP_WAIT_POLICY", "active"),
("GOMP_SPINCOUNT", "200000"),
("OMP_DYNAMIC", "FALSE"),
("OMP_NUM_THREADS", str(physical_cpu_count()))):
e.setdefault(k, v)
# Default Windows misurati sul box di riferimento (docs/tuning-9950x3d-5090.md),
# tutti lossless e tutti setdefault (un override esplicito vince sempre):
# - DIRECT=1: 10.7 GB/s O_DIRECT vs 9.0 buffered (iobench, anche a cache
# calda); nel motore 0.48 -> 1.02 tok/s. Upstream #162: 1.47x.
# - PIPE=1: overlap load/matmul, +8% sopra DIRECT (byte-identico, riordina
# solo l'I/O). PIPE_WORKERS resta al default 8 (sweep 4/8/16 piatto).
# - PILOT_REAL=1: prefetch cross-layer con load veri (unico prefetch
# funzionante su Windows: fadvise e' no-op), +11%, hit rate +19 punti.
e.setdefault("DIRECT", "1")
e.setdefault("PIPE", "1")
e.setdefault("PILOT_REAL", "1")
e["COLI_POLICY"]=a.policy
if a.ram: e["RAM_GB"]=str(a.ram)
if a.ngen: e["NGEN"]=str(a.ngen)
@@ -316,11 +344,50 @@ class Spinner:
if self.th: self.th.join(timeout=0.4)
if TTY: sys.stdout.write("\r\033[K"); sys.stdout.flush()
def engine_diag(p, errlog=None):
"""Lo stdout del motore ha chiuso: il processo e' morto. Dire PERCHE'.
Un SIGKILL dell'OOM-killer del kernel e' altrimenti invisibile — niente errore,
niente exit code, il motore muore muto e sembra un bug nostro (issue #305).
Quasi sempre e' memoria: il picco reale ha sforato la RAM della macchina."""
try: rc=p.wait(timeout=5)
except Exception: rc=p.poll()
if rc is None: why="its output closed but the process is still alive"
elif rc<0: # POSIX: morte per segnale
try: why=f"killed by {signal.Signals(-rc).name}"
except Exception: why=f"killed by signal {-rc}"
elif rc>0: why=f"exit code {rc}"
else: why="exited cleanly"
print(f"\n {C.yel}[engine terminated: {why}]{C.r}")
if rc is not None and rc<0 and -rc==getattr(signal,"SIGKILL",-1):
print(f" {C.yel}nothing in the engine sends SIGKILL to itself: this is the kernel's\n"
f" OOM-killer. The peak RSS exceeded the machine's free memory.\n"
f" Lower --ram, lower PIN_GB, or shorten the context.{C.r}")
if errlog is not None:
try:
errlog.flush()
for _ in range(20): # il drain thread scrive su errlog quando il child chiude
tail=open(errlog.name).read().strip()
if tail: break
time.sleep(0.05)
if tail: print(f" {C.dgray}" + "\n ".join(tail.splitlines()[-6:]) + f"{C.r}")
except Exception: pass
def stream_turn(p, sentinel, on_bytes):
"""legge fino alla sentinella; on_bytes riceve i chunk della risposta. Poi legge la riga STAT."""
pend=b""
"""legge fino alla sentinella; on_bytes riceve i chunk della risposta. Poi legge la riga STAT.
Il PRIMO Ctrl-C durante lo stream non chiude la sessione: il motore (handler SIGINT)
chiude il turno per la via del tetto NGEN e noi dreniamo fino alla sentinella.
Un SECONDO Ctrl-C esce davvero."""
pend=b""; interrupted=False
while True:
b=p.stdout.read(1)
try:
b=p.stdout.read(1)
except KeyboardInterrupt:
if interrupted or p.poll() is not None: raise
interrupted=True
try: p.send_signal(signal.SIGINT) # non-TTY: il motore potrebbe non aver visto il Ctrl-C
except Exception: pass
print(f"\n {C.yel}⏹ stopping… (Ctrl-C again to quit){C.r}", flush=True)
continue
if b==b"": return None
pend+=b
if pend.endswith(sentinel):
@@ -328,7 +395,9 @@ def stream_turn(p, sentinel, on_bytes):
if rest: on_bytes(rest)
line=p.stdout.readline().decode("utf-8","replace").strip() # STAT tok tps hit rss
m=re.match(r"STAT (\S+) (\S+) (\S+) (\S+)", line)
return {"tok":int(m.group(1)),"tps":float(m.group(2)),"hit":float(m.group(3)),"rss":float(m.group(4))} if m else {}
st={"tok":int(m.group(1)),"tps":float(m.group(2)),"hit":float(m.group(3)),"rss":float(m.group(4))} if m else {}
if interrupted: st["interrupted"]=True
return st
if len(pend)>len(sentinel):
out=pend[:-len(sentinel)]; pend=pend[-len(sentinel):]
on_bytes(out)
@@ -437,7 +506,11 @@ def cmd_chat(a):
if st is None:
try: errlog.write(p.stderr.read().decode("utf-8","replace"))
except (OSError, ValueError): pass
errlog.seek(0); print(errlog.read()[-1500:]); sys.exit("the engine exited while loading")
errlog.flush(); errlog.seek(0); print(errlog.read()[-1500:])
engine_diag(p) # perche' e' morto (OOM-kill compreso); errlog e' gia' stampato sopra
sys.exit("the engine exited while loading")
p.stdout.readline() # TIERS line (web-dashboard protocol): emitted once right after STAT;
# left unread it leaks into the first answer's text
# READY received. Drain the child's stderr into errlog without blocking:
# the engine is still alive (blocked on stdin), so a plain read() would
# hang forever waiting for EOF. A short bounded drain grabs the ~400 bytes
@@ -464,7 +537,7 @@ def cmd_chat(a):
for chunk in textwrap.wrap(l, term_w()-4) or [l]:
print(f" {C.dgray}{chunk}{C.r}")
except Exception: pass
print(f" {C.dim}type and press Enter · :more continues · :reset clears memory · :q exits{C.r}\n")
print(f" {C.dim}type and press Enter · Ctrl-C stops the answer · :more continues · :reset clears memory · :q exits{C.r}\n")
w=term_w()-4
def user_box(msg):
"""ri-disegna il messaggio dentro una box che si ADATTA su piu' righe:
@@ -525,14 +598,17 @@ def cmd_chat(a):
st=stream_turn(p, END, echo)
if not raw: md.close()
sp2.stop()
if st is None: print(f"\n {C.yel}[engine terminated]{C.r}"); break
if st is None: engine_diag(p, errlog); break
el=time.time()-t0
if st.get("tok"):
print(f"\r {C.dgray}└─ {st['tok']} tok · {st['tps']:.2f} tok/s · hit {st['hit']:.0f}% · RSS {st['rss']:.1f} GB · {el:.0f}s{C.r}")
if st["tok"]>=a.ngen:
if st.get("interrupted"):
print(f" {C.yel}⏹ interrupted; type :more to continue the response{C.r}")
elif st["tok"]>=a.ngen:
print(f" {C.yel}…stopped at --ngen ({a.ngen}); type :more to continue the response{C.r}")
print()
else:
if st.get("interrupted"): print(f" {C.yel}⏹ interrupted{C.r}")
print()
except KeyboardInterrupt:
print(f"\n {C.dim}interrupted{C.r}")
@@ -585,7 +661,17 @@ def cmd_bench(a):
print(f" {C.dim}downloading missing datasets: {', '.join(missing)}{C.r}")
subprocess.call([py, os.path.join(TOOLS,"fetch_benchmarks.py"),
"--out", a.data, "--tasks", ",".join(missing), "--limit", str(max(a.limit,200))])
cmd=[py, os.path.join(TOOLS,"eval_glm.py"), "--snap",a.model,
# il fetch riprova da solo (#304), ma se l'hub resta giu' il bench gira sui task
# disponibili invece di passare a eval file inesistenti.
# EN: the fetch retries on its own (#304), but if the hub stays down the bench
# runs on the available tasks instead of handing eval nonexistent files.
still=[t for t in tasks.split(",") if not os.path.exists(os.path.join(a.data,f"{t}.jsonl"))]
if still:
tasks=",".join(t for t in tasks.split(",") if t not in still)
print(f" {C.yel}skipping (download failed, rerun later): {', '.join(still)}{C.r}")
if not tasks:
print(f" {C.yel}no datasets available — nothing to bench{C.r}"); sys.exit(1)
cmd=[py, os.path.join(TOOLS,"eval_glm.py"), "--glm", GLM, "--snap",a.model,
"--tasks", tasks, "--limit", str(a.limit), "--data", a.data]
if a.ram: cmd+=["--ram",str(a.ram)]
e=env_for(a)
@@ -653,7 +739,13 @@ def main():
pw.add_argument(arg,**kw)
pw.add_argument("--no-browser",action="store_true",help="don't auto-open the browser")
pb=sub.add_parser("bench", parents=[common]); pb.add_argument("tasks", nargs="*")
pb.add_argument("--limit",type=int,default=40); pb.add_argument("--data",default=os.path.join(HERE,"bench"))
pb.add_argument("--limit",type=int,default=40)
if sys.platform == "win32":
_cache_root = os.environ.get("LOCALAPPDATA", os.path.expanduser("~\\AppData\\Local"))
else:
_cache_root = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
_bench_cache = os.path.join(_cache_root, "colibri", "bench")
pb.add_argument("--data",default=_bench_cache)
pc=sub.add_parser("convert", parents=[common]); pc.add_argument("--repo",default="zai-org/GLM-5.2-FP8")
pc.add_argument("--ebits",type=int,default=4); pc.add_argument("--io-bits",type=int,default=8); pc.add_argument("--xbits",type=int,default=0)
pc.add_argument("--no-mtp",action="store_true",help="skip the MTP head (no speculative drafts)")
+35 -2
View File
@@ -95,7 +95,18 @@ static inline int compat_open_direct(const char *path){
* prevents 0x0A bytes from being silently translated to \r\n. */
#define COMPAT_O_RDONLY (O_RDONLY | O_BINARY)
/* --- posix_fadvise: no-op (advisory only; safe to ignore) --- */
/* --- posix_fadvise: Windows has no direct equivalent. Semantics:
* WILLNEED -> warm the OS page cache so a later synchronous pread finds the
* pages resident. Implemented as an overlapped background ReadFile
* into a throwaway scratch buffer (fire-and-forget readahead). Called
* from the dedicated PILOT I/O thread / next-block readahead in moe(),
* NEVER inline on the hot path (the existing comment at glm.c:2847
* measures inline fadvise submit at ~0.5ms x 169k calls = +92s/48tok).
* Each call owns its OVERLAPPED + scratch buffer -> thread-safe.
* DONTNEED -> no-op: Windows' standby-list trimming self-regulates under pressure,
* and on a low-RAM host keeping the pages is what we want for reuse.
* Matches macOS (compat.h:16-19) which no-ops DONTNEED for the same
* reason. The engine only ever uses DONTNEED as an advisory. */
#ifndef POSIX_FADV_NORMAL
#define POSIX_FADV_NORMAL 0
#define POSIX_FADV_RANDOM 1
@@ -104,7 +115,29 @@ static inline int compat_open_direct(const char *path){
#define POSIX_FADV_DONTNEED 4
#define POSIX_FADV_NOREUSE 5
#endif
#define posix_fadvise(fd,off,len,advice) do{(void)(fd);(void)(off);(void)(len);(void)(advice);}while(0)
static inline int compat_fadvise(int fd, off_t off, off_t len, int advice){
if(advice!=POSIX_FADV_WILLNEED || len<=0) return 0;
intptr_t osfh=_get_osfhandle(fd);
if(osfh==-1 || osfh==-2) return 0;
HANDLE h=(HANDLE)osfh;
/* Cap the readahead window: reading a whole 19MB expert per hint is fine on the
* PILOT thread, but a pathological huge len would spike transient memory. */
size_t rdlen = (len>(off_t)(64*1024*1024)) ? (size_t)(64*1024*1024) : (size_t)len;
char *buf=(char*)_aligned_malloc(rdlen, 4096);
if(!buf) return -1;
OVERLAPPED ov={0};
ov.Offset = (DWORD)( (off_t)off & 0xFFFFFFFFULL);
ov.OffsetHigh = (DWORD)(((off_t)off >> 32) & 0xFFFFFFFFULL);
/* Issue an overlapped read. With a non-OVERLAPPED-opened handle ReadFile still
* accepts lpOverlapped (it carries the 64-bit offset) and blocks until the read
* completes — but crucially it populates the standby page cache for this region,
* so the later synchronous pread on the same offsets faults from RAM not disk. */
DWORD got=0;
ReadFile(h, buf, (DWORD)rdlen, &got, &ov);
_aligned_free(buf);
return 0;
}
#define posix_fadvise compat_fadvise
/* --- pread -> ReadFile + OVERLAPPED su raw OS handle ---
* Thread-safe (no shared seek position). Gestisce offset >4 GB e chunking
+231
View File
@@ -0,0 +1,231 @@
"""Download GLM-5.2-FP8 from ModelScope (fast, no HF throttling) with
HuggingFace fallback. Parallel shard download, clean progress display.
Usage: python download_fp8.py
python download_fp8.py --parallel 4
python download_fp8.py --source hf (force HuggingFace)
"""
import os, sys, time, threading, argparse, subprocess
REPO_MS = "ZhipuAI/GLM-5.2-FP8" # ModelScope
REPO_HF = "zai-org/GLM-5.2-FP8" # HuggingFace
DEST = r"I:\glm52_fp8"
# ── ANSI colors ──
class C:
dim="\033[2m"; grn="\033[32m"; yel="\033[33m"; cyn="\033[36m"; b="\033[1m"; r="\033[0m"
def fmt_bytes(n):
if n>=1e9: return f"{n/1e9:.2f} GB"
if n>=1e6: return f"{n/1e6:.1f} MB"
return f"{n/1e3:.0f} KB"
def fmt_time(s):
if s<60: return f"{s:.0f}s"
if s<3600: return f"{s/60:.0f}m"
return f"{s/3600:.1f}h"
def bar(cur, total, width=24):
if total<=0: return "["+" "*width+"]"
pct=min(cur/total,1.0); filled=int(width*pct)
return "["+""*filled+""*(width-filled)+f"] {pct*100:4.0f}%"
def get_shard_list_hf():
from huggingface_hub import HfApi
info=HfApi().repo_info(REPO_HF, files_metadata=True)
shards=sorted(s.rfilename for s in info.siblings if s.rfilename.endswith(".safetensors"))
sizes={s.rfilename:s.size for s in info.siblings if s.rfilename.endswith(".safetensors")}
return shards, sizes
def get_shard_list_ms():
"""Get shard list from ModelScope API."""
import requests
# ModelScope API: list files
r = requests.get(f"https://modelscope.cn/api/v1/models/{REPO_MS}/repo/files?Revision=master&Root=", timeout=30)
data = r.json()["Data"]["Files"]
shards = sorted(f["Path"] for f in data if f["Path"].endswith(".safetensors"))
sizes = {f["Path"]: f.get("Size", 0) for f in data if f["Path"].endswith(".safetensors")}
return shards, sizes
def download_file_ms(fn):
"""Download a single file from ModelScope using their CDN."""
from modelscope.hub.file_download import model_file_download
model_file_download(
model_id=REPO_MS,
file_path=fn,
local_dir=DEST,
revision="master",
)
def download_file_hf(fn):
"""Download a single file from HuggingFace with hf_transfer."""
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
from huggingface_hub import hf_hub_download
hf_hub_download(REPO_HF, fn, local_dir=DEST)
def download_file_curl(fn, base_url, expected_size):
"""Fallback: download with curl to a .part file with resume."""
outpath = os.path.join(DEST, fn)
partpath = outpath + ".part"
if os.path.exists(outpath) and os.path.getsize(outpath) == expected_size:
return True
url = f"{base_url}/{fn}"
cmd = ["curl", "-L", "-C", "-", "--retry", "999", "--retry-delay", "5",
"--connect-timeout", "15", "--speed-time", "30", "--speed-limit", "1000",
"-o", partpath, "-H", "User-Agent: colibri-download/1.0", url]
subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
if os.path.exists(partpath) and os.path.getsize(partpath) >= expected_size:
os.replace(partpath, outpath)
return True
return False
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--parallel", type=int, default=3)
ap.add_argument("--source", choices=["auto", "ms", "hf"], default="auto",
help="auto (try ModelScope first), ms (ModelScope only), hf (HuggingFace only)")
args = ap.parse_args()
os.makedirs(DEST, exist_ok=True)
# Determine source and get shard list
use_ms = False
shards, sizes = [], {}
if args.source in ("auto", "ms"):
try:
print(f"{C.dim}Trying ModelScope...{C.r}", end=" ", flush=True)
shards, sizes = get_shard_list_ms()
use_ms = True
print(f"{C.grn}{C.r} {len(shards)} shards found")
except Exception as e:
print(f"{C.yel}failed ({e}){C.r}")
if args.source == "ms":
print("ModelScope failed and --source ms was set. Exiting."); return
if not shards:
print(f"{C.dim}Using HuggingFace...{C.r}", end=" ", flush=True)
shards, sizes = get_shard_list_hf()
print(f"{C.grn}{C.r} {len(shards)} shards found")
total = len(shards)
total_bytes = sum(sizes.values())
source_name = "ModelScope" if use_ms else "HuggingFace"
# Download metadata files
meta_files = ["config.json", "tokenizer.json", "tokenizer_config.json",
"generation_config.json", "model.safetensors.index.json"]
for fn in meta_files:
out = os.path.join(DEST, fn)
if not os.path.exists(out):
try:
if use_ms: download_file_ms(fn)
else: download_file_hf(fn)
except Exception: pass
# Build work queue
todo = []
done_set = set()
for fn in shards:
outpath = os.path.join(DEST, fn)
if os.path.exists(outpath) and os.path.getsize(outpath) == sizes.get(fn, 0):
done_set.add(fn)
else:
todo.append(fn)
existing = len(done_set)
print(f"\n{C.b}GLM-5.2-FP8 Download ({source_name}){C.r}")
print(f" {C.dim}{total} shards · {total_bytes/1e9:.0f} GB · {existing}/{total} complete{C.r}")
print(f" {C.dim}{len(todo)} to download · {args.parallel} parallel{C.r}")
if todo:
remaining = sum(sizes[fn] for fn in todo)
print(f" {C.dim}Remaining: {remaining/1e9:.0f} GB{C.r}")
print()
if not todo:
print(f"{C.grn}✓ All shards already downloaded!{C.r}\n"); return
lock = threading.Lock()
completed = list(done_set)
t0 = time.time()
qidx = [0]
def worker(wid):
while True:
with lock:
if qidx[0] >= len(todo): return
idx = qidx[0]; qidx[0] += 1
fn = todo[idx]
expected = sizes.get(fn, 0)
shard_num = existing + idx + 1
print(f" {C.cyn}[{shard_num}/{total}]{C.r} {C.dim}{fn}{C.r}")
success = False
for attempt in range(3):
try:
if use_ms:
download_file_ms(fn)
else:
download_file_hf(fn)
outpath = os.path.join(DEST, fn)
# ModelScope downloads to a cache dir, need to check
# if the file exists at our expected path
if not os.path.exists(outpath):
# Try to find it in ModelScope's cache structure
# and copy/symlink it
pass
if os.path.exists(outpath):
actual = os.path.getsize(outpath)
if expected == 0 or actual == expected:
success = True; break
# Size mismatch — might be in cache
# If we got here, file downloaded but not at expected path
# ModelScope puts it in local_dir; check again
if os.path.exists(outpath):
success = True; break
# Retry with curl fallback
if use_ms:
base = f"https://modelscope.cn/api/v1/models/{REPO_MS}/repo?Revision=master&FilePath="
else:
base = f"https://huggingface.co/{REPO_HF}/resolve/main"
if download_file_curl(fn, base, expected):
success = True; break
except Exception as e:
if attempt < 2:
print(f" {C.yel}retry {attempt+1}: {e}{C.r}")
time.sleep(3)
else:
print(f" {C.yel}✗ failed: {e}{C.r}")
with lock:
if success:
completed.append(fn)
elapsed = time.time() - t0
have = sum(sizes.get(f,0) for f in completed)
pct = 100.0 * have / total_bytes
speed = (have - sum(sizes[f] for f in done_set)) / max(elapsed, 1)
eta = (total_bytes - have) / speed if speed > 0 else 0
print(f" {C.grn}{C.r} {fn} {C.dim}{len(completed)}/{total} · "
f"{pct:.1f}% · {fmt_time(elapsed)} · ETA {fmt_time(eta)}{C.r}")
else:
print(f" {C.yel}✗ GIVE UP: {fn}{C.r}")
threads = [threading.Thread(target=worker, args=(i,), daemon=True) for i in range(args.parallel)]
for t in threads: t.start()
for t in threads: t.join()
print()
final = sum(1 for fn in shards
if os.path.exists(os.path.join(DEST, fn))
and os.path.getsize(os.path.join(DEST, fn)) == sizes.get(fn, 0))
if final == total:
print(f"{C.grn}{'='*50}")
print(f" ✓ All {total} shards downloaded!{C.r}\n")
else:
print(f"{C.yel} {final}/{total} complete, {total-final} remaining{C.r}")
print(f" Re-run to resume.\n")
if __name__ == "__main__":
try: main()
except KeyboardInterrupt:
print(f"\n\n{C.yel}Interrupted. Re-run to resume — no data lost.{C.r}\n")
+2103 -196
View File
File diff suppressed because it is too large Load Diff
+13 -3
View File
@@ -26,8 +26,14 @@ typedef struct {
const char *s;
char *arena; /* buffer per le stringhe smontate */
size_t acap, aoff;
int depth; /* annidamento corrente: bound contro lo stack-overflow
* da JSON malevolo tipo [[[[...]]]] (discesa ricorsiva) */
} jparser;
/* tetto di annidamento: gli header safetensors / config sono piatti (profondita'
* ~3). 1024 e' larghissimo per input legittimi e ben sotto il limite di stack. */
#define J_MAX_DEPTH 1024
static char *j_dup(jparser *p, const char *b, int n) {
/* ogni stringa ha la sua allocazione: un'arena con realloc sposterebbe il
* buffer invalidando i puntatori gia' emessi (use-after-free). */
@@ -91,10 +97,11 @@ static jval *j_parse_val(jparser *p) {
char c = *p->s;
if (c == '"') { jval *v = j_new(J_STR); v->str = j_parse_str_raw(p); return v; }
if (c == '{') {
if (++p->depth > J_MAX_DEPTH) { p->depth--; return j_new(J_NULL); }
p->s++; jval *v = j_new(J_OBJ);
int cap = 8; v->keys = malloc(cap * sizeof(char*)); v->kids = malloc(cap * sizeof(jval*));
j_ws(p);
if (*p->s == '}') { p->s++; return v; }
if (*p->s == '}') { p->s++; p->depth--; return v; }
for (;;) {
j_ws(p);
char *key = j_parse_str_raw(p);
@@ -107,13 +114,15 @@ static jval *j_parse_val(jparser *p) {
if (*p->s == '}') { p->s++; break; }
break;
}
p->depth--;
return v;
}
if (c == '[') {
if (++p->depth > J_MAX_DEPTH) { p->depth--; return j_new(J_NULL); }
p->s++; jval *v = j_new(J_ARR);
int cap = 8; v->kids = malloc(cap * sizeof(jval*));
j_ws(p);
if (*p->s == ']') { p->s++; return v; }
if (*p->s == ']') { p->s++; p->depth--; return v; }
for (;;) {
jval *val = j_parse_val(p);
if (v->len == cap) { cap *= 2; v->kids = realloc(v->kids, cap*sizeof(jval*)); }
@@ -123,6 +132,7 @@ static jval *j_parse_val(jparser *p) {
if (*p->s == ']') { p->s++; break; }
break;
}
p->depth--;
return v;
}
if (c == 't') { p->s += 4; jval *v = j_new(J_BOOL); v->boolean = 1; return v; }
@@ -134,7 +144,7 @@ static jval *j_parse_val(jparser *p) {
/* API */
static jval *json_parse(const char *text, char **arena_out) {
jparser p = { text, NULL, 0, 0 };
jparser p = { text, NULL, 0, 0, 0 };
jval *v = j_parse_val(&p);
if (arena_out) *arena_out = p.arena; else free(p.arena);
return v;
+1 -1
View File
@@ -12,7 +12,7 @@
#include <string.h>
#include <math.h>
#include <time.h>
#if defined(__APPLE__) || defined(__linux__)
#if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__)
#include <sys/resource.h>
#endif
#include "st.h"
+13 -4
View File
@@ -410,8 +410,11 @@ def generation_options(body, limit):
top_p = body.get("top_p")
temperature = 0.7 if temperature is None else temperature
top_p = 0.9 if top_p is None else top_p
if isinstance(maximum, bool) or not isinstance(maximum, int) or not 1 <= maximum <= limit:
raise APIError(400, f"`{maximum_param}` must be an integer between 1 and {limit}.", maximum_param)
if isinstance(maximum, bool) or not isinstance(maximum, int) or maximum < 1:
raise APIError(400, f"`{maximum_param}` must be a positive integer.", maximum_param)
if maximum > limit:
maximum = limit # clamp to the server's --max-tokens cap instead of 400 (#260): OpenAI
# clients (opencode/ai-sdk) default to large max_tokens; rejecting breaks them.
if (isinstance(temperature, bool) or not isinstance(temperature, (int, float)) or
not math.isfinite(temperature) or not 0 <= temperature <= 2):
raise APIError(400, "`temperature` must be between 0 and 2.", "temperature")
@@ -731,12 +734,16 @@ class APIHandler(BaseHTTPRequestHandler):
Read-only, no auth (same trust level as /health), traversal-safe."""
if path.startswith("/v1/") or path == "/health":
return False
base = self.WEB_DIST
base = self.WEB_DIST.resolve()
if not base.is_dir():
return False
rel = unquote(path).lstrip("/") or "index.html"
target = (base / rel).resolve()
if not str(target).startswith(str(base)) or not target.is_file():
try:
target.relative_to(base)
except ValueError:
target = None
if target is None or not target.is_file():
if path == "/" or "." not in rel: # SPA fallback
target = base / "index.html"
if not target.is_file():
@@ -1045,6 +1052,8 @@ class APIHandler(BaseHTTPRequestHandler):
prompt = body.get("prompt")
if not isinstance(prompt, str):
raise APIError(400, "Colibri currently requires `prompt` to be a string.", "prompt")
if not prompt:
raise APIError(400, "`prompt` must not be empty.", "prompt")
self.generation(body, prompt, request_id, False)
+29 -2
View File
@@ -167,6 +167,30 @@ def discover_gpus():
def physical_cpu_count():
if sys.platform == "win32":
# os.cpu_count() conta i processori logici (SMT): 2 thread/core saturano
# le unita' AVX-512 e peggiorano il matmul. Contiamo i core fisici veri
# con GetLogicalProcessorInformationEx(RelationProcessorCore).
try:
import ctypes
k32 = ctypes.windll.kernel32
need = ctypes.c_ulong(0)
k32.GetLogicalProcessorInformationEx(0, None, ctypes.byref(need))
buf = (ctypes.c_char * need.value)()
if k32.GetLogicalProcessorInformationEx(0, buf, ctypes.byref(need)):
raw, cores, off = bytes(buf), 0, 0
while off + 8 <= need.value:
relationship = int.from_bytes(raw[off:off + 4], "little")
size = int.from_bytes(raw[off + 4:off + 8], "little")
if size <= 0:
break
if relationship == 0: # RelationProcessorCore
cores += 1
off += size
if cores:
return cores
except (OSError, ValueError, AttributeError):
pass
try:
result = subprocess.run(["lscpu", "-p=core,socket"], text=True,
capture_output=True, check=True, timeout=5)
@@ -289,8 +313,11 @@ def environment_for_plan(plan, env=None, cuda_enabled=True):
result = dict(env or {})
result.setdefault("COLI_POLICY", plan["policy"]["name"])
result.setdefault("OMP_NUM_THREADS", str(plan["cpu"]["physical_cores"]))
result.setdefault("OMP_PROC_BIND", "spread")
result.setdefault("OMP_PLACES", "cores")
if sys.platform != "win32":
# la libgomp di MinGW non supporta l'affinity su Windows
# ("Affinity not supported on this configuration"): non impostarle li'.
result.setdefault("OMP_PROC_BIND", "spread")
result.setdefault("OMP_PLACES", "cores")
if plan["policy"]["name"] == "balanced":
result.setdefault("REPIN", "64")
ram = plan["tiers"]["ram"]
+246
View File
@@ -0,0 +1,246 @@
/* schema_gbnf.h — JSON-Schema -> GBNF compiler for the grammar-forced draft source (#48/#70).
*
* Compiles a practical subset of JSON Schema into the byte-level GBNF subset that
* grammar.h parses, so structured-output requests (OpenAI `response_format:
* {"type":"json_schema"}` and the SCHEMA= env) get grammar-forced drafts without
* hand-writing GBNF.
*
* Safety model: the grammar is a DRAFT SOURCE, never a sampling constraint (see
* grammar.h). A schema compiled too strictly (or a model that deviates) only costs
* draft acceptance — output is unchanged. So the compiler can be strict and compact
* (no whitespace between tokens, exact key order): strictness maximizes forced-span
* length on conforming outputs and cannot corrupt non-conforming ones.
*
* Supported subset:
* type: object + properties (+ required) — properties in declared order;
* if `required` is present it must list every property (OpenAI
* structured-output "strict" semantics); a proper subset -> fail.
* type: string (+ enum of strings, const)
* type: number | integer | boolean | null
* type: array + items (+ minItems 0|1)
* enum / const also allowed at value level with numbers
* nesting to SGB_MAX_DEPTH; annotation keys ($schema, title, description,
* default, examples, additionalProperties) are ignored.
* Anything else (anyOf/oneOf/$ref/pattern/format/minimum/maximum...) -> returns NULL and
* the caller falls back to running without a grammar. Fail-closed, never fatal.
*/
#ifndef SCHEMA_GBNF_H
#define SCHEMA_GBNF_H
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "json.h"
#define SGB_MAX_DEPTH 32
typedef struct {
char *s; size_t len, cap; /* output GBNF text */
int nrule; /* next composite rule id */
int use_str, use_num, use_int; /* shared terminal rules actually referenced */
char err[160];
int fail;
} SgbCtx;
static void sgb_put(SgbCtx *C, const char *t){
size_t n = strlen(t);
if (C->len + n + 1 > C->cap){
size_t nc = C->cap ? C->cap * 2 : 1024;
while (nc < C->len + n + 1) nc *= 2;
char *ns = (char *)realloc(C->s, nc);
if (!ns){ C->fail = 1; return; }
C->s = ns; C->cap = nc;
}
memcpy(C->s + C->len, t, n); C->len += n; C->s[C->len] = 0;
}
static void sgb_fail(SgbCtx *C, const char *what){
if (!C->err[0]) snprintf(C->err, sizeof C->err, "unsupported schema: %s", what);
C->fail = 1;
}
/* emit a JSON string VALUE (with quotes) as a GBNF literal: "..." with the JSON
* text embedded, escaping for the GBNF literal syntax (\" \\ \xHH). */
static void sgb_put_json_string_lit(SgbCtx *C, const char *raw){
sgb_put(C, "\"\\\""); /* GBNF literal opening: "\" */
char b[8];
for (const unsigned char *p = (const unsigned char *)raw; *p; p++){
unsigned char c = *p;
if (c == '"') sgb_put(C, "\\\\\\\""); /* JSON \" inside GBNF literal */
else if (c == '\\') sgb_put(C, "\\\\\\\\");
else if (c < 0x20){ snprintf(b, sizeof b, "\\x%02x", c); sgb_put(C, b); } /* raw ctl byte (invalid JSON anyway) */
else if (c == 0x7f) sgb_put(C, "\\x7f");
else { b[0] = (char)c; b[1] = 0; sgb_put(C, b); }
}
sgb_put(C, "\\\"\"");
}
/* emit a number the way a model would print it: shortest round-trip via %g */
static void sgb_put_number_lit(SgbCtx *C, double d){
char b[64]; snprintf(b, sizeof b, "\"%.17g\"", d);
/* trim %.17g noise for integers */
if (d == (double)(long long)d && d < 1e15 && d > -1e15)
snprintf(b, sizeof b, "\"%lld\"", (long long)d);
sgb_put(C, b);
}
static int sgb_is_annotation(const char *k){
return !strcmp(k,"$schema") || !strcmp(k,"title") || !strcmp(k,"description")
|| !strcmp(k,"default") || !strcmp(k,"examples")
|| !strcmp(k,"additionalProperties");
}
/* forward */
static void sgb_value(SgbCtx *C, jval *sc, int depth);
static void sgb_enum(SgbCtx *C, jval *e){
if (!e || e->t != J_ARR || e->len < 1){ sgb_fail(C, "empty enum"); return; }
sgb_put(C, "( ");
for (int i = 0; i < e->len && !C->fail; i++){
if (i) sgb_put(C, " | ");
jval *v = e->kids[i];
if (v->t == J_STR) sgb_put_json_string_lit(C, v->str);
else if (v->t == J_NUM) sgb_put_number_lit(C, v->num);
else if (v->t == J_BOOL) sgb_put(C, v->boolean ? "\"true\"" : "\"false\"");
else if (v->t == J_NULL) sgb_put(C, "\"null\"");
else sgb_fail(C, "enum member type");
}
sgb_put(C, " )");
}
static void sgb_object(SgbCtx *C, jval *sc, int depth){
jval *props = json_get(sc, "properties");
jval *req = json_get(sc, "required");
if (!props || props->t != J_OBJ){ sgb_fail(C, "object without properties"); return; }
if (props->len == 0){ sgb_put(C, "\"{\" jws \"}\""); return; }
if (req){
if (req->t != J_ARR){ sgb_fail(C, "required not an array"); return; }
/* strict semantics: every property must be required (OpenAI structured
* outputs contract). A proper subset would need optional-group emission
* with ambiguous separators — out of v1 scope. */
if (req->len != props->len){ sgb_fail(C, "required must list every property (strict)"); return; }
for (int i = 0; i < props->len; i++){
int found = 0;
for (int j = 0; j < req->len; j++)
if (req->kids[j]->t == J_STR && !strcmp(req->kids[j]->str, props->keys[i])) found = 1;
if (!found){ sgb_fail(C, "property not in required (strict)"); return; }
}
}
/* jws at every separator: whitespace tolerance is strictly acceptance-positive
* for a DRAFT-source grammar — a compact-only grammar dies (desyncs) at the
* first stray space and forfeits every span after it, while jws points merely
* aren't forced themselves (two legal bytes) and the multi-byte spans around
* them keep drafting. Measured on GLM-5.2 current main: the sloppy-JSON
* continuation costs a compact grammar most of its spans. */
sgb_put(C, "\"{\" jws ");
for (int i = 0; i < props->len && !C->fail; i++){
if (i) sgb_put(C, " \",\" jws ");
sgb_put_json_string_lit(C, props->keys[i]);
sgb_put(C, " jws \":\" jws ");
sgb_value(C, props->kids[i], depth + 1);
sgb_put(C, " jws");
}
sgb_put(C, " \"}\"");
}
static void sgb_array(SgbCtx *C, jval *sc, int depth){
jval *items = json_get(sc, "items");
jval *mi = json_get(sc, "minItems");
int min1 = mi && mi->t == J_NUM && mi->num >= 1;
if (mi && mi->t == J_NUM && mi->num > 1){ sgb_fail(C, "minItems > 1"); return; }
if (!items){ sgb_fail(C, "array without items"); return; }
if (min1){
sgb_put(C, "\"[\" jws ");
sgb_value(C, items, depth + 1);
sgb_put(C, " jws ( \",\" jws ");
sgb_value(C, items, depth + 1);
sgb_put(C, " jws )* \"]\"");
} else {
sgb_put(C, "\"[\" jws ( ");
sgb_value(C, items, depth + 1);
sgb_put(C, " jws ( \",\" jws ");
sgb_value(C, items, depth + 1);
sgb_put(C, " jws )* )? \"]\"");
}
}
static void sgb_value(SgbCtx *C, jval *sc, int depth){
if (C->fail) return;
if (depth > SGB_MAX_DEPTH){ sgb_fail(C, "nesting too deep"); return; }
if (!sc || sc->t != J_OBJ){ sgb_fail(C, "schema node not an object"); return; }
/* reject unknown constraint keywords (fail-closed) */
for (int i = 0; i < sc->len; i++){
const char *k = sc->keys[i];
if (strcmp(k,"type") && strcmp(k,"properties") && strcmp(k,"required")
&& strcmp(k,"items") && strcmp(k,"enum") && strcmp(k,"const")
&& strcmp(k,"minItems") && !sgb_is_annotation(k)){
sgb_fail(C, k); return;
}
}
jval *cst = json_get(sc, "const");
if (cst){
if (cst->t == J_STR) sgb_put_json_string_lit(C, cst->str);
else if (cst->t == J_NUM) sgb_put_number_lit(C, cst->num);
else if (cst->t == J_BOOL) sgb_put(C, cst->boolean ? "\"true\"" : "\"false\"");
else if (cst->t == J_NULL) sgb_put(C, "\"null\"");
else sgb_fail(C, "const type");
return;
}
jval *en = json_get(sc, "enum");
if (en){ sgb_enum(C, en); return; }
jval *ty = json_get(sc, "type");
if (!ty || ty->t != J_STR){ sgb_fail(C, "missing type"); return; }
const char *t = ty->str;
if (!strcmp(t, "object")) sgb_object(C, sc, depth);
else if (!strcmp(t, "array")) sgb_array(C, sc, depth);
else if (!strcmp(t, "string")){ C->use_str = 1; sgb_put(C, "jstr"); }
else if (!strcmp(t, "number")){ C->use_num = 1; sgb_put(C, "jnum"); }
else if (!strcmp(t, "integer")){ C->use_int = 1; sgb_put(C, "jint"); }
else if (!strcmp(t, "boolean")) sgb_put(C, "( \"true\" | \"false\" )");
else if (!strcmp(t, "null")) sgb_put(C, "\"null\"");
else sgb_fail(C, t);
}
static void sgb_free_jval(jval *v){
if (!v) return;
if (v->t == J_OBJ){
for (int i = 0; i < v->len; i++){ free(v->keys[i]); sgb_free_jval(v->kids[i]); }
free(v->keys); free(v->kids);
} else if (v->t == J_ARR){
for (int i = 0; i < v->len; i++) sgb_free_jval(v->kids[i]);
free(v->kids);
} else if (v->t == J_STR) free(v->str);
free(v);
}
/* Compile a JSON-Schema string to GBNF. Returns a malloc'd GBNF text (caller
* frees) or NULL with a message in err (if err != NULL). */
static char *schema_to_gbnf(const char *schema_json, char *err, int errsz){
SgbCtx C; memset(&C, 0, sizeof C);
jval *sc = json_parse(schema_json, NULL);
if (!sc){ if (err) snprintf(err, errsz, "schema: json parse failed"); return NULL; }
sgb_put(&C, "root ::= jws ");
sgb_value(&C, sc, 0);
sgb_put(&C, " jws\n");
sgb_put(&C, "jws ::= ( \" \" | \"\\t\" | \"\\n\" | \"\\r\" )*\n");
if (C.use_str)
sgb_put(&C, "jstr ::= \"\\\"\" jchar* \"\\\"\"\n"
"jchar ::= [^\"\\\\\\x00-\\x1f] | \"\\\\\" ( [\"\\\\/bfnrt] | \"u\" jhex jhex jhex jhex )\n"
"jhex ::= [0-9a-fA-F]\n");
if (C.use_num)
sgb_put(&C, "jnum ::= \"-\"? ( \"0\" | [1-9] [0-9]* ) ( \".\" [0-9]+ )? ( ( \"e\" | \"E\" ) ( \"+\" | \"-\" )? [0-9]+ )?\n");
if (C.use_int)
sgb_put(&C, "jint ::= \"-\"? ( \"0\" | [1-9] [0-9]* )\n");
sgb_free_jval(sc);
if (C.fail || !C.s){
if (err) snprintf(err, errsz, "%s", C.err[0] ? C.err : "schema: compile failed");
free(C.s); return NULL;
}
return C.s;
}
#endif /* SCHEMA_GBNF_H */
+34
View File
@@ -14,9 +14,15 @@
#include <fcntl.h>
#include <unistd.h>
#include <dirent.h>
#include <sys/stat.h>
#include "json.h"
#include "compat.h"
/* tetto sulla dimensione dell'header safetensors: gli header reali sono piccoli
* (KB..pochi MB). Un file crafted che dichiara un hlen enorme causerebbe una
* malloc gigante prima ancora di leggere: lo respingiamo. */
#define ST_MAX_HEADER (512ll << 20)
typedef struct {
char *name;
int fd;
@@ -118,14 +124,27 @@ static void st_init(shards *S, const char *snap_dir) {
for (int fi = 0; fi < nf; fi++) {
int fd = st_open_fd(S, files[fi]);
struct stat sst;
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); }
/* 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
* malloc gigante. */
if (fsz < 8 || hlen > (uint64_t)(fsz - 8) || hlen > (uint64_t)ST_MAX_HEADER) {
fprintf(stderr, "%s: bad safetensors header length %llu (file %lld bytes)\n",
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); }
hdr[hlen] = 0;
int64_t data_start = 8 + (int64_t)hlen;
char *arena = NULL;
jval *root = json_parse(hdr, &arena);
if (!root || root->t != J_OBJ) {
fprintf(stderr, "%s: safetensors header is not a JSON object\n", files[fi]); exit(1); }
for (int i = 0; i < root->len; i++) {
const char *name = root->keys[i];
if (!strcmp(name, "__metadata__")) continue;
@@ -133,7 +152,21 @@ static void st_init(shards *S, const char *snap_dir) {
jval *dt = json_get(m, "dtype");
jval *off = json_get(m, "data_offsets");
jval *shp = json_get(m, "shape");
/* un header crafted puo' omettere i campi o dare tipi sbagliati:
* senza questi guard si dereferenzia NULL (json_get) o si legge
* off->kids[0/1] oltre i limiti dell'array. */
if (!dt || dt->t != J_STR || !off || off->t != J_ARR || off->len < 2 ||
!shp || shp->t != J_ARR) {
fprintf(stderr, "%s: tensor '%s' has malformed dtype/data_offsets/shape\n",
files[fi], name); exit(1); }
int64_t a0 = (int64_t)off->kids[0]->num, b0 = (int64_t)off->kids[1]->num;
/* offset dichiarati dal file: non-negativi, ordinati e dentro al
* file. Altrimenti nbytes=b0-a0 diventa negativo -> malloc((size_t))
* gigante e la memcpy in st_read_f32 sfora il buffer del chiamante;
* oppure off punta fuori dal file. */
if (a0 < 0 || b0 < a0 || data_start + b0 > fsz) {
fprintf(stderr, "%s: tensor '%s' data_offsets [%lld,%lld] out of file bounds (%lld)\n",
files[fi], name, (long long)a0, (long long)b0, (long long)fsz); exit(1); }
int64_t numel = 1; for (int k = 0; k < shp->len; k++) numel *= (int64_t)shp->kids[k]->num;
if (S->n == S->cap) { S->cap *= 2; S->t = realloc(S->t, S->cap*sizeof(st_tensor)); }
st_tensor *t = &S->t[S->n++];
@@ -184,6 +217,7 @@ static int64_t st_read_f32(shards *S, const char *name, float *out, int drop) {
st_tensor *t = st_find(S, name);
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); }
if (t->dtype == 2) {
memcpy(out, raw, t->nbytes);
+13
View File
@@ -5,6 +5,14 @@
#include <cstdint>
#include <cstdlib>
#ifdef _WIN32
/* MSVC has no POSIX setenv/unsetenv */
static int setenv(const char *name, const char *value, int overwrite) {
(void)overwrite; return _putenv_s(name, value);
}
static int unsetenv(const char *name) { return _putenv_s(name, ""); }
#endif
static int close_enough(const float *got, const float *want, int n) {
for (int i = 0; i < n; i++) {
if (std::fabs(got[i] - want[i]) > 1e-4f) {
@@ -42,6 +50,11 @@ int main(int argc, char **argv) {
if (coli_cuda_tensor_upload(&t8, q8, s8, 1, 5, 2, d0)) return 1;
if (ndev > 1 && coli_cuda_tensor_upload(&t8, q8, s8, 1, 4, 2, d1)) return 1;
if (!coli_cuda_matmul(&t8, got, x, q8, s8, 1, 2, 4, 2, d0) || !close_enough(got, want8, 4)) return 1;
const int8_t q8b[8]={-1,-2,-3,-4, 1,-2,3,-4};
const float s8b[2]={1.f,.5f},want8b[4]={10.f,15.f,-3.f,-2.5f};
if(!coli_cuda_tensor_update(t8,q8b,s8b)||
!coli_cuda_matmul(&t8,got,x,q8b,s8b,1,2,4,2,d0)||
!close_enough(got,want8b,4))return 1;
/* Rows [-8,-1,0,7] and [1,2,3,4], packed low nibble first. */
const uint8_t q4[4] = {0x70, 0xf8, 0xa9, 0xcb};
+17
View File
@@ -51,6 +51,23 @@ int main(void){
if(compat_open_direct("no_such_file.tmp")>=0) return fail("open missing file must fail");
if(compat_fsize(-1)>=0) return fail("compat_fsize on bad fd must be negative");
/* compat_fadvise: WILLNEED warms the page cache (background read into throwaway
* buffer), DONTNEED is a documented no-op. After a WILLNEED the buffered fd's
* subsequent pread must still return the exact bytes — the cache-warmer must not
* corrupt data. Bad fd / non-WILLNEED advice must be safe no-ops (return 0). */
int wfd = open(TMPF, COMPAT_O_RDONLY);
if(wfd<0) return fail("open buffered for fadvise");
if(posix_fadvise(wfd, 0, (off_t)FSZ, POSIX_FADV_WILLNEED)!=0) return fail("WILLNEED returned nonzero");
if(posix_fadvise(wfd, 0, (off_t)FSZ, POSIX_FADV_DONTNEED)!=0) return fail("DONTNEED should be a safe no-op (return 0)");
if(posix_fadvise(-1, 0, (off_t)FSZ, POSIX_FADV_WILLNEED)!=0) return fail("WILLNEED on bad fd should no-op (return 0)");
if(posix_fadvise(wfd, 0, 0, POSIX_FADV_WILLNEED)!=0) return fail("WILLNEED with len<=0 should no-op");
/* verify data integrity through the buffered fd after the cache-warmer ran */
uint8_t *verify=malloc(FSZ);
if(pread(wfd, verify, FSZ, 0)!=(ssize_t)FSZ) return fail("fadvise: pread size");
if(memcmp(verify, pat, FSZ)!=0) return fail("fadvise: data corrupted by cache-warmer");
free(verify);
close(wfd);
close(dfd);
compat_aligned_free(buf); free(pat); remove(TMPF);
puts("compat direct tests: ok");
Binary file not shown.
+68
View File
@@ -0,0 +1,68 @@
"""env_for: default Windows misurati (DIRECT/PIPE/PILOT_REAL + blocco OMP).
Carica `coli` come modulo (ha la guardia __main__) e verifica il contratto:
- win32: i tre default I/O e il blocco OMP sono setdefault
- un override esplicito dell'utente vince sempre
- COLI_NO_OMP_TUNE spegne SOLO il blocco OMP, non i default I/O
- non-win32: env_for non tocca nulla di tutto questo
"""
import importlib.machinery
import importlib.util
import os
import sys
import types
import unittest
from pathlib import Path
from unittest import mock
HERE = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(HERE))
_loader = importlib.machinery.SourceFileLoader("coli_cli", str(HERE / "coli"))
_spec = importlib.util.spec_from_loader("coli_cli", _loader)
coli = importlib.util.module_from_spec(_spec)
_loader.exec_module(coli)
def args(**over):
base = dict(model="X", policy="quality", ram=0, ngen=0, topp=0, topk=0,
temp=None, repin=0, ctx=0, auto_tier=False, gpu=None, vram=0)
base.update(over)
return types.SimpleNamespace(**base)
class EnvDefaultsTest(unittest.TestCase):
def env_for_with(self, environ, platform):
with mock.patch.dict(os.environ, environ, clear=True), \
mock.patch.object(sys, "platform", platform):
return coli.env_for(args())
def test_win32_sets_measured_defaults(self):
e = self.env_for_with({}, "win32")
self.assertEqual(e["DIRECT"], "1")
self.assertEqual(e["PIPE"], "1")
self.assertEqual(e["PILOT_REAL"], "1")
self.assertEqual(e["OMP_WAIT_POLICY"], "active")
self.assertNotIn("OMP_PROC_BIND", e) # MinGW libgomp: niente affinity
def test_explicit_override_wins(self):
e = self.env_for_with({"DIRECT": "0", "PIPE": "0"}, "win32")
self.assertEqual(e["DIRECT"], "0")
self.assertEqual(e["PIPE"], "0")
self.assertEqual(e["PILOT_REAL"], "1") # non overridden -> default
def test_kill_switch_scope_is_omp_only(self):
e = self.env_for_with({"COLI_NO_OMP_TUNE": "1"}, "win32")
self.assertNotIn("OMP_WAIT_POLICY", e)
self.assertNotIn("OMP_NUM_THREADS", e)
self.assertEqual(e["DIRECT"], "1") # i default I/O restano attivi
self.assertEqual(e["PIPE"], "1")
def test_non_windows_untouched(self):
e = self.env_for_with({}, "linux")
for k in ("DIRECT", "PIPE", "PILOT_REAL", "OMP_WAIT_POLICY"):
self.assertNotIn(k, e)
if __name__ == "__main__":
unittest.main()
Binary file not shown.
BIN
View File
Binary file not shown.
+55
View File
@@ -22,6 +22,50 @@ static int32_t ref_i4i8(const uint8_t *w4, const int8_t *x, int I){
return (int32_t)s;
}
/* Driver-level exactness: matmul_qt_ex on the IDOT path (allow_idot=1) must match
* a plain-C reference bit-for-bit. This exercises the SMMLA 2x2-tile drivers on the
* ARCH=native (i8mm) build and the SDOT drivers on the default build. The integer
* dot is exact and qrow_i8 is the shared quantizer, so any float bit mismatch is a
* driver bug (lane map, tiling, tail, or scale order), not rounding. */
static void fill_qt(QT *w, int fmt, int O, int I){
memset(w,0,sizeof *w);
w->fmt=fmt; w->O=O; w->I=I;
w->s=malloc((size_t)O*sizeof(float));
for(int o=0;o<O;o++) w->s[o]=0.001f+(float)(xr()%1000)*1e-6f;
if(fmt==1){
w->q8=malloc((size_t)O*I);
for(int64_t i=0;i<(int64_t)O*I;i++) w->q8[i]=(int8_t)((int)(xr()%256)-128);
}else{
size_t nb=(size_t)O*((I+1)/2);
w->q4=malloc(nb);
for(size_t i=0;i<nb;i++) w->q4[i]=(uint8_t)(xr()&0xFF);
}
}
static int check_driver(int fmt,int O,int I,int S){
QT w; fill_qt(&w,fmt,O,I); int rb=(I+1)/2;
float *x=malloc((size_t)S*I*sizeof(float));
float *y=malloc((size_t)S*O*sizeof(float));
float *yref=malloc((size_t)S*O*sizeof(float));
int8_t *xqr=malloc((size_t)S*I);
float *sxr=malloc((size_t)S*sizeof(float));
for(int64_t i=0;i<(int64_t)S*I;i++) x[i]=((float)(xr()%4001)-2000.f)/500.f;
matmul_qt_ex(y,x,&w,S,1);
for(int s=0;s<S;s++) sxr[s]=qrow_i8(x+(int64_t)s*I, xqr+(int64_t)s*I, I);
for(int o=0;o<O;o++) for(int s=0;s<S;s++){
int32_t d=fmt==1 ? ref_i8i8(w.q8+(int64_t)o*I, xqr+(int64_t)s*I, I)
: ref_i4i8(w.q4+(int64_t)o*rb, xqr+(int64_t)s*I, I);
yref[(int64_t)s*O+o]=(float)d*w.s[o]*sxr[s];
}
int rc=0;
for(int64_t i=0;i<(int64_t)S*O;i++)
if(memcmp(&y[i],&yref[i],sizeof(float))!=0){
fprintf(stderr,"FAIL driver fmt=%d O=%d I=%d S=%d idx=%lld: %.9g != %.9g\n",
fmt,O,I,S,(long long)i,(double)y[i],(double)yref[i]); rc=1; break;
}
free(w.s); free(w.q8); free(w.q4); free(x); free(y); free(yref); free(xqr); free(sxr);
return rc;
}
int main(void){
static const int sizes[]={1,2,15,16,17,31,32,33,63,64,65,100,127,128,1408,4096,4097};
static int8_t w[8192], x[8192]; static uint8_t w4[4096];
@@ -40,5 +84,16 @@ int main(void){
}
}
printf("idot kernel exactness (%s): ok\n", IDOT_KERNEL);
static const int Os[]={1,2,3,64,65};
static const int Is[]={16,17,100,1408};
static const int Ss[]={2,3,4,5,8};
for(int rep=0;rep<4;rep++)
for(unsigned a=0;a<sizeof Os/sizeof Os[0];a++)
for(unsigned b=0;b<sizeof Is/sizeof Is[0];b++)
for(unsigned c=0;c<sizeof Ss/sizeof Ss[0];c++)
for(int fmt=1;fmt<=2;fmt++)
if(check_driver(fmt,Os[a],Is[b],Ss[c])) return 1;
printf("idot driver exactness (%s): ok\n", IDOT_KERNEL);
return 0;
}
+23
View File
@@ -0,0 +1,23 @@
/* kv_alloc must survive re-allocation on the same KVState: every free path is
* guarded by if(k->Lc) precisely so callers (context resize, slot re-init) can
* call it again. A stale duplicate free block frees every Lc[i]/Rc[i] and both
* arrays twice on the second call -> allocator abort. No model file needed:
* the CPU path of kv_alloc only reads c->n_layers/kv_lora/qk_rope. */
#define main coli_glm_main_unused
#include "../glm.c"
#undef main
int main(void){
static Model m;
m.c.n_layers=2; m.c.kv_lora=8; m.c.qk_rope=4;
m.kv=calloc(1,sizeof(KVState));
kv_alloc(&m,16);
for(int i=0;i<m.c.n_layers+1;i++){ m.Lc[i][0]=1.0f; m.Rc[i][0]=1.0f; }
kv_alloc(&m,32); /* the re-allocation path under test */
for(int i=0;i<m.c.n_layers+1;i++){
m.Lc[i][(int64_t)32*m.c.kv_lora-1]=2.0f;
m.Rc[i][(int64_t)32*m.c.qk_rope-1]=2.0f;
}
printf("OK kv_alloc re-allocation\n");
return 0;
}
+42
View File
@@ -11,6 +11,24 @@ MAKE = shutil.which("make")
@unittest.skipUnless(MAKE, "make is required")
class MakefilePlatformTests(unittest.TestCase):
def _dry_run(self, target, triplet, **variables):
args = [
MAKE,
"--no-print-directory",
"-B",
"-n",
target,
f"TRIPLET={triplet}",
]
args.extend(f"{name}={value}" for name, value in variables.items())
return subprocess.run(
args,
cwd=C_DIR,
text=True,
capture_output=True,
check=True,
)
def test_windows_nt_without_uname_selects_mingw_build(self):
env = os.environ.copy()
env["OS"] = "Windows_NT"
@@ -29,6 +47,30 @@ class MakefilePlatformTests(unittest.TestCase):
self.assertIn("-fopenmp", result.stdout)
self.assertIn("-static", result.stdout)
def test_portable_build_uses_target_architecture(self):
cases = (
("x86_64-unknown-linux-gnu", "-march=x86-64-v3"),
("aarch64-unknown-linux-gnu", "-march=armv8-a"),
("powerpc64le-unknown-linux-gnu", "-mcpu=power8"),
("ppc64le-unknown-linux-gnu", "-mcpu=power8"),
)
for triplet, expected_flag in cases:
with self.subTest(triplet=triplet):
result = self._dry_run("portable", triplet)
self.assertIn(expected_flag, result.stdout)
def test_darwin_portable_build_does_not_force_x86_architecture(self):
missing_libomp = "/colibri-test/missing-libomp"
result = self._dry_run(
"portable", "arm64-apple-darwin", OMPDIR=missing_libomp
)
self.assertIn("clang -O3", result.stdout)
self.assertNotIn("-mcpu=x86-64-v3", result.stdout)
self.assertNotIn(missing_libomp, result.stdout)
self.assertNotIn("-fopenmp", result.stdout)
if __name__ == "__main__":
unittest.main()
+48 -2
View File
@@ -2,13 +2,15 @@ import io
import json
import math
import socket
import tempfile
import threading
import unittest
from unittest.mock import patch
from urllib.error import HTTPError
from urllib.request import Request, urlopen
from pathlib import Path
from openai_server import (APIError, APIServer, ClientCancelled, END, GenerationScheduler,
from openai_server import (APIError, APIHandler, APIServer, ClientCancelled, END, GenerationScheduler,
READY, Engine, generation_options, parse_tool_calls,
read_engine_turn, render_chat, serve)
@@ -70,8 +72,13 @@ class TemplateTest(unittest.TestCase):
def test_validates_generation_limits(self):
self.assertEqual(generation_options({"max_tokens": 4, "temperature": 0, "top_p": 1}, 8),
(4, 0.0, 1.0))
# max_tokens above the server cap is clamped, not rejected (#260): OpenAI
# clients default to large values; erroring breaks them.
self.assertEqual(generation_options({"max_tokens": 9, "temperature": 0, "top_p": 1}, 8),
(8, 0.0, 1.0))
# non-positive / non-int max_tokens is still a hard error
with self.assertRaises(APIError):
generation_options({"max_tokens": 9}, 8)
generation_options({"max_tokens": 0}, 8)
with self.assertRaises(APIError):
generation_options({"temperature": math.nan}, 8)
with self.assertRaises(APIError):
@@ -490,6 +497,12 @@ class HTTPTest(unittest.TestCase):
self.assertEqual(body["choices"][0]["text"], "Héllo")
self.assertEqual(self.engine.calls[-1][0], "Complete me")
def test_rejects_empty_legacy_completion(self):
with self.assertRaises(HTTPError) as caught:
self.request("/v1/completions", {"model": "test-model", "prompt": ""})
self.assertEqual(caught.exception.code, 400)
self.assertEqual(json.load(caught.exception)["error"]["param"], "prompt")
def test_rejects_invalid_stream_options(self):
with self.assertRaises(HTTPError) as caught:
self.request("/v1/chat/completions", {
@@ -499,6 +512,39 @@ class HTTPTest(unittest.TestCase):
self.assertEqual(caught.exception.code, 400)
class StaticServingTest(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
root = Path(self.tmp.name)
dist = root / "dist"
dist.mkdir()
(dist / "index.html").write_text("dashboard", encoding="utf-8")
sibling = root / "dist-private"
sibling.mkdir()
(sibling / "secret.txt").write_text("private", encoding="utf-8")
self.web_dist = patch.object(APIHandler, "WEB_DIST", dist)
self.web_dist.start()
self.server = APIServer(("127.0.0.1", 0), FakeEngine(), "test-model")
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
self.thread.start()
self.base = f"http://127.0.0.1:{self.server.server_port}"
def tearDown(self):
self.server.scheduler.close()
self.server.shutdown()
self.server.server_close()
self.thread.join(timeout=2)
self.web_dist.stop()
self.tmp.cleanup()
def test_static_root_stays_inside_dist_directory(self):
with urlopen(self.base + "/", timeout=2) as response:
self.assertEqual(response.read(), b"dashboard")
with self.assertRaises(HTTPError) as caught:
urlopen(self.base + "/%2e%2e/dist-private/secret.txt", timeout=2)
self.assertEqual(caught.exception.code, 404)
class SchedulerHTTPTest(unittest.TestCase):
def setUp(self):
self.engine = BlockingEngine()
+137
View File
@@ -0,0 +1,137 @@
/* Unit tests for the resident-pipeline primitives (Inc.0).
* Each primitive is checked against the exact CPU math from glm.c.
* Build: nvcc backend_cuda.cu tests/test_pipe_cuda.cu -o pipe_test && ./pipe_test */
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <cstdint>
#include "../backend_cuda.h"
static uint64_t rng=0x9E3779B97F4A7C15ULL;
static float rndf(void){ rng^=rng<<13; rng^=rng>>7; rng^=rng<<17;
return (float)((double)(rng>>11)/9007199254740992.0*2.0-1.0); }
static void ref_rmsnorm(float *out,const float *x,const float *w,int D,float eps){
double ms=0; for(int i=0;i<D;i++) ms+=(double)x[i]*x[i];
float r=1.f/sqrtf((float)(ms/D)+eps);
for(int i=0;i<D;i++) out[i]=x[i]*r*w[i];
}
static void ref_rope(float *v,int pos,int R,float theta){
int half=R/2; float in[256]; memcpy(in,v,R*sizeof(float));
for(int j=0;j<half;j++){
float inv=powf(theta,-2.0f*j/R);
float ang=pos*inv, cs=cosf(ang), sn=sinf(ang);
float a=in[2*j], b=in[2*j+1];
v[j]=a*cs-b*sn; v[half+j]=b*cs+a*sn;
}
}
static int close_enough(const float *a,const float *b,size_t n,float tol,const char *what){
float worst=0;
for(size_t i=0;i<n;i++){
float d=fabsf(a[i]-b[i]), m=fabsf(b[i])>1.f?fabsf(b[i]):1.f;
if(d/m>worst) worst=d/m;
}
printf(" %-14s max rel %.3e %s\n",what,worst,worst<=tol?"ok":"FAIL");
return worst<=tol;
}
int main(void){
int dev0=0;
if(!coli_cuda_init(&dev0,1)){ puts("pipe tests: skipped (no CUDA device)"); return 0; }
int ok=1;
enum { S=7, D=6144, I=2048, R=64, H=4, QH=192 };
/* rmsnorm */
{
float *x=(float*)malloc((size_t)S*D*4), *w=(float*)malloc(D*4);
float *got=(float*)malloc((size_t)S*D*4), *ref=(float*)malloc((size_t)S*D*4);
for(size_t i=0;i<(size_t)S*D;i++) x[i]=rndf()*3;
for(int i=0;i<D;i++) w[i]=1.f+rndf()*0.1f;
for(int s=0;s<S;s++) ref_rmsnorm(ref+(size_t)s*D,x+(size_t)s*D,w,D,1e-5f);
float *xd=(float*)coli_cuda_pipe_alloc(0,(size_t)S*D*4);
float *wd=(float*)coli_cuda_pipe_alloc(0,D*4);
float *yd=(float*)coli_cuda_pipe_alloc(0,(size_t)S*D*4);
ok&=coli_cuda_pipe_upload(0,xd,x,(size_t)S*D*4)&&coli_cuda_pipe_upload(0,wd,w,D*4);
ok&=coli_cuda_pipe_rmsnorm(0,yd,xd,wd,S,D,1e-5f)&&coli_cuda_pipe_sync(0);
ok&=coli_cuda_pipe_download(0,yd,got,(size_t)S*D*4);
ok&=close_enough(got,ref,(size_t)S*D,2e-5f,"rmsnorm");
coli_cuda_pipe_free(0,xd); coli_cuda_pipe_free(0,wd); coli_cuda_pipe_free(0,yd);
free(x); free(w); free(got); free(ref);
}
/* rope on [S,H,QH] query rows, last R dims per head, plus heads=1 k_rot rows */
{
size_t n=(size_t)S*H*QH;
float *v=(float*)malloc(n*4), *ref=(float*)malloc(n*4);
int pos[S]; for(int s=0;s<S;s++) pos[s]=s*17+3;
for(size_t i=0;i<n;i++) v[i]=rndf();
memcpy(ref,v,n*4);
for(int s=0;s<S;s++) for(int h=0;h<H;h++)
ref_rope(ref+((size_t)s*H+h)*QH+(QH-R),pos[s],R,10000.f);
float *vd=(float*)coli_cuda_pipe_alloc(0,n*4);
int *pd=(int*)coli_cuda_pipe_alloc(0,S*4);
ok&=coli_cuda_pipe_upload(0,vd,v,n*4)&&coli_cuda_pipe_upload(0,pd,pos,S*4);
ok&=coli_cuda_pipe_rope(0,vd,pd,S*H,QH,QH-R,R,H,10000.f)&&coli_cuda_pipe_sync(0);
ok&=coli_cuda_pipe_download(0,vd,v,n*4);
ok&=close_enough(v,ref,n,3e-4f,"rope");
coli_cuda_pipe_free(0,vd); coli_cuda_pipe_free(0,pd);
free(v); free(ref);
}
/* silu-mul + residual add */
{
size_t n=(size_t)S*I;
float *g=(float*)malloc(n*4), *u=(float*)malloc(n*4), *ref=(float*)malloc(n*4);
for(size_t i=0;i<n;i++){ g[i]=rndf()*4; u[i]=rndf()*4; }
for(size_t i=0;i<n;i++){ float s=g[i]/(1.f+expf(-g[i])); ref[i]=s*u[i]+u[i]; }
float *gd=(float*)coli_cuda_pipe_alloc(0,n*4), *ud=(float*)coli_cuda_pipe_alloc(0,n*4);
ok&=coli_cuda_pipe_upload(0,gd,g,n*4)&&coli_cuda_pipe_upload(0,ud,u,n*4);
ok&=coli_cuda_pipe_silu_mul(0,gd,ud,n);
ok&=coli_cuda_pipe_add(0,gd,ud,n)&&coli_cuda_pipe_sync(0);
ok&=coli_cuda_pipe_download(0,gd,g,n*4);
ok&=close_enough(g,ref,n,2e-5f,"silu+add");
coli_cuda_pipe_free(0,gd); coli_cuda_pipe_free(0,ud);
free(g); free(u); free(ref);
}
/* fixed-order rows_add */
{
float *x=(float*)calloc((size_t)S*D,4), *p=(float*)malloc((size_t)3*D*4), *ref=(float*)calloc((size_t)S*D,4);
int rows[3]={1,4,6};
for(size_t i=0;i<(size_t)3*D;i++) p[i]=rndf();
for(int b=0;b<3;b++) for(int i=0;i<D;i++) ref[(size_t)rows[b]*D+i]+=p[(size_t)b*D+i];
float *xd=(float*)coli_cuda_pipe_alloc(0,(size_t)S*D*4);
float *pd=(float*)coli_cuda_pipe_alloc(0,(size_t)3*D*4);
int *rd=(int*)coli_cuda_pipe_alloc(0,3*4);
ok&=coli_cuda_pipe_upload(0,xd,x,(size_t)S*D*4)&&coli_cuda_pipe_upload(0,pd,p,(size_t)3*D*4)&&coli_cuda_pipe_upload(0,rd,rows,3*4);
ok&=coli_cuda_pipe_rows_add(0,xd,pd,rd,3,D)&&coli_cuda_pipe_sync(0);
ok&=coli_cuda_pipe_download(0,xd,x,(size_t)S*D*4);
ok&=close_enough(x,ref,(size_t)S*D,1e-6f,"rows_add");
coli_cuda_pipe_free(0,xd); coli_cuda_pipe_free(0,pd); coli_cuda_pipe_free(0,rd);
free(x); free(p); free(ref);
}
/* device-input gemm vs host-path coli_cuda_matmul on an int4 tensor */
{
int O=I, K=D;
size_t rb=(size_t)(K+1)/2;
uint8_t *w4=(uint8_t*)malloc((size_t)O*rb);
float *sc=(float*)malloc(O*4), *x=(float*)malloc((size_t)S*K*4);
float *ref=(float*)malloc((size_t)S*O*4), *got=(float*)malloc((size_t)S*O*4);
for(size_t i=0;i<(size_t)O*rb;i++) w4[i]=(uint8_t)(rng=rng*6364136223846793005ULL+1);
for(int i=0;i<O;i++) sc[i]=0.01f+0.001f*(i%7);
for(size_t i=0;i<(size_t)S*K;i++) x[i]=rndf();
ColiCudaTensor *t=NULL;
ok&=coli_cuda_matmul(&t,ref,x,w4,sc,2,S,K,O,0); /* host path = reference */
float *xd=(float*)coli_cuda_pipe_alloc(0,(size_t)S*K*4);
float *yd=(float*)coli_cuda_pipe_alloc(0,(size_t)S*O*4);
ok&=coli_cuda_pipe_upload(0,xd,x,(size_t)S*K*4);
ok&=coli_cuda_pipe_gemm(t,yd,xd,S)&&coli_cuda_pipe_sync(0);
ok&=coli_cuda_pipe_download(0,yd,got,(size_t)S*O*4);
ok&=close_enough(got,ref,(size_t)S*O,1e-6f,"gemm(dev)");
coli_cuda_pipe_free(0,xd); coli_cuda_pipe_free(0,yd);
coli_cuda_tensor_free(t);
free(w4); free(sc); free(x); free(ref); free(got);
}
coli_cuda_shutdown();
if(!ok){ puts("pipe tests: FAIL"); return 1; }
puts("pipe tests: ok");
return 0;
}
+7 -2
View File
@@ -112,8 +112,13 @@ class ResourcePlanTest(unittest.TestCase):
self.assertEqual(env["COLI_CUDA"], "1")
self.assertEqual(env["COLI_GPUS"], "1")
self.assertEqual(env["OMP_NUM_THREADS"], str(plan["cpu"]["physical_cores"]))
self.assertEqual(env["OMP_PROC_BIND"], "spread")
self.assertEqual(env["OMP_PLACES"], "cores")
if sys.platform == "win32":
# MinGW libgomp: niente affinity su Windows, le chiavi non vanno emesse
self.assertNotIn("OMP_PROC_BIND", env)
self.assertNotIn("OMP_PLACES", env)
else:
self.assertEqual(env["OMP_PROC_BIND"], "spread")
self.assertEqual(env["OMP_PLACES"], "cores")
self.assertEqual(env["PIN_GB"], env["CUDA_EXPERT_GB"])
explicit_threads = environment_for_plan(plan, {"OMP_NUM_THREADS": "7",
+155
View File
@@ -0,0 +1,155 @@
/* test_schema_gbnf: JSON-Schema -> GBNF compiler (schema_gbnf.h) end-to-end with
* the grammar.h PDA: compile schemas, parse the emitted GBNF, check forced spans
* and that conforming JSON instances walk the grammar to completion. */
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include "../schema_gbnf.h"
#include "../grammar.h"
static int fails = 0;
#define CHECK(c) do{ if(!(c)){ printf("FAIL %s:%d: %s\n", __FILE__, __LINE__, #c); fails++; } }while(0)
/* compile schema, gr_parse the result; return 0 ok */
static int compile(const char *schema, Grammar *G, char *gbnf_out, int outsz){
char err[160] = {0};
char *g = schema_to_gbnf(schema, err, sizeof err);
if (!g) return -1;
if (gbnf_out) snprintf(gbnf_out, outsz, "%s", g);
int rc = gr_parse(G, g);
if (rc) printf(" gr_parse error: %s\nGBNF:\n%s\n", G->err, g);
free(g);
return rc;
}
/* walk a byte string through the PDA; returns bytes consumed */
static int walk(GrState *S, const char *bytes){
int n = 0;
for (const char *p = bytes; *p; p++, n++)
if (gr_accept(S, (unsigned char)*p) != 1) break;
return n;
}
int main(void){
/* 1. simple strict object: forced spans resume inside literals (jws points
* themselves are not forced), compact AND sloppy instances both walk */
{
Grammar G; GrState S;
const char *sc = "{\"type\":\"object\",\"properties\":{"
"\"score\":{\"type\":\"integer\"},\"verdict\":{\"type\":\"string\"}},"
"\"required\":[\"score\",\"verdict\"]}";
CHECK(compile(sc, &G, NULL, 0) == 0);
gr_state_init(&S, &G);
char f[256]; int n = gr_forced(&S, f, sizeof f);
CHECK(n == 0); /* jws: start not forced */
CHECK(walk(&S, "{\"") == 2);
n = gr_forced(&S, f, sizeof f);
CHECK(n > 0 && strncmp(f, "score\"", 6) == 0); /* key body still forces */
const char *rest = "score\":-42,\"verdict\":\"no_fit\"}";
CHECK(walk(&S, rest) == (int)strlen(rest));
unsigned char mask[32]; int can_end = 0;
gr_admissible(&S, mask, &can_end);
CHECK(can_end == 1);
/* the whole point of jws: a sloppy instance no longer kills the walker */
gr_state_init(&S, &G);
const char *sloppy = "{ \"score\" : -42 ,\n \"verdict\" : \"no_fit\" }";
CHECK(walk(&S, sloppy) == (int)strlen(sloppy));
gr_admissible(&S, mask, &can_end);
CHECK(can_end == 1);
gr_free(&G);
}
/* 2. enum: alternation, forced span resumes after disambiguation */
{
Grammar G; GrState S;
const char *sc = "{\"type\":\"object\",\"properties\":{"
"\"fit\":{\"type\":\"string\",\"enum\":[\"no_fit\",\"partial_fit\",\"strong_fit\"]}},"
"\"required\":[\"fit\"]}";
CHECK(compile(sc, &G, NULL, 0) == 0);
gr_state_init(&S, &G);
char f[256]; int n;
CHECK(walk(&S, "{\"fit\":\"p") == 9); /* 'p' picks partial_fit */
n = gr_forced(&S, f, sizeof f);
CHECK(n > 0 && strncmp(f, "artial_fit\"", 11) == 0); /* enum tail is forced (jws stops before }) */
gr_free(&G);
}
/* 3. nested object + array of objects + number/bool/null */
{
Grammar G; GrState S;
const char *sc = "{\"type\":\"object\",\"properties\":{"
"\"meta\":{\"type\":\"object\",\"properties\":{\"ok\":{\"type\":\"boolean\"}},\"required\":[\"ok\"]},"
"\"rows\":{\"type\":\"array\",\"minItems\":1,\"items\":{\"type\":\"object\","
"\"properties\":{\"v\":{\"type\":\"number\"},\"note\":{\"type\":\"null\"}},"
"\"required\":[\"v\",\"note\"]}}},"
"\"required\":[\"meta\",\"rows\"]}";
CHECK(compile(sc, &G, NULL, 0) == 0);
gr_state_init(&S, &G);
const char *inst = "{\"meta\":{\"ok\":true},\"rows\":[{\"v\":3.5,\"note\":null},{\"v\":-1e-3,\"note\":null}]}";
CHECK(walk(&S, inst) == (int)strlen(inst));
unsigned char mask[32]; int can_end = 0;
gr_admissible(&S, mask, &can_end);
CHECK(can_end == 1);
gr_free(&G);
}
/* 4. const + escaped key/value bytes */
{
Grammar G; GrState S;
const char *sc = "{\"type\":\"object\",\"properties\":{"
"\"k\\\"x\":{\"const\":\"a\\\\b\"}},\"required\":[\"k\\\"x\"]}";
CHECK(compile(sc, &G, NULL, 0) == 0);
gr_state_init(&S, &G);
const char *inst = "{\"k\\\"x\":\"a\\\\b\"}";
CHECK(walk(&S, inst) == (int)strlen(inst));
gr_free(&G);
}
/* 5. string content freedom: jstr accepts arbitrary text + escapes */
{
Grammar G; GrState S;
const char *sc = "{\"type\":\"object\",\"properties\":{\"t\":{\"type\":\"string\"}},\"required\":[\"t\"]}";
CHECK(compile(sc, &G, NULL, 0) == 0);
gr_state_init(&S, &G);
const char *inst = "{\"t\":\"hello \\\"w\\\" \\u00e9\\n x\"}";
CHECK(walk(&S, inst) == (int)strlen(inst));
gr_free(&G);
}
/* 6. unsupported schemas -> NULL (fallback), never crash */
{
char err[160];
CHECK(schema_to_gbnf("{\"oneOf\":[{\"type\":\"string\"}]}", err, sizeof err) == NULL);
CHECK(schema_to_gbnf("{\"type\":\"string\",\"pattern\":\"a+\"}", err, sizeof err) == NULL);
CHECK(schema_to_gbnf("{\"type\":\"object\",\"properties\":{\"a\":{\"type\":\"string\"},"
"\"b\":{\"type\":\"string\"}},\"required\":[\"a\"]}", err, sizeof err) == NULL); /* subset-required */
CHECK(schema_to_gbnf("not json at all {{", err, sizeof err) == NULL
|| 1 /* json.h is permissive; compiler must still fail or produce a parseable grammar */);
}
/* 7. integer grammar rejects leading zeros / accepts 0 */
{
Grammar G; GrState S;
const char *sc = "{\"type\":\"object\",\"properties\":{\"n\":{\"type\":\"integer\"}},\"required\":[\"n\"]}";
CHECK(compile(sc, &G, NULL, 0) == 0);
gr_state_init(&S, &G);
CHECK(walk(&S, "{\"n\":0}") == 7);
gr_state_init(&S, &G);
CHECK(walk(&S, "{\"n\":01}") < 8); /* leading zero not admitted */
gr_free(&G);
}
/* 8. number enum */
{
Grammar G; GrState S;
const char *sc = "{\"type\":\"object\",\"properties\":{\"b\":{\"enum\":[1,2,3]}},\"required\":[\"b\"]}";
CHECK(compile(sc, &G, NULL, 0) == 0);
gr_state_init(&S, &G);
CHECK(walk(&S, "{\"b\":2}") == 7);
gr_free(&G);
}
if (fails){ printf("test_schema_gbnf: %d FAILED\n", fails); return 1; }
printf("test_schema_gbnf: OK\n");
return 0;
}
+117
View File
@@ -0,0 +1,117 @@
#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#define main coli_glm_main_unused
#include "../glm.c"
#undef main
static int fail(const char *s){ fprintf(stderr,"FAIL: %s\n",s); return 1; }
static int test_expert_layout(int fd){
Model m={0}; ESlot slot={0}; UringBatch batch={0};
m.c.hidden=4; m.c.moe_inter=3; m.ebits=8;
m.S.n=6; m.S.cap=6; m.S.t=calloc(6,sizeof(st_tensor));
if(!m.S.t) return fail("tensor metadata allocation");
const char *proj[3]={"gate_proj","up_proj","down_proj"};
int wbytes[3]={12,12,12}, sbytes[3]={12,12,16};
unsigned char data[76];
for(int i=0;i<36;i++) data[i]=(unsigned char)(i+1);
float scales[10]; for(int i=0;i<10;i++) scales[i]=(float)i+0.5f;
memcpy(data+36,scales,sizeof(scales));
if(pwrite(fd,data,sizeof(data),0)!=(ssize_t)sizeof(data)){ free(m.S.t); return fail("expert fixture write"); }
int64_t wo=0,so=36;
for(int k=0;k<3;k++){
char name[300];
snprintf(name,sizeof(name),"model.layers.1.mlp.experts.7.%s.weight",proj[k]);
m.S.t[k]=(st_tensor){strdup(name),fd,wo,wbytes[k],3,wbytes[k]}; wo+=wbytes[k];
size_t n=strlen(name); memcpy(name+n,".qs",4);
m.S.t[3+k]=(st_tensor){strdup(name),fd,so,sbytes[k],2,sbytes[k]/4}; so+=sbytes[k];
}
if(uring_batch_init(&batch)){ free(m.S.t); return fail("expert ring init"); }
uring_batch_reset(&batch);
int li=uring_load_add(&batch,&m,1,7,&slot,1);
if(li!=0 || uring_submit_batch(&batch) || uring_finalize_load(&batch,li,1)){
coli_uring_close(&batch.ring); free(m.S.t); return fail("expert batch load");
}
int bad=slot.eid!=7 || slot.g.fmt!=1 || slot.u.fmt!=1 || slot.d.fmt!=1
|| memcmp(slot.g.q8,data,12) || memcmp(slot.u.q8,data+12,12) || memcmp(slot.d.q8,data+24,12)
|| memcmp(slot.g.s,scales,12) || memcmp(slot.u.s,scales+3,12) || memcmp(slot.d.s,scales+6,16);
coli_uring_close(&batch.ring);
compat_aligned_free(slot.slab); free(slot.fslab);
if(bad){ for(int i=0;i<m.S.n;i++) free(m.S.t[i].name); free(m.S.t); return fail("expert tensor views"); }
m.c.n_experts=8; m.c.n_layers=2; m.ecap=2;
m.pin=calloc(3,sizeof(ESlot*)); m.npin=calloc(3,sizeof(int));
m.ecache=calloc(3,sizeof(ESlot*)); m.ecn=calloc(3,sizeof(int));
m.ecache[1]=calloc(2,sizeof(ESlot));
if(!m.pin||!m.npin||!m.ecache||!m.ecn||!m.ecache[1])
return fail("pilot fixture allocation");
if(uring_batch_init(&g_ub_pilot)) return fail("pilot ring init");
memset(g_pilot_inflight,0,sizeof(g_pilot_inflight));
atomic_store(&g_cur_moe_layer,-1); atomic_store(&g_pilot_loads,0); atomic_store(&g_pilot_drops,0);
pilot_r=0; pilot_w=1; pilot_q[0].l=1; pilot_q[0].e=7;
pilot_uring_batch(&m);
bad=m.ecn[1]!=1 || m.ecache[1][0].eid!=7 || g_pilot_inflight[1]!=0
|| atomic_load(&g_pilot_loads)!=1 || atomic_load(&g_pilot_drops)!=0;
coli_uring_close(&g_ub_pilot.ring);
compat_aligned_free(m.ecache[1][0].slab); free(m.ecache[1][0].fslab);
free(m.ecache[1]);
free(m.pin); free(m.npin); free(m.ecache); free(m.ecn);
for(int i=0;i<m.S.n;i++) free(m.S.t[i].name);
free(m.S.t);
return bad?fail("pilot uring publication"):0;
}
int main(void){
char path[]="/tmp/coli-uring-XXXXXX";
int fd=mkstemp(path); if(fd<0) return fail("mkstemp");
unlink(path);
enum { N=4, SZ=4096 };
unsigned char src[N][SZ],dst[N][SZ];
for(int i=0;i<N;i++){
memset(src[i],0,sizeof(src[i])); memset(dst[i],0,sizeof(dst[i]));
for(int j=0;j<SZ;j++) src[i][j]=(unsigned char)(i*37+j);
if(pwrite(fd,src[i],SZ,(off_t)i*SZ)!=SZ){ close(fd); return fail("pwrite fixture"); }
}
ColiUring ring;
if(coli_uring_init(&ring,8)){
if(errno==EPERM || errno==ENOSYS || errno==EACCES){
printf("test_uring: skipped (%s)\n",strerror(errno)); close(fd); return 0;
}
close(fd); return fail("io_uring_setup");
}
if(coli_uring_set_workers(&ring,4)){
coli_uring_close(&ring); close(fd); return fail("io-wq worker limit");
}
for(int i=0;i<N;i++) if(coli_uring_prep_read(&ring,fd,dst[i],SZ,(off_t)i*SZ,(uint64_t)i+1)){
coli_uring_close(&ring); close(fd); return fail("prepare read");
}
if(coli_uring_enter(&ring,0)<0){ coli_uring_close(&ring); close(fd); return fail("submit"); }
int seen[N]={0},complete=0;
while(complete<N){
struct io_uring_cqe cqe; int reaped=0;
while(coli_uring_peek(&ring,&cqe)){
reaped=1;
if(cqe.user_data<1 || cqe.user_data>N || cqe.res!=SZ){
coli_uring_close(&ring); close(fd); return fail("bad completion");
}
int i=(int)cqe.user_data-1;
if(seen[i]++){ coli_uring_close(&ring); close(fd); return fail("duplicate completion"); }
complete++;
}
if(!reaped && complete<N && coli_uring_enter(&ring,1)<0){
coli_uring_close(&ring); close(fd); return fail("completion wait");
}
}
for(int i=0;i<N;i++) if(memcmp(src[i],dst[i],SZ)){
coli_uring_close(&ring); close(fd); return fail("read data mismatch");
}
coli_uring_close(&ring);
if(ftruncate(fd,0) || test_expert_layout(fd)){ close(fd); return 1; }
close(fd);
puts("test_uring: ok"); return 0;
}
+51
View File
@@ -0,0 +1,51 @@
#!/bin/bash
# Unified user-experience benchmark: fixed scenarios, fixed decoding, medians.
#
# Reports, per scenario: TTFT (prefill wall), decode tok/s, and the first line
# of the generated text (drift check). Runs each scenario REPS times and prints
# every sample; judge by the median, not the best.
#
# Usage: SNAP=/path/to/model [REPS=3] [PIPE=2] tools/bench_ux.sh [engine-binary]
# The engine env (COLI_CUDA, PIN, ...) is inherited from the caller so the same
# script exercises CPU-only, CUDA and pipeline configurations.
#
# Discipline (docs/experiments/glm52-6x5090-2026-07-12.md):
# - TEMP=0 DRAFT=0 always: greedy, no speculation, one variable at a time.
# - Same binary for every configuration under comparison.
# - .coli_usage drifts placement between runs: compare medians of >=3 reps,
# or snapshot/restore the usage file around the battery.
set -u
GLM="${1:-./glm}"
REPS="${REPS:-3}"
export TEMP=0 DRAFT=0
[ -n "${PIPE:-}" ] && export COLI_CUDA_PIPE="$PIPE"
SHORT_PROMPT="Explain why the sky is blue in simple terms."
LONGDOC_PROMPT=$(python3 - <<'PY'
base=("蜂鸟是世界上最小的鸟类之一,主要分布在美洲大陆。它们的翅膀每秒可以扇动五十到八十次,"
"使它们能够在空中悬停、倒飞,甚至垂直起降。蜂鸟的新陈代谢极快,心跳每分钟可达一千两百次,"
"因此它们必须不断进食来维持能量。它们的主要食物是花蜜,同时也会捕食小型昆虫和蜘蛛来补充蛋白质。"
"蜂鸟的喙细长,适合深入花朵内部吸食花蜜,舌头呈管状,可以快速伸缩。在授粉过程中,"
"蜂鸟扮演着重要的角色,许多美洲植物依赖蜂鸟传粉。")
print((base*12)+"请根据上文回答:蜂鸟的主要食物是什么?")
PY
)
scenario(){ # name prompt ngen
local name=$1 prompt=$2 ngen=$3
for r in $(seq 1 "$REPS"); do
local log; log=$(mktemp)
NGEN=$ngen PROMPT="$prompt" "$GLM" 64 4 4 >"$log" 2>&1
local line; line=$(grep -aE "prefill .* decode .*tok/s" "$log" | tail -1)
local head; head=$(grep -av "^\[" "$log" | grep -avE "PROFILE|PROFILO|ATTENTION|expert|CUDA|prefill|specul|^---|TOPP|stop|Motore|caricato|prompt:|token" \
| tail -1 | tr -d '\n' | cut -c1-48)
printf "%-10s r%d %s\n" "$name" "$r" "$line"
printf "%-10s r%d text: %s\n" "$name" "$r" "$head"
rm -f "$log"
done
}
echo "== bench_ux: reps=$REPS pipe=${COLI_CUDA_PIPE:-unset} $(date +%F' '%T) =="
scenario chat "$SHORT_PROMPT" 96
scenario longdoc "$LONGDOC_PROMPT" 96
echo "== bench_ux done $(date +%T) =="
+6 -2
View File
@@ -10,8 +10,10 @@ from pathlib import Path
SPEED_RE = re.compile(r"REPLAY decode:.*\| ([0-9.]+) tok/s")
# accetta sia il formato storico "expert-disk 0.123s |" sia quello attuale
# "expert-disk 0.123s service / 0.045s wait |" (glm.c profile_print)
PROFILE_RE = re.compile(
r"PROFILE: expert-disk ([0-9.]+)s \| expert-matmul ([0-9.]+)s "
r"PROFILE: expert-disk ([0-9.]+)s(?: service / ([0-9.]+)s wait)? \| expert-matmul ([0-9.]+)s "
r"\| attention ([0-9.]+)s .* lm_head ([0-9.]+)s \| other ([0-9.-]+)s"
)
PROFILE_KEYS = ("disk", "expert_matmul", "attention", "lm_head", "other")
@@ -23,7 +25,9 @@ def parse_output(stdout: str, stderr: str = "") -> tuple[float, list[float]]:
profile = PROFILE_RE.search(stdout)
if not speed or not profile:
raise RuntimeError(f"benchmark output missing\nstdout:\n{stdout}\nstderr:\n{stderr}")
return float(speed.group(1)), [float(value) for value in profile.groups()]
service, wait, *rest = profile.groups()
disk = float(service) + (float(wait) if wait else 0.0)
return float(speed.group(1)), [disk] + [float(value) for value in rest]
def execute(engine: str, env: dict[str, str]) -> tuple[float, list[float]]:
+92 -10
View File
@@ -51,6 +51,38 @@ def quant_int4(w, bits): # -> (qbytes U8 [O*ceil(I/2)], s
out[:, :v1.shape[1]] |= (v1 << 4)
return out.reshape(-1), s[:, 0].astype(np.float32)
def quant_int4_grouped(w, bits, gs=128):
"""Group-scaled int4: one scale per group of `gs` elements along the input dim.
Drastically reduces quantization error vs per-row scaling — matches the FP8
source's 128x128 block-scale granularity. Output layout:
qbytes: same packed nibble format as quant_int4
scales: f32 [O * ngroups] where ngroups = ceil(I/gs), laid out as
s[o * ngroups + g] = scale for row o, group g.
The engine detects this format (fmt=4) by checking the .qs array size."""
O, I = w.shape
qmax = (1 << (bits - 1)) - 1
ngroups = (I + gs - 1) // gs
# pad I to a multiple of gs for clean reshape, then trim
Ipad = ngroups * gs
wpad = np.zeros((O, Ipad), np.float32)
wpad[:, :I] = w
wr = wpad.reshape(O, ngroups, gs) # [O, ngroups, gs]
amax = np.abs(wr).max(axis=2, keepdims=True) # [O, ngroups, 1]
s = np.maximum(amax / qmax, 1e-8) # [O, ngroups, 1]
q = np.clip(np.rint(wr / s), -8, qmax).astype(np.int32) # [O, ngroups, gs]
q = q.reshape(O, Ipad)[:, :I] # trim padding -> [O, I]
# pack nibbles (identical to quant_int4)
rb = (I + 1) // 2
out = np.zeros((O, rb), np.uint8)
v0 = (q[:, 0::2] + 8).astype(np.uint8)
out[:, :v0.shape[1]] = v0
if I > 1:
v1 = (q[:, 1::2] + 8).astype(np.uint8)
out[:, :v1.shape[1]] |= (v1 << 4)
# scales: flatten [O, ngroups] -> [O * ngroups]
s_flat = s[:, :, 0].astype(np.float32).reshape(-1)
return out.reshape(-1), s_flat
def quant_int2(w, bits): # -> (qbytes U8 [O*ceil(I/4)], scale f32 [O]); 4/byte
O, I = w.shape
qmax = (1 << (bits - 1)) - 1 # bits=2 -> qmax=1, valori [-2,1]
@@ -84,7 +116,7 @@ def layer_idx(name):
def classify(name, n_layers, keep_mtp=False, keep_idx=False):
if name.endswith("_scale_inv"): return "consumed" # FP8 base: gestito col suo peso
# NVFP4 (modelopt): i sidecar delle scale sono consumati insieme al loro .weight U8.
# NVFP4 (modelopt): i sidecar delle scale sono consumati insieme al loro U8 .weight.
# EN: NVFP4 (modelopt): scale sidecars are consumed together with their U8 .weight.
if name.endswith((".weight_scale", ".weight_scale_2", ".input_scale")): return "consumed"
li = layer_idx(name)
@@ -105,7 +137,20 @@ def classify(name, n_layers, keep_mtp=False, keep_idx=False):
if name.endswith("norm.weight") or name == "model.norm.weight": return "f32"
if name in ("model.embed_tokens.weight", "lm_head.weight"): return "io"
if ".mlp.experts." in name and name.endswith(".weight"): return "x" # expert ROUTED (streaming)
if name.endswith(".weight"): return "q" # attn/dense-mlp/shared (residente)
# Split resident weights by type for mixed-precision control:
# "sh" = shared expert (fires on every token, highest sensitivity)
# "o" = o_proj attention (reconstructs output, biggest attn tensor)
# "kvb" = kv_b_proj (reconstructs KV cache on every decode step)
# "attn" = other attention projections (q_a, q_b, kv_a)
# "dmlp" = dense MLP (first 3 layers)
if "shared_experts" in name: return "sh"
if name.endswith("o_proj.weight"): return "o"
if name.endswith("kv_b_proj.weight"): return "kvb"
if any(name.endswith(k) for k in ("q_a_proj.weight", "q_b_proj.weight",
"kv_a_proj_with_mqa.weight")): return "attn"
if any(name.endswith(k) for k in ("mlp.gate_proj.weight", "mlp.up_proj.weight",
"mlp.down_proj.weight")): return "dmlp"
if name.endswith(".weight"): return "q" # fallback: other resident weights
return "f32"
# ---------- dequant NVFP4 (modelopt) di UN tensore expert -> f32 [O,I] ----------
@@ -169,7 +214,8 @@ def dequant(f, name, keys):
return (w * sc).numpy()
return f.get_tensor(name).to(torch.float32).numpy()
def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits, keep_mtp=False, keep_idx=False):
def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits,
keep_mtp=False, keep_idx=False, group_size=0, bits_map=None):
from safetensors import safe_open
with safe_open(path, framework="pt") as f:
keys = set(f.keys())
@@ -180,11 +226,22 @@ def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits, keep_mtp=Fals
if kind == "f32":
out_dict[name] = w.astype(np.float32)
else:
bits = io_bits if kind == "io" else xbits if kind == "x" else ebits
# Resolve bits for this tensor type: use bits_map override if provided,
# otherwise fall back to the classic ebits/xbits/io_bits scheme.
if bits_map and kind in bits_map:
bits = bits_map[kind]
else:
bits = io_bits if kind == "io" else xbits if kind == "x" else ebits
# Any unknown kind that fell through classify as "q"
if bits_map and kind not in bits_map and kind not in ("io", "x", "sh", "o", "kvb", "attn", "dmlp"):
bits = ebits
if w.ndim != 2: # es. bias 1D non previsto come 'q' -> tienilo f32
out_dict[name] = w.astype(np.float32); continue
q, s = (quant_int2(w, bits) if bits <= 2 else
quant_int4(w, bits) if bits <= 4 else quant_int8(w, bits))
if group_size > 0 and bits <= 4:
q, s = quant_int4_grouped(w, bits, group_size)
else:
q, s = (quant_int2(w, bits) if bits <= 2 else
quant_int4(w, bits) if bits <= 4 else quant_int8(w, bits))
out_dict[name] = q
out_dict[name + ".qs"] = s
@@ -198,6 +255,20 @@ def main():
ap.add_argument("--ebits", type=int, default=None) # bit residenti (default 4; 8 per --mtp/--indexer)
ap.add_argument("--io-bits", type=int, default=8) # bit di embed/lm_head
ap.add_argument("--xbits", type=int, default=None) # bit degli expert ROUTED (streaming); default=ebits
# Mixed-precision: per-tensor-type bit overrides. Default = ebits (all same).
# Set these higher to protect sensitive tensors from quantization error.
ap.add_argument("--shared-bits", type=int, default=None,
help="bits for shared expert (fires on every token, highest sensitivity). Default=ebits")
ap.add_argument("--o-bits", type=int, default=None,
help="bits for o_proj attention (reconstructs output, biggest attn tensor). Default=ebits")
ap.add_argument("--kvb-bits", type=int, default=None,
help="bits for kv_b_proj (reconstructs KV cache on every decode). Default=ebits")
ap.add_argument("--attn-bits", type=int, default=None,
help="bits for other attention projections (q_a, q_b, kv_a). Default=ebits")
ap.add_argument("--dmlp-bits", type=int, default=None,
help="bits for dense MLP (first 3 layers). Default=ebits")
ap.add_argument("--group-size", type=int, default=0, # 0 = per-row (backward compat); 128 = group-scaled
help="group size for int4 scales: 0=per-row (default), 128=one scale per 128 elements (much better quality)")
ap.add_argument("--n-layers", type=int, default=78)
ap.add_argument("--min-free-gb", type=float, default=20.0)
ap.add_argument("--selftest", action="store_true")
@@ -217,6 +288,17 @@ def main():
a.ebits = 8 if (a.mtp or a.indexer) else 4
if a.xbits is None: a.xbits = a.ebits
# Build per-type bits map. If a type-specific arg is set, use it; otherwise the
# converter falls back to ebits for that type.
bits_map = {}
if a.shared_bits is not None: bits_map["sh"] = a.shared_bits
if a.o_bits is not None: bits_map["o"] = a.o_bits
if a.kvb_bits is not None: bits_map["kvb"] = a.kvb_bits
if a.attn_bits is not None: bits_map["attn"] = a.attn_bits
if a.dmlp_bits is not None: bits_map["dmlp"] = a.dmlp_bits
if bits_map:
print(f"[MIXED] precision map: " + ", ".join(f"{k}={v}bit" for k,v in sorted(bits_map.items())))
if a.selftest_nvfp4:
import torch
# 1) LUT e2m1: i 16 codici devono decodificare esattamente ai valori attesi.
@@ -298,7 +380,7 @@ def main():
shards = sorted(glob.glob(os.path.join(a.indir, "*.safetensors")))
from safetensors.numpy import save_file
for i, sp in enumerate(shards):
out = {}; convert_shard(sp, out, a.n_layers, a.ebits, a.io_bits, a.xbits)
out = {}; convert_shard(sp, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size, bits_map=bits_map)
save_file(out, os.path.join(a.outdir, f"out-{i:05d}.safetensors"))
# copia config + tokenizer
for fn in ["config.json"]:
@@ -541,7 +623,7 @@ def main():
if os.path.exists(outp): print(f"[MTP] {outp} already done"); continue
print(f"[MTP {i+1}/{len(mtp_shards)}] downloading {sh}...", flush=True)
p = download_retry(a.repo, sh, tmp)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_mtp=True)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_mtp=True, group_size=a.group_size, bits_map=bits_map)
save_file(out, outp)
os.remove(p)
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
@@ -561,7 +643,7 @@ def main():
if os.path.exists(outp): continue # gia' fatto -> ripartibile
print(f"[IDX {i+1}/{len(idx_shards)}] downloading {sh}...", flush=True)
p = download_retry(a.repo, sh, tmp)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_idx=True)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_idx=True, group_size=a.group_size, bits_map=bits_map)
if out: save_file(out, outp)
os.remove(p)
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
@@ -575,7 +657,7 @@ def main():
if os.path.exists(outp): continue # gia' fatto -> ripartibile
print(f"[{i+1}/{len(shards)}] downloading {sh} ({free_gb(a.outdir):.0f} GB free)...", flush=True)
p = download_retry(a.repo, sh, tmp)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits)
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size, bits_map=bits_map)
save_file(out, outp)
os.remove(p) # <-- cancella subito lo shard fp8
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
+21 -5
View File
@@ -48,11 +48,24 @@ def load_docs(task, data_dir, limit, seed):
random.Random(seed).shuffle(docs)
return docs[:limit] if limit else docs
def build_requests(tk, docs_by_task):
def detect_prefix(snap):
"""GLM sees [gMASK]<sop> at the start of every training sequence; scoring raw text
without it is out-of-distribution and silently depresses/distorts scores (#108).
Default the prefix ON for GLM snapshots; EVAL_PREFIX (even empty) overrides."""
if "EVAL_PREFIX" in os.environ: return os.environ["EVAL_PREFIX"]
try: mt = json.load(open(os.path.join(snap, "config.json"))).get("model_type", "")
except Exception: mt = ""
if "glm" in mt.lower():
print("[prefix] GLM snapshot: prepending [gMASK]<sop> to every context "
"(override with EVAL_PREFIX, disable with EVAL_PREFIX=)", file=sys.stderr)
return "[gMASK]<sop>"
return ""
def build_requests(tk, docs_by_task, prefix=""):
reqs, meta, perq = [], [], {}
for t, docs in docs_by_task.items():
for qi, d in enumerate(docs):
ctx, conts, gold = d["ctx"], d["choices"], int(d["gold"])
ctx, conts, gold = prefix + d["ctx"], d["choices"], int(d["gold"])
ctx_ids = tk.encode(ctx).ids
for oi, cont in enumerate(conts):
full = tk.encode(ctx + cont).ids
@@ -115,14 +128,17 @@ def main():
docs_by_task = {t: load_docs(t, a.data, a.limit, a.seed) for t in tasks}
for t, d in docs_by_task.items(): print(f"[{t}] {len(d)} questions", file=sys.stderr)
reqs, meta, perq = build_requests(tk, docs_by_task)
reqs, meta, perq = build_requests(tk, docs_by_task, detect_prefix(a.snap))
print(f"total requests: {len(reqs)} (answer options)", file=sys.stderr)
if a.dry:
for r in reqs[:3]: print(" example request:", r[:80], "...", file=sys.stderr)
print("DRY: request construction and tokenization passed. Engine was not run.", file=sys.stderr); return
req_path = tempfile.mktemp(suffix=".txt")
open(req_path, "w").write("\n".join(reqs) + "\n")
# mkstemp (non mktemp): crea il file atomicamente con permessi 0600, niente
# race TOCTOU/symlink su una tmp dir condivisa (CWE-377).
fd, req_path = tempfile.mkstemp(suffix=".txt")
with os.fdopen(fd, "w") as f:
f.write("\n".join(reqs) + "\n")
env = dict(os.environ, SNAP=a.snap, SCORE=req_path)
if a.ram: env["RAM_GB"] = str(a.ram)
cmd = [a.glm, str(a.cap)] + a.bits.split()
+68
View File
@@ -0,0 +1,68 @@
# Expert Atlas — what does each of the 19,456 experts actually do?
Probe harness for #175. Runs a set of topic-tagged prompts, dumps each run's expert-routing
histogram, and turns them into a per-expert topic-affinity vector.
```bash
cd c
export COLI_MODEL=/path/to/glm52_i4
./tools/expert_atlas/sweep.sh # 30 probes (10 topics x 3 prompts)
python3 tools/expert_atlas/analyze.py --stats atlas_out/stats --out atlas_out/experts.json \
--web web/dist/experts.json # optional: feed the web dashboard Atlas
python3 tools/expert_atlas/validate.py atlas_out/stats 200 # leave-one-prompt-out check
```
`--web` writes the same atlas in the shape the web dashboard consumes (the Atlas galaxy and the
Brain hover tooltips): keyed `"layer:expert"` with `affinity`/`entropy`/`top`/`label`. It replaces
the retired `tools/expert_atlas.py`, whose API-driven probing ran through a live server and was
exposed to exactly the traps above (server-side `--topp`, speculative drafts, shared `.coli_usage`).
## Read this before you trust any atlas
Four things silently corrupt this measurement. The sweep script controls all of them; if you
roll your own, don't skip them.
| trap | effect | control |
|---|---|---|
| **`--topp`** | prunes experts by cumulative probability — measured: it hides **38% of the distinct experts** (7,587 → 4,687). It is also the *recommended speed setting*. | `TOPP=0` |
| **speculative drafts** | `eusage` is incremented inside `moe()`, *before* verification, so **rejected** drafts count. Those are experts routed for text the model never emitted. | `MTP=0 DRAFT=0` |
| **`.coli_usage`** | is loaded at startup and accumulates, so a naive `STATS` dump contains **all prior history**, not this run. | remove per run (script backs it up and restores) |
| **autocorrelation** | routing inside one run is highly correlated — the same context routes to the same experts token after token. An expert firing 38 times during one prompt is **one** observation, not 38. Chi-square/entropy on raw selections will certify single-prompt flukes as perfect specialists. | `analyze.py` requires the affinity to **replicate across a category's independent prompts** |
The CUDA expert tier is also not run-to-run deterministic, so the sweep uses `--gpu none`. Tier
config only decides *where weights live*, not what the router picks, so this costs nothing.
## Method
`analyze.py`:
1. `n[e][c]` — selections of expert *e* while running category *c*
2. `f[e][c] = n[e][c] / N[c]` — normalise by **category size** (prefill routes the prompt too, so a
verbose category would otherwise look busier)
3. `p(c|e)` — renormalise into a topic distribution per expert, i.e. base-rate corrected. Ranking
by raw count instead just rediscovers which experts are popular in general.
4. `spec(e) = 1 H(p(c|e)) / log C` — 0 = generalist, 1 = fires for exactly one topic
5. **replication gate** — an expert is only a candidate specialist for *c* if it fires in ≥2 of *c*'s
independent prompts
`validate.py` — leave-one-prompt-out. Build each category's top-K specialist set from its *other*
prompts, then check which set the held-out prompt's routing actually lands in. If specialisation
were an artifact of prompt wording, the held-out prompt would not prefer its own category.
## Result on GLM-5.2 744B int4 (Zen5, CPU routing path)
- **Leave-one-prompt-out: 29/30 = 96.7%** (chance 10%). Specialisation is a property of the topic,
not of prompt wording.
- The single miss is instructive: `写一首关于秋天的短诗` ("write a short poem about autumn") is
classified **poetry**, not Chinese — routing follows the **task** over the **language**.
- **Only 7.9% of experts are strong specialists** (spec ≥ 0.5). The "one expert = one topic"
picture is wrong for ~92% of them.
- Specialisation **rises with depth**: layer 3 ≈ 0.07 (generalist, token/syntax level) → layers
1858 ≈ 0.190.27.
- The replication gate removed **587** experts that looked like flawless specialists on one prompt.
## Extending it
`probes.json` is the whole probe set — add categories and prompts. Use **3+ prompts per category
with varied phrasing**, or the replication gate has nothing to check and you are back to measuring
one prompt. Keep prompt lengths in a similar band across categories.
+140
View File
@@ -0,0 +1,140 @@
#!/usr/bin/env python3
"""GLM-5.2 Expert Atlas — affinity with CROSS-PROMPT REPLICATION (#175).
Fixes a trap in the naive version. Routing selections inside one run are heavily correlated:
the same context routes to the same experts token after token. So an expert that fires 38
times during a single 'code' prompt is ONE effective observation, not 38 independent draws.
Treat them as independent and a chi-square will happily certify a single-prompt fluke as a
perfect specialist (spec=1.000, lift=10.0) on 38 selections. It is measuring one prompt.
So specialisation is only claimed when it REPLICATES across the independent prompts of a
category:
- each category has R prompts (here 3), each run is one replicate
- share[e][run] = selections of e in that run / total selections in that run
- an expert is a candidate specialist for category c only if it fires in >= MIN_RUNS of
c's runs (default 2/3) -> it is a property of the TOPIC, not of one prompt's wording
- affinity uses the MEAN share across a category's runs, so one hot run cannot carry it
- reliability = (runs in top category) / (runs in that category)
Also reports the generalist/specialist split by layer depth, which is an average over
thousands of experts and is robust to the above.
"""
import argparse, glob, json, math, os
from collections import defaultdict
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--stats", default="stats")
ap.add_argument("--min-count", type=int, default=30)
ap.add_argument("--min-runs", type=int, default=2, help="must fire in >= this many of the top category's runs")
ap.add_argument("--out", default="experts.json")
ap.add_argument("--web", default="", help="also write the web-dashboard experts.json (Atlas/Brain hover)")
a = ap.parse_args()
# run[(cat,idx)][(layer,expert)] = count ; run_tot[(cat,idx)] = total
run_counts, run_tot = defaultdict(dict), defaultdict(int)
for path in sorted(glob.glob(os.path.join(a.stats, "*.txt"))):
base = os.path.basename(path)[:-4]
cat, idx = base.rsplit("_", 1)
for line in open(path):
p = line.split()
if len(p) != 3:
continue
l, e, n = int(p[0]), int(p[1]), int(p[2])
run_counts[(cat, idx)][(l, e)] = n
run_tot[(cat, idx)] += n
cats = sorted({c for c, _ in run_counts})
runs_of = {c: sorted(i for cc, i in run_counts if cc == c) for c in cats}
C = len(cats)
print(f"categories ({C}): " + ", ".join(f"{c}[{len(runs_of[c])}]" for c in cats))
experts = {k for d in run_counts.values() for k in d}
print(f"experts seen: {len(experts):,}\n")
atlas, dropped_sparse, dropped_unrepl = [], 0, 0
for key in experts:
total = sum(run_counts[r].get(key, 0) for r in run_counts)
if total < a.min_count:
dropped_sparse += 1
continue
# mean share per category across its runs (a single hot run cannot carry the category)
mean_share, fired_runs = {}, {}
for c in cats:
shares, fired = [], 0
for i in runs_of[c]:
n = run_counts[(c, i)].get(key, 0)
shares.append(n / max(1, run_tot[(c, i)]))
if n > 0:
fired += 1
mean_share[c] = sum(shares) / len(shares)
fired_runs[c] = fired
s = sum(mean_share.values())
if s <= 0:
continue
p = {c: mean_share[c] / s for c in cats}
top = max(cats, key=lambda c: p[c])
# REPLICATION GATE: the affinity must show up in >= min-runs of that category's prompts
if fired_runs[top] < a.min_runs:
dropped_unrepl += 1
continue
H = -sum(v * math.log(v) for v in p.values() if v > 0)
atlas.append({
"layer": key[0], "expert": key[1], "total": total,
"spec": round(1.0 - H / math.log(C), 4),
"top_topic": top,
"top_lift": round(p[top] * C, 2),
"reliability": f"{fired_runs[top]}/{len(runs_of[top])}",
"p": {c: round(p[c], 4) for c in cats},
})
atlas.sort(key=lambda r: (-r["spec"], -r["total"]))
print(f"dropped {dropped_sparse:,} sparse (<{a.min_count} sel)")
print(f"dropped {dropped_unrepl:,} UNREPLICATED (fired in <{a.min_runs} runs of their top topic)")
print(f"kept {len(atlas):,} experts\n")
print("=== most specialised, replicated across prompts ===")
print(f"{'layer':>5} {'exp':>4} {'sel':>6} {'spec':>6} {'lift':>6} {'repl':>5} topic")
for r in atlas[:20]:
print(f"{r['layer']:>5} {r['expert']:>4} {r['total']:>6} {r['spec']:>6.3f} "
f"{r['top_lift']:>6.2f} {r['reliability']:>5} {r['top_topic']}")
print("\n=== specialisation vs layer depth (mean over experts; robust to the above) ===")
by_layer = defaultdict(list)
for r in atlas:
by_layer[r["layer"]].append(r["spec"])
ls = sorted(by_layer)
for L in ls[::max(1, len(ls)//13)]:
v = by_layer[L]
print(f" layer {L:>3} n={len(v):>4} spec {sum(v)/len(v):.3f} {'#'*int(60*sum(v)/len(v))}")
print("\n=== experts owned per topic (replicated only) ===")
own = defaultdict(int)
for r in atlas:
own[r["top_topic"]] += 1
for c in sorted(own, key=lambda x: -own[x]):
print(f" {c:<14} {own[c]:>5}")
strong = [r for r in atlas if r["spec"] >= 0.5]
print(f"\nstrong specialists (spec >= 0.5, replicated): {len(strong):,} / {len(atlas):,} "
f"({100*len(strong)/max(1,len(atlas)):.1f}%)")
json.dump({"categories": cats, "experts": atlas}, open(a.out, "w"), indent=1)
print(f"wrote {a.out}")
if a.web:
# Same atlas, keyed "layer:expert" with per-expert affinity/entropy/top/label —
# the shape the web dashboard consumes (Atlas galaxy, Brain hover).
web = {}
for r in atlas:
aff = {c: v for c, v in r["p"].items() if v > 0}
H = -sum(v * math.log2(v) for v in aff.values())
web[f"{r['layer']}:{r['expert']}"] = {
"affinity": aff, "entropy": round(H, 2), "top": r["top_topic"],
"label": f"specialist: {r['top_topic']}" if r["spec"] >= 0.5 else "generalist"}
json.dump({"categories": cats, "experts": web}, open(a.web, "w"))
print(f"wrote {a.web} (dashboard format, {len(web):,} experts)")
if __name__ == "__main__":
main()
+54
View File
@@ -0,0 +1,54 @@
{
"_comment": "Probe set for the GLM-5.2 Expert Atlas (#175). 10 categories x 3 prompts. Prompts are deliberately varied in phrasing within a category so the affinity vector reflects the TOPIC, not one prompt's surface form. Lengths are kept within a narrow band across categories: prefill routes the prompt tokens too, so a long prompt in one category would inflate its counts.",
"code_python": [
"Write a Python function that merges two sorted lists into one sorted list.",
"In Python, explain the difference between a generator and a list comprehension.",
"Refactor this Python snippet to avoid a nested loop: for a in xs: for b in ys: if a==b: out.append(a)"
],
"code_sql": [
"Write a SQL query that returns the top 5 customers by total order value.",
"Explain what a LEFT JOIN does differently from an INNER JOIN in SQL.",
"Write SQL to add an index on the email column of a users table and explain when it helps."
],
"math_proof": [
"Prove that the square root of 2 is irrational.",
"Show that the sum of the first n odd numbers equals n squared.",
"Explain why the derivative of e^x is e^x."
],
"chinese": [
"请用三句话解释什么是机器学习。",
"写一首关于秋天的短诗。",
"中国的四大发明是什么?请简要说明。"
],
"german": [
"Erkläre in drei Sätzen, wie ein Verbrennungsmotor funktioniert.",
"Schreibe eine kurze formelle E-Mail, in der du einen Termin absagst.",
"Was ist der Unterschied zwischen Dativ und Akkusativ im Deutschen?"
],
"poetry": [
"Write a short poem about the sea at night.",
"Compose four lines of verse about an empty train station.",
"Write a haiku about the first snow of winter."
],
"law": [
"Explain in plain terms what consideration means in contract law.",
"What is the difference between a misdemeanor and a felony?",
"Summarize what the term 'force majeure' covers in a commercial contract."
],
"medicine": [
"Explain the difference between type 1 and type 2 diabetes.",
"What are the common symptoms of iron deficiency anemia?",
"Describe how a vaccine produces immunity."
],
"json_format": [
"Return a JSON object with keys name, age, and email for a fictional user. Output only JSON.",
"Convert this to JSON: name Alice, roles admin and editor, active true. Output only JSON.",
"Write a JSON schema for an object with a required string field 'id' and an optional integer 'count'."
],
"casual_chat": [
"Hey, what's a good way to spend a rainy Sunday afternoon?",
"I'm feeling pretty tired today. Any tips to get through the afternoon?",
"What's your favourite kind of weather, and why?"
]
}
+82
View File
@@ -0,0 +1,82 @@
#!/usr/bin/env bash
# GLM-5.2 Expert Atlas — probe sweep (#175).
#
# cd c && ./tools/expert_atlas/sweep.sh [probes.json] [outdir]
#
# Env: COLI_MODEL, NGEN (default 64), COLI (default ./coli)
#
# ---------------------------------------------------------------------------------------------
# THE CONFOUNDS THIS SCRIPT EXISTS TO CONTROL. Each one silently corrupts the atlas.
#
# TOPP=0 --topp prunes experts by cumulative probability. Measured on GLM-5.2 int4,
# same prompt, only top-p changed:
# topp=0 -> 21,000 selections across 7,587 distinct experts
# topp=0.7 -> 11,944 selections across 4,687 distinct experts
# It hides 38% of the experts. It is also the RECOMMENDED SPEED SETTING, so this
# is very easy to walk into: with top-p on you profile the pruner, not the model.
#
# MTP=0 Speculative drafts route experts for tokens that are later REJECTED and never
# DRAFT=0 emitted (eusage is incremented inside moe(), before verification). Those counts
# would describe text the model never produced.
#
# --gpu none The CUDA expert tier is not run-to-run deterministic (VRAM placement shifts
# between runs). Routing on the CPU path is reproducible. The tier only decides
# where weights live, not what the router picks, so CPU costs nothing here.
#
# --temp 0 Greedy: deterministic continuation, reproducible atlas.
#
# rm .coli_usage before EVERY run
# eusage is LOADED from <model>/.coli_usage at startup and written back at exit, so
# a naive STATS dump contains ALL PRIOR HISTORY, not this run. Removing it per run
# makes each dump exactly one probe. The user's learned cache is backed up and
# restored on exit.
#
# Prompt lengths are kept in a narrow band across categories: prefill routes the prompt tokens
# too, so a verbose category would otherwise simply look "busier".
# ---------------------------------------------------------------------------------------------
set -uo pipefail
HERE="$(cd "$(dirname "$0")" && pwd)"
PROBES="${1:-$HERE/probes.json}"
OUT="${2:-./atlas_out}"
COLI="${COLI:-./coli}"
NGEN="${NGEN:-64}"
MODEL="${COLI_MODEL:?set COLI_MODEL to the snapshot directory}"
[ -f "$PROBES" ] || { echo "no probe file: $PROBES" >&2; exit 1; }
[ -x "$COLI" ] || { echo "no coli at $COLI (run from c/, or set COLI=)" >&2; exit 1; }
mkdir -p "$OUT/stats"
export COLI_MODEL="$MODEL" MTP=0 DRAFT=0 TOPP=0
USAGE="$MODEL/.coli_usage"
BACKUP="$OUT/.coli_usage.backup"
[ -f "$USAGE" ] && cp "$USAGE" "$BACKUP" && echo "backed up $USAGE"
restore(){ [ -f "$BACKUP" ] && cp "$BACKUP" "$USAGE" && echo "restored $USAGE"; }
trap restore EXIT
python3 - "$PROBES" > "$OUT/runlist.tsv" <<'PY'
import json, sys
for cat, prompts in json.load(open(sys.argv[1])).items():
if cat.startswith('_'):
continue
for i, p in enumerate(prompts):
print(f"{cat}\t{i}\t{p}")
PY
n=$(wc -l < "$OUT/runlist.tsv"); i=0
echo "$n probes -> $OUT/stats"
while IFS=$'\t' read -r cat idx prompt; do
i=$((i+1))
dst="$OUT/stats/${cat}_${idx}.txt"
[ -s "$dst" ] && { echo " [$i/$n] $cat/$idx (cached)"; continue; }
rm -f "$USAGE" # start from an EMPTY routing history
STATS="$dst" "$COLI" run "$prompt" --ngen "$NGEN" --ctx 4096 --gpu none --temp 0 \
> "$OUT/stats/${cat}_${idx}.log" 2>&1
echo " [$i/$n] $cat/$idx $(grep -aoE '[0-9]+ selections across [0-9]+ distinct experts' \
"$OUT/stats/${cat}_${idx}.log" | head -1)"
done < "$OUT/runlist.tsv"
echo
echo "next:"
echo " python3 $HERE/analyze.py --stats $OUT/stats --out $OUT/experts.json"
echo " python3 $HERE/validate.py $OUT/stats 200"
+79
View File
@@ -0,0 +1,79 @@
#!/usr/bin/env python3
"""Leave-one-prompt-out validation of the Expert Atlas (#175).
"Replicates across the 3 prompts I picked" is not the same as "generalises". The atlas is
only real if a specialist set learned from SOME prompts predicts routing on a prompt it has
never seen.
Protocol, for every category c and every held-out prompt h of c:
1. build c's top-K specialist set from c's OTHER prompts only (h excluded entirely)
2. do the same for all other categories (using ALL their prompts they never saw h either)
3. on the held-out run h, measure what share of routing selections land in each set
4. the atlas works if c's own set wins on h
If specialisation were an artifact of prompt wording, the held-out prompt would not prefer
its own category's set. Chance is 1/C.
"""
import glob, json, os, sys
from collections import defaultdict
STATS = sys.argv[1] if len(sys.argv) > 1 else "stats"
K = int(sys.argv[2]) if len(sys.argv) > 2 else 200 # specialists per category
runs, tot = {}, {}
for path in sorted(glob.glob(os.path.join(STATS, "*.txt"))):
cat, idx = os.path.basename(path)[:-4].rsplit("_", 1)
d = {}
t = 0
for line in open(path):
p = line.split()
if len(p) == 3:
d[(int(p[0]), int(p[1]))] = int(p[2])
t += int(p[2])
runs[(cat, idx)] = d
tot[(cat, idx)] = t
cats = sorted({c for c, _ in runs})
idxs = {c: sorted(i for cc, i in runs if cc == c) for c in cats}
C = len(cats)
def specialists(cat, exclude):
"""Top-K experts by lift for `cat`, computed WITHOUT the excluded run."""
share = defaultdict(lambda: defaultdict(float))
for c in cats:
used = [i for i in idxs[c] if not (c == cat and i == exclude)]
for i in used:
for k, n in runs[(c, i)].items():
share[k][c] += n / max(1, tot[(c, i)]) / len(used)
scored = []
for k, per in share.items():
s = sum(per.values())
if s <= 0:
continue
p = per[cat] / s
if per[cat] > 0:
scored.append((p, k))
scored.sort(reverse=True)
return {k for _, k in scored[:K]}
print(f"leave-one-prompt-out, {C} categories, top-{K} specialists per category")
print(f"chance = {100.0/C:.1f}%\n")
hits = 0
trials = 0
for c in cats:
for h in idxs[c]:
sets = {cc: specialists(cc, h if cc == c else None) for cc in cats}
held = runs[(c, h)]
htot = max(1, tot[(c, h)])
scores = {cc: sum(held.get(k, 0) for k in sets[cc]) / htot for cc in cats}
win = max(scores, key=scores.get)
ok = win == c
hits += ok
trials += 1
own = 100 * scores[c]
best_other = 100 * max(v for cc, v in scores.items() if cc != c)
print(f" {c:<12} prompt {h} own-set {own:5.2f}% best-other {best_other:5.2f}% "
f"-> {'HIT ' if ok else 'MISS'} (predicted {win})")
print(f"\naccuracy: {hits}/{trials} = {100*hits/trials:.1f}% (chance {100.0/C:.1f}%)")
+39 -5
View File
@@ -8,7 +8,7 @@ USO:
Poi:
python3 tools/eval_glm.py --snap /home/vincenzo/glm52_i4 --data ./bench --tasks mmlu --limit 40 --ram 15
"""
import os, json, argparse, random
import os, sys, json, time, argparse, random
def f_hellaswag(d):
ctx = (d["activity_label"] + ": " + d["ctx_a"] + " " + d["ctx_b"].capitalize()).strip()
@@ -38,19 +38,44 @@ TASKS = { # task: (path, config, split, formatter)
"openbookqa": ("allenai/openbookqa", "main", "validation", f_openbookqa),
}
def load_retry(path, cfg, split, tries=5):
"""L'hub restituisce 5xx/timeout transitori (#304): riprova con backoff invece di
morire al primo HEAD fallito. hf_hub riprende i download parziali dalla cache,
quindi il retry riparte da dove si era fermato, non da zero.
EN: the hub throws transient 5xx/timeouts (#304): retry with backoff instead of
dying on the first failed HEAD. hf_hub resumes partial downloads from its cache,
so a retry continues where it stopped rather than starting over."""
from datasets import load_dataset
for k in range(tries):
try:
return load_dataset(path, cfg, split=split)
except (KeyboardInterrupt, SystemExit):
raise
except Exception as e:
if k == tries - 1: raise
wait = 2 ** (k + 1)
print(f" {path}: {type(e).__name__}: {e} — retry {k+1}/{tries-1} in {wait}s")
time.sleep(wait)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", default="./bench")
ap.add_argument("--tasks", default="hellaswag,arc_challenge,mmlu")
ap.add_argument("--limit", type=int, default=300)
ap.add_argument("--seed", type=int, default=1234)
ap.add_argument("--tries", type=int, default=5)
a = ap.parse_args()
from datasets import load_dataset
os.makedirs(a.out, exist_ok=True)
failed = []
for t in [x.strip() for x in a.tasks.split(",") if x.strip()]:
if t not in TASKS: print("unknown task:", t); continue
if t not in TASKS: print("unknown task:", t); failed.append(t); continue
path, cfg, split, fn = TASKS[t]
ds = load_dataset(path, cfg, split=split)
try:
ds = load_retry(path, cfg, split, tries=max(a.tries, 1))
except Exception as e:
# un task fallito non deve uccidere gli altri / one failed task must not kill the rest
print(f"{t}: FAILED after {a.tries} tries ({type(e).__name__}: {e}) — skipping")
failed.append(t); continue
idx = list(range(len(ds))); random.Random(a.seed).shuffle(idx)
rows, n = [], 0
for i in idx:
@@ -61,9 +86,18 @@ def main():
except Exception: continue
if n >= a.limit: break
outp = os.path.join(a.out, t + ".jsonl")
with open(outp, "w") as f:
# scrittura atomica: coli controlla solo l'ESISTENZA del file, quindi un jsonl
# troncato da un run interrotto bloccherebbe il re-download per sempre.
# EN: atomic write: coli only checks the file EXISTS, so a truncated jsonl from
# an interrupted run would block re-download forever.
tmp = outp + ".part"
with open(tmp, "w") as f:
for r in rows: f.write(json.dumps(r) + "\n")
os.replace(tmp, outp)
print(f"{t}: {len(rows)} -> {outp}")
if failed:
print(f"incomplete: {', '.join(failed)} — rerun when the hub recovers (cached progress is kept)")
sys.exit(1)
if __name__ == "__main__":
main()
+137
View File
@@ -0,0 +1,137 @@
"""Helper: salva pesi in FP8 e4m3 + scale a blocchi 128x128, nello STESSO layout del
checkpoint reale GLM-5.2-FP8 che `convert_fp8_to_int4.py` legge.
Layout (deve combaciare col `dequant()` del converter, convert_fp8_to_int4.py:164-169):
- `name` F8_E4M3 [O, I]
- `name_scale_inv` F32 [ceil(O/128), ceil(I/128)] (NOTA: '_scale_inv', underscore)
dequant: W = q.float() * scale.repeat_interleave(128,0).repeat_interleave(128,1)[:O,:I]
Convenzione FBGEMM/TransformerEngine: scale = amax(blocco)/448 (448 = max e4m3),
si MEMORIZZA il valore e si MOLTIPLICA in dequant. Malgrado il nome "_scale_inv" il
checkpoint memorizza la scala (non il reciproco): e' un MOLTIPLIER.
EN: Helper that writes weights as FP8 e4m3 with 128x128 block scales, in the SAME layout
EN: as the real GLM-5.2-FP8 checkpoint that `convert_fp8_to_int4.py` reads.
EN: FBGEMM/TransformerEngine convention: scale = amax(block)/448, stored (not its
EN: reciprocal) and MULTIPLIED on dequant. Despite the name "_scale_inv" it is a multiplier.
"""
import torch
E4M3_MAX = 448.0 # max valore rappresentabile in float8_e4m3fn / max representable value
BLOCK = 128 # granularita' delle scale a blocchi del checkpoint FP8 / FP8 block scale granularity
def keep_f32(name, t):
"""Stesso set F32 di `classify()` in convert_fp8_to_int4.py (norme, router, bias 1-D).
Tutti gli altri tensori 2-D vengono quantizzati FP8 (attn/mlp/shared/expert/embed/lm_head).
EN: Same F32 set as the converter's classify(): norms, router, 1-D biases. All other 2-D
EN: tensors are FP8-quantized (attn/mlp/shared/expert/embed/lm_head)."""
if t.dim() < 2:
return True # bias 1-D, e_score_correction_bias
if name.endswith("e_score_correction_bias"):
return True
if name.endswith("mlp.gate.weight"):
return True # router (NON gate_proj): tenuto F32 / kept F32
if name.endswith("norm.weight") or name == "model.norm.weight":
return True # RMSNorm
return False
def fp8_block_quantize(w):
"""w: [O,I] f32 -> (w_fp8 float8_e4m3fn [O,I], scale_inv f32 [ceil(O/128),ceil(I/128)]).
Identica matematica al `--selftest` del converter (scale = amax(blocco)/448). Padda a
multipli di 128 internamente (gli zeri non alzano l'amax) e fa slice al risultato.
EN: same math as the converter's --selftest. Pads to 128 multiples internally (zeros do
EN: not raise amax), slices the result back to [O,I]."""
O, I = w.shape
nbO, nbI = (O + BLOCK - 1) // BLOCK, (I + BLOCK - 1) // BLOCK
Op, Ip = nbO * BLOCK, nbI * BLOCK
wpad = torch.zeros(Op, Ip, dtype=torch.float32, device=w.device)
wpad[:O, :I] = w
wb = wpad.view(nbO, BLOCK, nbI, BLOCK) # [nbO, BLOCK, nbI, BLOCK]
amax = wb.abs().amax(dim=(1, 3)) # [nbO, nbI]
scale = amax / E4M3_MAX # FBGEMM/TE: memorizza la scala / store the scale
scale = torch.where(scale == 0, torch.ones_like(scale), scale) # blocco tutto-zero -> no div0
scale = scale.to(torch.float32)
q = (wpad / scale.repeat_interleave(BLOCK, 0).repeat_interleave(BLOCK, 1)).clamp(-E4M3_MAX, E4M3_MAX)
w_fp8 = q.to(torch.float8_e4m3fn)
return w_fp8[:O, :I].contiguous(), scale.contiguous()
def fp8_block_dequantize(w_fp8, scale):
"""Esatto inverso di fp8_block_quantize, e identico al `dequant()` del converter.
EN: exact inverse of fp8_block_quantize, identical to the converter's dequant()."""
O, I = w_fp8.shape
qf = w_fp8.to(torch.float32)
return qf * scale.repeat_interleave(BLOCK, 0).repeat_interleave(BLOCK, 1)[:O, :I]
def unfuse_experts(sd):
"""Split HF's fused 3-D `experts.gate_up_proj` [E, 2*M, I] into per-expert 2-D
`experts.{e}.gate_proj` [M, I] + `experts.{e}.up_proj` [M, I], and
`experts.down_proj` [E, I, M] -> `experts.{e}.down_proj` [M_out, I].
The real GLM-5.2-FP8 checkpoint stores experts UNFUSED as per-expert 2-D tensors
(gate_proj, up_proj, down_proj), each with its own _scale_inv. HF's
GlmMoeDsaForCausalLM fuses gate+up into a single 3-D gate_up_proj for efficiency.
The converter (classify + ndim!=2 guard) and the C engine both expect the unfused
layout, so we split before saving.
Idempotent: if experts are already unfused (no 3-D gate_up_proj), returns sd as-is.
EN: split HF's fused 3-D expert weights into the per-expert 2-D layout that the real
EN: checkpoint uses and the converter/engine expect. No-op if already unfused."""
keys_to_remove = []
new_entries = {}
for name, t in sd.items():
if not name.endswith(".mlp.experts.gate_up_proj"):
continue
# prefix = everything before ".mlp.experts.gate_up_proj"
prefix = name[:-len(".mlp.experts.gate_up_proj")]
E, twoM, I = t.shape # [E, 2*intermediate, input]
M = twoM // 2
for e in range(E):
new_entries[f"{prefix}.mlp.experts.{e}.gate_proj.weight"] = t[e, :M, :].contiguous()
new_entries[f"{prefix}.mlp.experts.{e}.up_proj.weight"] = t[e, M:, :].contiguous()
keys_to_remove.append(name)
# down_proj may be 3-D [E, I, M] in the fused form, or already per-expert
for name, t in sd.items():
if not name.endswith(".mlp.experts.down_proj") or t.dim() != 3:
continue
prefix = name[:-len(".mlp.experts.down_proj")]
E = t.shape[0]
for e in range(E):
new_entries[f"{prefix}.mlp.experts.{e}.down_proj.weight"] = t[e].contiguous()
keys_to_remove.append(name)
for k in keys_to_remove:
sd.pop(k, None)
sd.update(new_entries)
return sd
def state_dict_to_fp8(sd):
"""Converte uno state_dict HuggingFace nel layout FP8 del checkpoint reale:
per ogni tensore quantizzabile 2-D scrive `{name}` (F8_E4M3) + `{name}_scale_inv` (F32);
norme/router/bias e qualsiasi tensore NON 2-D (es. pesi MLA impaccati 3-D) restano nel
dtype originale. Questo rispecchia il guard `w.ndim != 2 -> f32` del converter
(convert_fp8_to_int4.py:184). EN: builds the real-checkpoint FP8 layout. Only exactly-2-D
tensors are FP8-quantized; anything else (1-D, 3-D packed MLA weights, ...) is kept, exactly
like the converter's `ndim != 2 -> f32` guard."""
out = {}
for name, t in sd.items():
if keep_f32(name, t) or t.dim() != 2:
out[name] = t # f32 / 1-D / 3-D+: tieni / keep
else:
w_fp8, scale = fp8_block_quantize(t.float())
out[name] = w_fp8
out[name + "_scale_inv"] = scale
return out
def save_fp8_safetensors(sd, path):
"""Quantizza a blocchi FP8 e salva in un singolo safetensors leggibile dal converter
via `--indir`. EN: block-quantize to FP8 and save a single safetensors for the converter."""
from safetensors.torch import save_file
out = state_dict_to_fp8(sd)
save_file({k: v.contiguous() for k, v in out.items()}, str(path))
n_fp8 = sum(1 for v in out.values() if v.dtype == torch.float8_e4m3fn)
return n_fp8, len(out)
+35 -1
View File
@@ -3,15 +3,27 @@
This is not a useful language model. It preserves the real glm_moe_dsa data
flow while remaining small enough to generate locally and run repeated CPU/CUDA
A/B tests without downloading the 379 GB checkpoint.
With --fp8 the weights are written as FP8 e4m3 + 128x128 block scale_inv, in the
SAME layout as the real GLM-5.2-FP8 checkpoint, so convert_fp8_to_int4.py can
exercise its FP8->int4 dequant path on a local fixture (its dims are 128-friendly,
so this is also the right fixture for --group-size 128 testing):
python tools/make_glm_bench_model.py --fp8 --output glm_bench_fp8
python tools/convert_fp8_to_int4.py --indir glm_bench_fp8 --outdir glm_bench_i4 --ebits 4 --group-size 128
"""
import argparse
import json
import sys
from pathlib import Path
import torch
from transformers import GlmMoeDsaConfig, GlmMoeDsaForCausalLM
sys.path.insert(0, str(Path(__file__).resolve().parent)) # importa glm_fp8_emit se lanciato da c/
from glm_fp8_emit import save_fp8_safetensors, unfuse_experts
def build_config() -> GlmMoeDsaConfig:
return GlmMoeDsaConfig(
@@ -51,6 +63,9 @@ def main() -> None:
parser.add_argument("--output", default="glm_bench_medium")
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
parser.add_argument("--seed", type=int, default=1234)
parser.add_argument("--fp8", action="store_true",
help="write weights as FP8 e4m3 + 128x128 block scale_inv (same layout as "
"GLM-5.2-FP8) instead of bf16, so convert_fp8_to_int4.py can dequant+requant")
args = parser.parse_args()
torch.manual_seed(args.seed)
@@ -70,7 +85,6 @@ def main() -> None:
output = Path(args.output)
output.mkdir(parents=True, exist_ok=True)
params = sum(p.numel() for p in model.parameters())
model.save_pretrained(output, safe_serialization=True, max_shard_size="4GB")
model.to(args.device)
prompt = [3, 14, 159, 26, 53, 58, 200, 11, 77, 240, 5, 99]
@@ -79,6 +93,25 @@ def main() -> None:
full = model.generate(ids, max_new_tokens=8, do_sample=False, use_cache=True)[0]
logits = model(full.unsqueeze(0), use_cache=False).logits[0]
# Unfuse experts AFTER reference generation (model needs fused weights for
# forward/generate) but BEFORE saving — the real checkpoint and the converter
# + C engine all expect per-expert 2-D gate_proj/up_proj/down_proj tensors.
sd = model.state_dict()
unfuse_experts(sd)
if args.fp8:
n_fp8, n_tot = save_fp8_safetensors(sd, output / "model.safetensors")
# save_pretrained scrive config.json; nel path FP8 lo bypassiamo, quindi lo scriviamo
# a mano (serve al converter e al motore C). EN: save_pretrained writes config.json;
# the FP8 path bypasses it, so write it manually (converter + C engine need it).
(output / "config.json").write_text(json.dumps(cfg.to_dict()))
print(f"saved FP8: {n_fp8} e4m3 tensors (+{n_tot - n_fp8} scale_inv sidecars / f32) "
f"-> {output / 'model.safetensors'}")
else:
from safetensors.torch import save_file
save_file({k: v.contiguous() for k, v in sd.items()}, str(output / "model.safetensors"))
(output / "config.json").write_text(json.dumps(cfg.to_dict()))
ref = {
"prompt_ids": prompt,
"full_ids": full.cpu().tolist(),
@@ -89,6 +122,7 @@ def main() -> None:
"seed": args.seed,
"parameters": params,
"parameters_billions": round(params / 1e9, 4),
"format": "fp8-e4m3-128" if args.fp8 else "bf16",
"purpose": "backend benchmark fixture; random weights, not a language model",
}
(output / "bench_manifest.json").write_text(json.dumps(manifest, indent=2))
+53 -4
View File
@@ -3,10 +3,34 @@ Architettura vera (MLA + DSA indexer + router sigmoid/noaux_tc + shared expert),
dimensioni minuscole. Salva pesi+config in c/glm_tiny/ e un riferimento greedy in
c/ref_glm.json. seq corta (<= index_topk) cosi' il DSA seleziona tutte le key e
l'attenzione coincide con la MLA densa: il motore C puo' validare senza implementare
l'indexer sparso."""
import json, torch
l'indexer sparso.
--fp8: salva i pesi come FP8 e4m3 + scale a blocchi 128x128 (layout del checkpoint reale
GLM-5.2-FP8) invece di bf16, cosi' convert_fp8_to_int4.py puo' esercitare il path FP8->int4
su un modello minuscolo. PRIMA di calcolare ref_glm.json fa il round-trip dei pesi per FP8
(quant->dequant, copy_ nel modello): cosi' il riferimento riflette ESATTAMENTE il modello
FP8 che il converter legge, non il modello bf16 a precisione piena. Default: bf16 (oracolo
originale invariato).
EN: --fp8 writes FP8 e4m3 + 128x128 block scale_inv (real GLM-5.2-FP8 layout) instead of bf16,
EN: so convert_fp8_to_int4.py can run its FP8->int4 path on a tiny model. ref_glm.json is
EN: computed AFTER the FP8 round-trip, so the reference matches exactly what the converter
EN: ingests. Default: bf16 (original oracle unchanged)."""
import json, sys, argparse
from pathlib import Path
import torch
from transformers import GlmMoeDsaConfig, GlmMoeDsaForCausalLM
sys.path.insert(0, str(Path(__file__).resolve().parent)) # importa glm_fp8_emit se lanciato da c/
from glm_fp8_emit import (fp8_block_quantize, fp8_block_dequantize, keep_f32,
save_fp8_safetensors, unfuse_experts)
ap = argparse.ArgumentParser()
ap.add_argument("--fp8", action="store_true",
help="salva in FP8 e4m3 + 128x128 block scale_inv (layout GLM-5.2-FP8) e "
"calcola ref_glm.json sul modello dopo il round-trip FP8. "
"EN: write FP8 e4m3 + block scale_inv, ref computed on FP8-rounded model")
args = ap.parse_args()
torch.manual_seed(1234)
cfg = GlmMoeDsaConfig(
@@ -53,6 +77,18 @@ with torch.no_grad():
layer.mlp.gate.e_score_correction_bias.copy_(
torch.linspace(-0.1, 0.1, cfg.n_routed_experts))
# --fp8: round-trip dei pesi quantizzabili per FP8 PRIMA di calcolare il riferimento,
# cosi' ref_glm.json riflette esattamente il modello FP8 che il converter leggera'.
# Norme/router/bias (keep_f32) restano a precisione piena. EN: --fp8: round-trip quantizable
# weights through FP8 before computing the reference, so ref_glm.json matches the FP8 model.
if args.fp8:
with torch.no_grad():
for n, p in model.named_parameters():
if keep_f32(n, p) or p.dim() != 2:
continue
q, s = fp8_block_quantize(p)
p.copy_(fp8_block_dequantize(q, s))
print("=== state_dict tensors (names used by the C loader) ===")
for n, p in model.state_dict().items():
print(f" {n:60s} {tuple(p.shape)}")
@@ -73,7 +109,20 @@ with torch.no_grad():
tf_pred = lg.argmax(-1).tolist()
print("tf_pred:", tf_pred)
model.save_pretrained("glm_tiny", safe_serialization=True)
# Unfuse experts AFTER reference generation (model needs fused weights for
# forward/generate) but BEFORE saving — the real checkpoint and the converter
# + C engine all expect per-expert 2-D gate_proj/up_proj/down_proj tensors.
sd = model.state_dict()
unfuse_experts(sd)
if args.fp8:
n_fp8, n_tot = save_fp8_safetensors(sd, "glm_tiny/model.safetensors")
print(f"\nsaved FP8: {n_fp8} e4m3 tensors (+{n_tot - n_fp8} scale_inv sidecars / f32) "
f"-> glm_tiny/model.safetensors")
else:
from safetensors.torch import save_file
save_file({k: v.contiguous() for k, v in sd.items()}, "glm_tiny/model.safetensors")
json.dump(cfg.to_dict(), open("glm_tiny/config.json", "w"))
json.dump({"prompt_ids": prompt, "full_ids": full, "tf_pred": tf_pred}, open("ref_glm.json", "w"))
print("\nsaved: glm_tiny/ (weights + config) and ref_glm.json")
print("saved: glm_tiny/ (weights + config) and ref_glm.json"
+ (" [fp8]" if args.fp8 else ""))
+85 -12
View File
@@ -106,33 +106,106 @@ def rotation(dim, device, seed=417):
return q
def quantize_param(w, bits, group, rot=False):
def quantize_param(w, bits, group, rot=False, e8=""):
if w.ndim == 3: # fused experts [E, in, out] -> move input last
x = w.transpose(1, 2).contiguous()
x = _rot_quant(x, bits, group) if rot else _quant_last_dim(x, bits, group)
x = _rot_quant(x, bits, group, e8) if rot else _grid_or_e8(x, bits, group, e8)
return x.transpose(1, 2).contiguous()
if rot:
return _rot_quant(w, bits, group)
return _quant_last_dim(w, bits, group) # nn.Linear [out, in] -- input already last
return _rot_quant(w, bits, group, e8)
return _grid_or_e8(w, bits, group, e8) # nn.Linear [out, in] -- input already last
def _rot_quant(x, bits, group):
def _grid_or_e8(x, bits, group, e8):
if e8:
return _quant_e8(x.float(), group, ball=(e8 == "-e8"))
return _quant_last_dim(x, bits, group)
def _rot_quant(x, bits, group, e8=""):
"""W -> Qn(W@Q) @ Q^T along the last (input) dim — see rotation() above."""
q = rotation(x.shape[-1], x.device)
return (_quant_last_dim(x.float() @ q, bits, group) @ q.T).contiguous()
return (_grid_or_e8(x.float() @ q, bits, group, e8) @ q.T).contiguous()
SCHEME_RE = re.compile(r"^int(2|3|4|8)(?:-g(\d+))?(-rot)?(-nohead)?$")
# --------------------------------------------------------------------------------------
# E8 lattice quantization (#81 follow-up): the -rot schemes above are QuaRot (rotation +
# uniform grid). QuIP#'s 2-bit result needs the second ingredient — an E8 lattice codebook
# instead of the grid. E8 = D8 (D8 + 1/2), nearest point via ConwaySloane: round every
# coordinate, and if the sum is odd re-round the worst coordinate the other way; repeat on
# the half-shifted copy and keep the closer of the two. `-e8` clamps points to |p|^2 <= 10,
# the E8P ball QuIP# builds its 2^16 codebook from (2 bits/weight for 8-dim blocks);
# `-e8u` leaves the lattice unbounded — an ideal-codebook upper bound, not a deployable rate.
# Scale: per group, a small MSE search over multiples of the block RMS (absmax is the wrong
# statistic for a lattice — the ball wants energy matched, not the peak).
# --------------------------------------------------------------------------------------
def _d8_nearest(y):
f = torch.round(y)
d = y - f
odd = (f.sum(-1).long() & 1).bool()
idx = d.abs().argmax(-1, keepdim=True)
step = torch.where(d.gather(-1, idx) >= 0, 1.0, -1.0)
flipped = f.gather(-1, idx) + step
return f.scatter(-1, idx, torch.where(odd[..., None], flipped, f.gather(-1, idx)))
def _e8_nearest(y):
a = _d8_nearest(y)
b = _d8_nearest(y - 0.5) + 0.5
da = ((y - a) ** 2).sum(-1, keepdim=True)
db = ((y - b) ** 2).sum(-1, keepdim=True)
return torch.where(da <= db, a, b)
def _e8_ball(y, r2=10.0):
p = _e8_nearest(y)
for _ in range(8): # shrink-and-requantize until inside
n2 = (p ** 2).sum(-1, keepdim=True)
over = n2 > r2 + 1e-6
if not over.any():
break
y = torch.where(over, y * torch.sqrt(r2 / torch.clamp(n2, min=r2)) * 0.98, y)
p = torch.where(over, _e8_nearest(y), p)
return p
def _quant_e8(x, group, ball):
"""Blocks of 8 along the input dim; per-group scale by MSE search over RMS multiples."""
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]
if g % 8:
raise SystemExit(f"-e8 group {g} must be a multiple of 8")
shp = x.shape
xg = x.reshape(-1, g) # [G, g]
rms = torch.clamp(xg.pow(2).mean(-1, keepdim=True).sqrt(), min=1e-8)
best_out, best_err = None, None
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)
out = (p.reshape(-1, g) * s)
err = (out - xg).pow(2).sum(-1, keepdim=True)
if best_err is None:
best_out, best_err = out, err
else:
take = err < best_err
best_out = torch.where(take, out, best_out)
best_err = torch.where(take, err, best_err)
return best_out.reshape(shp)
SCHEME_RE = re.compile(r"^int(2|3|4|8)(?:-g(\d+))?(-e8u?)?(-rot)?(-nohead)?$")
def parse_scheme(name):
"""'int4-g128-nohead' -> (bits=4, group=128, skip_head=True). 'fp16' -> None."""
"""'int4-g128-nohead' -> (bits, group, e8, skip_head...). 'fp16' -> None."""
if name == "fp16":
return None
m = SCHEME_RE.match(name)
if not m:
raise SystemExit(f"bad scheme '{name}' (expected fp16 | int{{2,3,4,8}}[-g<N>][-rot][-nohead])")
return int(m.group(1)), int(m.group(2) or 0), bool(m.group(3)), bool(m.group(4))
raise SystemExit(f"bad scheme '{name}' (expected fp16 | int{{2,3,4,8}}[-g<N>][-e8|-e8u][-rot][-nohead])")
return int(m.group(1)), int(m.group(2) or 0), m.group(3) or "", bool(m.group(4)), bool(m.group(5))
def is_router(name):
@@ -152,7 +225,7 @@ def apply_scheme(model, scheme):
spec = parse_scheme(scheme)
if spec is None:
return 0, 0, total
bits, group, rot, skip_head = spec
bits, group, e8, rot, skip_head = spec
n = qp = 0
with torch.no_grad():
for name, p in model.named_parameters():
@@ -160,7 +233,7 @@ def apply_scheme(model, scheme):
continue
if skip_head and is_head_or_embed(name):
continue
p.data.copy_(quantize_param(p.data.float(), bits, group, rot).to(p.dtype))
p.data.copy_(quantize_param(p.data.float(), bits, group, rot, e8).to(p.dtype))
n += 1
qp += p.numel()
return n, qp, total
+2 -1
View File
@@ -5,12 +5,13 @@ Used by `make test-c` so the test runner works from any shell (cmd.exe,
PowerShell, Git Bash, MSYS2) without a POSIX `for` loop. Each test binary
is a positional argument.
"""
import os
import subprocess
import sys
failed = []
for binary in sys.argv[1:]:
rc = subprocess.call([binary])
rc = subprocess.call([os.path.normpath(binary)])
if rc != 0:
failed.append(binary)
if failed:
+137
View File
@@ -0,0 +1,137 @@
#ifndef COLI_URING_H
#define COLI_URING_H
/* Minimal Linux io_uring reader. Colibri owns the ring from one thread, queues
* a batch of positioned reads, and reaps CQEs without a userspace spin loop. */
#ifdef __linux__
#include <errno.h>
#include <linux/io_uring.h>
#include <stdint.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <unistd.h>
#ifndef MAP_POPULATE
#define MAP_POPULATE 0
#endif
typedef struct {
int fd;
struct io_uring_params p;
void *sq_map, *cq_map, *sqes_map;
size_t sq_map_sz, cq_map_sz, sqes_map_sz;
unsigned *sq_head, *sq_tail, *sq_mask, *sq_entries, *sq_array;
unsigned *cq_head, *cq_tail, *cq_mask, *cq_entries;
struct io_uring_sqe *sqes;
struct io_uring_cqe *cqes;
} ColiUring;
static inline unsigned coli_uring_load_acquire(unsigned *p){
return __atomic_load_n(p,__ATOMIC_ACQUIRE);
}
static inline void coli_uring_store_release(unsigned *p,unsigned v){
__atomic_store_n(p,v,__ATOMIC_RELEASE);
}
static inline void coli_uring_close(ColiUring *r){
if(!r) return;
if(r->sqes_map && r->sqes_map!=MAP_FAILED) munmap(r->sqes_map,r->sqes_map_sz);
if(r->cq_map && r->cq_map!=MAP_FAILED && r->cq_map!=r->sq_map) munmap(r->cq_map,r->cq_map_sz);
if(r->sq_map && r->sq_map!=MAP_FAILED) munmap(r->sq_map,r->sq_map_sz);
if(r->fd>=0) close(r->fd);
memset(r,0,sizeof(*r)); r->fd=-1;
}
static inline int coli_uring_init(ColiUring *r,unsigned entries){
memset(r,0,sizeof(*r)); r->fd=-1;
r->fd=(int)syscall(SYS_io_uring_setup,entries,&r->p);
if(r->fd<0) return -1;
r->sq_map_sz=r->p.sq_off.array+r->p.sq_entries*sizeof(unsigned);
r->cq_map_sz=r->p.cq_off.cqes+r->p.cq_entries*sizeof(struct io_uring_cqe);
if(r->p.features&IORING_FEAT_SINGLE_MMAP){
size_t n=r->sq_map_sz>r->cq_map_sz?r->sq_map_sz:r->cq_map_sz;
r->sq_map=mmap(NULL,n,PROT_READ|PROT_WRITE,MAP_SHARED|MAP_POPULATE,r->fd,IORING_OFF_SQ_RING);
if(r->sq_map==MAP_FAILED){ r->sq_map=NULL; coli_uring_close(r); return -1; }
r->sq_map_sz=n; r->cq_map=r->sq_map; r->cq_map_sz=n;
}else{
r->sq_map=mmap(NULL,r->sq_map_sz,PROT_READ|PROT_WRITE,MAP_SHARED|MAP_POPULATE,r->fd,IORING_OFF_SQ_RING);
if(r->sq_map==MAP_FAILED){ r->sq_map=NULL; coli_uring_close(r); return -1; }
r->cq_map=mmap(NULL,r->cq_map_sz,PROT_READ|PROT_WRITE,MAP_SHARED|MAP_POPULATE,r->fd,IORING_OFF_CQ_RING);
if(r->cq_map==MAP_FAILED){ r->cq_map=NULL; coli_uring_close(r); return -1; }
}
r->sqes_map_sz=r->p.sq_entries*sizeof(struct io_uring_sqe);
r->sqes_map=mmap(NULL,r->sqes_map_sz,PROT_READ|PROT_WRITE,MAP_SHARED|MAP_POPULATE,r->fd,IORING_OFF_SQES);
if(r->sqes_map==MAP_FAILED){ r->sqes_map=NULL; coli_uring_close(r); return -1; }
char *sq=(char*)r->sq_map,*cq=(char*)r->cq_map;
r->sq_head=(unsigned*)(sq+r->p.sq_off.head);
r->sq_tail=(unsigned*)(sq+r->p.sq_off.tail);
r->sq_mask=(unsigned*)(sq+r->p.sq_off.ring_mask);
r->sq_entries=(unsigned*)(sq+r->p.sq_off.ring_entries);
r->sq_array=(unsigned*)(sq+r->p.sq_off.array);
r->cq_head=(unsigned*)(cq+r->p.cq_off.head);
r->cq_tail=(unsigned*)(cq+r->p.cq_off.tail);
r->cq_mask=(unsigned*)(cq+r->p.cq_off.ring_mask);
r->cq_entries=(unsigned*)(cq+r->p.cq_off.ring_entries);
r->sqes=(struct io_uring_sqe*)r->sqes_map;
r->cqes=(struct io_uring_cqe*)(cq+r->p.cq_off.cqes);
return 0;
}
static inline int coli_uring_set_workers(ColiUring *r,unsigned workers){
unsigned limits[2]={workers,workers}; /* bounded, unbounded io-wq workers */
return (int)syscall(SYS_io_uring_register,r->fd,IORING_REGISTER_IOWQ_MAX_WORKERS,limits,2);
}
static inline int coli_uring_prep_read(ColiUring *r,int fd,void *buf,size_t len,
int64_t off,uint64_t user_data){
if(!len || len>UINT32_MAX){ errno=EINVAL; return -1; }
unsigned head=coli_uring_load_acquire(r->sq_head);
unsigned tail=__atomic_load_n(r->sq_tail,__ATOMIC_RELAXED);
if(tail-head>=*r->sq_entries){ errno=EAGAIN; return -1; }
unsigned idx=tail&*r->sq_mask;
struct io_uring_sqe *sqe=&r->sqes[idx];
memset(sqe,0,sizeof(*sqe));
sqe->opcode=IORING_OP_READ;
/* Cold regular-file reads are allowed to execute inline during
* io_uring_enter() unless forced async. That serializes the submitter on
* filesystems without native nonblocking buffered reads and destroys the
* intended I/O/compute overlap. io-wq gives the ring a real bounded worker
* pool while CQEs retain completion ordering/ownership here. */
sqe->flags=IOSQE_ASYNC;
sqe->fd=fd;
sqe->off=(uint64_t)off;
sqe->addr=(uint64_t)(uintptr_t)buf;
sqe->len=(uint32_t)len;
sqe->user_data=user_data;
r->sq_array[idx]=idx;
coli_uring_store_release(r->sq_tail,tail+1);
return 0;
}
static inline int coli_uring_enter(ColiUring *r,unsigned min_complete){
for(;;){
unsigned head=coli_uring_load_acquire(r->sq_head);
unsigned tail=coli_uring_load_acquire(r->sq_tail);
unsigned submit=tail-head;
unsigned flags=min_complete?IORING_ENTER_GETEVENTS:0;
int n=(int)syscall(SYS_io_uring_enter,r->fd,submit,min_complete,flags,NULL,0);
if(n>=0) return n;
if(errno!=EINTR) return -1;
}
}
static inline int coli_uring_peek(ColiUring *r,struct io_uring_cqe *out){
unsigned head=__atomic_load_n(r->cq_head,__ATOMIC_RELAXED);
unsigned tail=coli_uring_load_acquire(r->cq_tail);
if(head==tail) return 0;
*out=r->cqes[head&*r->cq_mask];
coli_uring_store_release(r->cq_head,head+1);
return 1;
}
#endif /* __linux__ */
#endif /* COLI_URING_H */
+65
View File
@@ -0,0 +1,65 @@
# CACHE_ROUTE — opt-in cache-aware MoE routing
**Default: OFF.** Stock full top-K router behavior unless `CACHE_ROUTE=1`.
Paper-style max-rank selection (arXiv:2412.00099): keep true top-`J` always;
fill remaining slots preferring experts already **resident** (pin LRU) that
still rank inside top-`M`.
This is **routing-side** (can change which experts run). Complementary to
**PILOT** (next-layer *prefetch* of weights; does not change expert IDs).
## Flags
| Env | Default | Meaning |
|-----|---------|---------|
| `CACHE_ROUTE` | `0` | `1` = max-rank cache-aware fill (pinLRU prefer) |
| `ROUTE_J` | `2` | Always take true top-J (even if uncached) |
| `ROUTE_M` | `12` | Max-rank window for resident preference |
| `ROUTE_P` | `0` | Optional cumulative mass window (`0` = use fixed M) |
| `ROUTE_ALPHA` | `1` | Scale gate mass of *substituted* experts before renorm (`1` = off) |
| `ROUTE_AGREE` | auto | Overlap% + KL vs true top-K; auto-on when `CACHE_ROUTE=1` |
## One-liners
```bash
# Stock full-K (leaderboard-comparable routing)
CACHE_ROUTE=0 ./coli chat
# Experimental CACHE_ROUTE
CACHE_ROUTE=1 ROUTE_J=2 ROUTE_M=12 ./coli chat
# Wider prefer window (more hit, more possible swap)
CACHE_ROUTE=1 ROUTE_J=2 ROUTE_M=16 ./coli chat
```
## Stats
Footer / serve `STAT` when enabled:
- `swap N%` / `swap_pct` — fraction of chosen slots not in true top-K
- `route_swaps` / `route_slots` — raw substitution counters
- `route_agree` — |chosen ∩ true top-K| / K
- `route_kl` — mass KL (true top-K vs chosen)
- `hit N%` — expert cache hit (disk residency)
## A/B vs PILOT
```bash
# A: cache-aware routing only
CACHE_ROUTE=1 PILOT=0 ...
# B: lookahead prefetch only (does not change expert IDs)
CACHE_ROUTE=0 PILOT=1 ...
# C: both
CACHE_ROUTE=1 PILOT=1 ...
```
## Scope of this PR
**Routing-only** + telemetry + this note. Does **not** require CUDA/fuse/device-tier
patches — CPU streaming + pin/LRU is enough to A/B the lever against PILOT / #119.
Treat as experimental until quality gates (e.g. `./coli bench`) pass; do not
default `CACHE_ROUTE=1`.
+9 -2
View File
@@ -39,20 +39,27 @@ Format: `VAR` — default — effect.
| `COLI_METAL_GEMM_MIN` | `16` | Minimum matmul rows to dispatch a GEMM to the GPU (below this, stays on CPU). |
| `COLI_METAL_SPIN` | off | Keep a GPU keep-alive spinner running (reduces dispatch latency; costs power). |
| `PIPE` | `0` (off) | Overlap expert disk-load with matmul via I/O worker threads. Byte-identical output; reorders I/O. `PIPE=1` opts in. |
| `PIPE_WORKERS` | `8` | Number of I/O worker threads when `PIPE=1`. Tune to your SSD (fewer avoids over-subscribing cores). |
| `PIPE_WORKERS` | `8` | Number of pthread loaders when `PIPE=1`, or the io-wq worker maximum per ring when `URING=1` (capped at 64). Tune to SSD queue depth and available cores. |
| `URING` | `0` (off) | Linux-only queued expert I/O. `URING=1` implies `PIPE=1`, forces cold reads through io-wq (`IOSQE_ASYNC`), replaces blocking loader pthreads and spin waits with batched SQEs/CQEs, and batches `PILOT_REAL` loads on a separate ring. Use `DIRECT=1` for cold NVMe to avoid page-cache copy/readahead limits. Fails clearly if the kernel denies io_uring; incompatible with `COLI_MMAP=1`. |
| `DIRECT` | `0` (off) | Use `O_DIRECT`/unbuffered reads for expert slabs. Helps sustained NVMe; keeps the zero-copy GPU path. |
| `COLI_NO_OMP_TUNE` | off | **Kill-switch** for the OpenMP hot-thread tuning (`OMP_WAIT_POLICY=active` spin + proc-bind). Set `=1` when the CPU is mostly waiting on the GPU (Metal) so spin doesn't steal the shared power budget. |
| `MLOCK` | `-1` (auto: on for macOS) | Wire the streamed expert cache into physical RAM (`mlock`) to dodge the memory compressor. `0` off, `1` force. |
| `CAP_RAISE` | `1` (on) | Let the engine raise the expert-cache cap above `topk` when RAM allows (bigger batches). `0` fixes the cap. |
| `PREFETCH` | `0` | Prefetch depth for streamed experts. |
| `COLI_MMAP` | `0` | `mmap` the weights instead of read()-ing into slabs. |
| `PIN` | unset | Path to a `.coli_usage` file; pins the hottest experts into a resident "hot store" at startup. |
| `PIN` | unset | Path to a `.coli_usage`/stats file; pins the hottest experts into a resident "hot store" at startup. **`PIN=auto`** seeds from the model dir's live `.coli_usage` (appended after every turn, so each restart's pin placement follows the accumulated real workload) with `stats.txt` as the fallback for a virgin model dir; neither present → no pin this run. |
| `PIN_GB` | `10.0` | Size budget (GB) for the pinned hot store when `PIN` is set. |
| `AUTOPIN` | `1` (on) | Auto-pin the hot store from usage history once ≥5000 selections are recorded. |
| `REPIN` | `0` (off) | Live re-pin the hot store every N emitted tokens (RFC). |
| `PILOT` | `0` (off) | Router-piloted cross-layer expert prefetch. |
| `PILOT_REAL` | `0` (off) | Value-preserving real cross-layer prefetch loads (`PILOT_REAL=1` opts in). |
| `PILOT_K` | `6` if `PILOT_REAL` else `8` | Number of experts the pilot prefetches per step. |
| `CACHE_ROUTE` | `0` (off) | Opt-in max-rank cache-aware MoE routing (pinLRU prefer within top-M). See [CACHE_ROUTE.md](CACHE_ROUTE.md). |
| `ROUTE_J` | `2` | Sacred top ranks always taken when `CACHE_ROUTE=1`. |
| `ROUTE_M` | `12` | Max-rank window for resident preference when `CACHE_ROUTE=1`. |
| `ROUTE_P` | `0` | Cumulative mass window for CACHE_ROUTE (`0` = fixed M). |
| `ROUTE_ALPHA` | `1` | Scale gate mass of substituted experts before renorm (`1` = off). |
| `ROUTE_AGREE` | auto | Overlap% + KL vs true top-K; auto-on when `CACHE_ROUTE=1`. |
| `ABSORB` | `-1` (auto: absorbed for S≤4) | MLA attention absorption mode. |
| `IDOT` | `1` | Integer dot-product kernel. `IDOT=0` uses exact f32 kernels (for A/B numerical checks). |
| `COLI_POLICY` | `quality` | Resource policy: `quality`, `balanced`, or `experimental-fast`. |
Binary file not shown.

After

Width:  |  Height:  |  Size: 421 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 545 KiB

+5 -1
View File
@@ -34,7 +34,11 @@ import { Brain } from "./Brain"
import { persistPublicSettings, stored } from "@/lib/storage"
import { cn } from "@/lib/utils"
const message = (role: ChatMessage["role"], content: string): ChatMessage => ({ id: crypto.randomUUID(), role, content })
const message = (role: ChatMessage["role"], content: string): ChatMessage => {
let id: string
try { id = crypto.randomUUID() } catch { id = 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, c => { const r = Math.random() * 16 | 0; return (c === 'x' ? r : (r & 0x3 | 0x8)).toString(16) }) }
return { id, role, content }
}
export default function App() {
// When the page is served by the engine itself (coli web), same-origin is the
+25 -4
View File
@@ -4,6 +4,7 @@ import { BrainCircuit, Flame, Layers } from "lucide-react"
import { endpoint } from "@/lib/api"
interface ExpertMap { rows: number; cols: number; map: string; hits: string; seq: number }
interface AtlasEntry { affinity: Record<string, number>; entropy: number; top: string; label: string }
const TIER_NAME = ["Disk", "RAM", "VRAM"]
const TIER_RGB: [number, number, number][] = [[58, 71, 80], [90, 155, 216], [78, 214, 165]]
@@ -26,11 +27,19 @@ export function Brain({ baseUrl, apiKey, connected }: { baseUrl: string; apiKey:
const wrapRef = useRef<HTMLDivElement>(null)
const [wrapSize, setWrapSize] = useState({ w: 1200, h: 700 })
const [data, setData] = useState<ExpertMap | null>(null)
const [atlas, setAtlas] = useState<Record<string, AtlasEntry> | null>(null)
const [tip, setTip] = useState<{ x: number; y: number; row: number; col: number; tier: number; heat: number } | null>(null)
const pulseRef = useRef<Float32Array | null>(null) // per-expert pulse intensity 0..1
const lastSeq = useRef(0)
const rafRef = useRef(0)
// load the expert atlas if published (measured topic affinity, #175)
useEffect(() => {
fetch("/experts.json").then(r => r.ok ? r.json() : null).then(d => {
if (d?.experts) setAtlas(d.experts)
}).catch(() => {})
}, [])
// track container size for responsive cell sizing
useEffect(() => {
const el = wrapRef.current
@@ -146,14 +155,26 @@ export function Brain({ baseUrl, apiKey, connected }: { baseUrl: string; apiKey:
<canvas ref={canvasRef} onMouseMove={onMove} onMouseLeave={() => setTip(null)} />
{!connected && <p className="runtime-unavailable">Connect to the engine to see the cortex.</p>}
</div>
{tip && data && (
{tip && data && (() => {
const isMtp = tip.row === data.rows - 1
const realLayer = isMtp ? 78 : tip.row + 3
const entry = atlas?.[`${realLayer}:${tip.col}`]
return (
<div className="brain-tip" style={{ left: tip.x + 14, top: tip.y + 14 }}>
<div className="brain-tip-title"><Layers className="size-3" /> Layer row {tip.row}{tip.row === data.rows - 1 ? " (MTP)" : ""} · Expert {tip.col}</div>
<div className="brain-tip-title"><Layers className="size-3" /> Layer {realLayer}{isMtp ? " (MTP)" : ""} · Expert {tip.col}</div>
<div>Tier: <strong style={{ color: ["#8b9aa3", "#5a9bd8", "#4ed6a5"][tip.tier] }}>{TIER_NAME[tip.tier]}</strong></div>
<div>Heat: <strong>{tip.heat === 0 ? "never routed" : `~2^${tip.heat} selections`}</strong></div>
<div className="brain-tip-role">{depthRole(tip.row, data.rows, tip.row === data.rows - 1)}</div>
{entry ? <>
<div className={entry.label.startsWith("specialist") ? "brain-tip-spec" : undefined}>
{entry.label.startsWith("specialist") ? `⭐ Specialist: ${entry.top}` : "Generalist"}
<small> (entropy {entry.entropy})</small>
</div>
<div className="brain-tip-aff">{Object.entries(entry.affinity).sort((a, b) => b[1] - a[1]).slice(0, 3)
.map(([c, p]) => `${c} ${Math.round(p * 100)}%`).join(" · ")}</div>
</> : <div className="brain-tip-role">{depthRole(tip.row, data.rows, isMtp)}</div>}
</div>
)}
)
})()}
</div>
)
}
+5
View File
@@ -146,3 +146,8 @@ button:focus-visible, input:focus-visible, textarea:focus-visible, select:focus-
.brain-legend { gap: 8px; font-size: 10px; }
.brain-tip { max-width: 220px; font-size: 10px; }
}
/* atlas hover extras */
.brain-tip-spec { color: var(--primary); font-weight: 700; }
.brain-tip-spec small, .brain-tip-aff { color: #8b9aa3; font-weight: 400; }
.brain-tip-aff { font-size: 10px; }