diff --git a/README.md b/README.md index 11696de..e5671ea 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,11 @@ **Tiny engine, immense model.** Run **GLM-5.2 (744B-parameter MoE)** on a consumer machine with ~25 GB of RAM — in pure C, with zero dependencies, by streaming experts from disk. +Colibrì is a lightweight, quality-preserving MoE runtime that treats VRAM, +RAM, and storage as one managed memory hierarchy. Insufficient fast memory may +reduce speed, but the default policy never silently changes model precision or +router semantics. + ``` $ ./coli chat 🐦 colibrì v1.0 — GLM-5.2 · 744B MoE · int4 · streaming CPU @@ -285,9 +290,13 @@ cross-compiling. Requesting CUDA with a CPU-only binary, an invalid device, or an unavailable runtime fails at startup instead of silently falling back. The normal `make` build and runtime behavior are unchanged. CUDA defaults to an -expert-only accelerator: resident dense/attention tensors stay on CPU because -fixture measurements show that moving them does not help while expert I/O is -the bottleneck. `CUDA_DENSE=1` keeps the earlier all-resident experimental path. +expert-only accelerator. `CUDA_DENSE=1` additionally distributes resident +dense/attention projection tensors round-robin across the selected devices; +their projected footprint is reserved before the expert tier is placed. On six +RTX 5090s with a 150 GB expert tier, a warmed two-request/64-token GLM-5.2 run +improved from 1.650 to 2.157 aggregate tok/s (+30.8%) while retaining the full +expert tier. Treat this as an opt-in until the projected dense set and the 2 GB +per-device runtime reserve fit the target GPUs. A measured `PIN` profile can promote its hottest experts into the persistent VRAM tier while keeping the rest in RAM: @@ -295,9 +304,10 @@ VRAM tier while keeping the rest in RAM: STATS=stats.txt SNAP=/nvme/glm52_i4 ./glm 64 4 4 # collect routing frequencies first COLI_CUDA=1 COLI_GPU=0 CUDA_EXPERT_GB=16 \ PIN=stats.txt PIN_GB=160 SNAP=/nvme/glm52_i4 ./glm 64 4 4 -# multi-GPU expert tier, 96 GB total budget across six devices -COLI_CUDA=1 COLI_GPUS=0,1,2,3,4,5 CUDA_EXPERT_GB=96 \ -PIN=stats.txt PIN_GB=160 SNAP=/nvme/glm52_i4 ./glm 64 4 4 +# multi-GPU expert tier, 150 GB total budget across six 32 GB devices +COLI_CUDA=1 COLI_GPUS=0,1,2,3,4,5 CUDA_EXPERT_GB=150 \ +CUDA_DENSE=1 PIN=stats.txt PIN_GB=300 RAM_GB=226 \ +SNAP=/nvme/glm52_i4 ./glm 64 4 4 ``` Selected experts are uploaded during startup, so capacity failures occur before @@ -305,14 +315,35 @@ inference and the log reports their exact tensor footprint. The budget is clampe against free VRAM after reserving the projected dense resident set and 2 GB of runtime headroom per selected device. With `COLI_GPUS`, `CUDA_EXPERT_GB` is a total budget across the device set; experts are assigned whole to the -least-loaded device that can hold them. A NUMA-local RAM backing store is not -implemented yet. +least-loaded device that can hold them. Multi-GPU runs also default to +`PIN_FILL=1`: the measured hot set is placed first, then unused VRAM is filled +with zero-heat experts. `CUDA_RELEASE_HOST=1` (the multi-GPU default) releases +the RAM copy after a successful upload and reloads it from disk only if CUDA +later fails. Set either variable to `0` to restore the conservative behavior. +When host backing is released, placement is disjoint and staged: the hottest +prefix is loaded, uploaded to VRAM, and freed before the next-ranked suffix is +loaded into RAM. `PIN_GB` therefore describes the combined ranked set rather +than duplicate RAM and VRAM copies. On a 256 GB dual-socket host, moving from a +150 GB VRAM + 130 GB RAM placement to 150 GB VRAM + 150 GB RAM raised fixed-token +replay from 1.87 to 2.16 tok/s (+15.7%), reduced expert disk wait from 5.144s to +3.948s, and kept the projected RAM peak below `RAM_GB=226`. The cache cap adjusts +down automatically (54 to 40 in that run) so the larger pinned tier does not exceed +the process budget. Start lower on hosts with less available RAM. +MTP speculation defaults off on CUDA because cold draft routes increase expert +traffic; an explicit `DRAFT=n` still overrides the default. + +On six RTX 5090 32 GB cards with GLM-5.2 int4, a 150 GB hot-first tier sustained +0.94 token/s over a 64-token varied prompt (87.8% expert hit rate), and reached +1.64 token/s on a warmed short prompt (99.3% hit rate). The same capacity filled +without routing heat managed only 0.29 token/s, so profile quality matters more +than raw VRAM capacity. These are single-run engineering measurements, not a +portable performance guarantee. Current limitations: devices use independent contexts and synchronous -host-staged activation copies—there is no P2P/NCCL dependency yet. The kernels -are correctness-first custom kernels rather than cuBLAS/Tensor Core kernels. -This draft intentionally makes no end-to-end speedup claim before the full model -is benchmarked. +host-staged activation copies—there is no P2P/NCCL dependency yet. Independent +expert groups execute concurrently across devices, but a single expert is not +sharded. The kernels are correctness-first custom kernels rather than +cuBLAS/Tensor Core kernels. For a reproducible backend A/B without the full checkpoint, generate the deterministic 313M-parameter `glm_moe_dsa` fixture and run fixed-token replay: @@ -344,6 +375,49 @@ compatible endpoint. Nothing leaves the endpoint you configure. The terminal Useful knobs (env or flags): `--temp T` token sampling temperature (default 0.7 + nucleus 0.90 — tuned for int4; 0 = greedy), `--topp 0.7` adaptive expert top-p (30–40% less disk), `--ngen N` max tokens per answer (`:more` in chat continues a truncated one), `--repin N` adapt RAM/VRAM hot experts every N emitted tokens, `AUTOPIN=0` disable the learning cache's auto-pin, `THINK=1` enable GLM-5.2's reasoning block, `DRAFT=n` MTP draft depth, `GRAMMAR=g.gbnf` grammar-forced drafts for constrained JSON/NDJSON output (`GRAMMAR_DRAFT=n` caps the forced span), `TF=1` teacher-forcing validation, `PILOT=1` router-lookahead disk prefetch (experimental — see below), `CAP_RAISE=0` don't auto-grow the expert cache. +### Resource policy + +`coli plan` reports the planned hot (VRAM), warm (RAM), and cold backing +(disk) tiers, the reason for each placement, and the expected bottleneck. The +default `--policy quality` and `--policy balanced` modes preserve checkpoint +quantization and router decisions unless `--topk` or `--topp` is passed; those +explicit lossy overrides print a warning and proceed. + +Auto-tier plans size OpenMP from physical cores and bind workers across cores. +Memory-bound quantized kernels can regress sharply when SMT siblings compete +for limited memory channels; explicit `OMP_*` settings always take precedence. + +```bash +coli plan --model /models/glm52_i4 --policy quality +coli run --auto-tier --policy quality "Explain MoE offloading" +# Explicit research-only router reduction: +coli run --policy experimental-fast --topk 4 "Benchmark prompt" +``` + +Disk is an immutable recovery source, not a normal decode target. If the plan +leaves cold expert bytes on disk, speed depends on cache hit rate; output +quality does not. + +Cold expert reads use a deferred pipeline: resident RAM/VRAM experts execute +while missing experts are loaded in a bounded background I/O pool, then the +cold results join before the layer completes. `IO_THREADS=n` overrides the +default eight loader threads when foreground work exists. Profiling reports +both disk service time and the smaller foreground-visible wait time so overlap +is explicit rather than credited as unexplained speedup. + +`--policy balanced` enables lossless live placement (`REPIN=64`). At safe +request boundaries, a per-layer LFRU score combines decaying session frequency +with recent access and replaces at most four sufficiently colder pinned +experts. `--policy quality` leaves live replacement off by default; `REPIN=0` +always disables it. Persistent `.coli_usage` history and session-local LFRU +state remain separate. + +For single-token q4 CPU experts, gate and up projections share one OpenMP +dispatch while retaining the same per-row AVX2/NEON arithmetic. This removes +one thread-team launch per RAM expert without activation requantization or a +lower-precision fallback. It is a stepping stone toward a persistent native +CPU expert pool, not a replacement for one. + **The expert cache auto-sizes to your RAM** (since 2026-07-10): the engine now *raises* the LRU cap to fill your `--ram` budget instead of only lowering it. Before this fix a 128 GB machine ran with the same 8-experts/layer cache as a 16 GB one (issue #12) — **if you benchmarked colibrì before this date, rerun: your numbers were capped.** **Router-lookahead prefetch** (`PILOT=1`, experimental): GLM-5.2's expert routing is measurably predictable *ahead of time* — applying layer L+1's router to layer L's post-attention state recalls **71.6%** of the true top-8 (vs 41.3% for "same experts as last token"). `PILOT=1` uses this to issue next-layer expert readahead from a dedicated I/O thread while the current layer computes. On our dev box the disk is already ~80% saturated, so it measures neutral; on machines where compute and disk are balanced (like the Ryzen AI 9 in issue #12: 43% disk / 46% matmul) it should overlap real work — measurements welcome. diff --git a/c/Makefile b/c/Makefile index 3833e22..d2da448 100644 --- a/c/Makefile +++ b/c/Makefile @@ -66,7 +66,7 @@ CUDA_ARCH ?= native NVCCFLAGS ?= -O3 -std=c++17 -arch=$(CUDA_ARCH) -Xcompiler=-Wall,-Wextra PYTHON ?= python3 CUDA_OBJ = -TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_idot$(EXE) +TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) ifeq ($(CUDA),1) ifeq ($(UNAME_S),Darwin) $(error CUDA=1 is supported only on Linux) @@ -118,6 +118,11 @@ cuda-test: backend_cuda.cu backend_cuda.h tests/test_backend_cuda.cu $(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) + ./backend_cuda_bench$(EXE) + olmoe$(EXE): olmoe.c st.h json.h compat.h $(CC) $(CFLAGS) olmoe.c -o olmoe$(EXE) $(LDFLAGS) @@ -140,6 +145,9 @@ tests/test_tier$(EXE): tests/test_tier.c tier.h tests/test_grammar$(EXE): tests/test_grammar.c grammar.h $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) +tests/test_decode_batch$(EXE): tests/test_decode_batch.c decode_batch.h + $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) + tests/test_idot$(EXE): tests/test_idot.c glm.c st.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) @@ -158,7 +166,7 @@ check: $(MAKE) test clean: - rm -f olmoe$(EXE) glm$(EXE) iobench$(EXE) backend_cuda.o backend_cuda_test$(EXE) backend_metal.o backend_metal_test $(TEST_BINS) + rm -f olmoe$(EXE) glm$(EXE) iobench$(EXE) backend_cuda.o backend_cuda_test$(EXE) backend_cuda_bench$(EXE) backend_metal.o backend_metal_test $(TEST_BINS) rm -rf tests/__pycache__ -.PHONY: all glm cuda-test portable test-c test-python test check clean +.PHONY: all glm cuda-test cuda-bench portable test-c test-python test check clean diff --git a/c/backend_cuda.cu b/c/backend_cuda.cu index b40d361..fda663a 100644 --- a/c/backend_cuda.cu +++ b/c/backend_cuda.cu @@ -1,9 +1,12 @@ #include "backend_cuda.h" #include +#include #include #include +#include +#include struct ColiCudaTensor { void *weights; @@ -15,13 +18,27 @@ struct ColiCudaTensor { typedef struct { int device; - float *x, *y; - size_t x_cap, y_cap; + float *x, *y, *gate, *up; + size_t x_cap, y_cap, gate_cap, up_cap; + uint8_t *qx; float *qscale; + size_t qx_cap, qscale_cap; + float *host_x,*host_y; size_t host_x_cap,host_y_cap; + float *aq,*al,*ar,*ac; size_t aq_cap,al_cap,ar_cap,ac_cap; + cudaStream_t stream; + void *group_desc; size_t group_desc_cap; size_t tensor_count, tensor_bytes; } DeviceContext; +typedef struct { + const void *g,*u,*d; const float *gs,*us,*ds; + int gf,uf,df,rows,offset; +} GroupDesc; + static DeviceContext g_ctx[COLI_CUDA_MAX_DEVICES]; static int g_nctx; +static uint64_t g_group_calls,g_group_experts,g_group_rows; +static double g_group_h2d_ms,g_group_kernel_ms,g_group_d2h_ms; +static std::mutex g_group_stats_mu; static int cuda_ok(cudaError_t err, const char *what) { if (err == cudaSuccess) return 1; @@ -38,7 +55,7 @@ static int select_ctx(DeviceContext *ctx) { return ctx && cuda_ok(cudaSetDevice(ctx->device), "select device"); } -static size_t row_bytes(int fmt, int I) { +__host__ __device__ static size_t row_bytes(int fmt, int I) { if (fmt == 0) return (size_t)I * sizeof(float); if (fmt == 1) return (size_t)I; if (fmt == 2) return (size_t)(I + 1) / 2; @@ -53,12 +70,16 @@ __device__ static float weight_at(const void *weights, int fmt, size_t row, int const uint8_t *q = base; if (fmt == 2) { uint8_t v = q[i >> 1]; - return static_cast(((i & 1) ? (v >> 4) : (v & 15)) - 8); + int n=(i&1)?(v>>4):(v&15); return static_cast(n&8?n-16:n); } uint8_t v = q[i >> 2]; return static_cast(((v >> ((i & 3) * 2)) & 3) - 2); } +__global__ static void offset_to_signed_s4(uint8_t *q,size_t n){ + size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x;if(i=S)return; const float *xs=x+(size_t)s*K; + float v=0; for(int i=threadIdx.x;i>=1){if(threadIdx.x0?m[0]/7.f:1.f; if(!threadIdx.x)scale[s]=sc; + uint8_t *dst=q+(size_t)s*((K+1)/2); + for(int b=threadIdx.x;b<(K+1)/2;b+=blockDim.x){ + int i=b*2,a=__float2int_rn(xs[i]/sc),c=i+1= 750 + using namespace nvcuda; + int warp=threadIdx.x/32,lane=threadIdx.x%32,tile=blockIdx.x*8+warp,c=blockIdx.y; + if(tile*8>=O)return; GroupDesc d=desc[c]; + const void *w=which==0?d.g:(which==1?d.u:d.d); + const float *ws=which==0?d.gs:(which==1?d.us:d.ds); + int fmt=which==0?d.gf:(which==1?d.uf:d.df); + if(fmt!=2)return; + wmma::fragment acc; wmma::fill_fragment(acc,0); + const uint8_t *a=x+(size_t)d.offset*((K+1)/2); + const uint8_t *b=(const uint8_t*)w+(size_t)(tile*8)*((K+1)/2); + for(int k=0;k af; + wmma::fragment bf; + wmma::load_matrix_sync(af,a+k/2,K); + wmma::load_matrix_sync(bf,b+k/2,K); + wmma::mma_sync(acc,af,bf,acc); + } + __shared__ int out[8][64]; wmma::store_matrix_sync(out[warp],acc,8,wmma::mem_row_major); + for(int i=lane;i<64;i+=32){int s=i/8,o=tile*8+i%8; + if(s=d.rows) return; + const void *w=which?d.u:d.g; const float *sc=which?d.us:d.gs; int fmt=which?d.uf:d.gf; + size_t rb=row_bytes(fmt,D),row=(size_t)o*rb; const float *xs=x+(size_t)(d.offset+s)*D; + float sum=0; for(int i=threadIdx.x;i>=1){ if(threadIdx.x=d.rows) return; + size_t rb=row_bytes(d.df,I),row=(size_t)o*rb; const float *xs=x+(size_t)(d.offset+s)*I; + float sum=0; for(int i=threadIdx.x;i>=1){ if(threadIdx.x>4; *lo=(float)(a&8?a-16:a); *hi=(float)(b&8?b-16:b); +} + +/* Exact low-row W4A32 path. It consumes each packed weight byte once instead + * of routing both nibbles through weight_at(), preserving FP32 activations. */ +__global__ static void grouped_hidden_w4(float *y,const float *x,const GroupDesc *desc, + int I,int D,int which){ + int o=blockIdx.x,s=blockIdx.y,c=blockIdx.z;GroupDesc d=desc[c];if(s>=d.rows)return; + const uint8_t *w=(const uint8_t*)(which?d.u:d.g);const float *sc=which?d.us:d.gs; + const uint8_t *row=w+(size_t)o*((D+1)/2);const float *xs=x+(size_t)(d.offset+s)*D; + float sum=0;for(int b=threadIdx.x;b<(D+1)/2;b+=blockDim.x){float a,z;unpack_s4(row[b],&a,&z); + int i=b*2;sum+=xs[i]*a;if(i+1>=1){if(threadIdx.x=d.rows)return; + const uint8_t *gr=(const uint8_t*)d.g+(size_t)o*((D+1)/2); + const uint8_t *ur=(const uint8_t*)d.u+(size_t)o*((D+1)/2); + const float *xs=x+(size_t)(d.offset+s)*D;float ga=0,ua=0; + for(int b=threadIdx.x;b<(D+1)/2;b+=blockDim.x){float g0,g1,u0,u1;unpack_s4(gr[b],&g0,&g1);unpack_s4(ur[b],&u0,&u1); + int i=b*2;ga+=xs[i]*g0;ua+=xs[i]*u0;if(i+1>=1){if(threadIdx.x=d.rows)return; + const uint8_t *row=(const uint8_t*)d.d+(size_t)o*((I+1)/2); + const float *xs=x+(size_t)(d.offset+s)*I;float sum=0; + for(int b=threadIdx.x;b<(I+1)/2;b+=blockDim.x){float a,z;unpack_s4(row[b],&a,&z); + int i=b*2;sum+=xs[i]*a;if(i+1>=1){if(threadIdx.x= bytes) return 1; if (*ptr) cudaFree(*ptr); @@ -91,6 +243,16 @@ static int reserve(float **ptr, size_t *cap, size_t bytes) { return 1; } +static int reserve_bytes(void **ptr,size_t *cap,size_t bytes){ + if(*cap>=bytes) return 1; if(*ptr) cudaFree(*ptr); *ptr=nullptr; *cap=0; + if(!cuda_ok(cudaMalloc(ptr,bytes),"descriptor allocation")) return 0; *cap=bytes; return 1; +} + +static int reserve_pinned(float **ptr,size_t *cap,size_t bytes){ + if(*cap>=bytes)return 1;if(*ptr)cudaFreeHost(*ptr);*ptr=nullptr;*cap=0; + if(!cuda_ok(cudaMallocHost(ptr,bytes),"pinned staging allocation"))return 0;*cap=bytes;return 1; +} + extern "C" int coli_cuda_init(const int *devices, int count) { int available = 0; if (!devices || count < 1 || count > COLI_CUDA_MAX_DEVICES) return 0; @@ -114,6 +276,9 @@ extern "C" int coli_cuda_init(const int *devices, int count) { if (!select_ctx(ctx)) { g_nctx = 0; return 0; } cudaDeviceProp prop{}; if (!cuda_ok(cudaGetDeviceProperties(&prop, device), "device properties")) { g_nctx = 0; return 0; } + if(!cuda_ok(cudaStreamCreateWithFlags(&ctx->stream,cudaStreamNonBlocking),"stream creation")){ + g_nctx=0;return 0; + } g_nctx++; std::fprintf(stderr, "[CUDA] device %d: %s, %.1f GB VRAM, sm_%d%d\n", device, prop.name, prop.totalGlobalMem / 1e9, prop.major, prop.minor); @@ -127,8 +292,24 @@ extern "C" void coli_cuda_shutdown(void) { if (!select_ctx(ctx)) continue; if (ctx->x) cudaFree(ctx->x); if (ctx->y) cudaFree(ctx->y); - ctx->x = ctx->y = nullptr; - ctx->x_cap = ctx->y_cap = 0; + if (ctx->gate) cudaFree(ctx->gate); + if (ctx->up) cudaFree(ctx->up); + if (ctx->qx) cudaFree(ctx->qx); + if (ctx->qscale) cudaFree(ctx->qscale); + if(ctx->aq)cudaFree(ctx->aq);if(ctx->al)cudaFree(ctx->al);if(ctx->ar)cudaFree(ctx->ar);if(ctx->ac)cudaFree(ctx->ac); + if (ctx->host_x) cudaFreeHost(ctx->host_x); + if (ctx->host_y) cudaFreeHost(ctx->host_y); + 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->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->group_desc=nullptr; ctx->group_desc_cap=0; } g_nctx = 0; } @@ -155,6 +336,13 @@ extern "C" void coli_cuda_stats(int device, size_t *tensor_count, size_t *tensor if (tensor_bytes) *tensor_bytes = bytes; } +extern "C" void coli_cuda_group_stats(uint64_t *calls, uint64_t *experts, uint64_t *rows, + double *h2d_ms, double *kernel_ms, double *d2h_ms) { + if(calls) *calls=g_group_calls; if(experts) *experts=g_group_experts; if(rows) *rows=g_group_rows; + if(h2d_ms) *h2d_ms=g_group_h2d_ms; if(kernel_ms) *kernel_ms=g_group_kernel_ms; + if(d2h_ms) *d2h_ms=g_group_d2h_ms; +} + extern "C" int coli_cuda_tensor_upload(ColiCudaTensor **tensor, const void *weights, const float *scales, int fmt, int I, int O, int device) { @@ -174,6 +362,8 @@ extern "C" int coli_cuda_tensor_upload(ColiCudaTensor **tensor, coli_cuda_tensor_free(t); return 0; } + if(fmt==2){offset_to_signed_s4<<<(unsigned)((t->weight_bytes+255)/256),256>>>((uint8_t*)t->weights,t->weight_bytes); + if(!cuda_ok(cudaGetLastError(),"int4 weight conversion")){coli_cuda_tensor_free(t);return 0;}} if (fmt) { if (!cuda_ok(cudaMalloc(&t->scales, (size_t)O * sizeof(float)), "scale allocation") || !cuda_ok(cudaMemcpy(t->scales, scales, (size_t)O * sizeof(float), cudaMemcpyHostToDevice), "scale upload")) { @@ -207,6 +397,151 @@ extern "C" int coli_cuda_matmul(ColiCudaTensor **tensor, return 1; } +extern "C" int coli_cuda_expert_mlp(ColiCudaTensor *gate, ColiCudaTensor *up, + ColiCudaTensor *down, float *y, + const float *x, int S) { + if (!gate || !up || !down || !x || !y || S < 1 || + gate->device != up->device || gate->device != down->device || + gate->I != up->I || gate->O != up->O || + down->I != gate->O || down->O != gate->I) return 0; + DeviceContext *ctx = find_ctx(gate->device); + if (!select_ctx(ctx)) return 0; + int D = gate->I, I = gate->O; + size_t xb=(size_t)S*D*sizeof(float), ib=(size_t)S*I*sizeof(float); + size_t yb=(size_t)S*D*sizeof(float); + if (!reserve(&ctx->x,&ctx->x_cap,xb) || !reserve(&ctx->y,&ctx->y_cap,yb) || + !reserve(&ctx->gate,&ctx->gate_cap,ib) || !reserve(&ctx->up,&ctx->up_cap,ib)) return 0; + if (!cuda_ok(cudaMemcpy(ctx->x,x,xb,cudaMemcpyHostToDevice),"expert input upload")) return 0; + dim3 hidden_grid((unsigned)I,(unsigned)S), output_grid((unsigned)D,(unsigned)S); + quant_matmul<<>>(ctx->gate,ctx->x,gate->weights,gate->scales, + gate->fmt,S,D,I,row_bytes(gate->fmt,D)); + quant_matmul<<>>(ctx->up,ctx->x,up->weights,up->scales, + up->fmt,S,D,I,row_bytes(up->fmt,D)); + size_t n=(size_t)S*I; + silu_mul<<<(unsigned)((n+255)/256),256>>>(ctx->gate,ctx->up,n); + quant_matmul<<>>(ctx->y,ctx->gate,down->weights,down->scales, + down->fmt,S,I,D,row_bytes(down->fmt,I)); + if (!cuda_ok(cudaGetLastError(),"expert MLP launch") || + !cuda_ok(cudaMemcpy(y,ctx->y,yb,cudaMemcpyDeviceToHost),"expert output download")) return 0; + return 1; +} + +extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates, + ColiCudaTensor *const *ups, + ColiCudaTensor *const *downs, + const int *rows, int count, + float *y, const float *x) { + if (!gates || !ups || !downs || !rows || !x || !y || count < 1) return 0; + ColiCudaTensor *first=gates[0]; + if (!first) return 0; + int device=first->device,D=first->I,I=first->O,total=0,max_rows=0; + GroupDesc host[64]; if(count>64) return 0; + int all_s4=1; + for(int c=0;cdevice!=device||u->device!=device||d->device!=device|| + g->I!=D||u->I!=D||g->O!=I||u->O!=I||d->I!=I||d->O!=D) return 0; + host[c]={g->weights,u->weights,d->weights,g->scales,u->scales,d->scales, + g->fmt,u->fmt,d->fmt,rows[c],total}; + all_s4&=g->fmt==2&&u->fmt==2&&d->fmt==2; + total+=rows[c]; if(rows[c]>max_rows) max_rows=rows[c]; + } + DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0; + size_t xb=(size_t)total*D*sizeof(float), ib=(size_t)total*I*sizeof(float); + if(!reserve(&ctx->x,&ctx->x_cap,xb)||!reserve(&ctx->y,&ctx->y_cap,xb)|| + !reserve(&ctx->gate,&ctx->gate_cap,ib)||!reserve(&ctx->up,&ctx->up_cap,ib)|| + !reserve_bytes(&ctx->group_desc,&ctx->group_desc_cap,(size_t)count*sizeof(GroupDesc))) return 0; + int async=!getenv("COLI_CUDA_ASYNC")||atoi(getenv("COLI_CUDA_ASYNC")); + if(async&&(!reserve_pinned(&ctx->host_x,&ctx->host_x_cap,xb)|| + !reserve_pinned(&ctx->host_y,&ctx->host_y_cap,xb)))return 0; + cudaError_t copy_desc=async?cudaMemcpyAsync(ctx->group_desc,host,(size_t)count*sizeof(GroupDesc), + cudaMemcpyHostToDevice,ctx->stream) + :cudaMemcpy(ctx->group_desc,host,(size_t)count*sizeof(GroupDesc),cudaMemcpyHostToDevice); + if(!cuda_ok(copy_desc,"expert group descriptors"))return 0; + int profile=getenv("COLI_CUDA_PROFILE")&&atoi(getenv("COLI_CUDA_PROFILE")); + cudaEvent_t ev[4]={}; + if(profile) for(int i=0;i<4;i++) if(!cuda_ok(cudaEventCreate(&ev[i]),"profile event")) profile=0; + if(profile) cudaEventRecord(ev[0],ctx->stream); + if(async)std::memcpy(ctx->host_x,x,xb); + cudaError_t copy_x=async?cudaMemcpyAsync(ctx->x,ctx->host_x,xb,cudaMemcpyHostToDevice,ctx->stream) + :cudaMemcpy(ctx->x,x,xb,cudaMemcpyHostToDevice); + if(!cuda_ok(copy_x,"expert group input upload")) return 0; + if(profile) cudaEventRecord(ev[1],ctx->stream); + GroupDesc *dev=(GroupDesc*)ctx->group_desc; + int tc=getenv("COLI_CUDA_TC_INT4")&&atoi(getenv("COLI_CUDA_TC_INT4")); + tc=tc&&all_s4&&D%32==0&&I%32==0&&D%8==0&&I%8==0; + int tc_min=getenv("COLI_CUDA_TC_MIN_ROWS")?atoi(getenv("COLI_CUDA_TC_MIN_ROWS")):8; + for(int c=0;c=tc_min; + if(tc){ + size_t qb=(size_t)(total+7)*(size_t)(D>I?D:I)/2; + if(!reserve_bytes((void**)&ctx->qx,&ctx->qx_cap,qb)|| + !reserve(&ctx->qscale,&ctx->qscale_cap,(size_t)(total+7)*sizeof(float)))return 0; + cudaMemsetAsync(ctx->qx,0,qb,ctx->stream); + quantize_s4_rows<<stream>>>(ctx->qx,ctx->qscale,ctx->x,total,D); + grouped_s4_wmma<<stream>>>(ctx->gate,ctx->qx,ctx->qscale,dev,D,I,0); + grouped_s4_wmma<<stream>>>(ctx->up,ctx->qx,ctx->qscale,dev,D,I,1); + silu_mul<<<(unsigned)(((size_t)total*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)total*I); + quantize_s4_rows<<stream>>>(ctx->qx,ctx->qscale,ctx->gate,total,I); + grouped_s4_wmma<<stream>>>(ctx->y,ctx->qx,ctx->qscale,dev,I,D,2); + }else if(all_s4&&(!getenv("COLI_CUDA_W4_PACKED")||atoi(getenv("COLI_CUDA_W4_PACKED")))){ + dim3 hg((unsigned)I,(unsigned)max_rows,(unsigned)count),og((unsigned)D,(unsigned)max_rows,(unsigned)count); + int dual=!getenv("COLI_CUDA_DUAL_PROJ")||atoi(getenv("COLI_CUDA_DUAL_PROJ")); + if(dual)grouped_hidden_w4_dual<<stream>>>(ctx->gate,ctx->up,ctx->x,dev,I,D); + else{ + grouped_hidden_w4<<stream>>>(ctx->gate,ctx->x,dev,I,D,0); + grouped_hidden_w4<<stream>>>(ctx->up,ctx->x,dev,I,D,1); + } + silu_mul<<<(unsigned)(((size_t)total*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)total*I); + grouped_down_w4<<stream>>>(ctx->y,ctx->gate,dev,D,I); + }else{ + dim3 hg((unsigned)I,(unsigned)max_rows,(unsigned)count),og((unsigned)D,(unsigned)max_rows,(unsigned)count); + grouped_hidden<<stream>>>(ctx->gate,ctx->x,dev,I,D,0); + grouped_hidden<<stream>>>(ctx->up,ctx->x,dev,I,D,1); + silu_mul<<<(unsigned)(((size_t)total*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)total*I); + grouped_down<<stream>>>(ctx->y,ctx->gate,dev,D,I); + } + if(profile) cudaEventRecord(ev[2],ctx->stream); + if(!async&&!cuda_ok(cudaStreamSynchronize(ctx->stream),"expert group synchronize"))return 0; + cudaError_t copy_y=async?cudaMemcpyAsync(ctx->host_y,ctx->y,xb,cudaMemcpyDeviceToHost,ctx->stream) + :cudaMemcpy(y,ctx->y,xb,cudaMemcpyDeviceToHost); + if(!cuda_ok(cudaGetLastError(),"expert group launch")||!cuda_ok(copy_y,"expert group output download"))return 0; + if(async){if(!cuda_ok(cudaStreamSynchronize(ctx->stream),"expert group synchronize"))return 0; + std::memcpy(y,ctx->host_y,xb);} + if(profile){ + cudaEventRecord(ev[3],ctx->stream); cudaEventSynchronize(ev[3]); float a=0,b=0,c=0; + cudaEventElapsedTime(&a,ev[0],ev[1]); cudaEventElapsedTime(&b,ev[1],ev[2]); + cudaEventElapsedTime(&c,ev[2],ev[3]); + { std::lock_guard lock(g_group_stats_mu); + g_group_h2d_ms+=a; g_group_kernel_ms+=b; g_group_d2h_ms+=c; } + for(int i=0;i<4;i++) cudaEventDestroy(ev[i]); + } + { std::lock_guard lock(g_group_stats_mu); + g_group_calls++; g_group_experts+=(uint64_t)count; g_group_rows+=(uint64_t)total; } + return 1; +} + +extern "C" int coli_cuda_attention_absorb(ColiCudaTensor *w,float *ctx,const float *q, + const float *latent,const float *rope,int H,int Q, + int R,int V,int K,int T,float scale){ + if(!w||!ctx||!q||!latent||!rope||H<1||Q<1||R<1||V<1||K<1||K>512||T<1||T>4096|| + w->I!=K||w->O!=H*(Q+V))return 0; + DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0; + size_t qb=(size_t)H*(Q+R)*sizeof(float),lb=(size_t)T*K*sizeof(float); + size_t rb=(size_t)T*R*sizeof(float),cb=(size_t)H*V*sizeof(float); + if(!reserve(&dc->aq,&dc->aq_cap,qb)||!reserve(&dc->al,&dc->al_cap,lb)|| + !reserve(&dc->ar,&dc->ar_cap,rb)||!reserve(&dc->ac,&dc->ac_cap,cb))return 0; + if(!cuda_ok(cudaMemcpyAsync(dc->aq,q,qb,cudaMemcpyHostToDevice,dc->stream),"attention q upload")|| + !cuda_ok(cudaMemcpyAsync(dc->al,latent,lb,cudaMemcpyHostToDevice,dc->stream),"attention latent upload")|| + !cuda_ok(cudaMemcpyAsync(dc->ar,rope,rb,cudaMemcpyHostToDevice,dc->stream),"attention rope upload"))return 0; + size_t shared=(size_t)(2*K+T)*sizeof(float); + attention_absorb_kernel<<stream>>>(dc->ac,dc->aq,dc->al,dc->ar,w->weights,w->scales, + w->fmt,H,Q,R,V,K,T,scale); + if(!cuda_ok(cudaGetLastError(),"attention absorb launch")|| + !cuda_ok(cudaMemcpyAsync(ctx,dc->ac,cb,cudaMemcpyDeviceToHost,dc->stream),"attention context download")|| + !cuda_ok(cudaStreamSynchronize(dc->stream),"attention synchronize"))return 0; + return 1; +} + extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) { if (!tensor) return; DeviceContext *ctx = find_ctx(tensor->device); diff --git a/c/backend_cuda.h b/c/backend_cuda.h index 5a7377b..2a1979c 100644 --- a/c/backend_cuda.h +++ b/c/backend_cuda.h @@ -21,6 +21,8 @@ int coli_cuda_device_at(int index); int coli_cuda_mem_info(int device, size_t *free_bytes, size_t *total_bytes); /* device < 0 returns aggregate statistics for all configured devices. */ void coli_cuda_stats(int device, size_t *tensor_count, size_t *tensor_bytes); +void coli_cuda_group_stats(uint64_t *calls, uint64_t *experts, uint64_t *rows, + double *h2d_ms, double *kernel_ms, double *d2h_ms); /* Upload without executing, so capacity failures happen during model startup. */ int coli_cuda_tensor_upload(ColiCudaTensor **tensor, @@ -38,6 +40,25 @@ int coli_cuda_matmul(ColiCudaTensor **tensor, const void *weights, const float *scales, int fmt, int S, int I, int O, int device); +/* Fused expert pipeline: y = down(silu(gate(x)) * up(x)). All three tensors + * must already be resident on one device. Activations cross PCIe once in + * each direction instead of once per matrix. */ +int coli_cuda_expert_mlp(ColiCudaTensor *gate, ColiCudaTensor *up, + ColiCudaTensor *down, float *y, const float *x, int S); + +/* Packed group of same-shaped experts. Inputs and outputs contain sum(rows) + * consecutive [D] rows in call order. */ +int coli_cuda_expert_group(ColiCudaTensor *const *gates, + ColiCudaTensor *const *ups, + ColiCudaTensor *const *downs, + const int *rows, int count, + float *y, const float *x); + +/* Decode-only MLA weight-absorption core for one token. kv_b is [H*(Q+V),K]. */ +int coli_cuda_attention_absorb(ColiCudaTensor *kv_b,float *ctx,const float *q, + const float *latent,const float *rope,int H,int Q, + int R,int V,int K,int T,float attention_scale); + void coli_cuda_tensor_free(ColiCudaTensor *tensor); size_t coli_cuda_tensor_bytes(const ColiCudaTensor *tensor); int coli_cuda_tensor_device(const ColiCudaTensor *tensor); diff --git a/c/coli b/c/coli index e4e89ea..05a14a7 100755 --- a/c/coli +++ b/c/coli @@ -127,6 +127,7 @@ def resource_request(a, env): def env_for(a): e = dict(os.environ, SNAP=a.model) + e["COLI_POLICY"]=a.policy if a.ram: e["RAM_GB"]=str(a.ram) if a.ngen: e["NGEN"]=str(a.ngen) if a.topp: e["TOPP"]=str(a.topp) @@ -147,7 +148,7 @@ def env_for(a): if a.vram and a.gpu!="none": e["CUDA_EXPERT_GB"]=str(a.vram) try: ram,ctx,devices,vram=resource_request(a,e) - plan=build_plan(a.model,ram,ctx,devices,vram) + plan=build_plan(a.model,ram,ctx,devices,vram,policy=a.policy) except (OSError,ValueError,json.JSONDecodeError) as error: sys.exit(f"{C.yel}invalid resource plan:{C.r} {error}") has_cuda=cuda_binary() @@ -325,7 +326,7 @@ def cmd_plan(a): ram,ctx,devices,vram=resource_request(a,os.environ) if ctx<1: raise ValueError("--ctx must be positive") if a.vram<0: raise ValueError("--vram cannot be negative") - plan=build_plan(a.model,ram,ctx,devices,vram) + plan=build_plan(a.model,ram,ctx,devices,vram,policy=a.policy) except (OSError, ValueError, json.JSONDecodeError) as error: sys.exit(f"{C.yel}cannot create resource plan:{C.r} {error}") if a.json: @@ -516,6 +517,9 @@ def main(): common.add_argument("--ctx",type=int,default=0) common.add_argument("--gpu",default=None,help="auto, none, or a device list such as 0,1") common.add_argument("--vram",type=float,default=0,help="total VRAM budget in GB (0=auto)") + common.add_argument("--policy",choices=("quality","balanced","experimental-fast"), + default=os.environ.get("COLI_POLICY","quality"), + help="resource policy (explicit --topk/--topp overrides warn and proceed)") common.add_argument("--repin", type=int, default=0, help="adapt RAM/VRAM experts every N tokens") common.add_argument("--cap", type=int, default=8); common.add_argument("--ngen", type=int, default=1024) # rete di sicurezza: la fine vera la decidono gli stop token common.add_argument("--topp", type=float, default=0); common.add_argument("--topk", type=int, default=0) diff --git a/c/decode_batch.h b/c/decode_batch.h new file mode 100644 index 0000000..ad3d2ed --- /dev/null +++ b/c/decode_batch.h @@ -0,0 +1,37 @@ +#ifndef COLIBRI_DECODE_BATCH_H +#define COLIBRI_DECODE_BATCH_H + +#include +#include +#include + +/* `base` belongs to one sequence's KV state. Keeping this arithmetic in a + * model-independent seam makes ragged decode row ownership directly testable. */ +static inline float *coli_kv_row(float *base, int position, int width) +{ + return base + (size_t)position * (size_t)width; +} + +typedef struct { + unsigned long long id, bytes; + int slot, max_tokens; + float temperature, top_p; +} ColiSubmit; + +/* Parse the textual header. The payload is read separately using `bytes`, so + * it may contain newlines. Reject trailing fields to keep framing unambiguous. */ +static inline int coli_submit_parse(const char *line, ColiSubmit *s) +{ + char tail; + if (!line || !s || + sscanf(line, "SUBMIT %llu %d %llu %d %f %f %c", &s->id, &s->slot, + &s->bytes, &s->max_tokens, &s->temperature, &s->top_p, + &tail) != 6) + return 0; + return s->id > 0 && s->bytes <= (16u << 20) && s->slot >= 0 && s->max_tokens >= 1 && + isfinite(s->temperature) && isfinite(s->top_p) && + s->temperature >= 0 && s->temperature <= 2 && + s->top_p > 0 && s->top_p <= 1; +} + +#endif diff --git a/c/glm.c b/c/glm.c index 981e6d7..1f1e603 100644 --- a/c/glm.c +++ b/c/glm.c @@ -27,6 +27,7 @@ #include /* PIPE ready-flags/job queue + PILOT_REAL cross-layer handshake */ #include /* sched_yield: PIPE spin / PILOT barrier */ #include +#include #if defined(__APPLE__) || defined(__linux__) #include #include /* mlock: inchioda le pagine in RAM / wire pages into RAM */ @@ -36,6 +37,7 @@ #include "tok.h" #include "tier.h" #include "grammar.h" /* metodo F: draft grammaticali (#48) */ +#include "decode_batch.h" #ifdef _OPENMP #include /* scratch per-thread nell'attention */ #else @@ -130,6 +132,11 @@ typedef struct { char disk_path[2048]; } KVState; +typedef struct { + KVState *kv; + int token, pos; +} DecodeRow; + typedef struct { Cfg c; shards S; int ebits, dbits; /* bit expert / bit densa */ @@ -146,6 +153,7 @@ typedef struct { ESlot **pin; int *npin; /* HOT-STORE: expert pinnati in RAM (mai evicted) */ uint32_t **eusage; /* contatori persistenti (per STATS/PIN) */ uint32_t **eheat; /* calore recente per promotion/demotion live */ + uint32_t **elast, eaccess_clock; /* recency per LFRU session-local */ /* DSA lightning indexer (attivo solo se i pesi out-idx-* sono presenti) */ int has_dsa; QT *ix_wq, *ix_wk, *ix_wp; /* per layer FULL: wq_b, wk, weights_proj */ @@ -161,7 +169,8 @@ typedef struct { uint64_t eclock, hits, miss, ereq; uint64_t gpu_expert_calls; int gpu_expert_count; int64_t gpu_expert_bytes; uint64_t n_fw, n_emit; /* metodo E: forward di decode / token emessi */ - double t_edisk, t_emm, t_attn, t_kvb, t_head;/* profiling: dove va il tempo (sempre attivo) */ + double t_edisk, t_ewait, t_emm, t_attn, t_kvb, t_head;/* profiling: dove va il tempo */ + double t_aproj,t_acore,t_aout; /* attention breakdown */ int64_t resident_bytes; } Model; @@ -170,6 +179,7 @@ static void usage_save(Model *m); /* cache che impara: definita accanto a static int g_cuda_enabled; static double g_cuda_expert_gb; static int g_cuda_dense; +static int g_cuda_release_host; static int g_cuda_devices[COLI_CUDA_MAX_DEVICES], g_cuda_ndev, g_cuda_rr; static int64_t g_cuda_dense_projected[COLI_CUDA_MAX_DEVICES]; static void qt_cuda_reset(QT *t){ @@ -188,6 +198,14 @@ static void cuda_stats_print(void){ coli_cuda_stats(g_cuda_devices[i],&n,&b); fprintf(stderr,"[CUDA] device %d: %zu tensors, %.2f GB\n",g_cuda_devices[i],n,b/1e9); } + uint64_t calls=0,experts=0,rows=0; double h2d=0,kernel=0,d2h=0; + coli_cuda_group_stats(&calls,&experts,&rows,&h2d,&kernel,&d2h); + if(calls) fprintf(stderr,"[CUDA] expert groups: %llu call, %llu expert, %llu righe " + "(%.2f expert/call)%s\n",(unsigned long long)calls,(unsigned long long)experts, + (unsigned long long)rows,(double)experts/calls, + getenv("COLI_CUDA_PROFILE")?"; timing sotto":""); + if(calls&&getenv("COLI_CUDA_PROFILE")) fprintf(stderr, + "[CUDA] expert groups timing: H2D %.1f ms | kernel %.1f ms | D2H %.1f ms\n",h2d,kernel,d2h); } static int parse_cuda_devices(const char *list, int *out){ if(!list||!*list) return 0; @@ -280,6 +298,53 @@ static void matmul_i4(float *y, const float *x, const uint8_t *q4, const float * if(i>1]; int lo=(int)(byte&0xF)-8; a += xs[i]*(float)lo; } y[(int64_t)s*O+o]=a*sc; } } } +/* Decode hot path for gate+up: same exact q4 dot products as matmul_i4, but one + * OpenMP dispatch covers both matrices. KTransformers uses persistent pools; + * this keeps colibri dependency-free while removing one team launch/expert. */ +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>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 + for(;i+1>1]; a+=x[i]*(float)((b&15)-8)+x[i+1]*(float)((b>>4)-8); } + if(i>1]&15)-8); + (zfmt==2&&wu->fmt==2&&wg->I==wu->I&&wg->O==wu->O) + matmul_i4_pair(g,u,x,wg->q4,wg->s,wu->q4,wu->s,wg->I,wg->O); + else { matmul_qt(g,x,wg,S); matmul_qt(u,x,wu,S); } +} /* y[S,O] = x[S,I] @ W^T con W int2 impacchettato (4 valori/byte) + scala[O]. nibble 2-bit -> [-2,1]. */ 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; @@ -861,6 +926,7 @@ static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits m->eroute=calloc(NR,sizeof(int*)); m->enr=calloc(NR,sizeof(int)); m->pin=calloc(NR,sizeof(ESlot*)); m->npin=calloc(NR,sizeof(int)); m->eusage=calloc(NR,sizeof(uint32_t*)); m->eheat=calloc(NR,sizeof(uint32_t*)); + m->elast=calloc(NR,sizeof(uint32_t*)); m->kv=calloc(1,sizeof(KVState)); m->kv_start=m->kv->kv_start=calloc(NR,sizeof(int)); for(int i=0;in_layers;i++){ @@ -891,6 +957,7 @@ static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits m->eroute[i]=calloc(c->topk,sizeof(int)); /* metodo C: ultimo routing del layer */ m->eusage[i]=calloc(c->n_experts,sizeof(uint32_t)); m->eheat[i]=calloc(c->n_experts,sizeof(uint32_t)); + m->elast[i]=calloc(c->n_experts,sizeof(uint32_t)); } #undef P } @@ -936,6 +1003,7 @@ static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits m->eroute[i]=calloc(c->topk,sizeof(int)); m->eusage[i]=calloc(c->n_experts,sizeof(uint32_t)); m->eheat[i]=calloc(c->n_experts,sizeof(uint32_t)); + m->elast[i]=calloc(c->n_experts,sizeof(uint32_t)); m->kv_start[i]=-1; /* KV MTP: parte dalla prima posizione di decode */ #undef PM } @@ -1279,6 +1347,24 @@ static inline void pipe_wait(int q){ while(!atomic_load_explicit(&g_pp.ready[q],memory_order_acquire)) sched_yield(); } +#ifdef COLI_CUDA +static void expert_host_release(Model *m, ESlot *s){ + if(!s->slab&&!s->fslab) return; +#if defined(__APPLE__) || defined(__linux__) + if(s->slab) munlock(s->slab,(size_t)s->slab_cap); + if(s->fslab) munlock(s->fslab,(size_t)s->fslab_cap*sizeof(float)); +#endif + int64_t bytes=qt_bytes(&s->g)+qt_bytes(&s->u)+qt_bytes(&s->d); + free(s->slab); free(s->fslab); s->slab=NULL; s->fslab=NULL; s->slab_cap=s->fslab_cap=0; + QT *q[3]={&s->g,&s->u,&s->d}; + for(int k=0;k<3;k++){ q[k]->qf=NULL; q[k]->q8=NULL; q[k]->q4=NULL; q[k]->s=NULL; } + m->resident_bytes-=bytes; if(m->resident_bytes<0) m->resident_bytes=0; +} +static void expert_host_ensure(Model *m, int layer, ESlot *s){ + if(!s->slab) expert_load(m,layer,s->eid,s,1); +} +#endif + /* prefetch asincrono dei pesi di un expert (e delle sue scale .qs): avvia il readahead * cosi' le letture sincrone successive trovano la page-cache calda. */ static void expert_prefetch(Model *m, int layer, int eid){ @@ -1324,7 +1410,10 @@ static int cmp_fdesc(const void *a,const void *b){ float x=*(const float*)a, y=*(const float*)b; return xy?-1:0; } /* attenzione MLA con KV-cache compressa, su token nuovi x[S,hidden], pos_base = pos del primo */ -static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_base, float *out){ +/* kvs/pos describe a ragged decode batch: each row may belong to a different + * sequence. NULL keeps the original contiguous, currently-bound KV path. */ +static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int pos_base, + KVState *const *kvs, const int *positions, float *out){ Cfg *c=&m->c; int H=c->n_heads, D=c->hidden, qh=c->qk_head, vh=c->v_head; int kvb_dim=H*(c->qk_nope+vh), Tk=pos_base+S; double ta0=now_s(); @@ -1361,14 +1450,16 @@ static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_ba /* 1) per ogni token nuovo: query roped + latente normato e k_rot roped -> in cache. * QR tiene il residuo q_a per TUTTE le posizioni: serve anche all'indexer DSA. */ for(int s=0;skv; + const float *xs=x+(int64_t)s*D; int pos=positions?positions[s]:pos_base+s; float *qresid=QR+(int64_t)s*c->q_lora; matmul_qt(qresid, xs, &l->q_a, 1); rmsnorm(qresid, qresid, l->q_a_ln, c->q_lora, c->eps); float *qfull=Q+(int64_t)s*H*qh; matmul_qt(qfull, qresid, &l->q_b, 1); for(int h=0;hqk_nope, pos, c); matmul_qt(comp, xs, &l->kv_a, 1); - float *Ldst=m->Lc[layer]+(int64_t)pos*c->kv_lora, *Rdst=m->Rc[layer]+(int64_t)pos*c->qk_rope; + float *Ldst=coli_kv_row(ks->Lc[layer],pos,c->kv_lora); + float *Rdst=coli_kv_row(ks->Rc[layer],pos,c->qk_rope); memcpy(Ldst, comp, c->kv_lora*sizeof(float)); rmsnorm(Ldst, Ldst, l->kv_a_ln, c->kv_lora, c->eps); /* latente normato */ memcpy(Rdst, comp+c->kv_lora, c->qk_rope*sizeof(float)); @@ -1379,12 +1470,13 @@ static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_ba * dai layer SHARED successivi). Selezione attiva solo con contesto > index_topk * (o DSA_FORCE=1 per il test: selezionare TUTTO deve dare l'output denso esatto). */ const int *dsel=NULL, *dnsel=NULL; int dtopk=0; - if(m->has_dsa && layern_layers && m->kv_start[layer]==0){ + if(m->has_dsa && layern_layers && ((!kvs && m->kv_start[layer]==0) || kvs)){ int nh=c->index_nh, hd=c->index_hd; dtopk=c->index_topk; if(c->idx_type[layer]){ for(int s=0;sIc[layer]+(int64_t)pos*hd; + KVState *ks=kvs?kvs[s]:m->kv; + const float *xs=x+(int64_t)s*D; int pos=positions?positions[s]:pos_base+s; + float *kd=coli_kv_row(ks->Ic[layer],pos,hd); matmul_qt(kd, xs, &m->ix_wk[layer], 1); layernorm(kd, m->ix_knw[layer], m->ix_knb[layer], hd, 1e-6f); rope_interleave(kd, pos, c); /* primi qk_rope dim, interleaved */ @@ -1397,18 +1489,20 @@ static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_ba } #pragma omp parallel for schedule(dynamic,1) for(int s=0;skv; + int pos=positions?positions[s]:pos_base+s, nk=pos+1; + if(ks->kv_start[layer]!=0){ m->dsa_nsel[s]=0; continue; } if(nk<=dtopk && !g_dsa_force){ m->dsa_nsel[s]=0; continue; } int keep = nkq_lora, &m->ix_wq[layer], 1); for(int h=0;hix_wp[layer], 1); float wsc=1.f/sqrtf((float)nh), rs=1.f/sqrtf((float)hd); float *isc=falloc(nk); for(int t=0;tIc[layer]+(int64_t)t*hd; + const float *kt=coli_kv_row(ks->Ic[layer],t,hd); float a=0; for(int h=0;hthr) dst[nd++]=t; for(int t=0;tdsa_nsel[s]=nd; - free(qi); free(isc); free(tmp); + free(qi); free(w32); free(isc); free(tmp); } } if(m->dsa_nsel){ dsel=m->dsa_sel; dnsel=m->dsa_nsel; } @@ -1433,31 +1527,48 @@ static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_ba * k/v per ogni token del contesto. Per linearita': * q·k_nope_t = (W_K^hT q_nope)·L_t ctx^h = W_V^h (Σ_t a_t L_t) * costo per step ~O(T·kv_lora) invece di O(T·H·(nope+vh)) del matmul kvb_all. */ - int absorb = g_absorb==1 || (g_absorb<0 && S<=4); + int absorb = kvs || g_absorb==1 || (g_absorb<0 && S<=4); if(absorb && c->kv_lora<=512){ + m->t_aproj+=now_s()-ta0; double tac=now_s(); int kvl=c->kv_lora, r0v=c->qk_nope; /* offset righe V dentro il blocco di testa */ - /* punteggi per-thread sul HEAP: un sc[8192] fisso sullo stack va in overflow quando - * il layer attende su tutto il contesto (nessuna selezione DSA: snapshot senza - * indexer, o layer MTP) e nt supera 8192 — scrittura oltre lo stack del worker - * OMP => segfault (e poco sotto il limite: corruzione SILENZIOSA dello stack). */ - int64_t sc_cap = Tk - m->kv_start[layer]; /* nt massimo (kv_start=-1 del MTP: +1, ok) */ + /* punteggi per-thread sul HEAP (vedi dev): cap Tk+1 copre anche il kv_start + * per-slot del percorso kvs (MTP: kv_start=-1 -> nt=Tk+1). */ + int64_t sc_cap = (int64_t)Tk+1; float *sc_all = falloc((int64_t)omp_get_max_threads()*sc_cap); + int cuda_core=0; +#ifdef COLI_CUDA + if(S<=4&&g_cuda_enabled&&getenv("COLI_CUDA_ATTN")&&atoi(getenv("COLI_CUDA_ATTN"))&& + l->kv_b.cuda_eligible&&qt_cuda_upload(&l->kv_b)){ + cuda_core=1; + for(int s=0;skv;int pos=positions?positions[s]:pos_base+s; + int st0=ks->kv_start[layer],nt=pos+1-st0; + if(dnsel&&dnsel[s]>0){cuda_core=0;break;} + cuda_core=coli_cuda_attention_absorb(l->kv_b.cuda,ctx+(int64_t)s*H*vh, + Q+(int64_t)s*H*qh,coli_kv_row(ks->Lc[layer],st0,kvl), + coli_kv_row(ks->Rc[layer],st0,c->qk_rope),H,c->qk_nope,c->qk_rope, + vh,kvl,nt,c->attn_scale); + } + } +#endif + if(!cuda_core){ #pragma omp parallel for collapse(2) schedule(static) for(int s=0;skv; + int pos=positions?positions[s]:pos_base+s; const float *qp=Q+(int64_t)s*H*qh+(int64_t)h*qh; const float *qr=qp+c->qk_nope; int rbase=h*(c->qk_nope+vh); float qabs[512]; memset(qabs,0,kvl*sizeof(float)); for(int d=0;dqk_nope;d++) qt_addrow(&l->kv_b, rbase+d, qp[d], qabs); float *sc = sc_all + (int64_t)omp_get_thread_num()*sc_cap; - int st0=m->kv_start[layer]; + int st0=ks->kv_start[layer]; int ns = (dnsel && dnsel[s]>0) ? dnsel[s] : 0; /* DSA: lista top-k o range pieno */ const int *tlist = ns ? dsel+(int64_t)s*dtopk : NULL; int nt = ns ? ns : pos+1-st0; for(int jj=0;jjLc[layer]+(int64_t)t*kvl; - const float *kr=m->Rc[layer]+(int64_t)t*c->qk_rope; + const float *Lt=coli_kv_row(ks->Lc[layer],t,kvl); + const float *kr=coli_kv_row(ks->Rc[layer],t,c->qk_rope); float a=0; for(int i=0;iqk_rope;d++) a+=qr[d]*kr[d]; sc[jj]=a*c->attn_scale; @@ -1465,17 +1576,19 @@ static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_ba softmax(sc,nt); float clat[512]; memset(clat,0,kvl*sizeof(float)); for(int jj=0;jjLc[layer]+(int64_t)t*kvl; + const float *Lt=coli_kv_row(ks->Lc[layer],t,kvl); float a=sc[jj]; for(int i=0;ikv_b, rbase+r0v, vh, clat, ctx+((int64_t)s*H+h)*vh); } - matmul_qt(out, ctx, &l->o, S); + } + m->t_acore+=now_s()-tac; double tao=now_s(); + matmul_qt(out, ctx, &l->o, S); m->t_aout+=now_s()-tao; free(ctx); free(Q); free(QR); free(comp); free(sc_all); m->t_attn += now_s()-ta0; return; } /* 2) ricostruzione di k_nope+value per TUTTI i token 0..Tk-1 (un solo matmul su kv_b) */ - double tk0=now_s(); + m->t_aproj+=now_s()-ta0; double tk0=now_s(); int stL=m->kv_start[layer]; float *kvb_all=falloc((int64_t)Tk*kvb_dim); matmul_qt(kvb_all+(int64_t)stL*kvb_dim, m->Lc[layer]+(int64_t)stL*c->kv_lora, &l->kv_b, Tk-stL); @@ -1484,6 +1597,7 @@ static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_ba * (punteggi sul heap, per-thread: vedi il commento nel ramo absorb) */ int64_t sc_cap = Tk - stL; float *sc_all = falloc((int64_t)omp_get_max_threads()*sc_cap); + double tac=now_s(); #pragma omp parallel for collapse(2) schedule(static) for(int s=0;sqk_nope+vh)+c->qk_nope; float a=sc[jj]; for(int d=0;do, S); + m->t_acore+=now_s()-tac; double tao=now_s(); + matmul_qt(out, ctx, &l->o, S); m->t_aout+=now_s()-tao; free(ctx); free(Q); free(QR); free(comp); free(kvb_all); free(sc_all); m->t_attn += now_s()-ta0; } +static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_base, float *out){ + attention_rows(m,l,layer,x,S,pos_base,NULL,NULL,out); +} + /* MoE GLM su x[S,hidden] -> out (router sigmoid/noaux_tc, n_group=1, + shared expert). * BATCH-UNION: per S>1 (prefill, verifica MTP) ogni expert UNICO del batch viene caricato * una volta sola e moltiplicato per tutte le posizioni che lo usano (pesi letti 1 volta); @@ -1574,6 +1693,7 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){ for(int kk=0;kkeusage[layer][idx[kk]]++; if(m->eheat[layer][idx[kk]]eheat[layer][idx[kk]]++; + m->elast[layer][idx[kk]]=++m->eaccess_clock; } if(c->norm_topk){ float sm=0; for(int kk=0;kkrouted_scale; @@ -1603,6 +1723,13 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){ /* ---- FASE C/D: risolvi (pin/cache/disco) e calcola, a blocchi di 64 unici ---- */ float *xg=falloc((int64_t)S*D), *gg=falloc((int64_t)S*I), *uu=falloc((int64_t)S*I), *hh=falloc((int64_t)S*D); int *rows=malloc(S*sizeof(int)); float *rw=malloc(S*sizeof(float)); +#ifdef COLI_CUDA + int group_enabled=S<=64; + float *group_x=group_enabled?falloc((int64_t)S*K*D):NULL; + float *group_y=group_enabled?falloc((int64_t)S*K*D):NULL; + int *group_row=group_enabled?malloc((size_t)64*S*sizeof(int)):NULL; + float *group_weight=group_enabled?malloc((size_t)64*S*sizeof(float)):NULL; +#endif int shared_on_gpu=0; (void)shared_on_gpu; /* set by the Metal path when Phase E was fused */ for(int base=0;baseg.cuda_eligible) m->gpu_expert_calls++; + if(group_enabled && g_cuda_enabled && e->g.cuda_eligible && e->u.cuda_eligible && e->d.cuda_eligible && + !omp_in_parallel()){ + group_e[ngroup]=e; group_n[ngroup]=nr; + for(int r=0;rg, nr); - matmul_qt(uu, xg, &e->u, nr); +#ifdef COLI_CUDA + if(!group_enabled && g_cuda_enabled && e->g.cuda_eligible && e->u.cuda_eligible && + e->d.cuda_eligible && !omp_in_parallel() && + coli_cuda_expert_mlp(e->g.cuda,e->u.cuda,e->d.cuda,hh,xg,nr)){ + for(int r=0;rt_emm+=now_s()-t0; continue; + } + if(!e->slab) expert_host_ensure(m,layer,e); +#endif + expert_gate_up(gg,uu,xg,&e->g,&e->u,nr); for(int64_t z=0;z<(int64_t)nr*I;z++) gg[z]=siluf(gg[z])*uu[z]; matmul_qt(hh, gg, &e->d, nr); for(int r=0;rt_emm += now_s()-t0; } +#ifdef COLI_CUDA + ColiCudaTensor *dev_g[COLI_CUDA_MAX_DEVICES][64],*dev_u[COLI_CUDA_MAX_DEVICES][64]; + ColiCudaTensor *dev_d[COLI_CUDA_MAX_DEVICES][64]; + int dev_rows[COLI_CUDA_MAX_DEVICES][64],dev_which[COLI_CUDA_MAX_DEVICES][64]; + int dev_nc[COLI_CUDA_MAX_DEVICES]={0},dev_total[COLI_CUDA_MAX_DEVICES]={0}; + int dev_off[COLI_CUDA_MAX_DEVICES]={0},dev_ok[COLI_CUDA_MAX_DEVICES]={0}; + for(int di=0;dig.cuda_device==g_cuda_devices[di]) dev_total[di]+=group_n[q]; + for(int di=1;dig.cuda_device==device){ + int nc=dev_nc[di]++; ESlot *e=group_e[q]; + dev_g[di][nc]=e->g.cuda; dev_u[di][nc]=e->u.cuda; dev_d[di][nc]=e->d.cuda; + dev_rows[di][nc]=group_n[q]; dev_which[di][nc]=q; + for(int r=0;r1) schedule(static) + for(int di=0;dig.cuda,e->u.cuda,e->d.cuda,hh,xg,nr)){ + expert_host_ensure(m,layer,e); + expert_gate_up(gg,uu,xg,&e->g,&e->u,nr); + for(int64_t z=0;z<(int64_t)nr*I;z++) gg[z]=siluf(gg[z])*uu[z]; + matmul_qt(hh,gg,&e->d,nr); + } + } + float *src=dev_ok[di]?group_y+(int64_t)off*D:hh; + for(int r=0;rt_emm+=now_s()-tg; +#endif /* No drain barrier: the per-expert pipe_wait(qof[j]) above (issued for every * dispatched miss slot, before the nr==0 skip) already waited on all ws[] loads * for this block, so they are complete before the LRU swap — and the gen-tagged @@ -1783,6 +1974,9 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){ } free(logit); free(choice); free(idxs); free(ws); free(keff); free(uniq); free(xg); free(gg); free(uu); free(hh); free(rows); free(rw); free(sg); free(su); +#ifdef COLI_CUDA + free(group_x); free(group_y); free(group_row); free(group_weight); +#endif } static void dense_mlp(Layer *l, float *x, int S, int D, int I, float *out){ @@ -1914,7 +2108,8 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S){ } /* forward di UN layer (usato dai 78 principali e dal layer MTP) */ -static void layer_forward(Model *m, Layer *l, int li, float *x, int S, int pos_base, float *nrm, float *tmp){ +static void layer_forward_rows(Model *m, Layer *l, int li, float *x, int S, int pos_base, + KVState *const *kvs, const int *positions, float *nrm, float *tmp){ Cfg *c=&m->c; int D=c->hidden; if(g_spec && g_prefetch && l->sparse && m->enr[li]>0) for(int z=0;zenr[li];z++) expert_prefetch(m,li,m->eroute[li][z]); @@ -1974,7 +2169,7 @@ static void layer_forward(Model *m, Layer *l, int li, float *x, int S, int pos_b } #endif for(int s=0;sin_ln, D, c->eps); - attention(m,l,li,nrm,S,pos_base,tmp); + attention_rows(m,l,li,nrm,S,pos_base,kvs,positions,tmp); for(int64_t j=0;j<(int64_t)S*D;j++) x[j]+=tmp[j]; if(g_pilot && S<=8 && li+1n_layers && m->L[li+1].sparse) pilot_prefetch(m,li+1,x,S); if(g_looka && S==1 && li+1n_layers && m->L[li+1].sparse) la_predict(m,li+1,x,1); @@ -1982,7 +2177,11 @@ static void layer_forward(Model *m, Layer *l, int li, float *x, int S, int pos_b if(l->sparse) moe(m,l,li,nrm,S,tmp); else dense_mlp(l,nrm,S,D,c->dense_inter,tmp); for(int64_t j=0;j<(int64_t)S*D;j++) x[j]+=tmp[j]; } -static void layers_forward(Model *m, float *x, int S, int pos_base){ +static void layer_forward(Model *m, Layer *l, int li, float *x, int S, int pos_base, float *nrm, float *tmp){ + layer_forward_rows(m,l,li,x,S,pos_base,NULL,NULL,nrm,tmp); +} +static void layers_forward_rows(Model *m, float *x, int S, int pos_base, + KVState *const *kvs, const int *positions){ Cfg *c=&m->c; int D=c->hidden; if(g_pilot_real){ /* nuovo forward: il possesso-layer riparte da -1 (i layer si rifanno da 0) */ pthread_mutex_lock(&g_pilot_mx); @@ -1995,10 +2194,13 @@ static void layers_forward(Model *m, float *x, int S, int pos_base){ * puo' arrivare dopo MINUTI di streaming — al buio sembra un blocco. */ if(S>=8 && (i%4==0 || i==c->n_layers-1)) fprintf(stderr,"[prefill] layer %d/%d · %d token\n", i+1, c->n_layers, S); - layer_forward(m,&m->L[i],i,x,S,pos_base,nrm,tmp); + layer_forward_rows(m,&m->L[i],i,x,S,pos_base,kvs,positions,nrm,tmp); } free(nrm); free(tmp); } +static void layers_forward(Model *m, float *x, int S, int pos_base){ + layers_forward_rows(m,x,S,pos_base,NULL,NULL); +} static void kv_alloc(Model *m, int max_t){ Cfg *c=&m->c; @@ -2068,6 +2270,42 @@ static float *step_all(Model *m, const int *ids, int S, int pos_base){ free(x); free(row); return lo; } +/* One decode token from each independent sequence, evaluated as a single MoE + * batch. Prefill and speculative batches retain their contiguous-KV path. */ +static float *step_decode_batch(Model *m, const DecodeRow *rows, int S){ + Cfg *c=&m->c; int D=c->hidden; + /* Ragged KV currently uses MLA absorption; the stack kernel is sized to 512. */ + if(!rows || S<1 || S>64 || c->kv_lora>512) return NULL; + KVState *kvs[64]; int positions[64]; + float *x=falloc((int64_t)S*D); + for(int s=0;sLc || !rows[s].kv->Rc || !rows[s].kv->kv_start || + rows[s].token<0 || rows[s].token>=c->vocab || + rows[s].pos<0 || rows[s].pos>=rows[s].kv->max_t){ + free(x); return NULL; + } + for(int l=0;ln_layers;l++){ + if(!rows[s].kv->Lc[l] || !rows[s].kv->Rc[l] || + rows[s].kv->kv_start[l]<0 || rows[s].kv->kv_start[l]>rows[s].pos || + (m->has_dsa && c->idx_type[l] && + (!rows[s].kv->Ic || !rows[s].kv->Ic[l]))){ free(x); return NULL; } + } + for(int p=0;pfinal_norm,D,c->eps); + double th0=now_s(); + float *logit=falloc((int64_t)S*c->vocab); + matmul_qt(logit,norm,&m->lm_head,S); + m->t_head+=now_s()-th0; + free(x); free(norm); + return logit; +} + /* METODO E — prompt-lookup: cerca l'occorrenza piu' recente dell'ultimo bigramma nel * contesto e propone i token che la seguirono. Zero pesi extra, zero costo: e' solo * un'ipotesi che il modello verifichera'. */ @@ -2420,10 +2658,12 @@ static void generate(Model *m, const int *prompt, int np, int n_new, int *out){ } static void profile_print(Model *m, double elapsed){ - double accounted=m->t_edisk+m->t_emm+m->t_attn+m->t_head; - printf("PROFILE: expert-disk %.3fs | expert-matmul %.3fs | attention %.3fs " + double accounted=m->t_ewait+m->t_emm+m->t_attn+m->t_head; + printf("PROFILE: expert-disk %.3fs service / %.3fs wait | expert-matmul %.3fs | attention %.3fs " "(including kvb %.3fs) | lm_head %.3fs | other %.3fs\n", - m->t_edisk,m->t_emm,m->t_attn,m->t_kvb,m->t_head,elapsed-accounted); + m->t_edisk,m->t_ewait,m->t_emm,m->t_attn,m->t_kvb,m->t_head,elapsed-accounted); + printf("ATTENTION: projection/RoPE %.3fs | score-softmax-value %.3fs | output projection %.3fs\n", + m->t_aproj,m->t_acore,m->t_aout); #ifdef COLI_METAL if(g_metal_enabled){ uint64_t ok=0,fb=0,ex=0; double su=0,gp=0,sc=0; coli_metal_moe_counts(&ok,&fb,&ex); coli_metal_moe_times(&su,&gp,&sc); @@ -2443,7 +2683,8 @@ static void run_replay(Model *m, const int *full, int nfull, int np){ kv_alloc(m,nfull+2); float *logit=step(m,full,np-1,0); free(logit); m->hits=m->miss=m->ereq=m->gpu_expert_calls=0; - m->t_edisk=m->t_emm=m->t_attn=m->t_kvb=m->t_head=0; + m->t_edisk=m->t_ewait=m->t_emm=m->t_attn=m->t_kvb=m->t_head=0; + m->t_aproj=m->t_acore=m->t_aout=0; double t0=now_s(); int steps=0; for(int i=np-1;ipin[l]; int ids[4096], zp, eu; long g; int np=m->npin[l]; if(np>4096) np=4096; for(int z=0;zeheat[l],c->n_experts,ids,np,&zp,&eu,&g)) continue; + if(!tier_pick_lfru(m->eheat[l],m->elast[l],m->eaccess_clock, + c->n_experts,ids,np,&zp,&eu,&g)) continue; if(nbout[w].gain){ out[w].gain=g; out[w].l=l; out[w].slot=zp; out[w].eid=eu; } } @@ -2571,6 +2813,7 @@ static void repin_pass(Model *m){ +(int64_t)coli_cuda_tensor_bytes(s->u.cuda) +(int64_t)coli_cuda_tensor_bytes(s->d.cuda); m->gpu_expert_bytes+=now_gpu-old_gpu; tier="VRAM"; + if(g_cuda_release_host) expert_host_release(m,s); } else { qt_cuda_reset(&s->g); qt_cuda_reset(&s->u); qt_cuda_reset(&s->d); s->g.cuda_eligible=s->u.cuda_eligible=s->d.cuda_eligible=0; @@ -2686,6 +2929,141 @@ static void serve_ctx_free(Model *m, ServeCtx *s){ free(k->Lc); free(k->Rc); free(k->Ic); free(k->kv_start); free(s->hist); } +typedef struct { + int active, pending, emitted, maximum, prompt_tokens, length_limited; + unsigned long long id; + float temp, top_p; + double started; + uint64_t hits0, miss0; +} ServeReq; + +static void mux_data(Tok *T, unsigned long long id, int token){ + char out[256]; int n=tok_decode(T,&token,1,out,sizeof(out)); + printf("DATA %llu %d\n",id,n); if(n>0) fwrite(out,1,(size_t)n,stdout); putchar('\n'); + fflush(stdout); +} + +static void mux_done(Model *m, ServeCtx *sc, ServeReq *r){ + double dt=now_s()-r->started; if(dt<1e-6) dt=1e-6; + double dh=(double)(m->hits-r->hits0), dm=(double)(m->miss-r->miss0); + printf("DONE %llu STAT %d %.2f %.1f %.2f %d %d\n",r->id,r->emitted, + r->emitted/dt,(dh+dm)>0?100.0*dh/(dh+dm):0.0,rss_gb(), + r->prompt_tokens,r->length_limited); + fflush(stdout); kv_bind(m,&sc->kv); kv_disk_append(m,sc->hist,sc->len); + r->active=0; +} + +/* Read and prefill one request. Returns -1 on EOF, 0 for a rejected frame and + * 1 for an accepted request. Prefill deliberately remains serial: continuous + * batching starts at decode, where every active slot contributes one row. */ +static int mux_submit(Model *m, Tok *T, ServeCtx *ctx, ServeReq *req, int nctx, + int maxctx, int eos){ + char *line=NULL; size_t cap=0; ssize_t nr=getline(&line,&cap,stdin); + if(nr<0){ free(line); return -1; } + if(nr && line[nr-1]=='\n') line[--nr]=0; + if(!strncmp(line,"CANCEL ",7)){ + unsigned long long id=0; char tail; + if(sscanf(line+7,"%llu %c",&id,&tail)!=1 || id==0){ + printf("ERROR 0 BAD_REQUEST\n"); fflush(stdout); free(line); return 0; + } + for(int i=0;i=nctx || memchr(raw,0,(size_t)sub.bytes)){ + printf("ERROR %llu BAD_REQUEST\n",sub.id); fflush(stdout); free(raw); free(line); return 0; + } + if(req[sub.slot].active){ + printf("ERROR %llu SLOT_BUSY\n",sub.id); fflush(stdout); free(raw); free(line); return 0; + } + for(int i=0;ikv); + int *tmp=malloc(maxctx*sizeof(int)); + int nt=tok_encode(T,raw,(int)sub.bytes,tmp,maxctx-2); + free(raw); free(line); + if(nt<1){ free(tmp); printf("ERROR %llu EMPTY_PROMPT\n",sub.id); fflush(stdout); return 0; } + int prefix=0; while(prefixlen && prefixhist[prefix]==tmp[prefix]) prefix++; + if(prefixlen){ sc->len=prefix; if(m->has_mtp) m->kv_start[m->c.n_layers]=-1; + kv_disk_truncate(m,sc->len); } + int add=nt-sc->len; + if(add>0) memcpy(sc->hist+sc->len,tmp+sc->len,(size_t)add*sizeof(int)); + free(tmp); + float *logit = add>0 ? step(m,sc->hist+sc->len,add,sc->len) + : step(m,sc->hist+sc->len-1,1,sc->len-1); + sc->len+=add; sc->first=0; + ServeReq *r=&req[sub.slot]; memset(r,0,sizeof(*r)); + r->id=sub.id; r->maximum=sub.max_tokens; r->temp=sub.temperature; r->top_p=sub.top_p; + r->prompt_tokens=nt; r->started=now_s(); r->hits0=m->hits; r->miss0=m->miss; + int room=maxctx-sc->len-1; if(r->maximum>room){r->maximum=room; r->length_limited=1;} + g_temp=r->temp; g_nuc=r->top_p; + int next=pick_tok(logit,m->c.vocab,-1); free(logit); + if(r->maximum<=0 || next==eos || is_stop(next)){ mux_done(m,sc,r); return 1; } + r->pending=next; r->emitted=1; r->active=1; sc->hist[sc->len]=next; m->n_emit++; + mux_data(T,r->id,next); + if(r->emitted>=r->maximum) mux_done(m,sc,r); + return 1; +} + +static void run_serve_mux(Model *m, const char *snap){ + char tkp[2048]; snprintf(tkp,sizeof(tkp),"%s/tokenizer.json",snap); + Tok T; tok_load(&T,tkp); int eos=tok_id_of(&T,"<|endoftext|>"); stops_arm(&m->c,eos); + g_draft=0; /* one scheduler owns every forward; MTP/speculation is not ragged-safe */ + int maxctx=getenv("CTX")?atoi(getenv("CTX")):4096; + int nctx=getenv("KV_SLOTS")?atoi(getenv("KV_SLOTS")):1; + if(nctx<1||nctx>16){fprintf(stderr,"KV_SLOTS deve essere tra 1 e 16\n");exit(2);} + g_kvsave=getenv("KVSAVE")?atoi(getenv("KVSAVE")):1; + KVState *initial=m->kv; free(initial->kv_start); free(initial); + ServeCtx *ctx=calloc(nctx,sizeof(*ctx)); ServeReq *req=calloc(nctx,sizeof(*req)); + for(int i=0;i0 && FD_ISSET(STDIN_FILENO,&rfds)) if(mux_submit(m,&T,ctx,req,nctx,maxctx,eos)<0) eof=1; + active=0; for(int i=0;in_fw++; + for(int s=0;slen++; g_temp=r->temp; g_nuc=r->top_p; + int next=pick_tok(lo+(int64_t)s*m->c.vocab,m->c.vocab,-1); + if(next==eos || is_stop(next)){mux_done(m,sc,r);continue;} + r->pending=next; sc->hist[sc->len]=next; r->emitted++; m->n_emit++; + mux_data(&T,r->id,next); + if(r->emitted>=r->maximum) mux_done(m,sc,r); + } + free(lo); + } + usage_save(m); + for(int i=0;ikv=NULL; m->Lc=m->Rc=m->Ic=NULL; m->kv_start=NULL; m->max_t=0; +} + static void run_serve(Model *m, const char *snap){ char tkp[2048]; snprintf(tkp,sizeof(tkp),"%s/tokenizer.json",snap); Tok T; tok_load(&T,tkp); @@ -2811,6 +3189,7 @@ static void run_serve(Model *m, const char *snap){ free(raw); g_temp=base_temp; g_nuc=base_nuc; usage_save(m); /* la cache che impara: storia aggiornata a ogni turno */ kv_disk_append(m,hist,len); /* KV su disco: il prossimo avvio riparte da qui */ + repin_pass(m); /* safe request boundary: adapt session-local hot tier */ } free(line); free(buf); usage_save(m); @@ -2922,60 +3301,72 @@ static void pin_wire(Model *m){ "(no compression) in %.0fs\n", wired/1e9, now_s()-t0); } +typedef struct { int l,e; uint32_t c; } PinRec; +static int pin_rec_cmp(const void *a,const void *b){ + const PinRec *x=a,*y=b; return x->cc?1:x->c>y->c?-1:0; +} static void pin_load(Model *m, const char *statspath, double gb){ FILE *f=fopen(statspath,"r"); if(!f){ perror(statspath); return; } - typedef struct { int l,e; uint32_t c; } Rec; Cfg *c=&m->c; int cap=(c->n_layers+1)*c->n_experts; - Rec *r=malloc((size_t)cap*sizeof(Rec)); int n=0; + PinRec *r=malloc((size_t)cap*sizeof(PinRec)); int n=0; + unsigned char *seen=calloc((size_t)(c->n_layers+1)*c->n_experts,1); int l,e; uint32_t cnt; while(n=0 && e>=0 && en_experts && ((ln_layers && m->L[l].sparse) || (l==c->n_layers && m->has_mtp)); - if(ok) r[n++]=(Rec){l,e,cnt}; + int64_t key=(int64_t)l*c->n_experts+e; + if(ok&&!seen[key]){ r[n++]=(PinRec){l,e,cnt}; seen[key]=1; } } fclose(f); - for(int a=0;ar[best].c) best=b; - Rec t=r[a]; r[a]=r[best]; r[best]=t; - if(a>4095) break; /* bastano i top ~4k */ + int fill=getenv("PIN_FILL")?atoi(getenv("PIN_FILL")):0; +#ifdef COLI_CUDA + if(!getenv("PIN_FILL")&&g_cuda_release_host) fill=1; +#endif + if(fill) for(int li=0;li<=c->n_layers;li++){ + int sparse=(lin_layers&&m->L[li].sparse)||(li==c->n_layers&&m->has_mtp); + if(sparse) for(int ei=0;ein_experts;ei++) if(!seen[(int64_t)li*c->n_experts+ei]) + r[n++]=(PinRec){li,ei,0}; } + free(seen); + qsort(r,(size_t)n,sizeof(*r),pin_rec_cmp); int64_t eb=expert_bytes_probe(m,m->ebits); - int npin=(int)(gb*1e9/eb); if(npin>n) npin=n; if(npin>4096) npin=4096; + int npin=(int)(gb*1e9/eb); if(npin>n) npin=n; if(npin<1){ free(r); return; } int *cnt_l=calloc(c->n_layers+1,sizeof(int)); /* +1: riga MTP */ for(int a=0;an_layers;i++) if(cnt_l[i]) m->pin[i]=calloc(cnt_l[i],sizeof(ESlot)); + int *slot_of=malloc((size_t)npin*sizeof(int)), *next=calloc(c->n_layers+1,sizeof(int)); + for(int a=0;an_layers;i++) m->npin[i]=cnt_l[i]; double t0=now_s(); - #pragma omp parallel for schedule(dynamic,1) - for(int a=0;anpin[li]++; - expert_load(m,li,r[a].e,&m->pin[li][slot],1); +#ifdef COLI_CUDA + double remaining[COLI_CUDA_MAX_DEVICES]={0}, placed_b[COLI_CUDA_MAX_DEVICES]={0}; + int placed_n[COLI_CUDA_MAX_DEVICES]={0}, gpu_prefix=0; + double budget=g_cuda_expert_gb*1e9, safe_total=0; + if(g_cuda_enabled&&g_cuda_expert_gb>0) for(int i=0;iresident_bytes += (int64_t)npin*eb; - fprintf(stderr,"[PIN] hot store: %d experts in RAM (%.1f GB) loaded in %.0fs from %s\n", - npin, npin*eb/1e9, now_s()-t0, statspath); + if(budget>safe_total) budget=safe_total; + if(g_cuda_enabled&&g_cuda_release_host&&budget>0){ gpu_prefix=(int)(budget/eb)+g_cuda_ndev; if(gpu_prefix>npin)gpu_prefix=npin; } +#else + int gpu_prefix=0; +#endif + /* Load the VRAM-ranked prefix first. Once uploaded its host backing is + * released before the disjoint RAM-ranked suffix is allocated. */ + #pragma omp parallel for schedule(dynamic,1) + for(int a=0;a<(gpu_prefix?gpu_prefix:npin);a++) + expert_load(m,r[a].l,r[a].e,&m->pin[r[a].l][slot_of[a]],1); + m->resident_bytes+=(int64_t)(gpu_prefix?gpu_prefix:npin)*eb; #ifdef COLI_CUDA if(g_cuda_enabled && g_cuda_expert_gb>0){ - double remaining[COLI_CUDA_MAX_DEVICES]={0}, placed_b[COLI_CUDA_MAX_DEVICES]={0}; - int placed_n[COLI_CUDA_MAX_DEVICES]={0}; - double budget=g_cuda_expert_gb*1e9, safe_total=0; - for(int i=0;isafe_total) budget=safe_total; - for(int a=0;agpu_expert_bytesgpu_expert_bytesnpin[li];z++) if(m->pin[li][z].eid==r[a].e){ - ESlot *s=&m->pin[li][z]; + { ESlot *s=&m->pin[li][slot_of[a]]; int64_t need=qt_bytes(&s->g)+qt_bytes(&s->u)+qt_bytes(&s->d); if(m->gpu_expert_bytes+need>budget) break; int tried[COLI_CUDA_MAX_DEVICES]={0}, placed=0; @@ -2993,6 +3384,7 @@ static void pin_load(Model *m, const char *statspath, double gb){ +(int64_t)coli_cuda_tensor_bytes(s->d.cuda); m->gpu_expert_count++; m->gpu_expert_bytes+=actual; remaining[best]-=actual; placed_b[best]+=actual; placed_n[best]++; + if(g_cuda_release_host) expert_host_release(m,s); placed=1; } else { qt_cuda_reset(&s->g); qt_cuda_reset(&s->u); qt_cuda_reset(&s->d); @@ -3000,7 +3392,6 @@ static void pin_load(Model *m, const char *statspath, double gb){ remaining[best]=0; /* device rejected its projected capacity */ } } - break; } } fprintf(stderr,"[CUDA] hot expert tier: %d/%d experts, VRAM %.2f GB (total budget %.1f GB)\n", @@ -3009,8 +3400,16 @@ static void pin_load(Model *m, const char *statspath, double gb){ g_cuda_devices[i],placed_n[i],placed_b[i]/1e9); } #endif + if(gpu_prefix>0&&gpu_prefixpin[r[a].l][slot_of[a]],1); + m->resident_bytes+=(int64_t)(npin-gpu_prefix)*eb; + } + fprintf(stderr,"[PIN] placement: %d VRAM + %d RAM expert (%.1f GB warm) in %.0fs da %s\n", + m->gpu_expert_count,npin-m->gpu_expert_count,(npin-m->gpu_expert_count)*eb/1e9,now_s()-t0,statspath); pin_wire(m); /* inchioda in RAM (no compressione) / wire in RAM (no compression) */ - free(r); free(cnt_l); + free(r); free(cnt_l); free(slot_of); free(next); } static double g_mem_avail_boot=0; /* MemAvailable all'avvio, prima di caricare il modello */ @@ -3158,6 +3557,14 @@ int main(int argc, char **argv){ if(g_mmap) fprintf(stderr,"[MMAP] expert = viste zero-copy nei file (page cache = cache)\n"); g_topk = getenv("TOPK")?atoi(getenv("TOPK")):0; g_topp = getenv("TOPP")?atof(getenv("TOPP")):0; + const char *policy=getenv("COLI_POLICY"); if(!policy) policy="quality"; + int experimental=!strcmp(policy,"experimental-fast"); + if(strcmp(policy,"quality")&&strcmp(policy,"balanced")&&!experimental){ + fprintf(stderr,"COLI_POLICY non valida: quality, balanced o experimental-fast\n"); return 2; + } + if(!experimental&&(g_topk>0||g_topp>0)){ + fprintf(stderr,"[policy] --topp/--topk drop low-weight experts (~1.6x fewer reads, small quality cost)\n"); + } g_mlock = getenv("MLOCK")?atoi(getenv("MLOCK")):-1; /* -1 auto (ON macOS), 0 off, 1 force / auto (ON macOS), 0 off, 1 force */ g_spec = getenv("SPEC")?atoi(getenv("SPEC")):1; g_draft = getenv("DRAFT")?atoi(getenv("DRAFT")):-1; /* -1 = auto: 3 se MTP, 0 senza */ @@ -3203,10 +3610,13 @@ int main(int argc, char **argv){ } g_cuda_dense=getenv("CUDA_DENSE")?atoi(getenv("CUDA_DENSE")):0; g_cuda_expert_gb=getenv("CUDA_EXPERT_GB")?atof(getenv("CUDA_EXPERT_GB")):0; + g_cuda_release_host=getenv("CUDA_RELEASE_HOST")?atoi(getenv("CUDA_RELEASE_HOST")):(g_cuda_ndev>1); if((getenv("COLI_GPU")||getenv("COLI_GPUS"))&&!g_cuda_enabled){ fprintf(stderr,"COLI_GPU(S) requires COLI_CUDA=1\n"); return 2; } if(g_cuda_dense&&!g_cuda_enabled){ fprintf(stderr,"CUDA_DENSE requires COLI_CUDA=1\n"); return 2; } if(g_cuda_expert_gb>0 && !g_cuda_enabled){ fprintf(stderr,"CUDA_EXPERT_GB requires COLI_CUDA=1\n"); return 2; } - if(g_cuda_enabled) fprintf(stderr,"[CUDA] mode: routed experts%s\n",g_cuda_dense?" + resident dense tensors":" only (resident dense on CPU)"); + if(g_cuda_enabled) fprintf(stderr,"[CUDA] mode: routed experts%s%s\n", + g_cuda_dense?" + resident dense tensors":" only (resident dense on CPU)", + g_cuda_release_host?"; VRAM experts without host backing":""); #else if((getenv("COLI_CUDA") && atoi(getenv("COLI_CUDA"))) || getenv("COLI_GPU") || getenv("COLI_GPUS") || @@ -3233,7 +3643,13 @@ int main(int argc, char **argv){ printf("== GLM C engine (glm_moe_dsa), cache=%d experts/layer | experts@%d-bit dense@%d-bit | idot: " IDOT_KERNEL " ==\n", cap, ebits, dbits); g_mem_avail_boot = mem_available_gb(); Model m; double t0=now_s(); model_init(&m,snap,cap,ebits,dbits); - if(g_draft<0) g_draft = m.has_mtp ? 3 : 0; + if(g_draft<0){ +#ifdef COLI_CUDA + g_draft = (m.has_mtp&&!g_cuda_enabled) ? 3 : 0; +#else + g_draft = m.has_mtp ? 3 : 0; +#endif + } if(getenv("DSA_TOPK")) m.c.index_topk=atoi(getenv("DSA_TOPK")); /* override per test */ printf("loaded in %.2fs | resident dense: %.2f MB | layers=%d experts=%d | MTP %s (draft=%d)\n", now_s()-t0, m.resident_bytes/(1024.0*1024.0), m.c.n_layers, m.c.n_experts, @@ -3272,7 +3688,11 @@ int main(int argc, char **argv){ if(getenv("SCORE")){ run_score(&m, getenv("SCORE")); if(stats) stats_dump(&m,stats); return 0; } /* modo serve persistente per la CLI 'coli': SERVE=1 */ - if(getenv("SERVE")){ run_serve(&m, snap); if(stats) stats_dump(&m,stats); return 0; } + if(getenv("SERVE")){ + if(getenv("SERVE_BATCH") && atoi(getenv("SERVE_BATCH"))) run_serve_mux(&m,snap); + else run_serve(&m,snap); + if(stats) stats_dump(&m,stats); return 0; + } /* modo testo reale: PROMPT="..." [NGEN=n] -> tokenizza, genera, detokenizza */ if(getenv("PROMPT")){ diff --git a/c/openai_server.py b/c/openai_server.py index 0c26b78..47672d0 100644 --- a/c/openai_server.py +++ b/c/openai_server.py @@ -6,8 +6,10 @@ import codecs import collections import contextlib import json +import math import os import select +import queue import signal import socket import subprocess @@ -54,18 +56,22 @@ def error_object(error): class GenerationScheduler: - """Bounded FIFO admission for the engine's single mutable KV context.""" + """Bounded FIFO admission for the engine's independent KV contexts.""" - def __init__(self, max_queue=8, queue_timeout=300): + def __init__(self, max_queue=8, queue_timeout=300, capacity=1): if max_queue < 0: raise ValueError("max_queue cannot be negative") if queue_timeout <= 0: raise ValueError("queue_timeout must be positive") + if capacity < 1: + raise ValueError("capacity must be positive") self.max_queue = max_queue self.queue_timeout = queue_timeout + self.capacity = capacity + self.free_slots = set(range(capacity)) self.condition = threading.Condition() self.queue = collections.deque() - self.active = False + self.active = 0 self.closed = False self.admitted = 0 self.completed = 0 @@ -74,14 +80,14 @@ class GenerationScheduler: self.cancelled = 0 @contextlib.contextmanager - def admit(self, cancelled=None): + def admit(self, cancelled=None, slot=None): ticket = object() queued_at = time.monotonic() with self.condition: if self.closed: raise APIError(503, "The inference scheduler is shutting down.", None, "scheduler_closed", "server_error") - if (self.active or self.queue) and len(self.queue) >= self.max_queue: + if (self.active >= self.capacity or self.queue) and len(self.queue) >= self.max_queue: self.rejected += 1 raise APIError(429, "The inference queue is full.", None, "queue_full", "rate_limit_error", {"Retry-After": "1"}) @@ -93,7 +99,8 @@ class GenerationScheduler: self.condition.notify_all() raise APIError(503, "The inference scheduler is shutting down.", None, "scheduler_closed", "server_error") - if not self.active and self.queue[0] is ticket: + available = min(self.free_slots) if slot is None and self.free_slots else slot + if self.queue[0] is ticket and available in self.free_slots: break if cancelled and cancelled(): self.queue.remove(ticket) @@ -109,20 +116,23 @@ class GenerationScheduler: "queue_timeout", "rate_limit_error", {"Retry-After": "1"}) self.condition.wait(min(remaining, 0.25)) self.queue.popleft() - self.active = True + self.free_slots.remove(available) + self.active += 1 self.admitted += 1 wait_seconds = time.monotonic() - queued_at try: - yield wait_seconds + yield wait_seconds, available finally: with self.condition: - self.active = False + self.active -= 1 + self.free_slots.add(available) self.completed += 1 self.condition.notify_all() def snapshot(self): with self.condition: return {"active": self.active, "queued": len(self.queue), + "capacity": self.capacity, "max_queue": self.max_queue, "queue_timeout_seconds": self.queue_timeout, "admitted": self.admitted, "completed": self.completed, "rejected": self.rejected, "timed_out": self.timed_out, @@ -325,9 +335,11 @@ def generation_options(body, limit): top_p = 0.9 if top_p is None else top_p if isinstance(maximum, bool) or not isinstance(maximum, int) or not 1 <= maximum <= limit: raise APIError(400, f"`{maximum_param}` must be an integer between 1 and {limit}.", maximum_param) - if isinstance(temperature, bool) or not isinstance(temperature, (int, float)) or not 0 <= temperature <= 2: + if (isinstance(temperature, bool) or not isinstance(temperature, (int, float)) or + not math.isfinite(temperature) or not 0 <= temperature <= 2): raise APIError(400, "`temperature` must be between 0 and 2.", "temperature") - if isinstance(top_p, bool) or not isinstance(top_p, (int, float)) or not 0 < top_p <= 1: + if (isinstance(top_p, bool) or not isinstance(top_p, (int, float)) or + not math.isfinite(top_p) or not 0 < top_p <= 1): raise APIError(400, "`top_p` must be greater than 0 and at most 1.", "top_p") return maximum, float(temperature), float(top_p) @@ -363,17 +375,100 @@ def read_engine_turn(stream, sentinel, on_bytes): class Engine: def __init__(self, executable, model, cap=8, max_tokens=1024, env=None, kv_slots=1): - child_env = dict(env or os.environ, SNAP=str(model), SERVE="1", NGEN=str(max_tokens), - KV_SLOTS=str(kv_slots)) + child_env = dict(env or os.environ, SNAP=str(model), SERVE="1", SERVE_BATCH="1", + NGEN=str(max_tokens), KV_SLOTS=str(kv_slots)) self.process = subprocess.Popen( [str(executable), str(cap)], env=child_env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, bufsize=0, ) - self.lock = threading.Lock() + self.write_lock = threading.Lock() + self.pending_lock = threading.Lock() + self.pending = {} + self.next_request_id = 1 + self.closed = False + self.dispatcher_error = None self.kv_slots = kv_slots read_engine_turn(self.process.stdout, READY, lambda _: None) + self.dispatcher = threading.Thread(target=self._dispatch_stdout, + name="colibri-stdout", daemon=True) + self.dispatcher.start() - def generate(self, prompt, max_tokens, temperature, top_p, on_text, cache_slot=0): + @staticmethod + def _stats(fields): + if len(fields) < 5 or fields[0] != "STAT": + raise RuntimeError(f"invalid engine status: {' '.join(fields)}") + return { + "completion_tokens": int(fields[1]), + "tokens_per_second": float(fields[2]), + "cache_hit_percent": float(fields[3]), + "rss_gb": float(fields[4]), + "prompt_tokens": int(fields[5]) if len(fields) > 5 else 0, + "length_limited": bool(int(fields[6])) if len(fields) > 6 else False, + } + + def _fail_pending(self, error): + with self.pending_lock: + requests = list(self.pending.values()) + self.pending.clear() + for events in requests: + events.put(("error", error)) + + def _read_exact(self, size): + chunks = [] + remaining = size + while remaining: + chunk = self.process.stdout.read(remaining) + if chunk == b"": + raise RuntimeError("truncated engine DATA payload") + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) + + def _dispatch_stdout(self): + try: + while True: + line = self.process.stdout.readline() + if line == b"": + raise RuntimeError("colibri engine exited unexpectedly") + fields = line.decode("utf-8", "replace").strip().split() + if not fields: + continue + kind = fields[0] + if kind == "DATA" and len(fields) == 3: + request_id = fields[1] + size = int(fields[2]) + if not 0 <= size <= 65536: + raise RuntimeError("invalid engine DATA size") + data = self._read_exact(size) + if self._read_exact(1) != b"\n": + raise RuntimeError("invalid engine DATA terminator") + with self.pending_lock: + events = self.pending.get(request_id) + if events is not None: + events.put(("data", data)) + elif kind == "DONE" and len(fields) >= 7: + request_id = fields[1] + stats = self._stats(fields[2:]) + with self.pending_lock: + events = self.pending.pop(request_id, None) + if events is not None: + events.put(("done", stats)) + elif kind == "ERROR" and len(fields) >= 2: + request_id = fields[1] + message = " ".join(fields[2:]) or "engine request failed" + with self.pending_lock: + events = self.pending.pop(request_id, None) + if events is not None: + events.put(("error", RuntimeError(message))) + else: + raise RuntimeError(f"invalid engine response: {' '.join(fields)}") + except Exception as error: + if not self.closed: + self.dispatcher_error = error + self._fail_pending(error) + + def generate(self, prompt, max_tokens, temperature, top_p, on_text, cache_slot=0, + cancelled=None): if isinstance(cache_slot, bool) or not isinstance(cache_slot, int) or not 0 <= cache_slot < self.kv_slots: raise APIError(400, "Invalid cache slot.", "cache_slot") payload = prompt.encode("utf-8") @@ -386,26 +481,66 @@ class Engine: if text: on_text(text) - with self.lock: + events = queue.Queue() + with self.pending_lock: + if self.closed: + raise RuntimeError("colibri engine is shutting down") + if self.dispatcher_error is not None: + raise RuntimeError("colibri engine dispatcher stopped") from self.dispatcher_error if self.process.poll() is not None: raise RuntimeError("colibri engine is not running") - header = (f"\x02PROMPT {len(payload)} {max_tokens} {temperature:.8g} " - f"{top_p:.8g} {cache_slot}\n").encode() - self.process.stdin.write(header + payload + b"\n") - self.process.stdin.flush() - stats = read_engine_turn(self.process.stdout, END, decode) - tail = decoder.decode(b"", final=True) - if tail: - on_text(tail) - return stats + request_id = str(self.next_request_id) + self.next_request_id += 1 + self.pending[request_id] = events + header = (f"SUBMIT {request_id} {cache_slot} {len(payload)} {max_tokens} " + f"{temperature:.8g} {top_p:.8g}\n").encode() + try: + with self.write_lock: + if self.process.poll() is not None: + raise RuntimeError("colibri engine is not running") + self.process.stdin.write(header + payload + b"\n") + self.process.stdin.flush() + except Exception: + with self.pending_lock: + self.pending.pop(request_id, None) + raise + + cancel_sent = False + while True: + kind, value = events.get() + if kind == "data": + if not cancel_sent: + decode(value) + if cancelled and cancelled(): + cancel_sent = True + with self.write_lock: + self.process.stdin.write(f"CANCEL {request_id}\n".encode()) + self.process.stdin.flush() + elif kind == "done": + tail = decoder.decode(b"", final=True) + if tail: + on_text(tail) + return value + elif cancel_sent and isinstance(value, RuntimeError) and str(value) == "CANCELLED": + raise ClientCancelled() + else: + raise value def close(self): + with self.pending_lock: + if self.closed: + return + self.closed = True + self._fail_pending(RuntimeError("colibri engine is shutting down")) if self.process.poll() is None: self.process.terminate() try: self.process.wait(timeout=5) except subprocess.TimeoutExpired: self.process.kill() + self.process.wait(timeout=5) + if self.dispatcher is not threading.current_thread(): + self.dispatcher.join(timeout=5) def model_object(model_id, created): @@ -423,7 +558,7 @@ class APIServer(ThreadingHTTPServer): self.model_id = model_id self.api_key = api_key self.max_tokens = max_tokens - self.scheduler = GenerationScheduler(max_queue, queue_timeout) + self.scheduler = GenerationScheduler(max_queue, queue_timeout, kv_slots) self.kv_slots = kv_slots self.cors_origins = tuple(cors_origins) self.created = int(time.time()) @@ -543,8 +678,10 @@ class APIHandler(BaseHTTPRequestHandler): def generation(self, body, prompt, request_id, chat): maximum, temperature, top_p = generation_options(body, self.server.max_tokens) tools = (body.get("tools") or body.get("functions") or None) if chat else None - cache_slot = body.get("cache_slot", 0) - if isinstance(cache_slot, bool) or not isinstance(cache_slot, int) or not 0 <= cache_slot < self.server.kv_slots: + cache_slot = body.get("cache_slot") + if (cache_slot is not None and + (isinstance(cache_slot, bool) or not isinstance(cache_slot, int) or + not 0 <= cache_slot < self.server.kv_slots)): raise APIError(400, f"`cache_slot` must be an integer between 0 and {self.server.kv_slots - 1}.", "cache_slot") stream = body.get("stream", False) @@ -559,12 +696,14 @@ class APIHandler(BaseHTTPRequestHandler): completion_id = id_prefix + uuid.uuid4().hex created = int(time.time()) - with self.server.scheduler.admit(self.client_disconnected) as queue_wait: + with self.server.scheduler.admit(self.client_disconnected, cache_slot) as admission: + queue_wait, cache_slot = admission queue_headers = {"x-colibri-queue-wait-ms": str(round(queue_wait * 1000))} if not stream: output = [] stats = self.server.engine.generate( - prompt, maximum, temperature, top_p, output.append, cache_slot) + prompt, maximum, temperature, top_p, output.append, cache_slot, + self.client_disconnected) text = "".join(output) length_finish = "length" if stats["length_limited"] else "stop" if chat and tools: @@ -670,7 +809,8 @@ class APIHandler(BaseHTTPRequestHandler): emit(sp["buf"][:flush]) sp["buf"] = sp["buf"][flush:] stats = self.server.engine.generate( - prompt, maximum, temperature, top_p, emit_tools, cache_slot) + prompt, maximum, temperature, top_p, emit_tools, cache_slot, + lambda: not connected) if not sp["tool"] and sp["buf"]: emit(sp["buf"]) # no tool call happened: flush held tail _content, calls = parse_tool_calls("".join(raw), tools) @@ -686,7 +826,8 @@ class APIHandler(BaseHTTPRequestHandler): sys.stderr.write(chunk); sys.stderr.flush() emit(chunk) stats = self.server.engine.generate( - prompt, maximum, temperature, top_p, emit_plain, cache_slot) + prompt, maximum, temperature, top_p, emit_plain, cache_slot, + lambda: not connected) finish = "length" if stats["length_limited"] else "stop" ka_stop.set() # generation done: stop the keepalive pump ka_thread.join(timeout=2) @@ -762,20 +903,25 @@ def serve(model, host="127.0.0.1", port=8000, model_id="glm-5.2-colibri", api_ke raise ValueError("kv_slots must be between 1 and 16") if host not in ("127.0.0.1", "localhost", "::1") and not api_key: print("WARNING: API is listening beyond localhost without COLI_API_KEY", file=sys.stderr) - runtime = Engine(engine,model,cap,max_tokens,env,kv_slots) origins = DEFAULT_CORS_ORIGINS if cors_origins is None else tuple(cors_origins) - server = APIServer((host, port), runtime, model_id, api_key, max_tokens, origins, + # Bind before starting the 744B engine. A stale/occupied port must fail in + # milliseconds rather than loading hundreds of GB and leaking a child. + server = APIServer((host, port), None, model_id, api_key, max_tokens, origins, max_queue, queue_timeout, kv_slots) - print(f"OpenAI-compatible API listening on http://{host}:{port}/v1", file=sys.stderr) + runtime = None previous_sigterm = signal.getsignal(signal.SIGTERM) - signal.signal(signal.SIGTERM, lambda *_: threading.Thread(target=server.shutdown, daemon=True).start()) try: + runtime = Engine(engine,model,cap,max_tokens,env,kv_slots) + server.engine = runtime + print(f"OpenAI-compatible API listening on http://{host}:{port}/v1", file=sys.stderr) + signal.signal(signal.SIGTERM, lambda *_: threading.Thread(target=server.shutdown, daemon=True).start()) server.serve_forever() finally: signal.signal(signal.SIGTERM, previous_sigterm) server.scheduler.close() server.server_close() - runtime.close() + if runtime is not None: + runtime.close() def main(): diff --git a/c/resource_plan.py b/c/resource_plan.py index 22ac79e..1c3a572 100644 --- a/c/resource_plan.py +++ b/c/resource_plan.py @@ -105,9 +105,33 @@ def discover_gpus(): return devices +def physical_cpu_count(): + try: + result = subprocess.run(["lscpu", "-p=core,socket"], text=True, + capture_output=True, check=True, timeout=5) + cores = {tuple(map(int, line.split(","))) for line in result.stdout.splitlines() + if line and not line.startswith("#")} + if cores: + return len(cores) + except (OSError, ValueError, subprocess.SubprocessError): + pass + return os.cpu_count() or 1 + + +POLICIES = { + "quality": {"preserve_quantization": True, "preserve_router": True}, + "balanced": {"preserve_quantization": True, "preserve_router": True}, + "experimental-fast": {"preserve_quantization": False, "preserve_router": False}, +} + + def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0, - available_memory=None, available_disk=None, gpus=None): + available_memory=None, available_disk=None, gpus=None, + policy="quality", physical_cpus=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 cfg = info["config"] available_memory = memory_available() if available_memory is None else available_memory if available_disk is None: @@ -146,8 +170,13 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0, safe_vram += usable gpu_plan.append(dict(gpu, reserve_bytes=reserve, usable_bytes=usable)) requested_vram = int(vram_gb * GB) if vram_gb > 0 else safe_vram - vram_budget = min(requested_vram, safe_vram, cache_bytes) + # VRAM-resident experts do not need duplicate RAM backing: the checkpoint is + # their recovery source. RAM is therefore an independent warm compute tier. + vram_budget = min(requested_vram, safe_vram, info["expert_bytes"]) vram_experts = int(vram_budget // typical) if typical else 0 + hot_bytes = min(info["expert_bytes"], vram_experts * typical) + warm_bytes = min(max(0, info["expert_bytes"] - hot_bytes), cache_bytes) + cold_bytes = max(0, info["expert_bytes"] - hot_bytes - warm_bytes) warnings = [] if cap < 1: @@ -155,21 +184,41 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0, if gpu_indices is not None and len(gpus) != len(set(gpu_indices)): warnings.append("one or more requested GPUs were not detected") if gpus and vram_budget < requested_vram: - warnings.append("VRAM tier was clamped by free VRAM or its required RAM backing") + warnings.append("VRAM tier was clamped by free VRAM or model expert size") + if cold_bytes: + warnings.append("cold expert misses may reach disk; normal decode speed depends on hit rate") + + if cold_bytes: + bottleneck = "disk expert misses" + elif warm_bytes: + bottleneck = "CPU expert compute and RAM bandwidth" + else: + bottleneck = "GPU compute and interconnect" return { - "version": 1, + "version": 2, + "policy": {"name": policy, **POLICIES[policy], + "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)), + "thread_policy": "physical-cores"}, "tiers": { - "disk": {"role": "backing", "model_bytes": info["model_bytes"], - "available_bytes": available_disk}, - "ram": {"role": "resident+cache", "available_bytes": available_memory, + "disk": {"role": "cold-backing", "model_bytes": info["model_bytes"], + "available_bytes": available_disk, "cold_expert_bytes": cold_bytes}, + "ram": {"role": "resident+warm-experts", "available_bytes": available_memory, "budget_bytes": ram_budget, "dense_bytes": info["dense_bytes"], "runtime_bytes": runtime_bytes, "expert_cache_bytes": cache_bytes, - "cache_slots_per_layer": cap}, + "warm_expert_bytes": warm_bytes, "cache_slots_per_layer": cap}, "vram": {"role": "hot-experts", "devices": gpu_plan, - "budget_bytes": vram_budget, "expert_capacity": vram_experts}, + "budget_bytes": vram_budget, "hot_expert_bytes": hot_bytes, + "expert_capacity": vram_experts, "requires_host_backing": False}, }, + "expected_bottleneck": bottleneck, + "decisions": [ + {"target": "VRAM", "reason": "profile-ranked hot experts"}, + {"target": "RAM", "reason": "warm experts execute on CPU without quality loss"}, + {"target": "Disk", "reason": "immutable recovery source for cold experts"}, + ], "warnings": warnings, } @@ -177,6 +226,12 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0, def environment_for_plan(plan, env=None, cuda_enabled=True): """Apply a plan without overriding explicit user environment settings.""" result = dict(env or {}) + result.setdefault("COLI_POLICY", plan["policy"]["name"]) + result.setdefault("OMP_NUM_THREADS", str(plan["cpu"]["physical_cores"])) + result.setdefault("OMP_PROC_BIND", "spread") + result.setdefault("OMP_PLACES", "cores") + if plan["policy"]["name"] == "balanced": + result.setdefault("REPIN", "64") ram = plan["tiers"]["ram"] result.setdefault("RAM_GB", f"{ram['budget_bytes'] / GB:.3f}") @@ -203,11 +258,15 @@ def format_bytes(value): def format_plan(plan): model, tiers = plan["model"], plan["tiers"] - lines = [f"model {model['shards']} shards · {format_bytes(model['model_bytes'])}", - f"disk backing store · {format_bytes(tiers['disk']['available_bytes'])} free", + policy=plan["policy"] + lines = [f"policy {policy['name']} · quality-preserving {'yes' if policy['quality_preserving'] else 'no'}", + f"model {model['shards']} shards · {format_bytes(model['model_bytes'])}", + f"disk {format_bytes(tiers['disk']['cold_expert_bytes'])} cold experts · " + f"{format_bytes(tiers['disk']['available_bytes'])} free", f"RAM {format_bytes(tiers['ram']['budget_bytes'])} budget · " f"{format_bytes(tiers['ram']['dense_bytes'])} dense · " f"{format_bytes(tiers['ram']['runtime_bytes'])} runtime · " + f"{format_bytes(tiers['ram']['warm_expert_bytes'])} warm experts · " f"cap {tiers['ram']['cache_slots_per_layer']}/layer"] vram = tiers["vram"] if vram["devices"]: @@ -216,5 +275,6 @@ def format_plan(plan): f"~{vram['expert_capacity']} experts · {names}") else: lines.append("VRAM no NVIDIA device detected · CPU path") + lines.append(f"limit {plan['expected_bottleneck']}") lines.extend(f"warn {warning}" for warning in plan["warnings"]) return "\n".join(lines) diff --git a/c/tests/bench_tensor_core.cu b/c/tests/bench_tensor_core.cu new file mode 100644 index 0000000..c8c4362 --- /dev/null +++ b/c/tests/bench_tensor_core.cu @@ -0,0 +1,45 @@ +#include "../backend_cuda.h" + +#include +#include +#include +#include +#include + +static double run(ColiCudaTensor *g,ColiCudaTensor *u,ColiCudaTensor *d, + const float *x,float *y,int rows,int iterations,int mode){ + ColiCudaTensor *gs[1]={g},*us[1]={u},*ds[1]={d}; int rs[1]={rows}; + if(mode==2){setenv("COLI_CUDA_TC_INT4","1",1);setenv("COLI_CUDA_TC_MIN_ROWS","1",1);} + else unsetenv("COLI_CUDA_TC_INT4"); + setenv("COLI_CUDA_W4_PACKED",mode==0?"0":"1",1); + if(!coli_cuda_expert_group(gs,us,ds,rs,1,y,x))std::exit(2); + auto begin=std::chrono::steady_clock::now(); + for(int i=0;i(end-begin).count()/iterations; +} + +int main(){ + constexpr int D=6144,I=2048,O=8; + int device=0;if(!coli_cuda_init(&device,1))return 77; + std::vector hidden((size_t)I*D/2),down((size_t)D*I/2); + std::vector hs(I),ds(D),x((size_t)O*D),a((size_t)O*D),b((size_t)O*D),c((size_t)O*D); + for(size_t i=0;ilimit){std::fprintf(stderr,"relative RMS %.5f exceeds %.5f\n",r,limit);return 0;} return 1; +} + int main(int argc, char **argv) { int devices[COLI_CUDA_MAX_DEVICES], ndev = argc > 1 ? argc - 1 : 1; if (ndev > COLI_CUDA_MAX_DEVICES) return 2; @@ -55,8 +61,67 @@ int main(int argc, char **argv) { ColiCudaTensor *tf = nullptr; if (!coli_cuda_matmul(&tf, got, x, wf, nullptr, 0, 1, 4, 2, d0) || !close_enough(got, wantf, 2)) return 1; + const float eg[8] = {1,0,0,0, 0,1,0,0}; + const float eu[8] = {1,0,0,0, 0,1,0,0}; + const float ed[8] = {1,0, 0,1, 1,1, 1,-1}; + ColiCudaTensor *tg=nullptr,*tu=nullptr,*td=nullptr; + if (!coli_cuda_tensor_upload(&tg,eg,nullptr,0,4,2,d0) || + !coli_cuda_tensor_upload(&tu,eu,nullptr,0,4,2,d0) || + !coli_cuda_tensor_upload(&td,ed,nullptr,0,2,4,d0)) return 1; + float expert[8], want_expert[8]; + for(int s=0;s<2;s++){ + float a=x[s*4], b=x[s*4+1]; + a=(a/(1.0f+std::exp(-a)))*a; b=(b/(1.0f+std::exp(-b)))*b; + want_expert[s*4]=a; want_expert[s*4+1]=b; + want_expert[s*4+2]=a+b; want_expert[s*4+3]=a-b; + } + if (!coli_cuda_expert_mlp(tg,tu,td,expert,x,2) || + !close_enough(expert,want_expert,8)) return 1; + ColiCudaTensor *gates[2]={tg,tg},*ups[2]={tu,tu},*downs[2]={td,td}; + int group_rows[2]={1,1}; float grouped[8]; + if (!coli_cuda_expert_group(gates,ups,downs,group_rows,2,grouped,x) || + !close_enough(grouped,want_expert,8)) return 1; + + const float aw[16]={1,0,0,0, 0,1,0,0, 0,0,1,0, 0,0,0,1}; + const float aq[4]={1,2,.5f,-.5f},al[12]={1,0,0,0, 0,1,0,0, 0,0,1,0}; + const float ar[6]={1,0, 0,1, 1,1};float actx[2],aref[2]; + ColiCudaTensor *at=nullptr;if(!coli_cuda_tensor_upload(&at,aw,nullptr,0,4,4,d0))return 1; + float score[3];for(int t=0;t<3;t++)score[t]=aq[0]*al[t*4]+aq[1]*al[t*4+1]+aq[2]*ar[t*2]+aq[3]*ar[t*2+1]; + float mx=score[0],z=0;for(int t=1;t<3;t++)mx=score[t]>mx?score[t]:mx; + for(int t=0;t<3;t++){score[t]=std::exp(score[t]-mx);z+=score[t];}for(int t=0;t<3;t++)score[t]/=z; + for(int v=0;v<2;v++){aref[v]=0;for(int t=0;t<3;t++)aref[v]+=score[t]*al[t*4+2+v];} + if(!coli_cuda_attention_absorb(at,actx,aq,al,ar,1,2,2,2,4,3,1.f)|| + !close_enough(actx,aref,2))return 1; + coli_cuda_tensor_free(at); + + /* Native s4 WMMA path: compare the quantized-activation result against the + existing FP32-activation/s4-weight grouped implementation. */ + uint8_t w4[32*32/2]; float ws4[32], gx4[64], scalar4[64], tensor4[64]; + for(int i=0;i<(int)sizeof(w4);i++){ + int lo=((i%15)-7)&15,hi=(((i*3)%15)-7)&15; + w4[i]=(uint8_t)(lo|(hi<<4)); + } + for(int i=0;i<32;i++)ws4[i]=0.01f+(i%5)*0.002f; + for(int i=0;i<64;i++)gx4[i]=std::sin((float)(i+1)*0.17f)*2.f; + ColiCudaTensor *g4=nullptr,*u4=nullptr,*d4=nullptr; + if(!coli_cuda_tensor_upload(&g4,w4,ws4,2,32,32,d0)|| + !coli_cuda_tensor_upload(&u4,w4,ws4,2,32,32,d0)|| + !coli_cuda_tensor_upload(&d4,w4,ws4,2,32,32,d0))return 1; + ColiCudaTensor *gg4[2]={g4,g4},*ug4[2]={u4,u4},*dg4[2]={d4,d4}; + if(!coli_cuda_expert_group(gg4,ug4,dg4,group_rows,2,scalar4,gx4))return 1; + setenv("COLI_CUDA_TC_INT4","1",1); + setenv("COLI_CUDA_TC_MIN_ROWS","1",1); + if(!coli_cuda_expert_group(gg4,ug4,dg4,group_rows,2,tensor4,gx4)|| + !relative_rms(tensor4,scalar4,64,0.30f))return 1; + unsetenv("COLI_CUDA_TC_INT4"); + unsetenv("COLI_CUDA_TC_MIN_ROWS"); + coli_cuda_tensor_free(g4);coli_cuda_tensor_free(u4);coli_cuda_tensor_free(d4); + uint64_t group_calls=0,group_experts=0,group_total_rows=0; + coli_cuda_group_stats(&group_calls,&group_experts,&group_total_rows,nullptr,nullptr,nullptr); + if(group_calls!=3||group_experts!=6||group_total_rows!=6) return 1; + coli_cuda_stats(-1, &count, &bytes); - if (count != 4 || bytes != 70) { + if (count != 7 || bytes != 166) { std::fprintf(stderr, "unexpected CUDA stats: %zu tensors, %zu bytes\n", count, bytes); return 1; } @@ -64,15 +129,18 @@ int main(int argc, char **argv) { coli_cuda_tensor_device(t4) != d1 || coli_cuda_tensor_device(t2) != d1) return 1; coli_cuda_stats(d0, &count, &bytes); if (ndev > 1) { - if (count != 2 || bytes != 48) return 1; + if (count != 5 || bytes != 144) return 1; coli_cuda_stats(d1, &count, &bytes); if (count != 2 || bytes != 22) return 1; - } else if (count != 4 || bytes != 70) return 1; + } else if (count != 7 || bytes != 166) return 1; coli_cuda_tensor_free(t8); coli_cuda_tensor_free(t4); coli_cuda_tensor_free(t2); coli_cuda_tensor_free(tf); + coli_cuda_tensor_free(tg); + coli_cuda_tensor_free(tu); + coli_cuda_tensor_free(td); coli_cuda_stats(-1, &count, &bytes); if (count || bytes) return 1; coli_cuda_shutdown(); diff --git a/c/tests/test_decode_batch b/c/tests/test_decode_batch new file mode 100755 index 0000000..989fe42 Binary files /dev/null and b/c/tests/test_decode_batch differ diff --git a/c/tests/test_decode_batch.c b/c/tests/test_decode_batch.c new file mode 100644 index 0000000..d370d2d --- /dev/null +++ b/c/tests/test_decode_batch.c @@ -0,0 +1,56 @@ +#include +#include + +#include "../decode_batch.h" + +static void test_rows_use_their_own_sequence_storage(void) +{ + float sequence_a[4 * 3] = {0}; + float sequence_b[4 * 3] = {0}; + + float *a2 = coli_kv_row(sequence_a, 2, 3); + float *b1 = coli_kv_row(sequence_b, 1, 3); + a2[0] = 20.0f; + b1[2] = 12.0f; + + assert(a2 == &sequence_a[6]); + assert(b1 == &sequence_b[3]); + assert(sequence_a[6] == 20.0f); + assert(sequence_b[5] == 12.0f); + assert(sequence_a[5] == 0.0f); + assert(sequence_b[6] == 0.0f); +} + +static void test_const_reader_selects_the_same_row(void) +{ + float storage[5 * 7] = {0}; + const float *row = coli_kv_row(storage, 4, 7); + + assert(row == &storage[28]); +} + +static void test_submit_header(void) +{ + ColiSubmit sub; + assert(coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95", &sub)); + assert(sub.id == 42 && sub.slot == 3 && sub.bytes == 17); + assert(sub.max_tokens == 64 && sub.temperature > .69f && sub.top_p > .94f); + assert(!coli_submit_parse("SUBMIT 1 -1 2 3 0.7 1", &sub)); + assert(!coli_submit_parse("SUBMIT 1 0 2 0 0.7 1", &sub)); + assert(!coli_submit_parse("SUBMIT 1 0 2 3 4 1", &sub)); + assert(!coli_submit_parse("SUBMIT 0 0 2 3 1 1", &sub)); + assert(!coli_submit_parse("SUBMIT 1 0 2 3 nan 1", &sub)); + assert(!coli_submit_parse("SUBMIT 1 0 2 3 1 inf", &sub)); + assert(coli_submit_parse("SUBMIT 1 0 16777216 3 1 1", &sub)); + assert(!coli_submit_parse("SUBMIT 1 0 16777217 3 1 1", &sub)); + assert(!coli_submit_parse("SUBMIT 1 0 2 3 1 1 trailing", &sub)); +} + +int main(void) +{ + test_rows_use_their_own_sequence_storage(); + test_const_reader_selects_the_same_row(); + test_submit_header(); + puts("decode batch helper tests: ok"); + return 0; +} diff --git a/c/tests/test_doctor.py b/c/tests/test_doctor.py index c618636..1647a0d 100644 --- a/c/tests/test_doctor.py +++ b/c/tests/test_doctor.py @@ -142,7 +142,7 @@ class DoctorTest(unittest.TestCase): output = format_doctor(self.report()) self.assertIn("model.path", output) - self.assertIn("disk backing store", output) + self.assertIn("disk 0.0 GB cold experts", output) self.assertTrue(output.endswith("result ok")) def test_cli_json_is_machine_readable_without_loading_model(self): diff --git a/c/tests/test_i4_acc512 b/c/tests/test_i4_acc512 new file mode 100755 index 0000000..76f68a3 Binary files /dev/null and b/c/tests/test_i4_acc512 differ diff --git a/c/tests/test_idot b/c/tests/test_idot new file mode 100755 index 0000000..74dd03c Binary files /dev/null and b/c/tests/test_idot differ diff --git a/c/tests/test_openai_server.py b/c/tests/test_openai_server.py index 06d3218..9792403 100644 --- a/c/tests/test_openai_server.py +++ b/c/tests/test_openai_server.py @@ -1,19 +1,24 @@ import io import json +import math +import socket import threading import unittest +from unittest.mock import patch from urllib.error import HTTPError from urllib.request import Request, urlopen from openai_server import (APIError, APIServer, ClientCancelled, END, GenerationScheduler, - generation_options, read_engine_turn, render_chat, serve) + READY, Engine, generation_options, read_engine_turn, render_chat, + serve) class FakeEngine: def __init__(self): self.calls = [] - def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0): + def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, + cancelled=None): self.calls.append((prompt, maximum, temperature, top_p, cache_slot)) on_text("Hé") on_text("llo") @@ -26,10 +31,12 @@ class BlockingEngine(FakeEngine): self.entered = threading.Event() self.release = threading.Event() - def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0): + def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0, + cancelled=None): self.entered.set() self.release.wait(2) - return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot) + return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot, + cancelled) class TemplateTest(unittest.TestCase): @@ -65,6 +72,10 @@ class TemplateTest(unittest.TestCase): (4, 0.0, 1.0)) with self.assertRaises(APIError): generation_options({"max_tokens": 9}, 8) + with self.assertRaises(APIError): + generation_options({"temperature": math.nan}, 8) + with self.assertRaises(APIError): + generation_options({"top_p": math.inf}, 8) self.assertEqual(generation_options({"temperature": None, "top_p": None}, 8), (8, 0.7, 0.9)) @@ -82,8 +93,27 @@ class ProtocolTest(unittest.TestCase): with self.assertRaisesRegex(ValueError, "kv_slots"): serve("/missing", kv_slots=0) + def test_occupied_port_fails_before_engine_start(self): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen() + try: + with patch("openai_server.subprocess.Popen") as popen: + with self.assertRaises(OSError): + serve("/missing", port=listener.getsockname()[1]) + popen.assert_not_called() + finally: + listener.close() + class SchedulerTest(unittest.TestCase): + def test_admits_up_to_capacity_without_serializing(self): + scheduler = GenerationScheduler(max_queue=0, queue_timeout=1, capacity=2) + with scheduler.admit() as first: + with scheduler.admit() as second: + self.assertEqual({first[1], second[1]}, {0, 1}) + self.assertEqual(scheduler.snapshot()["active"], 2) + def test_rejects_when_waiting_queue_is_full(self): scheduler = GenerationScheduler(max_queue=0, queue_timeout=1) with scheduler.admit(): @@ -160,6 +190,213 @@ class SchedulerTest(unittest.TestCase): self.assertEqual(errors, ["scheduler_closed"]) +class BlockingStream: + def __init__(self, initial=b""): + self.buffer = bytearray(initial) + self.closed = False + self.condition = threading.Condition() + + def feed(self, data): + with self.condition: + self.buffer.extend(data) + self.condition.notify_all() + + def read(self, size=1): + with self.condition: + while len(self.buffer) < size and not self.closed: + self.condition.wait() + if not self.buffer and self.closed: + return b"" + size = min(size, len(self.buffer)) + data = bytes(self.buffer[:size]) + del self.buffer[:size] + return data + + def readline(self): + with self.condition: + while b"\n" not in self.buffer and not self.closed: + self.condition.wait() + if not self.buffer and self.closed: + return b"" + end = self.buffer.find(b"\n") + size = len(self.buffer) if end < 0 else end + 1 + data = bytes(self.buffer[:size]) + del self.buffer[:size] + return data + + def close(self): + with self.condition: + self.closed = True + self.condition.notify_all() + + +class FakeProcess: + def __init__(self, on_write): + self.stdout = BlockingStream(READY + b"STAT 0 0 0 0\n") + self.stdin = self + self.on_write = on_write + self.writes = [] + self.returncode = None + + def write(self, data): + self.writes.append(data) + self.on_write(self, data) + return len(data) + + def flush(self): + pass + + def poll(self): + return self.returncode + + def terminate(self): + self.returncode = 0 + self.stdout.close() + + def wait(self, timeout=None): + return self.returncode + + def kill(self): + self.terminate() + + +class DispatcherTest(unittest.TestCase): + def test_dispatches_interleaved_requests_by_id(self): + submitted = [] + + def respond(process, frame): + fields = frame.split(b"\n", 1)[0].split() + self.assertEqual(fields[0], b"SUBMIT") + submitted.append(fields[1]) + if len(submitted) == 2: + first, second = submitted + process.stdout.feed(b"DATA " + second + b" 3\nB-2\n") + process.stdout.feed(b"DATA " + first + b" 3\nA-1\n") + process.stdout.feed(b"DONE " + second + b" STAT 1 2.5 0 1.0 4 0\n") + process.stdout.feed(b"DATA " + first + b" 3\nA-2\n") + process.stdout.feed(b"DONE " + first + b" STAT 2 3.5 0 1.0 5 1\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model", kv_slots=2) + results = {} + + def generate(name, prompt, slot): + chunks = [] + stats = engine.generate(prompt, 8, 0.7, 0.9, chunks.append, slot) + results[name] = ("".join(chunks), stats) + + threads = [threading.Thread(target=generate, args=("a", "alpha", 0)), + threading.Thread(target=generate, args=("b", "beta", 1))] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=2) + self.assertFalse(thread.is_alive()) + engine.close() + + self.assertEqual(results["a"][0], "A-1A-2") + self.assertTrue(results["a"][1]["length_limited"]) + self.assertEqual(results["b"][0], "B-2") + headers = [frame.split(b"\n", 1)[0].split() for frame in process.writes] + self.assertEqual({int(header[2]) for header in headers}, {0, 1}) + self.assertEqual({header[3] for header in headers}, {b"4", b"5"}) + + def test_routes_engine_error_to_request(self): + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"ERROR " + request_id + b" slot is busy\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "slot is busy"): + engine.generate("hello", 4, 0.7, 0.9, lambda _: None) + engine.close() + + def test_close_wakes_pending_generation_and_is_idempotent(self): + process = FakeProcess(lambda _process, _frame: None) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + errors = [] + + def generate(): + try: + engine.generate("hello", 4, 0.7, 0.9, lambda _: None) + except RuntimeError as error: + errors.append(str(error)) + + thread = threading.Thread(target=generate) + thread.start() + for _ in range(100): + with engine.pending_lock: + if engine.pending: + break + threading.Event().wait(0.01) + engine.close() + engine.close() + thread.join(timeout=2) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, ["colibri engine is shutting down"]) + self.assertFalse(engine.dispatcher.is_alive()) + with engine.pending_lock: + self.assertFalse(engine.pending) + with self.assertRaisesRegex(RuntimeError, "shutting down"): + engine.generate("again", 4, 0.7, 0.9, lambda _: None) + + def test_protocol_corruption_fails_request_and_stops_dispatcher(self): + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"DATA " + request_id + b" -1\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + with self.assertRaisesRegex(RuntimeError, "DATA size"): + engine.generate("hello", 4, 0.7, 0.9, lambda _: None) + with self.assertRaisesRegex(RuntimeError, "dispatcher stopped"): + engine.generate("again", 4, 0.7, 0.9, lambda _: None) + engine.close() + + def test_decodes_utf8_split_across_data_frames(self): + def respond(process, frame): + request_id = frame.split()[1] + process.stdout.feed(b"DATA " + request_id + b" 1\n\xc3\n") + process.stdout.feed(b"DATA " + request_id + b" 1\n\xa9\n") + process.stdout.feed(b"DONE " + request_id + b" STAT 1 1 0 1 1 0\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + chunks = [] + engine.generate("hello", 4, 0.7, 0.9, chunks.append) + engine.close() + self.assertEqual(chunks, ["é"]) + + def test_cancels_generation_after_consumer_disconnects(self): + request_id = None + + def respond(process, frame): + nonlocal request_id + fields = frame.split() + if fields[0] == b"SUBMIT": + request_id = fields[1] + process.stdout.feed(b"DATA " + request_id + b" 1\nx\n") + elif fields[0] == b"CANCEL": + self.assertEqual(fields[1], request_id) + process.stdout.feed(b"ERROR " + request_id + b" CANCELLED\n") + + process = FakeProcess(respond) + with patch("openai_server.subprocess.Popen", return_value=process): + engine = Engine("glm", "model") + output = [] + with self.assertRaises(ClientCancelled): + engine.generate("hello", 8, 0.7, 0.9, output.append, cancelled=lambda: True) + engine.close() + self.assertEqual(output, ["x"]) + self.assertEqual(process.writes[-1].split(), [b"CANCEL", request_id]) + + class HTTPTest(unittest.TestCase): @classmethod def setUpClass(cls): diff --git a/c/tests/test_resource_plan.py b/c/tests/test_resource_plan.py index 3b965fe..d44f661 100644 --- a/c/tests/test_resource_plan.py +++ b/c/tests/test_resource_plan.py @@ -57,11 +57,16 @@ class ResourcePlanTest(unittest.TestCase): 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) - self.assertEqual(plan["version"], 1) + available_memory=32 * GB, available_disk=100 * GB, gpus=gpus, + physical_cpus=24) + self.assertEqual(plan["version"], 2) + self.assertEqual(plan["policy"]["name"], "quality") + self.assertEqual(plan["cpu"]["physical_cores"], 24) + self.assertTrue(plan["policy"]["preserve_quantization"]) + self.assertFalse(plan["tiers"]["vram"]["requires_host_backing"]) self.assertEqual(plan["tiers"]["ram"]["budget_bytes"], 16 * GB) self.assertLessEqual(plan["tiers"]["vram"]["budget_bytes"], 8 * GB) - self.assertIn("required RAM backing", plan["warnings"][0]) + self.assertIn("clamped", plan["warnings"][0]) self.assertIn("0:test-gpu", format_plan(plan)) def test_filters_requested_devices(self): @@ -78,7 +83,7 @@ class ResourcePlanTest(unittest.TestCase): "--gpu", "none", "--json", ], text=True, capture_output=True, check=True) plan = json.loads(run.stdout) - self.assertEqual(plan["version"], 1) + self.assertEqual(plan["version"], 2) self.assertEqual(plan["model"]["expert_count"], 2) def test_applies_plan_without_overriding_explicit_settings(self): @@ -93,8 +98,15 @@ class ResourcePlanTest(unittest.TestCase): self.assertEqual(env["RAM_GB"], "12") self.assertEqual(env["COLI_CUDA"], "1") self.assertEqual(env["COLI_GPUS"], "1") + self.assertEqual(env["OMP_NUM_THREADS"], str(plan["cpu"]["physical_cores"])) + self.assertEqual(env["OMP_PROC_BIND"], "spread") + self.assertEqual(env["OMP_PLACES"], "cores") self.assertEqual(env["PIN_GB"], env["CUDA_EXPERT_GB"]) + explicit_threads = environment_for_plan(plan, {"OMP_NUM_THREADS": "7", + "OMP_PROC_BIND": "close"}) + self.assertEqual(explicit_threads["OMP_NUM_THREADS"], "7") + self.assertEqual(explicit_threads["OMP_PROC_BIND"], "close") 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, @@ -106,6 +118,32 @@ class ResourcePlanTest(unittest.TestCase): self.assertNotIn("COLI_GPU", disabled) self.assertNotIn("CUDA_EXPERT_GB", disabled) + def test_rejects_unknown_policy_and_marks_experimental_policy(self): + with self.assertRaisesRegex(ValueError, "unknown policy"): + build_plan(self.model, available_memory=16 * GB, available_disk=1, + gpus=[], policy="fast-ish") + plan = build_plan(self.model, available_memory=16 * GB, available_disk=1, + gpus=[], policy="experimental-fast") + self.assertFalse(plan["policy"]["quality_preserving"]) + self.assertFalse(plan["policy"]["preserve_router"]) + + def test_balanced_policy_enables_lossless_live_repin(self): + plan = build_plan(self.model, available_memory=16 * GB, available_disk=1, + gpus=[], policy="balanced") + env = environment_for_plan(plan) + self.assertEqual(env["COLI_POLICY"], "balanced") + self.assertEqual(env["REPIN"], "64") + explicit = environment_for_plan(plan, {"REPIN": "0"}) + self.assertEqual(explicit["REPIN"], "0") + + def test_plan_explains_hot_warm_and_cold_placement(self): + plan = build_plan(self.model, ram_gb=4, vram_gb=0, + available_memory=4 * GB, available_disk=1, gpus=[]) + self.assertEqual([item["target"] for item in plan["decisions"]], + ["VRAM", "RAM", "Disk"]) + self.assertIn("quality-preserving yes", format_plan(plan)) + self.assertIn("expected_bottleneck", plan) + if __name__ == "__main__": unittest.main() diff --git a/c/tests/test_tier.c b/c/tests/test_tier.c index 9b68981..eb6f1ac 100644 --- a/c/tests/test_tier.c +++ b/c/tests/test_tier.c @@ -17,6 +17,11 @@ int main(void){ tier_decay(heat,6); if(heat[0]!=10 || heat[1]!=1 || heat[4]!=15) return fail("heat decay"); + + uint32_t freq[5]={10,10,2,18,18}, last[5]={10,90,95,20,99}; + int live[2]={0,1}; + if(!tier_pick_lfru(freq,last,100,5,live,2,&slot,&eid,&gain)) return fail("LFRU promotion"); + if(slot!=0||eid!=4) return fail("LFRU did not prefer recent ties"); puts("tier tests: ok"); return 0; } diff --git a/c/tier.h b/c/tier.h index 8f66cdd..859cacf 100644 --- a/c/tier.h +++ b/c/tier.h @@ -24,6 +24,35 @@ static int tier_pick_swap(const uint32_t *heat, int nexpert, return 1; } +/* LFRU: frequency is the primary signal; recency breaks close calls. A recent + * access contributes at most 255 points while one frequency count is worth + * 256, so a merely recent expert cannot displace a genuinely hotter one. */ +static uint64_t tier_lfru_score(uint32_t heat, uint32_t last, uint32_t clock){ + uint32_t age=clock-last, recent=age<255?255-age:0; + return ((uint64_t)heat<<8)|recent; +} + +static int tier_pick_lfru(const uint32_t *heat, const uint32_t *last, uint32_t clock, + int nexpert, const int *pinned, int npin, + int *slot, int *eid, long *gain){ + if(!heat||!last||!pinned||npin<1||nexpert<1) return 0; + int cold=0; + for(int z=1;zhs)){ hot=e; hs=score; } + } + if(hot<0) return 0; + uint64_t cs=tier_lfru_score(heat[pinned[cold]],last[pinned[cold]],clock); + /* Retain the existing 25%+4-frequency hysteresis in score units. */ + if(hs<=cs+(cs>>2)+(4u<<8)) return 0; + *slot=cold; *eid=hot; *gain=(long)((hs-cs)>>8); return 1; +} + static void tier_decay(uint32_t *heat, int nexpert){ for(int e=0;e>=1; }