cleanup

src/exo/worker/engines/mlx/cache.py

@@ -338,6 +338,14 @@ def _measure_single_cache_bytes(
             for c in entry.caches
         )
 
+    total = 0
+    if isinstance(entry, ArraysCache):
+        state = entry.state  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        for arr in state:  # pyright: ignore[reportUnknownVariableType]
+            if isinstance(arr, mx.array):
+                total += arr.nbytes
+        return total
+
     total = 0
     for attr_name in ("keys", "values"):
         val: object = getattr(entry, attr_name, None)
@@ -346,16 +354,10 @@ def _measure_single_cache_bytes(
         if isinstance(val, mx.array):
             total += val.nbytes
         elif isinstance(val, (tuple, list)):
-            # QuantizedKVCache stores tuples of arrays (data, scales, biases)
             for arr in val:  # pyright: ignore[reportUnknownVariableType]
                 if isinstance(arr, mx.array):
                     total += arr.nbytes
 
-    if isinstance(entry, ArraysCache):
-        state = entry.state  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-        for arr in state:  # pyright: ignore[reportUnknownVariableType]
-            if isinstance(arr, mx.array):
-                total += arr.nbytes
     return total
 
 

src/exo/worker/engines/mlx/generator/generate.py

@@ -153,11 +153,10 @@ def warmup_inference(
     tokenizer: TokenizerWrapper,
     group: mx.distributed.Group | None,
 ) -> tuple[int, int]:
-    """Run warmup inference and measure KV cache cost per token.
+    """Run warmup inference and measure various metrics.
 
     Returns:
-        (tokens_generated, bytes_per_token) where bytes_per_token is the
-        measured KV cache memory consumed per token across all local layers.
+        tokens_generated, bytes_per_token
     """
     content = "Prompt to warm up the inference engine. Repeat this."
 
@@ -373,7 +372,6 @@ def mlx_generate(
                 f"KV cache hit: {prefix_hit_length}/{len(all_prompt_tokens)} tokens cached ({100 * prefix_hit_length / len(all_prompt_tokens):.1f}%)"
             )
 
-    # OOM prevention: check if the prompt will fit in memory
     if bytes_per_token > 0:
         oom_error = _check_memory_budget(
             bytes_per_token=bytes_per_token,