[object Object]

← back to Exo

fix: force gc + clear_cache after KV prefix cache eviction (#1832)

af9e847edbb1939872f79889cfe1617d9fe1362e · 2026-04-17 12:57:32 -0500 · Adam Durham

## Summary
- After `KVPrefixCache` evicts LRU entries, the MLX Metal buffers stay
allocated until Python's GC runs
- This leaks ~3-4 GB between long-context requests, reducing the
effective context ceiling for back-to-back requests
- Adding `gc.collect()` + `mx.clear_cache()` after eviction frees Metal
buffers promptly

## Test plan
- [x] Measured on 2-node PP cluster with Qwen3.5-397B-A17B-4bit at 63K
context
- [x] Before: 108.88 GB retained after eviction (3.78 GB above baseline)
- [x] After: 105.48 GB retained after eviction (0.38 GB above baseline —
draft model KV + minor overhead)
- [x] `gc.collect()` adds ~2-3ms latency, runs once per eviction cycle
(not per token)
- [ ] Verify with `uv run pytest`

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Adam Durham <adam@example.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>

Files touched

Diff

commit af9e847edbb1939872f79889cfe1617d9fe1362e
Author: Adam Durham <amdnative@gmail.com>
Date:   Fri Apr 17 12:57:32 2026 -0500

    fix: force gc + clear_cache after KV prefix cache eviction (#1832)
    
    ## Summary
    - After `KVPrefixCache` evicts LRU entries, the MLX Metal buffers stay
    allocated until Python's GC runs
    - This leaks ~3-4 GB between long-context requests, reducing the
    effective context ceiling for back-to-back requests
    - Adding `gc.collect()` + `mx.clear_cache()` after eviction frees Metal
    buffers promptly
    
    ## Test plan
    - [x] Measured on 2-node PP cluster with Qwen3.5-397B-A17B-4bit at 63K
    context
    - [x] Before: 108.88 GB retained after eviction (3.78 GB above baseline)
    - [x] After: 105.48 GB retained after eviction (0.38 GB above baseline —
    draft model KV + minor overhead)
    - [x] `gc.collect()` adds ~2-3ms latency, runs once per eviction cycle
    (not per token)
    - [ ] Verify with `uv run pytest`
    
    🤖 Generated with [Claude Code](https://claude.com/claude-code)
    
    ---------
    
    Co-authored-by: Adam Durham <adam@example.com>
    Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
    Co-authored-by: rltakashige <rl.takashige@gmail.com>
---
 src/exo/worker/engines/mlx/cache.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py
index 6d3fd3df..f2f197cd 100644
--- a/src/exo/worker/engines/mlx/cache.py
+++ b/src/exo/worker/engines/mlx/cache.py
@@ -1,3 +1,4 @@
+import gc
 import os
 from copy import deepcopy
 from typing import TYPE_CHECKING
@@ -307,6 +308,7 @@ class KVPrefixCache:
         if len(self.caches) == 0:
             return
 
+        evicted_any = False
         # Evict LRU entries until below threshold
         while (
             len(self.caches) > 0
@@ -320,10 +322,16 @@ class KVPrefixCache:
             self._media_regions.pop(lru_index)
             self._last_used.pop(lru_index)
             self.prefill_tps.pop(lru_index)
+
+            evicted_any = True
             logger.info(
                 f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory usage"
             )
 
+        if evicted_any:
+            gc.collect()
+            mx.clear_cache()
+
     def get_memory_used_percentage(self) -> float:
         local_pressure: float = get_memory_used_percentage()

← 01598960 Add model card for Qwen3.6-35B-A3B-8bit (#1917)  ·  back to Exo  ·  Improve build CI (#1920) bf8aacfd →