mirror of
https://github.com/exo-explore/exo.git
synced 2026-09-12 21:38:59 -04:00
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:
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)
|
||||
|
||||
|
||||
Reference in new issue
Block a user