Merge #192: serve stage 2 — response_format, per-request grammars, grammar-forced drafts

Per-request grammar constraints in serve mode: response_format plumbed through the API, per-request grammar compilation, and grammar-forced draft verification. Verified locally on top of current dev: clean build, token-exact tiny models unchanged, decode-batch helper tests ok, 40/40 openai_server tests. CI 8/8 green. Thanks @fabio-rovai for keeping this rebased through the refactor.

c/colibri.c

@@ -596,12 +596,6 @@ static int g_draft=0;    /* metodo E: DRAFT=n token auto-speculati per forward v
  * parte (#8). MAI un vincolo sul sampling: solo proposte, la verifica batch-union
  * decide — grammatica sbagliata = draft rifiutati, output identico.
  * GRAMMAR_DRAFT=n (default 24) limita i token forzati per forward. */
-static Grammar g_gram; static GrState g_gst;
-static Tok *g_gr_T=NULL;
-static int g_gr_on=0;     /* grammatica caricata e walker vivo */
-static int g_gr_armed=0;  /* lazy: parte dal primo byte ammesso dalla radice (salta i preamboli) */
-static int g_gr_max=24;
-static uint64_t g_gr_prop=0, g_gr_acc=0;
 static FILE *g_route_fp=NULL; /* ROUTE_TRACE=<path>: dump per-position top-K routing (ids:gates)
                                * per layer — offline co-activation / coupling analysis. Zero
                                * effect on computation; measurement only. */
@@ -621,6 +615,19 @@ static int g_couple=0, g_couple_k=8, g_couple_d=1;
 static int16_t *cp_pred=NULL;    /* [(L*2+(dL-1))*E + e]*CP_M + j -> target id (-1 none) */
 static float   *cp_cnt=NULL;
 static long g_cp_enq=0;
+/* All grammar-forced-draft state in one struct so it can become per-request
+ * in the multiplexed server. Fields (same semantics as the former globals):
+ * on = grammar loaded and walker alive; armed = lazy start from the first byte
+ * accepted at the root (skips preambles); max = forced-span cap per forward;
+ * prop/acc = proposed/accepted forced-draft counters. */
+typedef struct {
+    Grammar gram;
+    GrState st;
+    Tok *T;
+    int on, armed, max;
+    uint64_t prop, acc;
+} GrDraft;
+static GrDraft g_grd={.max=24};   /* process-level instance: PROMPT mode + run_serve keep using this */
 static void couple_prefetch(Model *m, int layer, const int *idx, int Ke);
 static int g_looka=0;    /* LOOKA=1: misura (solo contatori, zero effetti) quanto il routing MoE
                           * e' predicibile IN ANTICIPO — la domanda che decide se un prefetch
@@ -4092,10 +4099,36 @@ static void repin_pass(Model *m){ repin_pass_limit(m,16); }
  * grammar_draft propone lo span FORZATO successivo (un solo byte legale per posizione)
  * gia' tokenizzato. Il confine di tokenizzazione non e' garantito coincidere con quello
  * del modello: la verifica assorbe la differenza (al peggio l'ultimo draft e' rifiutato). */
-static void grammar_setup(Tok *T){
+/* Compile grammar TEXT into g: raw GBNF, or — if the first non-space byte is '{'
+ * — a JSON-Schema compiled via schema_gbnf.h. Takes ownership of txt (always
+ * freed). Fail-soft: returns -1 with g->on=0, engine runs without a grammar. */
+static int grammar_setup_text(GrDraft *g, Tok *T, char *txt, const char *label){
+    const char *p=txt; while(*p==' '||*p=='\t'||*p=='\n'||*p=='\r') p++;
+    if(*p=='{'){
+        char serr[160];
+        char *gbnf=schema_to_gbnf(txt,serr,sizeof serr);
+        free(txt);
+        if(!gbnf){ fprintf(stderr,"[SCHEMA] %s: %s (running without grammar)\n",label,serr); return -1; }
+        txt=gbnf;
+    }
+    if(gr_parse(&g->gram,txt)){ fprintf(stderr,"[GRAMMAR] %s: %s\n",label,g->gram.err); free(txt); return -1; }
+    free(txt);
+    gr_state_init(&g->st,&g->gram);
+    if(!g->st.alive){ fprintf(stderr,"[GRAMMAR] %s: grammar cannot be evaluated (left recursion?)\n",label); return -1; }
+    if(g->max<1) g->max=24;
+    if(g->max>48) g->max=48;
+    g->T=T; g->on=1; g->armed=0;
+    fprintf(stderr,"[GRAMMAR] %s: %d rules, forced span capped at %d tokens/forward\n",label,g->gram.n,g->max);
+    return 0;
+}
+/* Release a per-request grammar so the slot can host the next request (keeps max). */
+static void grammar_teardown(GrDraft *g){
+    if(g->gram.n) gr_free(&g->gram);
+    int max=g->max; memset(g,0,sizeof(*g)); g->max=max;
+}
+static void grammar_setup(GrDraft *g, Tok *T){
     /* GRAMMAR=<file.gbnf> takes precedence; SCHEMA=<file.json> compiles a JSON-Schema
-     * to GBNF (schema_gbnf.h) for the same draft source. Both fail soft: the engine
+     * to GBNF. Both fail soft: the engine runs without a grammar, output unchanged. */
-     * runs without a grammar and output is unchanged. */
     const char *gf=getenv("GRAMMAR");
     const char *sf=(gf&&*gf)?NULL:getenv("SCHEMA");
     if((!gf||!*gf)&&(!sf||!*sf)) return;
@@ -4107,59 +4140,45 @@ static void grammar_setup(Tok *T){
     if(!txt || fread(txt,1,(size_t)n,f)!=(size_t)n){
         fprintf(stderr,"[GRAMMAR] failed to read %s\n",path); fclose(f); free(txt); return; }
     fclose(f); txt[n]=0;
-    if(sf){ /* schema -> GBNF, then the same gr_parse as the GRAMMAR path */
+    if(getenv("GRAMMAR_DRAFT")) g->max=atoi(getenv("GRAMMAR_DRAFT"));
-        char serr[160];
+    grammar_setup_text(g,T,txt,path);
-        char *gbnf=schema_to_gbnf(txt,serr,sizeof serr);
-        free(txt);
-        if(!gbnf){ fprintf(stderr,"[SCHEMA] %s: %s (running without grammar)\n",sf,serr); return; }
-        txt=gbnf;
-    }
-    if(gr_parse(&g_gram,txt)){ fprintf(stderr,"[GRAMMAR] %s: %s\n",path,g_gram.err); free(txt); return; }
-    free(txt);
-    gr_state_init(&g_gst,&g_gram);
-    if(!g_gst.alive){ fprintf(stderr,"[GRAMMAR] %s: grammar cannot be evaluated (left recursion?)\n",path); return; }
-    if(getenv("GRAMMAR_DRAFT")) g_gr_max=atoi(getenv("GRAMMAR_DRAFT"));
-    if(g_gr_max<1) g_gr_max=1;
-    if(g_gr_max>48) g_gr_max=48;
-    g_gr_T=T; g_gr_on=1;
-    fprintf(stderr,"[GRAMMAR] %s: %d rules, forced span capped at %d tokens/forward\n",path,g_gram.n,g_gr_max);
 }
 /* stato pulito all'inizio di ogni RISPOSTA (non tra i \x02MORE, che continuano) */
-static void grammar_reset(void){
+static void grammar_reset(GrDraft *g){
-    if(!g_gr_on) return;
+    if(!g->on) return;
-    gr_state_init(&g_gst,&g_gram); g_gr_armed=0;
+    gr_state_init(&g->st,&g->gram); g->armed=0;
-    if(!g_gst.alive) g_gr_on=0;
+    if(!g->st.alive) g->on=0;
 }
 /* consuma i byte di un token emesso. Preambolo (prima dell'arming): ignorato.
  * Desync dopo l'arming: si riarma in attesa del prossimo inizio valido — al peggio
  * i draft vengono rifiutati dalla verifica, l'output non cambia MAI. */
-static void gr_feed(int t){
+static void gr_feed(GrDraft *g, int t){
-    if(!g_gr_on||!g_gr_T) return;
+    if(!g->on||!g->T) return;
-    char b[64]; int n=tok_decode(g_gr_T,&t,1,b,63);
+    char b[64]; int n=tok_decode(g->T,&t,1,b,63);
     for(int i=0;i<n;i++){
-        int r=gr_accept(&g_gst,(unsigned char)b[i]);
+        int r=gr_accept(&g->st,(unsigned char)b[i]);
-        if(r==1){ g_gr_armed=1; continue; }
+        if(r==1){ g->armed=1; continue; }
-        if(r<0){ g_gr_on=0; return; }                 /* walker spento: fine dei draft */
+        if(r<0){ g->on=0; return; }                   /* walker spento: fine dei draft */
-        if(!g_gr_armed) continue;                     /* preambolo: aspetta l'inizio */
+        if(!g->armed) continue;                       /* preambolo: aspetta l'inizio */
-        gr_state_init(&g_gst,&g_gram); g_gr_armed=0;  /* desync: riparti dalla radice */
+        gr_state_init(&g->st,&g->gram); g->armed=0;   /* desync: riparti dalla radice */
-        if(!g_gst.alive){ g_gr_on=0; return; }
+        if(!g->st.alive){ g->on=0; return; }
-        if(gr_accept(&g_gst,(unsigned char)b[i])==1) g_gr_armed=1;
+        if(gr_accept(&g->st,(unsigned char)b[i])==1) g->armed=1;
     }
 }
 /* propone lo span forzato come token (max cap); 0 se la grammatica dirama qui */
-static int grammar_draft(int *draft, int cap){
+static int grammar_draft(GrDraft *g, int *draft, int cap){
-    if(!g_gr_on||!g_gr_armed||!g_gr_T||cap<1) return 0;
+    if(!g->on||!g->armed||!g->T||cap<1) return 0;
-    if(g_gr_prop>=32 && g_gr_acc*2<g_gr_prop){        /* guardia adattiva, come per MTP:
+    if(g->prop>=32 && g->acc*2<g->prop){              /* guardia adattiva, come per MTP:
         acceptance sotto il 50% = tokenizzazione fuori asse, meglio spegnersi */
-        g_gr_on=0;
+        g->on=0;
         fprintf(stderr,"[GRAMMAR] %.0f%% acceptance after %llu proposals: grammar drafts disabled\n",
-            100.0*g_gr_acc/g_gr_prop,(unsigned long long)g_gr_prop);
+            100.0*g->acc/g->prop,(unsigned long long)g->prop);
         return 0;
     }
-    char fb[512]; int nb=gr_forced(&g_gst,fb,(int)sizeof fb-1);
+    char fb[512]; int nb=gr_forced(&g->st,fb,(int)sizeof fb-1);
     if(nb<=0) return 0;
-    int g=tok_encode(g_gr_T,fb,nb,draft,cap);
+    int nt=tok_encode(g->T,fb,nb,draft,cap);          /* renamed local: 'g' is now the state param */
-    return g>0?g:0;
+    return nt>0?nt:0;
 }
 
 /* ---- SAMPLING (temperatura + nucleus) con verifica speculativa LOSSLESS ----
@@ -4217,12 +4236,12 @@ static int spec_decode(Model *m, int *all, int kv, int n_new, int eos, float *lo
         int next=pick_tok(logit,V,carry_ban); carry_ban=-1; free(logit); logit=NULL;
         if((eos>=0 && next==eos) || is_stop(next)) break;
         emit(next,ud); all[kv]=next; emitted++; m->n_emit++;
-        gr_feed(next);                                  /* il walker segue l'output emesso */
+        gr_feed(&g_grd,next);                           /* il walker segue l'output emesso */
         if(emitted>=n_new) break;                       /* l'ultimo token non serve forwardarlo */
         int g = 0, gsrc = 0;                            /* sorgente: 1=grammatica 2=MTP/n-gram */
-        if(g_gr_on){                                    /* metodo F: prima la grammatica — dove
+        if(g_grd.on){                                   /* metodo F: prima la grammatica — dove
                                                          * forza, l'acceptance e' ~1 (#48) */
-            g=grammar_draft(draft,g_gr_max);
+            g=grammar_draft(&g_grd,draft,g_grd.max);
             if(g>0) gsrc=1;
         }
         if(!g && g_draft>0 && m->has_mtp){
@@ -4244,7 +4263,7 @@ static int spec_decode(Model *m, int *all, int kv, int n_new, int eos, float *lo
         if(g>n_new-emitted) g=n_new-emitted;
         if(kv+1+g+1>m->max_t) g=m->max_t-kv-2;
         if(g<0) g=0;
-        if(gsrc==1) g_gr_prop+=(uint64_t)g;
+        if(gsrc==1) g_grd.prop+=(uint64_t)g;
         int S=1+g; int batch[64]; batch[0]=next; memcpy(batch+1,draft,g*sizeof(int));
         double tf0=g_prof?now_s():0;
         float *lo=step_all(m,batch,S,kv); m->n_fw++;
@@ -4260,9 +4279,9 @@ static int spec_decode(Model *m, int *all, int kv, int n_new, int eos, float *lo
             if(!accept){ if(g_temp>0) carry_ban=draft[k]; break; }
             if((eos>=0 && draft[k]==eos) || is_stop(draft[k])){ done=1; break; }
             emit(draft[k],ud); all[kv+1+k]=draft[k]; emitted++; m->n_emit++;
-            gr_feed(draft[k]); k++;
+            gr_feed(&g_grd,draft[k]); k++;
         }
-        if(gsrc==1) g_gr_acc+=(uint64_t)k;
+        if(gsrc==1) g_grd.acc+=(uint64_t)k;
         else if(gsrc==2 && m->has_mtp) m->mtp_acc+=k;
         if(m->has_mtp && k>=1) mtp_absorb(m, all+kv+1, m->h_all, k, kv);   /* KV MTP in sync coi verificati */
         /* hlast deve corrispondere all'ultima posizione ACCETTATA (kv+k), non a fine batch */
@@ -4564,7 +4583,7 @@ static void run_text(Model *m, const char *snap, const char *prompt, int ngen){
     Tok T; tok_load(&T,tkp);
     int eos=tok_id_of(&T,"<|endoftext|>");
     stops_arm_tok(&m->c, eos, &T);
-    grammar_setup(&T);                   /* metodo F: GRAMMAR=file.gbnf (#48) */
+    grammar_setup(&g_grd,&T);                   /* metodo F: GRAMMAR=file.gbnf (#48) */
     if(g_temp<0) g_temp=0.7f;            /* auto: 0.7, NON l'1.0 ufficiale — la coda della
                                           * distribuzione int4 e' rumore di quantizzazione */
     int cap=(int)strlen(prompt)+16; int *pids=malloc((cap+2)*sizeof(int));
@@ -4605,7 +4624,7 @@ static void run_text(Model *m, const char *snap, const char *prompt, int ngen){
     ProfBase pb; prof_base(m,&pb);
     double t=now_s();
     EmitStream es={&T,m,t,0,0};
-    grammar_reset();
+    grammar_reset(&g_grd);
     int produced=spec_decode(m,all,np,ngen,eos,logit,emit_stream,&es,NULL);
     double dt=now_s()-t;
     double tot=m->hits+m->miss;
@@ -4632,8 +4651,8 @@ static void run_text(Model *m, const char *snap, const char *prompt, int ngen){
         m->n_fw?(double)m->n_emit/m->n_fw:1.0, (unsigned long long)m->n_fw, (unsigned long long)m->n_emit,
         m->mtp_prop?100.0*m->mtp_acc/m->mtp_prop:0.0, (unsigned long long)m->mtp_acc, (unsigned long long)m->mtp_prop);
     if(g_cp_enq) printf("couple: %ld cross-layer prefetch hints enqueued\n", g_cp_enq);
-    if(g_gr_prop) printf("grammar: %.0f%% acceptance (%llu/%llu forced drafts)\n",
+    if(g_grd.prop) printf("grammar: %.0f%% acceptance (%llu/%llu forced drafts)\n",
-        100.0*g_gr_acc/g_gr_prop, (unsigned long long)g_gr_acc, (unsigned long long)g_gr_prop);
+        100.0*g_grd.acc/g_grd.prop, (unsigned long long)g_grd.acc, (unsigned long long)g_grd.prop);
     if(g_disk_split) printf("disk-load split: draft %llu + absorb %llu + verify/main %llu misses | "
            "MTP-layer %llu loads %.2f GB | main-layers %llu loads %.2f GB (MTP %.1f%% of bytes)\n",
         (unsigned long long)m->miss_draft, (unsigned long long)m->miss_absorb,
@@ -4953,8 +4972,8 @@ static void mux_done(Model *m, ServeCtx *sc, ServeReq *r){
 /* 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,
+static int mux_submit(Model *m, Tok *T, ServeCtx *ctx, ServeReq *req, GrDraft *grd,
-                      int maxctx, int eos){
+                      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;
@@ -4975,21 +4994,31 @@ static int mux_submit(Model *m, Tok *T, ServeCtx *ctx, ServeReq *req, int nctx,
     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; }
+    char *gtxt=NULL;                     /* optional per-request grammar/schema text */
+    if(sub.gbytes){
+        gtxt=malloc((size_t)sub.gbytes+1);
+        if(!gtxt){ fprintf(stderr,"OOM multiplex payload\n"); exit(1); }
+        if(fread(gtxt,1,(size_t)sub.gbytes,stdin)!=(size_t)sub.gbytes){
+            free(gtxt); free(raw); free(line); return -1; }
+        gtxt[sub.gbytes]=0;
+    }
     int delim=fgetc(stdin);
     if(delim!='\n'){
         printf("ERROR %llu BAD_FRAME\n",sub.id); fflush(stdout);
-        free(raw); free(line); return -1;
+        free(gtxt); 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;
+        printf("ERROR %llu BAD_REQUEST\n",sub.id); fflush(stdout); free(gtxt); 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;
+        printf("ERROR %llu SLOT_BUSY\n",sub.id); fflush(stdout); free(gtxt); 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;
+        printf("ERROR %llu DUPLICATE_ID\n",sub.id); fflush(stdout); free(gtxt); free(raw); free(line); return 0;
     }
+    grammar_teardown(&grd[sub.slot]);    /* new request owns the slot's grammar state */
+    if(gtxt) grammar_setup_text(&grd[sub.slot],T,gtxt,"request");   /* fail-soft; owns gtxt */
     ServeCtx *sc=&ctx[sub.slot]; kv_bind(m,&sc->kv);
     int *tmp=malloc(maxctx*sizeof(int));
     if(!tmp){ fprintf(stderr,"OOM mux_submit tmp\n"); free(raw); free(line); exit(1); }
@@ -5015,6 +5044,7 @@ static int mux_submit(Model *m, Tok *T, ServeCtx *ctx, ServeReq *req, int nctx,
     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++;
+    if(grd[sub.slot].on){ grammar_reset(&grd[sub.slot]); gr_feed(&grd[sub.slot],next); }
     mux_data(T,r->id,next);
     if(r->emitted>=r->maximum) mux_done(m,sc,r);
     return 1;
@@ -5023,14 +5053,18 @@ static int mux_submit(Model *m, Tok *T, ServeCtx *ctx, ServeReq *req, int nctx,
 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_tok(&m->c,eos,&T);
-    g_draft=0; /* one scheduler owns every forward; MTP/speculation is not ragged-safe */
+    g_draft=0; /* one scheduler owns every forward; MTP/n-gram speculation is not ragged-safe.
+                * Grammar-forced drafts ARE mux-safe (below): a drafting slot leaves the shared
+                * batch for one forward and runs the proven single-sequence verify path
+                * (kv_bind + step_all), exactly like prefill already does per submission. */
     int maxctx=getenv("CTX")?atoi(getenv("CTX")):4096;
     int nctx=getenv("KV_SLOTS")?atoi(getenv("KV_SLOTS")):1;
     if(nctx<1||nctx>512){fprintf(stderr,"KV_SLOTS must be between 1 and 512\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);
+    GrDraft *grd=calloc(nctx,sizeof(*grd));   /* per-slot request grammars (SUBMIT 7th field) */
+    for(int i=0;i<nctx;i++){ serve_ctx_init(m,&ctx[i],snap,i,maxctx); grd[i].max=24; }
 #ifdef _WIN32
     /* Same byte-exact protocol as run_serve: in TEXT mode the CRT collapses CRLF in
      * fread() payloads (waits forever for the missing bytes) and expands LF on the
@@ -5078,14 +5112,47 @@ static void run_serve_mux(Model *m, const char *snap){
             else ready=(PeekNamedPipe(ih,NULL,0,NULL,&avail,NULL) && avail>0)?1:0;
             if(ready)
 #endif
-                if(mux_submit(m,&T,ctx,req,nctx,maxctx,eos)<0) eof=1;
+                if(mux_submit(m,&T,ctx,req,grd,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[512]; int slots[512], 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;
+            ServeCtx *sc=&ctx[i]; ServeReq *r=&req[i]; GrDraft *gd=&grd[i];
+            /* grammar-forced drafts (greedy requests only: verification under sampling
+             * needs rejection resampling, out of scope here). The slot leaves the shared
+             * batch for one forward and runs the single-sequence verify path. */
+            int draft[49], k=0;
+            if(gd->on && r->temp==0) k=grammar_draft(gd,draft,gd->max);
+            if(k>0 && sc->len+1+k<maxctx-1){
+                if(sc->len+1+k>=(int)sc->kv.max_t) k=(int)sc->kv.max_t-sc->len-2;
+                if(k<1){ rows[S]=(DecodeRow){&sc->kv,r->pending,sc->len}; slots[S++]=i; continue; }
+                kv_bind(m,&sc->kv);
+                int seq[50]; seq[0]=r->pending;
+                memcpy(seq+1,draft,(size_t)k*sizeof(int));
+                float *lo=step_all(m,seq,1+k,sc->len); m->n_fw++;
+                gd->prop+=(uint64_t)k;
+                int done=0;
+                for(int j=0;j<=k && !done;j++){
+                    g_temp=r->temp; g_nuc=r->top_p;
+                    int next=pick_tok(lo+(int64_t)j*m->c.vocab,m->c.vocab,-1);
+                    sc->len++;                       /* seq[j] joins the committed history */
+                    if(next==eos || is_stop(next)){ mux_done(m,sc,r); done=1; break; }
+                    r->pending=next; sc->hist[sc->len]=next; r->emitted++; m->n_emit++;
+                    if(gd->on) gr_feed(gd,next);
+                    mux_data(&T,r->id,next);
+                    if(r->emitted>=r->maximum){ mux_done(m,sc,r); done=1; break; }
+                    if(j<k){
+                        if(next!=draft[j]) break;    /* rejected: seq[j+1..] stale, overwritten next forward */
+                        gd->acc++;
+                    }
+                }
+                free(lo);
+                continue;                            /* handled outside the shared batch */
+            }
+            rows[S]=(DecodeRow){&sc->kv,r->pending,sc->len}; slots[S++]=i;
         }
+        if(S==0) continue;               /* every active slot drafted this round */
         double tf0=g_prof?now_s():0;
         float *lo=step_decode_batch(m,rows,S); if(!lo){fprintf(stderr,"decode batch failed\n");break;}
         m->n_fw++;
@@ -5096,13 +5163,15 @@ static void run_serve_mux(Model *m, const char *snap){
             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++;
+            if(grd[i].on) gr_feed(&grd[i],next);   /* walker stays in sync when not drafting */
             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);
+    for(int i=0;i<nctx;i++){ serve_ctx_free(m,&ctx[i]); grammar_teardown(&grd[i]); }
+    free(ctx); free(req); free(grd);
     m->kv=NULL; m->Lc=m->Rc=m->Ic=NULL; m->kv_start=NULL; m->max_t=0;
 }
 
@@ -5126,7 +5195,7 @@ static void run_serve(Model *m, const char *snap){
     Tok T; tok_load(&T,tkp);
     int eos=tok_id_of(&T,"<|endoftext|>");
     stops_arm_tok(&m->c, eos, &T);
-    grammar_setup(&T);                   /* metodo F: GRAMMAR=file.gbnf (#48) */
+    grammar_setup(&g_grd,&T);                   /* metodo F: GRAMMAR=file.gbnf (#48) */
     if(g_temp<0) g_temp=0.7f;            /* auto: 0.7, NON l'1.0 ufficiale — la coda della
                                           * distribuzione int4 e' rumore di quantizzazione */
     int ngen=getenv("NGEN")?atoi(getenv("NGEN")):256;
@@ -5245,7 +5314,7 @@ static void run_serve(Model *m, const char *snap){
         else logit=step(m,hist+len-1,1,len-1);   /* prompt identico/prefisso: rigenera i logits */
         EmitStream es={&T,m,now_s(),0,1};
         int prod=0;
-        grammar_reset();                         /* nuova risposta = nuovo documento (MORE invece continua) */
+        grammar_reset(&g_grd);                         /* nuova risposta = nuovo documento (MORE invece continua) */
         if(cur>0) prod=spec_decode(m,hist,len,cur,eos,logit,emit_stream,&es,&len);
         else free(logit);
         double tdt=now_s()-tt0; if(tdt<1e-6) tdt=1e-6;

c/decode_batch.h

@@ -13,22 +13,32 @@ static inline float *coli_kv_row(float *base, int position, int width)
 }
 
 typedef struct {
-    unsigned long long id, bytes;
+    unsigned long long id, bytes, gbytes;
     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. */
+ * it may contain newlines. Reject trailing fields to keep framing unambiguous.
+ * Optional 7th field `gbytes`: length of a per-request grammar (raw GBNF, or a
+ * JSON-Schema compiled engine-side) appended to the payload AFTER the prompt
+ * bytes. 6-field headers remain valid (gbytes = 0). */
 static inline int coli_submit_parse(const char *line, ColiSubmit *s)
 {
     char tail;
-    if (!line || !s ||
+    if (!line || !s) return 0;
-        sscanf(line, "SUBMIT %llu %d %llu %d %f %f %c", &s->id, &s->slot,
+    s->gbytes = 0;
+    if (sscanf(line, "SUBMIT %llu %d %llu %d %f %f %llu %c", &s->id, &s->slot,
                &s->bytes, &s->max_tokens, &s->temperature, &s->top_p,
-               &tail) != 6)
+               &s->gbytes, &tail) != 7) {
-        return 0;
+        s->gbytes = 0;
-    return s->id > 0 && s->bytes <= (16u << 20) && s->slot >= 0 && s->max_tokens >= 1 &&
+        if (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->gbytes <= (1u << 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;

c/openai_server.py

@@ -370,6 +370,21 @@ def render_chat(messages, enable_thinking=False, reasoning_effort=None, tools=No
     return "".join(prompt)
 
 
+# Generic whitespace-tolerant JSON grammar for response_format {"type": "json_object"}.
+# Draft-source semantics: positions with one legal byte draft; jws points just keep
+# the walker alive through the model's own spacing (see docs/grammar-draft.md).
+GENERIC_JSON_GBNF = (
+    'root ::= jws jval jws\n'
+    'jval ::= jobj | jarr | jstr | jnum | "true" | "false" | "null"\n'
+    'jobj ::= "{" jws ( jstr jws ":" jws jval jws ( "," jws jstr jws ":" jws jval jws )* )? "}"\n'
+    'jarr ::= "[" jws ( jval jws ( "," jws jval jws )* )? "]"\n'
+    'jstr ::= "\\"" jchar* "\\""\n'
+    'jchar ::= [^"\\\\\\x00-\\x1f] | "\\\\" ( ["\\\\/bfnrt] | "u" jhex jhex jhex jhex )\n'
+    'jhex ::= [0-9a-fA-F]\n'
+    'jnum ::= "-"? ( "0" | [1-9] [0-9]* ) ( "." [0-9]+ )? ( ( "e" | "E" ) ( "+" | "-" )? [0-9]+ )?\n'
+    'jws ::= ( " " | "\\t" | "\\n" | "\\r" )*\n'
+)
+
 def generation_options(body, limit):
     if body.get("n", 1) != 1:
         raise APIError(400, "Colibri currently supports `n=1` only.", "n", "unsupported_value")
@@ -424,10 +439,38 @@ def generation_options(body, limit):
         raise APIError(400, "Token penalties are not supported yet.", None, "unsupported_parameter")
     if body.get("seed") is not None:
         raise APIError(400, "Per-request seeds are not supported yet.", "seed", "unsupported_parameter")
+    # response_format -> optional per-request grammar for the engine's grammar-forced
+    # draft source (#70/#148). NEVER a sampling constraint: drafts are verified, so a
+    # schema the engine cannot compile degrades to "no speedup", not to an error and
+    # not to changed output. json_schema payloads are forwarded as-is (the engine
+    # compiles them via schema_gbnf.h); {"type": "gbnf"} is a raw-GBNF extension.
+    grammar = None
     response_format = body.get("response_format")
-    if response_format not in (None, {"type": "text"}):
+    if response_format is not None and response_format != {"type": "text"}:
-        raise APIError(400, "Only the default text response format is supported.",
+        if not isinstance(response_format, dict) or "type" not in response_format:
-                       "response_format", "unsupported_parameter")
+            raise APIError(400, "`response_format` must be an object with a `type`.",
+                           "response_format", "invalid_value")
+        ftype = response_format["type"]
+        if ftype == "json_object":
+            grammar = GENERIC_JSON_GBNF
+        elif ftype == "json_schema":
+            schema = (response_format.get("json_schema") or {}).get("schema")
+            if not isinstance(schema, dict):
+                raise APIError(400, "`response_format.json_schema.schema` must be an object.",
+                               "response_format", "invalid_value")
+            grammar = json.dumps(schema)
+        elif ftype == "gbnf":
+            grammar = response_format.get("grammar")
+            if not isinstance(grammar, str) or not grammar.strip():
+                raise APIError(400, "`response_format.grammar` must be a non-empty GBNF string.",
+                               "response_format", "invalid_value")
+        else:
+            raise APIError(400, "`response_format.type` must be \"text\", \"json_object\", "
+                                "\"json_schema\" or \"gbnf\".",
+                           "response_format", "unsupported_value")
+        if grammar is not None and len(grammar.encode("utf-8")) > (1 << 20):
+            raise APIError(400, "`response_format` grammar/schema exceeds 1 MiB.",
+                           "response_format", "invalid_value")
 
     maximum = body.get("max_completion_tokens")
     maximum_param = "max_completion_tokens"
@@ -454,7 +497,7 @@ def generation_options(body, limit):
     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)
+    return maximum, float(temperature), float(top_p), grammar
 
 
 def read_engine_turn(stream, sentinel, on_bytes):
@@ -618,12 +661,15 @@ class Engine:
                 self._fail_pending(error)
 
     def generate(self, prompt, max_tokens, temperature, top_p, on_text, cache_slot=0,
-                 cancelled=None):
+                 cancelled=None, grammar=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")
         if b"\0" in payload:
             raise APIError(400, "NUL bytes are not supported in prompts.", "messages")
+        gpayload = grammar.encode("utf-8") if grammar else b""
+        if b"\0" in gpayload:
+            raise APIError(400, "NUL bytes are not supported in grammars.", "response_format")
         decoder = codecs.getincrementaldecoder("utf-8")("replace")
 
         def decode(data):
@@ -643,12 +689,13 @@ class Engine:
             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()
+                  f"{temperature:.8g} {top_p:.8g}"
+                  + (f" {len(gpayload)}" if gpayload else "") + "\n").encode()
         try:
             with self.write_lock:
                 if self.process.poll() is not None:
                     raise RuntimeError("colibri engine is not running")
-                self.process.stdin.write(header + payload + b"\n")
+                self.process.stdin.write(header + payload + gpayload + b"\n")
                 self.process.stdin.flush()
         except Exception:
             with self.pending_lock:
@@ -895,7 +942,7 @@ class APIHandler(BaseHTTPRequestHandler):
         if dbg >= 2:
             sys.stderr.write(f"\n===== PROMPT [{request_id}] =====\n{prompt}\n===== OUTPUT [{request_id}] =====\n")
             sys.stderr.flush()
-        maximum, temperature, top_p = generation_options(body, self.server.max_tokens)
+        maximum, temperature, top_p, grammar = generation_options(body, self.server.max_tokens)
         # tools and tool_choice come from chat_completion() already processed/filtered
         if chat and tool_choice == "none":
             tools = None          # client forbade tools: never surface tool_calls
@@ -924,7 +971,7 @@ class APIHandler(BaseHTTPRequestHandler):
                 output = []
                 stats = self.server.engine.generate(
                     prompt, maximum, temperature, top_p, output.append, cache_slot,
-                    self.client_disconnected)
+                    self.client_disconnected, grammar=grammar)
                 text = "".join(output)
                 length_finish = "length" if stats["length_limited"] else "stop"
                 if chat and tools:
@@ -1031,7 +1078,7 @@ class APIHandler(BaseHTTPRequestHandler):
                         sp["buf"] = sp["buf"][flush:]
                 stats = self.server.engine.generate(
                     prompt, maximum, temperature, top_p, emit_tools, cache_slot,
-                    lambda: not connected)
+                    lambda: not connected, grammar=grammar)
                 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)
@@ -1048,7 +1095,7 @@ class APIHandler(BaseHTTPRequestHandler):
                     emit(chunk)
                 stats = self.server.engine.generate(
                     prompt, maximum, temperature, top_p, emit_plain, cache_slot,
-                    lambda: not connected)
+                    lambda: not connected, grammar=grammar)
                 finish = "length" if stats["length_limited"] else "stop"
             ka_stop.set()                          # generation done: stop the keepalive pump
             ka_thread.join(timeout=2)

c/tests/test_decode_batch.c

@@ -44,6 +44,12 @@ static void test_submit_header(void)
     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));
+    /* optional 7th field: per-request grammar length (0 when absent) */
+    assert(coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95", &sub) && sub.gbytes == 0);
+    assert(coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95 512", &sub) && sub.gbytes == 512);
+    assert(coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95 1048576", &sub));
+    assert(!coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95 1048577", &sub));
+    assert(!coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95 512 extra", &sub));
 }
 
 int main(void)

c/tests/test_openai_server.py

@@ -20,8 +20,8 @@ class FakeEngine:
         self.calls = []
 
     def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0,
-                 cancelled=None):
+                 cancelled=None, grammar=None):
-        self.calls.append((prompt, maximum, temperature, top_p, cache_slot))
+        self.calls.append((prompt, maximum, temperature, top_p, cache_slot, grammar))
         on_text("Hé")
         on_text("llo")
         return {"prompt_tokens": 7, "completion_tokens": 2, "length_limited": False}
@@ -34,7 +34,7 @@ class BlockingEngine(FakeEngine):
         self.release = threading.Event()
 
     def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0,
-                 cancelled=None):
+                 cancelled=None, grammar=None):
         self.entered.set()
         self.release.wait(2)
         return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot,
@@ -71,11 +71,11 @@ class TemplateTest(unittest.TestCase):
 
     def test_validates_generation_limits(self):
         self.assertEqual(generation_options({"max_tokens": 4, "temperature": 0, "top_p": 1}, 8),
-                         (4, 0.0, 1.0))
+                         (4, 0.0, 1.0, None))
         # max_tokens above the server cap is clamped, not rejected (#260): OpenAI
         # clients default to large values; erroring breaks them.
         self.assertEqual(generation_options({"max_tokens": 9, "temperature": 0, "top_p": 1}, 8),
-                         (8, 0.0, 1.0))
+                         (8, 0.0, 1.0, None))
         # non-positive / non-int max_tokens is still a hard error
         with self.assertRaises(APIError):
             generation_options({"max_tokens": 0}, 8)
@@ -84,7 +84,31 @@ class TemplateTest(unittest.TestCase):
         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))
+                         (8, 0.7, 0.9, None))
+        # response_format -> grammar plumbing (draft source, never a constraint)
+        opts = generation_options({"max_tokens": 4, "response_format": {"type": "json_object"}}, 8)
+        self.assertIn("root ::=", opts[3])
+        schema = {"type": "object", "properties": {"a": {"type": "string"}}, "required": ["a"]}
+        opts = generation_options({"max_tokens": 4, "response_format":
+                                   {"type": "json_schema", "json_schema": {"schema": schema}}}, 8)
+        self.assertEqual(json.loads(opts[3]), schema)
+        opts = generation_options({"max_tokens": 4, "response_format":
+                                   {"type": "gbnf", "grammar": 'root ::= "x"'}}, 8)
+        self.assertEqual(opts[3], 'root ::= "x"')
+        with self.assertRaises(APIError):
+            generation_options({"response_format": {"type": "yaml"}}, 8)
+        with self.assertRaises(APIError):
+            generation_options({"response_format": {"type": "json_schema", "json_schema": {}}}, 8)
+        with self.assertRaises(APIError):   # non-dict response_format
+            generation_options({"response_format": "json"}, 8)
+        with self.assertRaises(APIError):   # empty gbnf
+            generation_options({"response_format": {"type": "gbnf", "grammar": "  "}}, 8)
+        with self.assertRaises(APIError):   # oversized grammar (> 1 MiB pre-check)
+            generation_options({"response_format": {"type": "gbnf", "grammar": "x" * ((1 << 20) + 1)}}, 8)
+        # malformed GBNF passes the gateway by design: the ENGINE fail-softs it
+        # (draft source only — bad grammar costs the speedup, never the request)
+        opts = generation_options({"response_format": {"type": "gbnf", "grammar": "not a grammar ::="}}, 8)
+        self.assertEqual(opts[3], "not a grammar ::=")
 
 
 class ProtocolTest(unittest.TestCase):

docs/grammar-draft.md

@@ -83,3 +83,27 @@ this implementation adds: forced spans as a **draft source verified in the targe
 own forward** (lossless even under a wrong grammar, composes with MTP/n-gram in one
 union batch), deployed where the win is denominated in expert I/O rather than
 forward passes.
+
+## Server usage: `response_format` (OpenAI API)
+
+The gateway (`openai_server.py`) turns `response_format` into a **per-request**
+grammar carried to the engine over the `SUBMIT` protocol (optional 7th header
+field, `gbytes`; 6-field headers from older clients remain valid — the field is
+additive and back-compatible in both directions):
+
+```jsonc
+{"response_format": {"type": "json_object"}}                       // generic JSON grammar
+{"response_format": {"type": "json_schema",
+                     "json_schema": {"schema": { ... }}}}          // compiled by schema_gbnf.h
+{"response_format": {"type": "gbnf", "grammar": "root ::= ..."}}   // raw GBNF (extension)
+```
+
+Semantics are identical to `GRAMMAR=`/`SCHEMA=`: a **draft source, never a
+sampling constraint**. A schema outside the supported subset, or malformed GBNF,
+costs the speedup — never the request, never the output. Drafting engages for
+greedy requests (`temperature: 0`); sampled requests run undrafted. Compile
+overhead is negligible: ~8 µs/request for a typical schema, ~18 µs at the
+32-level nesting cap (measured, M3 Max). Grammar payloads are capped at 1 MiB.
+As with MTP (#100), a drafted greedy run may differ from an undrafted one in
+near-tie tokens (the verify forward has a different batch shape); each output
+is a valid greedy stream of its own forward shapes.