Skip to content

Commit 8e24fc5

Browse files
howard0suCopilot
andcommitted
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>
1 parent fc360e5 commit 8e24fc5

1 file changed

Lines changed: 8 additions & 0 deletions

File tree

dflash/src/gemma4/gemma4_backend.cpp

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -671,6 +671,14 @@ GenerateResult Gemma4Backend::restore_and_generate(int slot,
671671
cache_.cur_pos = committed;
672672
}
673673

674+
// Full feature mirror resync after restore: do_prefill only synced the
675+
// delta [snap_pos..committed). Re-sync the entire [0..committed) range so
676+
// the draft model sees correct features for the full context.
677+
if (feature_mirror_.target_feat && cache_.target_feat && !draft_parked_ && committed > 0) {
678+
draft_feature_mirror_sync_tail(cache_.target_feat, cache_.target_feat_cap,
679+
feature_mirror_, committed);
680+
}
681+
674682
// Generate
675683
if (req.n_gen > 0) {
676684
const bool can_spec = dflash_target_

0 commit comments

Comments
 (0)