Skip to content

[review-finding] DECODE_APPLY restore (PR #93): 3 remaining gaps vs STATE_PUT — failure-path slot poison, dropped checkpoint, n_prompt_tokens_cache source #94

Description

@ddvnguyen

Context

PR #93 ("DECODE_APPLY KV restore check was inverted") fixes the restore-status polarity and ports trailing-logits injection. Review of the updated head (0a0fa9a72) found 3 remaining gaps in the DECODE_APPLY restore flow, each a divergence from the proven STATE_PUT sibling (server-context.cpp).

Gap 1 — failure path leaves the slot poisoned (fix before merge)

STATE_PUT failure branch clears slot state explicitly:

  • server-context.cpp:3395-3404: on n_read == 0slot->prompt.tokens.clear(); slot->prompt.checkpoints.clear(); slot->n_prompt_tokens_cache = 0; llama_memory_seq_rm(...). Comment: "n_past > 0 with no KV cells → pos_min == -1 abort on the next decode that touches this slot."

DECODE_APPLY status == 0 branch (server-context.cpp:4541-4558) only resets reserved_for_decode_id and writes the error entry. But the v2 header parse has already re-populated slot->prompt.tokens (~4494-4497) before llama_state_seq_set_data runs. On failure the slot keeps tokens with no matching KV cells. update_slots computes n_past from prompt.tokens.get_common_prefix(input_tokens) (~6049), so a later request sharing the prefix gets n_past > 0 with empty KV → the exact pos_min == -1 abort STATE_PUT guards against. The poisoned slot persists until a fresh DECODE_APPLY/PREFILL calls prompt_clear.

Fix: mirror STATE_PUT cleanup in the status == 0 branch (clear tokens/checkpoints, reset n_prompt_tokens_cache, llama_memory_seq_rm).

Gap 2 — blob checkpoint silently dropped (latent; hybrid/recurrent models)

STATE_PUT restores the native checkpoint from the v2 blob into slot->prompt.checkpoints (server-context.cpp:3447-3460), with the note that the native checkpoint (pos_max at n-4) rewinds recurrent state cleanly, and the old fabricated checkpoint "corrupts hybrid/recurrent model output". STATE_PUT falls back to create_checkpoint(*slot, 0, 0, n-1) (~3458) when absent.

DECODE_APPLY skips the checkpoint bytes when advancing past the header (server-context.cpp:4507-4519) and never registers them. It relies on just_restored (~6326-6331) which skips checkpoint search for the first decode — OK for the current dense/MoE COMBINED target (dense-27b-combined), but for any recurrent/hybrid model the KV cache is restored without its memory checkpoint → silent corruption.

Fix: parse + register the native checkpoint (or at minimum create_checkpoint fallback), mirroring STATE_PUT.

Gap 3 — n_prompt_tokens_cache source inconsistency (minor)

  • STATE_PUT: slot->n_prompt_tokens_cache = hdr_n_tok (server-context.cpp:3437)
  • DECODE_APPLY: slot->n_prompt_tokens_cache = blob_n_past (server-context.cpp:4581)

Both are written by PREFILL as n_tokens vs tokens.size() (~4048-4049), equal in the normal flow today. Harmless now, but the two restore paths should read the same header field to avoid a future off-by-one.

Verification

Fork-side only. Requires a build + live-rig E2E per project workflow. The changes are localized to the DECODE_APPLY restore block, mirroring existing STATE_PUT code. Closes nothing on its own — it hardens PR #93.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    review-findingFinding created from code review

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions