fix(prefill): keep MLX worker compute thread-affine

Warm the model on its dedicated job thread so MLX 0.31 streams never cross thread boundaries, and emit LaunchAgent files with permissions launchd accepts during upgrades.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

deploy/install_prefill_worker_launchd.sh

@@ -87,6 +87,7 @@ cat > "$PLIST" <<EOF
 </dict></plist>
 EOF
 
+chmod 644 "$PLIST"
 launchctl bootout "gui/$(id -u)/$LABEL" 2>/dev/null || true
 launchctl bootstrap "gui/$(id -u)" "$PLIST"
 echo "installed $LABEL -> $PLIST"

inference_engine/distributed/prefill_worker.py

@@ -9,7 +9,7 @@
 from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import dataclass, field
 from enum import IntEnum
-from typing import Protocol, Sequence
+from typing import Callable, Protocol, Sequence
 
 import grpc
 
@@ -85,9 +85,10 @@ class PrefillJobStore:
 
     def __init__(
         self,
-        engine: PrefillComputeEngine,
+        engine: PrefillComputeEngine | None,
         cache_store: PrefixCacheStore,
         *,
+        engine_factory: Callable[[], PrefillComputeEngine] | None = None,
         max_concurrent_jobs: int = 1,
         max_jobs: int = 128,
         completed_ttl_s: float = 600.0,
@@ -100,7 +101,10 @@ def __init__(
             max_prompt_tokens,
         ) <= 0:
             raise ValueError("worker limits must be > 0")
+        if (engine is None) == (engine_factory is None):
+            raise ValueError("provide exactly one of engine or engine_factory")
         self.engine = engine
+        self.engine_factory = engine_factory
         self.cache_store = cache_store
         self.max_concurrent_jobs = int(max_concurrent_jobs)
         self.max_jobs = int(max_jobs)
@@ -109,11 +113,16 @@ def __init__(
         self._jobs: dict[str, PrefillJob] = {}
         self._requests: dict[tuple[str, str], str] = {}
         self._lock = threading.RLock()
+        self._thread_local = threading.local()
         self._executor = ThreadPoolExecutor(
             max_workers=self.max_concurrent_jobs,
             thread_name_prefix="kakeya-prefill-worker",
         )
 
+    def warmup(self) -> None:
+        """Construct a factory-backed engine on its eventual compute thread."""
+        self._executor.submit(self._engine_for_current_thread).result()
+
     def submit(
         self,
         *,
@@ -243,7 +252,7 @@ def _run(self, job_id: str) -> None:
                 timer.daemon = True
                 timer.start()
         try:
-            blocks = tuple(self.engine.compute_prefill(
+            blocks = tuple(self._engine_for_current_thread().compute_prefill(
                 job.token_ids,
                 job.block_hashes,
                 compression=job.compression,
@@ -289,6 +298,16 @@ def _run(self, job_id: str) -> None:
                 job.compute_ms = (time.perf_counter() - started) * 1000.0
                 job.finished_at = time.time()
 
+    def _engine_for_current_thread(self) -> PrefillComputeEngine:
+        if self.engine is not None:
+            return self.engine
+        engine = getattr(self._thread_local, "engine", None)
+        if engine is None:
+            assert self.engine_factory is not None
+            engine = self.engine_factory()
+            self._thread_local.engine = engine
+        return engine
+
     def _gc_locked(self) -> None:
         cutoff = time.time() - self.completed_ttl_s
         for job_id, job in list(self._jobs.items()):

scripts/start_prefill_worker_node.py

@@ -83,27 +83,37 @@ async def serve(args) -> None:
         )
         if args.fleet_psk_file else None
     )
-    verifier = MLXSinkWindowVerifier(VerifierConfig(
-        model_id=args.model_id,
-        sink_size=args.sink,
-        window_size=args.window,
-        dtype=torch.bfloat16,
-        device="cpu",
-    ))
+    if args.max_concurrent_jobs != 1:
+        raise SystemExit(
+            "MLX prefill workers require --max-concurrent-jobs 1 so the "
+            "model and its stream remain on one compute thread",
+        )
     store = PrefixCacheStore(
         compatibility,
         max_bytes=int(args.cache_gb * (1 << 30)),
         node_id=args.node_id,
     )
-    engine = MLXPrefillComputeEngine(verifier, compatibility)
+
+    def engine_factory() -> MLXPrefillComputeEngine:
+        verifier = MLXSinkWindowVerifier(VerifierConfig(
+            model_id=args.model_id,
+            sink_size=args.sink,
+            window_size=args.window,
+            dtype=torch.bfloat16,
+            device="cpu",
+        ))
+        return MLXPrefillComputeEngine(verifier, compatibility)
+
     jobs = PrefillJobStore(
-        engine,
+        None,
         store,
+        engine_factory=engine_factory,
         max_concurrent_jobs=args.max_concurrent_jobs,
         max_jobs=args.max_jobs,
         completed_ttl_s=args.job_ttl_s,
         max_prompt_tokens=args.max_prompt_tokens,
     )
+    jobs.warmup()
 
     def card() -> NodeCapability:
         inflight, queued, load, queued_tokens = jobs.stats()

tests/inference_engine/bridge/test_prefill_worker_launchd.py

@@ -22,6 +22,7 @@ def test_worker_installer_emits_full_cache_compatibility_contract():
         assert f"<string>{flag}</string>" in source
     assert 'PEER="${KAKEYA_WORKER_PEER:-}"' in source
     assert "<string>--peer</string>" in source
+    assert 'chmod 644 "$PLIST"' in source
 
 
 def test_head_runtime_discovers_and_uses_worker_cache_port():

tests/inference_engine/distributed/test_prefill_worker.py

@@ -162,6 +162,10 @@ def test_job_store_validation_queue_stats_and_gc():
     cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="w")
     with pytest.raises(ValueError):
         PrefillJobStore(_Engine(), cache, max_jobs=0)
+    with pytest.raises(ValueError, match="exactly one"):
+        PrefillJobStore(None, cache)
+    with pytest.raises(ValueError, match="exactly one"):
+        PrefillJobStore(_Engine(), cache, engine_factory=_Engine)
     blocking = _Engine()
     blocking.block.set()
     jobs = PrefillJobStore(blocking, cache, max_jobs=1, max_prompt_tokens=4)
@@ -230,6 +234,42 @@ def test_job_store_validation_queue_stats_and_gc():
         completed_jobs.close()
 
 
+def test_factory_engine_is_warmed_and_used_on_same_compute_thread():
+    cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="w")
+    created_on = []
+    computed_on = []
+
+    class ThreadBoundEngine(_Engine):
+        def compute_prefill(self, *args, **kwargs):
+            computed_on.append(threading.get_ident())
+            return super().compute_prefill(*args, **kwargs)
+
+    def factory():
+        created_on.append(threading.get_ident())
+        return ThreadBoundEngine()
+
+    jobs = PrefillJobStore(None, cache, engine_factory=factory)
+    try:
+        jobs.warmup()
+        job = jobs.submit(
+            request_id="thread-affinity",
+            tenant_id="tenant",
+            token_ids=[1, 2],
+            block_hashes=[b"a" * 32],
+            compatibility=COMPAT,
+            compression=CompressionCodec.NONE,
+        )
+        for _ in range(100):
+            if job.state == PrefillJobState.COMPLETED:
+                break
+            time.sleep(0.005)
+        assert job.state == PrefillJobState.COMPLETED
+        assert created_on == computed_on
+        assert created_on[0] != threading.get_ident()
+    finally:
+        jobs.close()
+
+
 def test_job_store_failure_modes_and_precancelled_run():
     cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="w")