url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/dev' into pr437

Browse Source 
# Conflicts:
#	c/colibri.c
This commit is contained in:
JustVugg
2026-07-20 17:04:18 +02:00
parent 26bd8b403a 4b704823db
commit 3455eeea0a
97 changed files with 9171 additions and 2086 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/.gitignore
+1
View File
@@ -2,3 +2,4 @@
glm_tiny/
olmoe_hf/
olmoe_i4/
.build-config
c/Makefile
+71 -25
View File
@@ -56,7 +56,7 @@ 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 ->
# which unlocks the i8mm SMMLA int8/int4 dot kernels in colibri.c. ARCH unset ->
# no -mcpu, default build byte-identical. Apple clang knows apple-m4 / native.
ifneq ($(ARCH),)
CFLAGS += -mcpu=$(ARCH)
@@ -70,7 +70,7 @@ else ifneq ($(IS_WIN),)
# ARCH default = x86-64-v3 (portable binary with AVX2). For max speed on THIS
# machine use ARCH=native: on AVX-VNNI CPUs (Intel Alder Lake+, Meteor Lake+)
# it also unlocks the 128-bit VPDPBUSD int8/int4 dot kernel (dot_i8i8/dot_i4i8),
# which the x86-64-v3 baseline does not define. The #ifdef guards in glm.c mean
# which the x86-64-v3 baseline does not define. The #ifdef guards in colibri.c mean
# a v3 build simply compiles out the VNNI path - safe on any x86-64.
CC = gcc
ARCH ?= x86-64-v3
@@ -166,7 +166,7 @@ else
PYTHON ?= python3
endif
CUDA_OBJ =
TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE)
TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_sample_nan$(EXE) tests/test_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE) tests/test_logit_nan$(EXE)
ifneq (,$(LINUX))
TEST_BINS += tests/test_uring$(EXE)
endif
@@ -207,16 +207,31 @@ LDFLAGS += -framework Metal -framework Foundation -lc++
METAL_OBJ = backend_metal.o
endif
all: glm$(EXE)
all: colibri$(EXE)
# phony 'glm' → 'glm.exe' on Windows (so 'make glm' and 'coli build' work on every platform)
glm: glm$(EXE)
# phony targets — 'glm' kept for backward compatibility
colibri: colibri$(EXE)
glm: colibri$(EXE)
glm$(EXE): glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h $(CUDA_OBJ) $(METAL_OBJ)
$(CC) $(CFLAGS) glm.c $(CUDA_OBJ) $(METAL_OBJ) -o glm$(EXE) $(LDFLAGS)
# Config stamp: make only tracks file timestamps, not flag changes. Without this,
# `make colibri.exe CUDA_DLL=1` after a prior CPU-only build reports "up to date"
# and silently keeps the CPU-only binary (no CUDA loader) — a build that looks like
# it worked but isn't. We record the build-affecting flags in .build-config and
# rewrite it ONLY when they change (evaluated here at parse time, so the file's
# timestamp moves exactly when the config moves). The binary and CUDA/loader objects depend
# on it, so they relink on a config change and stay put otherwise. (#306)
BUILD_CONFIG := $(CC)|$(CFLAGS)|$(LDFLAGS)|CUDA=$(CUDA)|CUDA_DLL=$(CUDA_DLL)|ARCH=$(ARCH)|CUDA_ARCH=$(CUDA_ARCH)|METAL=$(METAL)
BUILD_CONFIG_OLD := $(shell cat .build-config 2>/dev/null)
ifneq "$(BUILD_CONFIG)" "$(BUILD_CONFIG_OLD)"
$(shell printf '%s' '$(BUILD_CONFIG)' > .build-config)
endif
.build-config: ;
colibri$(EXE): colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h quant.h sample.h kv_persist.h telemetry.h $(CUDA_OBJ) $(METAL_OBJ) .build-config
$(CC) $(CFLAGS) colibri.c $(CUDA_OBJ) $(METAL_OBJ) -o colibri$(EXE) $(LDFLAGS)
# Windows runtime loader object: resolves coli_cuda_* from coli_cuda.dll.
backend_loader.o: backend_loader.c backend_cuda.h compat.h
backend_loader.o: backend_loader.c backend_cuda.h compat.h .build-config
$(CC) $(CFLAGS) -c backend_loader.c -o $@
# Windows CUDA DLL: compile backend_cuda.cu with nvcc (+MSVC cl.exe as host
@@ -225,15 +240,15 @@ backend_loader.o: backend_loader.c backend_cuda.h compat.h
# "x64 Native Tools Command Prompt"). COLI_CUDA_BUILDING_DLL enables
# __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 "$(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; }
$(NVCC) $(NVCCFLAGS) -shared -DCOLI_CUDA_BUILDING_DLL \
"$(NVCC)" $(NVCCFLAGS) -shared -DCOLI_CUDA_BUILDING_DLL \
-L"$(CUDA_HOME)/lib/x64" -lcudart \
backend_cuda.cu -o coli_cuda.dll
backend_cuda.o: 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; }
$(NVCC) $(NVCCFLAGS) -c backend_cuda.cu -o $@
backend_cuda.o: backend_cuda.cu backend_cuda.h .build-config
@command -v "$(NVCC)" >/dev/null 2>&1 || { echo "nvcc not found: set CUDA_HOME or NVCC" >&2; exit 1; }
"$(NVCC)" $(NVCCFLAGS) -c backend_cuda.cu -o $@
backend_metal.o: backend_metal.mm backend_metal.h
$(METALXX) -c backend_metal.mm -o $@
@@ -243,13 +258,13 @@ metal-test: tests/test_backend_metal.mm backend_metal.mm backend_metal.h
./backend_metal_test
cuda-test: backend_cuda.cu backend_cuda.h tests/test_backend_cuda.cu
@command -v $(NVCC) >/dev/null 2>&1 || { echo "nvcc not found: set CUDA_HOME or NVCC" >&2; exit 1; }
$(NVCC) $(NVCCFLAGS) backend_cuda.cu tests/test_backend_cuda.cu -o backend_cuda_test$(EXE)
@command -v "$(NVCC)" >/dev/null 2>&1 || { echo "nvcc not found: set CUDA_HOME or NVCC" >&2; exit 1; }
"$(NVCC)" $(NVCCFLAGS) backend_cuda.cu tests/test_backend_cuda.cu -o backend_cuda_test$(EXE)
./backend_cuda_test$(EXE)
cuda-bench: backend_cuda.cu backend_cuda.h tests/bench_tensor_core.cu
@command -v $(NVCC) >/dev/null 2>&1 || { echo "nvcc not found: set CUDA_HOME or NVCC" >&2; exit 1; }
$(NVCC) $(NVCCFLAGS) backend_cuda.cu tests/bench_tensor_core.cu -o backend_cuda_bench$(EXE)
@command -v "$(NVCC)" >/dev/null 2>&1 || { echo "nvcc not found: set CUDA_HOME or NVCC" >&2; exit 1; }
"$(NVCC)" $(NVCCFLAGS) backend_cuda.cu tests/bench_tensor_core.cu -o backend_cuda_bench$(EXE)
./backend_cuda_bench$(EXE)
olmoe$(EXE): olmoe.c st.h json.h compat.h
@@ -271,7 +286,7 @@ PORTABLE_ARCH = native
endif
portable:
$(MAKE) glm$(EXE) ARCH=$(PORTABLE_ARCH)
$(MAKE) colibri$(EXE) ARCH=$(PORTABLE_ARCH)
iobench$(EXE): iobench.c compat.h
$(CC) $(CFLAGS) iobench.c -o iobench$(EXE) $(LDFLAGS)
@@ -279,6 +294,9 @@ iobench$(EXE): iobench.c compat.h
tests/test_json$(EXE): tests/test_json.c json.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_st_pread$(EXE): tests/test_st_pread.c st.h json.h compat.h
$(CC) $(CFLAGS) -DST_PREAD_CHUNK=7 $< -o $@ $(LDFLAGS)
tests/test_st$(EXE): tests/test_st.c st.h json.h compat.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
@@ -294,10 +312,30 @@ tests/test_schema_gbnf$(EXE): tests/test_schema_gbnf.c schema_gbnf.h grammar.h j
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 uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
tests/test_idot$(EXE): tests/test_idot.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.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
tests/test_i4_grouped$(EXE): tests/test_i4_grouped.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_stops$(EXE): tests/test_stops.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_topp$(EXE): tests/test_topp.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
# bench_topp is a microbenchmark (old qsort vs new heap partial-select, #335), NOT a test
# gate -- intentionally absent from TEST_BINS. Build on demand: make tests/bench_topp
tests/bench_topp$(EXE): tests/bench_topp.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_sample_nan$(EXE): tests/test_sample_nan.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_kv_alloc$(EXE): tests/test_kv_alloc.c colibri.c st.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_logit_nan$(EXE): tests/test_logit_nan.c colibri.c st.h uring.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
@@ -306,7 +344,15 @@ 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
tests/test_dsa_select$(EXE): tests/test_dsa_select.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
# bench_dsa_select is a microbenchmark (old qsort vs new quickselect partial-select, #356),
# NOT a test gate -- intentionally absent from TEST_BINS. Build on demand: make tests/bench_dsa_select
tests/bench_dsa_select$(EXE): tests/bench_dsa_select.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_uring$(EXE): tests/test_uring.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
test-c: $(TEST_BINS)
@@ -323,12 +369,12 @@ check:
$(MAKE) portable
$(MAKE) test
install: glm$(EXE) olmoe$(EXE)
install: colibri$(EXE) olmoe$(EXE)
$(INSTALL) -d $(DESTDIR)$(BINDIR)
$(INSTALL) -d $(DESTDIR)$(LIBEXECDIR)
$(INSTALL) -d $(DESTDIR)$(LIBEXECDIR)/tools
$(INSTALL) -m 755 coli $(DESTDIR)$(BINDIR)/coli
$(INSTALL) -m 755 glm$(EXE) $(DESTDIR)$(LIBEXECDIR)/glm$(EXE)
$(INSTALL) -m 755 colibri$(EXE) $(DESTDIR)$(LIBEXECDIR)/colibri$(EXE)
$(INSTALL) -m 755 olmoe$(EXE) $(DESTDIR)$(LIBEXECDIR)/olmoe$(EXE)
$(INSTALL) -m 644 resource_plan.py doctor.py openai_server.py $(DESTDIR)$(LIBEXECDIR)/
$(INSTALL) -m 644 tools/*.py $(DESTDIR)$(LIBEXECDIR)/tools/
@@ -342,4 +388,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 colibri glm cuda-test cuda-bench cuda-dll portable test-c test-python test check clean install uninstall bench
c/backend_cuda.cu
+143 -3
View File
@@ -8,12 +8,21 @@
#include <cstring>
#include <mutex>
struct RaggedKVEntry {
const void *key;
const float *host_l,*host_r;
float *latent,*rope;
int length,capacity,K,R;
};
struct ColiCudaTensor {
void *weights;
float *scales;
size_t weight_bytes;
int fmt, I, O, device;
int tracked;
RaggedKVEntry ragged[512];
int ragged_count;
};
typedef struct {
@@ -23,7 +32,7 @@ typedef struct {
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 *host_x,*host_y,*host_kv; size_t host_x_cap,host_y_cap,host_kv_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;
@@ -338,6 +347,49 @@ __global__ static void attention_absorb_batch_kernel(float *ctx,const float *q,
ctx[((size_t)s*H+h)*V+v]=a*(fmt?wscale[row]:1.f);}
}
/* Independent device-resident KV sequence per row. lengths selects the valid
* prefix; latent/rope point at paged caches updated by the host wrapper. */
__global__ static void attention_absorb_ragged_kernel(float *ctx,const float *q,
const float *const *latent,const float *const *rope,const int *lengths,
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=lengths[s],rbase=h*(Q+V);
if(s>=S||nt<1||nt>T)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);
const float *ls=latent[s],*rs=rope[s];
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=ls+(size_t)t*K;
const float *rt=rs+(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]*ls[(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);}
}
__global__ static void ragged_kv_append(float *const *latent,float *const *rope,
const float *packed,const int *old_len,const int *add,const int *offset,int K,int R){
int s=blockIdx.x,n=add[s],base=offset[s];
for(int i=threadIdx.x;i<n*(K+R);i+=blockDim.x){
if(i<n*K)latent[s][(size_t)old_len[s]*K+i]=packed[base+i];
else rope[s][(size_t)old_len[s]*R+i-n*K]=packed[base+i];
}
}
static int reserve(float **ptr, size_t *cap, size_t bytes) {
if (*cap >= bytes) return 1;
if (*ptr) cudaFree(*ptr);
@@ -406,16 +458,17 @@ extern "C" void coli_cuda_shutdown(void) {
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->host_kv) cudaFreeHost(ctx->host_kv);
if (ctx->stream) cudaStreamDestroy(ctx->stream);
if (ctx->group_desc) cudaFree(ctx->group_desc);
ctx->x = ctx->y = ctx->gate = ctx->up = nullptr;
ctx->qx=nullptr; ctx->qscale=nullptr;
ctx->aq=ctx->al=ctx->ar=ctx->ac=nullptr;
ctx->host_x=ctx->host_y=nullptr;ctx->stream=nullptr;
ctx->host_x=ctx->host_y=ctx->host_kv=nullptr;ctx->stream=nullptr;
ctx->x_cap = ctx->y_cap = ctx->gate_cap = ctx->up_cap = 0;
ctx->qx_cap=ctx->qscale_cap=0;
ctx->aq_cap=ctx->al_cap=ctx->ar_cap=ctx->ac_cap=0;
ctx->host_x_cap=ctx->host_y_cap=0;
ctx->host_x_cap=ctx->host_y_cap=ctx->host_kv_cap=0;
ctx->group_desc=nullptr; ctx->group_desc_cap=0;
}
g_nctx = 0;
@@ -770,6 +823,89 @@ extern "C" int coli_cuda_attention_project_batch(ColiCudaTensor *w,ColiCudaTenso
return attention_absorb_batch_run(w,proj,out,q,latent,rope,S,H,Q,R,V,K,T,scale);
}
extern "C" int coli_cuda_attention_project_ragged(ColiCudaTensor *w,ColiCudaTensor *proj,
float *out,const float *q,const void *const *keys,
const float *const *latent,const float *const *rope,
const int *lengths,int S,int H,int Q,int R,int V,int K,int T,float scale){
if(!w||!proj||!out||!q||!keys||!latent||!rope||!lengths||S<1||S>512||T<1||T>8192||
H<1||Q<1||R<1||V<1||K<1||K>512||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;
float **dl=(float**)std::malloc((size_t)S*sizeof(*dl));
float **dr=(float**)std::malloc((size_t)S*sizeof(*dr));
int *old=(int*)std::malloc((size_t)S*sizeof(*old));
int *add=(int*)std::malloc((size_t)S*sizeof(*add));
int *off=(int*)std::malloc((size_t)S*sizeof(*off));int packed_n=0;
if(!dl||!dr||!old||!add||!off){std::free(dl);std::free(dr);std::free(old);std::free(add);std::free(off);return 0;}
for(int s=0;s<S;s++){
if(!keys[s]||lengths[s]<1||lengths[s]>T){std::free(dl);std::free(dr);std::free(old);std::free(add);std::free(off);return 0;}
RaggedKVEntry *e=nullptr;
for(int i=0;i<w->ragged_count;i++)if(w->ragged[i].key==keys[s]){e=&w->ragged[i];break;}
if(!e){
if(w->ragged_count>=512){std::free(dl);std::free(dr);std::free(old);std::free(add);std::free(off);return 0;}
e=&w->ragged[w->ragged_count++];std::memset(e,0,sizeof(*e));e->key=keys[s];
}
if(e->K!=K||e->R!=R||e->host_l!=latent[s]||e->host_r!=rope[s]||lengths[s]<e->length){
if(e->latent)cudaFree(e->latent);if(e->rope)cudaFree(e->rope);
e->latent=e->rope=nullptr;e->length=e->capacity=0;
e->K=K;e->R=R;e->host_l=latent[s];e->host_r=rope[s];
}
if(lengths[s]>e->capacity){
int cap=(lengths[s]+63)&~63;float *nl=nullptr,*nr=nullptr;
if(!cuda_ok(cudaMalloc(&nl,(size_t)cap*K*sizeof(float)),"ragged KV latent page")||
!cuda_ok(cudaMalloc(&nr,(size_t)cap*R*sizeof(float)),"ragged KV rope page")){
if(nl)cudaFree(nl);if(nr)cudaFree(nr);std::free(dl);std::free(dr);std::free(old);std::free(add);std::free(off);return 0;
}
if(e->length){
cudaMemcpyAsync(nl,e->latent,(size_t)e->length*K*sizeof(float),cudaMemcpyDeviceToDevice,dc->stream);
cudaMemcpyAsync(nr,e->rope,(size_t)e->length*R*sizeof(float),cudaMemcpyDeviceToDevice,dc->stream);
}
if(e->latent)cudaFree(e->latent);if(e->rope)cudaFree(e->rope);
e->latent=nl;e->rope=nr;e->capacity=cap;
}
dl[s]=e->latent;dr[s]=e->rope;old[s]=e->length;add[s]=lengths[s]-e->length;
off[s]=packed_n;packed_n+=add[s]*(K+R);
}
size_t qb=(size_t)S*H*(Q+R)*sizeof(float);
size_t cb=(size_t)S*H*V*sizeof(float),ob=(size_t)S*proj->O*sizeof(float);
size_t pb=(size_t)packed_n*sizeof(float);
size_t desc=(size_t)S*(2*sizeof(float*)+4*sizeof(int));
int ok=reserve(&dc->aq,&dc->aq_cap,qb)&&reserve(&dc->ac,&dc->ac_cap,cb)&&
reserve(&dc->y,&dc->y_cap,ob)&&reserve_bytes(&dc->group_desc,&dc->group_desc_cap,desc)&&
(!pb||(reserve(&dc->al,&dc->al_cap,pb)&&reserve_pinned(&dc->host_kv,&dc->host_kv_cap,pb)));
char *db=(char*)dc->group_desc;float **ddl=(float**)db,**ddr=ddl+S;
int *dn=(int*)(ddr+S),*dold=dn+S,*dadd=dold+S,*doff=dadd+S;
if(ok&&pb){
for(int s=0;s<S;s++)if(add[s]){
float *p=dc->host_kv+off[s];
std::memcpy(p,latent[s]+(size_t)old[s]*K,(size_t)add[s]*K*sizeof(float));
std::memcpy(p+(size_t)add[s]*K,rope[s]+(size_t)old[s]*R,(size_t)add[s]*R*sizeof(float));
}
ok=cuda_ok(cudaMemcpyAsync(dc->al,dc->host_kv,pb,cudaMemcpyHostToDevice,dc->stream),"ragged KV append upload");
}
if(ok)ok=cuda_ok(cudaMemcpyAsync(dc->aq,q,qb,cudaMemcpyHostToDevice,dc->stream),"ragged q upload")&&
cuda_ok(cudaMemcpyAsync(ddl,dl,(size_t)S*sizeof(float*),cudaMemcpyHostToDevice,dc->stream),"ragged latent pointers")&&
cuda_ok(cudaMemcpyAsync(ddr,dr,(size_t)S*sizeof(float*),cudaMemcpyHostToDevice,dc->stream),"ragged rope pointers")&&
cuda_ok(cudaMemcpyAsync(dn,lengths,(size_t)S*sizeof(int),cudaMemcpyHostToDevice,dc->stream),"ragged lengths upload")&&
cuda_ok(cudaMemcpyAsync(dold,old,(size_t)S*sizeof(int),cudaMemcpyHostToDevice,dc->stream),"ragged old lengths")&&
cuda_ok(cudaMemcpyAsync(dadd,add,(size_t)S*sizeof(int),cudaMemcpyHostToDevice,dc->stream),"ragged append lengths")&&
cuda_ok(cudaMemcpyAsync(doff,off,(size_t)S*sizeof(int),cudaMemcpyHostToDevice,dc->stream),"ragged append offsets");
if(ok&&pb)ragged_kv_append<<<S,256,0,dc->stream>>>(ddl,ddr,dc->al,dold,dadd,doff,K,R);
if(ok)for(int s=0;s<S;s++){
for(int i=0;i<w->ragged_count;i++)if(w->ragged[i].key==keys[s]){w->ragged[i].length=lengths[s];break;}
}
std::free(dl);std::free(dr);std::free(old);std::free(add);std::free(off);if(!ok)return 0;
size_t shared=(size_t)(2*K+T+256)*sizeof(float);
attention_absorb_ragged_kernel<<<dim3(H,S),256,shared,dc->stream>>>(dc->ac,dc->aq,ddl,ddr,
dn,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale);
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));
return cuda_ok(cudaGetLastError(),"ragged attention launch")&&
cuda_ok(cudaMemcpyAsync(out,dc->y,ob,cudaMemcpyDeviceToHost,dc->stream),"ragged output download")&&
cuda_ok(cudaStreamSynchronize(dc->stream),"ragged attention synchronize");
}
extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) {
if (!tensor) return;
DeviceContext *ctx = find_ctx(tensor->device);
@@ -781,6 +917,10 @@ extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) {
}
if (tensor->weights) cudaFree(tensor->weights);
if (tensor->scales) cudaFree(tensor->scales);
for(int i=0;i<tensor->ragged_count;i++){
if(tensor->ragged[i].latent)cudaFree(tensor->ragged[i].latent);
if(tensor->ragged[i].rope)cudaFree(tensor->ragged[i].rope);
}
std::free(tensor);
}
c/backend_cuda.h
+6 -1
View File
@@ -14,6 +14,7 @@
#define COLI_CUDA_DLLEXPORT
#endif
#ifdef __cplusplus
extern "C" {
#endif
@@ -92,6 +93,11 @@ COLI_CUDA_DLLEXPORT int coli_cuda_attention_project_batch(ColiCudaTensor *kv_b,C
const float *rope,int S,int H,int Q,int R,
int V,int K,int T,float attention_scale);
COLI_CUDA_DLLEXPORT int coli_cuda_attention_project_ragged(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj,
float *out,const float *q,const void *const *keys,
const float *const *latent,const float *const *rope,
const int *lengths,int S,int H,int Q,int R,int V,int K,int max_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);
@@ -143,4 +149,3 @@ COLI_CUDA_DLLEXPORT int coli_cuda_pipe_sync(int device);
#endif
#endif
c/backend_loader.c
+15
View File
@@ -61,6 +61,10 @@ typedef int (*fn_attention_absorb_batch)(ColiCudaTensor *kv_b,float *ctx,const f
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_ragged)(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj,
float *out,const float *q,const void *const *keys,
const float *const *latent,const float *const *rope,
const int *lengths,int S,int H,int Q,int R,int V,int K,int max_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);
@@ -107,6 +111,7 @@ static struct {
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_ragged attention_project_ragged;
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;
@@ -200,6 +205,7 @@ static int coli_cuda_load(void){
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_ragged, fn_attention_project_ragged)
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)
@@ -342,6 +348,15 @@ int coli_cuda_attention_project_batch(ColiCudaTensor *kv_b,ColiCudaTensor *o_pro
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_ragged(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj,
float *out,const float *q,const void *const *keys,
const float *const *latent,const float *const *rope,
const int *lengths,int S,int H,int Q,int R,int V,int K,int max_t,float attention_scale){
if(!coli_cuda_load()) return 0;
return g_cuda.attention_project_ragged(kv_b,o_proj,out,q,keys,latent,rope,lengths,
S,H,Q,R,V,K,max_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);
c/coli
+229 -20
View File
@@ -20,7 +20,7 @@ Configuration through environment variables or flags (also valid after the subco
--topp P adaptive expert top-p --topk N fixed top-k
--ngen N maximum response tokens --cap N cache slots/layer
"""
import os, sys, subprocess, argparse, json, time, signal, shutil, threading, re, codecs, tempfile, textwrap
import os, sys, subprocess, argparse, json, time, signal, shutil, threading, re, codecs, tempfile, textwrap, struct
# The engine mmaps every shard (144+ files); macOS default RLIMIT_NOFILE is 256.
if sys.platform != "win32":
@@ -40,6 +40,8 @@ if sys.platform == "win32":
except (AttributeError, OSError): pass
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from version import __version__ as _version
# Run-in-place (source checkout, "cd c && ./coli ..."): the engine, the
# support modules (resource_plan.py, doctor.py, openai_server.py) and
@@ -52,16 +54,20 @@ HERE = os.path.dirname(os.path.abspath(__file__))
# guess is right (e.g. a custom packaging layout).
_EXE = ".exe" if sys.platform == "win32" else ""
_LIBEXEC = os.path.join(os.path.dirname(HERE), "libexec", "colibri")
_here_colibri = os.path.join(HERE, "colibri" + _EXE)
_here_glm = os.path.join(HERE, "glm" + _EXE)
if os.environ.get("COLI_ENGINE"):
GLM = os.environ["COLI_ENGINE"]
TOOLS = os.path.join(os.path.dirname(GLM), "tools")
elif os.path.exists(_here_colibri):
GLM = _here_colibri
TOOLS = os.path.join(HERE, "tools")
elif os.path.exists(_here_glm):
GLM = _here_glm
TOOLS = os.path.join(HERE, "tools")
else:
GLM = os.path.join(_LIBEXEC, "glm" + _EXE)
GLM = os.path.join(_LIBEXEC, "colibri" + _EXE)
TOOLS = os.path.join(_LIBEXEC, "tools")
sys.path.insert(0, _LIBEXEC) # so `import resource_plan`, `doctor`, `openai_server` still resolve
@@ -115,7 +121,7 @@ def sprite_lines():
def banner(sub=""):
sp=sprite_lines()
txt=[
f"{C.teal}{C.b}colibrì{C.r} {C.dim}v1.0{C.r}",
f"{C.teal}{C.b}colibrì{C.r} {C.dim}v{_version}{C.r}",
f"{C.dim}tiny engine, immense model{C.r}",
f"{C.gray}GLM-5.2 · 744B MoE · int4 · streaming CPU{C.r}",
f"{C.dgray}{sub}{C.r}" if sub else "",
@@ -141,12 +147,24 @@ def need_model(model):
sys.exit(f"{C.yel}engine is not built.{C.r} Run: coli build")
def cuda_binary():
if not os.path.exists(GLM) or sys.platform != "linux": return False
try:
linked=subprocess.run(["ldd",GLM],capture_output=True,text=True,timeout=3)
return any("libcudart" in line and "not found" not in line
for line in linked.stdout.splitlines())
except (OSError,subprocess.SubprocessError): return False
if not os.path.exists(GLM): return False
if sys.platform == "linux":
try:
linked=subprocess.run(["ldd",GLM],capture_output=True,text=True,timeout=3)
return any("libcudart" in line and "not found" not in line
for line in linked.stdout.splitlines())
except (OSError,subprocess.SubprocessError): return False
if sys.platform == "win32":
# Windows CUDA_DLL=1 builds never link libcudart directly: glm.exe loads
# coli_cuda.dll at runtime via LoadLibrary (backend_loader.c), so there's no
# import-table entry for ldd/dumpbin to see. Detect the COLI_CUDA build via a
# marker string baked into glm.c's #ifdef COLI_CUDA block instead, and require
# coli_cuda.dll to actually sit next to glm.exe (else CUDA init fails at startup).
try:
with open(GLM,"rb") as f: built=b"[CUDA] mode: routed experts" in f.read()
except OSError: return False
return built and os.path.exists(os.path.join(os.path.dirname(GLM),"coli_cuda.dll"))
return False
def resource_request(a, env):
ctx=a.ctx or int(env.get("CTX",4096))
@@ -227,13 +245,13 @@ def env_for(a):
e["COLI_CUDA"]="0"; e.pop("CUDA_EXPERT_GB",None); e.pop("CUDA_DENSE",None)
else:
if not cuda_binary():
sys.exit(f"{C.yel}--gpu needs the CUDA build:{C.r} make glm CUDA=1 (this binary is CPU-only)")
sys.exit(f"{C.yel}--gpu needs the CUDA build:{C.r} make colibri CUDA=1 (this binary is CPU-only)")
e["COLI_CUDA"]="1"
if a.gpu!="auto": e["COLI_GPUS"]=a.gpu
e.setdefault("CUDA_DENSE","1")
if a.vram and a.gpu!="none":
if not cuda_binary():
sys.exit(f"{C.yel}--vram needs the CUDA build:{C.r} make glm CUDA=1 (this binary is CPU-only)")
sys.exit(f"{C.yel}--vram needs the CUDA build:{C.r} make colibri CUDA=1 (this binary is CPU-only)")
e["COLI_CUDA"]="1"; e["CUDA_EXPERT_GB"]=str(a.vram)
return e
@@ -407,8 +425,8 @@ def cmd_build(a):
banner("build")
if not os.path.exists(os.path.join(HERE, "Makefile")):
sys.exit(f"{C.yel}coli build{C.r} only works from a source checkout (this is an installed copy).\n"
f" Clone https://github.com/JustVugg/colibri and run ./setup.sh, or make -C c glm.")
sys.exit(subprocess.call(["make","-C",HERE,"glm"]))
f" Clone https://github.com/JustVugg/colibri and run ./setup.sh, or make -C c colibri.")
sys.exit(subprocess.call(["make","-C",HERE,"colibri"]))
def cmd_info(a):
banner("info")
@@ -480,13 +498,141 @@ def cmd_run(a):
need_model(a.model)
prompt=" ".join(a.prompt) if a.prompt else sys.exit('usage: coli run "your prompt"')
banner("run")
# template ufficiale GLM-5.2: niente \n dopo i ruoli; <think></think> = risposta diretta (nothink)
e=env_for(a); e["PROMPT"]=f"[gMASK]<sop><|user|>{prompt}<|assistant|><think></think>"
# template ufficiale GLM-5.2: niente \n dopo i ruoli; <think></think> = risposta diretta (nothink).
# THINK=1 lascia <think> aperto, stessa convenzione del serve mode (glm.c). EN: THINK=1 leaves
# <think> open so the engine emits its reasoning block; the default stays nothink.
tk="<think>" if os.environ.get("THINK","0")=="1" else "<think></think>"
e=env_for(a); e["PROMPT"]=f"[gMASK]<sop><|user|>{prompt}<|assistant|>{tk}"
sys.exit(subprocess.call([GLM, str(a.cap)], env=e))
def server_probe(base, api_key=None, timeout=1.5):
"""Is a coli serve alive at `base`? Returns its model_id, or None.
Probes /health then /v1/models — both cheap, neither touches the engine."""
import urllib.request, urllib.error
def get(path):
req=urllib.request.Request(base.rstrip("/")+path)
if api_key: req.add_header("Authorization", f"Bearer {api_key}")
with urllib.request.urlopen(req, timeout=timeout) as r:
return json.loads(r.read().decode("utf-8","replace"))
try:
if get("/health").get("status")!="ok": return None
data=get("/v1/models").get("data") or []
return data[0]["id"] if data else None
except Exception:
return None
def chat_attached(a, base, model_id):
"""The chat REPL over HTTP against a running `coli serve`.
Why this exists (the cold-chat cost, measured): spawning a private engine
pays 34-136 s of resident load on EVERY start, and begins with an empty
expert cache — hit rate 4% cold vs 55% warm, a ~10x on early decode. A
resident server pays load once and keeps the LRU warm across sessions;
its KV slots reuse the conversation prefix, so a continued chat skips
re-prefill too. The engine byte-protocol stays untouched — this is plain
OpenAI SSE over localhost, stdlib only."""
import urllib.request
print(f" {C.grn}✦ attached{C.r} {C.dim}to {base} · model {model_id} · the engine stays warm after you quit{C.r}")
print(f" {C.dim}type and press Enter · Ctrl-C stops the answer · :reset starts a new conversation · :q exits{C.r}\n")
msgs=[]
w=term_w()-4
while True:
if TTY:
print(f" {C.dgray}╭{'─'*w}╮{C.r}")
try: msg=input(f" {C.dgray}│{C.r} {C.teal}{C.b}›{C.r} ")
except EOFError: print(); break
print(f" {C.dgray}╰{'─'*w}╯{C.r}")
else:
try: msg=input()
except EOFError: break
msg=msg.strip()
if msg in (":q",":quit","exit"): break
if not msg: continue
if msg==":reset": msgs=[]; print(f" {C.dim}✦ new conversation{C.r}\n"); continue
msgs.append({"role":"user","content":msg})
body=json.dumps({"model":model_id,"messages":msgs,"stream":True,
"max_tokens":a.ngen}).encode()
req=urllib.request.Request(base.rstrip("/")+"/v1/chat/completions", data=body,
headers={"Content-Type":"application/json"})
if a.api_key: req.add_header("Authorization", f"Bearer {a.api_key}")
print(f"\n {C.teal}◆ colibrì{C.r}")
sp=Spinner("thinking…"); sp.start()
md=MDStream(" "); reply=[]; first=True; t0=time.time(); interrupted=False
try:
with urllib.request.urlopen(req) as r:
for raw in r:
line=raw.decode("utf-8","replace").strip()
if not line.startswith("data: "): continue
data=line[6:]
if data=="[DONE]": break
try: ev=json.loads(data)
except ValueError: continue
for ch in ev.get("choices",[]):
d=ch.get("delta",{})
txt=d.get("content")
if not txt: continue # ping/ruolo/reasoning: non è testo
if first: sp.stop(); first=False
md.feed(txt); reply.append(txt)
except KeyboardInterrupt:
interrupted=True # il server annulla la richiesta alla disconnessione
except OSError as e:
sp.stop()
print(f"\n {C.yel}[server unreachable: {e}]{C.r}"); break
sp.stop(); md.close()
if reply: msgs.append({"role":"assistant","content":"".join(reply)})
else: msgs.pop() # turno vuoto: non sporcare la history
el=time.time()-t0
note=" · ⏹ interrupted" if interrupted else ""
print(f"\r {C.dgray}└─ ~{len(''.join(reply))//4} tok · {el:.0f}s{note}{C.r}\n")
print(f" {C.dim}goodbye — the engine keeps running for the next chat 🐦{C.r}")
def kv_resume_notice(model_dir):
"""SERVE mode silently resumes .coli_kv from disk (glm.c kv_disk_load): a chat
started today continues a conversation from days ago, with `first=0` so the
turn is appended WITHOUT the [gMASK]<sop> prefix. The engine does announce it
on stderr — but nothing here ever shows that: the drain thread's
p.stderr.read() blocks until EOF, so on a healthy start errlog is still empty
when the status lines are printed. The warning only appeared once the engine
DIED, which is exactly when it no longer mattered.
Measured cost of the silence: a chat inherited 670 tokens of an old Italian
session ("il mio numero preferito e 7, ricordalo!"). Every later reply came
back in Italian, and "explain fibonacci in short" was answered about the
number 7 — the model was being coherent with a context nobody could see, and
it read as a quantization bug for a day.
So say it here, in Python, from the file itself: no pipe, no thread, no
Windows deadlock risk (see the stderr comment below)."""
p=os.path.join(model_dir, ".coli_kv")
try:
with open(p,"rb") as f:
if f.read(8)!=b"COLIKV1\0": return
h=struct.unpack("<8i", f.read(32))
n=h[6]
if n<1: return
age=time.time()-os.path.getmtime(p)
when=f"{age/86400:.0f}d ago" if age>86400 else f"{age/3600:.0f}h ago" if age>3600 else "just now"
print(f" {C.yel}↺ resuming a saved conversation: {n} tokens, last written {when}{C.r}")
print(f" {C.dgray} it steers tone, language and topic. :reset clears it · "
f"KVSAVE=0 disables saving · delete {p} to start clean{C.r}")
except (OSError, struct.error): pass
def cmd_chat(a):
# ATTACH: a running `coli serve` beats a private engine every time — the load
# (34-136 s) and the cache warmth survive between sessions. Explicit --attach
# wins; otherwise probe localhost quietly and use it if it's there. --no-attach
# forces the old behaviour. The probe costs ~1 ms when nothing is listening.
if not getattr(a,"no_attach",False):
base=getattr(a,"attach",None) or "http://127.0.0.1:8000"
mid=server_probe(base, getattr(a,"api_key",None))
if mid:
banner(f"chat · {mid} · attached")
chat_attached(a, base, mid); return
if getattr(a,"attach",None):
sys.exit(f"--attach: no coli serve answering at {base} (start one with: coli serve --model <dir>)")
need_model(a.model)
banner(f"chat · {os.path.basename(a.model)} · ram {a.ram or '-'}GB · topp {a.topp or 'off'}")
kv_resume_notice(a.model)
errlog=tempfile.NamedTemporaryFile(mode="w+", suffix=".log", delete=False)
e=env_for(a); e["SERVE"]="1"
# stderr -> PIPE, NOT stderr=errlog (file). On Windows/MinGW, pointing the
@@ -619,12 +765,65 @@ def cmd_chat(a):
except Exception: pass
print(f" {C.teal}goodbye{C.r} {C.dim}— the hummingbird returns to its nest{C.r} 🐦\n")
def serve_pidfile(port): return os.path.join(tempfile.gettempdir(), f"coli-serve-{port}.pid")
def cmd_serve(a):
need_model(a.model)
# pidfile: cosi' `coli stop` spegne tutto con un comando, senza pkill a mano.
# EN: pidfile so `coli stop` can shut everything down without manual pkill.
try:
with open(serve_pidfile(a.port),"w") as f: f.write(f"{os.getpid()} {a.model}\n")
except OSError: pass
from openai_server import serve
serve(a.model, a.host, a.port, a.model_id, a.api_key,
a.cap,a.ngen,GLM,env_for(a),a.cors_origin,
a.max_queue,a.queue_timeout,a.kv_slots)
try:
serve(a.model, a.host, a.port, a.model_id, a.api_key,
a.cap,a.ngen,GLM,env_for(a),a.cors_origin,
a.max_queue,a.queue_timeout,a.kv_slots)
finally:
try: os.unlink(serve_pidfile(a.port))
except OSError: pass
def cmd_stop(a):
"""Shut down a running `coli serve` AND its engine — one command, no pkill.
The engine re-execs itself for OMP tuning, so its process is named `exe`,
not `glm`: every `pkill -x glm` in history silently killed nothing (that is
how two 17+5 GB ghost engines OOM'd this box on 2026-07-16). This finds the
real processes: the pidfile first, then /proc by cmdline/environ — only
processes that are demonstrably ours (SERVE=1 + our SNAP, or `coli serve`
in the command line)."""
banner("stop")
targets=[] # (pid, descrizione)
pf=serve_pidfile(a.port)
try:
pid=int(open(pf).read().split()[0])
os.kill(pid,0); targets.append((pid,f"coli serve (pidfile, port {a.port})"))
except (OSError,ValueError,IndexError): pass
for pd in os.listdir("/proc"):
if not pd.isdigit(): continue
pid=int(pd)
try:
cmd=open(f"/proc/{pd}/cmdline","rb").read().replace(b"\0",b" ").decode("utf-8","replace")
if "coli" in cmd and " serve" in cmd and pid!=os.getpid():
if not any(p==pid for p,_ in targets): targets.append((pid,"coli serve (cmdline)"))
comm=open(f"/proc/{pd}/comm").read().strip()
if comm in ("colibri","glm","exe","olmoe"):
env=open(f"/proc/{pd}/environ","rb").read().replace(b"\0",b"\n").decode("utf-8","replace")
if "SERVE=1" in env: targets.append((pid,f"engine `{comm}` (SERVE=1)"))
except (OSError,PermissionError): continue
if not targets:
print(f" nothing running — no serve on port {a.port}, no SERVE engines"); return
for pid,desc in targets: print(f" {'would stop' if a.dry_run else 'stopping'} {pid}: {desc}")
if a.dry_run: return
for pid,_ in targets:
try: os.kill(pid, signal.SIGTERM)
except OSError: pass
time.sleep(2.0)
for pid,_ in targets:
try: os.kill(pid, signal.SIGKILL); print(f" {pid}: forced (SIGKILL)")
except OSError: pass # gia' morto: bene
try: os.unlink(pf)
except OSError: pass
print(f" {C.grn}✓ stopped{C.r} — RAM released")
def cmd_web(a):
"""serve + open the dashboard in the browser once the API answers."""
@@ -712,6 +911,7 @@ def main():
common.add_argument("--topp", type=float, default=0); common.add_argument("--topk", type=int, default=0)
common.add_argument("--temp", type=float, default=None) # temperatura token (0=greedy, default 1.0+nucleus .95)
ap=argparse.ArgumentParser(prog="coli", parents=[common], description="colibrì — run GLM-5.2 locally")
ap.add_argument("--version", action="version", version=f"colibrì {_version}")
sub=ap.add_subparsers(dest="cmd")
sub.add_parser("build", parents=[common]); sub.add_parser("info", parents=[common])
pp=sub.add_parser("plan",parents=[common])
@@ -719,7 +919,14 @@ def main():
pd=sub.add_parser("doctor",parents=[common])
pd.add_argument("--json",action="store_true",help="emit a versioned JSON report")
pr=sub.add_parser("run", parents=[common]); pr.add_argument("prompt", nargs="*")
sub.add_parser("chat", parents=[common])
pc=sub.add_parser("chat", parents=[common])
pc.add_argument("--attach", nargs="?", const="http://127.0.0.1:8000", default=None,
help="chat against a running `coli serve` instead of spawning an engine "
"(keeps the model loaded and the expert cache warm across chat sessions). "
"Bare --attach probes localhost:8000.")
pc.add_argument("--no-attach", action="store_true",
help="never auto-attach, always spawn a private engine")
pc.add_argument("--api-key", default=os.environ.get("COLI_API_KEY"))
ps=sub.add_parser("serve", parents=[common])
ps.add_argument("--host",default="127.0.0.1"); ps.add_argument("--port",type=int,default=8000)
ps.add_argument("--model-id",default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))
@@ -728,6 +935,8 @@ def main():
ps.add_argument("--max-queue",type=int,default=int(os.environ.get("COLI_MAX_QUEUE","8")))
ps.add_argument("--queue-timeout",type=float,default=float(os.environ.get("COLI_QUEUE_TIMEOUT","300")))
ps.add_argument("--kv-slots",type=int,default=int(os.environ.get("COLI_KV_SLOTS","1")))
pst=sub.add_parser("stop", parents=[common], help="shut down a running coli serve and its engine")
pst.add_argument("--port",type=int,default=8000); pst.add_argument("--dry-run",action="store_true")
pw=sub.add_parser("web", parents=[common], help="serve + open the dashboard in a browser")
for arg,kw in (("--host",dict(default="127.0.0.1")),("--port",dict(type=int,default=8000)),
("--model-id",dict(default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))),
@@ -751,7 +960,7 @@ def main():
pc.add_argument("--no-mtp",action="store_true",help="skip the MTP head (no speculative drafts)")
a=ap.parse_args()
handler={"build":cmd_build,"info":cmd_info,"plan":cmd_plan,"doctor":cmd_doctor,
"run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"bench":cmd_bench,
"run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"stop":cmd_stop,"bench":cmd_bench,
"convert":cmd_convert,"web":cmd_web}.get(a.cmd)
if handler: sys.exit(handler(a) or 0)
banner(); print(__doc__)
c/glm.c → c/colibri.c
+856 -1255
View File
File diff suppressed because it is too large Load Diff
c/compat.h
+55 -1
View File
@@ -80,6 +80,7 @@ static inline int compat_open_direct(const char *path){
#endif
#include <windows.h>
#include <io.h>
#include <direct.h> /* _mkdir (for the mkdtemp shim below) */
#include <process.h>
#include <malloc.h>
#include <fcntl.h>
@@ -143,6 +144,12 @@ static inline int compat_fadvise(int fd, off_t off, off_t len, int advice){
* Thread-safe (no shared seek position). Gestisce offset >4 GB e chunking
* per letture >2 GB (anche se i tensori individuali sono nell'ordine dei
* MB-centinaia di MB, il wrapper e' robusto per ogni taglia). */
/* Ultimo GetLastError() di una ReadFile fallita, per thread: il chiamante
* (pread_full in glm.c) lo stampa accanto a strerror. Senza questo, OGNI
* fallimento Windows collassa in "EIO -> Input/output error" e la diagnosi
* dal campo diventa un tirare a indovinare (#307: tre giri di ipotesi tra
* tre persone perche' il codice vero non compariva da nessuna parte). */
static __thread DWORD compat_pread_lasterr __attribute__((unused));
static inline ssize_t compat_pread(int fd, void *buf, size_t n, off_t off){
intptr_t osfh = _get_osfhandle(fd);
if(osfh == -1 || osfh == -2){ errno = EBADF; return -1; }
@@ -158,6 +165,7 @@ static inline ssize_t compat_pread(int fd, void *buf, size_t n, off_t off){
if(!ReadFile(h, (char*)buf + total, chunk32, &rd, &ov)){
DWORD err = GetLastError();
if(err == ERROR_HANDLE_EOF) break; /* past EOF → return bytes read (0 if none, matching POSIX pread) */
compat_pread_lasterr = err; /* preserva il codice VERO per il report (#307) */
if(err == ERROR_INVALID_HANDLE || err == ERROR_INVALID_FUNCTION) errno = EBADF;
else errno = EIO;
return -1;
@@ -237,7 +245,9 @@ static inline int compat_rename(const char *old, const char *new){
/* --- rss_gb: getrusage -> GetProcessMemoryInfo ---
* ru_maxrss in KB (come Linux): rss_gb() divide per 1e6 → GB corretti. */
#include <psapi.h>
#pragma comment(lib, "psapi.lib")
#ifdef _MSC_VER
#pragma comment(lib, "psapi.lib") /* MSVC: link psapi; MinGW/GCC uses -lpsapi */
#endif
struct rusage { long ru_maxrss; };
#define RUSAGE_SELF 0
static inline int getrusage(int who, struct rusage *r){
@@ -304,8 +314,52 @@ static inline int compat_setenv(const char *name, const char *value, int overwri
}
#define setenv(name,value,overwrite) compat_setenv(name,value,overwrite)
/* --- getenv_utf8: read an env var as UTF-8, not through the ANSI codepage ---
* Plain getenv()/_environ are populated by the CRT from the ANSI-codepage view
* of the process environment block, not UTF-8. A parent that hands the child a
* Unicode value via CreateProcessW's wide env block (e.g. Python's subprocess
* module, which coli uses to pass the chat prompt) round-trips correctly only
* through GetEnvironmentVariableW; going through narrow getenv() re-encodes it
* via CP_ACP first, so any non-ASCII prompt text (Cyrillic, CJK, ...) comes out
* corrupted before the byte-level tokenizer ever sees it. Read the wide value
* directly and convert straight to UTF-8, bypassing the ANSI codepage entirely.
* Returned buffer is intentionally leaked: called a handful of times at
* startup, lives for the process. */
static inline const char *compat_getenv_utf8(const char *name){
wchar_t wname[64];
if(MultiByteToWideChar(CP_UTF8, 0, name, -1, wname, 64) <= 0) return getenv(name);
DWORD need = GetEnvironmentVariableW(wname, NULL, 0);
if(!need) return NULL;
wchar_t *wval = (wchar_t*)malloc(need * sizeof(wchar_t));
if(!wval) return NULL;
GetEnvironmentVariableW(wname, wval, need);
int blen = WideCharToMultiByte(CP_UTF8, 0, wval, -1, NULL, 0, NULL, NULL);
char *val = blen>0 ? (char*)malloc((size_t)blen) : NULL;
if(val) WideCharToMultiByte(CP_UTF8, 0, wval, -1, val, blen, NULL, NULL);
free(wval);
return val;
}
#define getenv_utf8(name) compat_getenv_utf8(name)
/* --- mkdtemp -> _mktemp + _mkdir (POSIX mkdtemp assente su Windows) ---
* Test binaries (test_stops.c) create a scratch dir in the CWD via a
* "name_XXXXXX" template; POSIX mkdtemp fills the X's and mkdirs 0700. The
* Windows CRT has _mktemp (in-place, same XXXXXX contract) so we compose it.
* Returns the template pointer on success, NULL on failure — matching POSIX. */
static inline char *compat_mkdtemp(char *tmpl){
if(!tmpl) return NULL;
if(!_mktemp(tmpl)) return NULL; /* fills the trailing X's in place */
if(_mkdir(tmpl) != 0) return NULL; /* EEXIST is impossible post-_mktemp */
return tmpl;
}
#define mkdtemp(tmpl) compat_mkdtemp(tmpl)
#endif /* _WIN32 */
#ifndef getenv_utf8
#define getenv_utf8(name) getenv(name)
#endif
/* --- compat_aligned_free su piattaforme diverse da Windows ---
* Su Linux/macOS, posix_memalign usa free() normale. */
#ifndef compat_aligned_free
c/doctor.py
+26 -9
View File
@@ -19,16 +19,33 @@ def _check(identifier, status, summary, **details):
def cuda_linkage(engine_path):
"""Return CUDA linkage state without loading the executable or CUDA runtime."""
if not Path(engine_path).is_file() or os.name != "posix":
engine = Path(engine_path)
if not engine.is_file():
return {"linked": False, "missing": False}
try:
result = subprocess.run(["ldd", str(engine_path)], capture_output=True, text=True,
timeout=3, check=False)
except (OSError, subprocess.SubprocessError):
return {"linked": False, "missing": False}
lines = [line for line in result.stdout.splitlines() if "libcudart" in line]
return {"linked": any("not found" not in line for line in lines),
"missing": any("not found" in line for line in lines)}
if os.name == "posix":
try:
result = subprocess.run(["ldd", str(engine)], capture_output=True, text=True,
timeout=3, check=False)
except (OSError, subprocess.SubprocessError):
return {"linked": False, "missing": False}
lines = [line for line in result.stdout.splitlines() if "libcudart" in line]
return {"linked": any("not found" not in line for line in lines),
"missing": any("not found" in line for line in lines)}
if sys.platform == "win32":
# Windows CUDA_DLL=1 builds never link libcudart directly: glm.exe loads
# coli_cuda.dll at runtime via LoadLibrary (backend_loader.c), so there's no
# import-table entry for ldd/dumpbin to see. Detect the COLI_CUDA build via a
# marker string baked into glm.c's #ifdef COLI_CUDA block instead, and require
# coli_cuda.dll to actually sit next to glm.exe (else CUDA init fails at startup).
try:
built = b"[CUDA] mode: routed experts" in engine.read_bytes()
except OSError:
return {"linked": False, "missing": False}
if not built:
return {"linked": False, "missing": False}
dll_present = (engine.parent / "coli_cuda.dll").is_file()
return {"linked": dll_present, "missing": not dll_present}
return {"linked": False, "missing": False}
def run_doctor(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0, *,
c/kv_persist.h
+121
View File
@@ -0,0 +1,121 @@
/* kv_persist.h — .coli_kv on-disk KV cache persistence.
* Conversations reopen warm across engine restarts: the compressed MLA KV-cache
* is appended incrementally after every turn, crash-safe (nrec written last).
* Include after Model/KVState/Cfg are defined; requires now_s() and g_draft. */
#ifndef KV_PERSIST_H
#define KV_PERSIST_H
static int g_kvsave=1;
#define KV_MAGIC "COLIKV1\0"
static void kv_hdr(Model *m, int32_t *h, int nrec){
Cfg *c=&m->c; int nic=0;
for(int i=0;i<c->n_layers;i++) if(m->Ic && m->Ic[i]) nic++;
h[0]=c->n_layers; h[1]=c->kv_lora; h[2]=c->qk_rope;
h[3]=m->has_dsa?c->index_hd:0; h[4]=nic; h[5]=c->vocab; h[6]=nrec; h[7]=0;
}
static int64_t kv_rec_bytes(Model *m){
Cfg *c=&m->c;
int64_t rec = 4 + (int64_t)c->n_layers*(c->kv_lora+c->qk_rope)*4;
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i]) rec+=(int64_t)c->index_hd*4;
return rec;
}
static int kv_disk_open(Model *m){
KVState *k=m->kv;
if(k->disk_fp) return 1;
k->disk_fp=fopen(k->disk_path,"r+b");
if(!k->disk_fp){
k->disk_fp=fopen(k->disk_path,"wb");
if(!k->disk_fp) return 0;
int32_t h[8]; kv_hdr(m,h,0);
fwrite(KV_MAGIC,1,8,k->disk_fp); fwrite(h,4,8,k->disk_fp);
fflush(k->disk_fp);
fclose(k->disk_fp);
k->disk_fp=fopen(k->disk_path,"r+b");
if(!k->disk_fp) return 0;
}
return 1;
}
static void kv_disk_truncate(Model *m, int nrec){
if(!g_kvsave) return;
KVState *k=m->kv;
if(k->disk_fp){ fclose(k->disk_fp); k->disk_fp=NULL; }
FILE *f=fopen(k->disk_path,"r+b");
if(!f){ k->disk_nrec=0; return; }
k->disk_nrec=nrec;
int32_t nr=nrec; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f);
fflush(f); fclose(f);
}
static void kv_disk_reset(Model *m){ kv_disk_truncate(m,0); }
static void kv_disk_append(Model *m, const int *hist, int len){
KVState *k=m->kv;
if(!g_kvsave || len<=k->disk_nrec) return;
Cfg *c=&m->c;
if(!kv_disk_open(m)) return;
FILE *f=k->disk_fp;
int64_t rec = kv_rec_bytes(m);
if(rec > k->disk_buf_cap){
uint8_t *nb=realloc(k->disk_buf, rec);
if(!nb) return;
k->disk_buf=nb; k->disk_buf_cap=rec;
}
fseek(f, 8+8*4 + (int64_t)k->disk_nrec*rec, SEEK_SET);
for(int p=k->disk_nrec;p<len;p++){
uint8_t *b=k->disk_buf;
*(int32_t*)b = hist[p]; b+=4;
for(int i=0;i<c->n_layers;i++){
memcpy(b, m->Lc[i]+(int64_t)p*c->kv_lora, (size_t)c->kv_lora*4); b+=c->kv_lora*4;
memcpy(b, m->Rc[i]+(int64_t)p*c->qk_rope,(size_t)c->qk_rope*4); b+=c->qk_rope*4;
}
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i]){
memcpy(b, m->Ic[i]+(int64_t)p*c->index_hd, (size_t)c->index_hd*4); b+=c->index_hd*4;
}
fwrite(k->disk_buf, 1, (size_t)rec, f);
}
fflush(f);
int32_t nr=len; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f);
fflush(f);
k->disk_nrec=len;
}
static int kv_disk_load(Model *m, int *hist, int maxctx){
if(!g_kvsave) return 0;
KVState *k=m->kv;
Cfg *c=&m->c;
FILE *f=fopen(k->disk_path,"rb"); if(!f) return 0;
char mg[8]; int32_t h[8], w[8]; kv_hdr(m,w,0);
if(fread(mg,1,8,f)!=8 || memcmp(mg,KV_MAGIC,8) || fread(h,4,8,f)!=8 ||
h[0]!=w[0]||h[1]!=w[1]||h[2]!=w[2]||h[3]!=w[3]||h[4]!=w[4]||h[5]!=w[5]){
fprintf(stderr,"[KV] ignoring .coli_kv from a different model or version\n"); fclose(f); return 0; }
int nrec=h[6];
if(nrec<1){ fclose(f); return 0; }
if(nrec>=maxctx-8-g_draft){
fprintf(stderr,"[KV] saved conversation (%d tokens) exceeds the context: starting over\n",nrec);
fclose(f); return 0; }
double t0=now_s();
for(int p=0;p<nrec;p++){
int32_t tk; if(fread(&tk,4,1,f)!=1){ nrec=p; break; } hist[p]=tk;
for(int i=0;i<c->n_layers;i++){
if(fread(m->Lc[i]+(int64_t)p*c->kv_lora, 4, c->kv_lora, f)!=(size_t)c->kv_lora ||
fread(m->Rc[i]+(int64_t)p*c->qk_rope, 4, c->qk_rope, f)!=(size_t)c->qk_rope){ nrec=p; goto out; }
}
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i])
if(fread(m->Ic[i]+(int64_t)p*c->index_hd, 4, c->index_hd, f)!=(size_t)c->index_hd){ nrec=p; goto out; }
}
out:
fclose(f);
if(nrec>0){
if(m->has_mtp) m->kv_start[c->n_layers]=-1;
fprintf(stderr,"[KV] resumed conversation from disk: %d tokens in %.1fs (no re-prefill)\n",
nrec, now_s()-t0);
}
k->disk_nrec=nrec;
return nrec;
}
#endif /* KV_PERSIST_H */
c/olmoe.c
+573 -38
View File
@@ -5,6 +5,15 @@
* Densa (embed, attn, router, norme, lm_head) residente in RAM (float32).
* Expert letti dal disco on-demand via pread+fadvise(DONTNEED), cache LRU per-layer.
* Matmul multi-thread con OpenMP (niente BLAS).
*
* ENV VARS:
* PILOT=0/1/2/3 : 0=no prefetch, 1=1-layer lookahead, 2=2-layer, 3=3-layer lookahead
* HOT=N : pin top-N hot experts per layer permanently (never evict)
* WARMUP=N : tokens before hot pinning activates (default 5)
* WIDE=N : prefetch top-K*N candidates (default 1, try 2 or 3)
* SMOOTH=F : EMA coefficient for routing momentum (default 0.3, range 0.0-0.95)
* CONF_LIMIT=F : cumulative gate probability threshold for prefetch cutoff (default 0.92)
* (expert queue is sorted by eid for SSD read locality)
*/
#define _GNU_SOURCE
#include <stdio.h>
@@ -12,11 +21,22 @@
#include <string.h>
#include <math.h>
#include <time.h>
#include <pthread.h>
#if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__)
#include <sys/resource.h>
#include <unistd.h>
#endif
#include "st.h"
#ifdef _WIN32
#include <windows.h>
#define sleep_ms(ms) Sleep(ms)
#else
#define sleep_ms(ms) usleep((ms) * 1000)
#endif
/* ---------- config ---------- */
typedef struct {
int hidden, n_layers, n_heads, n_kv_heads, head_dim;
@@ -33,22 +53,61 @@ typedef struct {
* Ogni weight [out,in] tenuto come int8 (per-riga) + scala float per riga.
* Cosi' la RAM-cache scende da 4 byte/param (f32) a 1 byte/param: e' il
* meccanismo che fa stare GLM-5.2 nei 15 GB. dequant-on-use nel matmul. */
typedef struct { int eid; int8_t *g, *u, *d; float *gs, *us, *ds; uint64_t used; } Slot;
/* pinned=1 means this slot is strongly preferred to keep (hot expert); it will
* not be evicted during normal LRU eviction, but may be displaced under extreme
* cache pressure when all slots are pinned or in-flight. */
typedef struct { int eid; int pinned; int8_t *g, *u, *d; float *gs, *us, *ds; uint64_t used; } Slot;
typedef struct { Slot *slots; int n, cap; } LCache;
typedef struct {
Cfg c;
shards S;
int quant_bits; /* bit di quantizzazione degli expert (2..8); storage int8, niente f32 (#134) */
int quant_bits;
float *embed, *lm_head, *final_norm;
Layer *L;
LCache *cache; /* [n_layers] */
uint64_t clock, hits, miss;
/* kv-cache per-layer: K,V come [H * maxT * head_dim] */
float **K, **V; int kv_len, max_t;
double dense_load_s;
/* IMPROVEMENT 2: expert frequency heatmap */
uint32_t *freq;
int freq_token_count, hot_pinned, hot_n, warmup_tokens;
int token_count;
/* PREDICTION IMPROVEMENT A: per-layer EMA of gate logits across tokens.
* momentum_logits[l*E .. (l+1)*E-1] = EMA of gate outputs for layer l.
* Used exclusively by the PILOT prefetcher to stabilise routing predictions
* across tokens; does NOT affect actual MoE routing (pr is unchanged). */
float *momentum_logits; /* [n_layers * n_experts], EMA of gate logits */
float pilot_smooth; /* SMOOTH env: EMA coefficient 0.0-0.9 (default 0.3) */
uint8_t *is_pinned; /* [n_layers * n_experts], 1 if expert is globally pinned */
uint8_t *is_queued; /* [n_layers * n_experts], 1 if expert is currently in the prefetch queue */
float pilot_conf_limit; /* CONF_LIMIT env: cumulative gate probability threshold (e.g. 0.92) */
} Model;
static pthread_mutex_t g_pilot_mx = PTHREAD_MUTEX_INITIALIZER;
static struct { int l, e; } pilot_q[4096];
static volatile unsigned pilot_r = 0, pilot_w = 0;
static Model *pilot_m = NULL;
static int g_pilot = 0;
static int g_wide = 1; /* IMPROVEMENT 4: top-K * g_wide candidates prefetched */
static void pilot_prefetch(Model *m, int lnext, const float *x, int S);
static void *pilot_worker(void *arg);
static void ensure_pilot_worker_started(Model *m);
static void slot_ensure_allocated(Model *m, Slot *s);
static void ensure_pilot_worker_started(Model *m) {
if (!pilot_m) {
pilot_m = m;
pthread_t t;
if (pthread_create(&t, NULL, pilot_worker, NULL) != 0) {
fprintf(stderr, "Error: Failed to create pilot prefetch worker thread\n");
exit(1);
}
pthread_detach(t);
}
}
/* ---------- utility ---------- */
static double now_s(void) { struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec*1e-9; }
#if defined(__APPLE__)
@@ -210,51 +269,224 @@ static void model_init(Model *m, const char *snap, int cap, int bits) {
#undef LD
}
m->cache = calloc(c->n_layers, sizeof(LCache));
for (int i = 0; i < c->n_layers; i++) { m->cache[i].cap = cap; m->cache[i].slots = calloc(cap, sizeof(Slot)); }
for (int i = 0; i < c->n_layers; i++) {
m->cache[i].cap = cap;
m->cache[i].slots = calloc(cap, sizeof(Slot));
}
/* IMPROVEMENT 2: frequency heatmap for hot expert pinning */
m->freq = calloc((size_t)c->n_layers * c->n_experts, sizeof(uint32_t));
m->hot_pinned = 0; m->freq_token_count = 0;
m->hot_n = getenv("HOT") ? atoi(getenv("HOT")) : 0;
m->warmup_tokens = getenv("WARMUP") ? atoi(getenv("WARMUP")) : 5;
m->token_count = 0;
/* PREDICTION A: routing momentum — EMA of gate logits across tokens.
* Initialized to zero; first token sets EMA = fresh logits. */
m->momentum_logits = calloc((size_t)c->n_layers * c->n_experts, sizeof(float));
float sv = getenv("SMOOTH") ? (float)atof(getenv("SMOOTH")) : 0.3f;
if (sv < 0.f) sv = 0.f; if (sv > 0.95f) sv = 0.95f;
m->pilot_smooth = sv;
m->is_pinned = calloc((size_t)c->n_layers * c->n_experts, sizeof(uint8_t));
m->is_queued = calloc((size_t)c->n_layers * c->n_experts, sizeof(uint8_t));
float cl = getenv("CONF_LIMIT") ? (float)atof(getenv("CONF_LIMIT")) : 0.92f;
if (cl < 0.1f) cl = 0.1f; if (cl > 1.0f) cl = 1.0f;
m->pilot_conf_limit = cl;
m->dense_load_s = now_s() - t0;
// Persistent Hot Pinning: try to load hot_pinned.bin
char pinpath[512];
snprintf(pinpath, sizeof(pinpath), "%s/hot_pinned.bin", snap);
FILE *pinf = fopen(pinpath, "rb");
if (pinf) {
size_t expected_size = (size_t)c->n_layers * c->n_experts;
if (fread(m->is_pinned, 1, expected_size, pinf) == expected_size) {
m->hot_pinned = 1;
printf("[HOT] Loaded persistent pinning from %s\n", pinpath);
if (g_pilot) {
ensure_pilot_worker_started(m);
for (int l = 0; l < c->n_layers; l++) {
for (int e = 0; e < c->n_experts; e++) {
if (m->is_pinned[l * c->n_experts + e]) {
unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_RELAXED);
unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
if (w - r < 4096) {
pilot_q[w & 4095].l = l; pilot_q[w & 4095].e = e;
pthread_mutex_lock(&g_pilot_mx);
m->is_queued[l * c->n_experts + e] = 1;
pthread_mutex_unlock(&g_pilot_mx);
__atomic_store_n(&pilot_w, w + 1, __ATOMIC_RELEASE);
}
}
}
}
printf("[HOT] Pre-loading pinned experts into cache...\n");
double t_wait = now_s();
while (1) {
unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_ACQUIRE);
if (r == w) break;
sleep_ms(2);
}
printf("[HOT] Pre-loaded in %.1fs!\n", now_s() - t_wait);
}
}
fclose(pinf);
}
}
/* legge un weight dal disco (streaming) e lo quantizza in q[O,I]+scale[O].
* Container pre-quantizzato (convert_olmoe.py: int8 + scale f32 in "name.qs"):
* lettura raw diretta — meta' I/O e zero quantize_rows a runtime. Prima di
* questa patch il container int8 causava SIGBUS (st_read_f32 su tensori I8). */
static void load_expert_w(Model *m, const char *name, int8_t *q, float *scale, int O, int I, float *tmp) {
st_tensor *t = st_find(&m->S, name);
if (t && t->dtype == 3) { /* I8/U8: container colibri */
char qs[300]; snprintf(qs, sizeof(qs), "%s.qs", name);
st_read_raw(&m->S, name, q, 1);
st_read_f32(&m->S, qs, scale, 1);
return;
static void slot_ensure_allocated(Model *m, Slot *s) {
if (s->g) return;
Cfg *c = &m->c;
int64_t ng = (int64_t)c->inter * c->hidden;
int64_t nd = (int64_t)c->hidden * c->inter;
int8_t *w_block = malloc(ng + ng + nd);
if (!w_block) {
fprintf(stderr, "Error: Out of memory allocating slot weights block\n");
exit(1);
}
st_read_f32(&m->S, name, tmp, 1); /* pread + fadvise DONTNEED */
quantize_rows(tmp, q, scale, O, I, m->quant_bits);
s->g = w_block;
s->u = w_block + ng;
s->d = w_block + ng + ng;
float *s_block = falloc(c->inter + c->inter + c->hidden);
s->gs = s_block;
s->us = s_block + c->inter;
s->ds = s_block + c->inter + c->inter;
s->pinned = 0;
}
static void load_expert_merged(Model *m, int layer, int eid, Slot *s) {
char nm[256], qsnm[256];
snprintf(nm, sizeof(nm), "model.layers.%d.mlp.experts.%d.merged_weight", layer, eid);
snprintf(qsnm, sizeof(qsnm), "model.layers.%d.mlp.experts.%d.qs", layer, eid);
st_read_raw(&m->S, nm, s->g, 1);
st_read_f32(&m->S, qsnm, s->gs, 0); /* scales are F32; use typed reader for dtype safety */
}
/* ---------- cache expert: ritorna i pesi quantizzati (q+scale) da cache o disco ---------- */
static void expert_get(Model *m, int layer, int eid, Slot **out) {
LCache *lc = &m->cache[layer];
pthread_mutex_lock(&g_pilot_mx);
for (int i = 0; i < lc->n; i++) if (lc->slots[i].eid == eid) {
m->hits++; lc->slots[i].used = ++m->clock; *out = &lc->slots[i]; return;
m->hits++; lc->slots[i].used = ++m->clock; *out = &lc->slots[i];
pthread_mutex_unlock(&g_pilot_mx);
return;
}
m->miss++;
Cfg *c = &m->c;
int64_t ng = (int64_t)c->inter * c->hidden, nd = (int64_t)c->hidden * c->inter;
Slot *s;
if (lc->n < lc->cap) {
s = &lc->slots[lc->n++];
s->g = malloc(ng); s->u = malloc(ng); s->d = malloc(nd);
s->gs = falloc(c->inter); s->us = falloc(c->inter); s->ds = falloc(c->hidden);
} else { int lru = 0; for (int i = 1; i < lc->n; i++) if (lc->slots[i].used < lc->slots[lru].used) lru = i; s = &lc->slots[lru]; }
float *tmp = falloc(ng > nd ? ng : nd);
char nm[256];
snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.gate_proj.weight",layer,eid); load_expert_w(m,nm,s->g,s->gs,c->inter,c->hidden,tmp);
snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.up_proj.weight", layer,eid); load_expert_w(m,nm,s->u,s->us,c->inter,c->hidden,tmp);
snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.down_proj.weight",layer,eid); load_expert_w(m,nm,s->d,s->ds,c->hidden,c->inter,tmp);
free(tmp);
s->eid = eid; s->used = ++m->clock;
slot_ensure_allocated(m, s);
} else {
/* LRU eviction — skip pinned and in-flight (eid==-1) slots */
int lru = -1;
for (int i = 0; i < lc->n; i++) {
if (lc->slots[i].pinned || lc->slots[i].eid < 0) continue;
if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i;
}
if (lru < 0) {
/* All slots are pinned or in-flight; find oldest non-in-flight slot
* (may be pinned, but never select one currently being loaded). */
for (int i = 0; i < lc->n; i++) {
if (lc->slots[i].eid < 0) continue; /* never evict in-flight */
if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i;
}
}
if (lru < 0) lru = 0; /* absolute last resort: all in-flight, evict slot 0 */
s = &lc->slots[lru];
s->pinned = 0;
}
s->eid = -1;
s->used = ++m->clock;
pthread_mutex_unlock(&g_pilot_mx);
load_expert_merged(m, layer, eid, s);
pthread_mutex_lock(&g_pilot_mx);
s->eid = eid;
s->pinned = m->is_pinned[layer * c->n_experts + eid];
s->used = ++m->clock;
*out = s;
pthread_mutex_unlock(&g_pilot_mx);
}
/* ---------- IMPROVEMENT 2: pin top-N hot experts per layer ---------- */
static void pin_hot_experts(Model *m) {
Cfg *c = &m->c;
if (m->hot_n <= 0 || m->hot_pinned) return;
m->hot_pinned = 1;
int is_dynamic = (m->hot_n >= 100);
double thresh = is_dynamic ? (double)m->hot_n / 1000.0 : 0.0;
int pinned_total = 0;
for (int l = 0; l < c->n_layers; l++) {
uint32_t *freq_l = m->freq + (int64_t)l * c->n_experts;
uint64_t layer_total = 0;
for (int e = 0; e < c->n_experts; e++) layer_total += freq_l[e];
if (layer_total == 0) continue;
int max_pin = m->cache[l].cap - 8;
if (max_pin < 4) max_pin = 4;
int hn = is_dynamic ? max_pin : (m->hot_n < c->n_experts ? m->hot_n : c->n_experts);
if (hn > 256) hn = 256;
int hot_eids[256];
int actual_hn = 0;
for (int k = 0; k < hn; k++) {
int best = -1; uint32_t bv = 0;
for (int e = 0; e < c->n_experts; e++) {
int already = 0;
for (int j = 0; j < k; j++) if (hot_eids[j] == e) { already = 1; break; }
if (!already && freq_l[e] > bv) { bv = freq_l[e]; best = e; }
}
if (best < 0 || bv == 0) break;
if (is_dynamic && bv < thresh * layer_total) break;
hot_eids[k] = best;
actual_hn++;
}
for (int k = 0; k < actual_hn; k++) {
int eid = hot_eids[k];
m->is_pinned[l * c->n_experts + eid] = 1;
LCache *lc = &m->cache[l];
int found = 0;
pthread_mutex_lock(&g_pilot_mx);
for (int i = 0; i < lc->n; i++) {
if (lc->slots[i].eid == eid) { lc->slots[i].pinned = 1; found = 1; break; }
}
pthread_mutex_unlock(&g_pilot_mx);
if (!found && g_pilot > 0) {
/* Only enqueue when the prefetch worker is active (PILOT>0). */
ensure_pilot_worker_started(m);
unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_RELAXED);
unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
int gidx = l * c->n_experts + eid;
pthread_mutex_lock(&g_pilot_mx);
int already = m->is_queued[gidx];
if (!already && w - r < 4096) {
pilot_q[w & 4095].l = l; pilot_q[w & 4095].e = eid;
m->is_queued[gidx] = 1;
__atomic_store_n(&pilot_w, w + 1, __ATOMIC_RELEASE);
}
pthread_mutex_unlock(&g_pilot_mx);
}
pinned_total++;
}
}
if (is_dynamic) {
printf("[HOT] Dynamic Pinned %d experts total (thresh=%.1f%%) after %d warmup tokens\n",
pinned_total, thresh * 100.0, m->freq_token_count);
} else {
printf("[HOT] Pinned %d experts (top-%d/layer) after %d warmup tokens\n",
pinned_total, m->hot_n, m->freq_token_count);
}
}
/* ---------- RoPE su un vettore di una testa (head_dim) a posizione assoluta pos ---------- */
static void rope_head(float *x, int pos, const Cfg *c) {
int h = c->head_dim / 2;
@@ -325,6 +557,19 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
float *g = falloc(I), *u = falloc(I), *hh = falloc(D);
for (int s = 0; s < S; s++) {
float *pr = logits + (int64_t)s*E;
if (m->momentum_logits && m->pilot_smooth > 0.f) {
float *ema = m->momentum_logits + (int64_t)layer * E;
int is_zero = 1;
for (int e = 0; e < E; e++) { if (ema[e] != 0.f) { is_zero = 0; break; } }
if (is_zero) {
for (int e = 0; e < E; e++) ema[e] = pr[e];
} else {
for (int e = 0; e < E; e++) {
ema[e] = (1.f - m->pilot_smooth) * pr[e] + m->pilot_smooth * ema[e];
}
}
}
softmax_row(pr, E);
/* top-K indici (selezione parziale) */
int idx[64]; float val[64];
@@ -337,6 +582,11 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
idx[kk] = best; val[kk] = bv;
}
if (c->norm_topk) { float sm=0; for(int kk=0;kk<K;kk++) sm+=val[kk]; for(int kk=0;kk<K;kk++) val[kk]/=sm; }
/* IMPROVEMENT 2: update activation heatmap (before pinning activates) */
if (!m->hot_pinned && m->freq) {
uint32_t *freq_l = m->freq + (int64_t)layer * E;
for (int kk = 0; kk < K; kk++) if (idx[kk] >= 0) freq_l[idx[kk]]++;
}
const float *xs = x + (int64_t)s*D;
for (int kk = 0; kk < K; kk++) {
Slot *e; expert_get(m, layer, idx[kk], &e);
@@ -352,9 +602,20 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
free(logits); free(g); free(u); free(hh);
}
/* un passo: token nuovi ids[S] a posizione pos_base. Ritorna logits dell'ultimo token (malloc'd). */
static float *step(Model *m, const int *ids, int S, int pos_base) {
Cfg *c = &m->c; int D = c->hidden;
if (g_pilot && m->token_count > 0) {
/* Flush stale prefetch requests: clear is_queued so pilot_realload
* will skip any entries still sitting in pilot_q for the previous
* token. We deliberately do NOT move pilot_w backwards; that would
* break the ring-buffer invariant (pilot_r could exceed pilot_w if
* the worker consumed an entry concurrently). The worker will drain
* the stale slots harmlessly because pilot_realload already exits
* early when the expert is already cached or is_queued is clear. */
pthread_mutex_lock(&g_pilot_mx);
memset(m->is_queued, 0, (size_t)c->n_layers * c->n_experts);
pthread_mutex_unlock(&g_pilot_mx);
}
float *x = falloc((int64_t)S*D);
for (int s = 0; s < S; s++) memcpy(x + (int64_t)s*D, m->embed + (int64_t)ids[s]*D, D*sizeof(float));
float *nrm = falloc((int64_t)S*D), *tmp = falloc((int64_t)S*D);
@@ -363,12 +624,26 @@ static float *step(Model *m, const int *ids, int S, int pos_base) {
for (int s = 0; s < S; s++) rmsnorm_row(nrm + (int64_t)s*D, x + (int64_t)s*D, l->in_ln, D, c->eps);
attention(m, l, i, nrm, S, pos_base, tmp);
for (int64_t j = 0; j < (int64_t)S*D; j++) x[j] += tmp[j];
/* IMPROVEMENT 1: PILOT=1 -> 1-layer lookahead */
if (g_pilot >= 1 && S <= 8 && i + 1 < c->n_layers)
pilot_prefetch(m, i + 1, x, S);
for (int s = 0; s < S; s++) rmsnorm_row(nrm + (int64_t)s*D, x + (int64_t)s*D, l->post_ln, D, c->eps);
moe(m, l, i, nrm, S, tmp);
for (int64_t j = 0; j < (int64_t)S*D; j++) x[j] += tmp[j];
/* PREDICTION IMPROVEMENT C (Residual gate trick):
* PILOT=2 -> prefetch layer i+2 using completed state x (containing MoE residual). */
if (g_pilot >= 2 && S <= 8 && i + 2 < c->n_layers)
pilot_prefetch(m, i + 2, x, S);
if (g_pilot >= 3 && S <= 8 && i + 3 < c->n_layers)
pilot_prefetch(m, i + 3, x, S);
}
/* count actual tokens processed (S>1 during prefill) */
m->token_count += S; m->freq_token_count += S;
if (!m->hot_pinned && m->hot_n > 0 && m->freq_token_count >= m->warmup_tokens)
pin_hot_experts(m);
m->kv_len = pos_base + S;
/* solo l'ultimo token -> logits */
float *last = falloc(D);
rmsnorm_row(last, x + (int64_t)(S-1)*D, m->final_norm, D, c->eps);
float *logit = falloc(c->vocab);
@@ -377,6 +652,192 @@ static float *step(Model *m, const int *ids, int S, int pos_base) {
return logit;
}
static void pilot_realload(Model *m, int layer, int eid) {
LCache *lc = &m->cache[layer];
Cfg *c = &m->c;
pthread_mutex_lock(&g_pilot_mx);
/* Early-exit if entry was flushed (is_queued cleared) while waiting. */
if (!m->is_queued[layer * c->n_experts + eid]) {
pthread_mutex_unlock(&g_pilot_mx);
return;
}
for (int i = 0; i < lc->n; i++) {
if (lc->slots[i].eid == eid) {
m->is_queued[layer * c->n_experts + eid] = 0;
pthread_mutex_unlock(&g_pilot_mx);
return;
}
}
Slot *s;
if (lc->n < lc->cap) {
s = &lc->slots[lc->n++];
slot_ensure_allocated(m, s);
} else {
/* LRU eviction — skip pinned and in-flight (eid==-1) slots */
int lru = -1;
for (int i = 0; i < lc->n; i++) {
if (lc->slots[i].pinned || lc->slots[i].eid < 0) continue;
if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i;
}
if (lru < 0) {
m->is_queued[layer * c->n_experts + eid] = 0;
pthread_mutex_unlock(&g_pilot_mx);
return; /* all pinned/in-flight, skip */
}
s = &lc->slots[lru]; s->pinned = 0;
}
s->eid = -1; s->used = ++m->clock;
pthread_mutex_unlock(&g_pilot_mx);
load_expert_merged(m, layer, eid, s);
pthread_mutex_lock(&g_pilot_mx);
s->eid = eid;
s->pinned = m->is_pinned[layer * c->n_experts + eid];
s->used = ++m->clock;
m->is_queued[layer * c->n_experts + eid] = 0;
pthread_mutex_unlock(&g_pilot_mx);
}
static void *pilot_worker(void *arg) {
(void)arg;
while (1) {
unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_ACQUIRE);
if (r == w) {
sleep_ms(1);
continue;
}
int layer = pilot_q[r & 4095].l;
int eid = pilot_q[r & 4095].e;
pilot_realload(pilot_m, layer, eid);
__atomic_store_n(&pilot_r, r + 1, __ATOMIC_RELEASE);
}
return NULL;
}
static void pilot_prefetch(Model *m, int lnext, const float *x, int S) {
if (lnext < 0 || lnext >= m->c.n_layers) return;
Cfg *c = &m->c; int D = c->hidden, E = c->n_experts;
ensure_pilot_worker_started(m);
float *logits = falloc((int64_t)S * E);
Layer *l = &m->L[lnext];
// PREDICTION IMPROVEMENT B: Apply RMSNorm to x using destination layer's post_ln
// This scales inputs to the distribution expected by l->gate.
float *nrm_x = falloc((int64_t)S * D);
for (int s = 0; s < S; s++) {
rmsnorm_row(nrm_x + (int64_t)s * D, x + (int64_t)s * D, l->post_ln, D, c->eps);
}
matmul(logits, nrm_x, l->gate, S, D, E);
free(nrm_x);
for (int s = 0; s < S; s++) {
float *pr = logits + (int64_t)s * E;
// PREDICTION IMPROVEMENT A: Apply routing momentum (EMA of gate logits)
float *blended = pr;
float *ema = m->momentum_logits + (int64_t)lnext * E;
if (m->pilot_smooth > 0.f) {
blended = falloc(E);
int is_zero = 1;
for (int e = 0; e < E; e++) { if (ema[e] != 0.f) { is_zero = 0; break; } }
if (is_zero) {
for (int e = 0; e < E; e++) {
ema[e] = pr[e];
blended[e] = pr[e];
}
} else {
for (int e = 0; e < E; e++) {
blended[e] = (1.f - m->pilot_smooth) * pr[e] + m->pilot_smooth * ema[e];
ema[e] = blended[e]; // update EMA
}
}
}
int cand = 0;
int idx[128];
float max_logit = -1e30f;
for (int e = 0; e < E; e++) { if (blended[e] > max_logit) max_logit = blended[e]; }
float *exps = falloc(E);
float sum_exps = 0.f;
for (int e = 0; e < E; e++) {
exps[e] = expf(blended[e] - max_logit);
sum_exps += exps[e];
}
float cum_sum = 0.f;
int min_cand = c->topk;
int max_cand = c->topk * g_wide;
if (max_cand < min_cand) max_cand = min_cand;
if (max_cand > 128) max_cand = 128; /* idx[] buffer bound */
if (max_cand > E) max_cand = E;
for (int kk = 0; kk < max_cand; kk++) {
int best = -1; float bv = -1.f;
for (int e = 0; e < E; e++) {
int taken = 0; for (int j = 0; j < kk; j++) if (idx[j] == e) { taken=1; break; }
if (!taken && exps[e] > bv) { bv = exps[e]; best = e; }
}
if (best < 0) break;
idx[kk] = best;
cum_sum += bv;
cand++;
if (cum_sum >= m->pilot_conf_limit * sum_exps && cand >= min_cand) {
break;
}
}
free(exps);
if (blended != pr) free(blended);
/* IMPROVEMENT 5: sort candidates by eid for sequential SSD read locality */
for (int a = 0; a < cand-1; a++)
for (int b = a+1; b < cand; b++)
if (idx[b] >= 0 && (idx[a] < 0 || idx[a] > idx[b])) { int t = idx[a]; idx[a] = idx[b]; idx[b] = t; }
for (int kk = 0; kk < cand; kk++) {
int eid = idx[kk];
if (eid < 0) continue;
int found = 0;
pthread_mutex_lock(&g_pilot_mx);
LCache *lc = &m->cache[lnext];
for (int z = 0; z < lc->n; z++) {
if (lc->slots[z].eid == eid) { found = 1; break; }
}
pthread_mutex_unlock(&g_pilot_mx);
if (!found) {
int gidx = lnext * E + eid;
pthread_mutex_lock(&g_pilot_mx);
int already_queued = m->is_queued[gidx];
if (!already_queued) {
m->is_queued[gidx] = 1;
}
pthread_mutex_unlock(&g_pilot_mx);
if (!already_queued) {
unsigned w2 = __atomic_load_n(&pilot_w, __ATOMIC_RELAXED);
unsigned r2 = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
if (w2 - r2 < 4096) {
pilot_q[w2 & 4095].l = lnext;
pilot_q[w2 & 4095].e = eid;
__atomic_store_n(&pilot_w, w2 + 1, __ATOMIC_RELEASE);
} else {
pthread_mutex_lock(&g_pilot_mx);
m->is_queued[gidx] = 0;
pthread_mutex_unlock(&g_pilot_mx);
}
}
}
}
}
free(logits);
}
/* generazione greedy. prompt[np] -> riempie out[np+n_new] */
static void generate(Model *m, const int *prompt, int np, int n_new, int *out) {
Cfg *c = &m->c;
@@ -400,6 +861,37 @@ static void generate(Model *m, const int *prompt, int np, int n_new, int *out) {
}
}
/* teacher-forced NLL of full_ids[np..nfull): feed the REFERENCE token at each step
* (never the argmax), accumulate -log softmax(logits)[next_ref]. A loss meter for
* throughput experiments: same engine path as decode, so hit rate/speed stay
* comparable, but quality is measured as perplexity instead of exact-match.
* Cross-checked vs HF transformers bf16 on identical token ids: engine (int8
* experts) 12.11 ppl vs reference 12.25 (#108). Enabled by PPL=1. */
static int tf_nll(Model *m, const int *full, int nfull, int np, double *nll_out) {
Cfg *c = &m->c;
m->max_t = nfull;
m->K = calloc(c->n_layers, sizeof(float*)); m->V = calloc(c->n_layers, sizeof(float*));
for (int i = 0; i < c->n_layers; i++) {
m->K[i] = falloc((int64_t)c->n_heads * m->max_t * c->head_dim);
m->V[i] = falloc((int64_t)c->n_heads * m->max_t * c->head_dim);
}
double nll = 0; int scored = 0;
float *logit = step(m, full, np, 0); /* prefill on the prompt */
for (int i = np; i < nfull; i++) {
/* log softmax(logit)[full[i]] without materializing the softmax */
float mx = logit[0]; for (int v = 1; v < c->vocab; v++) if (logit[v] > mx) mx = logit[v];
double Z = 0; for (int v = 0; v < c->vocab; v++) Z += exp((double)logit[v] - mx);
nll += -((double)logit[full[i]] - mx - log(Z));
scored++;
free(logit); logit = NULL;
if (i == nfull - 1) break;
logit = step(m, &full[i], 1, i); /* teacher forcing */
}
if (logit) free(logit);
*nll_out = nll / scored;
return scored;
}
/* ---------- lettura ref.json ---------- */
static int *read_int_array(jval *o, const char *key, int *n_out) {
jval *a = json_get(o, key);
@@ -411,25 +903,48 @@ static int *read_int_array(jval *o, const char *key, int *n_out) {
int main(int argc, char **argv) {
const char *snap = getenv("SNAP");
if (!snap) { fprintf(stderr, "set SNAP=<snapshot directory>\n"); return 1; }
int cap = argc > 1 ? atoi(argv[1]) : 16;
int bits = argc > 2 ? atoi(argv[2]) : 8;
if (bits < 2 || bits > 8) { /* expert storage is int8_t: bits>8 truncates in quantize_rows (#134). f32 mode is not implemented here — int8 is already token-exact vs the oracle. */
fprintf(stderr, "quant_bits must be 2..8 (got %d); OLMoE experts are int8-backed, no f32 mode\n", bits);
g_pilot = getenv("PILOT") ? atoi(getenv("PILOT")) : 0;
g_wide = getenv("WIDE") ? atoi(getenv("WIDE")) : 1;
if (g_wide < 1) g_wide = 1;
if (g_wide > 4) g_wide = 4;
int hot_n = getenv("HOT") ? atoi(getenv("HOT")) : 0;
int cap = argc > 1 ? atoi(argv[1]) : 16;
int bits = argc > 2 ? atoi(argv[2]) : 8;
if (bits < 2 || bits > 8) {
fprintf(stderr, "quant_bits must be 2..8 (got %d)\n", bits);
return 1;
}
const char *refpath = argc > 3 ? argv[3] : "ref.json";
FILE *f = fopen(refpath, "rb"); if(!f){perror(refpath);return 1;}
float smooth = getenv("SMOOTH") ? (float)atof(getenv("SMOOTH")) : 0.3f;
float conf = getenv("CONF_LIMIT") ? (float)atof(getenv("CONF_LIMIT")) : 0.92f;
printf("== Streaming C engine v2.2 | cache=%d/layer bits=%d pilot=%d wide=%d hot=%d smooth=%.2f conf=%.2f ==\n",
cap, bits, g_pilot, g_wide, hot_n, smooth, conf);
FILE *f = fopen(refpath, "rb"); if (!f) { perror(refpath); return 1; }
fseek(f,0,SEEK_END); long n=ftell(f); fseek(f,0,SEEK_SET);
char *buf=malloc(n+1); if(fread(buf,1,n,f)!=(size_t)n){} buf[n]=0; fclose(f);
char *buf=malloc(n+1); if (fread(buf,1,n,f)!=(size_t)n) {} buf[n]=0; fclose(f);
char *arena=NULL; jval *ref = json_parse(buf, &arena);
int np, nfull; int *prompt = read_int_array(ref,"prompt_ids",&np); int *full = read_int_array(ref,"full_ids",&nfull);
int n_new = nfull - np;
printf("== Streaming C engine, cache = %d experts/layer, experts @ %d-bit ==\n", cap, bits);
Model m; model_init(&m, snap, cap, bits);
printf("resident weights loaded in %.1fs | RSS after load: %.2f GB\n", m.dense_load_s, rss_gb());
if (getenv("PPL") && atoi(getenv("PPL")) == 1) { /* loss-meter mode: teacher-forced NLL */
double nll; double t = now_s();
int scored = tf_nll(&m, full, nfull, np, &nll);
double dt = now_s() - t;
double tot = m.hits + m.miss;
printf("TF-NLL: %.4f nats/token over %d tokens | ppl = %.2f\n", nll, scored, exp(nll));
printf("Expert cache hit rate: %.1f%% (hit=%llu miss=%llu)\n", tot?100.0*m.hits/tot:0.0,
(unsigned long long)m.hits, (unsigned long long)m.miss);
printf("Speed: %.2f tok/s (%.1fs for %d tokens) | PEAK RSS: %.2f GB\n", scored/dt, dt, scored, rss_gb());
free(buf); free(arena);
return 0;
}
int *out = malloc((np + n_new) * sizeof(int));
double t = now_s();
generate(&m, prompt, np, n_new, out);
@@ -443,6 +958,26 @@ int main(int argc, char **argv) {
printf("\nPEAK RSS: %.2f GB\n", rss_gb());
printf("Expert cache hit rate: %.1f%% (hit=%llu miss=%llu)\n", tot?100.0*m.hits/tot:0.0,
(unsigned long long)m.hits, (unsigned long long)m.miss);
// Persistent Hot Pinning: save dynamic pinning if newly created
if (m.hot_pinned) {
char pinpath[512];
snprintf(pinpath, sizeof(pinpath), "%s/hot_pinned.bin", snap);
FILE *pinf_chk = fopen(pinpath, "rb");
if (!pinf_chk) {
FILE *pinf_save = fopen(pinpath, "wb");
if (pinf_save) {
size_t expected_size = (size_t)m.c.n_layers * m.c.n_experts;
fwrite(m.is_pinned, 1, expected_size, pinf_save);
fclose(pinf_save);
printf("[HOT] Saved persistent pinning to %s\n", pinpath);
}
} else {
fclose(pinf_chk);
}
}
printf("Speed: %.2f tok/s (%.1fs for %d tokens)\n", n_new/dt, dt, n_new);
free(buf); free(arena);
return 0;
c/openai_server.py
+35 -1
View File
@@ -27,6 +27,7 @@ HERE = Path(__file__).resolve().parent
END = b"\x01\x01END\x01\x01\n"
READY = b"\x01\x01READY\x01\x01\n"
MAX_BODY = 4 << 20
PROFILE_TURNS = 120 # rolling window of per-turn PROF snapshots kept for /profile
DEFAULT_CORS_ORIGINS = (
"http://127.0.0.1:8000",
"http://localhost:8000",
@@ -270,6 +271,14 @@ def parse_tool_calls(reply, tools=None):
salvaged.append(name)
calls.append({"id": "call_" + uuid.uuid4().hex[:24], "type": "function",
"function": {"name": name, "arguments": json.dumps(args, ensure_ascii=False)}})
if tools and not calls and re.search(r"</?tool_call>|</?arg_key>|</?arg_value>", reply):
# Diagnosi per la #401: il client ha dichiarato i tools e il modello ha PROVATO la
# sintassi, ma il parse rigoroso non ha agganciato nulla (tipico output int4 storpiato).
# EN: #401 field diagnosis: tools were declared and the model attempted the syntax,
# EN: but the strict parse matched nothing (typically quantization-mangled output).
sys.stderr.write("[api] tools declared and tool-call markers present, but no call "
"parsed -- output may be quantization-mangled; try COLI_TOOL_SALVAGE=1\n")
sys.stderr.flush()
text = _BOX_RE.sub("", reply)
if THINK_CLOSE in text:
text = text.split(THINK_CLOSE, 1)[1]
@@ -426,7 +435,10 @@ def generation_options(body, limit):
maximum = body.get("max_tokens")
maximum_param = "max_tokens"
if maximum is None:
maximum = min(256, limit)
# Client omitted max_tokens: honor the operator's configured budget (--max-tokens /
# --ngen), not an arbitrary 256 — `coli serve --ngen 32768` must mean 32768 (#382).
# Generation still ends at EOS, so this is a cap, not a target.
maximum = limit
temperature = body.get("temperature")
top_p = body.get("top_p")
temperature = 0.7 if temperature is None else temperature
@@ -494,6 +506,8 @@ class Engine:
self.emap = None
self.hits = None
self.hits_seq = 0 # latest "TIERS" snapshot from the engine
self.profile = collections.deque(maxlen=PROFILE_TURNS) # per-turn phase timings
self.profile_seq = 0
read_engine_turn(self.process.stdout, READY, lambda _: None)
self.dispatcher = threading.Thread(target=self._dispatch_stdout,
name="colibri-stdout", daemon=True)
@@ -571,6 +585,20 @@ class Engine:
elif kind == "HITS" and len(fields) == 4:
self.hits = fields[3]
self.hits_seq += 1
elif kind == "PROF" and len(fields) >= 10:
# per-turn phase timings: where the engine spent this turn's wall time
self.profile.append({
"wall_s": float(fields[1]),
"prompt_tokens": int(fields[2]),
"completion_tokens": int(fields[3]),
"expert_disk_s": float(fields[4]),
"expert_wait_s": float(fields[5]),
"expert_matmul_s": float(fields[6]),
"attention_s": float(fields[7]),
"lm_head_s": float(fields[8]),
"forwards": int(fields[9]),
})
self.profile_seq += 1
elif kind == "TIERS" and len(fields) >= 6:
self.tiers = {"vram": int(fields[1]), "ram": int(fields[2]),
"disk": int(fields[3]), "vram_gb": float(fields[4]),
@@ -803,6 +831,12 @@ class APIHandler(BaseHTTPRequestHandler):
payload["seq"] = eng.hits_seq
self.send_json(200, payload, request_id)
return
if path == "/profile":
eng = self.server.engine
payload = {"seq": getattr(eng, "profile_seq", 0) if eng else 0,
"turns": list(getattr(eng, "profile", ()) or ()) if eng else []}
self.send_json(200, payload, request_id)
return
if self.serve_static(path):
return
self.require_auth()
c/quant.h
+672
View File
@@ -0,0 +1,672 @@
/* quant.h — quantized matmul kernels (header-only, all functions static).
* Multi-architecture SIMD: AVX2 / AVX-512 / AVX-VNNI / ARM NEON / NEON-SDOT /
* NEON-i8mm / POWER VSX. Pure compute — no Model or QT dependency. */
#ifndef COLI_QUANT_H
#define COLI_QUANT_H
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
#ifdef _OPENMP
#include <omp.h>
#endif
/* ---- SIMD includes -------------------------------------------------------- */
#ifdef __AVX2__
#include <immintrin.h>
static inline float hsum256(__m256 v){
__m128 lo=_mm256_castps256_ps128(v), hi=_mm256_extractf128_ps(v,1);
lo=_mm_add_ps(lo,hi); __m128 sh=_mm_movehl_ps(lo,lo); lo=_mm_add_ps(lo,sh);
sh=_mm_shuffle_ps(lo,lo,1); lo=_mm_add_ss(lo,sh); return _mm_cvtss_f32(lo);
}
static inline int hsum256_i32(__m256i v){
__m128i lo=_mm256_castsi256_si128(v), hi=_mm256_extracti128_si256(v,1);
lo=_mm_add_epi32(lo,hi); lo=_mm_hadd_epi32(lo,lo); lo=_mm_hadd_epi32(lo,lo);
return _mm_cvtsi128_si32(lo);
}
#endif
#if defined(__AVXVNNI__) && defined(__AVX2__)
static inline int hsum128_i32(__m128i v){
v=_mm_hadd_epi32(v,v); v=_mm_hadd_epi32(v,v); return _mm_cvtsi128_si32(v);
}
#endif
#ifdef __ARM_NEON
#include <arm_neon.h>
#endif
#ifdef __VSX__
#include <altivec.h>
#undef vector
#undef pixel
#undef bool
#endif
/* ---- AVX-512 int4->float accumulator -------------------------------------- */
#if defined(__AVX512F__) && defined(__AVX512BW__)
static int g_i4_acc512=1;
static inline float dot_i4f_avx512(const uint8_t *w,const float *x,int I){
const __m128i m4=_mm_set1_epi8(0x0F); const __m512i b8=_mm512_set1_epi32(8);
__m512 acc0=_mm512_setzero_ps(),acc1=_mm512_setzero_ps(); int i=0;
for(;i+32<=I;i+=32){ __m128i by=_mm_loadu_si128((const __m128i*)(w+(i>>1)));
__m128i lo=_mm_and_si128(by,m4),hi=_mm_and_si128(_mm_srli_epi16(by,4),m4);
__m128i n0=_mm_unpacklo_epi8(lo,hi),n1=_mm_unpackhi_epi8(lo,hi);
__m512 w0=_mm512_cvtepi32_ps(_mm512_sub_epi32(_mm512_cvtepu8_epi32(n0),b8));
__m512 w1=_mm512_cvtepi32_ps(_mm512_sub_epi32(_mm512_cvtepu8_epi32(n1),b8));
acc0=_mm512_fmadd_ps(_mm512_loadu_ps(x+i),w0,acc0);
acc1=_mm512_fmadd_ps(_mm512_loadu_ps(x+i+16),w1,acc1);
}
return _mm512_reduce_add_ps(_mm512_add_ps(acc0,acc1));
}
static int i4_acc512_selftest(void){
enum { N=224 }; uint8_t w[(N+1)/2]; float x[N];
for(int i=0;i<N;i++){
int q=((i*13+5)&15)-8;
if(!(i&1)) w[i>>1]=(uint8_t)(q+8);
else w[i>>1]|=(uint8_t)((q+8)<<4);
x[i]=(float)(((i*29+7)%101)-50)/37.f;
}
for(int n=32;n<=N;n+=32){
float ref=0; for(int i=0;i<n;i++) ref+=x[i]*(float)(((w[i>>1]>>((i&1)*4))&15)-8);
float got=dot_i4f_avx512(w,x,n),tol=2e-5f*(1.f+fabsf(ref));
if(fabsf(got-ref)>tol){ fprintf(stderr,"AVX512 i4 selftest n=%d: %.9g != %.9g\n",n,got,ref); return 0; }
}
return 1;
}
#endif
/* ---- y[S,O] = x[S,I] @ W^T, W[O,I] f32 ---------------------------------- */
static void matmul(float *y, const float *x, const float *W, int S, int I, int O){
#pragma omp parallel for schedule(static)
for (int o=0;o<O;o++){ const float *w=W+(int64_t)o*I;
for (int s=0;s<S;s++){ const float *xs=x+(int64_t)s*I; float a=0; for(int i=0;i<I;i++) a+=xs[i]*w[i]; y[(int64_t)s*O+o]=a; } }
}
/* ---- y[S,O] = x[S,I] @ W^T, W int8 per-row + scale[O] ------------------- */
static void matmul_q(float *y, const float *x, const int8_t *q, const float *scale, int S, int I, int O){
#pragma omp parallel for schedule(static)
for (int o=0;o<O;o++){ const int8_t *w=q+(int64_t)o*I; float sc=scale[o];
for (int s=0;s<S;s++){ const float *xs=x+(int64_t)s*I; float a=0; int i=0;
#ifdef __AVX2__
__m256 acc=_mm256_setzero_ps();
for(;i+8<=I;i+=8){ __m256i wi=_mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i*)(w+i)));
acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i), _mm256_cvtepi32_ps(wi), acc); }
a=hsum256(acc);
#elif defined(__ARM_NEON)
float32x4_t ac0=vdupq_n_f32(0), ac1=vdupq_n_f32(0);
for(;i+8<=I;i+=8){ int16x8_t w16=vmovl_s8(vld1_s8(w+i));
ac0=vfmaq_f32(ac0, vld1q_f32(xs+i), vcvtq_f32_s32(vmovl_s16(vget_low_s16(w16))));
ac1=vfmaq_f32(ac1, vld1q_f32(xs+i+4), vcvtq_f32_s32(vmovl_s16(vget_high_s16(w16)))); }
a=vaddvq_f32(vaddq_f32(ac0,ac1));
#endif
for(;i<I;i++) a+=xs[i]*(float)w[i]; y[(int64_t)s*O+o]=a*sc; } }
}
/* ---- y[S,O] = x[S,I] @ W^T, W int4 packed (2/byte) + scale[O] ------------ */
static void matmul_i4(float *y, const float *x, const uint8_t *q4, const float *scale, int S, int I, int O){
int rb=(I+1)/2;
#pragma omp parallel for schedule(static)
for (int o=0;o<O;o++){ const uint8_t *w=q4+(int64_t)o*rb; float sc=scale[o];
for (int s=0;s<S;s++){ const float *xs=x+(int64_t)s*I; float a=0; int i=0;
#if defined(__AVX512F__) && defined(__AVX512BW__)
if(g_i4_acc512){ a=dot_i4f_avx512(w,xs,I); i=I&~31; }
else {
#endif
#ifdef __AVX2__
const __m128i m4=_mm_set1_epi8(0x0F); const __m256i b8=_mm256_set1_epi32(8);
__m256 acc=_mm256_setzero_ps();
for(;i+16<=I;i+=16){ __m128i by=_mm_loadl_epi64((const __m128i*)(w+(i>>1)));
__m128i lo=_mm_and_si128(by,m4), hi=_mm_and_si128(_mm_srli_epi16(by,4),m4);
__m128i nib=_mm_unpacklo_epi8(lo,hi);
__m256 w0=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nib),b8));
__m256 w1=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nib,8)),b8));
acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i), w0, acc);
acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i+8), w1, acc); }
a=hsum256(acc);
#elif defined(__ARM_NEON)
const uint8x8_t m4=vdup_n_u8(0x0F); const int8x8_t b8=vdup_n_s8(8);
float32x4_t ac0=vdupq_n_f32(0), ac1=vdupq_n_f32(0);
for(;i+16<=I;i+=16){ uint8x8_t by=vld1_u8(w+(i>>1));
uint8x8x2_t z=vzip_u8(vand_u8(by,m4), vshr_n_u8(by,4));
int16x8_t w0=vmovl_s8(vsub_s8(vreinterpret_s8_u8(z.val[0]),b8));
int16x8_t w1=vmovl_s8(vsub_s8(vreinterpret_s8_u8(z.val[1]),b8));
ac0=vfmaq_f32(ac0, vld1q_f32(xs+i), vcvtq_f32_s32(vmovl_s16(vget_low_s16(w0))));
ac1=vfmaq_f32(ac1, vld1q_f32(xs+i+4), vcvtq_f32_s32(vmovl_s16(vget_high_s16(w0))));
ac0=vfmaq_f32(ac0, vld1q_f32(xs+i+8), vcvtq_f32_s32(vmovl_s16(vget_low_s16(w1))));
ac1=vfmaq_f32(ac1, vld1q_f32(xs+i+12), vcvtq_f32_s32(vmovl_s16(vget_high_s16(w1)))); }
a=vaddvq_f32(vaddq_f32(ac0,ac1));
#endif
#if defined(__AVX512F__) && defined(__AVX512BW__)
}
#endif
for(;i+1<I;i+=2){ uint8_t byte=w[i>>1]; int lo=(int)(byte&0xF)-8, hi=(int)(byte>>4)-8;
a += xs[i]*(float)lo + xs[i+1]*(float)hi; }
if(i<I){ uint8_t byte=w[i>>1]; int lo=(int)(byte&0xF)-8; a += xs[i]*(float)lo; }
y[(int64_t)s*O+o]=a*sc; } }
}
/* ---- y[S,O] = x[S,I] @ W^T, W int4 packed + per-GROUP scales (fmt=4) ----- */
static void matmul_i4_grouped(float *y, const float *x, const uint8_t *q4, const float *scale,
int S, int I, int O, int gs){
int rb=(I+1)/2; int ng=(I+gs-1)/gs;
#pragma omp parallel for schedule(static)
for(int o=0;o<O;o++){
const uint8_t *w=q4+(int64_t)o*rb;
const float *scl=scale+(int64_t)o*ng;
for(int s=0;s<S;s++){
const float *xs=x+(int64_t)s*I; float a=0;
for(int g=0; g*gs<I; g++){
int base=g*gs; int glen=gs; if(base+glen>I) glen=I-base;
float sc=scl[g];
int i=base;
#ifdef __AVX2__
const __m128i m4=_mm_set1_epi8(0x0F); const __m256i b8=_mm256_set1_epi32(8);
__m256 acc=_mm256_setzero_ps();
for(; i+16<=base+glen; i+=16){ __m128i by=_mm_loadl_epi64((const __m128i*)(w+(i>>1)));
__m128i lo=_mm_and_si128(by,m4),hi=_mm_and_si128(_mm_srli_epi16(by,4),m4);
__m128i nib=_mm_unpacklo_epi8(lo,hi);
__m256 w0=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nib),b8));
__m256 w1=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nib,8)),b8));
acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i), w0, acc);
acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i+8), w1, acc); }
a+=hsum256(acc)*sc;
#endif
for(; i<base+glen; i+=2){
if(i+1<base+glen){ uint8_t byte=w[i>>1];
a+=(xs[i]*(float)((int)(byte&0xF)-8)+xs[i+1]*(float)((int)(byte>>4)-8))*sc; }
else { uint8_t byte=w[i>>1]; a+=xs[i]*(float)((int)(byte&0xF)-8)*sc; }
}
}
y[(int64_t)s*O+o]=a;
}
}
}
/* ---- fused gate+up: one OMP dispatch for both matrices -------------------- */
static void matmul_i4_pair(float *yg, float *yu, const float *x,
const uint8_t *qg, const float *sg,
const uint8_t *qu, const float *su, int I, int O){
int rb=(I+1)/2;
#pragma omp parallel for schedule(static)
for(int z=0;z<2*O;z++){
int o=z<O?z:z-O; const uint8_t *w=(z<O?qg:qu)+(int64_t)o*rb;
float a=0; int i=0;
#if defined(__AVX512F__) && defined(__AVX512BW__)
if(g_i4_acc512){ a=dot_i4f_avx512(w,x,I); i=I&~31; }
else {
#endif
#ifdef __AVX2__
const __m128i m4=_mm_set1_epi8(0x0F); const __m256i b8=_mm256_set1_epi32(8);
__m256 acc=_mm256_setzero_ps();
for(;i+16<=I;i+=16){ __m128i by=_mm_loadl_epi64((const __m128i*)(w+(i>>1)));
__m128i lo=_mm_and_si128(by,m4),hi=_mm_and_si128(_mm_srli_epi16(by,4),m4);
__m128i nib=_mm_unpacklo_epi8(lo,hi);
__m256 w0=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nib),b8));
__m256 w1=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nib,8)),b8));
acc=_mm256_fmadd_ps(_mm256_loadu_ps(x+i),w0,acc);
acc=_mm256_fmadd_ps(_mm256_loadu_ps(x+i+8),w1,acc); }
a=hsum256(acc);
#elif defined(__ARM_NEON)
const uint8x8_t m4=vdup_n_u8(0x0F); const int8x8_t b8=vdup_n_s8(8);
float32x4_t ac0=vdupq_n_f32(0),ac1=vdupq_n_f32(0);
for(;i+16<=I;i+=16){ uint8x8_t by=vld1_u8(w+(i>>1));
uint8x8x2_t n=vzip_u8(vand_u8(by,m4),vshr_n_u8(by,4));
int16x8_t w0=vmovl_s8(vsub_s8(vreinterpret_s8_u8(n.val[0]),b8));
int16x8_t w1=vmovl_s8(vsub_s8(vreinterpret_s8_u8(n.val[1]),b8));
ac0=vfmaq_f32(ac0,vld1q_f32(x+i),vcvtq_f32_s32(vmovl_s16(vget_low_s16(w0))));
ac1=vfmaq_f32(ac1,vld1q_f32(x+i+4),vcvtq_f32_s32(vmovl_s16(vget_high_s16(w0))));
ac0=vfmaq_f32(ac0,vld1q_f32(x+i+8),vcvtq_f32_s32(vmovl_s16(vget_low_s16(w1))));
ac1=vfmaq_f32(ac1,vld1q_f32(x+i+12),vcvtq_f32_s32(vmovl_s16(vget_high_s16(w1)))); }
a=vaddvq_f32(vaddq_f32(ac0,ac1));
#endif
#if defined(__AVX512F__) && defined(__AVX512BW__)
}
#endif
for(;i+1<I;i+=2){ uint8_t b=w[i>>1]; a+=x[i]*(float)((b&15)-8)+x[i+1]*(float)((b>>4)-8); }
if(i<I) a+=x[i]*(float)((w[i>>1]&15)-8);
(z<O?yg:yu)[o]=a*(z<O?sg:su)[o];
}
}
/* ---- y[S,O] = x[S,I] @ W^T, W int2 packed (4/byte) + scale[O] ------------ */
static void matmul_i2(float *y, const float *x, const uint8_t *q2, const float *scale, int S, int I, int O){
int rb=(I+3)/4;
#pragma omp parallel for schedule(static)
for (int o=0;o<O;o++){ const uint8_t *w=q2+(int64_t)o*rb; float sc=scale[o];
for (int s=0;s<S;s++){ const float *xs=x+(int64_t)s*I; float a=0; int i=0;
#ifdef __AVX2__
const __m128i m2=_mm_set1_epi8(0x03); const __m256i b2=_mm256_set1_epi32(2);
__m256 acc=_mm256_setzero_ps();
for(;i+16<=I;i+=16){ __m128i by=_mm_cvtsi32_si128(*(const int*)(w+(i>>2)));
__m128i p0=_mm_and_si128(by,m2), p1=_mm_and_si128(_mm_srli_epi16(by,2),m2);
__m128i p2=_mm_and_si128(_mm_srli_epi16(by,4),m2), p3=_mm_and_si128(_mm_srli_epi16(by,6),m2);
__m128i lo=_mm_unpacklo_epi8(p0,p1), hi=_mm_unpacklo_epi8(p2,p3);
__m128i nib=_mm_unpacklo_epi16(lo,hi);
__m256 w0=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nib),b2));
__m256 w1=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nib,8)),b2));
acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i), w0, acc);
acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i+8), w1, acc); }
a=hsum256(acc);
#elif defined(__ARM_NEON)
const uint8x8_t m2v=vdup_n_u8(3); const int8x8_t b2v=vdup_n_s8(2);
float32x4_t ac0=vdupq_n_f32(0), ac1=vdupq_n_f32(0);
for(;i+16<=I;i+=16){ uint32_t wd; memcpy(&wd, w+(i>>2), 4);
uint8x8_t by=vreinterpret_u8_u32(vdup_n_u32(wd));
uint8x8x2_t z01=vzip_u8(vand_u8(by,m2v), vand_u8(vshr_n_u8(by,2),m2v));
uint8x8x2_t z23=vzip_u8(vand_u8(vshr_n_u8(by,4),m2v), vshr_n_u8(by,6));
uint16x4x2_t zz=vzip_u16(vreinterpret_u16_u8(z01.val[0]), vreinterpret_u16_u8(z23.val[0]));
int16x8_t w0=vmovl_s8(vsub_s8(vreinterpret_s8_u16(zz.val[0]),b2v));
int16x8_t w1=vmovl_s8(vsub_s8(vreinterpret_s8_u16(zz.val[1]),b2v));
ac0=vfmaq_f32(ac0, vld1q_f32(xs+i), vcvtq_f32_s32(vmovl_s16(vget_low_s16(w0))));
ac1=vfmaq_f32(ac1, vld1q_f32(xs+i+4), vcvtq_f32_s32(vmovl_s16(vget_high_s16(w0))));
ac0=vfmaq_f32(ac0, vld1q_f32(xs+i+8), vcvtq_f32_s32(vmovl_s16(vget_low_s16(w1))));
ac1=vfmaq_f32(ac1, vld1q_f32(xs+i+12), vcvtq_f32_s32(vmovl_s16(vget_high_s16(w1)))); }
a=vaddvq_f32(vaddq_f32(ac0,ac1));
#endif
for(;i<I;i++){ uint8_t byte=w[i>>2]; int sh=(i&3)*2; a += xs[i]*(float)((int)((byte>>sh)&3)-2); }
y[(int64_t)s*O+o]=a*sc; } }
}
/* ---- IDOT: integer dot kernels (int8-quantized activations) --------------- */
#if defined(__AVX512VNNI__) && defined(__AVX512BW__)
#define IDOT_KERNEL "avx512-vnni"
#elif defined(__AVXVNNI__) && defined(__AVX2__)
#define IDOT_KERNEL "avx-vnni"
#elif defined(__AVX2__)
#define IDOT_KERNEL "avx2"
#elif defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8)
#define IDOT_KERNEL "neon-i8mm"
#elif defined(__ARM_NEON)
#define IDOT_KERNEL "neon"
#elif defined(__VSX__)
#define IDOT_KERNEL "vsx"
#else
#define IDOT_KERNEL "scalar"
#endif
static int g_idot=1;
#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD)
static int g_i4s=1;
#elif defined(__VSX__)
static int g_i4s=1;
#else
static int g_i4s=2;
#endif
static inline float qrow_i8(const float *x, int8_t *q, int I){
float amax=0; for(int i=0;i<I;i++){ float a=fabsf(x[i]); if(a>amax)amax=a; }
float s=amax/127.f; if(s<1e-12f) s=1e-12f; float inv=1.f/s;
for(int i=0;i<I;i++) q[i]=(int8_t)lrintf(x[i]*inv);
return s;
}
/* dot int8*int8 */
static inline int32_t dot_i8i8(const int8_t *w, const int8_t *x, int I){
int32_t sum=0; int i=0;
#if defined(__AVX512VNNI__) && defined(__AVX512BW__)
__m512i acc=_mm512_setzero_si512();
for(;i+64<=I;i+=64){
__m512i wv=_mm512_loadu_si512((const void*)(w+i));
__m512i xv=_mm512_loadu_si512((const void*)(x+i));
__mmask64 neg=_mm512_movepi8_mask(wv);
__m512i xs=_mm512_mask_sub_epi8(xv,neg,_mm512_setzero_si512(),xv);
acc=_mm512_dpbusd_epi32(acc,_mm512_abs_epi8(wv),xs);
}
sum=_mm512_reduce_add_epi32(acc);
#elif defined(__AVXVNNI__) && defined(__AVX2__)
__m128i acc=_mm_setzero_si128();
for(;i+16<=I;i+=16){
__m128i wv=_mm_loadu_si128((const __m128i*)(w+i));
__m128i xv=_mm_loadu_si128((const __m128i*)(x+i));
__m128i xs=_mm_sign_epi8(xv,wv);
acc=_mm_dpbusd_epi32(acc,_mm_abs_epi8(wv),xs);
}
sum=hsum128_i32(acc);
#elif defined(__AVX2__)
__m256i acc=_mm256_setzero_si256(); const __m256i ones=_mm256_set1_epi16(1);
for(;i+32<=I;i+=32){
__m256i wv=_mm256_loadu_si256((const __m256i*)(w+i));
__m256i xv=_mm256_loadu_si256((const __m256i*)(x+i));
__m256i p=_mm256_maddubs_epi16(_mm256_sign_epi8(wv,wv),_mm256_sign_epi8(xv,wv));
acc=_mm256_add_epi32(acc,_mm256_madd_epi16(p,ones));
}
sum=hsum256_i32(acc);
#elif defined(__ARM_NEON)
#if defined(__ARM_FEATURE_DOTPROD)
int32x4_t a0=vdupq_n_s32(0),a1=vdupq_n_s32(0),a2=vdupq_n_s32(0),a3=vdupq_n_s32(0);
for(;i+64<=I;i+=64){
a0=vdotq_s32(a0,vld1q_s8(w+i), vld1q_s8(x+i));
a1=vdotq_s32(a1,vld1q_s8(w+i+16),vld1q_s8(x+i+16));
a2=vdotq_s32(a2,vld1q_s8(w+i+32),vld1q_s8(x+i+32));
a3=vdotq_s32(a3,vld1q_s8(w+i+48),vld1q_s8(x+i+48));
}
int32x4_t acc=vaddq_s32(vaddq_s32(a0,a1),vaddq_s32(a2,a3));
for(;i+16<=I;i+=16) acc=vdotq_s32(acc,vld1q_s8(w+i),vld1q_s8(x+i));
sum=vaddvq_s32(acc);
#else
int32x4_t acc=vdupq_n_s32(0);
for(;i+16<=I;i+=16){
int8x16_t wv=vld1q_s8(w+i), xv=vld1q_s8(x+i);
int16x8_t p=vmull_s8(vget_low_s8(wv),vget_low_s8(xv));
p=vmlal_s8(p,vget_high_s8(wv),vget_high_s8(xv));
acc=vpadalq_s16(acc,p);
}
sum=vaddvq_s32(acc);
#endif
#elif defined(__VSX__)
__vector signed int acc=vec_splats(0);
const __vector signed char vz=vec_splats((signed char)0);
for(;i+16<=I;i+=16){
__vector signed char wv=vec_xl(0,(const signed char*)(w+i));
__vector signed char xv=vec_xl(0,(const signed char*)(x+i));
__vector __bool char neg=vec_cmplt(wv,vz);
__vector signed char xs=vec_sel(xv,vec_sub(vz,xv),neg);
__vector unsigned char wa=(__vector unsigned char)vec_sel(wv,vec_sub(vz,wv),neg);
acc=vec_msum(xs,wa,acc);
}
sum=vec_extract(acc,0)+vec_extract(acc,1)+vec_extract(acc,2)+vec_extract(acc,3);
#endif
for(;i<I;i++) sum+=(int32_t)w[i]*x[i];
return sum;
}
/* dot int4(packed)*int8 */
static inline int32_t dot_i4i8(const uint8_t *w4, const int8_t *x, int I){
int32_t sum=0; int i=0;
#if defined(__AVX512VNNI__) && defined(__AVX512BW__)
const __m256i m4v=_mm256_set1_epi8(0x0F);
const __m512i b8v=_mm512_set1_epi8(8);
const __m512i xidx=_mm512_setr_epi64(0,1,4,5,2,3,6,7);
__m512i acc=_mm512_setzero_si512();
for(;i+64<=I;i+=64){
__m256i by=_mm256_loadu_si256((const __m256i*)(w4+(i>>1)));
__m256i lo=_mm256_and_si256(by,m4v), hi=_mm256_and_si256(_mm256_srli_epi16(by,4),m4v);
__m256i z0=_mm256_unpacklo_epi8(lo,hi), z1=_mm256_unpackhi_epi8(lo,hi);
__m512i wv=_mm512_sub_epi8(_mm512_inserti64x4(_mm512_castsi256_si512(z0),z1,1),b8v);
__m512i xv=_mm512_permutexvar_epi64(xidx,_mm512_loadu_si512((const void*)(x+i)));
__mmask64 neg=_mm512_movepi8_mask(wv);
__m512i xs=_mm512_mask_sub_epi8(xv,neg,_mm512_setzero_si512(),xv);
acc=_mm512_dpbusd_epi32(acc,_mm512_abs_epi8(wv),xs);
}
sum=_mm512_reduce_add_epi32(acc);
#elif defined(__AVXVNNI__) && defined(__AVX2__)
const __m128i m4=_mm_set1_epi8(0x0F); const __m128i b8=_mm_set1_epi8(8);
__m128i acc=_mm_setzero_si128();
for(;i+32<=I;i+=32){
__m128i by=_mm_loadu_si128((const __m128i*)(w4+(i>>1)));
__m128i lo=_mm_and_si128(by,m4), hi=_mm_and_si128(_mm_srli_epi16(by,4),m4);
__m128i n0=_mm_unpacklo_epi8(lo,hi), n1=_mm_unpackhi_epi8(lo,hi);
__m128i w0=_mm_sub_epi8(n0,b8), w1=_mm_sub_epi8(n1,b8);
__m128i x0=_mm_loadu_si128((const __m128i*)(x+i));
__m128i x1=_mm_loadu_si128((const __m128i*)(x+i+16));
acc=_mm_dpbusd_epi32(acc,_mm_abs_epi8(w0),_mm_sign_epi8(x0,w0));
acc=_mm_dpbusd_epi32(acc,_mm_abs_epi8(w1),_mm_sign_epi8(x1,w1));
}
sum=hsum128_i32(acc);
#elif defined(__AVX2__)
const __m128i m4=_mm_set1_epi8(0x0F); const __m256i b8=_mm256_set1_epi8(8);
const __m256i ones=_mm256_set1_epi16(1);
__m256i acc=_mm256_setzero_si256();
for(;i+32<=I;i+=32){
__m128i by=_mm_loadu_si128((const __m128i*)(w4+(i>>1)));
__m128i lo=_mm_and_si128(by,m4), hi=_mm_and_si128(_mm_srli_epi16(by,4),m4);
__m128i n0=_mm_unpacklo_epi8(lo,hi), n1=_mm_unpackhi_epi8(lo,hi);
__m256i wv=_mm256_sub_epi8(_mm256_set_m128i(n1,n0),b8);
__m256i xv=_mm256_loadu_si256((const __m256i*)(x+i));
__m256i p=_mm256_maddubs_epi16(_mm256_sign_epi8(wv,wv),_mm256_sign_epi8(xv,wv));
acc=_mm256_add_epi32(acc,_mm256_madd_epi16(p,ones));
}
sum=hsum256_i32(acc);
#elif defined(__ARM_NEON)
const uint8x16_t m4q=vdupq_n_u8(0x0F); const int8x16_t b8q=vdupq_n_s8(8);
#if defined(__ARM_FEATURE_DOTPROD)
int32x4_t a0=vdupq_n_s32(0),a1=vdupq_n_s32(0),a2=vdupq_n_s32(0),a3=vdupq_n_s32(0);
for(;i+64<=I;i+=64){
uint8x16_t byA=vld1q_u8(w4+(i>>1)), byB=vld1q_u8(w4+(i>>1)+16);
uint8x16x2_t zA=vzipq_u8(vandq_u8(byA,m4q), vshrq_n_u8(byA,4));
uint8x16x2_t zB=vzipq_u8(vandq_u8(byB,m4q), vshrq_n_u8(byB,4));
a0=vdotq_s32(a0,vsubq_s8(vreinterpretq_s8_u8(zA.val[0]),b8q),vld1q_s8(x+i));
a1=vdotq_s32(a1,vsubq_s8(vreinterpretq_s8_u8(zA.val[1]),b8q),vld1q_s8(x+i+16));
a2=vdotq_s32(a2,vsubq_s8(vreinterpretq_s8_u8(zB.val[0]),b8q),vld1q_s8(x+i+32));
a3=vdotq_s32(a3,vsubq_s8(vreinterpretq_s8_u8(zB.val[1]),b8q),vld1q_s8(x+i+48));
}
int32x4_t acc=vaddq_s32(vaddq_s32(a0,a1),vaddq_s32(a2,a3));
for(;i+32<=I;i+=32){
uint8x16_t by=vld1q_u8(w4+(i>>1));
uint8x16x2_t z=vzipq_u8(vandq_u8(by,m4q), vshrq_n_u8(by,4));
acc=vdotq_s32(acc,vsubq_s8(vreinterpretq_s8_u8(z.val[0]),b8q),vld1q_s8(x+i));
acc=vdotq_s32(acc,vsubq_s8(vreinterpretq_s8_u8(z.val[1]),b8q),vld1q_s8(x+i+16));
}
sum=vaddvq_s32(acc);
#else
int32x4_t acc=vdupq_n_s32(0);
for(;i+32<=I;i+=32){
uint8x16_t by=vld1q_u8(w4+(i>>1));
uint8x16x2_t z=vzipq_u8(vandq_u8(by,m4q), vshrq_n_u8(by,4));
int8x16_t w0=vsubq_s8(vreinterpretq_s8_u8(z.val[0]),b8q);
int8x16_t w1=vsubq_s8(vreinterpretq_s8_u8(z.val[1]),b8q);
int8x16_t x0=vld1q_s8(x+i), x1=vld1q_s8(x+i+16);
int16x8_t p=vmull_s8(vget_low_s8(w0),vget_low_s8(x0));
p=vmlal_s8(p,vget_high_s8(w0),vget_high_s8(x0));
acc=vpadalq_s16(acc,p);
p=vmull_s8(vget_low_s8(w1),vget_low_s8(x1));
p=vmlal_s8(p,vget_high_s8(w1),vget_high_s8(x1));
acc=vpadalq_s16(acc,p);
}
sum=vaddvq_s32(acc);
#endif
#elif defined(__VSX__)
const __vector unsigned char m4v=vec_splats((unsigned char)0x0F);
const __vector unsigned char sh4=vec_splats((unsigned char)4);
const __vector signed char b8v=vec_splats((signed char)8);
const __vector signed char vz=vec_splats((signed char)0);
__vector signed int acc=vec_splats(0);
for(;i+32<=I;i+=32){
__vector unsigned char by=vec_xl(0,w4+(i>>1));
__vector unsigned char lo=vec_and(by,m4v), hi=vec_sr(by,sh4);
__vector signed char w0=vec_sub((__vector signed char)vec_mergeh(lo,hi),b8v);
__vector signed char w1=vec_sub((__vector signed char)vec_mergel(lo,hi),b8v);
__vector signed char x0=vec_xl(0,(const signed char*)(x+i));
__vector signed char x1=vec_xl(0,(const signed char*)(x+i+16));
__vector __bool char n0=vec_cmplt(w0,vz), n1=vec_cmplt(w1,vz);
acc=vec_msum(vec_sel(x0,vec_sub(vz,x0),n0),
(__vector unsigned char)vec_sel(w0,vec_sub(vz,w0),n0),acc);
acc=vec_msum(vec_sel(x1,vec_sub(vz,x1),n1),
(__vector unsigned char)vec_sel(w1,vec_sub(vz,w1),n1),acc);
}
sum=vec_extract(acc,0)+vec_extract(acc,1)+vec_extract(acc,2)+vec_extract(acc,3);
#endif
for(;i+1<I;i+=2){ uint8_t b=w4[i>>1]; sum+=((int)(b&0xF)-8)*x[i]+((int)(b>>4)-8)*x[i+1]; }
if(i<I){ uint8_t b=w4[i>>1]; sum+=((int)(b&0xF)-8)*x[i]; }
return sum;
}
/* ---- ARM i8mm SMMLA tiled kernels ---------------------------------------- */
#if defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8)
static inline int32x4_t mm_tile16(int32x4_t acc, int8x16_t wo, int8x16_t wo1,
int8x16_t xs, int8x16_t xs1){
acc=vmmlaq_s32(acc, vcombine_s8(vget_low_s8(wo), vget_low_s8(wo1)),
vcombine_s8(vget_low_s8(xs), vget_low_s8(xs1)));
return vmmlaq_s32(acc, vcombine_s8(vget_high_s8(wo), vget_high_s8(wo1)),
vcombine_s8(vget_high_s8(xs), vget_high_s8(xs1)));
}
static void matmul_q_idot_mm(float *y, const int8_t *xq, const float *sx, const int8_t *q,
const float *scale, int S, int I, int O){
#pragma omp parallel for schedule(static)
for(int o=0;o<(O&~1);o+=2){
const int8_t *wo=q+(int64_t)o*I, *wo1=q+(int64_t)(o+1)*I;
float sc0=scale[o], sc1=scale[o+1];
for(int s=0;s<(S&~1);s+=2){
const int8_t *xs=xq+(int64_t)s*I, *xs1=xq+(int64_t)(s+1)*I;
int32x4_t a0=vdupq_n_s32(0),a1=vdupq_n_s32(0),a2=vdupq_n_s32(0),a3=vdupq_n_s32(0); int i=0;
for(;i+64<=I;i+=64){
a0=mm_tile16(a0,vld1q_s8(wo+i), vld1q_s8(wo1+i), vld1q_s8(xs+i), vld1q_s8(xs1+i));
a1=mm_tile16(a1,vld1q_s8(wo+i+16),vld1q_s8(wo1+i+16),vld1q_s8(xs+i+16),vld1q_s8(xs1+i+16));
a2=mm_tile16(a2,vld1q_s8(wo+i+32),vld1q_s8(wo1+i+32),vld1q_s8(xs+i+32),vld1q_s8(xs1+i+32));
a3=mm_tile16(a3,vld1q_s8(wo+i+48),vld1q_s8(wo1+i+48),vld1q_s8(xs+i+48),vld1q_s8(xs1+i+48));
}
for(;i+16<=I;i+=16)
a0=mm_tile16(a0,vld1q_s8(wo+i),vld1q_s8(wo1+i),vld1q_s8(xs+i),vld1q_s8(xs1+i));
int32x4_t acc=vaddq_s32(vaddq_s32(a0,a1),vaddq_s32(a2,a3));
int32_t d00=vgetq_lane_s32(acc,0), d01=vgetq_lane_s32(acc,1);
int32_t d10=vgetq_lane_s32(acc,2), d11=vgetq_lane_s32(acc,3);
for(;i<I;i++){ int a=wo[i],b=wo1[i],u=xs[i],v=xs1[i];
d00+=a*u; d01+=a*v; d10+=b*u; d11+=b*v; }
y[(int64_t)s*O+o] =(float)d00*sc0*sx[s];
y[(int64_t)s*O+(o+1)] =(float)d10*sc1*sx[s];
y[(int64_t)(s+1)*O+o] =(float)d01*sc0*sx[s+1];
y[(int64_t)(s+1)*O+(o+1)]=(float)d11*sc1*sx[s+1];
}
if(S&1){ int s=S-1; const int8_t *xs=xq+(int64_t)s*I;
y[(int64_t)s*O+o] =(float)dot_i8i8(wo, xs,I)*sc0*sx[s];
y[(int64_t)s*O+(o+1)]=(float)dot_i8i8(wo1,xs,I)*sc1*sx[s]; }
}
if(O&1){ int o=O-1; const int8_t *w=q+(int64_t)o*I; float sc=scale[o];
#pragma omp parallel for schedule(static)
for(int s=0;s<S;s++) y[(int64_t)s*O+o]=(float)dot_i8i8(w,xq+(int64_t)s*I,I)*sc*sx[s]; }
}
static void matmul_i4_idot_mm(float *y, const int8_t *xq, const float *sx, const uint8_t *q4,
const float *scale, int S, int I, int O){
int rb=(I+1)/2;
#pragma omp parallel for schedule(static)
for(int o=0;o<(O&~1);o+=2){
const uint8x16_t m4q=vdupq_n_u8(0x0F); const int8x16_t b8q=vdupq_n_s8(8);
const uint8_t *wo=q4+(int64_t)o*rb, *wo1=q4+(int64_t)(o+1)*rb;
float sc0=scale[o], sc1=scale[o+1];
for(int s=0;s<(S&~1);s+=2){
const int8_t *xs=xq+(int64_t)s*I, *xs1=xq+(int64_t)(s+1)*I;
int32x4_t a0=vdupq_n_s32(0),a1=vdupq_n_s32(0),a2=vdupq_n_s32(0),a3=vdupq_n_s32(0); int i=0;
for(;i+64<=I;i+=64){
uint8x16_t byo=vld1q_u8(wo+(i>>1)), byo1=vld1q_u8(wo1+(i>>1));
uint8x16_t cyo=vld1q_u8(wo+(i>>1)+16), cyo1=vld1q_u8(wo1+(i>>1)+16);
uint8x16x2_t zo =vzipq_u8(vandq_u8(byo, m4q), vshrq_n_u8(byo, 4));
uint8x16x2_t zo1=vzipq_u8(vandq_u8(byo1,m4q), vshrq_n_u8(byo1,4));
uint8x16x2_t ko =vzipq_u8(vandq_u8(cyo, m4q), vshrq_n_u8(cyo, 4));
uint8x16x2_t ko1=vzipq_u8(vandq_u8(cyo1,m4q), vshrq_n_u8(cyo1,4));
a0=mm_tile16(a0, vsubq_s8(vreinterpretq_s8_u8(zo.val[0]),b8q),
vsubq_s8(vreinterpretq_s8_u8(zo1.val[0]),b8q),
vld1q_s8(xs+i), vld1q_s8(xs1+i));
a1=mm_tile16(a1, vsubq_s8(vreinterpretq_s8_u8(zo.val[1]),b8q),
vsubq_s8(vreinterpretq_s8_u8(zo1.val[1]),b8q),
vld1q_s8(xs+i+16), vld1q_s8(xs1+i+16));
a2=mm_tile16(a2, vsubq_s8(vreinterpretq_s8_u8(ko.val[0]),b8q),
vsubq_s8(vreinterpretq_s8_u8(ko1.val[0]),b8q),
vld1q_s8(xs+i+32), vld1q_s8(xs1+i+32));
a3=mm_tile16(a3, vsubq_s8(vreinterpretq_s8_u8(ko.val[1]),b8q),
vsubq_s8(vreinterpretq_s8_u8(ko1.val[1]),b8q),
vld1q_s8(xs+i+48), vld1q_s8(xs1+i+48));
}
for(;i+32<=I;i+=32){
uint8x16_t byo=vld1q_u8(wo+(i>>1)), byo1=vld1q_u8(wo1+(i>>1));
uint8x16x2_t zo =vzipq_u8(vandq_u8(byo, m4q), vshrq_n_u8(byo, 4));
uint8x16x2_t zo1=vzipq_u8(vandq_u8(byo1,m4q), vshrq_n_u8(byo1,4));
a0=mm_tile16(a0, vsubq_s8(vreinterpretq_s8_u8(zo.val[0]),b8q),
vsubq_s8(vreinterpretq_s8_u8(zo1.val[0]),b8q),
vld1q_s8(xs+i), vld1q_s8(xs1+i));
a1=mm_tile16(a1, vsubq_s8(vreinterpretq_s8_u8(zo.val[1]),b8q),
vsubq_s8(vreinterpretq_s8_u8(zo1.val[1]),b8q),
vld1q_s8(xs+i+16), vld1q_s8(xs1+i+16));
}
int32x4_t acc=vaddq_s32(vaddq_s32(a0,a1),vaddq_s32(a2,a3));
int32_t d00=vgetq_lane_s32(acc,0), d01=vgetq_lane_s32(acc,1);
int32_t d10=vgetq_lane_s32(acc,2), d11=vgetq_lane_s32(acc,3);
for(;i+1<I;i+=2){ uint8_t bo=wo[i>>1], bo1=wo1[i>>1];
int a0=(int)(bo&0xF)-8, a1=(int)(bo>>4)-8, b0=(int)(bo1&0xF)-8, b1=(int)(bo1>>4)-8;
int u0=xs[i],u1=xs[i+1],v0=xs1[i],v1=xs1[i+1];
d00+=a0*u0+a1*u1; d01+=a0*v0+a1*v1; d10+=b0*u0+b1*u1; d11+=b0*v0+b1*v1; }
if(i<I){ uint8_t bo=wo[i>>1], bo1=wo1[i>>1];
int a0=(int)(bo&0xF)-8, b0=(int)(bo1&0xF)-8;
d00+=a0*xs[i]; d01+=a0*xs1[i]; d10+=b0*xs[i]; d11+=b0*xs1[i]; }
y[(int64_t)s*O+o] =(float)d00*sc0*sx[s];
y[(int64_t)s*O+(o+1)] =(float)d10*sc1*sx[s];
y[(int64_t)(s+1)*O+o] =(float)d01*sc0*sx[s+1];
y[(int64_t)(s+1)*O+(o+1)]=(float)d11*sc1*sx[s+1];
}
if(S&1){ int s=S-1; const int8_t *xs=xq+(int64_t)s*I;
y[(int64_t)s*O+o] =(float)dot_i4i8(wo, xs,I)*sc0*sx[s];
y[(int64_t)s*O+(o+1)]=(float)dot_i4i8(wo1,xs,I)*sc1*sx[s]; }
}
if(O&1){ int o=O-1; const uint8_t *w=q4+(int64_t)o*rb; float sc=scale[o];
#pragma omp parallel for schedule(static)
for(int s=0;s<S;s++) y[(int64_t)s*O+o]=(float)dot_i4i8(w,xq+(int64_t)s*I,I)*sc*sx[s]; }
}
#endif
/* ---- IDOT dispatch (int8-quantized activations) --------------------------- */
static void matmul_q_idot(float *y, const int8_t *xq, const float *sx, const int8_t *q,
const float *scale, int S, int I, int O){
#if defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8)
if(S>=2){ matmul_q_idot_mm(y,xq,sx,q,scale,S,I,O); return; }
#endif
#pragma omp parallel for schedule(static)
for(int o=0;o<O;o++){ const int8_t *w=q+(int64_t)o*I; float sc=scale[o];
for(int s=0;s<S;s++) y[(int64_t)s*O+o]=(float)dot_i8i8(w,xq+(int64_t)s*I,I)*sc*sx[s]; }
}
static void matmul_i4_idot(float *y, const int8_t *xq, const float *sx, const uint8_t *q4,
const float *scale, int S, int I, int O){
int rb=(I+1)/2;
#if defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8)
if(S>=2){ matmul_i4_idot_mm(y,xq,sx,q4,scale,S,I,O); return; }
#endif
#pragma omp parallel for schedule(static)
for(int o=0;o<O;o++){ const uint8_t *w=q4+(int64_t)o*rb; float sc=scale[o];
for(int s=0;s<S;s++) y[(int64_t)s*O+o]=(float)dot_i4i8(w,xq+(int64_t)s*I,I)*sc*sx[s]; }
}
/* ---- per-thread quantization scratch -------------------------------------- */
typedef struct { int8_t *xq; size_t xq_cap; float *sx; size_t sx_cap; } QScratch;
static _Thread_local QScratch g_qscratch;
static void quant_scratch(size_t xn, size_t sn, int8_t **xq, float **sx){
if(xn>g_qscratch.xq_cap){
int8_t *p=realloc(g_qscratch.xq,xn);
if(!p){ fprintf(stderr,"OOM quant scratch\n"); exit(1); }
g_qscratch.xq=p; g_qscratch.xq_cap=xn;
}
if(sn>g_qscratch.sx_cap){
float *p=realloc(g_qscratch.sx,sn*sizeof(float));
if(!p){ fprintf(stderr,"OOM quant scales\n"); exit(1); }
g_qscratch.sx=p; g_qscratch.sx_cap=sn;
}
*xq=g_qscratch.xq; *sx=g_qscratch.sx;
}
/* ---- f32 -> quantized packing --------------------------------------------- */
static void quantize_rows(const float *w, int8_t *q, float *scale, int O, int I, int bits){
int qmax=(1<<(bits-1))-1;
#pragma omp parallel for schedule(static)
for(int o=0;o<O;o++){ const float *wr=w+(int64_t)o*I; float amax=0;
for(int i=0;i<I;i++){ float a=fabsf(wr[i]); if(a>amax)amax=a; }
float s=amax/qmax; if(s<1e-8f)s=1e-8f; scale[o]=s;
int8_t *qr=q+(int64_t)o*I;
for(int i=0;i<I;i++){ int v=(int)lrintf(wr[i]/s); if(v>qmax)v=qmax; if(v<-qmax-1)v=-qmax-1; qr[i]=(int8_t)v; }
}
}
static void pack_int4(const float *w, uint8_t *q4, float *scale, int O, int I, int bits){
int qmax=(1<<(bits-1))-1, rb=(I+1)/2;
#pragma omp parallel for schedule(static)
for(int o=0;o<O;o++){ const float *wr=w+(int64_t)o*I; float amax=0;
for(int i=0;i<I;i++){ float a=fabsf(wr[i]); if(a>amax)amax=a; }
float s=amax/qmax; if(s<1e-8f)s=1e-8f; scale[o]=s;
uint8_t *qr=q4+(int64_t)o*rb;
for(int i=0;i<I;i+=2){
int v0=(int)lrintf(wr[i]/s); if(v0>qmax)v0=qmax; if(v0<-8)v0=-8;
int v1=0; if(i+1<I){ v1=(int)lrintf(wr[i+1]/s); if(v1>qmax)v1=qmax; if(v1<-8)v1=-8; }
qr[i>>1] = (uint8_t)((v0+8) | ((v1+8)<<4));
}
}
}
static void pack_int2(const float *w, uint8_t *q2, float *scale, int O, int I, int bits){
int qmax=(1<<(bits-1))-1, rb=(I+3)/4;
#pragma omp parallel for schedule(static)
for(int o=0;o<O;o++){ const float *wr=w+(int64_t)o*I; float amax=0;
for(int i=0;i<I;i++){ float a=fabsf(wr[i]); if(a>amax)amax=a; }
float s=amax/qmax; if(s<1e-8f)s=1e-8f; scale[o]=s;
uint8_t *qr=q2+(int64_t)o*rb;
for(int i=0;i<I;i+=4){ uint8_t byte=0;
for(int k=0;k<4 && i+k<I;k++){ int v=(int)lrintf(wr[i+k]/s); if(v>qmax)v=qmax; if(v<-2)v=-2; byte|=(uint8_t)((v+2)<<(k*2)); }
qr[i>>2]=byte;
}
}
}
#endif /* COLI_QUANT_H */
c/ref_olmoe_real.json
+28
View File
@@ -0,0 +1,28 @@
{
"prompt_ids": [
510,
5347,
273,
6181,
310
],
"full_ids": [
510,
5347,
273,
6181,
310,
7785,
15,
187,
187,
510,
3565,
3448,
273,
6181,
310,
5112,
15
]
}
c/resource_plan.py
+107 -4
View File
@@ -203,6 +203,70 @@ def physical_cpu_count():
return os.cpu_count() or 1
def cpu_socket_count():
"""Return the number of physical CPU sockets visible to this process."""
if not sys.platform.startswith("linux"):
return 1
try:
result = subprocess.run(["lscpu", "-p=socket"], text=True,
capture_output=True, check=True, timeout=5)
sockets = {int(line) for line in result.stdout.splitlines()
if line and not line.startswith("#")}
if sockets:
return len(sockets)
except (OSError, ValueError, subprocess.SubprocessError):
pass
return 1
def _auto_tune(bottleneck_class, projected_hit, gpus, cpu_sockets, plan_has_metal):
"""Derive tuning knobs from the bottleneck classification."""
tune = {}
has_gpu = bool(gpus)
n_gpu = len(gpus)
# MTP: costs more than it saves when compute-bound (#389 measured 42% loss)
if bottleneck_class == "compute":
tune["DRAFT"] = {"value": "0",
"reason": "compute-bound: MTP batch overhead exceeds yield"}
elif bottleneck_class == "disk" and projected_hit < 0.90:
tune["DRAFT"] = {"value": "0",
"reason": "low hit rate: MTP widens expert union, adds disk reads"}
# otherwise leave DRAFT unset (engine default: auto)
# PIPE: resident pipeline mode depends on GPU count
if has_gpu and n_gpu == 1:
tune["COLI_CUDA_PIPE"] = {"value": "1",
"reason": "single GPU: S=1 pipeline gate"}
elif has_gpu and n_gpu > 1:
tune["COLI_CUDA_PIPE"] = {"value": "2",
"reason": "multi-GPU: residual stays on-device across layers"}
elif not has_gpu and bottleneck_class == "disk":
tune["PIPE"] = {"value": "1",
"reason": "overlap disk reads with resident expert compute"}
# NUMA: selective interleave for GPU hosts, blanket hint for CPU-only
if cpu_sockets > 1 and has_gpu:
tune["COLI_NUMA"] = {"value": "1",
"reason": "multi-socket + GPU: interleave expert slabs, protect DMA buffers"}
elif cpu_sockets > 1 and not has_gpu:
tune["COLI_NUMA"] = {"value": "1",
"reason": "multi-socket CPU-only: interleave expert slabs across nodes"}
tune["_numa_hint"] = "numactl --interleave=all may perform better on CPU-only hosts"
# OMP: kill hot-thread spin when GPU/Metal owns the power budget
if plan_has_metal:
tune["COLI_NO_OMP_TUNE"] = {"value": "1",
"reason": "Metal: OMP spin-wait steals GPU power budget"}
# PIN: fully resident if RAM allows and no GPU tier competes
if projected_hit >= 0.99 and not has_gpu:
tune["PIN_GB"] = {"value": "all",
"reason": "enough RAM for full expert residency"}
return tune
POLICIES = {
"quality": {"preserve_quantization": True, "preserve_router": True},
"balanced": {"preserve_quantization": True, "preserve_router": True},
@@ -212,11 +276,12 @@ POLICIES = {
def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
available_memory=None, available_disk=None, gpus=None,
policy="quality", physical_cpus=None):
policy="quality", physical_cpus=None, cpu_sockets=None):
if policy not in POLICIES:
raise ValueError(f"unknown policy: {policy}")
info = analyze_model(model)
physical_cpus = physical_cpu_count() if physical_cpus is None else physical_cpus
cpu_sockets = cpu_socket_count() if cpu_sockets is None else cpu_sockets
cfg = info["config"]
available_memory = memory_available() if available_memory is None else available_memory
if available_disk is None:
@@ -273,12 +338,28 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
if cold_bytes:
warnings.append("cold expert misses may reach disk; normal decode speed depends on hit rate")
total_expert = info["expert_bytes"]
resident_expert = hot_bytes + warm_bytes
projected_hit = resident_expert / total_expert if total_expert else 1.0
if cold_bytes:
bottleneck = "disk expert misses"
elif warm_bytes:
bottleneck = "CPU expert compute and RAM bandwidth"
bottleneck_class = "disk"
elif warm_bytes and gpus:
bottleneck = "CPU expert tail and GPU compute"
bottleneck_class = "mixed"
elif projected_hit >= 0.99:
if gpus:
bottleneck = "GPU compute and interconnect"
else:
bottleneck = "CPU expert compute (fully resident)"
bottleneck_class = "compute"
else:
bottleneck = "GPU compute and interconnect"
bottleneck = "CPU expert compute and RAM bandwidth"
bottleneck_class = "memory"
tune = _auto_tune(bottleneck_class, projected_hit, gpus, cpu_sockets,
plan_has_metal=False)
return {
"version": 2,
@@ -286,6 +367,7 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
"quality_preserving": policy != "experimental-fast"},
"model": {key: value for key, value in info.items() if key != "config"},
"cpu": {"physical_cores": max(1, int(physical_cpus)),
"sockets": max(1, int(cpu_sockets)),
"thread_policy": "physical-cores"},
"tiers": {
"disk": {"role": "cold-backing", "model_bytes": info["model_bytes"],
@@ -299,6 +381,9 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0,
"expert_capacity": vram_experts, "requires_host_backing": False},
},
"expected_bottleneck": bottleneck,
"bottleneck_class": bottleneck_class,
"projected_hit_rate": round(projected_hit, 4),
"tune": tune,
"decisions": [
{"target": "VRAM", "reason": "profile-ranked hot experts"},
{"target": "RAM", "reason": "warm experts execute on CPU without quality loss"},
@@ -318,6 +403,11 @@ def environment_for_plan(plan, env=None, cuda_enabled=True):
# ("Affinity not supported on this configuration"): non impostarle li'.
result.setdefault("OMP_PROC_BIND", "spread")
result.setdefault("OMP_PLACES", "cores")
tune = plan.get("tune", {})
for key, entry in tune.items():
if key.startswith("_"):
continue
result.setdefault(key, entry["value"])
if plan["policy"]["name"] == "balanced":
result.setdefault("REPIN", "64")
ram = plan["tiers"]["ram"]
@@ -364,5 +454,18 @@ def format_plan(plan):
else:
lines.append("VRAM no NVIDIA device detected · CPU path")
lines.append(f"limit {plan['expected_bottleneck']}")
hit = plan.get("projected_hit_rate", 0)
lines.append(f"hit {hit:.0%} projected expert residency")
tune = plan.get("tune", {})
if tune:
lines.append("")
lines.append("auto-tune:")
for key, entry in tune.items():
if key.startswith("_"):
continue
lines.append(f" {key}={entry['value']:12s} {entry['reason']}")
hint = tune.get("_numa_hint")
if hint:
lines.append(f" hint: {hint}")
lines.extend(f"warn {warning}" for warning in plan["warnings"])
return "\n".join(lines)
c/sample.h
+153
View File
@@ -0,0 +1,153 @@
/* sample.h — sampling (temperature + nucleus) and stop-set management.
* Header-only: all functions are static — include from the main engine file. */
#ifndef SAMPLE_H
#define SAMPLE_H
#include <math.h>
#include <stdio.h>
#include <string.h>
#include "tok.h"
/* ---- RNG (xorshift64*) -------------------------------------------------- */
static uint64_t g_rng = 0x9E3779B97F4A7C15ULL;
static inline double rndu(void){
g_rng ^= g_rng << 13; g_rng ^= g_rng >> 7; g_rng ^= g_rng << 17;
return (double)(g_rng >> 11) * (1.0 / 9007199254740992.0);
}
/* ---- argmax over a float vector ----------------------------------------- */
static inline int argmax_v(const float *lo, int V){
int b=-1; float bv=-INFINITY;
for(int i=0;i<V;i++){ float x=lo[i]; if(x==x && x>bv){ bv=x; b=i; } }
return b<0?0:b;
}
/* ---- distribution buffers (reused, single-threaded decode) --------------- */
static float *g_pbuf = NULL;
static int *g_pidx = NULL;
/* sift-down on max-heap in h[0..n), key = g_pbuf[h[i]] (#335: partial top-p).
* "hole" variant: carries the root value and deposits only at the end, so
* heapify is O(V) and each pop is O(log n) without qsort on the full vocab. */
static void topp_siftdown(int *h, int n, int i){
int iv = h[i]; float kv = g_pbuf[iv];
for (;;) {
int l = 2*i + 1;
if (l >= n) break;
int b = l; if (l+1 < n && g_pbuf[h[l+1]] > g_pbuf[h[l]]) b = l+1;
if (g_pbuf[h[b]] <= kv) break;
h[i] = h[b]; i = b;
}
h[i] = iv;
}
/* build the target distribution in g_pbuf: softmax(lo/temp) truncated to
* top-p g_nuc. Invariant: g_pbuf stays indexed by token-id (never reordered);
* the truncated tail is zeroed (dist_sample reads by id directly).
* Requires: g_temp, g_nuc, falloc() — declared in the main engine file. */
static void dist_build(const float *lo, int V){
if (!g_pbuf) { g_pbuf = falloc(V); g_pidx = malloc(V * sizeof(int)); }
int mxi = -1; float mx = 0;
for (int i = 0; i < V; i++)
if (isfinite(lo[i]) && (mxi < 0 || lo[i] > mx)) { mx = lo[i]; mxi = i; }
double s = 0; float invt = 1.f / (g_temp > 1e-4f ? g_temp : 1e-4f);
if (mxi >= 0) {
for (int i = 0; i < V; i++) {
g_pbuf[i] = isfinite(lo[i]) ? expf((lo[i] - mx) * invt) : 0.f;
s += g_pbuf[i];
}
}
if (mxi < 0 || !isfinite(s) || s <= 0.0) {
static int warned = 0;
if (!warned) { warned = 1; fprintf(stderr,
"[SAMPLE] warning: non-finite logits (NaN/Inf) — falling back to argmax; "
"output may be degraded. This usually means a numerical blow-up upstream.\n"); }
int a = (mxi >= 0) ? mxi : 0;
for (int i = 0; i < V; i++) g_pbuf[i] = 0.f;
g_pbuf[a] = 1.f;
return;
}
for (int i = 0; i < V; i++) g_pbuf[i] /= (float)s;
if (g_nuc > 0 && g_nuc < 1.f) {
for (int i = 0; i < V; i++) g_pidx[i] = i;
for (int i = V/2-1; i >= 0; i--) topp_siftdown(g_pidx, V, i);
double s2 = 0, cum = 0; int out = V;
do {
int root = g_pidx[0];
g_pidx[0] = g_pidx[--out]; g_pidx[out] = root;
s2 += g_pbuf[root]; cum += g_pbuf[root];
if (out > 0) topp_siftdown(g_pidx, out, 0);
} while (cum < g_nuc && out > 0);
for (int i = 0; i < out; i++) g_pbuf[g_pidx[i]] = 0;
float s2f = (float)s2;
for (int i = out; i < V; i++) g_pbuf[g_pidx[i]] /= s2f;
}
}
/* sample from g_pbuf; ban>=0 excludes that token (renormalizing on the fly) */
static int dist_sample(int V, int ban){
double z = 1.0 - (ban >= 0 ? g_pbuf[ban] : 0.0);
if (z <= 1e-12) z = 1e-12;
double u = rndu() * z, cum = 0;
for (int i = 0; i < V; i++) { if (i == ban) continue; cum += g_pbuf[i]; if (cum >= u) return i; }
for (int i = V-1; i >= 0; i--) if (i != ban && g_pbuf[i] > 0) return i;
return 0;
}
/* next token from logits: greedy if g_temp<=0, sampling otherwise.
* ban = token excluded because it was rejected by speculative verification. */
static int pick_tok(const float *lo, int V, int ban){
if (g_temp <= 0) return argmax_v(lo, V);
dist_build(lo, V);
return dist_sample(V, ban);
}
/* ---- stop set ----------------------------------------------------------- */
static int g_stop[64], g_nstop = 0;
static inline int is_stop(int t){
for (int i = 0; i < g_nstop; i++) if (t == g_stop[i]) return 1;
return 0;
}
/* T=NULL -> config stops only (validation/oracle, where the tokenizer is not needed). */
static void stops_arm_tok(const Cfg *c, int tok_eos, Tok *T){
g_nstop = 0;
for (int i = 0; i < c->n_stop && g_nstop < 64; i++) g_stop[g_nstop++] = c->stop_ids[i];
if (tok_eos >= 0 && !is_stop(tok_eos) && g_nstop < 64) g_stop[g_nstop++] = tok_eos;
int nsp = 0;
if (T) for (int id = 0; id < T->n_ids && g_nstop < 64; id++)
if (T->id_special[id] && !is_stop(id)) { g_stop[g_nstop++] = id; nsp++; }
/* #401: in serve mode keep ONLY <|endoftext|>. Role markers <|user|>/<|observation|>
* (config stops + tokenizer special set) are boundaries the Python server owns; as
* hard stops they cut generation the moment the model opens a <tool_call> block,
* because int4 argmax noise picks a stop-token ID over the correct '<' token. */
if (getenv("SERVE") && tok_eos >= 0) {
int kept = 0;
for (int i = 0; i < g_nstop; i++) if (g_stop[i] == tok_eos) g_stop[kept++] = g_stop[i];
if (kept < g_nstop) fprintf(stderr, "[stop] serve mode: filtered %d non-EOS stop tokens (tool-call safety, #401)\n", g_nstop - kept);
g_nstop = kept; nsp = 0;
}
fprintf(stderr, "[stop] %d stop tokens:", g_nstop);
for (int i = 0; i < g_nstop; i++) fprintf(stderr, " %d", g_stop[i]);
if (nsp) fprintf(stderr, " (%d from the tokenizer's special set)", nsp);
fprintf(stderr, "\n");
}
static void stops_arm(const Cfg *c, int tok_eos){ stops_arm_tok(c, tok_eos, NULL); }
/* ---- log-prob of a target token given the logit vector ------------------- */
static double logprob_target(const float *lo, int V, int target, int *am){
float mx = lo[0]; int best = 0;
for (int i = 1; i < V; i++) if (lo[i] > mx) { mx = lo[i]; best = i; }
double se = 0;
for (int i = 0; i < V; i++) se += exp((double)lo[i] - mx);
if (am) *am = (best == target);
return (double)(lo[target] - mx) - log(se);
}
/* "glm" in model_type, case-insensitive */
static int mt_is_glm(const char *s){
if (s) for (; *s; s++)
if ((s[0]|32) == 'g' && (s[1]|32) == 'l' && (s[2]|32) == 'm') return 1;
return 0;
}
#endif /* SAMPLE_H */
c/setup.sh
+2 -2
View File
@@ -32,11 +32,11 @@ esac
# 2) build: nativa (veloce, per QUESTA macchina). Per un binario da distribuire: make portable
echo " building (ARCH=${ARCH:-native})…"
make -s glm ARCH="${ARCH:-native}"
make -s colibri ARCH="${ARCH:-native}"
# 3) self-test sull'oracolo tiny, se presente
if [ -d glm_tiny ] && [ -f ref_glm.json ]; then
r=$(SNAP=./glm_tiny TF=1 ./glm 64 16 16 2>/dev/null | grep -oE "[0-9]+/[0-9]+ positions" || true)
r=$(SNAP=./glm_tiny TF=1 ./colibri 64 16 16 2>/dev/null | grep -oE "[0-9]+/[0-9]+ positions" || true)
echo " engine self-test: ${r:-?} (expected 32/32)"
fi
c/st.h
+60 -6
View File
@@ -12,6 +12,7 @@
#include <string.h>
#include <stdint.h>
#include <fcntl.h>
#include <errno.h>
#include <unistd.h>
#include <dirent.h>
#include <sys/stat.h>
@@ -104,6 +105,38 @@ static int st_direct_fd(shards *S, int fd) {
}
/* indicizza tutti i model-*.safetensors in snap_dir */
/* pread completo: chunk-loop (una singola pread si ferma a ~2^31 byte su Linux
* — i tensori bf16 grandi la superano), riprova su EINTR e riporta un errore
* ONESTO: perror stampava "Success" su una short-read (errno resta 0), lo
* stesso sintomo corretto in glm.c per #236. ST_PREAD_CHUNK e' sovrascrivibile
* per i test. EN: full pread — chunk loop (one pread caps at ~2^31 bytes and
* big bf16 tensors exceed it), EINTR retry, honest short-read errors.
* Exits on failure, like every st.h reader. */
#ifndef ST_PREAD_CHUNK
#define ST_PREAD_CHUNK (1u << 30)
#endif
static void st_pread_full(int fd, void *buf, int64_t n, int64_t off, const char *tag) {
char *p = (char *)buf;
int64_t got = 0;
while (got < n) {
int64_t want = n - got;
if (want > (int64_t)ST_PREAD_CHUNK) want = ST_PREAD_CHUNK;
ssize_t r = pread(fd, p + got, (size_t)want, off + got);
if (r < 0) {
if (errno == EINTR) continue;
fprintf(stderr, "%s: %s (off %lld, %lld/%lld bytes)\n", tag, strerror(errno),
(long long)off, (long long)got, (long long)n);
exit(1);
}
if (r == 0) {
fprintf(stderr, "%s: short read at EOF (off %lld, %lld/%lld bytes) — truncated file?\n",
tag, (long long)off, (long long)got, (long long)n);
exit(1);
}
got += r;
}
}
static void st_init(shards *S, const char *snap_dir) {
memset(S, 0, sizeof(*S));
S->cap = 4096; S->t = calloc(S->cap, sizeof(st_tensor));
@@ -128,7 +161,7 @@ static void st_init(shards *S, const char *snap_dir) {
if (fstat(fd, &sst) != 0) { perror("fstat shard"); exit(1); }
int64_t fsz = (int64_t)sst.st_size;
uint64_t hlen;
if (pread(fd, &hlen, 8, 0) != 8) { perror("pread hlen"); exit(1); }
st_pread_full(fd, &hlen, 8, 0, "pread hlen");
/* file malevolo/troncato: hlen deve stare nel file dopo gli 8 byte di
* prefisso e sotto il tetto. Senza questo bound hlen+1 puo' andare in
* overflow (malloc(0) e poi hdr[hlen]=0 fuori limiti) o forzare una
@@ -138,7 +171,7 @@ static void st_init(shards *S, const char *snap_dir) {
files[fi], (unsigned long long)hlen, (long long)fsz); exit(1); }
char *hdr = malloc(hlen + 1);
if (!hdr) { perror("malloc safetensors header"); exit(1); }
if (pread(fd, hdr, hlen, 8) != (ssize_t)hlen) { perror("pread hdr"); exit(1); }
st_pread_full(fd, hdr, (int64_t)hlen, 8, "pread hdr");
hdr[hlen] = 0;
int64_t data_start = 8 + (int64_t)hlen;
char *arena = NULL;
@@ -167,7 +200,19 @@ static void st_init(shards *S, const char *snap_dir) {
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;
/* SEC: lo shape viene da un file non fidato (mirror). Senza il guard
* di overflow, uno shape tipo [65535,65535,65535,...] fa avvolgere
* numel a un valore piccolo/negativo che poi passerebbe il cross-check
* numel*esz==nbytes in st_read_f32, riaprendo l'OOB. */
int64_t numel = 1; int bad_shape = 0;
for (int k = 0; k < shp->len; k++) {
int64_t d = (int64_t)shp->kids[k]->num;
if (d < 0 || (d != 0 && numel > INT64_MAX / d)) { bad_shape = 1; break; }
numel *= d;
}
if (bad_shape) {
fprintf(stderr, "%s: tensor '%s' shape overflows int64 — refusing (hostile or corrupt file)\n",
files[fi], name); exit(1); }
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++];
t->name = strdup(name); t->fd = fd; t->off = data_start + a0;
@@ -216,9 +261,18 @@ static void st_prefetch(shards *S, const char *name) {
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); }
/* SEC: numel viene dallo shape, nbytes dagli offset — due campi indipendenti
* del file. Se non concordano, la memcpy F32 (nbytes) o i loop BF16/F16
* (numel elementi da un raw di soli nbytes) sforano il buffer del chiamante,
* che e' dimensionato sul config, non sul file. Il chiamante che alloca su
* st_numel resta coerente; questo blocca l'ingresso ostile a monte. */
int esz = (t->dtype == 2) ? 4 : 2;
if (t->numel < 0 || t->numel > t->nbytes / esz || t->numel * (int64_t)esz != t->nbytes) {
fprintf(stderr, "%s: tensor '%s' shape/bytes mismatch (numel %lld, %lld bytes, dtype %d) — refusing (hostile or corrupt file)\n",
name, name, (long long)t->numel, (long long)t->nbytes, t->dtype); exit(1); }
void *raw = malloc(t->nbytes);
if (!raw) { fprintf(stderr, "malloc %lld bytes for tensor %s failed\n", (long long)t->nbytes, name); exit(1); }
if (pread(t->fd, raw, t->nbytes, t->off) != t->nbytes) { perror("pread data"); exit(1); }
st_pread_full(t->fd, raw, t->nbytes, t->off, "pread data");
if (t->dtype == 2) {
memcpy(out, raw, t->nbytes);
} else if (t->dtype == 0) {
@@ -243,7 +297,7 @@ static int64_t st_nbytes(shards *S, const char *name) {
static void st_read_raw(shards *S, const char *name, void *out, int drop) {
st_tensor *t = st_find(S, name);
if (!t) { fprintf(stderr, "missing tensor: %s\n", name); exit(1); }
if (pread(t->fd, out, t->nbytes, t->off) != t->nbytes) { perror("pread raw"); exit(1); }
st_pread_full(t->fd, out, t->nbytes, t->off, "pread raw");
if (drop) posix_fadvise(t->fd, t->off, t->nbytes, POSIX_FADV_DONTNEED);
}
@@ -256,7 +310,7 @@ static void st_read_slice_f32(shards *S, const char *name, int64_t elem_off, int
int esz = (t->dtype == 2) ? 4 : 2;
int64_t boff = t->off + elem_off * esz, nb = n_elems * esz;
void *raw = malloc(nb);
if (pread(t->fd, raw, nb, boff) != nb) { perror("pread slice"); exit(1); }
st_pread_full(t->fd, raw, nb, boff, "pread slice");
if (t->dtype == 2) memcpy(out, raw, nb);
else if (t->dtype == 0) { uint16_t *p = raw; for (int64_t i = 0; i < n_elems; i++) out[i] = bf16_to_f32(p[i]); }
else { uint16_t *p = raw; for (int64_t i = 0; i < n_elems; i++) out[i] = f16_to_f32(p[i]); }
c/telemetry.h
+189
View File
@@ -0,0 +1,189 @@
/* telemetry.h — dashboard protocol lines, stats/usage persistence, hardware probe.
* Include after Model/Cfg/QT/ESlot/shards and st.h are defined; requires
* qt_bytes(), now_s(), rss_gb(), edisk_s(), and the g_cuda_* globals (ifdef). */
#ifndef TELEMETRY_H
#define TELEMETRY_H
static int64_t tbytes(int O,int I,int bits){
if(bits>=16) return (int64_t)O*I*4;
if(bits>=5) return (int64_t)O*I + (int64_t)O*4;
return (int64_t)O*((I+1)/2) + (int64_t)O*4;
}
static int64_t expert_bytes_probe(Model *m, int ebits){
Cfg *c=&m->c; int64_t eb=0; char nm[256];
snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.0.gate_proj.weight",c->first_dense);
if(st_nbytes(&m->S,nm)>0){
const char *suf[3]={"gate_proj","up_proj","down_proj"};
for(int k=0;k<3;k++){
snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.0.%s.weight",c->first_dense,suf[k]);
eb+=st_nbytes(&m->S,nm);
snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.0.%s.weight.qs",c->first_dense,suf[k]);
int64_t q=st_nbytes(&m->S,nm); if(q>0) eb+=q;
}
}
if(eb<=0) eb = tbytes(c->moe_inter,c->hidden,ebits)*2 + tbytes(c->hidden,c->moe_inter,ebits);
return eb;
}
/* BRAIN MAP: per-turn expert hit bitmap for the dashboard. */
static uint8_t **g_ehit;
static void ehit_mark(Model *m, int layer, int eid){
if(!g_ehit){ Cfg *c=&m->c;
g_ehit=calloc(c->n_layers+1,sizeof(uint8_t*));
for(int i=0;i<=c->n_layers;i++) g_ehit[i]=calloc(c->n_experts,1);
}
g_ehit[layer][eid]=1;
}
/* CPU model + cores + RAM (GB); empty/zero where unavailable. */
static void hw_probe(char *cpu, size_t cn, int *cores, double *ram_total, double *ram_avail){
cpu[0]=0;
#ifdef _WIN32
#if defined(__x86_64__) || defined(__i386__)
{ unsigned int r[12]={0}; unsigned int *w=r;
for(unsigned int f=0x80000002u; f<=0x80000004u; f++,w+=4)
__get_cpuid(f,&w[0],&w[1],&w[2],&w[3]);
char *b=(char*)r; b[47]=0; while(*b==' ')b++;
snprintf(cpu,cn,"%s",b); }
#endif
#else
FILE *ci=fopen("/proc/cpuinfo","r");
if(ci){ char ln[256];
while(fgets(ln,sizeof(ln),ci)) if(!strncmp(ln,"model name",10)){
char *p=strchr(ln,':'); if(p){ p++; while(*p==' ')p++;
int n=(int)strlen(p); if(n>0&&p[n-1]=='\n')p[--n]=0;
snprintf(cpu,cn,"%s",p); } break; }
fclose(ci); }
#endif
*cores=0;
#ifdef _WIN32
{ SYSTEM_INFO si; GetSystemInfo(&si); *cores=(int)si.dwNumberOfProcessors; }
#elif defined(_SC_NPROCESSORS_ONLN)
*cores=(int)sysconf(_SC_NPROCESSORS_ONLN);
#endif
*ram_total=*ram_avail=0;
#ifdef _WIN32
compat_meminfo(ram_total,ram_avail);
#else
FILE *mi=fopen("/proc/meminfo","r");
if(mi){ char ln[256]; double mt=0,ma=0;
while(fgets(ln,sizeof(ln),mi)){
if(sscanf(ln,"MemTotal: %lf",&mt)==1) *ram_total=mt/1e6;
if(sscanf(ln,"MemAvailable: %lf",&ma)==1) *ram_avail=ma/1e6;
} fclose(mi); }
#endif
}
static void hwinfo_emit(Model *m){
Cfg *c=&m->c; (void)c;
char cpu[256]; int cores; double ram_total,ram_avail;
hw_probe(cpu,sizeof(cpu),&cores,&ram_total,&ram_avail);
int ngpu=0; double vram_total=0;
char gpu_name[128]="";
#ifdef COLI_CUDA
ngpu=g_cuda_ndev; vram_total=m->gpu_expert_bytes/1e9;
for(int i=0;i<g_cuda_ndev;i++){
size_t fr=0,to=0; coli_cuda_mem_info(g_cuda_devices[i],&fr,&to);
if(!i) vram_total=(double)to*g_cuda_ndev/1e9;
}
if(g_cuda_ndev>0)
snprintf(gpu_name,sizeof(gpu_name),"CUDA device x%d",g_cuda_ndev);
#endif
printf("HWINFO %d %.1f %.1f %d %.1f %s|%s\n",
cores,ram_total,ram_avail,ngpu,vram_total,cpu,gpu_name);
fflush(stdout);
}
static void tiers_emit(Model *m){
Cfg *c=&m->c; int nsp=0;
for(int i=0;i<c->n_layers;i++) if(m->L[i].sparse) nsp++;
int total=(nsp+(m->has_mtp?1:0))*c->n_experts;
int pinned=0,lru=0;
for(int i=0;i<=c->n_layers;i++){ pinned+=m->npin?m->npin[i]:0; lru+=m->ecn?m->ecn[i]:0; }
int vram=0; double vram_gb=0;
#ifdef COLI_CUDA
vram=m->gpu_expert_count; vram_gb=m->gpu_expert_bytes/1e9;
#endif
int ram=pinned-vram+lru; if(ram<0) ram=0;
int disk=total-vram-ram; if(disk<0) disk=0;
double eb=(double)expert_bytes_probe(m,m->ebits);
printf("TIERS %d %d %d %.2f %.2f\n",vram,ram,disk,vram_gb,ram*eb/1e9);
fflush(stdout);
}
static void emap_emit(Model *m){
Cfg *c=&m->c;
int rows=0;
for(int i=0;i<c->n_layers;i++) if(m->L[i].sparse) rows++;
int has_mtp = m->has_mtp && m->eusage[c->n_layers];
if(has_mtp) rows++;
int cols=c->n_experts;
char *hex=malloc((size_t)rows*cols*2+1); int w=0;
for(int i=0;i<=c->n_layers;i++){
int is_row = (i<c->n_layers && m->L[i].sparse) || (i==c->n_layers && has_mtp);
if(!is_row) continue;
for(int e=0;e<cols;e++){
int tier=0;
ESlot *P=m->pin[i];
for(int z=0;z<m->npin[i];z++) if(P[z].eid==e){
#ifdef COLI_CUDA
tier = P[z].g.cuda?2:1;
#else
tier = 1;
#endif
break; }
if(!tier && m->ecache && m->ecache[i])
for(int z=0;z<m->ecn[i];z++) if(m->ecache[i][z].eid==e){ tier=1; break; }
uint32_t u = m->eusage[i]?m->eusage[i][e]:0;
int heat=0; while(u){ heat++; u>>=1; } if(heat>63) heat=63;
int b=(tier<<6)|heat;
hex[w++]="0123456789abcdef"[b>>4]; hex[w++]="0123456789abcdef"[b&15];
}
}
hex[w]=0;
printf("EMAP %d %d %s\n",rows,cols,hex); fflush(stdout); free(hex);
}
static void hits_emit(Model *m){
Cfg *c=&m->c; if(!g_ehit) return;
int rows=0;
for(int i=0;i<c->n_layers;i++) if(m->L[i].sparse) rows++;
int has_mtp = m->has_mtp && m->eusage[c->n_layers];
if(has_mtp) rows++;
int cols=c->n_experts, nb=(rows*cols+7)/8;
uint8_t *bm=calloc(nb,1); int bit=0;
for(int i=0;i<=c->n_layers;i++){
int is_row = (i<c->n_layers && m->L[i].sparse) || (i==c->n_layers && has_mtp);
if(!is_row) continue;
for(int e=0;e<cols;e++,bit++)
if(g_ehit[i][e]){ bm[bit>>3]|=1<<(bit&7); g_ehit[i][e]=0; }
}
char *hex=malloc((size_t)nb*2+1); int w=0;
for(int b=0;b<nb;b++){ hex[w++]="0123456789abcdef"[bm[b]>>4]; hex[w++]="0123456789abcdef"[bm[b]&15]; }
hex[w]=0;
printf("HITS %d %d %s\n",rows,cols,hex); fflush(stdout); free(hex); free(bm);
}
static void stats_dump_q(Model *m, const char *path, int quiet){
char tmp[2100]; snprintf(tmp,sizeof(tmp),"%s.tmp",path);
FILE *f=fopen(tmp,"w"); if(!f){ if(!quiet) perror(tmp); return; }
Cfg *c=&m->c; int64_t tot=0, nz=0;
for(int i=0;i<=c->n_layers;i++){ if(!m->eusage[i]) continue;
for(int e=0;e<c->n_experts;e++) if(m->eusage[i][e]){ fprintf(f,"%d %d %u\n",i,e,m->eusage[i][e]); tot+=m->eusage[i][e]; nz++; } }
fclose(f); rename(tmp,path);
if(!quiet) fprintf(stderr,"[STATS] %lld selections across %lld distinct experts -> %s\n",(long long)tot,(long long)nz,path);
}
static void stats_dump(Model *m, const char *path){ stats_dump_q(m,path,0); }
static char g_usage_path[2100]="";
static int64_t usage_load(Model *m, const char *path){
FILE *f=fopen(path,"r"); if(!f) return 0;
Cfg *c=&m->c; int l,e; uint32_t cnt; int64_t tot=0;
while(fscanf(f,"%d %d %u",&l,&e,&cnt)==3)
if(l>=0&&l<=c->n_layers&&e>=0&&e<c->n_experts&&m->eusage[l]){ m->eusage[l][e]+=cnt; tot+=cnt; }
fclose(f); return tot;
}
static void usage_save(Model *m){ if(g_usage_path[0]) stats_dump_q(m,g_usage_path,1); }
#endif /* TELEMETRY_H */
c/tests/bench_dsa_select.c
+165
View File
@@ -0,0 +1,165 @@
/* Microbenchmark: old (full-qsort) vs new (quickselect partial-select) DSA top-keep.
*
* This is NOT a unit test -- test_dsa_select.c proves correctness. This measures the
* headline claim of #356: that replacing the O(nk log nk) qsort over all nk context
* scores with an O(nk) partial_select_desc is materially faster per call, which is
* the win the issue was opened for -- and that the win GROWS with context length
* (because quickselect is linear average, qsort is n-log-n).
*
* It re-implements the OLD top-keep inline (qsort + threshold + scans) on a private
* buffer so the A/B runs in one process, same inputs, same warm caches -- a controlled
* comparison. It calls the REAL (new) partial_select_desc via the include-glm.c
* pattern, replicating the production threshold derivation + scans.
*
* Methodology (chosen to be honest, not to flatter the change):
* - keep = 2048 (the real GLM-5.2 index_topk), nk swept across context lengths from
* the 2049 activation boundary up to 65536 (a long conversation).
* - Three score shapes: (a) realistic peaked -- a few hot keys, long tail, the shape
* real DSA attention scores take; (b) uniform random -- no structure; (c) a plateau
* of ties, to exercise the boundary-membership path.
* - Each (shape, nk) is timed over N_REPEAT=2000 iterations, with the scores frozen
* so both algorithms do IDENTICAL work. We report median ns/call and the new/old
* ratio. A warmup pass primes caches before timing.
*
* Run: make tests/bench_dsa_select && ./tests/bench_dsa_select (not in TEST_BINS)
*/
#define main coli_glm_main_unused
#include "../colibri.c"
#undef main
#include <math.h>
#include <stdint.h>
/* ---- the OLD algorithm, verbatim from dev before #356, on a private buffer ---- */
static int cmp_pdesc_old(const void *a, const void *b){
float x=*(const float*)a, y=*(const float*)b; return x<y?1:x>y?-1:0; }
static void keep_old(const float *isc, int nk, int keep, int *dst, int *nd_out){
float *tmp=malloc((size_t)nk*sizeof(float));
memcpy(tmp,isc,(size_t)nk*sizeof(float));
qsort(tmp,(size_t)nk,sizeof(float),cmp_pdesc_old);
float thr=tmp[keep-1]; int nd=0;
for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
free(tmp); *nd_out=nd;
}
/* ---- timing: median of N_REPEAT runs in ns/call, sorted ascending ---- */
#define N_REPEAT 2000
static double bench_ns(void (*fn)(const float*,int,int,int*,int*),
const float *isc, int nk, int keep){
static double ts[N_REPEAT]; int *dst=malloc((size_t)nk*sizeof(int)); int nd;
for(int r=0;r<N_REPEAT;r++){
double t0=now_s();
fn(isc,nk,keep,dst,&nd);
ts[r]=(now_s()-t0)*1e9;
}
for(int a=1;a<N_REPEAT;a++){ double k=ts[a]; int b=a-1;
while(b>=0 && ts[b]>k){ ts[b+1]=ts[b]; b--; } ts[b+1]=k; }
free(dst);
return ts[N_REPEAT/2];
}
/* Sort an array of doubles ascending (median-of-medians aggregation below). */
static void dsort(double *a, int n){
for(int s=1;s<n;s++){ double k=a[s]; int b=s-1;
while(b>=0 && a[b]>k){ a[b+1]=a[b]; b--; } a[b+1]=k; }
}
/* the NEW algorithm calls the real partial_select_desc + replicates the production
* threshold derivation and position scans. */
static void keep_new(const float *isc, int nk, int keep, int *dst, int *nd_out){
float *tmp=malloc((size_t)nk*sizeof(float));
memcpy(tmp,isc,(size_t)nk*sizeof(float));
partial_select_desc(tmp,nk,keep);
float thr=tmp[0]; for(int t=1;t<keep;t++) if(tmp[t]<thr) thr=tmp[t];
int nd=0;
for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
free(tmp); *nd_out=nd;
}
/* deterministic score fill for three shapes */
static uint32_t brng = 0xA5A5A5A5u;
static void brng_seed(uint32_t s){ brng = s; }
static double brand(void){ brng ^= brng << 13; brng ^= brng >> 17; brng ^= brng << 5;
return (double)(brng >> 8) * (1.0 / 16777216.0); }
static void fill_realistic(float *isc, int nk){ /* few hot, long distinct tail */
for(int i=0;i<nk;i++) isc[i]=(float)(-1.0 - brand()*4.0);
isc[0]=3.f; isc[nk/50<nk?nk/50:nk-1]=1.f; isc[nk/200<nk?nk/200:nk-1]=0.5f;
}
static void fill_uniform(float *isc, int nk){ /* no structure */
for(int i=0;i<nk;i++) isc[i]=(float)(brand()*1000.0);
}
static void fill_plateau(float *isc, int nk){ /* tie blocks -> boundary path */
for(int i=0;i<nk;i++) isc[i]=(float)(-(double)(i/7));
}
/* Multi-seed aggregation. Per @KingIcyCreamProjects (#357 thread): with a single frozen
* input per cell, quickselect's deterministic median-of-three pivot means one lucky input
* can spike a single nk row (an observed ~75x at nk=8192 on a 9950X3D that was really a
* ~13-40x algorithm). Two bugs compounded it: (1) the old bench drew ONE input per cell;
* (2) brng was never reset, so each cell's input depended on every prior cell's draws --
* reordering nks[] silently shifted all later inputs. Both fixed here: brng is reseeded
* per (shape, nk, seed), and we take the MEDIAN of N_SEEDS independent inputs, each itself
* a median over N_REPEAT timing reps. A lucky pivot now moves one of the N_SEEDS samples,
* not the reported number. */
int main(void){
int keep = 2048; /* GLM-5.2 index_topk */
int nks[] = {2049, 4096, 8192, 16384, 32768, 65536};
const int N_SEEDS = 11; /* odd so the median is a real sample, not interpolated */
float *isc = malloc((size_t)65536*sizeof(float));
int *dst = malloc((size_t)65536*sizeof(int)); int nd;
double *old_seeds = malloc((size_t)N_SEEDS*sizeof(double));
double *new_seeds = malloc((size_t)N_SEEDS*sizeof(double));
struct { const char *name; void (*fill)(float*,int); } shapes[] = {
{ "realistic", fill_realistic },
{ "uniform", fill_uniform },
{ "plateau", fill_plateau },
};
printf("bench_dsa_select: DSA top-keep, old (qsort) vs new (partial-select) keep=%d (median of %d seeds x %d reps)\n",
keep, N_SEEDS, N_REPEAT);
printf("%-12s %7s %14s %14s %9s\n", "shape", "nk", "old ns/call", "new ns/call", "speedup");
printf("------------------------------------------------------------------------\n");
for(size_t sh=0; sh<sizeof(shapes)/sizeof(shapes[0]); sh++){
for(size_t ni=0; ni<sizeof(nks)/sizeof(nks[0]); ni++){
int nk=nks[ni];
int bad = 0;
for(int sd=0; sd<N_SEEDS; sd++){
/* reseed per (shape,nk,seed) so each cell's input is reproducible and
* independent of cell ordering, and so lucky pivots are sampled, not fixed. */
brng_seed(0xA5A5A5A5u + (uint32_t)(sd*0x9E3779B9u));
shapes[sh].fill(isc,nk);
/* warmup both paths so caches/branch predictors are primed */
for(int w=0; w<50; w++){ keep_old(isc,nk,keep,dst,&nd); keep_new(isc,nk,keep,dst,&nd); }
/* sanity: both must keep exactly `keep` (correctness is test_dsa_select's
* job, but a count divergence here would make the timing meaningless) */
keep_old(isc,nk,keep,dst,&nd); int na=nd;
keep_new(isc,nk,keep,dst,&nd); int nb=nd;
if(na!=keep || nb!=keep){ bad++; continue; }
old_seeds[sd] = bench_ns(keep_old,isc,nk,keep);
new_seeds[sd] = bench_ns(keep_new,isc,nk,keep);
}
if(bad == N_SEEDS){
printf("%-12s %7d (BAD COUNTS on all %d seeds, skipped)\n",
shapes[sh].name, nk, bad);
continue;
}
/* report median-of-seed-medians: robust to a single lucky/unlucky pivot */
dsort(old_seeds, N_SEEDS);
dsort(new_seeds, N_SEEDS);
double t_old = old_seeds[N_SEEDS/2];
double t_new = new_seeds[N_SEEDS/2];
printf("%-12s %7d %14.0f %14.0f %8.2fx\n",
shapes[sh].name, nk, t_old, t_new, t_old/t_new);
}
printf("\n");
}
printf("bench_dsa_select: done (lower ns is better; speedup = old/new)\n");
free(isc); free(dst); free(old_seeds); free(new_seeds);
return 0;
}
c/tests/bench_topp.c
+131
View File
@@ -0,0 +1,131 @@
/* Microbenchmark: old (full-vocab qsort) vs new (heap partial-select) top-p truncation.
*
* This is NOT a unit test -- test_topp.c proves correctness. This measures the headline
* claim of #335: that replacing the O(V log V) qsort over all 151936 vocab entries with
* an O(V) heapify + k*log-V pops is materially faster per call, which is the win the issue
* was opened for.
*
* It re-implements the OLD dist_build inline (qsort + scan) on a private buffer so the A/B
* runs in one process, same inputs, same warm caches -- a controlled comparison. It calls
* the REAL (new) dist_build via the include-glm.c pattern on the global g_pbuf.
*
* Methodology (chosen to be honest, not to flatter the change):
* - V = 151936 (the actual GLM-5.2 vocab), g_nuc swept across the values that matter
* for serving: 0.5 / 0.9 (serve default) / 0.95 / 0.99.
* - Three logit shapes: (a) realistic peaked -- one hot token, long exponential tail,
* the shape real language-model logits take; (b) uniform -- worst case for the heap,
* maximum pop count; (c) a plateau of ties, to exercise the tie path.
* - Each (shape, nuc) is timed over N_REPEAT=2000 iterations, with the RNG/logits frozen
* so both algorithms do IDENTICAL work. We report median ns/call and the new/old ratio.
* - A warmup pass primes caches before timing.
*
* Run: make tests/bench_topp && ./tests/bench_topp (not in TEST_BINS -- not a gate)
*/
#define main coli_glm_main_unused
#include "../colibri.c"
#undef main
#include <math.h>
#include <stdint.h>
/* ---- the OLD algorithm, verbatim from dev before #354, on a private buffer ---- */
static float *s_pbuf; static int *s_pidx; static double *s_ref;
static int cmp_pdesc_old(const void *a, const void *b){
double pa = s_ref[*(const int*)a], pb = s_ref[*(const int*)b];
return pa < pb ? 1 : pa > pb ? -1 : 0; }
static void dist_build_old(const float *lo, int V, double temp, double nuc){
double mx = lo[0]; for (int i = 1; i < V; i++) if (lo[i] > mx) mx = lo[i];
double s = 0, invt = 1.0 / (temp > 1e-4 ? temp : 1e-4);
for (int i = 0; i < V; i++){ s_ref[i] = exp((lo[i]-mx)*invt); s += s_ref[i]; }
for (int i = 0; i < V; i++) s_ref[i] /= s;
if (nuc > 0 && nuc < 1.0){
for (int i = 0; i < V; i++) s_pidx[i] = i;
qsort(s_pidx, V, sizeof(int), cmp_pdesc_old);
double cum = 0; int keep = V;
for (int i = 0; i < V; i++){ cum += s_ref[s_pidx[i]]; if (cum >= nuc){ keep = i+1; break; } }
double s2 = 0;
for (int i = keep; i < V; i++) s_ref[s_pidx[i]] = 0;
for (int i = 0; i < keep; i++) s2 += s_ref[s_pidx[i]];
for (int i = 0; i < keep; i++) s_ref[s_pidx[i]] /= s2;
}
(void)s_pbuf;
}
/* ---- timing: median of N_REPEAT runs in ns/call, sorted ascending ---- */
#define N_REPEAT 2000
static double bench_ns(void (*fn)(const float*,int,double,double),
const float *lo, int V, double temp, double nuc){
static double ts[N_REPEAT];
for (int r = 0; r < N_REPEAT; r++){
double t0 = now_s();
fn(lo, V, temp, nuc);
ts[r] = (now_s() - t0) * 1e9;
}
/* insertion sort the N_REPEAT samples (small), take median */
for (int a = 1; a < N_REPEAT; a++){ double k = ts[a]; int b = a-1;
while (b >= 0 && ts[b] > k){ ts[b+1] = ts[b]; b--; } ts[b+1] = k; }
return ts[N_REPEAT/2];
}
/* the NEW algorithm is the real dist_build, but it writes g_pbuf (not a private buf).
* Wrap it so the bench signature matches, and set the globals it reads. */
static void dist_build_new(const float *lo, int V, double temp, double nuc){
g_temp = (float)temp; g_nuc = (float)nuc;
dist_build(lo, V);
}
/* deterministic logit fill for three shapes */
static uint32_t brng = 0xA5A5A5A5u;
static double brand(void){ brng ^= brng << 13; brng ^= brng >> 17; brng ^= brng << 5;
return (double)(brng >> 8) * (1.0 / 16777216.0); }
static void fill_realistic(float *lo, int V){ /* one hot, exponential tail -- like real logits */
for (int i = 0; i < V; i++) lo[i] = (float)(-4.0 * brand() - (double)i * 0.0001);
lo[0] = 6.f; lo[V/50] = 4.f; lo[V/200] = 3.f;
}
static void fill_uniform(float *lo, int V){ /* worst case for the heap: max pop count */
for (int i = 0; i < V; i++) lo[i] = 0.f;
}
static void fill_plateau(float *lo, int V){ /* ties: blocks of equal value */
for (int i = 0; i < V; i++) lo[i] = (float)(-(double)(i / 50));
}
int main(void){
int V = 151936;
float *lo = malloc((size_t)V * sizeof(float));
s_ref = malloc((size_t)V * sizeof(double));
s_pidx = malloc((size_t)V * sizeof(int));
/* force the new dist_build to allocate g_pbuf/g_pidx at full V once */
g_temp = 0.7f; g_nuc = 0.9f; dist_build(lo, V);
double temp = 0.7;
struct { const char *name; void (*fill)(float*,int); } shapes[] = {
{ "realistic", fill_realistic },
{ "uniform", fill_uniform },
{ "plateau", fill_plateau },
};
double nucs[] = { 0.5, 0.9, 0.95, 0.99 };
printf("bench_topp: top-p truncation, old (qsort) vs new (heap) V=%d temp=%.2f\n", V, temp);
printf("%-12s %6s %14s %14s %9s %9s\n", "shape", "nuc", "old ns/call", "new ns/call", "speedup", "keep");
printf("-----------------------------------------------------------------------------\n");
for (size_t sh = 0; sh < sizeof(shapes)/sizeof(shapes[0]); sh++){
shapes[sh].fill(lo, V);
for (size_t ni = 0; ni < sizeof(nucs)/sizeof(nucs[0]); ni++){
double nuc = nucs[ni];
/* warmup both paths so caches/branch predictors are primed */
for (int w = 0; w < 50; w++){ dist_build_old(lo, V, temp, nuc); dist_build_new(lo, V, temp, nuc); }
double t_old = bench_ns(dist_build_old, lo, V, temp, nuc);
double t_new = bench_ns(dist_build_new, lo, V, temp, nuc);
/* keep count = non-zero entries the new path leaves (== old's keep) */
int keep = 0; for (int i = 0; i < V; i++) if (g_pbuf[i] != 0.f) keep++;
printf("%-12s %6.2f %14.0f %14.0f %8.2fx %9d\n",
shapes[sh].name, nuc, t_old, t_new, t_old / t_new, keep);
}
printf("\n");
}
printf("bench_topp: done (lower ns is better; speedup = old/new)\n");
free(lo); free(s_ref); free(s_pidx);
return 0;
}
c/tests/test_benchmark_cuda_fixture.py
+5 -1
View File
@@ -1,11 +1,12 @@
import unittest
from tools.benchmark_cuda_fixture import parse_output
from tools.benchmark_cuda_fixture import parse_output, parse_p0
SAMPLE = """
REPLAY decode: 4 tokens | 12.34 tok/s
PROFILE: expert-disk 1.25s | expert-matmul 2.50s | attention 0.75s | lm_head 0.10s | other -0.05s
P0-EXEC: routed CPU 1.200s / 123.40 GB/s (456 row) | routed GPU critical 0.150s | router 0.200s | residual P2P 0.030s / 75 hop | orchestration 0.100s
"""
@@ -19,6 +20,9 @@ class ParseOutputTest(unittest.TestCase):
with self.assertRaisesRegex(RuntimeError, "benchmark output missing"):
parse_output("REPLAY decode: 4 tokens | 12.34 tok/s", "engine failed")
def test_extracts_p0_profile(self):
self.assertEqual(parse_p0(SAMPLE), [1.2, 123.4, 456.0, 0.15, 0.2, 0.03, 75.0, 0.1])
if __name__ == "__main__":
unittest.main()
c/tests/test_dsa_select.c
+223
View File
@@ -0,0 +1,223 @@
/* DSA top-keep partial-select: the quickselect rewrite (#356) must produce a
* BIT-IDENTICAL kept-position set to the old full-vocab qsort, for every score
* shape the attention indexer can see.
*
* Why this test exists (#356): attention_rows() selects the top-`keep` context
* keys (index_topk=2048 on GLM-5.2) to attend to. It previously did this by
* full-qsorting all `nk` scores (O(nk log nk)) and reading tmp[keep-1] as the
* threshold. It now does a partial_select_desc (quickselect, O(nk) average) and
* takes the threshold as the min of the selected top-keep block. The contract is
* subtle but STRONGER than test_topp's:
*
* The two position-order scans that build dst[] --
* for t: if isc[t] > thr -> keep (strictly above threshold)
* for t: if isc[t] == thr -> keep (ties, in position order)
* -- are UNCHANGED by the rewrite. So if the threshold value is identical,
* the kept-position set is identical element-by-element (not just as a
* multiset, which is all the unstable sampling heap in #335 could promise).
*
* Strategy: drive the REAL partial_select_desc (via the include-glm.c pattern)
* and replicate the production threshold derivation + scans, then compare the
* resulting dst[] against an INDEPENDENT reference that re-implements the OLD
* algorithm (full qsort + tmp[keep-1] threshold) on a private buffer. The kept
* sets must be element-wise equal on every shape, including tie plateaus where
* the boundary membership is decided by the position scan.
*
* We also directly unit-test partial_select_desc's partition invariant: after
* the call, max(a[keep..n)) <= min(a[0..keep)) -- i.e. the keep largest really
* did land in the prefix. This catches a broken quickselect even before the
* end-to-end comparison.
*
* In-memory only (no scratch files), so it builds clean on the Windows MinGW CI
* job without the unmerged compat shim. */
#define main coli_glm_main_unused
#include "../colibri.c"
#undef main
#include <math.h>
static int g_nfails = 0;
#define FAIL(fmt, ...) do { \
fprintf(stderr, " FAIL [%s nk=%d keep=%d shape=%s]: " fmt "\n", \
label, nk, keep, shape_name, ##__VA_ARGS__); \
g_nfails++; \
return; \
} while (0)
/* ---- independent reference: the OLD algorithm (full qsort + tmp[keep-1]) ---- */
/* qsort comparator matching the production cmp_fdesc exactly (desc, unstable). */
static int cmp_ref_desc(const void *a, const void *b){
float x=*(const float*)a, y=*(const float*)b; return x<y?1:x>y?-1:0; }
/* Reproduce the OLD glm.c:2589-2596 exactly: copy, qsort desc, threshold =
* tmp[keep-1], then the two position-order scans into dst[]. Returns nd. */
static int keep_old(const float *isc, int nk, int keep, int *dst){
float *tmp=malloc((size_t)nk*sizeof(float));
memcpy(tmp,isc,(size_t)nk*sizeof(float));
qsort(tmp,(size_t)nk,sizeof(float),cmp_ref_desc);
float thr=tmp[keep-1];
int nd=0;
for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
free(tmp);
return nd;
}
/* Reproduce the NEW glm.c path: partial_select desc, threshold = min of the
* selected block, same two scans. Uses the REAL partial_select_desc from glm.c. */
static int keep_new(const float *isc, int nk, int keep, int *dst){
float *tmp=malloc((size_t)nk*sizeof(float));
memcpy(tmp,isc,(size_t)nk*sizeof(float));
partial_select_desc(tmp,nk,keep);
float thr=tmp[0]; for(int t=1;t<keep;t++) if(tmp[t]<thr) thr=tmp[t];
int nd=0;
for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
free(tmp);
return nd;
}
/* ---- direct unit test of the partition invariant ---- */
static void check_partition(const char *label, const float *isc, int nk, int keep,
const char *shape_name){
float *tmp=malloc((size_t)nk*sizeof(float));
memcpy(tmp,isc,(size_t)nk*sizeof(float));
partial_select_desc(tmp,nk,keep);
/* invariant: every element of tmp[0..keep) is >= every element of tmp[keep..n).
* (>=, not >: equal values may sit on either side of the partition boundary,
* which is fine -- the threshold is the MIN of the prefix, and the position
* scan handles ties.) */
float top_min=INFINITY, tail_max=-INFINITY;
for(int i=0;i<keep;i++) if(tmp[i]<top_min) top_min=tmp[i];
for(int i=keep;i<nk;i++) if(tmp[i]>tail_max) tail_max=tmp[i];
if(!(top_min >= tail_max))
FAIL("partition invariant violated: top_min=%.9g < tail_max=%.9g", top_min, tail_max);
free(tmp);
}
/* ---- end-to-end: old vs new kept-set must be element-wise identical ---- */
static void check_case(const char *label, int nk, int keep, const char *shape_name,
const float *isc){
int *da=malloc((size_t)nk*sizeof(int));
int *db=malloc((size_t)nk*sizeof(int));
int na=keep_old(isc,nk,keep,da);
int nb=keep_new(isc,nk,keep,db);
/* 1. both keep exactly `keep` positions (the contract: keep the top-keep by
* count). A count mismatch is a real bug, not a tie artifact. */
if(na!=keep) FAIL("old kept %d, expected %d (old path is the reference)", na, keep);
if(nb!=keep) FAIL("new kept %d, expected %d", nb, keep);
if(na!=nb) FAIL("keep-count mismatch: old=%d new=%d", na, nb);
/* 2. element-wise identical dst[]. This is the strong contract: because the
* threshold is derived identically and the position-order scans are byte-
* for-byte the same, the kept SET and its ORDER must match exactly. (This
* is what makes #356 cleaner than #335, which was multiset-only.) */
int first_diff=-1;
for(int i=0;i<na;i++){ if(da[i]!=db[i]){ first_diff=i; break; } }
if(first_diff>=0)
FAIL("kept-set differs at index %d: old dst[%d]=%d new dst[%d]=%d",
first_diff, first_diff, da[first_diff], first_diff, db[first_diff]);
/* 3. also check the partition invariant directly (catches a subtly broken
* quickselect even if the threshold happened to come out right). */
check_partition(label,isc,nk,keep,shape_name);
free(da); free(db);
printf(" ok [nk=%d keep=%d shape=%s]\n", nk, keep, shape_name);
}
#undef FAIL
/* deterministic xorshift32 RNG (matches the test_i4_grouped.c / test_topp.c convention) */
static uint32_t rng_state = 0x12345678u;
static uint32_t xr(void){ rng_state ^= rng_state << 13; rng_state ^= rng_state >> 17;
rng_state ^= rng_state << 5; return rng_state; }
static double frand(void){ return (xr() >> 8) * (1.0 / 16777216.0); } /* [0,1) */
/* fill scores for a given shape. Shapes stress the threshold boundary and the
* quickselect's median-of-three pivot differently. */
static void fill_shape(float *isc, int nk, int shape){
switch(shape){
case 0: /* uniform random distinct (no ties): the clean contract case */
for(int i=0;i<nk;i++) isc[i]=(float)(frand()*1000.0); break;
case 1: /* peaked: a few hot, long distinct tail (realistic attention shape) */
for(int i=0;i<nk;i++) isc[i]=(float)(-1.0 - frand()*4.0);
isc[0]=3.f; if(nk>3) isc[nk/3]=1.f; if(nk>2) isc[nk/2]=0.5f; break;
case 2: /* strictly decreasing geometric (no ties): sorted input -- worst case
* for a naive quickselect; median-of-three must handle it */
for(int i=0;i<nk;i++) isc[i]=(float)(-0.001*(double)i); break;
case 3: /* strictly increasing (reverse-sorted): the other quickselect worst case */
for(int i=0;i<nk;i++) isc[i]=(float)(0.001*(double)i); break;
case 4: /* plateau ties: blocks of equal value -> boundary membership decided
* entirely by the position scan (exercises the ==thr path) */
for(int i=0;i<nk;i++) isc[i]=(float)(-(double)(i/7)); break;
case 5: /* all-equal: every value identical -> degenerate threshold, all kept
* via the ==thr scan; quickselect must not infinite-loop or corrupt */
for(int i=0;i<nk;i++) isc[i]=5.f; break;
}
}
int main(void){
/* Sizes around the real index_topk=2048 boundary, plus small cases that
* exercise the k>=n / k==1 / k==n edges. */
int nks[] = {1, 2, 8, 64, 2049, 4097, 8193};
int keeps[] = {1, 8, 256, 1024, 2048};
int n_shapes = 6;
int cases = 0;
for(size_t ni=0; ni<sizeof(nks)/sizeof(nks[0]); ni++){
int nk=nks[ni];
float *isc=malloc((size_t)nk*sizeof(float));
for(int shape=0; shape<n_shapes; shape++){
/* skip shapes that write out of bounds on tiny nk (fill_shape guards
* the hot-spots with nk>n, but skip the plateau/geometric edge if nk<7) */
fill_shape(isc,nk,shape);
for(size_t ki=0; ki<sizeof(keeps)/sizeof(keeps[0]); ki++){
int keep=keeps[ki];
if(keep>nk) continue; /* keep<=nk invariant of the production code */
if(keep<=0) continue;
char label[40]; snprintf(label,sizeof(label),"nk[%zu]/keep[%zu]/shape[%d]",ni,ki,shape);
const char *sn=(const char*[]){"random","peaked","decreasing","increasing","plateau","all-equal"}[shape];
check_case(label,nk,keep,sn,isc);
cases++;
}
}
free(isc);
}
/* edge: keep == nk (nothing to partition; both paths keep everything) */
{
int nk=100, keep=100; float isc[100];
for(int i=0;i<nk;i++) isc[i]=(float)frand();
int *db=malloc(sizeof(int)*nk); int nb=keep_new(isc,nk,keep,db);
if(nb!=nk){ fprintf(stderr," FAIL [keep==nk]: kept %d expected %d\n",nb,nk); g_nfails++; }
else printf(" ok [keep==nk nk=%d]\n",nk);
free(db); cases++;
}
/* edge: keep == 1 (threshold = the single max; quickselect must find it) */
{
int nk=500, keep=1; float isc[500];
for(int i=0;i<nk;i++) isc[i]=(float)frand();
int da[1],db[1]; int na=keep_old(isc,nk,keep,da), nb=keep_new(isc,nk,keep,db);
if(na!=1||nb!=1||da[0]!=db[0]){
fprintf(stderr," FAIL [keep==1]: old={%d (n=%d)} new={%d (n=%d)}\n",da[0],na,db[0],nb); g_nfails++; }
else printf(" ok [keep==1 argmax=%d]\n",db[0]); cases++;
}
/* edge: all-equal scores, keep in the middle -> every kept slot is a tie;
* the position scan must pick positions 0..keep-1 deterministically */
{
int nk=1000, keep=500; float isc[1000];
for(int i=0;i<nk;i++) isc[i]=3.14f;
int *db=malloc(sizeof(int)*nk); int nb=keep_new(isc,nk,keep,db);
int bad=0; for(int i=0;i<nb;i++) if(db[i]!=i) bad=1;
if(nb!=keep||bad){ fprintf(stderr," FAIL [all-equal keep=%d]: nb=%d bad=%d\n",keep,nb,bad); g_nfails++; }
else printf(" ok [all-equal keep=%d -> positions 0..%d]\n",keep,keep-1); cases++;
free(db);
}
printf("\ntest_dsa_select: %d cases run, %d failure(s)\n", cases, g_nfails);
if(g_nfails){ printf("test_dsa_select: FAIL\n"); return 1; }
printf("test_dsa_select: ok\n");
return 0;
}
c/tests/test_i4_grouped.c
+201
View File
@@ -0,0 +1,201 @@
/* Exactness test for the grouped-int4 kernel (fmt=4, one f32 scale per `gs`
* elements along I) against a plain-C reference that dequantizes and multiplies
* in double.
*
* Why this test exists: matmul_i4_grouped is the REFERENCE the CUDA fmt=4 path
* (#298) is expected to reproduce, and it had no test of its own. Debugging a
* backend against an unverified oracle means two moving targets. Anyone porting
* fmt=4 to a new backend can now diff against a kernel that is known exact here.
*
* Covers: I a clean multiple of gs, I with a partial last group (the `glen`
* clamp), odd I (the nibble tail), gs larger than I (single group), and the
* nibble edges 0 and 15 (which decode to -8 and +7 — an offset encoding, NOT
* two's complement; getting this backwards is silent and looks like noise).
*
* FP note: the kernel sums each group in f32 (AVX2 accumulator + scalar tail)
* while the reference sums in double, so we compare against a relative epsilon
* rather than bit-exactly. The tolerance is tight enough that a wrong scale
* index, a wrong group boundary or a swapped nibble cannot hide under it —
* those are O(1) relative errors, not O(1e-6). */
#define main coli_glm_main_unused
#include "../colibri.c"
#undef main
static uint32_t rng_state=0xC0FFEEu;
static uint32_t xr(void){ rng_state^=rng_state<<13; rng_state^=rng_state>>17; rng_state^=rng_state<<5; return rng_state; }
static float frand(void){ return (float)((int)(xr()%2001)-1000)/1000.0f; }
/* Reference: dequantize nibble -> (v-8)*scale[group], accumulate in double.
* Deliberately the dumbest possible expression of the format. */
static void ref_grouped(double *y, double *mag, const float *x, const uint8_t *q4,
const float *scale, int S, int I, int O, int gs){
int rb=(I+1)/2, ng=(I+gs-1)/gs;
for(int o=0;o<O;o++){
const uint8_t *w=q4+(int64_t)o*rb;
const float *scl=scale+(int64_t)o*ng;
for(int s=0;s<S;s++){
const float *xs=x+(int64_t)s*I; double a=0, m=0;
for(int i=0;i<I;i++){
uint8_t byte=w[i>>1];
int nib=(i&1)?(int)(byte>>4):(int)(byte&0xF);
double term=(double)xs[i] * (double)(nib-8) * (double)scl[i/gs];
a += term; m += fabs(term);
}
y[(int64_t)s*O+o]=a;
/* Sum of |terms|: the scale the f32 rounding error actually lives on.
* Comparing against |result| instead would flag pure cancellation --
* a dot product of signed terms can land near zero, and then a 1e-6
* absolute error reads as a 1e-3 relative one. That is the accumulator's
* precision, not a kernel defect. A wrong scale index or group boundary
* shifts the result by a fraction OF THE TERMS, so it is caught here. */
mag[(int64_t)s*O+o]=m;
}
}
}
static int check(const char *name, int S, int I, int O, int gs, int fill_edges){
int rb=(I+1)/2, ng=(I+gs-1)/gs;
uint8_t *q4=malloc((size_t)O*rb);
float *scale=malloc((size_t)O*ng*sizeof(float));
float *x=malloc((size_t)S*I*sizeof(float));
float *y=malloc((size_t)S*O*sizeof(float));
double *yr=malloc((size_t)S*O*sizeof(double));
double *ym=malloc((size_t)S*O*sizeof(double));
if(!q4||!scale||!x||!y||!yr||!ym){ fprintf(stderr,"%s: OOM\n",name); return 1; }
for(size_t i=0;i<(size_t)O*rb;i++) q4[i]=(uint8_t)(xr()&0xFF);
if(fill_edges){
/* nibble extremes: 0x0F -> +7, 0x00 -> -8. A two's-complement misread
* turns 15 into -1 instead of +7 and the error is data-dependent noise. */
for(size_t i=0;i<(size_t)O*rb && i<64;i++) q4[i]=(i&1)?0x00:0xFF;
}
/* scales span a few orders of magnitude: a wrong group index shows up big */
for(int i=0;i<O*ng;i++) scale[i]=(0.001f+(float)(xr()%1000)/1000.0f)*((xr()&1)?1.f:-1.f);
for(int i=0;i<S*I;i++) x[i]=frand();
matmul_i4_grouped(y,x,q4,scale,S,I,O,gs);
ref_grouped(yr,ym,x,q4,scale,S,I,O,gs);
int bad=0; double worst=0;
for(int i=0;i<S*O;i++){
double d=fabs((double)y[i]-yr[i]);
double rel = ym[i]>1e-30 ? d/ym[i] : d; /* error relative to the summed magnitude */
if(rel>worst) worst=rel;
if(rel>1e-6){
if(bad<3) fprintf(stderr,"%s: [%d] got %.9g want %.9g (|terms| %.3g, rel %.3g)\n",
name,i,(double)y[i],yr[i],ym[i],rel);
bad++;
}
}
free(q4);free(scale);free(x);free(y);free(yr);free(ym);
if(bad){ fprintf(stderr,"%s: FAIL (%d/%d mismatched, worst rel %.3g)\n",name,bad,S*O,worst); return 1; }
printf(" %-42s ok (S=%d I=%d O=%d gs=%d ng=%d, worst rel %.2g)\n",name,S,I,O,gs,ng,worst);
return 0;
}
/* matmul_i4_grouped_pair (fused gate+up, #298) reads x once instead of twice.
* Checked two ways, because "identical" is only true where it can be:
*
* - Correctness, always: both outputs must match the double reference within
* the same magnitude-relative epsilon as the unfused kernel.
* - Bit-exactness, only when I % gs == 0: then every group is covered by the
* AVX2 body, whose accumulation order is identical to the unfused kernel, so
* the results agree to the last bit. This is the shape the real g64
* checkpoints have (I = 2048 / 6144, gs = 64), i.e. the production path.
*
* With a PARTIAL last group the group tail falls to scalar code, and the
* compiler is free to contract/reassociate the fused body differently from the
* single-matrix one. The results then differ by ~1e-7 -- rounding, not logic
* (which of gate/up "differs" is arbitrary, the tell that it is FP luck).
* Demanding bit-exactness there would report a compiler artifact as a bug. */
#ifdef COLI_HAVE_GROUPED_PAIR
static int check_pair(const char *name, int S, int I, int O, int gs){
int rb=(I+1)/2, ng=(I+gs-1)/gs;
uint8_t *qg=malloc((size_t)O*rb), *qu=malloc((size_t)O*rb);
float *sg=malloc((size_t)O*ng*sizeof(float)), *su=malloc((size_t)O*ng*sizeof(float));
float *x=malloc((size_t)S*I*sizeof(float));
float *yg=malloc((size_t)S*O*sizeof(float)), *yu=malloc((size_t)S*O*sizeof(float));
float *rg=malloc((size_t)S*O*sizeof(float)), *ru=malloc((size_t)S*O*sizeof(float));
double *dg=malloc((size_t)S*O*sizeof(double)), *du=malloc((size_t)S*O*sizeof(double));
double *mg=malloc((size_t)S*O*sizeof(double)), *mu=malloc((size_t)S*O*sizeof(double));
if(!qg||!qu||!sg||!su||!x||!yg||!yu||!rg||!ru||!dg||!du||!mg||!mu){ fprintf(stderr,"%s: OOM\n",name); return 1; }
for(size_t i=0;i<(size_t)O*rb;i++){ qg[i]=(uint8_t)(xr()&0xFF); qu[i]=(uint8_t)(xr()&0xFF); }
for(int i=0;i<O*ng;i++){ sg[i]=frand(); su[i]=frand(); }
for(int i=0;i<S*I;i++) x[i]=frand();
matmul_i4_grouped_pair(yg,yu,x,qg,sg,qu,su,S,I,O,gs);
matmul_i4_grouped(rg,x,qg,sg,S,I,O,gs);
matmul_i4_grouped(ru,x,qu,su,S,I,O,gs);
ref_grouped(dg,mg,x,qg,sg,S,I,O,gs);
ref_grouped(du,mu,x,qu,su,S,I,O,gs);
int bad=0, exact=1; double worst=0;
for(int i=0;i<S*O;i++){
double eg = mg[i]>1e-30 ? fabs((double)yg[i]-dg[i])/mg[i] : fabs((double)yg[i]-dg[i]);
double eu = mu[i]>1e-30 ? fabs((double)yu[i]-du[i])/mu[i] : fabs((double)yu[i]-du[i]);
if(eg>worst) worst=eg;
if(eu>worst) worst=eu;
if(eg>1e-6||eu>1e-6){
if(bad<3) fprintf(stderr,"%s: [%d] gate %.9g/%.9g up %.9g/%.9g (rel %.3g/%.3g)\n",
name,i,(double)yg[i],dg[i],(double)yu[i],du[i],eg,eu);
bad++;
}
if(yg[i]!=rg[i]||yu[i]!=ru[i]) exact=0;
}
/* Aligned shapes run entirely through the AVX2 body: same order as unfused,
* so bit-exactness is a real invariant there and worth asserting. */
if(I%gs==0 && !exact){
fprintf(stderr,"%s: FAIL fused != unfused bitwise on an ALIGNED shape "
"(no scalar tail runs here; the orders must match)\n",name);
bad++;
}
free(qg);free(qu);free(sg);free(su);free(x);free(yg);free(yu);free(rg);free(ru);
free(dg);free(du);free(mg);free(mu);
if(bad){ fprintf(stderr,"%s: FAIL (%d mismatched)\n",name,bad); return 1; }
printf(" %-42s ok (S=%d I=%d O=%d gs=%d, worst rel %.2g%s)\n",name,S,I,O,gs,worst,
I%gs==0?", bit-exact vs unfused":"");
return 0;
}
#endif
int main(void){
int fail=0;
printf("test_i4_grouped: matmul_i4_grouped vs plain-C dequant reference\n");
/* the shape the g64 checkpoints actually use */
fail|=check("gs=64, I multiple of gs", 2, 512, 8, 64, 0);
fail|=check("gs=64, single row single token", 1, 128, 1, 64, 0);
fail|=check("gs=64, nibble edges (0x00/0xFF)", 1, 256, 4, 64, 1);
/* partial last group: glen clamp, the classic off-by-one */
fail|=check("gs=64, partial last group (I=200)", 2, 200, 4, 64, 0);
fail|=check("gs=64, I just over a group (I=65)", 1, 65, 3, 64, 0);
fail|=check("gs=64, I one under a group (I=63)", 1, 63, 3, 64, 0);
/* odd I: the scalar nibble tail (i+1 == I) */
fail|=check("gs=64, odd I (I=201)", 2, 201, 4, 64, 0);
fail|=check("gs=16, odd I (I=33)", 1, 33, 2, 16, 0);
/* gs > I: everything in one group */
fail|=check("gs=128 > I=64 (single group)", 1, 64, 4, 128, 0);
/* the other documented group size */
fail|=check("gs=128, I multiple of gs", 2, 512, 4, 128, 0);
/* batch: S>1 exercises the per-s inner loop against a shared scale row */
fail|=check("gs=64, batch S=8", 8, 320, 6, 64, 0);
#ifdef COLI_HAVE_GROUPED_PAIR
printf("test_i4_grouped: matmul_i4_grouped_pair (fused gate+up) vs two separate calls\n");
fail|=check_pair("pair: gs=64, I multiple of gs", 2, 512, 8, 64);
fail|=check_pair("pair: gs=64, partial last group",2, 200, 4, 64);
fail|=check_pair("pair: gs=64, odd I (I=201)", 2, 201, 4, 64);
fail|=check_pair("pair: gs=64, decode S=1", 1, 320, 6, 64);
fail|=check_pair("pair: gs=128, I=512", 2, 512, 4, 128);
#endif
if(fail){ printf("test_i4_grouped: FAIL\n"); return 1; }
printf("test_i4_grouped: ok\n");
return 0;
}
c/tests/test_idot.c
+1 -1
View File
@@ -7,7 +7,7 @@
* (sign-trick kernels must treat |−128| as 128 unsigned, not saturate to 127),
* and random data at qrow_i8's contract (|x| <= 127, w full int8 range). */
#define main coli_glm_main_unused
#include "../glm.c"
#include "../colibri.c"
#undef main
static uint32_t rng_state=0x12345678u;
c/tests/test_kv_alloc.c
+1 -1
View File
@@ -4,7 +4,7 @@
* 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"
#include "../colibri.c"
#undef main
int main(void){
c/tests/test_logit_nan.c
+70
View File
@@ -0,0 +1,70 @@
/* Regression for non-finite-logit poisoning of sampling.
*
* A single NaN or +Inf in the logits (a bad streamed expert tile, or an fp
* overflow in the matmul at a low-RAM eviction boundary) used to make softmax
* produce an all-NaN g_pbuf; dist_sample then never satisfied cum>=u and fell
* through to return token 0 — so the engine silently emitted an unbroken run of
* token 0 with no error, on the DEFAULT serve path (TEMP>0, 0<NUCLEUS<1).
*
* Fix under test: argmax_v() skips NaN (picks the max finite/+Inf entry instead
* of being pinned to index 0), and dist_build() detects a non-finite softmax sum
* and collapses to a one-hot over the finite argmax (warning once) rather than
* dividing every entry into NaN. Degrade + diagnose, never silently corrupt.
*
* No model file needed: exercises argmax_v / dist_build / dist_sample directly. */
#include <assert.h>
#include <math.h>
#define main coli_glm_main_unused
#include "../colibri.c"
#undef main
static int approx1(double x){ return x > 0.999 && x < 1.001; }
int main(void){
/* --- argmax_v must skip NaN (greedy decode + speculative-verify paths) --- */
{ float lo[8]={NAN,1.f,5.f,2.f,NAN,-3.f,4.f,0.f};
assert(argmax_v(lo,8)==2 && "pick max finite (idx2=5.0), not NaN-pinned idx0"); }
{ float lo[8]={3.f,INFINITY,1.f,2.f,0.f,-1.f,2.5f,1.5f};
assert(argmax_v(lo,8)==1 && "pick the +Inf position"); }
{ float lo[8]; for(int i=0;i<8;i++) lo[i]=NAN;
assert(argmax_v(lo,8)==0 && "all-NaN: no crash, defined fallback"); }
g_temp=0.7f; g_nuc=0.9f; /* the default serve/chat sampling path */
/* --- dist_build: a NaN logit must yield a finite one-hot, not all-NaN --- */
{ float lo[8]={0.5f,1.f,NAN,8.f,0.2f,-1.f,0.f,0.3f}; /* max finite = idx3 (8.0) */
dist_build(lo,8);
double sum=0; int nan=0;
for(int i=0;i<8;i++){ if(!(g_pbuf[i]==g_pbuf[i])) nan=1; sum+=g_pbuf[i]; }
assert(!nan && "g_pbuf must be finite after a NaN logit");
assert(approx1(sum) && "g_pbuf must normalize to 1");
assert(approx1(g_pbuf[3]) && "mass must land on the max finite logit (idx3)");
assert(dist_sample(8,-1)==3 && "sampler emits the finite argmax, not token 0"); }
/* --- NaN at index 0: poisons the max scan itself (the old mx=lo[0] seed),
* the failure mode that starts before the sum (review note on #369) --- */
{ float lo[8]={NAN,1.f,0.5f,6.f,0.2f,-1.f,0.f,0.3f}; /* max finite = idx3 (6.0) */
dist_build(lo,8);
double sum=0; int nan=0;
for(int i=0;i<8;i++){ if(!(g_pbuf[i]==g_pbuf[i])) nan=1; sum+=g_pbuf[i]; }
assert(!nan && "NaN at lo[0] must not poison via the mx seed");
assert(approx1(sum) && "still normalizes to 1");
assert(dist_sample(8,-1)==3 && "emits the max finite logit, not token 0"); }
/* --- all-NaN logits: worst case — must stay finite, no crash --- */
{ float lo[8]; for(int i=0;i<8;i++) lo[i]=NAN;
dist_build(lo,8);
for(int i=0;i<8;i++) assert(g_pbuf[i]==g_pbuf[i] && "all-NaN: g_pbuf stays finite"); }
/* --- regression: clean logits still produce a valid distribution --- */
{ float lo[8]={0.1f,0.2f,3.0f,0.4f,0.5f,0.6f,0.7f,0.8f}; /* peak = idx2 */
dist_build(lo,8);
double sum=0; int nan=0;
for(int i=0;i<8;i++){ if(!(g_pbuf[i]==g_pbuf[i])) nan=1; sum+=g_pbuf[i]; }
assert(!nan && "clean softmax stays finite");
assert(approx1(sum) && "clean softmax must sum to 1");
assert(g_pbuf[2]>=g_pbuf[0] && "peak token keeps the most mass"); }
printf("OK test_logit_nan: argmax_v NaN-skip + dist_build finite-collapse\n");
return 0;
}
c/tests/test_makefile_platform.py
+2 -2
View File
@@ -35,7 +35,7 @@ class MakefilePlatformTests(unittest.TestCase):
env["PATH"] = ""
result = subprocess.run(
[MAKE, "--no-print-directory", "-B", "-n", "glm"],
[MAKE, "--no-print-directory", "-B", "-n", "colibri"],
cwd=C_DIR,
env=env,
text=True,
@@ -43,7 +43,7 @@ class MakefilePlatformTests(unittest.TestCase):
check=True,
)
self.assertIn("-o glm.exe", result.stdout)
self.assertIn("-o colibri.exe", result.stdout)
self.assertIn("-fopenmp", result.stdout)
self.assertIn("-static", result.stdout)
c/tests/test_openai_server.py
+33
View File
@@ -380,6 +380,25 @@ class DispatcherTest(unittest.TestCase):
engine.close()
self.assertEqual(chunks, ["é"])
def test_records_profile_snapshots_from_prof_lines(self):
def respond(process, frame):
request_id = frame.split()[1]
process.stdout.feed(b"DATA " + request_id + b" 2\nok\n")
process.stdout.feed(b"PROF 2.500 7 12 0.400 0.100 0.900 0.600 0.200 15\n")
process.stdout.feed(b"DONE " + request_id + b" STAT 12 4.8 0 1.0 7 0\n")
process = FakeProcess(respond)
with patch("openai_server.subprocess.Popen", return_value=process):
engine = Engine("glm", "model")
engine.generate("hello", 16, 0.7, 0.9, lambda _: None)
engine.close()
self.assertEqual(engine.profile_seq, 1)
self.assertEqual(list(engine.profile), [{
"wall_s": 2.5, "prompt_tokens": 7, "completion_tokens": 12,
"expert_disk_s": 0.4, "expert_wait_s": 0.1, "expert_matmul_s": 0.9,
"attention_s": 0.6, "lm_head_s": 0.2, "forwards": 15,
}])
def test_cancels_generation_after_consumer_disconnects(self):
request_id = None
@@ -443,6 +462,20 @@ class HTTPTest(unittest.TestCase):
self.assertIn("queued", scheduler)
self.assertEqual(health["kv_slots"], 2)
def test_profile_reports_recent_turns_without_auth(self):
with urlopen(self.base + "/profile", timeout=2) as response:
self.assertEqual(json.load(response), {"seq": 0, "turns": []})
turn = {"wall_s": 2.5, "prompt_tokens": 7, "completion_tokens": 12,
"expert_disk_s": 0.4, "expert_wait_s": 0.1, "expert_matmul_s": 0.9,
"attention_s": 0.6, "lm_head_s": 0.2, "forwards": 15}
self.engine.profile = [turn]
self.engine.profile_seq = 1
try:
with urlopen(self.base + "/profile", timeout=2) as response:
self.assertEqual(json.load(response), {"seq": 1, "turns": [turn]})
finally:
del self.engine.profile, self.engine.profile_seq
def test_browser_preflight(self):
request = Request(self.base + "/v1/chat/completions", method="OPTIONS", headers={
"Origin": "http://localhost:5173",
c/tests/test_openai_tools_e2e.py
+182
View File
@@ -0,0 +1,182 @@
"""End-to-end tool-calling test for the OpenAI gateway (#401).
Unlike the unit tests in test_openai_server.py (which call parse_tool_calls /
render_chat directly), this suite runs openai_server.py as a real subprocess
against a mock engine that speaks the actual SERVE wire protocol
(READY / SUBMIT / DATA / DONE), then talks to it over real HTTP. It pins down
the full path a coding client exercises: tool declaration rendering, marker
suppression in streamed deltas (across chunk boundaries), tool_calls in both
response shapes, and the <|observation|><tool_response> round trip.
"""
import json
import os
import socket
import subprocess
import sys
import tempfile
import unittest
import urllib.request
from pathlib import Path
SERVER = Path(__file__).resolve().parent.parent / "openai_server.py"
MODEL_ID = "glm-5.2-colibri"
# Mock engine: replies are keyed on the prompt so one process covers every case.
# Prompts received are appended to MOCK_LOG for assertions on the rendering.
MOCK_ENGINE = r'''#!/usr/bin/env python3
import sys, os
out, inp = sys.stdout.buffer, sys.stdin.buffer
out.write(b"\x01\x01READY\x01\x01\n" + b"STAT 0 0 0 0 0\n"); out.flush()
def reply(rid, text, chunks=1):
data = text.encode("utf-8")
n = max(1, len(data) // chunks)
for i in range(0, len(data), n):
part = data[i:i+n]
out.write(("DATA %s %d\n" % (rid, len(part))).encode() + part + b"\n"); out.flush()
out.write(("DONE %s STAT %d 1.0 50.0 10.0 42 0\n" % (rid, len(text.split()))).encode())
out.flush()
while True:
line = inp.readline()
if not line: break
f = line.decode().strip().split()
if not f or f[0] != "SUBMIT": continue
rid, plen = f[1], int(f[3])
prompt = inp.read(plen).decode("utf-8", "replace"); inp.read(1)
with open(os.environ["MOCK_LOG"], "a") as log:
log.write(prompt + "\n\x00\n")
if "<tool_response>" in prompt:
reply(rid, "25 degrees and sunny in Rome.")
elif "weather in Rome" in prompt:
reply(rid, "<tool_call>get_weather<arg_key>city</arg_key>"
"<arg_value>Rome</arg_value></tool_call>")
elif "weather in Milan" in prompt:
# split across many tiny DATA chunks: streamed marker suppression must
# hold even when a marker straddles a chunk boundary
reply(rid, "Checking. <tool_call>get_weather<arg_key>city</arg_key>"
"<arg_value>Milan</arg_value></tool_call>", chunks=20)
else:
reply(rid, "Hello from the mock engine.")
'''
TOOLS = [{"type": "function", "function": {
"name": "get_weather",
"description": "Current weather for a city",
"parameters": {"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"]}}}]
@unittest.skipUnless(os.name == "posix",
"the mock engine is a shebang script the gateway execs directly; "
"Windows CreateProcess cannot run it. The gateway logic under test "
"is platform-independent and covered by the POSIX CI jobs.")
class ToolCallingE2E(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.tmp = tempfile.TemporaryDirectory()
mock = Path(cls.tmp.name) / "mock_engine.py"
mock.write_text(MOCK_ENGINE)
mock.chmod(0o755)
cls.mock_log = Path(cls.tmp.name) / "prompts.log"
cls.mock_log.touch()
with socket.socket() as probe: # free port, then hand it to the server
probe.bind(("127.0.0.1", 0))
cls.port = probe.getsockname()[1]
env = dict(os.environ, MOCK_LOG=str(cls.mock_log))
cls.server = subprocess.Popen(
[sys.executable, str(SERVER), "--model", cls.tmp.name,
"--engine", str(mock), "--port", str(cls.port)],
env=env, stderr=subprocess.DEVNULL)
cls.base = f"http://127.0.0.1:{cls.port}/v1"
for _ in range(100):
try:
urllib.request.urlopen(cls.base + "/models", timeout=2)
return
except OSError:
if cls.server.poll() is not None:
raise RuntimeError("gateway exited during startup")
import time
time.sleep(0.1)
raise RuntimeError("gateway did not come up")
@classmethod
def tearDownClass(cls):
cls.server.terminate()
cls.server.wait(timeout=5)
cls.tmp.cleanup()
def post(self, body, stream=False):
req = urllib.request.Request(
self.base + "/chat/completions", json.dumps(body).encode(),
{"Content-Type": "application/json"})
resp = urllib.request.urlopen(req, timeout=30)
if not stream:
return json.loads(resp.read())
events = []
for raw in resp:
line = raw.decode().strip()
if line.startswith("data: ") and line != "data: [DONE]":
events.append(json.loads(line[6:]))
return events
def test_tool_call_non_stream(self):
r = self.post({"model": MODEL_ID, "tools": TOOLS,
"messages": [{"role": "user",
"content": "What is the weather in Rome?"}]})
choice = r["choices"][0]
self.assertEqual(choice["finish_reason"], "tool_calls")
calls = choice["message"]["tool_calls"]
self.assertEqual(len(calls), 1)
self.assertEqual(calls[0]["function"]["name"], "get_weather")
self.assertEqual(json.loads(calls[0]["function"]["arguments"]), {"city": "Rome"})
self.assertNotIn("<tool_call>", choice["message"].get("content") or "")
def test_tool_call_streamed_markers_suppressed(self):
events = self.post({"model": MODEL_ID, "tools": TOOLS, "stream": True,
"messages": [{"role": "user",
"content": "What is the weather in Milan?"}]},
stream=True)
deltas = [e["choices"][0]["delta"] for e in events if e["choices"]]
text = "".join(d.get("content") or "" for d in deltas)
self.assertNotIn("<tool_call>", text)
self.assertNotIn("<arg_key>", text)
calls = [d["tool_calls"] for d in deltas if d.get("tool_calls")]
self.assertEqual(len(calls), 1)
self.assertEqual(calls[0][0]["function"]["name"], "get_weather")
self.assertEqual(json.loads(calls[0][0]["function"]["arguments"]),
{"city": "Milan"})
finish = [e["choices"][0]["finish_reason"] for e in events
if e["choices"] and e["choices"][0].get("finish_reason")]
self.assertEqual(finish, ["tool_calls"])
def test_tool_result_round_trip(self):
r = self.post({"model": MODEL_ID, "tools": TOOLS, "messages": [
{"role": "user", "content": "What is the weather in Rome?"},
{"role": "assistant", "content": None, "tool_calls": [
{"id": "call_x", "type": "function",
"function": {"name": "get_weather",
"arguments": "{\"city\": \"Rome\"}"}}]},
{"role": "tool", "tool_call_id": "call_x",
"content": "25 degrees, sunny"}]})
choice = r["choices"][0]
self.assertEqual(choice["finish_reason"], "stop")
self.assertFalse(choice["message"].get("tool_calls"))
self.assertIn("25 degrees", choice["message"]["content"])
rendered = self.mock_log.read_text().split("\x00")[-2]
self.assertIn("<|observation|><tool_response>25 degrees, sunny</tool_response>",
rendered)
self.assertIn("# Tools", rendered)
self.assertIn('"get_weather"', rendered)
def test_no_tools_plain_text(self):
r = self.post({"model": MODEL_ID,
"messages": [{"role": "user", "content": "Hi!"}]})
choice = r["choices"][0]
self.assertEqual(choice["finish_reason"], "stop")
self.assertIn("mock engine", choice["message"]["content"])
if __name__ == "__main__":
unittest.main()
c/tests/test_ragged_attention.cu
+39
View File
@@ -0,0 +1,39 @@
#include "../backend_cuda.h"
#include <cmath>
#include <cstdio>
#include <vector>
int main(){
int dev=0;if(!coli_cuda_init(&dev,1))return 77;
constexpr int S=3,H=2,Q=2,R=1,V=2,K=3,D=H*V,O=3,T=3;
std::vector<float> w(H*(Q+V)*K),p(O*D),q(S*H*(Q+R));
for(size_t i=0;i<w.size();i++)w[i]=((int)(i%11)-5)*.07f;
for(size_t i=0;i<p.size();i++)p[i]=((int)(i%7)-3)*.09f;
for(size_t i=0;i<q.size();i++)q[i]=((int)(i%13)-6)*.05f;
ColiCudaTensor *tw=nullptr,*tp=nullptr;
if(!coli_cuda_tensor_upload(&tw,w.data(),nullptr,0,K,H*(Q+V),dev)||
!coli_cuda_tensor_upload(&tp,p.data(),nullptr,0,D,O,dev))return 1;
int n[S]={1,2,3};std::vector<std::vector<float>> l(S),r(S);
const float *lp[S],*rp[S];
const void *keys[S];
for(int s=0;s<S;s++){
l[s].resize(n[s]*K);r[s].resize(n[s]*R);
for(size_t i=0;i<l[s].size();i++)l[s][i]=((int)((i+s*3)%9)-4)*.08f;
for(size_t i=0;i<r[s].size();i++)r[s][i]=((int)((i+s)%5)-2)*.06f;
lp[s]=l[s].data();rp[s]=r[s].data();keys[s]=&l[s];
}
float got[S*O],ref[S*O],warm[S*O];
int first[S]={1,1,1};
if(!coli_cuda_attention_project_ragged(tw,tp,warm,q.data(),keys,lp,rp,first,
S,H,Q,R,V,K,1,.2f))return 2;
if(!coli_cuda_attention_project_ragged(tw,tp,got,q.data(),keys,lp,rp,n,S,H,Q,R,V,K,T,.2f))return 2;
for(int s=0;s<S;s++)if(!coli_cuda_attention_project_batch(tw,tp,ref+s*O,
q.data()+s*H*(Q+R),lp[s],rp[s],1,H,Q,R,V,K,n[s],.2f))return 3;
double e=0,z=0;for(int i=0;i<S*O;i++){
double d=got[i]-ref[i];e+=d*d;z+=(double)ref[i]*ref[i];
}
double rms=std::sqrt(e/(z+1e-30));std::printf("ragged_relative_rms=%.9g\n",rms);
coli_cuda_tensor_free(tw);coli_cuda_tensor_free(tp);coli_cuda_shutdown();
return rms<1e-6?0:4;
}
c/tests/test_resource_plan.py
+92 -2
View File
@@ -10,6 +10,7 @@ from resource_plan import (
GB,
analyze_model,
build_plan,
cpu_socket_count,
environment_for_plan,
format_plan,
memory_available,
@@ -66,15 +67,19 @@ class ResourcePlanTest(unittest.TestCase):
# 0 slots/layer. The value must be a sane positive number of bytes.
self.assertGreater(memory_available(), 0)
def test_cpu_socket_count_is_positive(self):
self.assertGreaterEqual(cpu_socket_count(), 1)
def test_builds_bounded_three_tier_plan(self):
gpus = [{"index": 0, "name": "test-gpu", "total_bytes": 12 * GB,
"free_bytes": 10 * GB}]
plan = build_plan(self.model, ram_gb=16, context=32, vram_gb=20,
available_memory=32 * GB, available_disk=100 * GB, gpus=gpus,
physical_cpus=24)
physical_cpus=24, cpu_sockets=2)
self.assertEqual(plan["version"], 2)
self.assertEqual(plan["policy"]["name"], "quality")
self.assertEqual(plan["cpu"]["physical_cores"], 24)
self.assertEqual(plan["cpu"]["sockets"], 2)
self.assertTrue(plan["policy"]["preserve_quantization"])
self.assertFalse(plan["tiers"]["vram"]["requires_host_backing"])
self.assertEqual(plan["tiers"]["ram"]["budget_bytes"], 16 * GB)
@@ -105,7 +110,7 @@ class ResourcePlanTest(unittest.TestCase):
{"index": 1, "name": "b", "total_bytes": 12 * GB, "free_bytes": 10 * GB},
]
plan = build_plan(self.model, ram_gb=16, available_memory=32 * GB,
available_disk=1, gpus=gpus)
available_disk=1, gpus=gpus, cpu_sockets=2)
env = environment_for_plan(plan, {"RAM_GB": "12", "PIN": "stats.txt",
"COLI_GPUS": "1"})
self.assertEqual(env["RAM_GB"], "12")
@@ -125,6 +130,91 @@ class ResourcePlanTest(unittest.TestCase):
"OMP_PROC_BIND": "close"})
self.assertEqual(explicit_threads["OMP_NUM_THREADS"], "7")
self.assertEqual(explicit_threads["OMP_PROC_BIND"], "close")
if sys.platform.startswith("linux"):
self.assertEqual(env["COLI_NUMA"], "1")
explicit_numa = environment_for_plan(plan, {"COLI_NUMA": "0"})
self.assertEqual(explicit_numa["COLI_NUMA"], "0")
def test_single_socket_plan_does_not_enable_numa(self):
plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
gpus=[], physical_cpus=8, cpu_sockets=1)
self.assertNotIn("COLI_NUMA", environment_for_plan(plan))
def test_auto_tune_mtp_off_when_compute_bound(self):
# Tiny model with 64 GB RAM and no GPU: all experts fit in RAM with no
# warm tier, so the plan classifies as compute-bound.
plan = build_plan(self.model, ram_gb=64, available_memory=64 * GB,
available_disk=100 * GB, gpus=[], physical_cpus=24,
cpu_sockets=2)
# With such a small model fully in RAM and no GPU, bottleneck is compute
self.assertEqual(plan["bottleneck_class"], "compute")
self.assertIn("DRAFT", plan["tune"])
self.assertEqual(plan["tune"]["DRAFT"]["value"], "0")
env = environment_for_plan(plan)
self.assertEqual(env["DRAFT"], "0")
explicit = environment_for_plan(plan, {"DRAFT": "3"})
self.assertEqual(explicit["DRAFT"], "3")
def test_auto_tune_mtp_off_when_disk_low_hit(self):
# Use a model large enough that 8 GB RAM can't hold all experts.
big = tempfile.TemporaryDirectory()
bigmodel = Path(big.name)
(bigmodel / "config.json").write_text(json.dumps({
"num_hidden_layers": 2, "n_routed_experts": 4,
"kv_lora_rank": 4, "qk_rope_head_dim": 2,
"qk_nope_head_dim": 3, "v_head_dim": 5, "num_attention_heads": 2,
}))
expert_size = 3 * GB # each expert 3 GB → 12 GB total, won't fit in 8 GB budget
write_shard(bigmodel / "out-00000.safetensors", [
("model.embed_tokens.weight", 100),
("model.layers.0.self_attn.q_a_proj.weight", 200),
])
for i in range(4):
write_shard(bigmodel / f"out-{i+1:05d}.safetensors", [
(f"model.layers.1.mlp.experts.{i}.gate_proj.weight", expert_size),
])
plan = build_plan(bigmodel, ram_gb=0, available_memory=4 * GB,
available_disk=100 * GB, gpus=[], physical_cpus=8,
cpu_sockets=1)
big.cleanup()
self.assertEqual(plan["bottleneck_class"], "disk")
self.assertLess(plan["projected_hit_rate"], 0.90)
self.assertEqual(plan["tune"]["DRAFT"]["value"], "0")
def test_auto_tune_pipe_multi_gpu(self):
gpus = [
{"index": 0, "name": "a", "total_bytes": 32 * GB, "free_bytes": 30 * GB},
{"index": 1, "name": "b", "total_bytes": 32 * GB, "free_bytes": 30 * GB},
]
plan = build_plan(self.model, ram_gb=16, available_memory=32 * GB,
available_disk=1, gpus=gpus, cpu_sockets=2)
self.assertEqual(plan["tune"]["COLI_CUDA_PIPE"]["value"], "2")
env = environment_for_plan(plan)
self.assertEqual(env["COLI_CUDA_PIPE"], "2")
def test_auto_tune_pipe_single_gpu(self):
gpus = [{"index": 0, "name": "a", "total_bytes": 12 * GB, "free_bytes": 10 * GB}]
plan = build_plan(self.model, ram_gb=16, available_memory=32 * GB,
available_disk=1, gpus=gpus, cpu_sockets=1)
self.assertEqual(plan["tune"]["COLI_CUDA_PIPE"]["value"], "1")
def test_auto_tune_numa_hint_for_cpu_only(self):
plan = build_plan(self.model, ram_gb=64, available_memory=64 * GB,
available_disk=1, gpus=[], physical_cpus=64, cpu_sockets=2)
self.assertIn("_numa_hint", plan["tune"])
self.assertIn("numactl", plan["tune"]["_numa_hint"])
self.assertIn("auto-tune", format_plan(plan))
def test_format_plan_shows_tune_and_hit_rate(self):
plan = build_plan(self.model, ram_gb=64, available_memory=64 * GB,
available_disk=100 * GB, gpus=[], physical_cpus=24,
cpu_sockets=1)
text = format_plan(plan)
self.assertIn("hit", text)
self.assertIn("auto-tune", text)
self.assertIn("DRAFT", text)
def test_cpu_binary_does_not_apply_gpu_tier(self):
plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
gpus=[{"index": 0, "name": "a", "total_bytes": 8 * GB,
c/tests/test_sample_nan.c
+62
View File
@@ -0,0 +1,62 @@
/* Regressione #369: un solo logit NaN o +Inf faceva emettere a dist_build/dist_sample
* il token 0 PER SEMPRE, in silenzio (il fallback `g_pbuf[i]>0` e' falso su NaN ovunque).
* Ora la distribuzione degenere ripiega sull'argmax dei logit FINITI e avvisa una volta.
*
* Il test verifica: (a) con logit sani il campionamento resta corretto; (b) con NaN/+Inf
* iniettato il token scelto e' l'argmax dei FINITI (mai 0 per default), su ogni posizione
* del NaN inclusa lo[0]; (c) nessun NaN sopravvive in g_pbuf. */
#define main coli_glm_main_unused
#include "../colibri.c"
#undef main
#include <stdio.h>
static uint32_t rs=0x1234abcd; static uint32_t xr(){rs^=rs<<13;rs^=rs>>17;rs^=rs<<5;return rs;}
static int pbuf_has_nan(int V){ for(int i=0;i<V;i++) if(!isfinite(g_pbuf[i])) return 1; return 0; }
int main(void){
int fail=0, V=2000;
float *lo=malloc(V*sizeof(float));
g_temp=0.7f; g_nuc=0.90f; /* il path serve di default */
/* (a) logit sani: il campionato deve avere prob > 0 e nessun NaN nel buffer */
for(int i=0;i<V;i++) lo[i]=(float)((int)(xr()%2000)-1000)/100.f;
int known=1337; lo[known]=50.f; /* picco netto */
dist_build(lo,V);
if(pbuf_has_nan(V)){ printf(" FAIL: NaN in g_pbuf con logit sani\n"); fail=1; }
if(g_pbuf[known]<=0.f){ printf(" FAIL: il picco ha prob 0\n"); fail=1; }
if(!fail) printf(" logit sani: distribuzione valida, picco vivo ok\n");
/* (b) NaN/+Inf iniettato in varie posizioni; il finito-argmax deve vincere */
float bad[3]; bad[0]=NAN; bad[1]=INFINITY; bad[2]=-INFINITY;
const char *bn[3]={"NaN","+Inf","-Inf"};
for(int b=0;b<3;b++){
for(int pos=0;pos<3;pos++){ /* lo[0], meta', ultimo */
for(int i=0;i<V;i++) lo[i]=(float)((int)(xr()%400)-200)/100.f;
int amax=777; lo[amax]=9.0f; /* massimo FINITO atteso */
int at = pos==0?0 : pos==1?V/2 : V-1;
if(at==amax) amax=amax+1; /* non sovrapporre */
lo[amax]=9.0f;
lo[at]=bad[b]; /* veleno */
dist_build(lo,V);
if(pbuf_has_nan(V)){ printf(" FAIL: NaN sopravvive (%s @ %d)\n",bn[b],at); fail=1; continue; }
/* con -Inf il fallback puo' non scattare (max finito resta), ma il buffer
* deve restare valido e sommare ~1: con +Inf/NaN scatta il delta su amax */
int picked=-1; float pv=-1;
for(int i=0;i<V;i++) if(g_pbuf[i]>pv){pv=g_pbuf[i];picked=i;}
if(b<2 && picked!=amax){ /* NaN e +Inf: delta esatto su amax */
printf(" FAIL: %s @ %d -> picked %d, atteso argmax finito %d\n",bn[b],at,picked,amax); fail=1;
}
}
}
if(!fail) printf(" NaN/+Inf iniettato: argmax dei finiti vince, mai 0/NaN ok\n");
/* (c) caso estremo: TUTTI non finiti -> non deve crashare, buffer valido */
for(int i=0;i<V;i++) lo[i]=NAN;
dist_build(lo,V);
if(pbuf_has_nan(V)){ printf(" FAIL: tutti-NaN lascia NaN nel buffer\n"); fail=1; }
else printf(" tutti non-finiti: nessun crash, buffer valido ok\n");
printf(fail?"test_sample_nan: FAIL\n":"test_sample_nan: ok\n");
return fail;
}
c/tests/test_st_pread.c
+89
View File
@@ -0,0 +1,89 @@
/* st_pread_full: chunk loop + honest truncation errors.
* Built with -DST_PREAD_CHUNK=7 so a ~100-byte tensor takes many pread calls —
* exercising the loop that production only needs past 2^31 bytes (one pread
* caps there on Linux; big bf16 tensors exceed it). Also forks a child against
* a truncated shard and requires exit(1) with a "short read" message instead
* of the old perror("... : Success"). */
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#ifndef _WIN32
#include <sys/wait.h>
#include <unistd.h>
#endif
#include "../st.h"
#define CHECK(condition) do { \
if (!(condition)) { \
fprintf(stderr, "%s:%d: check failed: %s\n", __FILE__, __LINE__, #condition); \
return 1; \
} \
} while (0)
static void write_snap(const char *dir, int truncate_bytes) {
char path[512];
snprintf(path, sizeof(path), "%s/model.safetensors", dir);
unsigned char data[96];
for (int i = 0; i < 96; i++) data[i] = (unsigned char)(i * 7 + 3);
const char *hdr = "{\"t\":{\"dtype\":\"U8\",\"shape\":[96],\"data_offsets\":[0,96]}}";
uint64_t hlen = strlen(hdr);
FILE *f = fopen(path, "wb");
fwrite(&hlen, 8, 1, f);
fwrite(hdr, 1, hlen, f);
fwrite(data, 1, (size_t)(96 - truncate_bytes), f);
fclose(f);
}
int main(void) {
/* relative to the CWD, per test_stops: MinGW .exe files resolve Windows
* paths and "/tmp" is not one */
char dir[] = "test_st_pread_XXXXXX";
if (!mkdtemp(dir)) { perror("mkdtemp"); return 1; }
/* 1) chunk loop: 96-byte tensor read 7 bytes at a time, content exact */
write_snap(dir, 0);
shards S; st_init(&S, dir);
unsigned char out[96] = {0};
st_read_raw(&S, "t", out, 0);
for (int i = 0; i < 96; i++) CHECK(out[i] == (unsigned char)(i * 7 + 3));
#ifndef _WIN32
/* 2) shard truncated AFTER st_init (init validates static bounds, so the
* pread path only fires when the file shrinks underneath a live handle):
* child must exit(1) with an honest message, not perror's "Success" */
char shard[512]; snprintf(shard, sizeof(shard), "%s/model.safetensors", dir);
struct stat sb; CHECK(stat(shard, &sb) == 0);
CHECK(truncate(shard, sb.st_size - 40) == 0);
int pipefd[2]; CHECK(pipe(pipefd) == 0);
pid_t pid = fork(); CHECK(pid >= 0);
if (pid == 0) {
dup2(pipefd[1], 2); close(pipefd[0]); close(pipefd[1]);
unsigned char buf[96];
st_read_raw(&S, "t", buf, 0); /* inherited handles; must exit(1) inside */
_exit(42); /* reaching here = bug */
}
close(pipefd[1]);
char err[512] = {0};
ssize_t n = read(pipefd[0], err, sizeof(err)-1); (void)n;
close(pipefd[0]);
int status = 0; waitpid(pid, &status, 0);
CHECK(WIFEXITED(status) && WEXITSTATUS(status) == 1);
CHECK(strstr(err, "short read") != NULL);
CHECK(strstr(err, "Success") == NULL);
#else
/* fork/pipe/truncate are POSIX; Windows still runs the chunk-loop check */
printf("test_st_pread: truncation subtest skipped on Windows\n");
#endif
char cmd[600];
#ifdef _WIN32
snprintf(cmd, sizeof(cmd), "rmdir /s /q %s", dir);
#else
snprintf(cmd, sizeof(cmd), "rm -rf %s", dir);
#endif
if (system(cmd)) {}
printf("test_st_pread: chunk loop + honest truncation error: ok\n");
return 0;
}
c/tests/test_stops.c
+136
View File
@@ -0,0 +1,136 @@
/* Stop-token arming: the engine must not depend on a converted checkpoint's
* config.json being complete.
*
* Why this test exists (#298, #307): GLM-5.2 declares THREE eos ids
* (<|endoftext|>, <|user|>, <|observation|>). A conversion tool that rewrites
* config.json with a reduced eos_token_id list leaves the engine stopping on
* fewer tokens than the model actually emits -- and the ones it misses get
* DETOKENIZED AND PRINTED INTO THE CHAT as literal text, with generation
* continuing past the real end of the turn. @woolcoxm hit exactly this on a
* g64 checkpoint and confirmed it by comparing token ids.
*
* Two independent defenses, one test each:
* 1. eos_token_id is unioned from generation_config.json (HuggingFace's
* authority for generation) as well as config.json.
* 2. every added-token the TOKENIZER marks "special":true is a stop, whatever
* the configs say. Those are control tokens (<|user|>, <sop>, ...) and are
* never legitimate content. <think>/<tool_call> are "special":false and
* must NOT be swept up -- they are real output.
*
* Defense 2 is what makes this robust against checkpoints we don't control:
* even with BOTH configs mutilated, a control token cannot leak into a reply. */
#define main coli_glm_main_unused
#include "../colibri.c"
#undef main
static const char *TOKJSON =
"{\"model\":{\"vocab\":{\"a\":0,\"b\":1,\"c\":2},\"merges\":[[\"a\",\"b\"]]},"
" \"added_tokens\":["
" {\"id\":100,\"content\":\"<|endoftext|>\",\"special\":true},"
" {\"id\":101,\"content\":\"<|user|>\",\"special\":true},"
" {\"id\":102,\"content\":\"<|observation|>\",\"special\":true},"
" {\"id\":103,\"content\":\"<|assistant|>\",\"special\":true},"
" {\"id\":104,\"content\":\"<sop>\",\"special\":true},"
" {\"id\":110,\"content\":\"<think>\",\"special\":false},"
" {\"id\":111,\"content\":\"<tool_call>\",\"special\":false}"
"]}";
/* minimal config.json that survives load_cfg's range validation */
static void write_cfg(const char *dir, const char *fname, const char *eos_json){
char p[512]; snprintf(p,sizeof(p),"%s/%s",dir,fname);
FILE *f=fopen(p,"w"); if(!f){ perror(p); exit(1); }
fprintf(f,"{\"hidden_size\":64,\"num_hidden_layers\":2,\"num_attention_heads\":4,"
"\"n_routed_experts\":8,\"num_experts_per_tok\":2,\"moe_intermediate_size\":32,"
"\"intermediate_size\":64,\"first_k_dense_replace\":1,\"q_lora_rank\":0,"
"\"kv_lora_rank\":16,\"qk_nope_head_dim\":8,\"qk_rope_head_dim\":8,"
"\"v_head_dim\":8,\"n_shared_experts\":1,\"vocab_size\":200,"
"\"n_group\":1,\"topk_group\":1,\"rope_theta\":10000.0");
if(eos_json) fprintf(f,",\"eos_token_id\":%s",eos_json);
fprintf(f,"}\n"); fclose(f);
}
static void write_tok(const char *dir){
char p[512]; snprintf(p,sizeof(p),"%s/tokenizer.json",dir);
FILE *f=fopen(p,"w"); if(!f){ perror(p); exit(1); }
fputs(TOKJSON,f); fclose(f);
}
static void rm_file(const char *dir, const char *fname){
char p[512]; snprintf(p,sizeof(p),"%s/%s",dir,fname); remove(p);
}
static int expect(const char *what, int id, int want){
int got=is_stop(id);
if(got!=want){ fprintf(stderr," FAIL %s: token %d stop=%d, expected %d\n",what,id,got,want); return 1; }
return 0;
}
int main(void){
int fail=0;
/* Relative to the CWD, like test_compat_direct's TMPF — NOT "/tmp/...".
* These binaries are built by MinGW into native Windows .exe files, which
* resolve Windows paths: "/tmp" is not one, so mkdtemp there fails ENOENT
* and the whole `make check` goes red on the windows job (and only there). */
char dir[]="test_stops_XXXXXX";
if(!mkdtemp(dir)){ perror("mkdtemp"); return 1; }
write_tok(dir);
char tkp[512]; snprintf(tkp,sizeof(tkp),"%s/tokenizer.json",dir);
Tok T; tok_load(&T,tkp);
printf("test_stops: stop arming vs incomplete checkpoint metadata\n");
/* 1. the tokenizer's "special" flag must survive loading, and must NOT
* swallow <think>/<tool_call>, which are real content. */
if(!T.id_special[101]){ fprintf(stderr," FAIL <|user|> not marked special\n"); fail=1; }
if(!T.id_special[104]){ fprintf(stderr," FAIL <sop> not marked special\n"); fail=1; }
if(T.id_special[110]){ fprintf(stderr," FAIL <think> wrongly marked special\n"); fail=1; }
if(T.id_special[111]){ fprintf(stderr," FAIL <tool_call> wrongly marked special\n"); fail=1; }
if(!fail) printf(" tokenizer: special flag parsed, <think>/<tool_call> excluded ok\n");
/* 2. config.json mutilated to one eos, generation_config.json intact:
* the union must recover all three. This is @woolcoxm's case. */
{ Cfg c; memset(&c,0,sizeof c);
write_cfg(dir,"config.json","[100]");
write_cfg(dir,"generation_config.json","[100,101,102]");
load_cfg(&c,dir);
if(c.n_stop!=3){ fprintf(stderr," FAIL union: expected 3 stops, got %d\n",c.n_stop); fail=1; }
else printf(" config eos=[100] + generation_config=[100,101,102] -> 3 stops ok\n"); }
/* 3. no generation_config.json at all: must not crash, config alone wins */
{ Cfg c; memset(&c,0,sizeof c);
write_cfg(dir,"config.json","[100,101,102]");
rm_file(dir,"generation_config.json");
load_cfg(&c,dir);
if(c.n_stop!=3){ fprintf(stderr," FAIL no-genconfig: expected 3, got %d\n",c.n_stop); fail=1; }
else printf(" generation_config.json absent -> config alone, no crash ok\n"); }
/* 4. BOTH configs mutilated: the tokenizer's special set must still stop the
* turn on every control token, and must still leave <think> alone. */
{ Cfg c; memset(&c,0,sizeof c);
write_cfg(dir,"config.json","[100]");
write_cfg(dir,"generation_config.json","[100]");
load_cfg(&c,dir);
stops_arm_tok(&c,-1,&T);
fail|=expect("endoftext (config)", 100,1);
fail|=expect("<|user|> (tokenizer)", 101,1);
fail|=expect("<|observation|>", 102,1);
fail|=expect("<|assistant|>", 103,1);
fail|=expect("<sop>", 104,1);
fail|=expect("<think> must NOT stop",110,0);
fail|=expect("<tool_call> must NOT stop",111,0);
fail|=expect("plain token must NOT stop",2,0);
if(!fail) printf(" both configs mutilated -> tokenizer still stops all 5 control tokens ok\n"); }
/* 5. T=NULL (validation/oracle path): config stops only, unchanged behaviour */
{ Cfg c; memset(&c,0,sizeof c);
write_cfg(dir,"config.json","[100]");
rm_file(dir,"generation_config.json");
load_cfg(&c,dir);
stops_arm_tok(&c,-1,NULL);
fail|=expect("T=NULL: config stop",100,1);
fail|=expect("T=NULL: no tokenizer sweep",101,0);
if(!fail) printf(" T=NULL -> config stops only (validation path untouched) ok\n"); }
rm_file(dir,"config.json"); rm_file(dir,"generation_config.json"); rm_file(dir,"tokenizer.json");
rmdir(dir);
if(fail){ printf("test_stops: FAIL\n"); return 1; }
printf("test_stops: ok\n");
return 0;
}
c/tests/test_topp.c
+280
View File
@@ -0,0 +1,280 @@
/* Top-p (nucleus) truncation in dist_build: the partial-select rewrite (#335) must be
* indistinguishable from the old full-vocab qsort for every shape dist_sample can see.
*
* Why this test exists (#335): dist_build() previously qsort-ed the entire 151936-entry
* vocab on every sampled token to find the few-hundred-token head whose cumulative mass
* reaches g_nuc. It now heapifies (O(V)) and pops only the head (k * O(log V)). The win
* is structural; the risk is a silent sampling-distribution change, because the contract
* is subtle:
*
* dist_sample() iterates g_pbuf[0..V-1] BY TOKEN ID and sums probabilities directly.
* So dist_build MUST leave g_pbuf indexed by id (never reordered) AND must zero every
* truncated tail entry -- merely excluding the tail from the head would leave mass on
* it and the sampled distribution would drift with no crash and no error.
*
* Strategy: drive the REAL dist_build (via the test_stops.c include-glm.c pattern) on a
* sweep of distributions and g_nuc values, and compare against an INDEPENDENT reference
* that re-implements the OLD algorithm (full qsort + zero-tail + renorm) in double on a
* private buffer. On shapes with no ties the renormalized head must be BIT-IDENTICAL to
* the reference (the issue's stated invariant: s2 accumulates in the same descending
* order). On tie shapes, where the unstable qsort already left ordering unspecified, we
* check multiset equality instead. Every shape also checks: exact-zero tails, head sums
* to 1.0, and a sane keep-count.
*
* No scratch files: the test runs entirely in memory (no mkdtemp), so it builds clean on
* the Windows MinGW CI job without the unmerged compat shim (#352). */
#define main coli_glm_main_unused
#include "../colibri.c"
#undef main
#include <math.h>
static int g_nfails = 0;
/* pointer set by ref_build so cmp_ref_desc can read the current reference buffer
* (the qsort comparator gets no user-data argument in C). */
static const double *g_ref_p = NULL;
#define FAIL(fmt, ...) do { \
fprintf(stderr, " FAIL [%s V=%d nuc=%.3f shape=%s]: " fmt "\n", \
label, V, nuc, shape_name, ##__VA_ARGS__); \
g_nfails++; \
return; \
} while (0)
/* ---- independent reference: the OLD algorithm, in double, on a private buffer ------- */
/* Stable qsort by descending probability (ties broken by ascending index, which makes
* the reference deterministic regardless of the production comparator). */
static int cmp_ref_desc(const void *a, const void *b){
double pa = ((const double *)g_ref_p)[*(const int*)a];
double pb = ((const double *)g_ref_p)[*(const int*)b];
if (pa < pb) return 1;
if (pa > pb) return -1;
/* tie -> lower index first (stable, unlike the production comparator) */
return *(const int*)a - *(const int*)b;
}
/* Build the reference distribution into out[0..V-1] (indexed by token id), mirroring the
* old dist_build: softmax(lo/temp) truncated to top-p nuc, tail zeroed, head renormalized.
* Returns the keep-count through *keep_out. */
static void ref_build(const float *lo, int V, double temp, double nuc,
double *out, int *pidx, int *keep_out){
double mx = lo[0]; for (int i = 1; i < V; i++) if (lo[i] > mx) mx = lo[i];
double s = 0, invt = 1.0 / (temp > 1e-4 ? temp : 1e-4);
for (int i = 0; i < V; i++){ out[i] = exp((lo[i]-mx)*invt); s += out[i]; }
for (int i = 0; i < V; i++) out[i] /= s;
if (nuc > 0 && nuc < 1.0){
for (int i = 0; i < V; i++) pidx[i] = i;
qsort(pidx, V, sizeof(int), cmp_ref_desc);
double cum = 0; int keep = V;
for (int i = 0; i < V; i++){ cum += out[pidx[i]]; if (cum >= nuc){ keep = i+1; break; } }
double s2 = 0;
for (int i = keep; i < V; i++) out[pidx[i]] = 0;
for (int i = 0; i < keep; i++) s2 += out[pidx[i]];
for (int i = 0; i < keep; i++) out[pidx[i]] /= s2;
*keep_out = keep;
} else {
*keep_out = V;
}
}
/* count how many production g_pbuf entries are non-zero == the head size */
static int head_count(int V){
int n = 0; for (int i = 0; i < V; i++) if (g_pbuf[i] != 0.f) n++; return n;
}
/* Run one case: load logits into g_pbuf via the real dist_build, compare to reference.
* shape_name is for diagnostics only. */
static void check_case(const char *label, int V, double nuc, const char *shape_name,
const float *lo){
/* reference on a private buffer */
double *ref = malloc((size_t)V * sizeof(double));
int *ridx = malloc((size_t)V * sizeof(int));
int ref_keep = 0;
g_ref_p = ref; /* cmp_ref_desc reads this */
ref_build(lo, V, g_temp, nuc, ref, ridx, &ref_keep);
/* production: drive the real dist_build (writes the global g_pbuf) */
g_nuc = (float)nuc;
dist_build(lo, V);
int got_keep = head_count(V);
/* 1. keep-count must match the reference exactly. The partial select and the old
* qsort keep the same NUMBER of tokens by construction (same cumulative-mass rule);
* a count divergence is a real bug, not a tie artifact. */
if (got_keep != ref_keep)
FAIL("keep-count mismatch: got %d, ref %d", got_keep, ref_keep);
/* 2. Detect ties across the WHOLE pre-truncation distribution, not just the kept set.
* A tie at the head/tail boundary makes which-side-a-token-lands-on interchangeable:
* both algorithms keep the right count but may keep different MEMBERS. So any input
* with a duplicated softmax value needs the relaxed multiset comparison below. We
* detect this on the reference softmax (pre-truncation) by sorting all V values. */
int has_ties = 0;
{
double *all = malloc((size_t)V * sizeof(double));
/* reconstruct the pre-truncation softmax the same way ref_build does */
double mx = lo[0]; for (int i = 1; i < V; i++) if (lo[i] > mx) mx = lo[i];
double s = 0, invt = 1.0 / (g_temp > 1e-4 ? g_temp : 1e-4);
for (int i = 0; i < V; i++){ all[i] = exp((lo[i]-mx)*invt); s += all[i]; }
for (int i = 0; i < V; i++) all[i] /= s;
for (int a = 1; a < V; a++){ double k = all[a]; int b = a-1;
while (b >= 0 && all[b] > k){ all[b+1] = all[b]; b--; } all[b+1] = k; }
for (int a = 1; a < V; a++) if (all[a] == all[a-1]){ has_ties = 1; break; }
free(all);
}
if (has_ties){
/* Multiset equality of the non-zero (head) values. Ties make membership
* interchangeable, so we compare sorted value-multisets, not id-aligned values.
* Tolerance is 1e-6 relative -- the engine uses float arithmetic, the reference
* double, so sub-ULP noise is expected (matches test_i4_grouped.c's convention). */
double *got = malloc((size_t)ref_keep * sizeof(double));
int gm = 0;
for (int i = 0; i < V; i++) if (g_pbuf[i] != 0.f) got[gm++] = (double)g_pbuf[i];
if (gm != ref_keep)
FAIL("tie-shape head size mismatch: got %d non-zero, ref %d", gm, ref_keep);
for (int a = 1; a < gm; a++){ double k = got[a]; int b = a-1;
while (b >= 0 && got[b] > k){ got[b+1] = got[b]; b--; } got[b+1] = k; }
double *rsort = malloc((size_t)ref_keep * sizeof(double));
int rm = 0;
for (int i = 0; i < V; i++) if (ref[i] != 0.0) rsort[rm++] = ref[i];
for (int a = 1; a < rm; a++){ double k = rsort[a]; int b = a-1;
while (b >= 0 && rsort[b] > k){ rsort[b+1] = rsort[b]; b--; } rsort[b+1] = k; }
int mm = 0; double worst = 0;
for (int i = 0; i < gm; i++){
double d = fabs(got[i] - rsort[i]);
double rel = rsort[i] > 1e-30 ? d / rsort[i] : d;
if (rel > worst) worst = rel;
if (rel > 1e-6) mm++;
}
free(got); free(rsort);
if (mm) FAIL("tie-shape multiset mismatch: %d/%d head values differ beyond 1e-6 rel (worst %.3g)",
mm, ref_keep, worst);
} else {
/* No ties anywhere: membership is forced, so compare id-aligned head values. The
* engine computes in float (g_pbuf /= (float)s2) while the reference uses double,
* so the comparison is relative-tolerance (1e-6), not bit-exact -- the partial
* select and qsort accumulate s2 in the same descending order, so any difference
* is pure float-rounding noise, not an ordering bug. */
int bad = 0; int first_id = -1; float gv = 0, rv = 0; double worst = 0;
for (int i = 0; i < V; i++){
if (ref[i] == 0.0) continue; /* tail */
float want = (float)ref[i];
double d = fabs((double)g_pbuf[i] - (double)want);
double rel = fabs((double)want) > 1e-30 ? d / fabs((double)want) : d;
if (rel > worst) worst = rel;
if (rel > 1e-6){
bad++; if (first_id < 0){ first_id = i; gv = g_pbuf[i]; rv = want; }
if (bad > 3) break;
}
}
if (bad)
FAIL("head not within 1e-6 rel of reference: %d entries differ (first id %d: got %.9g want %.9g, worst %.3g)",
bad, first_id, (double)gv, (double)rv, worst);
}
/* 3. head must renormalize to 1.0 (within float epsilon) */
double sum = 0; for (int i = 0; i < V; i++) sum += g_pbuf[i];
if (fabs(sum - 1.0) > 1e-5)
FAIL("head does not sum to 1.0: sum=%.12g (keep=%d)", sum, got_keep);
free(ref); free(ridx);
printf(" ok [V=%d nuc=%.3f shape=%s keep=%d%s sum=%.10f]\n",
V, nuc, shape_name, got_keep, has_ties ? " (ties)" : "", sum);
}
#undef FAIL
/* deterministic xorshift32 RNG (matches the test_i4_grouped.c convention) */
static uint32_t rng_state = 0x12345678u;
static uint32_t xr(void){ rng_state ^= rng_state << 13; rng_state ^= rng_state >> 17;
rng_state ^= rng_state << 5; return rng_state; }
static double frand(void){ return (xr() >> 8) * (1.0 / 16777216.0); } /* [0,1) */
/* fill logits for a given shape. Shapes chosen to stress the comparator and the head/tail
* boundary differently. */
static void fill_shape(float *lo, int V, int shape){
switch (shape){
case 0: /* uniform -> every token equal probability -> massive tie plateau */
for (int i = 0; i < V; i++) lo[i] = 0.f; break;
case 1: /* peaked: one dominant token, rest small and distinct (no ties).
* The fixed hot-spots are clamped to V-1 so small V (incl. V=1) doesn't
* write out of bounds and corrupt heap metadata on the later free(lo). */
for (int i = 0; i < V; i++) lo[i] = (float)(-1.0 - frand()*4.0);
lo[0] = 3.f; lo[V/3<V?V/3:V-1] = 1.f; lo[V/2<V?V/2:V-1] = 0.5f; break;
case 2: /* all-equal distinct decay (no ties): geometric, strictly decreasing */
for (int i = 0; i < V; i++) lo[i] = (float)(-0.001 * (double)i); break;
case 3: /* plateau ties: blocks of equal value -> comparator tie handling */
for (int i = 0; i < V; i++) lo[i] = (float)(-(double)(i / 7)); /* 7-wide plateaus */
break;
case 4: /* sharp-tail: a few hot, then a long flat floor (small tie at the floor).
* Hot count is min(12,V) so V<12 (incl. V=1) stays in bounds. */
for (int i = 0; i < V; i++) lo[i] = -8.f;
{ int hot = V<12 ? V : 12; for (int i = 0; i < hot; i++) lo[i] = (float)(2.0 - frand()); } break;
}
}
int main(void){
/* sizes: small for exhaustive tie detection up to near-production scale */
int sizes[] = {1, 2, 8, 64, 257, 1519}; /* 1519 ~= V/100 of GLM-5.2 */
double nucs[] = {0.001, 0.5, 0.9, 0.999}; /* tight -> almost-everything */
int n_shapes = 5;
/* temperature used by dist_build: pick a normal serving value */
g_temp = 0.7f;
int cases = 0;
for (size_t si = 0; si < sizeof(sizes)/sizeof(sizes[0]); si++){
int V = sizes[si];
/* dist_build allocates g_pbuf/g_pidx ONCE and reuses them (single-V invariant in
* real serving, where V is the constant model vocab). This sweep varies V, so free
* and force a reallocation per size -- otherwise a later, larger V would overflow
* the buffer sized for the first (smallest) V. */
free(g_pbuf); g_pbuf = NULL; free(g_pidx); g_pidx = NULL;
float *lo = malloc((size_t)V * sizeof(float));
for (int shape = 0; shape < n_shapes; shape++){
fill_shape(lo, V, shape);
for (size_t ni = 0; ni < sizeof(nucs)/sizeof(nucs[0]); ni++){
char label[32]; snprintf(label, sizeof(label), "size[%zu]/shape[%d]", si, shape);
const char *sn = (const char*[]){"uniform","peaked","geometric","plateau","sharptail"}[shape];
check_case(label, V, nucs[ni], sn, lo);
cases++;
}
}
free(lo);
}
/* guard-off path: g_nuc >= 1 must skip truncation entirely (full softmax kept) */
{
int V = 256; float lo[256];
for (int i = 0; i < V; i++) lo[i] = (float)(frand()*4 - 2);
g_nuc = 1.0f; dist_build(lo, V);
int nz = 0; for (int i = 0; i < V; i++) if (g_pbuf[i] != 0.f) nz++;
if (nz != V){ fprintf(stderr, " FAIL [guard-off nuc=1.0]: %d/%d entries kept, expected all\n", nz, V); g_nfails++; }
else printf(" ok [guard-off nuc=1.0 keep=%d]\n", nz);
cases++;
g_nuc = 0.0f; dist_build(lo, V);
nz = 0; for (int i = 0; i < V; i++) if (g_pbuf[i] != 0.f) nz++;
if (nz != V){ fprintf(stderr, " FAIL [guard-off nuc=0.0]: %d/%d entries kept, expected all\n", nz, V); g_nfails++; }
else printf(" ok [guard-off nuc=0.0 keep=%d]\n", nz);
cases++;
}
/* extreme tie edge case: V=1, single token -> keep=1 regardless of nuc */
{
float lo[1] = {5.f};
g_nuc = 0.5f; dist_build(lo, 1);
if (g_pbuf[0] == 0.f || !(fabs((double)g_pbuf[0] - 1.0) < 1e-6)){
fprintf(stderr, " FAIL [V=1]: g_pbuf[0]=%.9g, expected 1.0\n", (double)g_pbuf[0]); g_nfails++;
} else printf(" ok [V=1 keep=1]\n");
cases++;
}
printf("\ntest_topp: %d cases run, %d failure(s)\n", cases, g_nfails);
if (g_nfails){ printf("test_topp: FAIL\n"); return 1; }
printf("test_topp: ok\n");
return 0;
}
c/tests/test_uring.c
+1 -1
View File
@@ -6,7 +6,7 @@
#include <string.h>
#include <unistd.h>
#define main coli_glm_main_unused
#include "../glm.c"
#include "../colibri.c"
#undef main
static int fail(const char *s){ fprintf(stderr,"FAIL: %s\n",s); return 1; }
c/tok.h
+8
View File
@@ -42,6 +42,11 @@ typedef struct {
hmap vocab; /* stringa byte-level -> id */
hmap merges; /* "left\0right" -> rank */
char **id2str; int *id_added; int n_ids; /* id -> stringa; id_added=1 se added-token (output letterale) */
int *id_special; /* 1 = added-token con "special":true nel tokenizer:
* token di CONTROLLO (<|user|>, <|assistant|>, <sop>, ...),
* mai contenuto legittimo di una risposta. Distinto da
* id_added, che copre anche <think>/<tool_call> ("special"
* false), i quali sono testo vero e vanno renderizzati. */
Special *sp; int nsp; /* added tokens, ordinati per lunghezza decrescente */
uint32_t byte2cp[256]; int byte2cp_len[256]; char byte2str[256][3];
int16_t cp2byte[1024];
@@ -106,6 +111,7 @@ static void tok_load(Tok *T, const char *path){
T->n_ids=maxid+1;
T->id2str=calloc(T->n_ids,sizeof(char*));
T->id_added=calloc(T->n_ids,sizeof(int));
T->id_special=calloc(T->n_ids,sizeof(int));
/* vocab: stringa -> id (capacita' potenza di 2, ~2-3x) */
int vc=1; while(vc < vocab->len*2) vc<<=1;
@@ -133,6 +139,8 @@ static void tok_load(Tok *T, const char *path){
char *content=json_get(a,"content")->str; int id=(int)json_get(a,"id")->num;
T->sp[i].str=content; T->sp[i].len=(int)strlen(content); T->sp[i].id=id;
T->id2str[id]=content; T->id_added[id]=1;
jval *sf=json_get(a,"special"); /* "special": true/false */
if(sf && sf->t==J_BOOL && sf->boolean) T->id_special[id]=1;
}
qsort(T->sp,T->nsp,sizeof(Special),cmp_sp_len); /* match piu' lungo per primo */
}
c/tools/benchmark_cuda_fixture.py
+26 -3
View File
@@ -17,6 +17,12 @@ PROFILE_RE = re.compile(
r"\| attention ([0-9.]+)s .* lm_head ([0-9.]+)s \| other ([0-9.-]+)s"
)
PROFILE_KEYS = ("disk", "expert_matmul", "attention", "lm_head", "other")
P0_RE = re.compile(
r"P0-EXEC: routed CPU ([0-9.]+)s / ([0-9.]+) GB/s \(([0-9]+) row\) \| routed GPU critical ([0-9.]+)s \| "
r"router ([0-9.]+)s \| residual P2P ([0-9.]+)s / ([0-9]+) hop \| orchestration ([0-9.]+)s"
)
P0_KEYS = ("routed_cpu", "routed_cpu_gb_s", "routed_cpu_rows", "routed_gpu_critical",
"router", "p2p", "p2p_hops", "orchestration")
def parse_output(stdout: str, stderr: str = "") -> tuple[float, list[float]]:
@@ -30,11 +36,20 @@ def parse_output(stdout: str, stderr: str = "") -> tuple[float, list[float]]:
return float(speed.group(1)), [disk] + [float(value) for value in rest]
def execute(engine: str, env: dict[str, str]) -> tuple[float, list[float]]:
def parse_p0(stdout: str) -> list[float]:
"""Extract the optional PROF=1 execution-layer breakdown."""
row = P0_RE.search(stdout)
if not row:
raise RuntimeError("benchmark output missing P0-EXEC profile")
return [float(value) for value in row.groups()]
def execute(engine: str, env: dict[str, str]) -> tuple[float, list[float], list[float]]:
run = subprocess.run(
[engine, "4", "4", "4"], env=env, text=True, capture_output=True, check=True
)
return parse_output(run.stdout, run.stderr)
speed, profile = parse_output(run.stdout, run.stderr)
return speed, profile, parse_p0(run.stdout)
def main() -> None:
@@ -63,6 +78,8 @@ def main() -> None:
OMP_NUM_THREADS=str(args.threads),
OMP_PROC_BIND="spread",
OMP_PLACES="cores",
DRAFT="0",
PROF="1",
)
execute(args.engine, base | {"STATS": str(stats)})
@@ -86,13 +103,15 @@ def main() -> None:
execute(args.engine, base | extra) # warm-up
speeds = {name: [] for name in modes}
profiles = {name: [] for name in modes}
p0_profiles = {name: [] for name in modes}
names = list(modes)
for run_index in range(args.runs):
order = names[run_index % len(names):] + names[:run_index % len(names)]
for name in order:
speed, profile = execute(args.engine, base | modes[name])
speed, profile, p0 = execute(args.engine, base | modes[name])
speeds[name].append(speed)
profiles[name].append(profile)
p0_profiles[name].append(p0)
result = {}
for name in names:
@@ -103,6 +122,10 @@ def main() -> None:
key: statistics.median(row[index] for row in profiles[name])
for index, key in enumerate(PROFILE_KEYS)
},
"median_p0": {
key: statistics.median(row[index] for row in p0_profiles[name])
for index, key in enumerate(P0_KEYS)
},
}
print(json.dumps(result, indent=2))
c/tools/convert_fp8_to_int4.py
+113 -7
View File
@@ -286,6 +286,17 @@ def main():
# testa MTP a int4 = acceptance ~0-4% (misurato, issue #8): il draft sbaglia sempre
# e la speculazione non parte mai. A int8: 39-59%, 2.2-2.8 token/forward.
a.ebits = 8 if (a.mtp or a.indexer) else 4
if a.mtp and a.ebits < 8 and a.group_size <= 0:
# Non solo lossy: eh_proj ha ~20-30x di asimmetria di scala fra le due meta' di
# colonna, quindi l'int4 per-riga (UNA scala per riga) arrotonda a ZERO l'intera
# meta' embedding -> il draft non vede il token -> acceptance ~0% (issue #8).
# EN: not merely lossy: eh_proj has ~20-30x column-scale asymmetry, so per-row
# EN: int4 rounds its ENTIRE embedding half to exact zeros -> the draft cannot
# EN: see the input token -> ~0% acceptance (issue #8). A container converted
# EN: this way is repairable in place with tools/repair_mtp_int8.py.
print(f"WARNING: --mtp with --ebits {a.ebits} and per-row scales ZEROES eh_proj's "
"embedding half -> MTP acceptance ~0% (issue #8). Use the default --ebits 8, "
"or add --group-size 128 for group-scaled int4.")
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
@@ -299,6 +310,16 @@ def main():
if bits_map:
print(f"[MIXED] precision map: " + ", ".join(f"{k}={v}bit" for k,v in sorted(bits_map.items())))
# Il PIANO risolto, PRIMA di toccare qualunque cosa (#383): --mtp/--indexer cambiano il
# default di ebits a 8 (testa int4 = acceptance ~0%, issue #8) e il ramo grouped e'
# gated su bits<=4 — combinazioni sorprendenti devono mostrarsi al secondo 1 di un job
# da ore, non nel size-check dopo. EN: print the RESOLVED plan before doing anything.
mode = "MTP head only" if a.mtp else "DSA indexer only" if a.indexer else "main model"
grp = f"grouped gs={a.group_size} (fmt=4)" if (a.group_size and a.ebits <= 4) else \
(f"PER-ROW (grouped branch needs bits<=4; ebits={a.ebits} disables it)" if a.group_size else "per-row")
print(f"[PLAN] mode: {mode} | source: {'local ' + a.indir if a.indir else 'download ' + a.repo} | "
f"experts {a.ebits}-bit, embed/lm_head {a.io_bits}-bit, x {a.xbits}-bit | {grp}")
if a.selftest_nvfp4:
import torch
# 1) LUT e2m1: i 16 codici devono decodificare esattamente ai valori attesi.
@@ -379,14 +400,99 @@ def main():
if a.indir: # conversione locale (test)
shards = sorted(glob.glob(os.path.join(a.indir, "*.safetensors")))
from safetensors.numpy import save_file
# #383: se l'indice c'e', i passaggi --mtp/--indexer convertono SOLO gli shard
# che contengono i tensori richiesti (3 invece di scandire tutti i 141 — ogni
# scansione a vuoto apre comunque uno shard da 5 GB). Senza indice: scansione
# completa come prima.
# EN: #383: when the index is present, the --mtp/--indexer passes convert ONLY
# the shards that hold the requested tensors (3 instead of scanning all 141 —
# every empty scan still opens a 5 GB shard). Without the index: full scan as
# before.
if a.mtp or a.indexer:
idxp = os.path.join(a.indir, "model.safetensors.index.json")
if os.path.exists(idxp):
wmap = json.load(open(idxp))["weight_map"]
if a.mtp:
want = {v for k, v in wmap.items() if k.startswith(f"model.layers.{a.n_layers}.")}
else:
want = {v for k, v in wmap.items() if "indexer" in k and 0 <= layer_idx(k) < a.n_layers}
keep = [sp for sp in shards if os.path.basename(sp) in want]
print(f"[PLAN] index: {len(keep)}/{len(shards)} local shard(s) hold the requested tensors")
shards = keep
# BUG #355: questo ramo ignorava --mtp/--indexer. Con --mtp scriveva
# out-NNNNN (gli STESSI nomi di una conversione normale) in ebits=8 e
# keep_mtp=False -> il "secondo passaggio MTP" nella stessa outdir
# SOVRASCRIVEVA il container gia' finito con una riconversione int8
# completa, in silenzio (137/141 shard distrutti prima di accorgersene).
# Ora il ramo locale rispecchia il download path: prefisso corretto,
# flag passate, shard vuoti saltati.
prefix = "out-mtp-" if a.mtp else "out-idx-" if a.indexer else "out-"
# RIPRESA (#383): i nomi out-NNNNN contano gli shard EMESSI, non l'indice di
# input (gli shard senza tensori rilevanti non producono file), quindi "il
# file esiste" non basta per saltare il lavoro gia' fatto. Un manifest
# sidecar ricorda input -> output (o "vuoto") e con quali parametri: la
# ripresa salta solo cio' che combacia, e parametri diversi sulla stessa
# outdir vengono rifiutati invece di mescolare container (il modo #355).
# EN: RESUME (#383): out-NNNNN names count EMITTED shards, not the input
# EN: index (shards with no relevant tensors emit no file), so "the file
# EN: exists" is not enough to skip completed work. A sidecar manifest
# EN: records input -> output (or "empty") plus the conversion parameters:
# EN: resume skips only what matches, and different parameters on the same
# EN: outdir are refused instead of mixing containers (the #355 failure mode).
params = {"ebits": a.ebits, "io_bits": a.io_bits, "xbits": a.xbits,
"group_size": a.group_size, "n_layers": a.n_layers, "bits_map": bits_map}
prog_path = os.path.join(a.outdir, f".{prefix}progress.json")
prog = {}
if os.path.exists(prog_path):
try: prog = json.loads(open(prog_path).read())
except (OSError, ValueError): prog = {}
if prog and prog.get("params") != params:
print(f"ERROR: {prog_path} records a conversion with {prog.get('params')};\n"
f" this run uses {params}. Refusing to mix conversions in the same "
f"outdir — use a fresh --outdir (or delete the manifest and the "
f"{prefix}*.safetensors shards to redo).")
return
done = prog.setdefault("shards", {}); prog["params"] = params
n = 0; fresh = 0; skipped = 0
for i, sp in enumerate(shards):
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"]:
src = os.path.join(a.indir, fn)
if os.path.exists(src): shutil.copy(src, a.outdir)
print(f"converted {len(shards)} shards -> {a.outdir}")
key = os.path.basename(sp)
prev = done.get(key) # None = mai visto; "" = visto, vuoto; nome = emesso
if prev is not None and (prev == "" or os.path.exists(os.path.join(a.outdir, prev))):
if prev: n += 1
skipped += 1
continue
out = {}
convert_shard(sp, out, a.n_layers, a.ebits, a.io_bits, a.xbits,
keep_mtp=a.mtp, keep_idx=a.indexer,
group_size=a.group_size, bits_map=bits_map)
if not out: # shard senza MTP/idx: niente file (come il download path)
done[key] = ""
else:
name = f"{prefix}{n:05d}.safetensors"
save_file(out, os.path.join(a.outdir, name))
done[key] = name; n += 1; fresh += 1
tmp_prog = prog_path + ".tmp" # scrittura atomica: una ripresa non vede mai un manifest mezzo scritto
with open(tmp_prog, "w") as f: json.dump(prog, f, indent=1) # EN: atomic write: a resume never sees a half-written manifest
os.replace(tmp_prog, prog_path)
if skipped: print(f"[RESUME] {skipped} shard(s) already done in {a.outdir}, skipped")
# metadati per la conversione principale: gli stessi quattro file del download
# path — senza tokenizer.json chat/serve non partono. I passaggi mtp/idx vanno
# nella stessa outdir di un container gia' completo di metadati.
# EN: metadata for the main pass: the same four files as the download path —
# EN: chat/serve won't start without tokenizer.json. The mtp/idx passes target
# EN: an outdir whose container already has its metadata.
if not a.mtp and not a.indexer:
copied, missing = [], []
for fn in ["config.json", "tokenizer.json", "tokenizer_config.json", "generation_config.json"]:
src = os.path.join(a.indir, fn)
if os.path.exists(src): shutil.copy(src, a.outdir); copied.append(fn)
else: missing.append(fn)
print(f"[META] copied from {a.indir}: {', '.join(copied) if copied else 'nothing'}")
if missing:
print(f"[META] WARNING: not found in {a.indir}: {', '.join(missing)}"
+ (" — chat/serve need tokenizer.json" if "tokenizer.json" in missing else ""))
tag = "MTP" if a.mtp else "indexer" if a.indexer else "main"
print(f"converted {fresh} {tag} shard(s), {n} in container -> {a.outdir} ({prefix}NNNNN)")
return
# reale: scarica shard per shard, converte, cancella
c/tools/convert_olmoe.py
+23 -8
View File
@@ -3,7 +3,10 @@
Downloads or converts a local OLMoE checkpoint (e.g., allenai/OLMoE-1B-7B-0125-Instruct).
Dense weights stay as-is (engine reads BF16/F16 → F32 on load).
Expert weights get row-wise int8 quantization with float32 scales.
Expert weights get row-wise symmetric quantization to --ebits bits (default 4)
with float32 scales. Storage stays one value per int8 byte regardless of bits,
matching the engine's expert layout (olmoe.c quantize_rows) — for 4 bits the
values are simply confined to [-8, 7] with scales computed against qmax=7.
Usage:
python tools/convert_olmoe.py --repo allenai/OLMoE-1B-7B-0125-Instruct --out ./olmoe_i4
@@ -29,12 +32,21 @@ except ImportError as exc:
EXPERT_KEY_RE = r"model\.layers\.\d+\.mlp\.experts\.\d+\.(gate_proj|up_proj|down_proj)\.weight"
def quantize_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Row-wise int8 quantization. Returns (int8_weights, float32_scales)."""
def quantize_row(w: torch.Tensor, bits: int = 8) -> tuple[torch.Tensor, torch.Tensor]:
"""Row-wise symmetric quantization to `bits` (2..8).
Returns (int8_weights, float32_scales). Storage is one value per int8 byte
for every bit width — the engine dequantizes as q*scale and never assumes
the full int8 range — mirroring olmoe.c quantize_rows():
qmax = 2**(bits-1) - 1 (8 -> 127, 4 -> 7, 2 -> 1)
scale = amax(|w|, row) / qmax
q = clamp(round(w / scale), -qmax-1, qmax)
"""
qmax = (1 << (bits - 1)) - 1
w_f32 = w.float()
row_max = w_f32.abs().amax(dim=1, keepdim=True).clamp(min=1e-12)
scales = row_max / 127.0
q = (w_f32 / scales).round().clamp(-128, 127).to(torch.int8)
scales = row_max / qmax
q = (w_f32 / scales).round().clamp(-qmax - 1, qmax).to(torch.int8)
return q, scales.squeeze(1)
@@ -50,9 +62,12 @@ def main():
src.add_argument("--model", help="Local HF checkpoint directory")
ap.add_argument("--out", required=True, help="Output directory for int4 model")
ap.add_argument("--ebits", type=int, default=4,
help="Expert quant bits (4 or 8, default 4)")
help="Expert quant bits (2..8, default 4)")
args = ap.parse_args()
if not 2 <= args.ebits <= 8: # storage is int8_t; engine rejects the same range (olmoe.c)
sys.exit(f"--ebits must be 2..8 (got {args.ebits})")
if args.repo:
from huggingface_hub import snapshot_download
from huggingface_hub.errors import LocalEntryNotFoundError
@@ -96,7 +111,7 @@ def main():
for name, tensor in tensors.items():
if is_expert_weight(name):
expert_count += 1
q, scales = quantize_row(tensor)
q, scales = quantize_row(tensor, args.ebits)
total_expert_f32 += tensor.numel() * tensor.element_size()
total_expert_q += q.numel() * 1 + scales.numel() * 4
out_tensors[name] = q
@@ -109,7 +124,7 @@ def main():
ratio = total_expert_q / max(total_expert_f32, 1) * 100
print(f"ok")
print(f"\nDone. {expert_count} expert tensors quantized.")
print(f"\nDone. {expert_count} expert tensors quantized to int{args.ebits}.")
print(f"Expert storage: {total_expert_f32/1e9:.1f} GB -> {total_expert_q/1e9:.1f} GB ({ratio:.0f}%)")
print(f"Model ready at: {out}")
print(f"\nRun: SNAP={out} ./olmoe.exe 32 4 16")
c/tools/convert_olmoe_merged.py
+145
View File
@@ -0,0 +1,145 @@
#!/usr/bin/env python3
"""Convert OLMoE HuggingFace checkpoint to colibri merged int8 format.
Consolidates gate_proj, up_proj, and down_proj into a single merged tensor per expert.
This allows olmoe.c to load an expert in a single disk read call instead of 3.
Usage:
python tools/convert_olmoe_merged.py --repo allenai/OLMoE-1B-7B-0125-Instruct --out ./olmoe_merged
"""
import argparse, json, os, sys, re
from pathlib import Path
# Windows: force UTF-8 output
if sys.platform == "win32":
for s in (sys.stdout, sys.stderr):
try: s.reconfigure(encoding="utf-8")
except (AttributeError, OSError): pass
try:
import torch
from safetensors.torch import load_file, save_file
import huggingface_hub
except ImportError as exc:
sys.exit(f"Missing dependencies: {exc}. Install: pip install torch safetensors huggingface_hub")
EXPERT_KEY_RE = r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight"
def quantize_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Row-wise int8 quantization. Returns (int8_weights, float32_scales)."""
w_f32 = w.float()
row_max = w_f32.abs().amax(dim=1, keepdim=True).clamp(min=1e-12)
scales = row_max / 127.0
q = (w_f32 / scales).round().clamp(-128, 127).to(torch.int8)
return q, scales.squeeze(1)
def main():
ap = argparse.ArgumentParser(description="Convert OLMoE HF checkpoint -> colibri merged int8")
src = ap.add_mutually_exclusive_group(required=True)
src.add_argument("--repo", help="HuggingFace repo ID")
src.add_argument("--model", help="Local HF checkpoint directory")
ap.add_argument("--out", required=True, help="Output directory for merged model")
args = ap.parse_args()
if args.repo:
from huggingface_hub import snapshot_download
from huggingface_hub.errors import LocalEntryNotFoundError
print(f"Downloading/Resolving {args.repo}...")
try:
src_dir = snapshot_download(args.repo, local_files_only=True, max_workers=4)
except LocalEntryNotFoundError:
src_dir = None
if src_dir is None or not any(Path(src_dir).glob("*.safetensors")):
print("Downloading safetensors...")
src_dir = snapshot_download(args.repo, max_workers=4)
else:
src_dir = args.model
src = Path(src_dir)
if not src.is_dir():
sys.exit(f"Model directory not found: {src}")
if not (src / "config.json").is_file():
sys.exit(f"config.json missing in {src}")
out = Path(args.out)
out.mkdir(parents=True, exist_ok=True)
# Copy config.json
import shutil
shutil.copy2(src / "config.json", out / "config.json")
print(f"config.json -> {out}")
# Process safetensors
shards = sorted(src.glob("*.safetensors"))
if not shards:
sys.exit(f"No safetensors found in {src}")
print("Loading all shards to build complete state dict...")
state_dict = {}
for si, shard in enumerate(shards, 1):
print(f"Loading shard {si}/{len(shards)}: {shard.name}...")
tensors = load_file(str(shard))
state_dict.update(tensors)
# Gather experts
experts = {}
for name in list(state_dict.keys()):
m = re.match(EXPERT_KEY_RE, name)
if m:
layer_idx, expert_idx, proj = m.groups()
layer_idx = int(layer_idx)
expert_idx = int(expert_idx)
key = (layer_idx, expert_idx)
if key not in experts:
experts[key] = {}
experts[key][proj] = state_dict.pop(name)
print(f"Found {len(experts)} experts to merge.")
# Process and merge experts
out_tensors = {}
total_expert_f32 = 0
total_expert_q = 0
for (layer, expert), projs in sorted(experts.items()):
if not ("gate_proj" in projs and "up_proj" in projs and "down_proj" in projs):
sys.exit(f"Missing projection for layer {layer} expert {expert}!")
gate = projs["gate_proj"]
up = projs["up_proj"]
down = projs["down_proj"]
total_expert_f32 += (gate.numel() + up.numel() + down.numel()) * gate.element_size()
# Quantize each projection separately
q_gate, s_gate = quantize_row(gate)
q_up, s_up = quantize_row(up)
q_down, s_down = quantize_row(down)
# Merge weights and scales contiguously
merged_q = torch.cat([q_gate.flatten(), q_up.flatten(), q_down.flatten()])
merged_scales = torch.cat([s_gate, s_up, s_down])
total_expert_q += merged_q.numel() * 1 + merged_scales.numel() * 4
# Save to output
out_tensors[f"model.layers.{layer}.mlp.experts.{expert}.merged_weight"] = merged_q
out_tensors[f"model.layers.{layer}.mlp.experts.{expert}.qs"] = merged_scales
# Copy remaining dense tensors
print(f"Adding remaining {len(state_dict)} dense tensors...")
out_tensors.update(state_dict)
# Save to a single output safetensors file for simpler loading
out_file = out / "model.safetensors"
print(f"Saving merged safetensors model to {out_file}...")
save_file(out_tensors, str(out_file))
ratio = total_expert_q / max(total_expert_f32, 1) * 100
print(f"\nDone. {len(experts)} experts successfully merged and saved.")
print(f"Expert storage: {total_expert_f32/1e9:.1f} GB -> {total_expert_q/1e9:.1f} GB ({ratio:.0f}%)")
print(f"Model ready at: {out}")
if __name__ == "__main__":
main()
c/tools/make_glm_oracle.py
+40
View File
@@ -17,6 +17,45 @@ EN: computed AFTER the FP8 round-trip, so the reference matches exactly what the
EN: ingests. Default: bf16 (original oracle unchanged)."""
import json, sys, argparse
from pathlib import Path
# --- Version gate (must run BEFORE the heavy `from transformers import ...`,
# which triggers transformers' lazy-loading and can in turn reset the in-memory
# __version__ attribute; importlib.metadata reads the installed package version
# directly and is immune to that). ---
#
# GLM-5.2's MLA attention uses interleaved (DeepSeek-style) RoPE, which is what
# the C engine implements. transformers < 5.11.0 applied split-half (Llama-style)
# RoPE in GlmMoeDsa* instead; an oracle built on those versions drifts and the
# engine then scores 25/32 instead of the documented 32/32 (issue #281). Weights
# come out identical across versions — only the forward pass differs — so there
# is no safe "partial" run: a too-old transformers silently produces an invalid
# ref_glm.json. Hard-fail rather than warn. EN: same.
import transformers
from importlib.metadata import version as _pkg_version, PackageNotFoundError
_MIN_TRANSFORMERS = (5, 11)
def _tf_version_tuple():
try:
v = _pkg_version("transformers") # authoritative: installed dist metadata
except PackageNotFoundError:
v = getattr(transformers, "__version__", "0") # fallback (editable/src installs)
out = []
for part in v.split(".")[:2]: # major.minor only
out.append(int("".join(c for c in part if c.isdigit()) or "0"))
while len(out) < 2:
out.append(0)
return tuple(out[:2])
_tf_ver = _tf_version_tuple()
if _tf_ver < _MIN_TRANSFORMERS:
sys.exit(
f"\nERROR: make_glm_oracle.py requires transformers >= "
f"{'.'.join(map(str, _MIN_TRANSFORMERS))}.0 (found {_tf_ver[0]}.{_tf_ver[1]}). "
f"GLM-5.2 MLA uses interleaved RoPE; older versions apply split-half RoPE "
f"and silently produce an oracle the engine scores 25/32 against (issue #281). "
f"Upgrade: pip install -U 'transformers>=5.11'\n"
)
import torch
from transformers import GlmMoeDsaConfig, GlmMoeDsaForCausalLM
@@ -115,6 +154,7 @@ print("tf_pred:", tf_pred)
sd = model.state_dict()
unfuse_experts(sd)
Path("glm_tiny").mkdir(parents=True, exist_ok=True) # safetensors/json won't create the dir themselves
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) "
c/tools/make_olmoe_real_oracle.py
+70
View File
@@ -0,0 +1,70 @@
"""Generate reference token IDs for the real OLMoE-1B-7B model.
Uses the HF model loaded from the local cache to produce a small
reference output for olmoe.exe validation. Saves to ref_olmoe_real.json.
Usage: python tools/make_olmoe_real_oracle.py
"""
import json
import sys
from pathlib import Path
if sys.platform == "win32":
for s in (sys.stdout, sys.stderr):
try:
s.reconfigure(encoding="utf-8")
except (AttributeError, OSError):
pass
try:
import torch
from transformers import AutoTokenizer, OlmoeForCausalLM
except ImportError as exc:
sys.exit(f"Missing deps: {exc}. Run: pip install torch transformers")
MODEL_ID = "allenai/OLMoE-1B-7B-0125-Instruct"
OUT_JSON = Path(__file__).resolve().parent.parent / "ref_olmoe_real.json"
PROMPT = "The capital of France is"
MAX_NEW_TOKENS = 12
print(f"Loading tokenizer from {MODEL_ID} ...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
print("Encoding prompt ...")
enc = tokenizer(PROMPT, return_tensors="pt")
prompt_ids = enc["input_ids"][0].tolist()
print(f" Prompt IDs ({len(prompt_ids)}): {prompt_ids}")
print(f"Loading OLMoE model from {MODEL_ID} ...")
print(" (this will use ~14 GB RAM — please be patient)")
model = OlmoeForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
device_map="cpu",
low_cpu_mem_usage=True,
)
model.eval()
print(" Model loaded!")
print(f"Generating {MAX_NEW_TOKENS} tokens ...")
with torch.no_grad():
out = model.generate(
enc["input_ids"],
max_new_tokens=MAX_NEW_TOKENS,
do_sample=False,
use_cache=True,
)
full_ids = out[0].tolist()
gen_ids = full_ids[len(prompt_ids):]
print(f"Prompt IDs : {prompt_ids}")
print(f"Full IDs : {full_ids}")
print(f"Generated : {gen_ids}")
print(f"Text : {tokenizer.decode(gen_ids, skip_special_tokens=True)!r}")
payload = {"prompt_ids": prompt_ids, "full_ids": full_ids}
OUT_JSON.write_text(json.dumps(payload, indent=2))
print(f"\nSaved reference to {OUT_JSON}")
c/tools/quant_ablation.py
+15 -4
View File
@@ -118,7 +118,7 @@ def quantize_param(w, bits, group, rot=False, e8=""):
def _grid_or_e8(x, bits, group, e8):
if e8:
return _quant_e8(x.float(), group, ball=(e8 == "-e8"))
return _quant_e8(x.float(), group, bits, ball=(e8 == "-e8"))
return _quant_last_dim(x, bits, group)
@@ -169,8 +169,15 @@ def _e8_ball(y, r2=10.0):
return p
def _quant_e8(x, group, ball):
"""Blocks of 8 along the input dim; per-group scale by MSE search over RMS multiples."""
_E8_R2_REPORTED = set()
def _e8_radius(bits):
# E8 lattice: points within |p|^2<=r2 grow ~r2^4, so +1 bit (x256 codebook) needs r2 x4.
# Anchor: r2=10 is the ~2^16 E8P ball (2 bits over 8 dims). Scale from there.
return 10.0 * (4.0 ** (bits - 2))
def _quant_e8(x, group, bits, ball):
"""Blocks of 8 along the input dim; per-group scale by MSE search over RMS multiples.
ball=True clamps to the rate-scaled E8 ball for `bits`; ball=False is the unbounded ideal."""
if x.shape[-1] % 8:
raise SystemExit(f"-e8 needs input dim divisible by 8 (got {x.shape[-1]})")
g = group or x.shape[-1]
@@ -183,7 +190,11 @@ def _quant_e8(x, group, ball):
for k in (0.5, 0.7, 0.9, 1.1, 1.4, 1.8, 2.4):
s = rms * k
yb = (xg / s).reshape(-1, g // 8, 8)
p = _e8_ball(yb) if ball else _e8_nearest(yb)
p = _e8_ball(yb, _e8_radius(bits)) if ball else _e8_nearest(yb)
if ball and bits not in _E8_R2_REPORTED:
_E8_R2_REPORTED.add(bits)
import sys as _sys
_sys.stderr.write(f"[e8] bits={bits}: ball r2={_e8_radius(bits):.1f}\n")
out = (p.reshape(-1, g) * s)
err = (out - xg).pow(2).sum(-1, keepdim=True)
if best_err is None:
c/tools/repair_mtp_int8.py
+229
View File
@@ -0,0 +1,229 @@
"""Repair an existing colibri int4 container whose MTP head was quantized at int4.
WHY THIS EXISTS
`model.layers.<N>.eh_proj.weight` [D, 2D] multiplies the MTP concat
[embedding_norm ; hidden_norm], and its two column halves differ in scale by
~20-30x per row (embedding-half absmax ~0.05, hidden-half ~1.5 on GLM-5.2).
Per-row int4 uses ONE scale (= absmax/7) per row, so every embedding-half
weight lands below half a quantization step and np.rint rounds the ENTIRE
embedding half to exact zeros (packed bytes 0x88). The MTP head then drafts
garbage: acceptance ~0% (issue #8 measured 0-4% at int4; 39-59% at int8 —
which is why `convert_fp8_to_int4.py --mtp` defaults to --ebits 8).
A container converted (or downloaded) with an int4 MTP head does not need a
full re-conversion: this script re-downloads ONLY the affected dense tensors
(~355 MB of HTTP range reads against the FP8 source repo), requantizes them
at int8 with the converter's exact math, and patches the local shards in
place. Originals are kept beside as *.bak-int4.
WHAT IT TOUCHES
The MTP layer's dense tensors only (eh_proj, q_a/q_b/kv_a/kv_b/o_proj,
shared_experts.*): the ones that stream into RAM once and stay resident.
Routed experts (model.layers.<N>.mlp.experts.*) are NOT touched — they are
statistically like the main layers' experts and int4 is acceptable there.
The engine auto-detects int8 vs int4 by blob size (qt_from_disk), so no
engine or config change is needed. Cost: ~+133 MB on disk / resident RAM.
USAGE
python3 tools/repair_mtp_int8.py --snap /path/to/glm52_i4 # repair
python3 tools/repair_mtp_int8.py --snap /path/to/glm52_i4 --dry-run # inspect only
--source-repo defaults to zai-org/GLM-5.2-FP8 (the checkpoint the public
int4 containers were converted from). Requires numpy and network access;
no torch, no HF token (public repo, anonymous range reads).
"""
import argparse, glob, json, os, ssl, struct, sys, urllib.request
import numpy as np
# macOS python.org builds ship no CA bundle: use certifi when available (Linux
# system Pythons generally have working system certs and skip this).
try:
import certifi
_SSL_CTX = ssl.create_default_context(cafile=certifi.where())
except ImportError:
_SSL_CTX = ssl.create_default_context()
# ---------- HTTP range reads against the source repo ----------
def http_range(url, start, length, tries=5):
req = urllib.request.Request(url, headers={"User-Agent": "colibri-mtp-repair",
"Range": f"bytes={start}-{start+length-1}"})
for attempt in range(tries):
try:
with urllib.request.urlopen(req, timeout=30, context=_SSL_CTX) as r:
data = r.read()
if len(data) == length:
return data
except KeyboardInterrupt:
raise
except Exception as ex:
if attempt == tries - 1:
raise RuntimeError(f"range read failed for {url}: {ex}")
raise RuntimeError(f"short range read for {url}")
class SourceRepo:
def __init__(self, repo, revision="main"):
self.base = f"https://huggingface.co/{repo}/resolve/{revision}/"
with urllib.request.urlopen(self.base + "model.safetensors.index.json", timeout=30, context=_SSL_CTX) as r:
self.wmap = json.loads(r.read())["weight_map"]
self._hdr = {}
def _shard_header(self, shard):
if shard not in self._hdr:
n = struct.unpack("<Q", http_range(self.base + shard, 0, 8))[0]
self._hdr[shard] = (json.loads(http_range(self.base + shard, 8, n)), 8 + n)
return self._hdr[shard]
def meta(self, name):
shard = self.wmap.get(name)
if not shard:
return None
hdr, _ = self._shard_header(shard)
return hdr[name]
def fetch_f32(self, name):
"""Download one tensor and dequantize to f32 [O, I] (BF16 or FP8+block scales)."""
shard = self.wmap[name]
hdr, base = self._shard_header(shard)
m = hdr[name]
o0, o1 = m["data_offsets"]
raw = http_range(self.base + shard, base + o0, o1 - o0)
if m["dtype"] == "BF16":
u = np.frombuffer(raw, dtype=np.uint16).astype(np.uint32) << 16
return u.view(np.float32).reshape(m["shape"]).astype(np.float32)
if m["dtype"] == "F8_E4M3":
b = np.frombuffer(raw, dtype=np.uint8).astype(np.uint16)
sign = np.where(b & 0x80, -1.0, 1.0)
e = (b >> 3) & 0xF
mant = (b & 7).astype(np.float64)
v = (sign * np.where(e > 0, (1 + mant / 8) * np.exp2(e.astype(np.float64) - 7),
mant / 8 * np.exp2(-6.0))).reshape(m["shape"])
sn = name + "_scale_inv"
sshard = self.wmap[sn]
shdr, sbase = self._shard_header(sshard)
sm = shdr[sn]
so0, so1 = sm["data_offsets"]
sc = np.frombuffer(http_range(self.base + sshard, sbase + so0, so1 - so0),
dtype=np.float32).reshape(sm["shape"])
O, I = m["shape"]
scf = np.repeat(np.repeat(sc, 128, axis=0)[:O], 128, axis=1)[:, :I]
return (v * scf).astype(np.float32)
raise ValueError(f"{name}: unsupported source dtype {m['dtype']}")
# ---------- quantization: identical to convert_fp8_to_int4.quant_int8 ----------
def quant_int8(w):
amax = np.abs(w).max(axis=1, keepdims=True)
s = np.maximum(amax / 127, 1e-8)
q = np.clip(np.rint(w / s), -128, 127).astype(np.int8)
return q.reshape(-1).view(np.uint8).copy(), s[:, 0].astype(np.float32)
# ---------- local safetensors IO (no deps; preserves byte-identity of untouched tensors) ----------
def read_shard(path):
with open(path, "rb") as fh:
n = struct.unpack("<Q", fh.read(8))[0]
hdr = json.loads(fh.read(n))
base = 8 + n
order = sorted(((k, v) for k, v in hdr.items() if k != "__metadata__"),
key=lambda kv: kv[1]["data_offsets"][0])
out = {}
for k, v in order:
fh.seek(base + v["data_offsets"][0])
out[k] = (v["dtype"], v["shape"], fh.read(v["data_offsets"][1] - v["data_offsets"][0]))
return out, hdr.get("__metadata__")
def write_shard(path, tensors, meta):
hdr = {}
off = 0
for k, (dt, shape, raw) in tensors.items():
hdr[k] = {"dtype": dt, "shape": shape, "data_offsets": [off, off + len(raw)]}
off += len(raw)
if meta:
hdr["__metadata__"] = meta
hj = json.dumps(hdr).encode()
hj += b" " * ((8 - len(hj) % 8) % 8)
with open(path, "wb") as fh:
fh.write(struct.pack("<Q", len(hj)))
fh.write(hj)
for _, (_, _, raw) in tensors.items():
fh.write(raw)
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--snap", required=True, help="local colibri int4 container directory")
ap.add_argument("--source-repo", default="zai-org/GLM-5.2-FP8")
ap.add_argument("--revision", default="main")
ap.add_argument("--dry-run", action="store_true", help="report what would change, touch nothing")
a = ap.parse_args()
cfg = json.load(open(os.path.join(a.snap, "config.json")))
L = cfg["num_hidden_layers"]
pref = f"model.layers.{L}."
src = SourceRepo(a.source_repo, a.revision)
# find the MTP layer's dense int4 tensors across the local shards
plan = {} # shard path -> [tensor names to repair]
already_ok, skipped_experts = [], 0
for f in sorted(glob.glob(os.path.join(a.snap, "*.safetensors"))):
with open(f, "rb") as fh:
n = struct.unpack("<Q", fh.read(8))[0]
hdr = json.loads(fh.read(n))
for name, v in hdr.items():
if not name.startswith(pref) or name.endswith(".qs") or v.get("dtype") != "U8":
continue
if ".mlp.experts." in name:
skipped_experts += 1
continue
m = src.meta(name)
if m is None:
print(f" ?? {name}: not in source repo, skipping")
continue
O, I = m["shape"]
nb = v["data_offsets"][1] - v["data_offsets"][0]
if nb == O * I:
already_ok.append(name)
elif nb == O * ((I + 1) // 2):
plan.setdefault(f, []).append((name, O, I))
else:
print(f" ?? {name}: unexpected blob size {nb} for shape {O}x{I}, skipping")
n_fix = sum(len(v) for v in plan.values())
print(f"MTP layer {L}: {n_fix} dense tensor(s) at per-row int4 to repair, "
f"{len(already_ok)} already int8, {skipped_experts} routed-expert tensors left as-is")
if not n_fix:
print("nothing to do."); return
if a.dry_run:
for f, names in plan.items():
for nm, O, I in names:
print(f" would repair {nm} [{O},{I}] in {os.path.basename(f)}")
return
for f, names in plan.items():
print(f"patching {os.path.basename(f)} ({len(names)} tensors)")
tensors, meta = read_shard(f)
for nm, O, I in names:
print(f" {nm}: fetching source + requantizing at int8...", flush=True)
w = src.fetch_f32(nm)
assert w.shape == (O, I), f"{nm}: source shape {w.shape} != container {O}x{I}"
q, s = quant_int8(w)
tensors[nm] = ("U8", [len(q)], q.tobytes())
tensors[nm + ".qs"] = ("F32", [O], s.tobytes())
# verify first row against the source
loc = q[:I].view(np.int8).astype(np.float64) * s[0]
ref = w[0].astype(np.float64)
cos = float(ref @ loc / (np.linalg.norm(ref) * np.linalg.norm(loc) + 1e-30))
print(f" row-0 cosine vs source: {cos:.5f}")
bak = f + ".bak-int4"
write_shard(f + ".new", tensors, meta)
os.replace(f, bak)
os.replace(f + ".new", f)
print(f" saved; original kept as {os.path.basename(bak)}")
print("done. Re-run with --dry-run to confirm (should report 'already int8').")
if __name__ == "__main__":
main()
c/tools/test_olmoe_real.py
+91
View File
@@ -0,0 +1,91 @@
"""Bootstrap ref_olmoe_real.json by running olmoe.exe once and capturing output.
Step 1: Creates a temp ref with only prompt_ids (no full_ids).
Step 2: Runs olmoe.exe, parses the generated IDs from stdout.
Step 3: Saves {prompt_ids, full_ids} as ref_olmoe_real.json.
Step 4: Runs olmoe.exe again against the saved ref to verify determinism.
No RAM loading of the full model -- the engine streams from SSD as designed.
"""
import json
import os
import re
import subprocess
import sys
from pathlib import Path
if sys.platform == "win32":
for s in (sys.stdout, sys.stderr):
try:
s.reconfigure(encoding="utf-8")
except (AttributeError, OSError):
pass
HERE = Path(__file__).resolve().parent.parent
ext = ".exe" if sys.platform == "win32" else ""
ENGINE = HERE / f"olmoe{ext}"
SNAP = os.getenv("SNAP", str(HERE.parent / "olmoe_merged"))
REF_OUT = HERE / "ref_olmoe_real.json"
BOOTSTRAP_REF = HERE / "ref_olmoe_bootstrap.json"
PROMPT_IDS = [510, 5347, 273, 6181, 310] # "The capital of France is"
MAX_NEW = 12
CACHE_SIZE = 32 # experts cached per layer
QUANT_BITS = 8 # engine supports 2-8; 8 = int8 (lossless vs our quant)
# ── Step 1: Write bootstrap ref with dummy full_ids = prompt_ids ──────────
# olmoe.exe needs full_ids to know how many tokens to generate (nfull - np).
# We extend with MAX_NEW zeros so the engine generates MAX_NEW tokens.
bootstrap = {
"prompt_ids": PROMPT_IDS,
"full_ids": PROMPT_IDS + [0] * MAX_NEW,
}
BOOTSTRAP_REF.write_text(json.dumps(bootstrap))
print(f"Bootstrap ref written to {BOOTSTRAP_REF}")
env = {**os.environ, "SNAP": str(SNAP)}
# ── Step 2: Run engine once to capture generated IDs ─────────────────────
print(f"\n{'='*60}")
print(f"Run 1/2 — capturing engine output (cache={CACHE_SIZE}, bits={QUANT_BITS}) ...")
print(f"{'='*60}")
cmd = [str(ENGINE), str(CACHE_SIZE), str(QUANT_BITS), str(BOOTSTRAP_REF)]
r1 = subprocess.run(cmd, env=env, capture_output=True, text=True, cwd=str(HERE))
print(r1.stdout)
if r1.returncode != 0:
print("STDERR:", r1.stderr, file=sys.stderr)
sys.exit(r1.returncode)
# Parse "C engine : <id> <id> ..." line
m = re.search(r"C engine\s*:\s*([\d ]+)", r1.stdout)
if not m:
sys.exit("Could not parse 'C engine :' line from output")
gen_ids = [int(x) for x in m.group(1).split()]
print(f"Captured generated IDs: {gen_ids}")
full_ids = PROMPT_IDS + gen_ids
real_ref = {"prompt_ids": PROMPT_IDS, "full_ids": full_ids}
REF_OUT.write_text(json.dumps(real_ref, indent=2))
print(f"\nReal reference saved to {REF_OUT}")
# ── Step 3: Run engine again against real ref — verify determinism ────────
print(f"\n{'='*60}")
print("Run 2/2 — verifying determinism ...")
print(f"{'='*60}")
cmd2 = [str(ENGINE), str(CACHE_SIZE), str(QUANT_BITS), str(REF_OUT)]
r2 = subprocess.run(cmd2, env=env, capture_output=True, text=True, cwd=str(HERE))
print(r2.stdout)
if r2.returncode != 0:
print("STDERR:", r2.stderr, file=sys.stderr)
sys.exit(r2.returncode)
if "Matching tokens: 12/12" in r2.stdout or f"Matching tokens: {MAX_NEW}/{MAX_NEW}" in r2.stdout:
print("✓ Engine is DETERMINISTIC — same output on both runs!")
else:
m2 = re.search(r"Matching tokens: (\d+)/(\d+)", r2.stdout)
if m2:
print(f"⚠ Partial match: {m2.group(0)} — engine may be non-deterministic")
else:
print("⚠ Could not find matching tokens line")
BOOTSTRAP_REF.unlink(missing_ok=True)
c/version.py
+3
View File
@@ -0,0 +1,3 @@
"""Single source of truth for the colibrì version number."""
__version__ = "1.0.0"