Fix pipeline prefill short-prompt wedge + dashboard kv-backend display - #104
Conversation
…nd display Two small bugs surfaced during the rotorquant rollout, neither caused by rotorquant itself but both blocking real usage: 1. Pipeline-parallel prefill wedges when ``n_real == 1`` (the prompt fits inside a single per-rank chunk). PR #101's ``cb2dc68e`` removed the prior ``num_tokens >= prefill_step_size`` gate to fix a Gemma warmup hang in stream_generate; the unintended consequence was that *every* short prompt now routed through pipeline_parallel_prefill and hit a different wedge there. Reproduced on Gemma 4 26B during 23-token warmup. Fix: route a prompt through pipeline_parallel_prefill only when it produces at least two effective per-rank chunks. Long prompts still use the pipeline path PR #101 was preserving; short prompts go through stream_generate, which is known to work for that shape. 2. ``GET /config`` reads only ``EXO_KV_CACHE_BACKEND`` when ``skulk.yaml`` exists on disk, ignoring the newer ``SKULK_KV_CACHE_BACKEND`` env var. The ``fileExists=False`` branch already handles both, so the dashboard silently displayed "default" even though the runtime was honoring the SKULK_ override correctly. One-line fix to match the other branch. Tests: - Updated test_prefill_uses_pipeline_parallel_path_for_long_pipeline_prompts to use a 5000-token prompt (4 chunks at group_size=3, qualifies for the pipeline path) - Added test_prefill_uses_stream_generate_for_short_pipeline_prompts as a regression test for the Gemma 4 23-token warmup wedge Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 32593224d8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Fixes two rollout-blocking issues: a pipeline-parallel prefill hang on short prompts, and an incorrect KV-backend value displayed in the dashboard when skulk.yaml exists.
Changes:
- Add a new prefill path gate intended to avoid
pipeline_parallel_prefillwhen prompts fit in a single per-rank chunk. - Update
/configto consistently preferSKULK_KV_CACHE_BACKENDoverEXO_KV_CACHE_BACKENDregardless of whether the config file exists. - Update/add unit tests to cover long-prompt pipeline prefill selection and short-prompt stream_generate selection.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/exo/worker/engines/mlx/generator/generate.py |
Adds a pipeline_chunks/use_pipeline_prefill gate and updates logging to reflect the selected prefill path. |
src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py |
Updates the existing test to use a long prompt for pipeline prefill and adds a regression test ensuring short prompts use stream_generate. |
src/exo/api/main.py |
Makes /config’s fileExists=True branch use the same env var fallback chain as fileExists=False for kv_cache_backend. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: edf7145a0d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 70e4192adc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cc0c0332ac
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b91a99b62d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4c6000f09e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 66392dfffb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 097685eec1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 19 out of 19 changed files in this pull request and generated 6 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6ae70786e1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 23 out of 23 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2302c410b9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 23 out of 23 changed files in this pull request and generated 5 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e7a1d591a4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9163d8d659
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 33f8c0f358
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 24 out of 24 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 26 out of 26 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (1)
src/exo/worker/engines/mlx/cache.py:596
- The warning message hard-codes
EXO_KV_CACHE_BACKEND, but this code now also readsSKULK_KV_CACHE_BACKENDviapreferred_env_value. This can mislead operators when the invalid value actually came from the SKULK env var. Consider updating the log text to refer to the effective KV backend env (or mention both SKULK/EXO) so the warning points users to the right knob.
if backend not in VALID_KV_CACHE_BACKENDS:
logger.warning(
f"Unknown EXO_KV_CACHE_BACKEND={backend!r}; falling back to {DEFAULT_KV_CACHE_BACKEND!r}"
)
return DEFAULT_KV_CACHE_BACKEND
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Summary
Two small bugs surfaced during the rotorquant rollout (PR #103) testing. Neither is caused by rotorquant itself, but both block real usage of the default backend on pipeline-parallel clusters.
Bug 1 — pipeline_parallel_prefill wedge on short prompts (severity 5, blocks gemma 4)
PR #101's commit cb2dc68 removed the prior `num_tokens >= prefill_step_size` gate in `prefill()` to fix a Gemma warmup hang in `stream_generate`. The unintended consequence was that every short prompt now routed through `pipeline_parallel_prefill`, which has its own wedge when `n_real == 1` (the prompt fits entirely inside a single per-rank chunk).
Reproduced on: `mlx-community/gemma-4-26b-a4b-it-4bit` on a 3-node pipeline cluster with the default KV cache backend, hanging in warmup on the 23-token warmup prompt. Stack trace ends inside `mx.eval(x)` at `auto_parallel.py:168`, called from gemma4's first language-model layer, called from the first `pipeline_parallel_prefill` iteration. `n_real == 1` because 23 tokens fits in one chunk (`effective_prefill_step_size = 4096 // 3 = 1365`).
This was originally surfaced during testing of PR #103 (rotorquant), where rotorquant appeared to work and default appeared to break — but it turned out rotorquant just dodged the wedge entirely because quantized backends force `SequentialGenerator`, which never calls `pipeline_parallel_prefill`. I flagged the prefill-gate removal as severity 4 in my PR #101 review and asked for an inline comment about the trade-off; nobody knew it would actually wedge until we hit it on hardware.
Fix
Restore the short-prompt guard, but with a smarter threshold: route a prompt through `pipeline_parallel_prefill` only when it would produce at least two effective per-rank chunks. This preserves PR #101's intent (pipeline_parallel_prefill is the right path for non-trivial prompts) while routing warmup and short conversational turns back through `stream_generate`, which is known to work for that shape.
```python
pipeline_chunks = (num_tokens + effective_prefill_step_size - 1) // effective_prefill_step_size
use_pipeline_prefill = is_pipeline and pipeline_chunks >= 2
```
For a 3-node cluster (`group_size=3`, `effective_step=1365`):
The cutoff is conservative — we'd rather route a borderline prompt through the slower-but-safe `stream_generate` than into the wedge.
Tests
Bug 2 — dashboard reports wrong KV backend when skulk.yaml exists (severity 4, cosmetic but confusing)
`GET /config` has two branches: `fileExists=False` reads both `SKULK_KV_CACHE_BACKEND` and `EXO_KV_CACHE_BACKEND`, but `fileExists=True` only reads `EXO_KV_CACHE_BACKEND`. So if you have a `skulk.yaml` and set `SKULK_KV_CACHE_BACKEND=rotorquant_adaptive`, the runtime correctly honors it (`constants.py` checks both), but the dashboard Settings panel silently displays "default".
One-line fix: make the `fileExists=True` branch use the same fallback chain as the `fileExists=False` branch.
Testing
Relationship to PR #103 (rotorquant)
This PR is independent and can land in any order:
If both land, gemma 4 will work on both default and rotorquant_adaptive; today only rotorquant_adaptive works on this hardware.