Unify continuous batching + heterogeneous runtime: decode batching, physical-core planning, disjoint VRAM/RAM placement, topp-policy warning (CPU-validated, CUDA on 6x5090) (#68)
* Fuse CUDA expert MLP execution * Group CUDA expert transfers by device * Instrument grouped CUDA expert execution * Bound grouped CUDA decode scratch * Execute expert groups across GPUs in parallel * Release host backing for multi-GPU experts * Define quality-preserving memory policies * Overlap cold expert loading with resident compute * Adapt expert placement with session LFRU * Fuse q4 expert gate and up dispatch * Plan CPU work on physical cores * Batch grouped expert CUDA kernels * Separate VRAM and RAM expert placement * Add ragged multi-sequence decode forward * feat(runtime): add continuous decode scheduler * Route concurrent API requests through batch scheduler * Harden multiplex request lifecycle and framing * Cancel disconnected multiplex requests * Bind API port before starting the engine * fix automatic KV slot allocation * add native int4 Tensor Core grouped GEMM * add Tensor Core throughput benchmark * optimize packed int4 low-row kernels * add asynchronous CUDA staging streams * document validated six-GPU dense acceleration * tune six-GPU expert hot set * raise validated expert hot-set target * add CUDA MLA absorption core * fuse grouped expert gate and up projections * Warn for explicit lossy routing flags
This commit is contained in:
ZacharyZcR
@@ -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.
|
||||
|
||||
+11
-3
@@ -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
|
||||
|
||||
+341
-6
@@ -1,9 +1,12 @@
|
||||
#include "backend_cuda.h"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <mma.h>
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <mutex>
|
||||
|
||||
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<float>(((i & 1) ? (v >> 4) : (v & 15)) - 8);
|
||||
int n=(i&1)?(v>>4):(v&15); return static_cast<float>(n&8?n-16:n);
|
||||
}
|
||||
uint8_t v = q[i >> 2];
|
||||
return static_cast<float>(((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<n)q[i]^=0x88;
|
||||
}
|
||||
|
||||
__global__ static void quant_matmul(float *y, const float *x, const void *weights,
|
||||
const float *scales, int fmt, int S, int I, int O,
|
||||
size_t rb) {
|
||||
@@ -81,6 +102,137 @@ __global__ static void quant_matmul(float *y, const float *x, const void *weight
|
||||
y[(size_t)s * O + o] = partial[0] * (fmt ? scales[o] : 1.0f);
|
||||
}
|
||||
|
||||
__global__ static void silu_mul(float *gate, const float *up, size_t n) {
|
||||
size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < n) {
|
||||
float v = gate[i];
|
||||
gate[i] = (v / (1.0f + expf(-v))) * up[i];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void quantize_s4_rows(uint8_t *q,float *scale,const float *x,int S,int K){
|
||||
int s=blockIdx.x; if(s>=S)return; const float *xs=x+(size_t)s*K;
|
||||
float v=0; for(int i=threadIdx.x;i<K;i+=blockDim.x)v=fmaxf(v,fabsf(xs[i]));
|
||||
__shared__ float m[256]; m[threadIdx.x]=v; __syncthreads();
|
||||
for(int n=128;n;n>>=1){if(threadIdx.x<n)m[threadIdx.x]=fmaxf(m[threadIdx.x],m[threadIdx.x+n]);__syncthreads();}
|
||||
float sc=m[0]>0?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<K?__float2int_rn(xs[i+1]/sc):0;
|
||||
a=max(-8,min(7,a)); c=max(-8,min(7,c)); dst[b]=(uint8_t)((a&15)|((c&15)<<4));
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void grouped_s4_wmma(float *y,const uint8_t *x,const float *xscale,
|
||||
const GroupDesc *desc,int K,int O,int which){
|
||||
#if __CUDA_ARCH__ >= 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<wmma::accumulator,8,8,32,int> 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<K;k+=32){
|
||||
wmma::fragment<wmma::matrix_a,8,8,32,wmma::experimental::precision::s4,wmma::row_major> af;
|
||||
wmma::fragment<wmma::matrix_b,8,8,32,wmma::experimental::precision::s4,wmma::col_major> 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&&o<O)y[(size_t)(d.offset+s)*O+o]=(float)out[warp][i]*xscale[d.offset+s]*ws[o];}
|
||||
#endif
|
||||
}
|
||||
|
||||
__global__ static void grouped_hidden(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 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<D;i+=blockDim.x) sum+=xs[i]*weight_at(w,fmt,row,i);
|
||||
__shared__ float p[256]; p[threadIdx.x]=sum; __syncthreads();
|
||||
for(int n=128;n;n>>=1){ if(threadIdx.x<n)p[threadIdx.x]+=p[threadIdx.x+n]; __syncthreads(); }
|
||||
if(!threadIdx.x) y[(size_t)(d.offset+s)*I+o]=p[0]*(fmt?sc[o]:1.f);
|
||||
}
|
||||
|
||||
__global__ static void grouped_down(float *y,const float *x,const GroupDesc *desc,int D,int I){
|
||||
int o=blockIdx.x,s=blockIdx.y,c=blockIdx.z; GroupDesc d=desc[c];
|
||||
if(s>=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<I;i+=blockDim.x) sum+=xs[i]*weight_at(d.d,d.df,row,i);
|
||||
__shared__ float p[256]; p[threadIdx.x]=sum; __syncthreads();
|
||||
for(int n=128;n;n>>=1){ if(threadIdx.x<n)p[threadIdx.x]+=p[threadIdx.x+n]; __syncthreads(); }
|
||||
if(!threadIdx.x) y[(size_t)(d.offset+s)*D+o]=p[0]*(d.df?d.ds[o]:1.f);
|
||||
}
|
||||
|
||||
__device__ static void unpack_s4(uint8_t v,float *lo,float *hi){
|
||||
int a=v&15,b=v>>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<D)sum+=xs[i+1]*z;}
|
||||
__shared__ float p[256];p[threadIdx.x]=sum;__syncthreads();
|
||||
for(int n=128;n;n>>=1){if(threadIdx.x<n)p[threadIdx.x]+=p[threadIdx.x+n];__syncthreads();}
|
||||
if(!threadIdx.x)y[(size_t)(d.offset+s)*I+o]=p[0]*sc[o];
|
||||
}
|
||||
|
||||
__global__ static void grouped_hidden_w4_dual(float *gate,float *up,const float *x,
|
||||
const GroupDesc *desc,int I,int D){
|
||||
int o=blockIdx.x,s=blockIdx.y,c=blockIdx.z;GroupDesc d=desc[c];if(s>=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<D){ga+=xs[i+1]*g1;ua+=xs[i+1]*u1;}}
|
||||
__shared__ float gp[256],upv[256];gp[threadIdx.x]=ga;upv[threadIdx.x]=ua;__syncthreads();
|
||||
for(int n=128;n;n>>=1){if(threadIdx.x<n){gp[threadIdx.x]+=gp[threadIdx.x+n];upv[threadIdx.x]+=upv[threadIdx.x+n];}__syncthreads();}
|
||||
if(!threadIdx.x){size_t z=(size_t)(d.offset+s)*I+o;gate[z]=gp[0]*d.gs[o];up[z]=upv[0]*d.us[o];}
|
||||
}
|
||||
|
||||
__global__ static void grouped_down_w4(float *y,const float *x,const GroupDesc *desc,int D,int I){
|
||||
int o=blockIdx.x,s=blockIdx.y,c=blockIdx.z;GroupDesc d=desc[c];if(s>=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<I)sum+=xs[i+1]*z;}
|
||||
__shared__ float p[256];p[threadIdx.x]=sum;__syncthreads();
|
||||
for(int n=128;n;n>>=1){if(threadIdx.x<n)p[threadIdx.x]+=p[threadIdx.x+n];__syncthreads();}
|
||||
if(!threadIdx.x)y[(size_t)(d.offset+s)*D+o]=p[0]*d.ds[o];
|
||||
}
|
||||
|
||||
__global__ static void attention_absorb_kernel(float *ctx,const float *q,const float *latent,
|
||||
const float *rope,const void *weights,const float *wscale,
|
||||
int fmt,int H,int Q,int R,int V,int K,int T,float scale){
|
||||
int h=blockIdx.x,tid=threadIdx.x,rbase=h*(Q+V);extern __shared__ float sm[];
|
||||
float *qa=sm,*cl=qa+K,*scores=cl+K;
|
||||
for(int k=tid;k<K;k+=blockDim.x){float a=0;for(int d=0;d<Q;d++)
|
||||
a+=q[(size_t)h*(Q+R)+d]*weight_at(weights,fmt,(size_t)(rbase+d)*row_bytes(fmt,K),k)*(fmt?wscale[rbase+d]:1.f);qa[k]=a;}
|
||||
__syncthreads();
|
||||
for(int t=tid;t<T;t+=blockDim.x){float a=0;const float *lt=latent+(size_t)t*K,*rt=rope+(size_t)t*R;
|
||||
for(int k=0;k<K;k++)a+=qa[k]*lt[k];for(int d=0;d<R;d++)a+=q[(size_t)h*(Q+R)+Q+d]*rt[d];scores[t]=a*scale;}
|
||||
__syncthreads();
|
||||
if(!tid){float mx=scores[0];for(int t=1;t<T;t++)mx=fmaxf(mx,scores[t]);float z=0;
|
||||
for(int t=0;t<T;t++){scores[t]=expf(scores[t]-mx);z+=scores[t];}for(int t=0;t<T;t++)scores[t]/=z;}
|
||||
__syncthreads();
|
||||
for(int k=tid;k<K;k+=blockDim.x){float a=0;for(int t=0;t<T;t++)a+=scores[t]*latent[(size_t)t*K+k];cl[k]=a;}
|
||||
__syncthreads();
|
||||
for(int v=tid;v<V;v+=blockDim.x){int row=rbase+Q+v;float a=0;size_t rb=row_bytes(fmt,K);
|
||||
for(int k=0;k<K;k++)a+=cl[k]*weight_at(weights,fmt,(size_t)row*rb,k);ctx[(size_t)h*V+v]=a*(fmt?wscale[row]:1.f);}
|
||||
}
|
||||
|
||||
static int reserve(float **ptr, size_t *cap, size_t bytes) {
|
||||
if (*cap >= bytes) return 1;
|
||||
if (*ptr) cudaFree(*ptr);
|
||||
@@ -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<<<hidden_grid,256>>>(ctx->gate,ctx->x,gate->weights,gate->scales,
|
||||
gate->fmt,S,D,I,row_bytes(gate->fmt,D));
|
||||
quant_matmul<<<hidden_grid,256>>>(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<<<output_grid,256>>>(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;c<count;c++){
|
||||
ColiCudaTensor *g=gates[c],*u=ups[c],*d=downs[c];
|
||||
if(!g||!u||!d||rows[c]<1||g->device!=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<count&&tc;c++)tc=rows[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<<<total,256,0,ctx->stream>>>(ctx->qx,ctx->qscale,ctx->x,total,D);
|
||||
grouped_s4_wmma<<<dim3((unsigned)((I+63)/64),(unsigned)count),256,0,ctx->stream>>>(ctx->gate,ctx->qx,ctx->qscale,dev,D,I,0);
|
||||
grouped_s4_wmma<<<dim3((unsigned)((I+63)/64),(unsigned)count),256,0,ctx->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<<<total,256,0,ctx->stream>>>(ctx->qx,ctx->qscale,ctx->gate,total,I);
|
||||
grouped_s4_wmma<<<dim3((unsigned)((D+63)/64),(unsigned)count),256,0,ctx->stream>>>(ctx->y,ctx->qx,ctx->qscale,dev,I,D,2);
|
||||
}else if(all_s4&&(!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<<<hg,256,0,ctx->stream>>>(ctx->gate,ctx->up,ctx->x,dev,I,D);
|
||||
else{
|
||||
grouped_hidden_w4<<<hg,256,0,ctx->stream>>>(ctx->gate,ctx->x,dev,I,D,0);
|
||||
grouped_hidden_w4<<<hg,256,0,ctx->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<<<og,256,0,ctx->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<<<hg,256,0,ctx->stream>>>(ctx->gate,ctx->x,dev,I,D,0);
|
||||
grouped_hidden<<<hg,256,0,ctx->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<<<og,256,0,ctx->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<std::mutex> 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<std::mutex> 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<<<H,256,shared,dc->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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
#ifndef COLIBRI_DECODE_BATCH_H
|
||||
#define COLIBRI_DECODE_BATCH_H
|
||||
|
||||
#include <stddef.h>
|
||||
#include <stdio.h>
|
||||
#include <math.h>
|
||||
|
||||
/* `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
|
||||
@@ -27,6 +27,7 @@
|
||||
#include <stdatomic.h> /* PIPE ready-flags/job queue + PILOT_REAL cross-layer handshake */
|
||||
#include <sched.h> /* sched_yield: PIPE spin / PILOT barrier */
|
||||
#include <unistd.h>
|
||||
#include <sys/select.h>
|
||||
#if defined(__APPLE__) || defined(__linux__)
|
||||
#include <sys/resource.h>
|
||||
#include <sys/mman.h> /* 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 <omp.h> /* 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<I){ uint8_t byte=w[i>>1]; int lo=(int)(byte&0xF)-8; a += xs[i]*(float)lo; }
|
||||
y[(int64_t)s*O+o]=a*sc; } }
|
||||
}
|
||||
/* 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<O?z:z-O; const uint8_t *w=(z<O?qg:qu)+(int64_t)o*rb;
|
||||
float a=0; int i=0;
|
||||
#ifdef __AVX2__
|
||||
const __m128i m4=_mm_set1_epi8(0x0F); const __m256i b8=_mm256_set1_epi32(8);
|
||||
__m256 acc=_mm256_setzero_ps();
|
||||
for(;i+16<=I;i+=16){ __m128i by=_mm_loadl_epi64((const __m128i*)(w+(i>>1)));
|
||||
__m128i lo=_mm_and_si128(by,m4),hi=_mm_and_si128(_mm_srli_epi16(by,4),m4);
|
||||
__m128i nib=_mm_unpacklo_epi8(lo,hi);
|
||||
__m256 w0=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nib),b8));
|
||||
__m256 w1=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nib,8)),b8));
|
||||
acc=_mm256_fmadd_ps(_mm256_loadu_ps(x+i),w0,acc);
|
||||
acc=_mm256_fmadd_ps(_mm256_loadu_ps(x+i+8),w1,acc); }
|
||||
a=hsum256(acc);
|
||||
#elif defined(__ARM_NEON)
|
||||
const uint8x8_t m4=vdup_n_u8(0x0F); const int8x8_t b8=vdup_n_s8(8);
|
||||
float32x4_t ac0=vdupq_n_f32(0),ac1=vdupq_n_f32(0);
|
||||
for(;i+16<=I;i+=16){ uint8x8_t by=vld1_u8(w+(i>>1));
|
||||
uint8x8x2_t n=vzip_u8(vand_u8(by,m4),vshr_n_u8(by,4));
|
||||
int16x8_t w0=vmovl_s8(vsub_s8(vreinterpret_s8_u8(n.val[0]),b8));
|
||||
int16x8_t w1=vmovl_s8(vsub_s8(vreinterpret_s8_u8(n.val[1]),b8));
|
||||
ac0=vfmaq_f32(ac0,vld1q_f32(x+i),vcvtq_f32_s32(vmovl_s16(vget_low_s16(w0))));
|
||||
ac1=vfmaq_f32(ac1,vld1q_f32(x+i+4),vcvtq_f32_s32(vmovl_s16(vget_high_s16(w0))));
|
||||
ac0=vfmaq_f32(ac0,vld1q_f32(x+i+8),vcvtq_f32_s32(vmovl_s16(vget_low_s16(w1))));
|
||||
ac1=vfmaq_f32(ac1,vld1q_f32(x+i+12),vcvtq_f32_s32(vmovl_s16(vget_high_s16(w1)))); }
|
||||
a=vaddvq_f32(vaddq_f32(ac0,ac1));
|
||||
#endif
|
||||
for(;i+1<I;i+=2){ uint8_t b=w[i>>1]; a+=x[i]*(float)((b&15)-8)+x[i+1]*(float)((b>>4)-8); }
|
||||
if(i<I) a+=x[i]*(float)((w[i>>1]&15)-8);
|
||||
(z<O?yg:yu)[o]=a*(z<O?sg:su)[o];
|
||||
}
|
||||
}
|
||||
|
||||
static void matmul_qt(float *y,const float *x,QT *w,int S);
|
||||
static void expert_gate_up(float *g,float *u,const float *x,QT *wg,QT *wu,int S){
|
||||
if(S==1&&wg->fmt==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;i<c->n_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 x<y?1:x>y?-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;s<S;s++){
|
||||
const float *xs=x+(int64_t)s*D; int pos=pos_base+s;
|
||||
KVState *ks=kvs?kvs[s]:m->kv;
|
||||
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;h<H;h++) rope_interleave(qfull+(int64_t)h*qh+c->qk_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 && layer<c->n_layers && m->kv_start[layer]==0){
|
||||
if(m->has_dsa && layer<c->n_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;s<S;s++){
|
||||
const float *xs=x+(int64_t)s*D; int pos=pos_base+s;
|
||||
float *kd=m->Ic[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;s<S;s++){
|
||||
int pos=pos_base+s, nk=pos+1;
|
||||
KVState *ks=kvs?kvs[s]:m->kv;
|
||||
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 = nk<dtopk ? nk : dtopk;
|
||||
float *qi=falloc((int64_t)nh*hd);
|
||||
matmul_qt(qi, QR+(int64_t)s*c->q_lora, &m->ix_wq[layer], 1);
|
||||
for(int h=0;h<nh;h++) rope_interleave(qi+(int64_t)h*hd, pos, c);
|
||||
float w32[64];
|
||||
float *w32=falloc(nh);
|
||||
matmul_qt(w32, x+(int64_t)s*D, &m->ix_wp[layer], 1);
|
||||
float wsc=1.f/sqrtf((float)nh), rs=1.f/sqrtf((float)hd);
|
||||
float *isc=falloc(nk);
|
||||
for(int t=0;t<nk;t++){
|
||||
const float *kt=m->Ic[layer]+(int64_t)t*hd;
|
||||
const float *kt=coli_kv_row(ks->Ic[layer],t,hd);
|
||||
float a=0;
|
||||
for(int h=0;h<nh;h++){ const float *qhp=qi+(int64_t)h*hd;
|
||||
float d0=0; for(int i=0;i<hd;i++) d0+=qhp[i]*kt[i];
|
||||
@@ -1424,7 +1518,7 @@ static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_ba
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
|
||||
for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
|
||||
m->dsa_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;s<S&&cuda_core;s++){
|
||||
KVState *ks=kvs?kvs[s]:m->kv;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;s<S;s++) for(int h=0;h<H;h++){
|
||||
int pos=pos_base+s;
|
||||
KVState *ks=kvs?kvs[s]:m->kv;
|
||||
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;d<c->qk_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;jj<nt;jj++){ int t = tlist ? tlist[jj] : st0+jj;
|
||||
const float *Lt=m->Lc[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;i<kvl;i++) a+=qabs[i]*Lt[i];
|
||||
for(int d=0;d<c->qk_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;jj<nt;jj++){ int t = tlist ? tlist[jj] : st0+jj;
|
||||
const float *Lt=m->Lc[layer]+(int64_t)t*kvl;
|
||||
const float *Lt=coli_kv_row(ks->Lc[layer],t,kvl);
|
||||
float a=sc[jj]; for(int i=0;i<kvl;i++) clat[i]+=a*Lt[i]; }
|
||||
qt_matvec_rows(&l->kv_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;s<S;s++) for(int h=0;h<H;h++){
|
||||
int pos=pos_base+s;
|
||||
@@ -1507,11 +1621,16 @@ static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_ba
|
||||
const float *vv=kvb_all+(int64_t)t*kvb_dim+(int64_t)h*(c->qk_nope+vh)+c->qk_nope;
|
||||
float a=sc[jj]; for(int d=0;d<vh;d++) cx[d]+=a*vv[d]; }
|
||||
}
|
||||
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(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;kk<Ke;kk++){
|
||||
m->eusage[layer][idx[kk]]++;
|
||||
if(m->eheat[layer][idx[kk]]<UINT32_MAX) m->eheat[layer][idx[kk]]++;
|
||||
m->elast[layer][idx[kk]]=++m->eaccess_clock;
|
||||
}
|
||||
if(c->norm_topk){ float sm=0; for(int kk=0;kk<Ke;kk++) sm+=w[kk]; sm+=1e-20f; for(int kk=0;kk<Ke;kk++) w[kk]/=sm; }
|
||||
for(int kk=0;kk<Ke;kk++) w[kk]*=c->routed_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;base<nu;base+=64){
|
||||
int nb = nu-base<64 ? nu-base : 64;
|
||||
@@ -1699,6 +1826,9 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){
|
||||
if(!found) expert_prefetch(m,layer,eid);
|
||||
}
|
||||
}
|
||||
#ifdef COLI_CUDA
|
||||
ESlot *group_e[64]; int group_n[64]; int ngroup=0;
|
||||
#endif
|
||||
#ifdef COLI_METAL
|
||||
if(g_metal_enabled){
|
||||
/* PIPE drain. Two reasons this barrier is mandatory here, and not optional:
|
||||
@@ -1746,17 +1876,78 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){
|
||||
if(!nr) continue;
|
||||
#ifdef COLI_CUDA
|
||||
if(g_cuda_enabled && e->g.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;r<nr;r++){ group_row[(int64_t)ngroup*S+r]=rows[r]; group_weight[(int64_t)ngroup*S+r]=rw[r]; }
|
||||
ngroup++; continue;
|
||||
}
|
||||
#endif
|
||||
for(int r=0;r<nr;r++) memcpy(xg+(int64_t)r*D, x+(int64_t)rows[r]*D, D*sizeof(float));
|
||||
double t0=now_s();
|
||||
matmul_qt(gg, xg, &e->g, 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;r<nr;r++){ float *os=out+(int64_t)rows[r]*D,wgt=rw[r],*hr=hh+(int64_t)r*D;
|
||||
for(int d=0;d<D;d++) os[d]+=wgt*hr[d]; }
|
||||
m->t_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;r<nr;r++){ float *os=out+(int64_t)rows[r]*D, wgt=rw[r], *hr=hh+(int64_t)r*D;
|
||||
for(int d=0;d<D;d++) os[d]+=wgt*hr[d]; }
|
||||
m->t_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;di<g_cuda_ndev;di++) for(int q=0;q<ngroup;q++)
|
||||
if(group_e[q]->g.cuda_device==g_cuda_devices[di]) dev_total[di]+=group_n[q];
|
||||
for(int di=1;di<g_cuda_ndev;di++) dev_off[di]=dev_off[di-1]+dev_total[di-1];
|
||||
for(int di=0;di<g_cuda_ndev;di++){
|
||||
int cursor=0,device=g_cuda_devices[di];
|
||||
for(int q=0;q<ngroup;q++) if(group_e[q]->g.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;r<group_n[q];r++) memcpy(group_x+(int64_t)(dev_off[di]+cursor+r)*D,
|
||||
x+(int64_t)group_row[(int64_t)q*S+r]*D,D*sizeof(float));
|
||||
cursor+=group_n[q];
|
||||
}
|
||||
}
|
||||
double tg=now_s();
|
||||
#pragma omp parallel for if(g_cuda_ndev>1) schedule(static)
|
||||
for(int di=0;di<g_cuda_ndev;di++) if(dev_nc[di])
|
||||
dev_ok[di]=coli_cuda_expert_group(dev_g[di],dev_u[di],dev_d[di],dev_rows[di],dev_nc[di],
|
||||
group_y+(int64_t)dev_off[di]*D,group_x+(int64_t)dev_off[di]*D);
|
||||
for(int di=0;di<g_cuda_ndev;di++){
|
||||
int off=dev_off[di];
|
||||
for(int q=0;q<dev_nc[di];q++){
|
||||
int gi=dev_which[di][q],nr=group_n[gi]; ESlot *e=group_e[gi];
|
||||
if(!dev_ok[di]){
|
||||
for(int r=0;r<nr;r++) memcpy(xg+(int64_t)r*D,x+(int64_t)group_row[(int64_t)gi*S+r]*D,D*sizeof(float));
|
||||
if(!coli_cuda_expert_mlp(e->g.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;r<nr;r++){ float *os=out+(int64_t)group_row[(int64_t)gi*S+r]*D,wgt=group_weight[(int64_t)gi*S+r];
|
||||
for(int d=0;d<D;d++) os[d]+=wgt*src[(int64_t)r*D+d]; }
|
||||
off+=nr;
|
||||
}
|
||||
}
|
||||
m->t_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;z<m->enr[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;s<S;s++) rmsnorm(nrm+(int64_t)s*D, x+(int64_t)s*D, l->in_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+1<c->n_layers && m->L[li+1].sparse) pilot_prefetch(m,li+1,x,S);
|
||||
if(g_looka && S==1 && li+1<c->n_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;s<S;s++){
|
||||
if(!rows[s].kv || !rows[s].kv->Lc || !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;l<c->n_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;p<s;p++) if(rows[p].kv==rows[s].kv){ free(x); return NULL; }
|
||||
kvs[s]=rows[s].kv; positions[s]=rows[s].pos;
|
||||
embed_row(m,rows[s].token,x+(int64_t)s*D);
|
||||
}
|
||||
layers_forward_rows(m,x,S,0,kvs,positions);
|
||||
float *norm=falloc((int64_t)S*D);
|
||||
for(int s=0;s<S;s++)
|
||||
rmsnorm(norm+(int64_t)s*D,x+(int64_t)s*D,m->final_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;i<nfull-1;i++){
|
||||
logit=step(m,full+i,1,i); free(logit); steps++;
|
||||
@@ -2539,7 +2780,8 @@ static int repin_pick(Model *m, RepinCand *out, int maxc){
|
||||
ESlot *P=m->pin[l]; int ids[4096], zp, eu; long g;
|
||||
int np=m->npin[l]; if(np>4096) np=4096;
|
||||
for(int z=0;z<np;z++) ids[z]=P[z].eid;
|
||||
if(!tier_pick_swap(m->eheat[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(nb<maxc){ out[nb].gain=g; out[nb].l=l; out[nb].slot=zp; out[nb].eid=eu; nb++; }
|
||||
else { int w=0; for(int b=1;b<maxc;b++) if(out[b].gain<out[w].gain) w=b;
|
||||
if(g>out[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;i++) if(req[i].active && req[i].id==id){
|
||||
req[i].active=0; kv_bind(m,&ctx[i].kv);
|
||||
kv_disk_append(m,ctx[i].hist,ctx[i].len);
|
||||
printf("ERROR %llu CANCELLED\n",id); fflush(stdout); free(line); return 0;
|
||||
}
|
||||
printf("ERROR %llu NOT_FOUND\n",id); fflush(stdout); free(line); return 0;
|
||||
}
|
||||
ColiSubmit sub; int valid=coli_submit_parse(line,&sub);
|
||||
if(!valid){ printf("ERROR 0 BAD_REQUEST\n"); fflush(stdout); free(line); return 0; }
|
||||
char *raw=malloc((size_t)sub.bytes+1);
|
||||
if(!raw){ fprintf(stderr,"OOM multiplex payload\n"); exit(1); }
|
||||
if(fread(raw,1,(size_t)sub.bytes,stdin)!=(size_t)sub.bytes){ free(raw); free(line); return -1; }
|
||||
int delim=fgetc(stdin);
|
||||
if(delim!='\n'){
|
||||
printf("ERROR %llu BAD_FRAME\n",sub.id); fflush(stdout);
|
||||
free(raw); free(line); return -1;
|
||||
}
|
||||
raw[sub.bytes]=0;
|
||||
if(sub.slot>=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;i<nctx;i++) if(req[i].active && req[i].id==sub.id){
|
||||
printf("ERROR %llu DUPLICATE_ID\n",sub.id); fflush(stdout); free(raw); free(line); return 0;
|
||||
}
|
||||
ServeCtx *sc=&ctx[sub.slot]; kv_bind(m,&sc->kv);
|
||||
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(prefix<sc->len && prefix<nt && sc->hist[prefix]==tmp[prefix]) prefix++;
|
||||
if(prefix<sc->len){ 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;i<nctx;i++) serve_ctx_init(m,&ctx[i],snap,i,maxctx);
|
||||
setvbuf(stdin,NULL,_IONBF,0);
|
||||
printf("\x01\x01READY\x01\x01\nSTAT 0 0.00 0.0 %.2f\n",rss_gb()); fflush(stdout);
|
||||
int eof=0;
|
||||
for(;;){
|
||||
int active=0; for(int i=0;i<nctx;i++) active+=req[i].active;
|
||||
fd_set rfds; FD_ZERO(&rfds); FD_SET(STDIN_FILENO,&rfds);
|
||||
struct timeval tv={0,0}, *ptv=active?&tv:NULL;
|
||||
int ready=eof?0:select(STDIN_FILENO+1,&rfds,NULL,NULL,ptv);
|
||||
if(ready>0 && FD_ISSET(STDIN_FILENO,&rfds)) if(mux_submit(m,&T,ctx,req,nctx,maxctx,eos)<0) eof=1;
|
||||
active=0; for(int i=0;i<nctx;i++) active+=req[i].active;
|
||||
if(!active){ if(eof) break; continue; }
|
||||
DecodeRow rows[16]; int slots[16], S=0;
|
||||
for(int i=0;i<nctx;i++) if(req[i].active){
|
||||
rows[S]=(DecodeRow){&ctx[i].kv,req[i].pending,ctx[i].len}; slots[S++]=i;
|
||||
}
|
||||
float *lo=step_decode_batch(m,rows,S); if(!lo){fprintf(stderr,"decode batch failed\n");break;}
|
||||
m->n_fw++;
|
||||
for(int s=0;s<S;s++){
|
||||
int i=slots[s]; ServeCtx *sc=&ctx[i]; ServeReq *r=&req[i];
|
||||
sc->len++; 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;i<nctx;i++) serve_ctx_free(m,&ctx[i]); free(ctx); free(req);
|
||||
m->kv=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->c<y->c?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<cap && fscanf(f,"%d %d %u",&l,&e,&cnt)==3){
|
||||
int ok = l>=0 && e>=0 && e<c->n_experts &&
|
||||
((l<c->n_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;a<n;a++){ int best=a; /* selection sort parziale, poi taglio */
|
||||
for(int b=a+1;b<n;b++) if(r[b].c>r[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=(li<c->n_layers&&m->L[li].sparse)||(li==c->n_layers&&m->has_mtp);
|
||||
if(sparse) for(int ei=0;ei<c->n_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;a<npin;a++) cnt_l[r[a].l]++;
|
||||
for(int i=0;i<=c->n_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;a<npin;a++) slot_of[a]=next[r[a].l]++;
|
||||
for(int i=0;i<=c->n_layers;i++) m->npin[i]=cnt_l[i];
|
||||
double t0=now_s();
|
||||
#pragma omp parallel for schedule(dynamic,1)
|
||||
for(int a=0;a<npin;a++){
|
||||
int li=r[a].l, slot;
|
||||
#pragma omp critical
|
||||
slot=m->npin[li]++;
|
||||
expert_load(m,li,r[a].e,&m->pin[li][slot],1);
|
||||
}
|
||||
m->resident_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);
|
||||
#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};
|
||||
int placed_n[COLI_CUDA_MAX_DEVICES]={0}, gpu_prefix=0;
|
||||
double budget=g_cuda_expert_gb*1e9, safe_total=0;
|
||||
for(int i=0;i<g_cuda_ndev;i++){
|
||||
if(g_cuda_enabled&&g_cuda_expert_gb>0) for(int i=0;i<g_cuda_ndev;i++){
|
||||
size_t free_b=0,total_b=0;
|
||||
if(coli_cuda_mem_info(g_cuda_devices[i],&free_b,&total_b)){
|
||||
/* Dense tensors are assigned round-robin and upload lazily.
|
||||
* Reserve their projected footprint plus 2 GB per device. */
|
||||
remaining[i]=(double)free_b-(double)g_cuda_dense_projected[i]-2e9;
|
||||
if(remaining[i]<0) remaining[i]=0;
|
||||
safe_total+=remaining[i];
|
||||
if(remaining[i]<0) remaining[i]=0; safe_total+=remaining[i];
|
||||
}
|
||||
}
|
||||
if(budget>safe_total) budget=safe_total;
|
||||
for(int a=0;a<npin && m->gpu_expert_bytes<budget;a++){
|
||||
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){
|
||||
int gpu_limit=gpu_prefix?gpu_prefix:npin;
|
||||
for(int a=0;a<gpu_limit && m->gpu_expert_bytes<budget;a++){
|
||||
int li=r[a].l;
|
||||
for(int z=0;z<m->npin[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_prefix<npin){
|
||||
#pragma omp parallel for schedule(dynamic,1)
|
||||
for(int a=gpu_prefix;a<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)(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")){
|
||||
|
||||
+177
-31
@@ -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")
|
||||
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")
|
||||
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)
|
||||
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 stats
|
||||
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,19 +903,24 @@ 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()
|
||||
if runtime is not None:
|
||||
runtime.close()
|
||||
|
||||
|
||||
|
||||
+71
-11
@@ -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)
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
#include "../backend_cuda.h"
|
||||
|
||||
#include <chrono>
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <vector>
|
||||
|
||||
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<iterations;i++)if(!coli_cuda_expert_group(gs,us,ds,rs,1,y,x))std::exit(2);
|
||||
auto end=std::chrono::steady_clock::now();
|
||||
return std::chrono::duration<double,std::milli>(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<unsigned char> hidden((size_t)I*D/2),down((size_t)D*I/2);
|
||||
std::vector<float> 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;i<hidden.size();i++)hidden[i]=(unsigned char)((i*17+29)&255);
|
||||
for(size_t i=0;i<down.size();i++)down[i]=(unsigned char)((i*13+41)&255);
|
||||
for(int i=0;i<I;i++)hs[i]=0.006f+(i%11)*0.0002f;
|
||||
for(int i=0;i<D;i++)ds[i]=0.006f+(i%7)*0.0002f;
|
||||
for(size_t i=0;i<x.size();i++)x[i]=std::sin((float)(i+1)*0.013f)*2.f;
|
||||
ColiCudaTensor *g=nullptr,*u=nullptr,*d=nullptr;
|
||||
if(!coli_cuda_tensor_upload(&g,hidden.data(),hs.data(),2,D,I,device)||
|
||||
!coli_cuda_tensor_upload(&u,hidden.data(),hs.data(),2,D,I,device)||
|
||||
!coli_cuda_tensor_upload(&d,down.data(),ds.data(),2,I,D,device))return 2;
|
||||
for(int rows: {1,2,4,8}){
|
||||
double scalar=run(g,u,d,x.data(),a.data(),rows,3,0);
|
||||
double packed=run(g,u,d,x.data(),b.data(),rows,3,1);
|
||||
double tc=run(g,u,d,x.data(),c.data(),rows,3,2);
|
||||
double pe=0,te=0,ref=0;for(int i=0;i<rows*D;i++){double p=b[i]-a[i],t=c[i]-a[i];pe+=p*p;te+=t*t;ref+=(double)a[i]*a[i];}
|
||||
std::printf("rows=%d scalar_ms=%.3f packed_ms=%.3f packed_speedup=%.3fx packed_rms=%.7f tensor_ms=%.3f tensor_speedup=%.3fx tensor_rms=%.5f\n",
|
||||
rows,scalar,packed,scalar/packed,std::sqrt(pe/(ref+1e-20)),tc,scalar/tc,std::sqrt(te/(ref+1e-20)));
|
||||
}
|
||||
coli_cuda_tensor_free(g);coli_cuda_tensor_free(u);coli_cuda_tensor_free(d);coli_cuda_shutdown();
|
||||
}
|
||||
@@ -15,6 +15,12 @@ static int close_enough(const float *got, const float *want, int n) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
static int relative_rms(const float *got,const float *want,int n,float limit){
|
||||
double err=0,ref=0; for(int i=0;i<n;i++){double d=got[i]-want[i];err+=d*d;ref+=(double)want[i]*want[i];}
|
||||
float r=(float)std::sqrt(err/(ref+1e-20));
|
||||
if(r>limit){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();
|
||||
|
||||
Executable
BIN
Binary file not shown.
@@ -0,0 +1,56 @@
|
||||
#include <assert.h>
|
||||
#include <stdio.h>
|
||||
|
||||
#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;
|
||||
}
|
||||
@@ -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):
|
||||
|
||||
Executable
BIN
Binary file not shown.
Executable
BIN
Binary file not shown.
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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;z<npin;z++)
|
||||
if(tier_lfru_score(heat[pinned[z]],last[pinned[z]],clock)<
|
||||
tier_lfru_score(heat[pinned[cold]],last[pinned[cold]],clock)) cold=z;
|
||||
int hot=-1; uint64_t hs=0;
|
||||
for(int e=0;e<nexpert;e++){
|
||||
int resident=0; for(int z=0;z<npin;z++) if(pinned[z]==e){resident=1;break;}
|
||||
uint64_t score=tier_lfru_score(heat[e],last[e],clock);
|
||||
if(!resident&&(hot<0||score>hs)){ 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<nexpert;e++) heat[e]>>=1;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user