gemma4: full feature mirror resync after prefix cache restore

After restoring KV from a snapshot, do_prefill only syncs the feature mirror
for the delta tokens [snap_pos..committed). The positions [0..snap_pos) in
the mirror retain stale data from the previous request's decode phase (which
may have diverged from the current prompt context after the ring buffer
wraps).

Fix: call draft_feature_mirror_sync_tail after restore to resync the entire
[0..committed) feature range from cache_.target_feat to the mirror. This
ensures the draft model sees consistent features and maintains high
acceptance rate (AL) during speculative decoding.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

gemma4: save/restore target_feat in prefix cache snapshot

Matching Qwen35's approach: save target_feat (BF16 feature ring buffer) and
last_tok as part of the KV snapshot. On restore, target_feat is copied back
to GPU before the delta prefill + feature mirror resync.

Previously, only K/V tensors were snapshotted. After restore, the feature
mirror contained stale data from the previous request's decode phase, causing
the draft model to make poor predictions and halving speculative decode
acceptance rate (52% → 24%).

With this fix, the full feature state is correctly restored, and the
subsequent draft_feature_mirror_sync_tail ensures the mirror matches.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: howard0su

dflash/src/gemma4/gemma4_backend.cpp

@@ -511,6 +511,16 @@ GenerateResult Gemma4Backend::generate(const GenerateRequest & req,
         return result;
     }
 
+    // Inline snapshot at snap_pos for prefix cache
+    if (req.snap_slot >= 0 && req.snap_pos > 0 && req.snap_pos <= committed) {
+        cache_.cur_pos = req.snap_pos;
+        if (snapshot_save(req.snap_slot)) {
+            std::fprintf(stderr, "[gemma4] inline-snap slot=%d cur_pos=%d\n",
+                         req.snap_slot, req.snap_pos);
+        }
+        cache_.cur_pos = committed;
+    }
+
     if (req.n_gen > 0) {
         // Try speculative decode if draft is available and temp==0
         const bool can_spec = dflash_target_
@@ -623,8 +633,15 @@ GenerateResult Gemma4Backend::restore_and_generate(int slot,
         }
     }
 
+    // Restore target_feat from snapshot
+    if (snap.feat_snap && cache_.target_feat) {
+        const size_t feat_nbytes = ggml_nbytes(snap.feat_snap);
+        ggml_backend_tensor_set(cache_.target_feat, snap.feat_snap->data, 0, feat_nbytes);
+    }
+
     const int snap_pos = snap.cur_pos;
     cache_.cur_pos = snap_pos;
+    cache_.last_tok = snap.last_tok;
 
     // Set up sampler
     sampler_ = req.sampler;
@@ -651,6 +668,24 @@ GenerateResult Gemma4Backend::restore_and_generate(int slot,
     }
     // else: prompt_len == snap_pos → no delta, committed stays at snap_pos
 
+    // Inline snapshot at snap_pos for prefix cache (new snap from this request)
+    if (req.snap_slot >= 0 && req.snap_pos > 0 && req.snap_pos <= committed) {
+        cache_.cur_pos = req.snap_pos;
+        if (snapshot_save(req.snap_slot)) {
+            std::fprintf(stderr, "[gemma4] inline-snap slot=%d cur_pos=%d\n",
+                         req.snap_slot, req.snap_pos);
+        }
+        cache_.cur_pos = committed;
+    }
+
+    // Full feature mirror resync after restore: do_prefill only synced the
+    // delta [snap_pos..committed). Re-sync the entire [0..committed) range so
+    // the draft model sees correct features for the full context.
+    if (feature_mirror_.target_feat && cache_.target_feat && !draft_parked_ && committed > 0) {
+        draft_feature_mirror_sync_tail(cache_.target_feat, cache_.target_feat_cap,
+                                       feature_mirror_, committed);
+    }
+
     // Generate
     if (req.n_gen > 0) {
         const bool can_spec = dflash_target_
@@ -738,8 +773,9 @@ bool Gemma4Backend::snapshot_save(int slot) {
     if (needs_alloc) {
         free_gemma4_snapshot(snap);
 
+        const int n_feat_tensors = (cache_.target_feat && cache_.target_feat_cap > 0) ? 1 : 0;
         ggml_init_params ip{};
-        ip.mem_size = ggml_tensor_overhead() * (size_t)(n_layer * 2 + 4) + 4096;
+        ip.mem_size = ggml_tensor_overhead() * (size_t)(n_layer * 2 + n_feat_tensors + 4) + 4096;
         ip.no_alloc = true;
         snap.ctx = ggml_init(ip);
         if (!snap.ctx) return false;
@@ -759,10 +795,21 @@ bool Gemma4Backend::snapshot_save(int slot) {
             }
         }
 
+        // target_feat: save min(snap_pos, target_feat_cap) positions
+        snap.feat_snap = nullptr;
+        snap.feat_cap  = 0;
+        if (cache_.target_feat && cache_.target_feat_cap > 0) {
+            const int feat_len = std::min(snap_pos, cache_.target_feat_cap);
+            snap.feat_snap = ggml_new_tensor_2d(snap.ctx, cache_.target_feat->type,
+                                                 cache_.target_feat->ne[0], feat_len);
+            snap.feat_cap = cache_.target_feat_cap;
+        }
+
         snap.buf = ggml_backend_alloc_ctx_tensors(snap.ctx, snap_backend_);
         if (!snap.buf) {
             ggml_free(snap.ctx); snap.ctx = nullptr;
             snap.k_snap.clear(); snap.v_snap.clear();
+            snap.feat_snap = nullptr;
             return false;
         }
     }
@@ -792,9 +839,15 @@ bool Gemma4Backend::snapshot_save(int slot) {
         }
     }
     snap.cur_pos = snap_pos;
+    snap.last_tok = cache_.last_tok;
 
-    std::printf("[gemma4] snapshot saved slot=%d pos=%d\n", slot, snap.cur_pos);
-    std::fflush(stdout);
+    // target_feat: copy min(snap_pos, cap) positions from GPU to snapshot
+    if (snap.feat_snap && cache_.target_feat) {
+        const size_t feat_nbytes = ggml_nbytes(snap.feat_snap);
+        ggml_backend_tensor_get(cache_.target_feat, snap.feat_snap->data, 0, feat_nbytes);
+    }
+
+    std::fprintf(stderr, "[gemma4] snapshot saved slot=%d pos=%d\n", slot, snap.cur_pos);
     return true;
 }
 

dflash/src/gemma4/gemma4_internal.h

@@ -193,8 +193,11 @@ bool  create_gemma4_target_feat(ggml_backend_t backend, Gemma4Cache & cache,
 // Snapshot
 struct Gemma4Snapshot {
     int cur_pos = 0;
+    int32_t last_tok = -1;
     std::vector<ggml_tensor *> k_snap;
     std::vector<ggml_tensor *> v_snap;
+    ggml_tensor *             feat_snap = nullptr;  // [fc_in, feat_len]
+    int                       feat_cap  = 0;
     ggml_context *        ctx = nullptr;
     ggml_backend_buffer_t buf = nullptr;
 };

dflash/src/gemma4/gemma4_loader.cpp

@@ -478,7 +478,10 @@ void free_gemma4_snapshot(Gemma4Snapshot & s) {
     if (s.buf) { ggml_backend_buffer_free(s.buf); s.buf = nullptr; }
     if (s.ctx) { ggml_free(s.ctx); s.ctx = nullptr; }
     s.k_snap.clear(); s.v_snap.clear();
-    s.cur_pos = 0;
+    s.feat_snap = nullptr;
+    s.feat_cap  = 0;
+    s.cur_pos   = 0;
+    s.last_tok  = -1;
 }
 
 }  // namespace dflash27b