Skip to content

server: support live prefix rewind for GLM-5.3 and preserve common prefix on forks - #1005

Open
aaa2015 wants to merge 1 commit into
antirez:mainfrom
aaa2015:fix/server-live-prefix-rewind
Open

server: support live prefix rewind for GLM-5.3 and preserve common prefix on forks#1005
aaa2015 wants to merge 1 commit into
antirez:mainfrom
aaa2015:fix/server-live-prefix-rewind

Conversation

@aaa2015

@aaa2015 aaa2015 commented Sep 9, 2026

Copy link
Copy Markdown

Summary

In ds4-server, live session memory can often be reused when a new request shares a prefix with the active session. However, under high-intensity multi-turn agent sessions (such as tool-calling workflows), we observed frequent reason=token-mismatch cache misses and massive disk spills (e.g. live kv cache miss live=165755 prompt=4977 common=4973 reason=token-mismatch).

Investigation uncovered two root causes:

  1. Engine capability gatekeeper lockout: ds4_server.c:12371 strictly checked ds4_engine_is_glm_dsa(s->engine), omitting GLM-5.3, which shares DS4_MODEL_FAMILY_GLM_DSA but is a distinct variant.
  2. Strict equality requirement on fork: live_prefix_rewind_target() strictly asserted if (common != prompt_len) return -1;. When a multi-turn conversation or agent tool-call branched with common < prompt_len (e.g. 4,973 tokens matched out of 4,977, with 4 trailing new tokens), the entire resident session of 165K tokens was rejected, triggering an expensive disk cache eviction and a full prompt re-prefill.

Proposed Changes

  1. ds4.h & ds4.c: Added bool ds4_engine_can_rewind(ds4_engine *e) to export whether the engine supports session rewinds, so the scheduler asks a question about rollback rather than about GLM.
  2. ds4_server.c:
    • Generalized live_prefix_rewind_target(): when common < prompt_len, it rewinds to common, keeping the shared tokens resident and requiring ds4_session_sync() to only evaluate the suffix tokens. When common == prompt_len, it continues to rewind to prompt_len - 1 to re-evaluate the final token for logits generation.
    • Updated test_live_prefix_rewind_target() with regression cases covering fork rollbacks and matching the observed multi-turn token counts.

Note on Metal DSpark

An earlier revision of this PR admitted Metal DSpark through the new helper. That is removed here, because ds4_session_rewind() has no DSpark path yet:

  • the rewind clears checkpoint_valid, so rewind_valid fails and the request falls back to a full rebuild — the exact cost the helper is meant to avoid;
  • kv_cache_store_current(s, slot, "evict") then runs on the already-truncated session, so the persisted snapshot is shorter than the one that was resident.

Both are regressions relative to the current behaviour, where DSpark never enters the rewind branch. The DSpark branch is worth enabling together with the engine-side snapshot reuse that makes ds4_session_rewind() able to restore its state.

Verification

Machine: Apple M4, macOS 26.4.1, 16 GB. Backend: Metal.

make ds4_test && ./ds4_test --server
→ server: OK
→ ds4 tests: ok

make test-session-state
→ session state tests: ok
→ TP command tests: ok

make tests/test_session_state_gpu && ./tests/test_session_state_gpu
→ session state tests: ok

test_live_prefix_rewind_target() covers the fork-rollback cases added here,
including the observed multi-turn token counts (4973 of 4977 shared, 165755
resident). The model-backed suites (--logprob-vectors, --long-context, ...)
were not run because this machine has no GGUF checked out.

@aaa2015
aaa2015 force-pushed the fix/server-live-prefix-rewind branch from 352ac5b to cdcd890 Compare September 9, 2026 23:15
…efix on forks

A live session could only be rewound when the prompt was fully cached, so any
fork that shared a long prefix but diverged near the end threw away the whole
resident checkpoint and paid a full prefill.  Accept a partial rewind: when the
common prefix is shorter than the prompt, rewind to it and let the scheduler
evaluate only the suffix, and when the whole prompt is cached keep rewinding to
the last token so it is resampled.

Expose the capability as ds4_engine_can_rewind() instead of testing the engine
family at the call site, so the scheduler asks a question about rollback rather
than about GLM.

Metal DSpark is deliberately not admitted yet: ds4_session_rewind() has no
DSpark path, so a rewind there would drop the checkpoint, log a rebuild, and
persist a shorter snapshot than the one that was already resident.  Enable it
together with the engine-side snapshot reuse.
@aaa2015
aaa2015 force-pushed the fix/server-live-prefix-rewind branch from cdcd890 to e4cb4a7 Compare September 10, 2026 00:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant