Description
Running gemma-4-31b-it-bf16 in tensor-parallel across 2x Mac Studio M4 Ultra (64GB each), the runner crashes with signal 6 (SIGABRT) during prefill. The Metal error is:
libc++abi: terminating due to uncaught exception of type std::runtime_error: [METAL] Command buffer execution failed: Insufficient Memory (00000008:kIOGPUCommandBufferCallbackErrorOutOfMemory)
Root Cause
KVPrefixCache._evict_if_needed() only runs when a new cache entry is ADDED (after prefill completes), but the OOM happens DURING the prefill forward pass — before eviction gets a chance to free memory. Under sustained load with many requests, the KV prefix cache accumulates entries until Metal cannot allocate memory for the next prefill step's intermediate activations.
Reproduction
- Model: mlx-community/gemma-4-31b-it-bf16 (62GB model, ~31GB per node)
- Setup: 2-node tensor-parallel over Thunderbolt
- Workload: sustained stream of requests with varying prompt lengths (500-2000 tokens)
- Crashes occur after 2-4 hours of sustained load, always during prefill
- All three observed crashes had identical Metal error messages
- Stack traces show crashes in
generate_step → stream_generate → prefill → submit
Environment
- macOS, 2x Mac Studio M4 Ultra 64GB
- MLX with Metal backend
- EXO_MEMORY_THRESHOLD default (0.80 for 64GB)
Proposed Fix
Call evict_if_needed() pre-emptively before starting the prefill forward pass in batch_generate.py's submit() method, after the KV cache lookup. This frees stale cache entries and reclaims GPU memory before the large allocation.
Description
Running gemma-4-31b-it-bf16 in tensor-parallel across 2x Mac Studio M4 Ultra (64GB each), the runner crashes with signal 6 (SIGABRT) during prefill. The Metal error is:
Root Cause
KVPrefixCache._evict_if_needed()only runs when a new cache entry is ADDED (after prefill completes), but the OOM happens DURING the prefill forward pass — before eviction gets a chance to free memory. Under sustained load with many requests, the KV prefix cache accumulates entries until Metal cannot allocate memory for the next prefill step's intermediate activations.Reproduction
generate_step→stream_generate→prefill→submitEnvironment
Proposed Fix
Call
evict_if_needed()pre-emptively before starting the prefill forward pass inbatch_generate.py'ssubmit()method, after the KV cache lookup. This frees stale cache entries and reclaims GPU memory before the large allocation.