SIGN IN SIGN UP

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