fix: per-layer KV cache buffers — full-tensor copy per update made Fish S2 Pro decode unusably slow

[ComicBit]

Problem

KVCache (used by Fish S2 Pro decode) stores all layers' K/V in a single 5-D tensor (num_layers, batch, heads, max_length, head_dim). MLX's __setitem__ is a functional scatter, so every per-layer slice write copies the entire tensor. With 36 slow-AR layers × (K+V), that's 72 full-tensor copies per generated token — multi-GB of memcopy per token, with transient allocations pushing past physical RAM into swap on 16 GB machines.

Observed on an M2 Pro 16 GB with appautomaton/fishaudio-s2-pro-8bit-mlx: per-token decode time grew monotonically (38s → 69s → 87s+), and a single short sentence never finished generating (killed after 50+ min).

Fix

Per-layer buffers grown in fixed chunks (_ALLOC_STEP = 256) — the same pattern mlx-lm uses. A slice update now touches only that layer's ~2 MB buffer and MLX can donate it in place.

Public API unchanged: update, get(layer_idx), offset, _offsets (read directly by layers.py for the RoPE offset), reset, trim_to all behave identically. get(None) stacks per-layer slices to preserve the previous 5-D return shape.

Verification

• All 48 fish_s2 unit tests pass (test_fish_s2_model, test_fish_s2_generation, test_fish_s2_config, test_fish_s2_checkpoint_remap).
• Instrumented decode benchmark on the same machine: per-step cost is now flat (2–3 s/step under heavy ambient load; the machine was simultaneously running a rustc build), MLX active memory constant at 6.73 GB / peak 6.79 GB for the int8 model. Before the fix, step cost grew without bound and memory transients forced swap.
• The 2–3 s/step number was taken on a heavily loaded 16 GB machine; on an idle machine the expected steady state is in line with the ~21 tok/s the docs advertise.