From dc1ce2a2cf2c81417cdd5084bd7f5caee6c3071a Mon Sep 17 00:00:00 2001
From: Ryuichi Leo Takashige <leo@exolabs.net>
Date: Wed, 25 Feb 2026 20:54:02 +0000
Subject: [PATCH] cleanup

---
 src/exo/worker/engines/mlx/cache.py              | 14 ++++++++------
 src/exo/worker/engines/mlx/generator/generate.py |  6 ++----
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py
index b1383350..cc43d2e0 100644
--- a/src/exo/worker/engines/mlx/cache.py
+++ b/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
 
 
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 94912f38..afb1ad39 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/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,