Laya-MLX Metal Memory Growth Fix

#5
by finche - opened

Issue

After running for several hours, the Laya-MLX sidecar consumed approximately 21.7 GB of memory on macOS. The system also reported approximately 39 GB of compressed memory. Restarting the Laya process immediately released the memory.

Root Cause

vmmap confirmed that the memory was allocated by the Metal/MLX graphics subsystem rather than by the Python heap:

Physical footprint:         21.7 GB
IOAccelerator (graphics):   21.4 GB
Normal malloc regions:      approximately 65 MB

The expanded Laya model files occupy approximately 647 MB, so the issue was not oversized model weights and was not a conventional Python object leak.

The growth was caused by MLX/Metal inference buffers and allocator cache accumulating over time, amplified by:

  1. A default batch_size of 16 although normal decisions use one sample;
  2. Variable input lengths and question counts producing different Metal buffer shapes;
  3. No explicit cleanup of completed MLX allocator buffers after inference.

Fix

Updated scripts/laya_service.py to:

  • Set the default batch_size to 1;
  • Add the LAYA_BATCH_SIZE environment variable for controlled tuning;
  • Limit the MLX allocator cache with LAYA_CACHE_LIMIT_MB, defaulting to 1024 MB;
  • Call mx.clear_cache() after each completed inference;
  • Expose request count, last inference latency, batch size, and cache limit through /health.

Verification

Before the fix:

Physical footprint:         21.7 GB
IOAccelerator (graphics):   21.4 GB

After restarting Laya with the fix and processing 48 requests:

Physical footprint:         834.6 MB
Peak footprint:             1.3 GB
IOAccelerator (graphics):   620.2 MB
Inference latency:           34.5 ms

The existing :9005/predict API contract remains unchanged. Jev, JevCache, and the deterministic fallback path are unaffected.

Rollback

Restore the previous Laya startup and inference-cache behavior in scripts/laya_service.py, then restart the Laya sidecar.

Runtime tuning is available through:

LAYA_BATCH_SIZE=1 LAYA_CACHE_LIMIT_MB=1024

Conclusion

This fix changes Laya from a process that requires manual restarts to release accumulated Metal memory into a managed sidecar with bounded batching, a bounded allocator cache, post-inference cleanup, and runtime observability.

Sign up or log in to comment