From 67a80cfb148e655fa22a2a3d25c87f20ced66bfa Mon Sep 17 00:00:00 2001 From: dmcc73 Date: Tue, 21 Apr 2026 11:13:15 +0100 Subject: [PATCH] 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) --- .../mlx/patches/qwen3_5_moe_split/dflash_split.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe_split/dflash_split.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe_split/dflash_split.py index 6b08c4530..8733b06bd 100644 --- a/src/exo/worker/engines/mlx/patches/qwen3_5_moe_split/dflash_split.py +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe_split/dflash_split.py @@ -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)