llama : fix K/V and recurrent state cleanup after failed restores

[CHIPMUNK-T0T]

Overview

A failed per-sequence state restore rolls back sequence metadata with seq_rm(), but does not clear K/V or recurrent-state tensor data already written by the failed restore. In contrast, the corresponding whole-context failure path calls clear(true), so the existing rollback is asymmetric.

This PR adds failure containment for failed state restores in both attention and recurrent memory.

• llama_kv_cache and llama_memory_recurrent clear data written by a failed restore.
• Buffer-backed readers discard deferred tensor writes to prevent post-rollback flushes.
• Whole-context buffer restore uses the same deferred-write cleanup.

Per-sequence failures before modifying the target sequence leave it unchanged. Failures after modification remove the affected state and prevent tensor writes from that restore attempt from surviving rollback.

A failed per-sequence restore leaves other sequences and subsequent inference unaffected. A failed whole-context restore still clears the whole memory through clear(true) as before, so all sequences are emptied rather than left holding stale data.

This does not provide transactional rollback to the previous sequence state. Return values and server HTTP error behavior are unchanged.

Fixes #27068

Additional information

Failure containment for file and buffer restores

File and buffer restores require different cleanup on failure.

File restores may already have written tensor data when validation fails, so the affected K/V or recurrent-state ranges must be cleared. Buffer restores can still have deferred tensor writes pending at that point, so those writes must be discarded before the reader is destroyed after failure handling.

The memory cleanup and reader discard() therefore address separate parts of the same failure path rather than providing redundant rollback.

Mutation testing confirms that each part is required: removing K/V or recurrent cleanup breaks the corresponding file-restore case, removing host-reader discard() breaks buffer restore for both memory types, and removing whole-context discard() breaks whole-context restore for both. Removing device-reader discard() reaches GGML_ASSERT(buf_size == 0) during reader destruction.

Why recurrent memory is included

The original report focused on the K/V cache, but the Qwen3.5 hybrid reproduction can fail in llama_memory_recurrent, where the same per-sequence rollback asymmetry exists.

The fix therefore applies the same containment rule to both attention and recurrent memory, with independent regression coverage for each implementation.

The recurrent cleanup takes a (head, cell_count) range, and head is not reliable after a failure: state_read_meta() removes the target sequence first, then returns early when cell_count == 0 and returns false when find_slot() fails, leaving head on another sequence's cells. The cleanup is therefore gated on state_read_meta() having succeeded and uses the head captured before the rollback; on that path head + cell_count <= size is already asserted in state_read_meta().

Why whole-context buffer restore is also changed

The original report treated whole-context restore as contained by clear(true). That handles tensor data already applied to memory, but not deferred buffer writes that can still be flushed after the rollback.

Whole-context buffer restore therefore uses the same discard() handling so that pending writes cannot overwrite the cleanup performed on failure. The corresponding mutation test fails when this discard() is removed.

Tensor clearing

Failed file restores need to clear backend tensor ranges, but ggml_backend_tensor_memset() is not supported by every buffer backend and can assert on unsupported implementations.

The cleanup instead uses ggml_backend_tensor_set() through the shared internal llama_clear_tensor_data() helper. The helper lives in llama-impl because the same operation is required by both attention and recurrent memory.

The cleanup only runs on restore failure and adds no cost to normal decoding or successful restore.

Tests

Added two regression tests covering llama_kv_cache and llama_memory_recurrent.

They exercise failed per-sequence file/buffer restore, fragmented restore with another live sequence, the guarded ON_DEVICE restore path, and whole-context buffer restore, with Flash Attention both disabled and enabled.

Without the changes in this PR the new tests fail: the file and buffer restore cases produce logits that do not match the clean baseline, and the ON_DEVICE case aborts in ~llama_io_read_device() at GGML_ASSERT(buf_size == 0). They pass with the changes applied.

On the current d775b8967 base:

ctest -R "state|recurrent|generate-models"
8/8 passed

The 8 matching tests include the two new regression tests and existing related state tests.

CUDA ctest validation also passed 8/8 on the same base, with the K/V cache and the recurrent state allocated on the device.

Server-level fault injection was verified on an RTX 4070 with CUDA and Flash Attention enabled using Gemma 3 (llama_kv_cache_iswa) and Qwen3.5 (llama_memory_hybrid): the corrupted restore poisons subsequent inference before the fix, while inference remains clean after the fix.

Non-goals

This PR provides failure containment, not transactional restore. If a failed restore has already modified the target sequence, its previous contents are not restored.

It also does not provide atomic rollback across multiple components of hybrid/iSWA memory.

DSV4-specific restore paths are not changed here; related DSV4 rollback work is tracked separately in #26756.

Related ON_DEVICE issue

The separate ON_DEVICE pre-validation failure is tracked in #27439, with a fix proposed in draft PR #27487 from another contributor.

Known limitations

The regression fixtures directly cover llama_kv_cache and llama_memory_recurrent. Hybrid and iSWA paths are covered by the server-level Gemma 3 and Qwen3.5 validation above.

Requirements

• I have read and agree with the contributing guidelines

• AI usage disclosure: YES

  • Motivation & Design: I found this failure through slot save / restore fault injection against a stock llama-server build, reported it in Eval bug: failed slot restore leaves corrupted K/V data that breaks subsequent inference #27068, and defined the scope as failure containment rather than transactional restore.
  • Investigation & Code Review (Fable 5 / Opus 5 / GPT-5.6 Sol): AI agents were used to investigate the state-restore paths of the attention and recurrent memory implementations and assist with code review under my instructions.
  • Implementation (Opus 5 / GPT-5.6 Sol): AI assistance was used to draft parts of the implementation and tests based on my technical instructions.
  • Draft Translation (Opus 5): The PR text was written by me in Japanese first, with AI used to assist with the English translation.
  • Validation & Responsibility: I manually reviewed the submitted changes, ran the local tests, mutation tests, and server-level fault injection, and take full responsibility for the submitted changes.