qwen3_5_moe_split: rollback BatchKVCache._idx too, not just .offset

BatchKVCache.update_and_fetch writes at [_idx : _idx + new_S] and
returns keys[:_idx]. Stock DFlash's rollback only decrements .offset,
leaving _idx at the post-verify value. Result: stale rejected-draft
keys stay in the cache, and the next verify's attention attends to
them, corrupting the output.

Fix: decrement _idx by the same amount as offset. Now
update_and_fetch overwrites stale entries cleanly.

This matches observed behavior: first speculative verify produces
fine output, second+ verifies produce gibberish due to stale cache
contamination.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
dmcc73andClaude Opus 4.6 committed 2026-04-21 11:13:15 +01:00
1 parent 945ba5bb2a
commit 67a80cfb14
1 file changed
+8 -1
@@ -152,12 +152,19 @@ def make_split_speculative_next(group): # type: ignore[no-untyped-def]
)
n_accepted = int(gathered_n[MOE_RANK].item())
# 7. Rollback — same n_accepted on both ranks keeps caches consistent
# 7. Rollback — same n_accepted on both ranks keeps caches consistent.
# For BatchKVCache, both `offset` AND `_idx` must be decremented.
# `_idx` is the buffer-length tracker that `update_and_fetch` uses to
# decide where to write next and `fetch` uses to slice returned keys.
# Stock DFlash only decrements offset, which leaves stale rejected-draft
# entries in the cache that corrupt subsequent verifies.
rollback = verify_len - n_accepted
if rollback > 0:
for c in batch.cache:
if hasattr(c, "offset"):
c.offset -= rollback
if hasattr(c, "_idx"):
c._idx -= rollback
elif hasattr(c, "rollback"):
c.rollback(n_accepted)