Merge branch 'main' into fix/windows-cuda-dll-build
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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 (30–40% 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 (30–40% 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 11–22).
|
||||
|
||||
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
@@ -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
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
@@ -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
@@ -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 */
|
||||
|
||||
@@ -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%
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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")
|
||||
@@ -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;
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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 */
|
||||
@@ -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);
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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.
@@ -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.
Binary file not shown.
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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",
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
Executable
+51
@@ -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) =="
|
||||
@@ -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]]:
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
@@ -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
|
||||
18–58 ≈ 0.19–0.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.
|
||||
@@ -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()
|
||||
@@ -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?"
|
||||
]
|
||||
}
|
||||
Executable
+82
@@ -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"
|
||||
@@ -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}%)")
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
@@ -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))
|
||||
|
||||
@@ -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
@@ -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 Conway–Sloane: 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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 */
|
||||
@@ -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 (pin∪LRU 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
@@ -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 (pin∪LRU 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
@@ -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
@@ -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>
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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; }
|
||||
|
||||
Reference in New Issue
Block a user