qwen3_5_moe_split: use BatchKVCache._idx for mask slicing offset

BatchKVCache.update_and_fetch returns keys[:, :, :_idx, :] — the actual
K buffer length. BatchKVCache.make_mask also uses _idx as the offset to
create_causal_mask. Our pipelined mask slicer was using .offset (which
is _idx - left_padding for positional encodings), producing a mask too
short by left_padding.

Fix: prefer _idx if the cache has it (BatchKVCache), else fall back to
.offset (plain KVCache, where they're equal).

Shape error was "(1,1,4,65) vs (1,24,4,72)" — 72 = actual _idx + mid,
65 = our offset + mid with offset=61. Difference 7 = left_padding of
one of the batch entries.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/exo/worker/engines/mlx/patches/qwen3_5_moe_split/model_forward.py

@@ -174,13 +174,20 @@ def pipelined_layer_loop(
     capture_set: set[int] = set(capture_layers) if capture_layers is not None else set()
 
     # Compute cache offset for mask slicing (same for all GQA layers).
+    # For BatchKVCache, update_and_fetch returns keys[:_idx], so the mask must
+    # match _idx (actual K buffer length), NOT `offset` (which is _idx minus
+    # left_padding for positional encodings). Plain KVCache uses `offset`.
     fa_cache = cache[inner.fa_idx] if cache[inner.fa_idx] is not None else None
-    raw_offset = fa_cache.offset if (fa_cache is not None and hasattr(fa_cache, "offset")) else 0
-    # BatchKVCache stores offset as an mx.array (1,), KVCache stores it as a Python int.
-    if isinstance(raw_offset, mx.array):
-        offset = int(raw_offset.max().item())  # max so mask covers all batch entries
+    if fa_cache is None:
+        offset = 0
+    elif hasattr(fa_cache, "_idx"):
+        # BatchKVCache: _idx is Python int tracking actual K length
+        offset = int(fa_cache._idx)
+    elif hasattr(fa_cache, "offset"):
+        raw = fa_cache.offset
+        offset = int(raw.max().item()) if isinstance(raw, mx.array) else int(raw)
     else:
-        offset = int(raw_offset)
+        offset = 0
 
     # Slice masks once (reused every stage)
     fa_mask_H0, fa_mask_H1 = slice_fa_mask(fa_mask, mid, offset)