fix(store): numpy view imports dragged full base buffers into Metal
The post-turn store's per-block extraction called mx.array() on NON-CONTIGUOUS numpy views (np_k[:, :, start:end, :] of the full-layer mirror). MLX imports such a view with a buffer sized like the view's BASE array -- measured 30.6MB of Metal for a 0.26MB 64-token block slice. Over an 11k-token qwen hybrid store that is 2200+ block-layer iterations, each retaining a full-layer-sized buffer until the store completes: live Metal climbed linearly to ~158GB (instrumented curve: act 20.1 -> 157.9GB across the block loop), RSS to ~100GB, and the machine died -- as a silent, uncatchable command-buffer abort at the default wired limit, and TWICE as a machine-level watchdog kernel panic with the limit raised. This was the actual cause of the "mystery reboots"; the same idiom also ran (smaller) in the M3 and rotating restore paths. Fix: _mx_from_np_slice() copies the slice host-side with np.ascontiguousarray before mx.array, so the import is slice-sized. Applied at all 14 numpy-slice-to-MLX sites in prefix_cache.py and the 4 restore-rollback sites in scheduler.py. Tests: Metal-measured import-cost regression (fails if a full base buffer is dragged again) + a source pin that no bare mx.array(np_...[slice]) site returns. The instrumented store peaked at 106-109GB RSS on every pre-fix run of the 11k-token diagnostic.
J
Jinho Jang committed
21fd6f90909c0ac4dcfc58558b7317772c6ce62f
Parent: e8b2b60