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

* Add isolated sequence KV contexts

* Log projected multi-slot KV footprint

README.md

@@ -159,6 +159,27 @@ requests receive OpenAI-shaped HTTP 429 errors before streaming headers are sent
 `GET /health` exposes active/queued/completed/rejected counters, and successful
 generation responses include `x-colibri-queue-wait-ms`.
 
+### Isolated KV contexts
+
+`coli serve --kv-slots N` allocates up to 16 independent sequence contexts. Requests
+select one with the optional integer `cache_slot` field; ordinary OpenAI clients omit
+it and keep the original slot 0 behavior.
+
+```json
+{
+  "model": "glm-5.2-colibri",
+  "messages": [{"role": "user", "content": "Continue this conversation"}],
+  "cache_slot": 1
+}
+```
+
+Each slot owns its token history, compressed MLA/DSA KV memory, MTP window, and
+crash-safe persistence file (`.coli_kv`, `.coli_kv.1`, ...). The engine still executes
+one sequence at a time; this establishes explicit KV ownership without pretending that
+threaded HTTP is continuous batching. RAM admission accounts for every configured slot.
+Use `COLI_KV_SLOTS=N` as the environment equivalent. Start with a small value: at the
+default 4096-token context, every slot costs hundreds of MB.
+
 ### Experimental resident CUDA backend
 
 colibrì includes an opt-in CUDA backend for model-resident tensors. Streaming

c/coli

@@ -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

@@ -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); }
+    KVState *k=m->kv;
-    if(m->Ic){ for(int i=0;i<c->n_layers;i++) free(m->Ic[i]); free(m->Ic); m->Ic=NULL; }
+    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*));
+        k->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);
+        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*));
+    k->Lc=calloc(NR,sizeof(float*)); k->Rc=calloc(NR,sizeof(float*));
-    for(int i=0;i<NR;i++){ m->Lc[i]=falloc((int64_t)max_t*c->kv_lora);
+    for(int i=0;i<NR;i++){ k->Lc[i]=falloc((int64_t)max_t*c->kv_lora);
-        m->Rc[i]=falloc((int64_t)max_t*c->qk_rope); }
+        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");
+    KVState *k=m->kv;
-    if(!f){ g_kv_nrec=0; return; }
+    FILE *f=fopen(k->disk_path,"r+b");
-    g_kv_nrec=nrec;
+    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");
+    FILE *f=fopen(k->disk_path,"r+b");
-    if(!f){ f=fopen(g_kv_path,"wb"); if(!f) return;
+    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);
+    fseek(f, 8+8*4 + (int64_t)k->disk_nrec*rec, SEEK_SET);
-    for(int p=g_kv_nrec;p<len;p++){
+    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 nctx=getenv("KV_SLOTS")?atoi(getenv("KV_SLOTS")):1;
-    { int r=kv_disk_load(m,hist,maxctx); if(r>0){ len=r; first=0; } }
+    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;
+            unsigned long long nb=0; double rt=0, rp=0; int slot=0;
-            if(sscanf(line+8,"%llu %d %lf %lf",&nb,&req_ngen,&rt,&rp)!=4 ||
+            int nf=sscanf(line+8,"%llu %d %lf %lf %d",&nb,&req_ngen,&rt,&rp,&slot);
-               nb>(16u<<20) || req_ngen<1 || rt<0 || rt>2 || rp<=0 || rp>1){
+            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

@@ -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):
+    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))
+        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

@@ -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):
+    def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0):
-        self.calls.append((prompt, maximum, temperature, top_p))
+        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", {