danielkrause/colibri-strix
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Isolated sequence KV contexts: KVState extraction, KV_SLOTS serve contexts with per-slot persistence, budget-aware pool accounting (#29)

Browse Source 
* Add isolated sequence KV contexts

* Log projected multi-slot KV footprint
This commit is contained in:
ZacharyZcR
2026-07-10 22:19:43 +08:00
committed by GitHub
parent 76e858a80f
commit 9e6e1ac327
5 changed files with 192 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/coli
+3 -1
View File
@@ -440,7 +440,8 @@ def cmd_serve(a):
need_model(a.model)
from openai_server import serve
serve(a.model, a.host, a.port, a.model_id, a.api_key,
a.cap,a.ngen,GLM,env_for(a),a.cors_origin,a.max_queue,a.queue_timeout)
a.cap,a.ngen,GLM,env_for(a),a.cors_origin,
a.max_queue,a.queue_timeout,a.kv_slots)
def cmd_bench(a):
need_model(a.model)
@@ -506,6 +507,7 @@ def main():
ps.add_argument("--cors-origin",action="append",default=None)
ps.add_argument("--max-queue",type=int,default=int(os.environ.get("COLI_MAX_QUEUE","8")))
ps.add_argument("--queue-timeout",type=float,default=float(os.environ.get("COLI_QUEUE_TIMEOUT","300")))
ps.add_argument("--kv-slots",type=int,default=int(os.environ.get("COLI_KV_SLOTS","1")))
pb=sub.add_parser("bench", parents=[common]); pb.add_argument("tasks", nargs="*")
pb.add_argument("--limit",type=int,default=40); pb.add_argument("--data",default=os.path.join(HERE,"bench"))
pc=sub.add_parser("convert", parents=[common]); pc.add_argument("--repo",default="zai-org/GLM-5.2-FP8")
c/glm.c
+114 -45
View File
@@ -100,6 +100,13 @@ typedef struct {
typedef struct { int eid; QT g,u,d; uint8_t *slab; float *fslab;
int64_t slab_cap, fslab_cap; uint64_t used; } ESlot;
typedef struct {
float **Lc, **Rc, **Ic;
int *kv_start, max_t;
int disk_nrec;
char disk_path[2048];
} KVState;
typedef struct {
Cfg c; shards S;
int ebits, dbits; /* bit expert / bit densa */
@@ -108,8 +115,9 @@ typedef struct {
/* KV-cache MLA COMPRESSA: per token si tiene solo il latente normato [kv_lora] e
* k_rot [qk_rope] (576 vs 32768 valori/token). k_nope e value si ricostruiscono al
* volo con kv_b. E' cio' che rende gestibile il contesto su 15 GB (64 teste, no GQA). */
float **Lc, **Rc; int max_t;
float **Lc, **Rc; int max_t; /* alias della KVState attiva */
int *kv_start; /* prima pos valida nella KV del layer (MTP: parziale) */
KVState *kv;
ESlot **ecache; int *ecn; int ecap; /* LRU expert per-layer */
ESlot ws[64]; /* working set del layer corrente (load paralleli) */
ESlot **pin; int *npin; /* HOT-STORE: expert pinnati in RAM (mai evicted) */
@@ -119,7 +127,7 @@ typedef struct {
int has_dsa;
QT *ix_wq, *ix_wk, *ix_wp; /* per layer FULL: wq_b, wk, weights_proj */
float **ix_knw, **ix_knb; /* k_norm (LayerNorm, eps 1e-6) */
float **Ic; /* cache k dell'indexer per layer full [max_t*hd] */
float **Ic; /* alias KVState: cache indexer [max_t*hd] */
int *dsa_sel, *dsa_nsel; int dsa_scap; /* selezione per posizione del batch corrente */
/* testa MTP (layer n_layers, stile DeepSeek-V3): draft nativi ad alta acceptance */
int has_mtp; Layer mtpL; QT eh_proj;
@@ -711,7 +719,8 @@ 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->kv_start=calloc(NR,sizeof(int));
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++){
Layer *l=&m->L[i];
#define P(s) (snprintf(nm,sizeof(nm),"model.layers.%d." s,i),nm)
@@ -1366,17 +1375,24 @@ static void layers_forward(Model *m, float *x, int S, int pos_base){
static void kv_alloc(Model *m, int max_t){
Cfg *c=&m->c;
if(m->Lc){ for(int i=0;i<c->n_layers+1;i++){ free(m->Lc[i]); free(m->Rc[i]); } free(m->Lc); free(m->Rc); }
if(m->Ic){ for(int i=0;i<c->n_layers;i++) free(m->Ic[i]); free(m->Ic); m->Ic=NULL; }
KVState *k=m->kv;
if(k->Lc){ for(int i=0;i<c->n_layers+1;i++){ free(k->Lc[i]); free(k->Rc[i]); } free(k->Lc); free(k->Rc); }
if(k->Ic){ for(int i=0;i<c->n_layers;i++) free(k->Ic[i]); free(k->Ic); k->Ic=NULL; }
if(m->has_dsa){
m->Ic=calloc(c->n_layers,sizeof(float*));
for(int i=0;i<c->n_layers;i++) if(c->idx_type[i]) m->Ic[i]=falloc((int64_t)max_t*c->index_hd);
k->Ic=calloc(c->n_layers,sizeof(float*));
for(int i=0;i<c->n_layers;i++) if(c->idx_type[i]) k->Ic[i]=falloc((int64_t)max_t*c->index_hd);
}
m->max_t=max_t;
k->max_t=max_t;
int NR=c->n_layers+1; /* riga extra: KV del layer MTP */
m->Lc=calloc(NR,sizeof(float*)); m->Rc=calloc(NR,sizeof(float*));
for(int i=0;i<NR;i++){ m->Lc[i]=falloc((int64_t)max_t*c->kv_lora);
m->Rc[i]=falloc((int64_t)max_t*c->qk_rope); }
k->Lc=calloc(NR,sizeof(float*)); k->Rc=calloc(NR,sizeof(float*));
for(int i=0;i<NR;i++){ k->Lc[i]=falloc((int64_t)max_t*c->kv_lora);
k->Rc[i]=falloc((int64_t)max_t*c->qk_rope); }
m->Lc=k->Lc; m->Rc=k->Rc; m->Ic=k->Ic; m->max_t=k->max_t; m->kv_start=k->kv_start;
}
static void kv_bind(Model *m, KVState *k){
m->kv=k; m->Lc=k->Lc; m->Rc=k->Rc; m->Ic=k->Ic;
m->max_t=k->max_t; m->kv_start=k->kv_start;
}
static void mtp_absorb(Model *m, const int *next_ids, const float *x, int S, int pos_base);
@@ -1844,7 +1860,7 @@ static void repin_pass(Model *m){
* si appendono SOLO le posizioni nuove e si riscrive nrec per ultimo: un crash a meta'
* append lascia nrec vecchio = file coerente. La riga KV del layer MTP non si salva:
* al resume kv_start=-1 e la finestra di draft riparte da sola. */
static char g_kv_path[2048]; static int g_kvsave=1; static int g_kv_nrec=0;
static int g_kvsave=1;
#define KV_MAGIC "COLIKV1\0"
static void kv_hdr(Model *m, int32_t *h, int nrec){
Cfg *c=&m->c; int nic=0;
@@ -1852,24 +1868,26 @@ static void kv_hdr(Model *m, int32_t *h, int nrec){
h[0]=c->n_layers; h[1]=c->kv_lora; h[2]=c->qk_rope;
h[3]=m->has_dsa?c->index_hd:0; h[4]=nic; h[5]=c->vocab; h[6]=nrec; h[7]=0;
}
static void kv_disk_truncate(int nrec){
static void kv_disk_truncate(Model *m, int nrec){
if(!g_kvsave) return;
FILE *f=fopen(g_kv_path,"r+b");
if(!f){ g_kv_nrec=0; return; }
g_kv_nrec=nrec;
KVState *k=m->kv;
FILE *f=fopen(k->disk_path,"r+b");
if(!f){ k->disk_nrec=0; return; }
k->disk_nrec=nrec;
int32_t nr=nrec; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); fclose(f);
}
static void kv_disk_reset(void){ kv_disk_truncate(0); }
static void kv_disk_reset(Model *m){ kv_disk_truncate(m,0); }
static void kv_disk_append(Model *m, const int *hist, int len){
if(!g_kvsave || len<=g_kv_nrec) return;
KVState *k=m->kv;
if(!g_kvsave || len<=k->disk_nrec) return;
Cfg *c=&m->c;
FILE *f=fopen(g_kv_path,"r+b");
if(!f){ f=fopen(g_kv_path,"wb"); if(!f) return;
FILE *f=fopen(k->disk_path,"r+b");
if(!f){ f=fopen(k->disk_path,"wb"); if(!f) return;
int32_t h[8]; kv_hdr(m,h,0); fwrite(KV_MAGIC,1,8,f); fwrite(h,4,8,f); }
int64_t rec = 4 + (int64_t)c->n_layers*(c->kv_lora+c->qk_rope)*4;
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i]) rec+=(int64_t)c->index_hd*4;
fseek(f, 8+8*4 + (int64_t)g_kv_nrec*rec, SEEK_SET);
for(int p=g_kv_nrec;p<len;p++){
fseek(f, 8+8*4 + (int64_t)k->disk_nrec*rec, SEEK_SET);
for(int p=k->disk_nrec;p<len;p++){
int32_t tk=hist[p]; fwrite(&tk,4,1,f);
for(int i=0;i<c->n_layers;i++){
fwrite(m->Lc[i]+(int64_t)p*c->kv_lora, 4, c->kv_lora, f);
@@ -1880,12 +1898,13 @@ static void kv_disk_append(Model *m, const int *hist, int len){
}
fflush(f); /* dati prima, contatore poi */
int32_t nr=len; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); fclose(f);
g_kv_nrec=len;
k->disk_nrec=len;
}
static int kv_disk_load(Model *m, int *hist, int maxctx){
if(!g_kvsave) return 0;
KVState *k=m->kv;
Cfg *c=&m->c;
FILE *f=fopen(g_kv_path,"rb"); if(!f) return 0;
FILE *f=fopen(k->disk_path,"rb"); if(!f) return 0;
char mg[8]; int32_t h[8], w[8]; kv_hdr(m,w,0);
if(fread(mg,1,8,f)!=8 || memcmp(mg,KV_MAGIC,8) || fread(h,4,8,f)!=8 ||
h[0]!=w[0]||h[1]!=w[1]||h[2]!=w[2]||h[3]!=w[3]||h[4]!=w[4]||h[5]!=w[5]){
@@ -1912,10 +1931,30 @@ out:
fprintf(stderr,"[KV] conversazione ripresa dal disco: %d token in %.1fs (niente re-prefill)\n",
nrec, now_s()-t0);
}
g_kv_nrec=nrec;
k->disk_nrec=nrec;
return nrec;
}
typedef struct { KVState kv; int *hist, len, first; } ServeCtx;
static double kv_pool_bytes(Model *m, int max_ctx);
static void serve_ctx_init(Model *m, ServeCtx *s, const char *snap, int slot, int maxctx){
s->kv.kv_start=calloc(m->c.n_layers+1,sizeof(int));
if(m->has_mtp) s->kv.kv_start[m->c.n_layers]=-1;
kv_bind(m,&s->kv); kv_alloc(m,maxctx);
s->hist=malloc(maxctx*sizeof(int)); s->first=1;
if(slot==0) snprintf(s->kv.disk_path,sizeof(s->kv.disk_path),"%s/.coli_kv",snap);
else snprintf(s->kv.disk_path,sizeof(s->kv.disk_path),"%s/.coli_kv.%d",snap,slot);
s->len=kv_disk_load(m,s->hist,maxctx); if(s->len>0) s->first=0;
}
static void serve_ctx_free(Model *m, ServeCtx *s){
KVState *k=&s->kv; int NR=m->c.n_layers+1;
if(k->Lc) for(int i=0;i<NR;i++){ free(k->Lc[i]); free(k->Rc[i]); }
if(k->Ic) for(int i=0;i<m->c.n_layers;i++) free(k->Ic[i]);
free(k->Lc); free(k->Rc); free(k->Ic); free(k->kv_start); free(s->hist);
}
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);
@@ -1926,19 +1965,24 @@ static void run_serve(Model *m, const char *snap){
int ngen=getenv("NGEN")?atoi(getenv("NGEN")):256;
int maxctx=getenv("CTX")?atoi(getenv("CTX")):4096;
int templ=getenv("CHAT_TEMPLATE")?atoi(getenv("CHAT_TEMPLATE")):1;
kv_alloc(m,maxctx);
int len=0, first=1; /* len = contesto gia' in KV (persiste tra turni) */
int *hist=malloc(maxctx*sizeof(int)); /* storia token (= contenuto della KV): serve
* al lookup n-gram e resta allineata a len */
g_kvsave = getenv("KVSAVE")?atoi(getenv("KVSAVE")):1;
snprintf(g_kv_path,sizeof(g_kv_path),"%s/.coli_kv",snap);
{ int r=kv_disk_load(m,hist,maxctx); if(r>0){ len=r; first=0; } }
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); }
KVState *initial=m->kv; free(initial->kv_start); free(initial);
ServeCtx *ctx=calloc(nctx,sizeof(ServeCtx));
for(int i=0;i<nctx;i++) serve_ctx_init(m,&ctx[i],snap,i,maxctx);
int active=0; ServeCtx *sc=&ctx[0]; kv_bind(m,&sc->kv);
fprintf(stderr,"[KV] context slots: %d x %d token, projected pool %.2f GB\n",
nctx,maxctx,kv_pool_bytes(m,maxctx)/1e9);
#define hist (sc->hist)
#define len (sc->len)
#define first (sc->first)
char *line=NULL; size_t cap=0; ssize_t nr; char *buf=malloc(1<<16);
printf("\x01\x01" "READY" "\x01\x01\n"); printf("STAT 0 0.00 0.0 %.2f\n", rss_gb()); fflush(stdout);
while((nr=getline(&line,&cap,stdin))>0){
if(nr>0 && line[nr-1]=='\n') line[--nr]=0;
if(!strcmp(line,"\x02RESET")){ len=0; first=1; if(m->has_mtp) m->kv_start[m->c.n_layers]=-1;
kv_disk_reset();
kv_disk_reset(m);
printf("\x01\x01" "END" "\x01\x01\n"); printf("STAT 0 0.00 0.0 %.2f\n", rss_gb()); fflush(stdout); continue; }
if(!strcmp(line,"\x02MORE")){ /* continua la risposta troncata da NGEN:
la storia e' gia' in KV, basta ri-forwardare l'ULTIMO token per riavere i logits */
@@ -1960,16 +2004,18 @@ static void run_serve(Model *m, const char *snap){
* line protocol this accepts newlines. The tokenized prompt is matched
* against hist so the common KV prefix survives stateless HTTP turns.
* Per-request generation controls follow the byte count:
* \x02PROMPT <bytes> <max_tokens> <temperature> <top_p>\n<prompt>\n */
* \x02PROMPT <bytes> <max_tokens> <temperature> <top_p> [kv_slot]\n<prompt>\n */
char *raw=NULL, *input=line;
int input_n=(int)nr, raw_mode=0, req_ngen=ngen, prompt_tokens=0;
float base_temp=g_temp, base_nuc=g_nuc;
if(!strncmp(line,"\x02PROMPT ",8)){
unsigned long long nb=0; double rt=0, rp=0;
if(sscanf(line+8,"%llu %d %lf %lf",&nb,&req_ngen,&rt,&rp)!=4 ||
nb>(16u<<20) || req_ngen<1 || rt<0 || rt>2 || rp<=0 || rp>1){
unsigned long long nb=0; double rt=0, rp=0; int slot=0;
int nf=sscanf(line+8,"%llu %d %lf %lf %d",&nb,&req_ngen,&rt,&rp,&slot);
if(nf<4 || nb>(16u<<20) || req_ngen<1 || rt<0 || rt>2 || rp<=0 || rp>1 ||
slot<0 || slot>=nctx){
printf("\x01\x01" "END" "\x01\x01\n"); printf("STAT 0 0.00 0.0 %.2f 0 0\n",rss_gb()); fflush(stdout); continue;
}
active=slot; sc=&ctx[active]; kv_bind(m,&sc->kv);
raw=malloc((size_t)nb+1); if(!raw){fprintf(stderr,"OOM raw prompt\n");exit(1);}
if(fread(raw,1,(size_t)nb,stdin)!=(size_t)nb){free(raw);break;}
int delim=fgetc(stdin); if(delim!='\n' && delim!=EOF) ungetc(delim,stdin);
@@ -1978,7 +2024,7 @@ static void run_serve(Model *m, const char *snap){
raw[nb]=0; input=raw; input_n=(int)nb; raw_mode=1;
if(req_ngen>ngen) req_ngen=ngen;
g_temp=(float)rt; g_nuc=(float)rp;
}
} else { active=0; sc=&ctx[0]; kv_bind(m,&sc->kv); }
int bl=0, k=0; /* costruisce/tokenizza il turno */
/* template UFFICIALE GLM-5.2 (chat_template.jinja): niente \n dopo i ruoli, e dopo
* <|assistant|> serve SEMPRE il blocco think — <think></think> lo DISATTIVA (nothink):
@@ -1992,18 +2038,19 @@ static void run_serve(Model *m, const char *snap){
if(prefix<old_len){
len=prefix;
if(m->has_mtp) m->kv_start[m->c.n_layers]=-1;
kv_disk_truncate(len); /* il prossimo append sovrascrive solo la coda */
kv_disk_truncate(m,len); /* il prossimo append sovrascrive solo la coda */
}
k=prompt_tokens-len;
if(k>0) memcpy(hist+len,tmp+len,k*sizeof(int));
fprintf(stderr,"[API] KV prefix %d/%d token, prefill %d\n",len,prompt_tokens,k);
fprintf(stderr,"[API] KV slot %d prefix %d/%d token, prefill %d\n",
active,len,prompt_tokens,k);
free(tmp);
} else {
if(templ){ if(first) bl+=snprintf(buf+bl,(1<<16)-bl,"[gMASK]<sop>");
bl+=snprintf(buf+bl,(1<<16)-bl,"<|user|>%s<|assistant|>%s",input,tk); }
else bl+=snprintf(buf+bl,(1<<16)-bl,"%s",input);
k=tok_encode(&T,buf,bl,hist+len,maxctx-len); prompt_tokens=k;
if(len+k+8+g_draft>=maxctx){ len=0; first=1; kv_disk_reset();
if(len+k+8+g_draft>=maxctx){ len=0; first=1; kv_disk_reset(m);
bl=0; if(templ){ bl+=snprintf(buf+bl,(1<<16)-bl,"[gMASK]<sop><|user|>%s<|assistant|>%s",input,tk); }
else bl+=snprintf(buf+bl,(1<<16)-bl,"%s",input);
k=tok_encode(&T,buf,bl,hist,maxctx); if(k>maxctx-8-g_draft) k=maxctx-8-g_draft;
@@ -2032,8 +2079,13 @@ static void run_serve(Model *m, const char *snap){
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 */
}
free(line); free(hist); free(buf);
free(line); free(buf);
usage_save(m);
#undef hist
#undef len
#undef first
for(int i=0;i<nctx;i++) serve_ctx_free(m,&ctx[i]);
free(ctx); m->kv=NULL; m->Lc=m->Rc=m->Ic=NULL; m->kv_start=NULL; m->max_t=0;
}
static int *read_arr(jval*o,const char*k,int*n){ jval*a=json_get(o,k); int*r=malloc(a->len*sizeof(int));
@@ -2248,12 +2300,25 @@ static double mem_available_gb(void){
#endif
}
static int kv_slot_count(void){
if(!getenv("SERVE")) return 1;
return getenv("KV_SLOTS")?atoi(getenv("KV_SLOTS")):1;
}
static double kv_pool_bytes(Model *m, int max_ctx){
Cfg *c=&m->c; double one=(double)(c->n_layers+1)*max_ctx*(c->kv_lora+c->qk_rope)*4.0;
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(c->idx_type[i])
one+=(double)max_ctx*c->index_hd*4.0;
int slots=kv_slot_count(); if(slots<1||slots>16) slots=1;
return one*slots;
}
/* byte disponibili per gli expert (pin + LRU) nel budget — specchio del conto di cap_for_ram */
static double expert_avail(Model *m, double ram_gb, int ebits, int max_ctx){
Cfg *c=&m->c; int64_t eb=expert_bytes_probe(m,ebits);
if(ram_gb<=0){ ram_gb=g_mem_avail_boot*0.88; if(ram_gb<4) ram_gb=8; }
double slack = 1.2e9 + 2.5e9 + 64.0*(double)eb
+ (double)(c->n_layers+1)*max_ctx*(c->kv_lora+c->qk_rope)*4.0
+ kv_pool_bytes(m,max_ctx)
+ (double)max_ctx*c->n_heads*(c->qk_nope+c->v_head)*4.0;
return ram_gb*1e9 - (double)m->resident_bytes - slack;
}
@@ -2274,7 +2339,7 @@ static void cap_for_ram(Model *m, double ram_gb, int ebits, int max_ctx){
* KV cache a max_ctx, kvb_all della ricostruzione k/v in attention,
* attivazioni+logits+overhead ~1.2 GB */
double ws_b = 64.0*(double)eb;
double kv_b = (double)(c->n_layers+1)*max_ctx*(c->kv_lora+c->qk_rope)*4.0;
double kv_b = kv_pool_bytes(m,max_ctx);
double kvb_b = (double)max_ctx*c->n_heads*(c->qk_nope+c->v_head)*4.0;
/* RISERVA PAGE-CACHE (misurato 2026-07-06): strangolarla fa crollare le pread
* buffered da ~800 a ~180 MB/s — gli ultimi GB di LRU rendono MENO di quanto
@@ -2285,9 +2350,10 @@ static void cap_for_ram(Model *m, double ram_gb, int ebits, int max_ctx){
int capmax = (avail>0 && nsp>0) ? (int)(avail/((double)nsp*eb)) : 0;
if(capmax<1) capmax=1;
if(capmax < m->ecap){
fprintf(stderr,"[RAM_GB=%.1f%s] residente %.1f GB + slack %.1f GB (ws %.1f, KV@%d %.1f, kvb %.1f), "
fprintf(stderr,"[RAM_GB=%.1f%s] residente %.1f GB + slack %.1f GB (ws %.1f, KV %dx%d %.1f, kvb %.1f), "
"expert %.1f MB x %d layer -> cap abbassato %d->%d (proiezione picco %.1f GB)\n",
ram_gb, auto_b?" auto":"", m->resident_bytes/1e9, slack/1e9, ws_b/1e9, max_ctx, kv_b/1e9, kvb_b/1e9,
ram_gb,auto_b?" auto":"",m->resident_bytes/1e9,slack/1e9,ws_b/1e9,
kv_slot_count(),max_ctx,kv_b/1e9,kvb_b/1e9,
eb/1e6, nsp, m->ecap, capmax,
(m->resident_bytes + (double)capmax*nsp*eb + slack)/1e9);
m->ecap=capmax;
@@ -2344,6 +2410,9 @@ int main(int argc, char **argv){
int cap = argc>1?atoi(argv[1]):64;
int ebits= argc>2?atoi(argv[2]):8;
int dbits= argc>3?atoi(argv[3]):ebits;
if(getenv("SERVE") && (kv_slot_count()<1 || kv_slot_count()>16)){
fprintf(stderr,"KV_SLOTS deve essere tra 1 e 16\n"); return 2;
}
#ifdef COLI_CUDA
if(getenv("COLI_CUDA") && atoi(getenv("COLI_CUDA"))){
const char *one=getenv("COLI_GPU"), *many=getenv("COLI_GPUS");
c/openai_server.py
+30 -12
View File
@@ -247,16 +247,20 @@ def read_engine_turn(stream, sentinel, on_bytes):
class Engine:
def __init__(self, executable, model, cap=8, max_tokens=1024, env=None):
child_env = dict(env or os.environ, SNAP=str(model), SERVE="1", NGEN=str(max_tokens))
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))
self.process = subprocess.Popen(
[str(executable), str(cap)], env=child_env, stdin=subprocess.PIPE,
stdout=subprocess.PIPE, bufsize=0,
)
self.lock = threading.Lock()
self.kv_slots = kv_slots
read_engine_turn(self.process.stdout, READY, lambda _: None)
def generate(self, prompt, max_tokens, temperature, top_p, on_text):
def generate(self, prompt, max_tokens, temperature, top_p, on_text, cache_slot=0):
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")
if b"\0" in payload:
raise APIError(400, "NUL bytes are not supported in prompts.", "messages")
@@ -270,7 +274,8 @@ class Engine:
with self.lock:
if self.process.poll() is not None:
raise RuntimeError("colibri engine is not running")
header = f"\x02PROMPT {len(payload)} {max_tokens} {temperature:.8g} {top_p:.8g}\n".encode()
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)
@@ -296,13 +301,15 @@ class APIServer(ThreadingHTTPServer):
daemon_threads = True
def __init__(self, address, engine, model_id, api_key=None, max_tokens=1024,
cors_origins=DEFAULT_CORS_ORIGINS, max_queue=8, queue_timeout=300):
cors_origins=DEFAULT_CORS_ORIGINS, max_queue=8, queue_timeout=300,
kv_slots=1):
super().__init__(address, APIHandler)
self.engine = engine
self.model_id = model_id
self.api_key = api_key
self.max_tokens = max_tokens
self.scheduler = GenerationScheduler(max_queue, queue_timeout)
self.kv_slots = kv_slots
self.cors_origins = tuple(cors_origins)
self.created = int(time.time())
@@ -370,7 +377,8 @@ class APIHandler(BaseHTTPRequestHandler):
try:
path = urlsplit(self.path).path
if path == "/health":
self.send_json(200, {"status": "ok", "scheduler": self.server.scheduler.snapshot()}, request_id)
self.send_json(200, {"status": "ok", "scheduler": self.server.scheduler.snapshot(),
"kv_slots": self.server.kv_slots}, request_id)
return
self.require_auth()
if path == "/v1/models":
@@ -419,6 +427,10 @@ class APIHandler(BaseHTTPRequestHandler):
def generation(self, body, prompt, request_id, chat):
maximum, temperature, top_p = generation_options(body, self.server.max_tokens)
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:
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)
if not isinstance(stream, bool):
raise APIError(400, "`stream` must be a boolean.", "stream")
@@ -435,7 +447,8 @@ class APIHandler(BaseHTTPRequestHandler):
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)
stats = self.server.engine.generate(
prompt, maximum, temperature, top_p, output.append, cache_slot)
text = "".join(output)
finish = "length" if stats["length_limited"] else "stop"
choice = ({"index": 0, "message": {"role": "assistant", "content": text,
@@ -482,7 +495,8 @@ class APIHandler(BaseHTTPRequestHandler):
{"index": 0, "text": text, "logprobs": None, "finish_reason": None})
event([choice])
stats = self.server.engine.generate(prompt, maximum, temperature, top_p, emit)
stats = self.server.engine.generate(
prompt, maximum, temperature, top_p, emit, cache_slot)
finish = "length" if stats["length_limited"] else "stop"
final_choice = ({"index": 0, "delta": {}, "logprobs": None, "finish_reason": finish}
if chat else {"index": 0, "text": "", "logprobs": None,
@@ -536,7 +550,7 @@ class APIHandler(BaseHTTPRequestHandler):
def serve(model, host="127.0.0.1", port=8000, model_id="glm-5.2-colibri", api_key=None,
cap=8, max_tokens=1024, engine=HERE / "glm", env=None, cors_origins=None,
max_queue=8, queue_timeout=300):
max_queue=8, queue_timeout=300, kv_slots=1):
if not 1 <= max_tokens:
raise ValueError("max_tokens must be positive")
if not 1 <= port <= 65535:
@@ -545,11 +559,14 @@ def serve(model, host="127.0.0.1", port=8000, model_id="glm-5.2-colibri", api_ke
raise ValueError("max_queue cannot be negative")
if queue_timeout <= 0:
raise ValueError("queue_timeout must be positive")
if not 1 <= kv_slots <= 16:
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)
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,max_queue,queue_timeout)
server = APIServer((host, port), runtime, 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)
previous_sigterm = signal.getsignal(signal.SIGTERM)
signal.signal(signal.SIGTERM, lambda *_: threading.Thread(target=server.shutdown, daemon=True).start())
@@ -577,10 +594,11 @@ def main():
parser.add_argument("--max-queue", type=int, default=int(os.environ.get("COLI_MAX_QUEUE", "8")))
parser.add_argument("--queue-timeout", type=float,
default=float(os.environ.get("COLI_QUEUE_TIMEOUT", "300")))
parser.add_argument("--kv-slots", type=int, default=int(os.environ.get("COLI_KV_SLOTS", "1")))
args = parser.parse_args()
serve(args.model, args.host, args.port, args.model_id, args.api_key,
args.cap,args.max_tokens,args.engine,cors_origins=args.cors_origin,
max_queue=args.max_queue,queue_timeout=args.queue_timeout)
max_queue=args.max_queue,queue_timeout=args.queue_timeout,kv_slots=args.kv_slots)
if __name__ == "__main__":
c/tests/test_openai_server.py
+24 -9
View File
@@ -6,15 +6,15 @@ 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)
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):
self.calls.append((prompt, maximum, temperature, top_p))
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0):
self.calls.append((prompt, maximum, temperature, top_p, cache_slot))
on_text("Hé")
on_text("llo")
return {"prompt_tokens": 7, "completion_tokens": 2, "length_limited": False}
@@ -26,10 +26,10 @@ class BlockingEngine(FakeEngine):
self.entered = threading.Event()
self.release = threading.Event()
def generate(self, prompt, maximum, temperature, top_p, on_text):
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0):
self.entered.set()
self.release.wait(2)
return super().generate(prompt, maximum, temperature, top_p, on_text)
return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot)
class TemplateTest(unittest.TestCase):
@@ -78,6 +78,10 @@ class ProtocolTest(unittest.TestCase):
self.assertEqual(stats["prompt_tokens"], 7)
self.assertTrue(stats["length_limited"])
def test_rejects_invalid_kv_pool_before_engine_start(self):
with self.assertRaisesRegex(ValueError, "kv_slots"):
serve("/missing", kv_slots=0)
class SchedulerTest(unittest.TestCase):
def test_rejects_when_waiting_queue_is_full(self):
@@ -160,7 +164,7 @@ class HTTPTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.engine = FakeEngine()
cls.server = APIServer(("127.0.0.1", 0), cls.engine, "test-model", "secret", 16)
cls.server = APIServer(("127.0.0.1", 0),cls.engine,"test-model","secret",16,kv_slots=2)
cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
cls.thread.start()
cls.base = f"http://127.0.0.1:{cls.server.server_port}"
@@ -187,11 +191,13 @@ class HTTPTest(unittest.TestCase):
self.request("/v1/models", key="wrong")
self.assertEqual(caught.exception.code, 401)
def test_health_reports_scheduler(self):
def test_health_reports_scheduler_and_kv_slots(self):
with self.request("/health") as response:
scheduler = json.load(response)["scheduler"]
health = json.load(response)
scheduler = health["scheduler"]
self.assertEqual(scheduler["max_queue"], 8)
self.assertIn("queued", scheduler)
self.assertEqual(health["kv_slots"], 2)
def test_browser_preflight(self):
request = Request(self.base + "/v1/chat/completions", method="OPTIONS", headers={
@@ -207,7 +213,7 @@ class HTTPTest(unittest.TestCase):
def test_chat_completion(self):
with self.request("/v1/chat/completions", {
"model": "test-model", "messages": [{"role": "user", "content": "Hi"}],
"max_tokens": 4,
"max_tokens": 4, "cache_slot": 1,
}) as response:
body = json.load(response)
queue_wait = response.headers.get("x-colibri-queue-wait-ms")
@@ -216,6 +222,15 @@ class HTTPTest(unittest.TestCase):
self.assertEqual(body["usage"], {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9})
self.assertIsNotNone(queue_wait)
self.assertIn("<|user|>Hi<|assistant|><think></think>", self.engine.calls[-1][0])
self.assertEqual(self.engine.calls[-1][4], 1)
def test_rejects_invalid_cache_slot(self):
with self.assertRaises(HTTPError) as caught:
self.request("/v1/chat/completions", {
"model": "test-model", "messages": [{"role": "user", "content": "Hi"}],
"cache_slot": 2,
})
self.assertEqual(caught.exception.code, 400)
def test_streaming_chat_completion(self):
with self.request("/v1/chat/completions", {