Reuse the KV prefix in thinking mode on the continuous path

[fabiopili]

I've been chasing this issue after playing with v0.2.3 yesterday and v0.4.0 today. My sessions on pi and Open WebUI were clearly re-processing the whole conversation on each turn.

This PR addresses the issue following an approach already used upstream in antirez/ds4, which this fork's newer continuous-batching path never inherited. I think it also explains Entrpi/ds4-on-spark#4, opened earlier today with little detail and no logs.

Code was prepared with assistance from Claude, but led and reviewed by me. It's worth a second look, though.

What happens today

On the continuous batching path, a conversation in thinking mode re-prefills its entire prompt on every single turn. The prefix cache never hits.

It shows up in the serve log as an unbroken run of cold admissions, where the conversation only ever grows by a few tokens per turn:

cont admit bank=1 pos_base=0 suffix=29234 wall=41.44s
cont admit bank=1 pos_base=0 suffix=29308 wall=40.57s
cont admit bank=1 pos_base=0 suffix=31022 wall=43.10s

pos_base=0 on every turn means nothing was reused. The same conversation sent to the non-thinking deepseek-chat alias reuses normally:

partial fork admit src=0 dst=1 cut=29123 suffix=145
cont admit bank=1 pos_base=29056 suffix=212 wall=0.80s

That is 40.86s versus 0.80s for turn two, on byte-identical prompts, with only the thinking mode differing. Measured on a DGX Spark (GB10) at ctx 524288 with a real coding agent whose preamble is 29,234 tokens.

There is a second symptom with the same cause. A bank that never gets a warm record never becomes reuse-trustworthy, so it is never written to the disk KV cache either, not even on a clean shutdown. With --kv-disk-dir configured and a 64 GiB budget, the directory stayed completely empty across a full day of thinking-mode traffic and four clean stops. This is very likely what Entrpi/ds4-on-spark#4 ("nothing is getting written to disk, and it's causing a lot of prompt reprocessing") is actually reporting: two symptoms, one cause.

Cause

cont_warm_retire() in ds4_server.c bailed out for every thinking row:

if (ds4_think_mode_enabled(j->req.think_mode)) return;

so no warm record was ever built and cont_warm_pick() had nothing to match.

The session path already solves this problem with visible-live checkpoints (build_toolless_thinking_visible_text, thinking_live_visible_prefix_prompt): it keys the checkpoint on the transcript the client will replay, while the payload stays the exact sampled frontier including the hidden reasoning. The continuous batching path, added later, never inherited that idea.

It is not a recent regression. git tag --contains puts the guard in v0.1.0 onward.

The fix

A raw text key is only wrong when the next rendering drops the reasoning bytes, and prompt_preserves_reasoning already tells us when that happens:

• Reasoning preserved (the history carries tool context, so render_chat_prompt_text replays it): the raw key is exactly what comes back, so retire normally.
• Otherwise: key the record on the visible transcript, exactly as the session path does. The payload stays the bank's committed frontier, which cont_warm_pick already pairs with a tokenized text suffix, so the hidden reasoning stays resident in KV and is never re-prefilled.

The visible-key branch takes the same conservative gates as should_remember_thinking_checkpoint(): chat only, tool-less, a clean finish and a closed thinking block. Anything that does not qualify retires without a record, exactly as before.

48 insertions, 13 deletions, one file.

Testing

Machine: NVIDIA DGX Spark (GB10, sm_121, 128 GB unified memory), CUDA backend, driver 580.159.03 / CUDA 13.0, Ubuntu.
Model: DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf (IQ2_XXS/Q2K experts, Q8 dense), DSpark drafter armed, ctx 524288.
Build: make cuda CUDA_ARCH=sm_121 with -DDS4_CUDA_SPARK_HBM_CACHE=1.

DS4_TEST_MODEL=.../DeepSeek-V4-Flash-...-imatrix.gguf make test

passes in full against the patched build, with the model loaded and the GPU otherwise idle (4m28s). Tensor equivalence is exact across all five cases: top1_mismatch=0 worst_rms=0 worst_max_abs=0.

Turn two now reuses in every thinking shape that was previously broken:

Scenario	Before	After
thinking, no tools	pos_base=0	fork admit cached=3419 suffix=9
thinking, 35 tools	pos_base=0	partial truncate admit cut=5267 suffix=21
thinking, reasoning echoed back	pos_base=0	warm admit cached=3443 suffix=9
non-thinking (control)	reuses	reuses, unchanged

The disk tier now receives thinking-mode checkpoints, which it never did before:

kv cache stored tokens=6827 reason=bank-shutdown size=51.42 MiB save=11.9 ms bank persisted bank=0 tokens=6827
kv cache stored tokens=8681 reason=bank-shutdown size=59.18 MiB save=12.7 ms bank persisted bank=1 tokens=8681

Correctness was checked rather than assumed: a needle placed deep in the prompt and asked for only on the reused turn (pos_base=8448, so genuinely served from reused KV) came back exact.

Verified live on the serving stack

The numbers above swap the model id to isolate thinking mode. This is the cleaner comparison: the same real coding agent, the same 29,244 token preamble, thinking mode on both runs, only the binary differs.

Before, on the stock build:

turn 1  cont admit bank=0 pos_base=0 suffix=29244 wall=45.82s
turn 2  cont admit bank=0 pos_base=0 suffix=29321 wall=39.40s

After, with the patch, same client and same conversation shape:

turn 1  cont admit bank=0 pos_base=0 suffix=29244 wall=40.61s
turn 2  partial fork admit src=0 dst=1 cut=29133 suffix=171
        cont admit bank=1 pos_base=29056 suffix=248 wall=0.80s
turn 3  partial truncate admit bank=1 cut=29192 suffix=355
        cont admit bank=1 pos_base=29184 suffix=363 wall=0.93s

Turn one is identical to the token, as it should be. Turn two goes from 39.40s to 0.80s, and reuse holds as the conversation grows, through both admit paths (fork into a free bank, then truncate in place). This client sends 35 tool schemas, so it exercises the prompt_preserves_reasoning branch; the tool-less visible-transcript branch is covered by the table above.

Notes on the other regression tracks

make cpu fails on this branch, but it already fails identically at the unmodified base commit (4880163), so it is not something this change introduced:

ds4.c:28903: error: unknown type name 'ds4_gpu_graph'
ds4.c:28903: error: unknown type name 'payload_region'

I did not run the ds4-bench speed sweep. This change touches server-side warm record bookkeeping at job retirement, not any inference backend or kernel: the only added work is one detokenization pass over the generated tokens for thinking rows, which the non-thinking path already performed on every turn. Happy to run a before/after sweep if you would like it on the record.

Scope and limits

• Streaming, mid-stream client aborts and tool round-trips were all verified separately and were never part of this bug.
• Turns that stop on length retire without a record, matching the session path's existing conservatism. Loosening that is a separate question.
• The engine still validates every warm admit against its own committed history, so a stale or mismatched record degrades to a cold reset rather than serving wrong state.

[Entrpi]

Excellent find, and a genuinely well-built fix. Thank you for the depth of
the report, the two-symptom diagnosis (cold re-prefills plus the empty
disk tier, which indeed explains ds4-on-spark#4), and for testing both
branches before sending it.

I verified the change independently before merging:

Code review. The mechanism checks out end to end against the tree.
The visible key your toolless branch builds is byte-identical to what
render_chat_prompt_text emits when it replays a prior toolless thinking
turn (</think> + content + end marker), and the preserved branch's raw
key matches the verbatim reasoning replay the renderer does under tool
context. The part I most wanted to convince myself of was safety of the
shorter-key/longer-payload pairing, and it holds by construction: a full
record match reuses the bank's exact committed tokens under the engine's
own prefix validation, and the partial path cuts at the token-LCP against
committed history, so a cut can never land inside the hidden reasoning
region. Clients that do not echo reasoning_content back degrade to
reusing the shared preamble, which is still the bulk of the win. The
finish gate matches the session path's conservatism (EOS only, engine
passes hit_eos at every site). One pathological corner I noted for the
record: a generation containing a second literal </think> keys on the
first close while the response parser splits on the last, which can only
produce a record that never matches. Harmless, and not worth complicating
the code for.

Behavioral gate (DGX Spark GB10, your branch at 6f220b4, fresh
--kv-disk-dir): four legs, all green.

• Toolless thinking: turn 2 fork admit cached=4131 suffix=16, and a
  needle planted deep in the turn-1 preamble came back exact from the
  reused KV.
• Thinking + tool schema + reasoning_content echoed: fork admit cached=4468 suffix=17, needle exact.
• Non-thinking control: reuses as before. A nice by-design bonus your
  change enables: a deepseek-chat continuation of a deepseek-reasoner
  conversation byte-matched the thinking bank's visible record and reused
  it directly.
• Clean SIGTERM stop: 12 bank persist lines, 6 files in the disk dir
  (previously always empty for thinking traffic), and after a reboot the
  next turn restored the bank from disk (bank restore hit tokens=4222 load=66.9 ms) and answered from it correctly.

Zero cont failures, zero serial fallbacks, zero illegal accesses across
all legs.

Merging now that the v0.4.1 tag has landed (a speculation-controller
recalibration touching only ds4.c, so no conflict) and cutting this as
v0.4.2 with the gate above folded into the release battery. I will
credit you in the changelog and release notes.

[fabiopili]

Happy to help. Thanks for the work done here. I'm having a lot of fun running DeepSeek v4 Flash on a GB10 system thanks to this project. The performance and quality of the model is incredibly good.