Fix pipeline prefill short-prompt wedge + dashboard kv-backend display

src/exo/api/main.py

@@ -2592,7 +2592,8 @@ async def get_config(self) -> JSONResponse:
                 "fileExists": True,
                 "effective": {
                     "kv_cache_backend": os.environ.get(
-                        "EXO_KV_CACHE_BACKEND", "default"
+                        "SKULK_KV_CACHE_BACKEND",
+                        os.environ.get("EXO_KV_CACHE_BACKEND", "default"),
                     ),
                     "has_hf_token": has_hf_token or "HF_TOKEN" in os.environ,
                 },

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

@@ -1,9 +1,14 @@
+import contextlib
 import os
+import sys
 import threading
+import time
+import traceback
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Generator
 from functools import partial
 from inspect import signature
+from types import FrameType
 from typing import TYPE_CHECKING, Any, Literal, Protocol, cast
 
 import mlx.core as mx
@@ -70,10 +75,83 @@
 _pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = []
 
 
+def _mlx_hang_debug_enabled() -> bool:
+    value = os.environ.get("SKULK_MLX_HANG_DEBUG") or os.environ.get(
+        "EXO_MLX_HANG_DEBUG"
+    )
+    if value is None:
+        return False
+    return value.lower() not in {"", "0", "false", "no", "off"}
+
+
+def _mlx_hang_debug_interval_seconds() -> float:
+    raw = os.environ.get("SKULK_MLX_HANG_DEBUG_INTERVAL_SECONDS") or os.environ.get(
+        "EXO_MLX_HANG_DEBUG_INTERVAL_SECONDS"
+    )
+    if raw is None:
+        return 30.0
+    with contextlib.suppress(ValueError):
+        return max(float(raw), 1.0)
+    return 30.0
+
+
+@contextlib.contextmanager
+def _hang_debug_watch(label: str) -> Generator[None]:
+    """Emit periodic stack-rich logs while a pipeline wrapper stage is stuck."""
+    if not _mlx_hang_debug_enabled():
+        yield
+        return
+
+    interval_seconds = _mlx_hang_debug_interval_seconds()
+    started_at = time.monotonic()
+    monitored_thread_id = threading.get_ident()
+    finished = threading.Event()
+
+    logger.info(
+        f"[hang-debug] Entering {label} (watchdog interval={interval_seconds:.0f}s)"
+    )
+
+    def watchdog() -> None:
+        while not finished.wait(timeout=interval_seconds):
+            elapsed = time.monotonic() - started_at
+            current_frames = cast(
+                Callable[[], dict[int, FrameType]] | None,
+                getattr(sys, "_current_frames", None),
+            )
+            frame = (
+                None
+                if current_frames is None
+                else current_frames().get(monitored_thread_id)
+            )
+            if frame is None:
+                stack_text = "<no Python frame available>"
+            else:
+                stack_text = "".join(traceback.format_stack(frame))
+            logger.warning(
+                f"[hang-debug] Still in {label} after {elapsed:.1f}s\n{stack_text}"
+            )
+
+    watchdog_thread = threading.Thread(
+        target=watchdog,
+        name=f"hang-debug:{label}",
+        daemon=True,
+    )
+    watchdog_thread.start()
+
+    try:
+        yield
+    finally:
+        finished.set()
+        elapsed = time.monotonic() - started_at
+        logger.info(f"[hang-debug] Leaving {label} after {elapsed:.1f}s")
+
+
 def flush_prefill_sends() -> None:
     for output, dst, group in _pending_prefill_sends:
-        sent = mx.distributed.send(output, dst, group=group)
-        mx.async_eval(sent)
+        with _hang_debug_watch(f"flush_prefill_sends send dst={dst}"):
+            sent = mx.distributed.send(output, dst, group=group)
+        with _hang_debug_watch(f"flush_prefill_sends async_eval dst={dst}"):
+            mx.async_eval(sent)
     _pending_prefill_sends.clear()
 
 
@@ -163,10 +241,20 @@ def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
         if self.r != 0:
             # We want to avoid GPU timeout errors by evalling the distributed operation
             # so that it stays on CPU, which does not have a timeout.
-            mx.eval(x)
-            x = mx.distributed.recv_like(x, (self.r - 1), group=self.group)
-            mx.eval(x)
-        return self.original_layer(x, *args, **kwargs)
+            with _hang_debug_watch(
+                f"pipeline_first eval_input rank={self.r} src={self.r - 1}"
+            ):
+                mx.eval(x)
+            with _hang_debug_watch(
+                f"pipeline_first recv_like rank={self.r} src={self.r - 1}"
+            ):
+                x = mx.distributed.recv_like(x, (self.r - 1), group=self.group)
+            with _hang_debug_watch(
+                f"pipeline_first eval_recv rank={self.r} src={self.r - 1}"
+            ):
+                mx.eval(x)
+        with _hang_debug_watch(f"pipeline_first original_layer rank={self.r}"):
+            return self.original_layer(x, *args, **kwargs)
 
 
 class PipelineLastLayer(CustomMlxLayer):
@@ -190,36 +278,57 @@ def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
             x, *args, **kwargs
         ).arguments.get("cache", None)
 
-        output: mx.array = self.original_layer(x, *args, **kwargs)
+        with _hang_debug_watch(
+            f"pipeline_last original_layer rank={self.r} world={self.s} prefill={self.is_prefill}"
+        ):
+            output: mx.array = self.original_layer(x, *args, **kwargs)
 
         # Eval layer output to materialize it before send — this splits the graph
         # so the send is isolated and the receiving rank's recv can complete.
-        mx.eval(output)
+        with _hang_debug_watch(
+            f"pipeline_last eval_output rank={self.r} world={self.s} prefill={self.is_prefill}"
+        ):
+            mx.eval(output)
 
         if self.r != self.s - 1:
             if self.queue_sends:
                 _pending_prefill_sends.append(
                     (output, (self.r + 1) % self.s, self.group)
                 )
             else:
-                output = mx.distributed.send(
-                    output, (self.r + 1) % self.s, group=self.group
-                )
+                with _hang_debug_watch(
+                    f"pipeline_last send rank={self.r} dst={(self.r + 1) % self.s} prefill={self.is_prefill}"
+                ):
+                    output = mx.distributed.send(
+                        output, (self.r + 1) % self.s, group=self.group
+                    )
             if cache is not None:
                 # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA)
                 # doesn't have .keys directly; access via first sub-cache.
                 _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
                 if hasattr(_cache, "keys"):  # pyright: ignore[reportAny]
                     _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
-            mx.eval(output)
+            with _hang_debug_watch(
+                f"pipeline_last eval_send rank={self.r} dst={(self.r + 1) % self.s} prefill={self.is_prefill}"
+            ):
+                mx.eval(output)
             if cache is not None and hasattr(_cache, "keys"):  # type: ignore
-                mx.eval(_cache.keys)  # type: ignore
+                with _hang_debug_watch(
+                    f"pipeline_last eval_cache_dep rank={self.r} dst={(self.r + 1) % self.s} prefill={self.is_prefill}"
+                ):
+                    mx.eval(_cache.keys)  # type: ignore
 
         if not self.is_prefill:
-            output = mx.distributed.all_gather(output, group=self.group)[
-                -output.shape[0] :
-            ]
-            mx.eval(output)
+            with _hang_debug_watch(
+                f"pipeline_last all_gather rank={self.r} world={self.s}"
+            ):
+                output = mx.distributed.all_gather(output, group=self.group)[
+                    -output.shape[0] :
+                ]
+            with _hang_debug_watch(
+                f"pipeline_last eval_all_gather rank={self.r} world={self.s}"
+            ):
+                mx.eval(output)
 
         return output
 
@@ -436,8 +545,8 @@ def _patch_one(target: object) -> None:
         if getattr(cls, "_exo_pipeline_patched", False):
             return
 
-        original_call = cls.__call__  # type: ignore[attr-defined]
-        call_signature = signature(original_call)  # type: ignore[arg-type]
+        original_call = cls.__call__
+        call_signature = signature(original_call)
 
         def patched_call(
             self: object,
@@ -463,7 +572,7 @@ def patched_call(
 
             return logits
 
-        cls.__call__ = patched_call  # type: ignore[method-assign]
+        cls.__call__ = patched_call
         cls._exo_pipeline_patched = True  # type: ignore[attr-defined]
 
     _patch_one(model)

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

@@ -466,16 +466,37 @@ def combined_progress_callback(processed: int, total: int) -> None:
     effective_prefill_step_size = (
         prefill_step_size // min(4, group_size) if is_pipeline else prefill_step_size
     )
+
+    # Pipeline-parallel prefill currently wedges when the prompt fits inside a
+    # single per-rank chunk (i.e. ``n_real == 1`` in the prefill loop). This
+    # was observed during Gemma 4 warmup with a 23-token prompt and reproduced
+    # in the issue tracker for any short prompt under pipeline parallelism.
+    # PR #101's pipeline-prefill widening (cb2dc68e) removed the prior
+    # ``num_tokens >= prefill_step_size`` gate to fix a different Gemma warmup
+    # hang in stream_generate; the result was that *every* short prompt now
+    # routed through pipeline_parallel_prefill and hit the new wedge. The
+    # narrower gate below preserves PR #101's fix for non-trivial prompts
+    # (where pipeline_parallel_prefill is the right path) while routing
+    # short / warmup prompts back through stream_generate, which is known to
+    # work for that shape. See PR review thread on PR #103 for the full
+    # diagnosis.
+    pipeline_chunks = (
+        (num_tokens + effective_prefill_step_size - 1) // effective_prefill_step_size
+        if effective_prefill_step_size > 0
+        else 0
+    )
+    use_pipeline_prefill = is_pipeline and pipeline_chunks >= 2
     logger.info(
         "Prefill path selected: "
-        f"{'pipeline_parallel_prefill' if is_pipeline else 'stream_generate'} "
-        f"(rank={rank}, prompt_tokens={num_tokens}, "
+        f"{'pipeline_parallel_prefill' if use_pipeline_prefill else 'stream_generate'} "
+        f"(rank={rank}, prompt_tokens={num_tokens}, is_pipeline={is_pipeline}, "
         f"prefill_step_size_input={prefill_step_size}, "
-        f"prefill_step_size_effective={effective_prefill_step_size})"
+        f"prefill_step_size_effective={effective_prefill_step_size}, "
+        f"pipeline_chunks={pipeline_chunks})"
     )
 
     try:
-        if is_pipeline:
+        if use_pipeline_prefill:
             set_pipeline_queue_sends(model, queue_sends=True)
             assert group is not None, "Pipeline prefill requires a distributed group"
             with _hang_debug_watch(

src/exo/worker/runner/llm_inference/runner.py

@@ -69,6 +69,13 @@
 from .tool_parsers import make_mlx_parser
 
 
+def _should_skip_llm_warmup() -> bool:
+    # Temporary experiment for Gemma 4 pipeline bring-up: let runners report
+    # ready without synthetic warmup so we can probe the first real prompt path.
+    # Set SKULK_FORCE_LLM_WARMUP=1 to restore the normal behavior.
+    return os.environ.get("SKULK_FORCE_LLM_WARMUP") != "1"
+
+
 class ExitCode(str, Enum):
     AllTasksComplete = "AllTasksComplete"
     Shutdown = "Shutdown"
@@ -243,7 +250,14 @@ def on_layer_loaded(layers_loaded: int, total: int) -> None:
                 self.update_status(RunnerWarmingUp())
                 self.acknowledge_task(task)
 
-                self.generator.warmup()
+                if _should_skip_llm_warmup():
+                    logger.warning(
+                        "Skipping LLM warmup and marking runner ready "
+                        "(temporary debug bypass; set SKULK_FORCE_LLM_WARMUP=1 "
+                        "to restore synthetic warmup)"
+                    )
+                else:
+                    self.generator.warmup()
 
                 logger.info(
                     f"runner initialized in {time.time() - self.setup_start_time} seconds"

src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py

@@ -34,9 +34,15 @@ def __init__(self) -> None:
         self.layers: list[object] = []
 
 
-def test_prefill_uses_pipeline_parallel_path_for_short_pipeline_prompts(
+def test_prefill_uses_pipeline_parallel_path_for_long_pipeline_prompts(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
+    """Long prompts on pipeline-parallel models should use pipeline_parallel_prefill.
+
+    The threshold is two effective per-rank chunks. With group_size=3 and
+    prefill_step_size=4096, the effective per-rank chunk size is 1365, so a
+    5000-token prompt produces 4 chunks and qualifies.
+    """
     calls: list[str] = []
     fake_cache = _FakeCache()
 
@@ -70,15 +76,71 @@ def _pipeline_parallel_prefill(**kwargs: object) -> None:
         model=_FakeModel(),
         tokenizer=object(),
         sampler=object(),
-        prompt_tokens=list(range(128)),
+        prompt_tokens=list(range(5000)),
         cache=[fake_cache],
         group=_FakeGroup(),
         on_prefill_progress=None,
         distributed_prompt_progress_callback=None,
     )
 
     assert calls == ["pipeline_parallel_prefill"]
-    assert prefill_tokens == 128
+    assert prefill_tokens == 5000
+    assert snapshots == []
+    assert prefill_tps >= 0.0
+    assert fake_cache.trim_calls == [2]
+
+
+def test_prefill_uses_stream_generate_for_short_pipeline_prompts(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Short prompts that fit in a single per-rank chunk must avoid pipeline prefill.
+
+    Regression for the Gemma 4 warmup wedge: pipeline_parallel_prefill hangs
+    when ``n_real == 1`` (the prompt fits inside a single per-rank chunk),
+    so any such prompt — including the 23-token warmup prompt — must route
+    through stream_generate even on pipeline-parallel models. PR #101's
+    pipeline-prefill widening removed this guard and reintroduced the hang
+    on Gemma 4; this test locks the guard back in.
+    """
+    calls: list[str] = []
+    fake_cache = _FakeCache()
+
+    monkeypatch.setattr(generate_module, "mx_barrier", lambda _group: None)
+    monkeypatch.setattr(
+        generate_module,
+        "_has_pipeline_communication_layer",
+        lambda _model: True,
+    )
+
+    def _pipeline_parallel_prefill(**kwargs: object) -> None:
+        calls.append("pipeline_parallel_prefill")
+
+    monkeypatch.setattr(
+        generate_module,
+        "pipeline_parallel_prefill",
+        _pipeline_parallel_prefill,
+    )
+
+    def _stream_generate(*args: object, **kwargs: object) -> Iterator[object]:
+        calls.append("stream_generate")
+        yield object()
+
+    monkeypatch.setattr(generate_module, "stream_generate", _stream_generate)
+
+    prefill_tps, prefill_tokens, snapshots = generate_module.prefill(
+        model=_FakeModel(),
+        tokenizer=object(),
+        sampler=object(),
+        prompt_tokens=list(range(23)),  # exactly the gemma 4 warmup prompt size
+        cache=[fake_cache],
+        group=_FakeGroup(),
+        on_prefill_progress=None,
+        distributed_prompt_progress_callback=None,
+    )
+
+    assert calls == ["stream_generate"]
+    assert "pipeline_parallel_prefill" not in calls
+    assert prefill_tokens == 23
     assert snapshots == []
     assert prefill_tps >= 0.0
     assert fake_cache.trim_calls == [2]