From d14b09dbf468ca0d77853b6c8e40135329cba6bc Mon Sep 17 00:00:00 2001 From: JustVugg Date: Sun, 2 Aug 2026 21:19:06 +0200 Subject: [PATCH 1/2] feat(kimi): reuse the attention state a previous turn already built; stop leaking Lc/Rc MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Same defect as the Inkling PR this stacks on: a chat client resends the whole transcript each turn and kimi_k3.c re-processed turns 1..N-1 from scratch, so the cost of a message grew with the conversation and every replayed position pulled its experts off disk again. K3 needs the shared record more than the other engines do, not less. There is no single KV to inspect: 69 of its layers are KDA (Kimi Delta Attention, a RECURRENT linear-attention state) and only the 24 MLA layers keep Lc/Rc. So "how much of this prompt does the current state already cover" is not something that can be read off any buffer — it has to be recorded where the tokens are fed, which is what kv_prefix.h does. The prefill loop already carried absolute positions (step_chunk(m, ids+i, i, C)), so reuse is starting it at `reuse` instead of 0. SEPARATELY, kv_alloc leaked on every request. Serve calls it per turn, and it callocs Lc/Rc over the previous pointers without freeing them: m->Lc=calloc(c->n_layers,sizeof(float*)); m->Rc=calloc(c->n_layers,sizeof(float*)); That is n_layers x max_t x (kv_lora + qk_rope) floats per turn — hundreds of MB over a conversation on the 24 MLA layers at 4k context. The fix belongs here rather than in its own PR because prefix reuse REQUIRES keeping those buffers across turns: growing them discards the positions the record describes, so the lifetime question had to be answered either way. NOT MEASURED HERE. Kimi K3 is ~1.6 TB and I do not have the checkpoint, so there is no speed number and no end-to-end identity run for this engine. What is verified: it compiles, tests/test_kv_prefix.c covers the decision logic exhaustively, and the same mechanism measured 5.23x on DeepSeek V4 with byte-identical output. Anyone with the weights: the numbers would be welcome, and K3_PREFIX_LOG=1 prints what was reused. Stacked on feat/inkling-kv-prefix, which adds kv_prefix.h. Co-Authored-By: Claude Opus 5 (1M context) --- c/kimi_k3.c | 46 ++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 42 insertions(+), 4 deletions(-) diff --git a/c/kimi_k3.c b/c/kimi_k3.c index 78668e9c6..ca3a59cef 100644 --- a/c/kimi_k3.c +++ b/c/kimi_k3.c @@ -81,7 +81,8 @@ #include "tok.h" #include "quant.h" #include "omp_tune.h" -#include "route_trace.h" /* shared routing telemetry (#700) */ +#include "route_trace.h" +#include "kv_prefix.h" /* KV prefix reuse (shared) */ /* ---------- config ---------- */ typedef struct { @@ -154,6 +155,11 @@ typedef struct { float **cwq, **cwk, **cwv; /* conv windows [proj*conv_k], oldest first */ /* MLA cache */ float **Lc, **Rc; int max_t; + /* KV prefix reuse: what the current state was built from (kv_prefix.h). + * K3 has no single KV to inspect — 69 KDA layers carry a RECURRENT state + * and only the 24 MLA layers keep Lc/Rc — so an explicit record of the + * tokens fed is the only description of it that cannot drift. */ + kv_prefix kvp; /* experts */ ERef *eref; /* [n_layers][n_experts] (dense rows zeroed) */ LCache *ecache; @@ -1250,14 +1256,28 @@ static float *step_chunk(Model *m, const int *ids, int pos0, int C){ } m->t_head+=now_s()-t0; } + /* record what was just fed, at the positions it went to (kv_prefix.h) */ + kv_prefix_record(&m->kvp, ids, pos0, C); free(hidden);free(bres);free(prefix);free(nrm);free(att);free(mix);free(mlp); return logits; } static void kv_alloc(Model *m, int max_t){ - Cfg *c=&m->c; m->max_t=max_t; + Cfg *c=&m->c; + /* Serve calls this once per request. It used to calloc Lc/Rc over the old + * pointers without freeing them, leaking n_layers x max_t x (kv_lora + + * qk_rope) floats every turn — on K3's 24 MLA layers at 4k context that is + * hundreds of MB per conversation. Keep the buffers when they are already + * big enough, which is also what makes prefix reuse possible: growing them + * discards the positions fed[] describes. */ + if(m->Lc && max_t<=m->max_t) return; + if(m->Lc) for(int i=0;in_layers;i++){ free(m->Lc[i]); free(m->Rc[i]); } + free(m->Lc); free(m->Rc); + m->max_t=max_t; m->Lc=calloc(c->n_layers,sizeof(float*)); m->Rc=calloc(c->n_layers,sizeof(float*)); + /* the record is sized with the KV it describes; growing discards it */ + kv_prefix_alloc(&m->kvp,max_t); for(int i=0;in_layers;i++) if(!m->L[i].kda){ m->Lc[i]=falloc((int64_t)max_t*c->kv_lora); m->Rc[i]=falloc((int64_t)max_t*c->qk_rope); @@ -1416,6 +1436,7 @@ typedef struct { static void model_state_reset(Model *m){ Cfg *c=&m->c; + kv_prefix_clear(&m->kvp); /* the record describes the state we are dropping */ for(int i=0;in_layers;i++){ if(m->L[i].kda){ memset(m->kstate[i],0,(size_t)c->kda_heads*c->kda_hd*c->kda_hd*sizeof(float)); @@ -1483,14 +1504,31 @@ static void serve_one(Model *m, Tok *T, ServeReq *q){ fflush(stdout); free(ids); return; } printf("ACCEPT %s %d\n",q->id,np); fflush(stdout); - model_state_reset(m); + /* KV PREFIX REUSE (#639 for GLM; this engine re-prefilled every turn). + * A chat client resends the whole transcript each turn, so turn N used to + * re-process turns 1..N-1 from scratch — the cost of a message grew with + * the conversation, and every replayed position pulled its experts off + * disk again. When this prompt begins with the sequence the state already + * holds, that state IS the state at that position: keep it and prefill + * only the tail. This is why kv_alloc above must run FIRST and must keep + * its buffers: growing them discards the positions fed[] describes. + * At least one new token is required, since the state cannot be rewound. + * Either the reused positions are token-identical or nothing is reused; + * the emitted tokens are unchanged in both cases. */ kv_alloc(m,np+q->max_tok+8); + int reuse=kv_prefix_reuse(&m->kvp, ids, np); + if(!reuse) model_state_reset(m); + else if(getenv("K3_PREFIX_LOG")) + fprintf(stderr,"[PREFIX] reusing %d of %d prompt tokens (%.0f%%)\n", + reuse,np,100.0*reuse/np); int chunk=getenv("K3_CHUNK")?atoi(getenv("K3_CHUNK")):32; if(chunk<1) chunk=1; if(chunk>512) chunk=512; double t0=now_s(), a0=m->t_attn, e0=m->t_moe, d0=m->t_eload, h0=m->t_head; uint64_t hit0=m->hits, miss0=m->miss; float *lo=NULL; - for(int i=0;i Date: Sun, 2 Aug 2026 23:44:39 +0200 Subject: [PATCH 2/2] fix(kimi): grow Lc/Rc instead of restarting them, or reuse never fires Same defect CI caught on the Inkling side, in the same place. Freeing and re-allocating the MLA buffers on a longer prompt discards every position already computed -- and a conversation's prompt is longer every turn, so the state was thrown away immediately before the point of using it. Reuse could never fire in the one case it exists for. Lc/Rc are laid out [position][kv_lora] and [position][qk_rope], so unlike inkling's head-major K/V this grow is a straight prefix copy with no re-layout. The 69 KDA layers need nothing here: their recurrent state does not scale with max_t and survives on its own. K3_PREFIX_LOG now also reports a refusal and the state behind it (held, cap, prompt, diverged). On the Inkling side that diagnostic is what distinguished 'the engine decided not to reuse' from 'the harness lost the output', and it turned an unreadable CI failure into a one-line diagnosis. Co-Authored-By: Claude Opus 5 (1M context) --- c/kimi_k3.c | 55 ++++++++++++++++++++++++++++++++++++++++------------- 1 file changed, 42 insertions(+), 13 deletions(-) diff --git a/c/kimi_k3.c b/c/kimi_k3.c index ca3a59cef..d93602b31 100644 --- a/c/kimi_k3.c +++ b/c/kimi_k3.c @@ -1264,24 +1264,43 @@ static float *step_chunk(Model *m, const int *ids, int pos0, int C){ static void kv_alloc(Model *m, int max_t){ Cfg *c=&m->c; - /* Serve calls this once per request. It used to calloc Lc/Rc over the old - * pointers without freeing them, leaking n_layers x max_t x (kv_lora + - * qk_rope) floats every turn — on K3's 24 MLA layers at 4k context that is - * hundreds of MB per conversation. Keep the buffers when they are already - * big enough, which is also what makes prefix reuse possible: growing them - * discards the positions fed[] describes. */ + /* Serve calls this once per request, and it used to calloc Lc/Rc over the + * old pointers without freeing them: n_layers x max_t x (kv_lora + + * qk_rope) floats leaked every turn, hundreds of MB over a conversation on + * the 24 MLA layers at 4k context. + * + * GROW, DO NOT RESTART. Freeing and re-allocating also discards every + * position already computed, which defeats KV prefix reuse in the one case + * it exists for: a conversation whose prompt is longer every turn asks for + * a larger max_t every turn, so the state would be thrown away immediately + * before the point of using it. (Caught by CI on the Inkling side, where + * the same shape of bug sat in the same place.) + * + * Lc/Rc are laid out [position][kv_lora] and [position][qk_rope], so unlike + * inkling's head-major K/V a grow is a straight prefix copy — no re-layout. + * The 69 KDA layers are untouched here: their recurrent state does not + * scale with max_t and survives on its own. */ if(m->Lc && max_t<=m->max_t) return; - if(m->Lc) for(int i=0;in_layers;i++){ free(m->Lc[i]); free(m->Rc[i]); } - free(m->Lc); free(m->Rc); + + float **oldL=m->Lc, **oldR=m->Rc; + int keep=(m->Lc && m->kvp.len>0 && m->kvp.len<=max_t) ? m->kvp.len : 0; + m->max_t=max_t; m->Lc=calloc(c->n_layers,sizeof(float*)); m->Rc=calloc(c->n_layers,sizeof(float*)); - /* the record is sized with the KV it describes; growing discards it */ - kv_prefix_alloc(&m->kvp,max_t); for(int i=0;in_layers;i++) if(!m->L[i].kda){ m->Lc[i]=falloc((int64_t)max_t*c->kv_lora); m->Rc[i]=falloc((int64_t)max_t*c->qk_rope); + if(keep){ + memcpy(m->Lc[i], oldL[i], (size_t)keep*c->kv_lora*sizeof(float)); + memcpy(m->Rc[i], oldR[i], (size_t)keep*c->qk_rope*sizeof(float)); + } } + if(oldL) for(int i=0;in_layers;i++){ free(oldL[i]); free(oldR[i]); } + free(oldL); free(oldR); + + /* the record describes those same positions, so it survives with them */ + if(!kv_prefix_grow(&m->kvp,max_t,keep)) kv_prefix_clear(&m->kvp); } typedef struct { float p; int id; } SampleProb; @@ -1517,10 +1536,20 @@ static void serve_one(Model *m, Tok *T, ServeReq *q){ * the emitted tokens are unchanged in both cases. */ kv_alloc(m,np+q->max_tok+8); int reuse=kv_prefix_reuse(&m->kvp, ids, np); + if(getenv("K3_PREFIX_LOG")){ + /* Report the decision either way, with the state behind a "no". + * "It did not get faster" is otherwise the same observation as + * "reuse is not wired up" — for a user as much as for a test. */ + if(reuse) + fprintf(stderr,"[PREFIX] reusing %d of %d prompt tokens (%.0f%%)\n", + reuse,np,100.0*reuse/np); + else + fprintf(stderr,"[PREFIX] no reuse: held=%d cap=%d prompt=%d%s\n", + m->kvp.len,m->kvp.cap,np, + (m->kvp.len>0 && m->kvp.len512) chunk=512; double t0=now_s(), a0=m->t_attn, e0=m->t_moe, d0=m->t_eload, h0=m->t_head;