fix(dflash): stabilize long-prompt KVarN serving

Preserve the target model's absolute prompt position when DFlash flushes only a long-prompt prefill suffix into the cross-attention ring. This keeps drafter RoPE and masks aligned after full prompt reprocessing while retaining the bounded cross window.

Bound SWA/hybrid prompt checkpoint memory by the prompt-cache RAM limit, prune checkpoints copied into prompt-cache entries, allow oversized cache entries to be fully evicted, and make explicit prompt clears release retained checkpoint buffers.

Add regression coverage for absolute prefill commits, checkpoint-budget pruning, prompt clear ownership, and the server flush contract. Update the v0.3.2 changelog with the Qwen long-prompt garbage-output fix and Gemma follow-up memory/crash fix.

Author: Anbeeld

CHANGELOG.md

@@ -8,6 +8,8 @@
 - Added experimental KVarN KV-cache compression as cache-type names: `--cache-type-k kvarn2|kvarn3|kvarn4` and `--cache-type-v kvarn2|kvarn3|kvarn4`. KVarN is kept to the target context, validates unsupported placements at runtime, uses non-unified CUDA streams, preserves stream-scoped backend state, and is wired for the current Qwen3.6 and Gemma4 paths.
 - Hardened KVarN memory ownership and server rollback behavior. Full prompt-cache sequence clears are now allowed, arbitrary old suffix trims are rejected before mutation, composite memory preflights child removals before changing either side, and DFlash prompt-checkpoint reuse now skips or fully reprocesses prompts when the active memory cannot safely roll back.
 - Fixed the Gemma4 and Qwen3.6 DFlash + KVarN crash class seen after prompt-cache reuse. The server no longer asks KVarN-backed target memory to delete compressed historical groups during checkpoint restore, so repeated chat requests fall back to safe prefill instead of aborting in `seq_rm`.
+- Fixed long-prompt DFlash prefill after full prompt reprocessing. DFlash now keeps the drafter on the target model's absolute prompt timeline when flushing only the final cross-attention window, preventing invalid reduced-logit drafts and Qwen3.6 garbage output after long prompts.
+- Bounded prompt checkpoint memory for SWA/hybrid prompt-cache paths. Active checkpoints are byte-budgeted, prompt-cache saves prune or skip over-limit checkpoint state, cache eviction may remove every oversized entry, and explicit slot clears now release retained checkpoint buffers. This fixes the Gemma4 DFlash follow-up crash/memory spike where multi-GiB checkpoint lists were copied into prompt cache.
 - Reduced KVarN side effects outside KVarN runs. KVarN is disabled for draft and auxiliary contexts, iSWA now keeps the SWA side on the compact upstream cache instead of allocating a full-size KVarN SWA cache, and DFlash backup streams are gated by target architecture so Gemma recurrent paths do not allocate unused backup state.
 - Simplified the public KVarN surface while it is still experimental. The extra tuning flags were removed; the only user-facing selection is the cache type itself. Internal defaults such as sink tokens and Sinkhorn iterations now come from the selected `kvarnN` type.
 - Improved KVarN CUDA behavior and validation. Cross-stream state copies are explicitly ordered, CUDA materialization group counts are precomputed, `llama-bench` accepts KVarN cache-type names for benchmarking, and regression coverage now checks stream synchronization, sequence-removal capability, prompt-cache rollback safety, iSWA sizing policy, unsupported runtime policy, and argument parsing.

common/speculative.cpp

@@ -241,6 +241,21 @@ common_dflash_ring_write common_dflash_ring_write_plan(int ring_size, int ring_p
     return { normalized_pos, ring_size, skip };
 }
 
+int common_dflash_prefill_committed_after_flush(
+        int current_committed,
+        int n_written,
+        int commit_end_pos) {
+    if (n_written <= 0) {
+        return current_committed;
+    }
+
+    if (commit_end_pos > 0) {
+        return commit_end_pos;
+    }
+
+    return current_committed + n_written;
+}
+
 static bool common_speculative_are_compatible(
     const llama_model * model_tgt,
     const llama_model * model_dft) {
@@ -369,7 +384,7 @@ struct common_speculative_impl {
 
     virtual void update_logits_by_indices(llama_context * /*ctx*/, const std::vector<int> & /*capture_indices*/) {}
 
-    virtual int flush_prefill(int /*src_offset*/ = 0, int /*n_tokens*/ = 0) { return 0; }
+    virtual int flush_prefill(int /*src_offset*/ = 0, int /*n_tokens*/ = 0, int /*commit_end_pos*/ = 0) { return 0; }
 
     virtual int prepare_batch_draft(llama_context * /*ctx_dft_ext*/) { return -1; }
 
@@ -2663,7 +2678,7 @@ struct common_speculative_impl_dflash : public common_speculative_impl {
         return true;
     }
 
-    int flush_prefill(int src_offset = 0, int n_tokens = 0) override {
+    int flush_prefill(int src_offset = 0, int n_tokens = 0, int commit_end_pos = 0) override {
         llama_dflash_set_active_slot(ctx_tgt, seq_id);
 
         prefill_flush_called = true;
@@ -2785,10 +2800,10 @@ struct common_speculative_impl_dflash : public common_speculative_impl {
                     ? (writes_cpu_ring ? "cpu_ring+gpu_ring" : "gpu_ring")
                     : (writes_cpu_ring ? "cpu_ring" : "none");
 
-            LOG_INF("dflash prefill flush: capture_source=%s ring_dst=%s captured=%lld offset=%d n_tokens=%d to_write=%d n_src_layers=%d prefill_flushed=%d ring_write_pos=%d ring_filled=%d committed_len=%d\n",
+            LOG_INF("dflash prefill flush: capture_source=%s ring_dst=%s captured=%lld offset=%d n_tokens=%d to_write=%d n_src_layers=%d prefill_flushed=%d ring_write_pos=%d ring_filled=%d committed_len=%d commit_end_pos=%d\n",
                     capture_source, ring_dst, (long long) captured, offset, n_tokens,
                     to_write, n_src_layers, (int) prefill_flushed, ring_write_pos,
-                    ring_filled, committed_len);
+                    ring_filled, committed_len, commit_end_pos);
         }
 
         if (!prefill_flushed) {
@@ -2808,7 +2823,8 @@ struct common_speculative_impl_dflash : public common_speculative_impl {
             return 0;
         }
 
-        committed_len += actual_written;
+        committed_len = common_dflash_prefill_committed_after_flush(
+                committed_len, actual_written, commit_end_pos);
         update_drafter_kv_cache(actual_written);
         prefill_flushed = true;
         prefill_flush_written += actual_written;
@@ -4849,13 +4865,13 @@ void common_speculative_rollback_dft(common_speculative * spec, llama_seq_id seq
     }
 }
 
-int common_speculative_flush_prefill(common_speculative * spec, int src_offset, int n_tokens) {
+int common_speculative_flush_prefill(common_speculative * spec, int src_offset, int n_tokens, int commit_end_pos) {
     if (spec == nullptr) {
         return 0;
     }
     int total_written = 0;
     for (auto & impl : spec->impls) {
-        total_written += impl->flush_prefill(src_offset, n_tokens);
+        total_written += impl->flush_prefill(src_offset, n_tokens, commit_end_pos);
     }
     return total_written;
 }

common/speculative.h

@@ -114,6 +114,11 @@ void common_speculative_print_stats(const common_speculative * spec);
 
 common_dflash_ring_write common_dflash_ring_write_plan(int ring_size, int ring_pos, int n_tokens);
 
+int common_dflash_prefill_committed_after_flush(
+        int current_committed,
+        int n_written,
+        int commit_end_pos);
+
 bool common_dflash_prefill_capture_complete_for_test(int captured, int requested);
 
 bool common_dflash_cpu_ring_valid_after_write_for_test(
@@ -189,7 +194,7 @@ void common_speculative_draft_batch(
 void   common_speculative_update_logits(common_speculative * spec, llama_context * ctx, const llama_tokens & batch_tokens, int n_accepted);
 void   common_speculative_update_logits_deferred_dflash_kv(common_speculative * spec, llama_context * ctx, const llama_tokens & batch_tokens, int n_accepted);
 void   common_speculative_update_logits_by_indices(common_speculative * spec, llama_context * ctx, const std::vector<int> & capture_indices);
-int    common_speculative_flush_prefill(common_speculative * spec, int src_offset = 0, int n_tokens = 0);
+int    common_speculative_flush_prefill(common_speculative * spec, int src_offset = 0, int n_tokens = 0, int commit_end_pos = 0);
 void   common_speculative_set_prefill_capture_enabled(common_speculative * spec, bool enabled);
 void   common_speculative_discard_dflash_state(common_speculative * spec, const char * reason);
 void   common_speculative_note_prefill_suffix_scheduled(common_speculative * spec);

tests/CMakeLists.txt

@@ -270,6 +270,7 @@ set_property(TEST test-kvarn PROPERTY LABELS main)
 llama_build_and_test(test-adaptive-dm.cpp)
 llama_build_and_test(test-server-prompt-checkpoint.cpp)
 target_include_directories(test-server-prompt-checkpoint PRIVATE ${PROJECT_SOURCE_DIR}/tools/server ${PROJECT_SOURCE_DIR}/tools/mtmd)
+target_link_libraries(test-server-prompt-checkpoint PRIVATE server-context)
 llama_build_and_test(test-dflash-ring.cpp)
 llama_build_and_test(test-cpu-out-prod.cpp)
 llama_build_and_test(test-cpu-fattn-support.cpp)

tests/test-dflash-plumbing.cpp

@@ -1599,8 +1599,15 @@ int main(int argc, char ** argv) {
     // Server should pass the planned source offset. GPU staging normalizes to
     // window-relative offset 0 inside flush_prefill(), while CPU fallback needs
     // the original sub-batch offset.
-    ok &= expect(server_context.find("common_speculative_flush_prefill(pf.spec, pf.span.src_offset,") != std::string::npos,
-        "server must pass planned source offset to prefill flush");
+    const size_t flush_call = server_context.find("common_speculative_flush_prefill(");
+    ok &= expect(flush_call != std::string::npos &&
+                 server_context.find("pf.span.src_offset", flush_call) != std::string::npos &&
+                 server_context.find("pf.span.n_tokens", flush_call) != std::string::npos &&
+                 server_context.find("pf.span.capture_end", flush_call) != std::string::npos,
+        "server must pass planned source offset and absolute capture end to prefill flush");
+    ok &= expect(speculative_h.find("common_dflash_prefill_committed_after_flush") != std::string::npos &&
+                 speculative.find("common_dflash_prefill_committed_after_flush(") != std::string::npos,
+        "DFlash prefill flush must preserve absolute committed positions for long prompts");
 
     // Decode loop must compute intersection offsets
     ok &= expect(context_cpp.find("inter_begin") != std::string::npos,

tests/test-dflash-ring.cpp

@@ -46,5 +46,11 @@ int main() {
         assert(plan.src_token_offset == 0);
     }
 
+    {
+        assert(common_dflash_prefill_committed_after_flush(0, 508, 54366) == 54366);
+        assert(common_dflash_prefill_committed_after_flush(508, 512, 54878) == 54878);
+        assert(common_dflash_prefill_committed_after_flush(1020, 4, 0) == 1024);
+    }
+
     return 0;
 }

tests/test-server-prompt-checkpoint.cpp

@@ -22,5 +22,37 @@ int main() {
     assert(ckpt.ring_data.empty());
     assert(ckpt.ring_data.capacity() == 0);
 
+    server_prompt prompt;
+    for (int i = 0; i < 3; ++i) {
+        auto & cur = prompt.checkpoints.emplace_back();
+        cur.n_tokens = i + 1;
+        cur.data_tgt.resize((size_t) (i + 1) * 100);
+    }
+
+    {
+        server_prompt budgeted = server_prompt_clone_with_checkpoint_budget(prompt, 250, 600);
+        assert(budgeted.checkpoints.size() == 1);
+        assert(budgeted.checkpoints.front().n_tokens == 3);
+        assert(server_prompt_checkpoints_size(budgeted.checkpoints) == 300);
+    }
+
+    {
+        server_prompt budgeted = server_prompt_clone_with_checkpoint_budget(prompt, 600, 600);
+        assert(budgeted.checkpoints.empty());
+    }
+
+    {
+        server_prompt budgeted = server_prompt_clone_with_checkpoint_budget(prompt, 0, 0);
+        assert(budgeted.checkpoints.size() == 3);
+        assert(server_prompt_checkpoints_size(budgeted.checkpoints) == 600);
+    }
+
+    prompt.data.main.resize(64);
+    prompt.data.drft.resize(32);
+    prompt.clear();
+    assert(prompt.n_tokens() == 0);
+    assert(prompt.data.size() == 0);
+    assert(prompt.checkpoints.empty());
+
     return 0;
 }

tools/server/server-context.cpp

@@ -906,7 +906,7 @@ struct server_slot : server_adaptive_dm_state {
             common_context_seq_rm(ctx_dft, id, -1, -1);
         }
 
-        prompt.tokens.clear();
+        prompt.clear();
     }
 
     std::vector<common_adapter_lora_info> lora;
@@ -2224,6 +2224,28 @@ struct server_context_impl {
         GGML_ABORT("failed to expand recurrent state after prompt cache restore; continuing would make scheduler reservation inconsistent\n");
     }
 
+    size_t prompt_cache_limit_bytes() const {
+        return params_base.cache_ram_mib > 0
+            ? (size_t) params_base.cache_ram_mib * 1024ull * 1024ull
+            : 0;
+    }
+
+    size_t context_checkpoint_budget_bytes() const {
+        const size_t cache_limit = prompt_cache_limit_bytes();
+        return cache_limit > 0 ? cache_limit / 2 : 0;
+    }
+
+    void erase_oldest_checkpoint(server_slot & slot, const char * reason) {
+        const auto & cur = slot.prompt.checkpoints.front();
+
+        SLT_WRN(slot,
+                "erasing old context checkpoint (%s, pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n",
+                reason && reason[0] ? reason : "limit", cur.pos_min, cur.pos_max, cur.n_tokens,
+                (float) cur.size() / 1024 / 1024);
+
+        slot.prompt.checkpoints.erase(slot.prompt.checkpoints.begin());
+    }
+
     bool dflash_shared_drafter_batch_allowed(int n_drafting) {
         if (dflash_shared_drafter_batch_disabled()) {
             return false;
@@ -4016,39 +4038,76 @@ struct server_context_impl {
 
     // n_tokens_cur: the number of tokens added to the batch for the current slot
     void create_checkpoint(server_slot & slot, const int64_t n_tokens_cur, llama_pos pos_min, llama_pos pos_max) {
-        while (slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) {
-            // make room for the new checkpoint, if needed
-            const auto & cur = slot.prompt.checkpoints.front();
+        // Save DFlash ring buffer alongside the recurrent state checkpoint.
+        const bool dflash_checkpoint_slot =
+            slot.can_speculate() &&
+            params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH;
+        const size_t cur_size_tgt = ctx_tgt
+            ? llama_state_seq_get_size_ext(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)
+            : 0;
+        const size_t cur_size_dft = ctx_dft
+            ? llama_state_seq_get_size_ext(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)
+            : 0;
+        size_t ring_size = dflash_checkpoint_slot
+            ? common_speculative_ring_state_size(slot.get_spec())
+            : 0;
+        const size_t new_size = cur_size_tgt + cur_size_dft + ring_size;
+        const size_t checkpoint_budget = context_checkpoint_budget_bytes();
 
-            SLT_WRN(slot, "erasing old context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n",
-                    cur.pos_min, cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024);
+        if (checkpoint_budget > 0 && new_size > checkpoint_budget) {
+            SLT_WRN(slot,
+                    "skipping context checkpoint: state size %.3f MiB exceeds checkpoint budget %.3f MiB\n",
+                    new_size / (1024.0 * 1024.0), checkpoint_budget / (1024.0 * 1024.0));
+            return;
+        }
 
-            slot.prompt.checkpoints.erase(slot.prompt.checkpoints.begin());
+        while (slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) {
+            erase_oldest_checkpoint(slot, "count limit");
         }
 
-        auto & cur = slot.prompt.checkpoints.emplace_back();
+        while (checkpoint_budget > 0 &&
+                !slot.prompt.checkpoints.empty() &&
+                server_prompt_checkpoints_size(slot.prompt.checkpoints) + new_size > checkpoint_budget) {
+            erase_oldest_checkpoint(slot, "RAM budget");
+        }
 
+        common_prompt_checkpoint cur;
         cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max);
 
-        cur.update_tgt(ctx_tgt,       slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
-        cur.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+        try {
+            cur.data_tgt.resize(cur_size_tgt);
+            cur.data_dft.resize(cur_size_dft);
+            cur.ring_data.resize(ring_size);
+        } catch (const std::bad_alloc & e) {
+            SLT_WRN(slot,
+                    "skipping context checkpoint: failed to allocate %.3f MiB state (%s)\n",
+                    new_size / (1024.0 * 1024.0), e.what());
+            return;
+        }
+
+        if (cur_size_tgt > 0) {
+            const size_t n = llama_state_seq_get_data_ext(
+                    ctx_tgt, cur.data_tgt.data(), cur_size_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+            if (n != cur_size_tgt) {
+                GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", cur_size_tgt, n);
+            }
+        }
+
+        if (ctx_dft && cur_size_dft > 0) {
+            const size_t n = llama_state_seq_get_data_ext(
+                    ctx_dft.get(), cur.data_dft.data(), cur_size_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+            if (n != cur_size_dft) {
+                GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", cur_size_dft, n);
+            }
+        }
 
-        // Save DFlash ring buffer alongside the recurrent state checkpoint.
-        const bool dflash_checkpoint_slot =
-            slot.can_speculate() &&
-            params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH;
-        size_t ring_size = 0;
         bool ring_saved = false;
-        if (dflash_checkpoint_slot) {
-            ring_size = common_speculative_ring_state_size(slot.get_spec());
-            if (ring_size > 0) {
-                cur.ring_data.resize(ring_size);
-                ring_saved = common_speculative_ring_state_save(
-                        slot.get_spec(), cur.ring_data.data(), ring_size);
-                if (!ring_saved) {
-                    cur.ring_data.clear();
-                    ring_size = 0;
-                }
+        if (dflash_checkpoint_slot && ring_size > 0) {
+            ring_saved = common_speculative_ring_state_save(
+                    slot.get_spec(), cur.ring_data.data(), ring_size);
+            if (!ring_saved) {
+                cur.ring_data.clear();
+                ring_size = 0;
             }
         }
         if (dflash_checkpoint_slot &&
@@ -4059,10 +4118,13 @@ struct server_context_impl {
                     cur.pos_min, cur.pos_max, cur.n_tokens, ring_size, ring_saved ? 1 : 0);
         }
 
+        slot.prompt.checkpoints.push_back(std::move(cur));
+        const auto & saved = slot.prompt.checkpoints.back();
+
         SLT_INF(slot,
                 "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n",
-                (int) slot.prompt.checkpoints.size(), params_base.n_ctx_checkpoints, cur.pos_min,
-                cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024);
+                (int) slot.prompt.checkpoints.size(), params_base.n_ctx_checkpoints, saved.pos_min,
+                saved.pos_max, saved.n_tokens, (float) saved.size() / 1024 / 1024);
     }
 
     void process_single_task(server_task && task) {
@@ -6362,7 +6424,8 @@ struct server_context_impl {
             // DFlash: flush captured hidden states into the ring buffer before
             // the next llama_decode resets the capture buffer.
             for (auto & pf : pending_prefill_flushes) {
-                const int written = common_speculative_flush_prefill(pf.spec, pf.span.src_offset, pf.span.n_tokens);
+                const int written = common_speculative_flush_prefill(
+                        pf.spec, pf.span.src_offset, pf.span.n_tokens, pf.span.capture_end);
                 if (written != pf.span.n_tokens) {
                     SRV_ERR("dflash prefill flush mismatch: slot=%d requested=%d written=%d src_offset=%d; disabling DFlash drafting until fresh hiddens are available\n",
                             pf.slot_id, pf.span.n_tokens, written, pf.span.src_offset);