Add RotorQuant KV cache backend with deferred prefill on Metal - #103
Add RotorQuant KV cache backend with deferred prefill on Metal#103ttupper92618 wants to merge 22 commits into
Conversation
Pure-MLX port of IsoQuant 3-bit (block-diagonal quaternion rotations + Lloyd-Max centroids) with deferred prefill — a contribution that does not exist in any upstream project, since the llama.cpp fork ships the deferred-prefill flush CUDA-only. GQA-native, no fallback for grouped-query models. Also adds the "Everything Different About Skulk" living section to README.md and website/docs/everything-different.md so the divergence from upstream exo lives in one canonical place going forward. Tracked by #102. 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: 455759834a
ℹ️ 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
This PR adds a new MLX/Metal-native KV-cache backend family (“RotorQuant”) implementing IsoQuant 3-bit compression with a deferred-prefill path on Metal, and wires it into Skulk’s runtime/config/UI/docs so it can be selected cluster-wide.
Changes:
- Introduces
rotorquant/rotorquant_adaptiveKV cache backends (IsoQuant 3-bit with quaternion block rotations + norm correction), including a deferred-prefill state machine. - Integrates the new backends into MLX engine cache selection, runner constraints, and cluster config validation.
- Updates docs and dashboard UI to expose/describe the new backend options and adds a new “Everything Different About Skulk” canonical page.
Reviewed changes
Copilot reviewed 17 out of 18 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| website/sidebars.ts | Adds the new “Everything Different…” doc to the Docusaurus sidebar. |
| website/docs/kv-cache-backends.md | Documents RotorQuant backends, env vars, and updates the backend comparison table. |
| website/docs/intro.md | Links to the new “Everything Different…” page for onboarding. |
| website/docs/everything-different.md | New canonical living doc enumerating Skulk divergences from upstream exo. |
| src/exo/worker/runner/llm_inference/runner.py | Forces sequential generation mode for the new quantized backends. |
| src/exo/worker/engines/mlx/rotorquant/tests/test_iso_quantizer.py | Adds unit tests for rotation + quantizer correctness properties. |
| src/exo/worker/engines/mlx/rotorquant/tests/test_cache.py | Adds tests for deferred prefill behavior and cache lifecycle operations. |
| src/exo/worker/engines/mlx/rotorquant/tests/init.py | Initializes the rotorquant test package. |
| src/exo/worker/engines/mlx/rotorquant/tables.py | Vendors quaternion + centroid lookup tables and exposes accessors. |
| src/exo/worker/engines/mlx/rotorquant/rotation.py | Implements quaternion-based block-diagonal rotations in MLX. |
| src/exo/worker/engines/mlx/rotorquant/quantizer.py | Implements IsoQuantizer quantize/dequantize with norm correction. |
| src/exo/worker/engines/mlx/rotorquant/cache.py | Implements RotorQuantKVCache and factory helpers (including deferred-prefill). |
| src/exo/worker/engines/mlx/rotorquant/init.py | Exposes public rotorquant backend APIs. |
| src/exo/worker/engines/mlx/constants.py | Extends KV backend literals and adds RotorQuant env-var settings. |
| src/exo/worker/engines/mlx/cache.py | Wires rotorquant backends into make_kv_cache selection logic. |
| src/exo/store/config.py | Extends InferenceConfig.kv_cache_backend allowed values. |
| README.md | Adds a new “Everything Different…” section and RotorQuant env var entries. |
| dashboard-react/src/components/layout/SettingsPanel.tsx | Adds RotorQuant options to the KV backend dropdown + tooltip text updates. |
Comments suppressed due to low confidence (1)
dashboard-react/src/components/layout/SettingsPanel.tsx:475
- The tooltip/hint claims “Incompatible models fall back to Default automatically” (and specifically mentions GQA/non-power-of-two head_dim). In code, only the
optiqbackend has an explicit compatibility check + fallback;turboquant/rotorquantpaths raise on incompatibility (e.g., unsupported cache layouts or head_dim not divisible by 128). Either implement a similar fallback for these backends in make_kv_cache, or soften the UI copy to reflect that some incompatibilities will error instead of silently falling back.
`Takes effect on next model launch. Incompatible models fall back to Default automatically.`
}
/>
</FieldLabel>
<Select value={kvBackend} onChange={(e) => setKvBackend(e.target.value)} disabled={!!envOverride}>
<option value="default">Default (no quantization)</option>
<option value="rotorquant_adaptive">RotorQuant Adaptive (recommended)</option>
<option value="rotorquant">RotorQuant</option>
<option value="optiq">OptiQ (rotation-based)</option>
<option value="turboquant_adaptive">TurboQuant Adaptive</option>
<option value="turboquant">TurboQuant</option>
<option value="mlx_quantized">MLX Quantized (requires SKULK_KV_CACHE_BITS env)</option>
</Select>
{envOverride ? (
<HintText>Overridden by SKULK_KV_CACHE_BACKEND environment variable. Remove the env var to configure here.</HintText>
) : (
<HintText>Changes take effect on the next model launch. Models with incompatible architectures (GQA, non-power-of-two head_dim) will automatically fall back to default.</HintText>
)}
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- _flush_deferred no longer asserts on empty pending buffers; a decode-shaped first call (1-token prompt) is now a clean no-op that exits the deferred phase and lets the normal quantize-on- insert path handle the new token. Adds test_first_call_decode_ shaped_does_not_crash as a regression. - RotorQuantKVCache.state now returns a pytree of mx.array leaves only — the previous str tag would have raised inside mx.eval and mx.save_safetensors. Phase is disambiguated by tuple length (2 = pending, 4 = live) and tracked in meta_state. Adds test_state_round_trips_in_both_phases. - README example commands now use uv run skulk consistently instead of the deprecated uv run exo alias. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 53a0907f79
ℹ️ 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".
# Conflicts: # src/exo/worker/engines/mlx/cache.py
When deferred prefill is active and a prompt is split across multiple update_and_fetch calls (num_steps > 1), the append-to-pending branch returned without advancing cache.offset. Downstream attention reads cache.offset for RoPE position, so any prompt longer than the prefill chunk size (e.g. 4096) had every chunk after the first positioned as if it began at token 0, producing wrong attention logits exactly at the long-context regime RotorQuant is meant to serve. Fix: advance offset on every append in both phases. _flush_deferred now writes the quantized rows into [offset - num_pending, offset) instead of double-counting. trim() and the state setter were updated to keep offset consistent in the deferred phase as well, and size() collapses to a single offset read. Adds test_deferred_prefill_advances_offset_across_chunks as a regression for the multi-chunk position bug.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 17 out of 18 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| ### Inference and KV cache | ||
|
|
||
| - **RotorQuant KV cache backend** — pure-MLX port of IsoQuant 3-bit (block-diagonal quaternion rotations + Lloyd-Max centroids) with **deferred prefill on Metal**, a contribution that does not exist in any upstream project (the llama.cpp fork ships it CUDA-only). GQA-native; no fallback for grouped-query models. See [docs/kv-cache-backends.md](docs/kv-cache-backends.md). |
There was a problem hiding this comment.
This bullet links to docs/kv-cache-backends.md, but that page currently doesn’t mention RotorQuant (it still lists only default/mlx_quantized/turboquant/optiq). Either update docs/kv-cache-backends.md in this PR to include RotorQuant, or change the link to the up-to-date page under website/docs/kv-cache-backends.md / the published docs URL so readers don’t land on stale information.
| - **RotorQuant KV cache backend** — pure-MLX port of IsoQuant 3-bit (block-diagonal quaternion rotations + Lloyd-Max centroids) with **deferred prefill on Metal**, a contribution that does not exist in any upstream project (the llama.cpp fork ships it CUDA-only). GQA-native; no fallback for grouped-query models. See [docs/kv-cache-backends.md](docs/kv-cache-backends.md). | |
| - **RotorQuant KV cache backend** — pure-MLX port of IsoQuant 3-bit (block-diagonal quaternion rotations + Lloyd-Max centroids) with **deferred prefill on Metal**, a contribution that does not exist in any upstream project (the llama.cpp fork ships it CUDA-only). GQA-native; no fallback for grouped-query models. See [KV cache backends documentation](https://foxlight-foundation.github.io/Skulk/kv-cache-backends/). |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ff21598872
ℹ️ 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".
- make_kv_cache now checks head_dim % ISO3_BLOCK_SIZE before selecting rotorquant/rotorquant_adaptive; models with incompatible heads (e.g. 64-d) fall back to the default cache with a warning instead of crashing at first token. - Extract _make_default_cache helper so fallback paths don't duplicate the make_cache/KVCache logic. - Fix kv-cache-backends.md claiming 64-d heads are supported (they are not — ISO3_BLOCK_SIZE is 128). - Soften SettingsPanel tooltip/hint: only OptiQ auto-falls-back; other backends will error on incompatible models.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f40e9d4143
ℹ️ 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".
| first_layer = model.layers[0] | ||
| attn = getattr(first_layer, "self_attn", first_layer) | ||
| model_head_dim: int = getattr(attn, "head_dim", ISO3_BLOCK_SIZE) | ||
| else: |
There was a problem hiding this comment.
The rotorquant preflight reads head_dim only from model.layers[0] and defaults to ISO3_BLOCK_SIZE when that layer has no attention metadata, so mixed-cache models can pass this gate even if later KV attention layers are not divisible by 128. In that case backend selection succeeds but inference still crashes later when RotorQuantKVCache._ensure_quantizer rejects the real layer shape on first token; preflight should inspect the actual KV-bearing layers (or all attention layers) rather than a single first layer.
Useful? React with 👍 / 👎.
With is_prefill=False on the stream_generate fallback path for short prompts, PipelineLastLayer runs mx.distributed.all_gather on every forward pass. generate_step discards prefill logits and only evaluates cache states, leaving the all_gather as an unevaluated zombie collective. When a later mx.eval triggers it, the ranks are desynchronized and the collective deadlocks. Observed as a warmup hang on a 3-rank pipeline with Gemma 4 26B. The KV backend is irrelevant — the hang reproduces with any backend (logs confirmed the backend fell back to default before the hang). Fix: set is_prefill=True for ALL pipeline models during prefill, including the stream_generate fallback. This disables the all_gather and relies on point-to-point send/recv only. The first decode token is wrong on non-last ranks (they see rank-local activations), but prefill() trims those tokens before real generation begins and the finally block resets is_prefill=False for subsequent decode.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 18 out of 19 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| <HintText>Overridden by SKULK_KV_CACHE_BACKEND environment variable. Remove the env var to configure here.</HintText> | ||
| ) : ( | ||
| <HintText>Changes take effect on the next model launch. Models with incompatible architectures (GQA, non-power-of-two head_dim) will automatically fall back to default.</HintText> | ||
| <HintText>Changes take effect on the next model launch. OptiQ falls back to Default for unsupported architectures; other backends will error on incompatible models.</HintText> |
There was a problem hiding this comment.
Same issue as the tooltip: this hint claims only OptiQ falls back and “other backends will error”, but RotorQuant can also fall back to Default when it detects an unsupported head_dim (not divisible by 128). Align the hint text with the backend behavior.
Three changes: 1. get_kv_cache_backend() now reads os.environ at call time via preferred_env_value() instead of the frozen import-time constant. This ensures runtime updates (from config sync, dashboard, or coordinator) are always visible to the runner. 2. When the user sets SKULK_KV_CACHE_BACKEND at launch, main.py writes it back to exo.yaml so the config file stays in sync. Previously the config file could silently override the env var because constants.py read at import time, before main.py had a chance to apply the config. 3. Update tests to use patch.dict(os.environ) instead of patching the removed module-level KV_CACHE_BACKEND constant.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8545421141
ℹ️ 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".
| if launch_backend: | ||
| if update_config_field("inference", "kv_cache_backend", launch_backend): | ||
| logger.info( |
There was a problem hiding this comment.
This persists launch_backend directly from the environment into config without checking it against supported values. A typo like SKULK_KV_CACHE_BACKEND=rotorqunt is tolerated at runtime (the cache selector falls back to default), but it gets written to disk and then load_exo_config() will fail on the next restart because InferenceConfig.kv_cache_backend is a Literal[...]; that turns a temporary env mistake into a persistent startup failure until the file is manually fixed.
Useful? React with 👍 / 👎.
| with path.open("w") as f: | ||
| yaml.dump(raw, f, default_flow_style=False, sort_keys=False) |
There was a problem hiding this comment.
Writing the synced KV backend uses an unconditional file write with no error handling, and the startup path calls this before the node is fully running. If skulk.yaml exists but is read-only (for example, common config-map style mounts) or the filesystem is temporarily unwritable, this raises and aborts startup even though the env var already provides a valid backend; the sync should fail gracefully instead of taking the node down.
Useful? React with 👍 / 👎.
When SKULK_KV_CACHE_BACKEND is set at launch, the coordinator's _sync_config was overwriting both the env var AND the local exo.yaml with whatever the cluster broadcast (typically the old value from another node). This silently undid the user's intent. Fix: when _SKULK_KV_BACKEND_USER_SET is active, patch the incoming config YAML to preserve the user's KV backend before writing it to disk, and skip the env var update. The user's launch-time value remains the source of truth until the process exits.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 22 out of 23 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| `• Default — No cache quantization. Best baseline quality, highest memory use.\n` + | ||
| `• OptiQ — Rotation-based quantization via mlx-optiq. Best long-context quality.\n` + | ||
| `• RotorQuant Adaptive — IsoQuant 3-bit with deferred prefill, FP16 edge layers. Recommended.\n` + | ||
| `• RotorQuant — IsoQuant 3-bit on all KV layers. Most aggressive compression with deferred prefill.\n` + | ||
| `• OptiQ — Rotation-based quantization via mlx-optiq. Good long-context quality, no GQA support.\n` + | ||
| `• TurboQuant Adaptive — Quantizes middle KV layers, keeps edge layers in FP16. Proven stable.\n` + | ||
| `• TurboQuant — Quantizes all KV layers. Most aggressive compression, higher quality risk.\n` + | ||
| `• TurboQuant — Quantizes all KV layers. Most aggressive non-rotorquant compression.\n` + | ||
| `• MLX Quantized — MLX's built-in cache quantization.\n\n` + | ||
| `Takes effect on next model launch. Incompatible models fall back to Default automatically.` | ||
| `Takes effect on next model launch. OptiQ falls back to Default for unsupported architectures; other backends will error on incompatible models.` | ||
| } |
There was a problem hiding this comment.
The tooltip text says “other backends will error on incompatible models”, but the Python backend selection now explicitly falls back to the default cache for RotorQuant when head_dim is not divisible by 128 (src/exo/worker/engines/mlx/cache.py). Please update this copy to reflect the actual behavior (e.g., mention RotorQuant may also fall back in some cases, while other incompatibilities still raise).
| <HintText>Overridden by SKULK_KV_CACHE_BACKEND environment variable. Remove the env var to configure here.</HintText> | ||
| ) : ( | ||
| <HintText>Changes take effect on the next model launch. Models with incompatible architectures (GQA, non-power-of-two head_dim) will automatically fall back to default.</HintText> | ||
| <HintText>Changes take effect on the next model launch. OptiQ falls back to Default for unsupported architectures; other backends will error on incompatible models.</HintText> | ||
| )} |
There was a problem hiding this comment.
Same mismatch as the tooltip: this hint says “other backends will error on incompatible models”, but RotorQuant can fall back to Default for head_dim misalignment. Please align the hint text with the backend’s real fallback/error behavior so operators know what to expect.
| """Update a single field in the config file, preserving all other content. | ||
|
|
||
| Reads the raw YAML, patches ``raw[section][key] = value``, and writes | ||
| it back. Returns ``True`` if the file was updated, ``False`` if no | ||
| config file exists. |
There was a problem hiding this comment.
The docstring claims this function “preserv[es] all other content”, but round-tripping through PyYAML load/dump will drop comments and can reformat the file. Please either adjust the wording to reflect what is actually preserved (data only), or switch to a YAML library that can preserve formatting/comments if that guarantee is important.
| """Update a single field in the config file, preserving all other content. | |
| Reads the raw YAML, patches ``raw[section][key] = value``, and writes | |
| it back. Returns ``True`` if the file was updated, ``False`` if no | |
| config file exists. | |
| """Update a single field in the config file, preserving other config data. | |
| Reads the YAML into Python data, patches ``raw[section][key] = value``, | |
| and writes it back. This preserves other parsed configuration values, | |
| but comments and original formatting may be rewritten by PyYAML. | |
| Returns ``True`` if the file was updated, ``False`` if no config file | |
| exists. |
| raw[section] = sec | ||
| sec[key] = value | ||
| with path.open("w") as f: | ||
| yaml.dump(raw, f, default_flow_style=False, sort_keys=False) |
There was a problem hiding this comment.
This writes the config back with yaml.dump(), which can emit non-safe YAML tags and is generally discouraged compared to yaml.safe_dump() when persisting user-editable config. Prefer safe_dump (and consider setting explicit encoding) to avoid writing Python-specific tags into skulk.yaml/exo.yaml.
| yaml.dump(raw, f, default_flow_style=False, sort_keys=False) | |
| yaml.safe_dump(raw, f, default_flow_style=False, sort_keys=False) |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 032593b32b
ℹ️ 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".
| local_backend = os.environ.get("SKULK_KV_CACHE_BACKEND") | ||
| if local_backend: | ||
| inference = raw.get("inference") | ||
| if isinstance(inference, dict): | ||
| inference["kv_cache_backend"] = local_backend |
There was a problem hiding this comment.
When a launch-time env override is active, _sync_config copies SKULK_KV_CACHE_BACKEND directly into inference.kv_cache_backend without validating it. A typo such as rotorqunt will be persisted the next time config sync runs, and on restart load_exo_config() can fail Literal validation before the node finishes booting, requiring manual file repair. Validate local_backend against supported backends (or skip persisting invalid values) before writing YAML.
Useful? React with 👍 / 👎.
| set_pipeline_prefill(model, is_prefill=is_pipeline) | ||
|
|
||
| with _hang_debug_watch(f"prefill barrier rank={rank} group_size={group_size}"): | ||
| mx_barrier(group) |
There was a problem hiding this comment.
The prefill flag is set to True for pipeline models before entering the try/finally that resets it. If mx_barrier(group) raises (for example during a transient distributed failure), execution exits before the reset runs, leaving layers in prefill mode for subsequent requests; that can skip PipelineLastLayer's all_gather and produce rank-local logits during decode. Move the barrier into the protected block or add a reset path around barrier failures.
Useful? React with 👍 / 👎.
The last rank's PipelineLastLayer sliced all_gather output before eval, letting MLX satisfy the slice from local data without executing the collective — leaving other ranks blocked forever. Now eval the unsliced gathered tensor so every rank participates, then slice. Also adds generic _cache_dep_anchor()/_set_cache_dep_anchor() helpers so mx.depends works with RotorQuantKVCache (which uses _key_indices or _pending_keys instead of .keys). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f587d2b8f1
ℹ️ 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".
| for attr in ("_key_indices", "_pending_keys"): | ||
| if hasattr(cache, attr) and isinstance(getattr(cache, attr, None), mx.array): | ||
| setattr(cache, attr, value) |
There was a problem hiding this comment.
_set_cache_dep_anchor can assign a dependency-wrapped slice back to cache._key_indices when used with RotorQuantKVCache: _cache_dep_anchor reads the first leaf from cache.state, and RotorQuantKVCache.state exposes self._key_indices[..., :self.offset, :] (a truncated live view). Writing that view back to _key_indices drops the preallocated tail, so each subsequent decode step fails the capacity check and re-enters _expand_storage, causing repeated concat/realloc work in pipeline mode. On long generations this turns decode into avoidable O(T²)-style memory churn and significant slowdown.
Useful? React with 👍 / 👎.
…ampling The pipeline decode path used all_gather to broadcast the last rank's output to all ranks. This collective was prone to deadlock during warmup: generate_step prefetches one decode step ahead, and when the warmup generator is abandoned mid-iteration, the orphaned collective can mismatch with the post-warmup barrier, leaving one rank blocked indefinitely. Replace all_gather with point-to-point send/recv from the last rank to each other rank. Bilateral operations cannot deadlock from collective-ordering mismatches. Also switch distributed warmup to greedy sampling (temperature=0) so every rank samples the same token from the shared logits — stochastic sampling with per-rank random states can diverge. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 26 out of 27 changed files in this pull request and generated 4 comments.
Comments suppressed due to low confidence (1)
src/exo/worker/engines/mlx/constants.py:47
KV_CACHE_BACKENDis computed at import time. With the new design that readsSKULK_KV_CACHE_BACKENDdynamically (seeget_kv_cache_backend()), this constant can become stale after config sync / dashboard updates, but it’s still used (e.g.,generator/generate.pygate formlx_quantized). Consider removing/privatizing this constant and migrating call sites toget_kv_cache_backend()(or documenting clearly thatKV_CACHE_BACKENDis a startup-time snapshot).
_kv_cache_backend_value = preferred_env_value(
"SKULK_KV_CACHE_BACKEND",
"EXO_KV_CACHE_BACKEND",
DEFAULT_KV_CACHE_BACKEND,
)
KV_CACHE_BACKEND: KVCacheBackend = cast(
KVCacheBackend,
_kv_cache_backend_value if _kv_cache_backend_value else DEFAULT_KV_CACHE_BACKEND,
)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| n_tasks = len(tasks) | ||
| logger.info(f"mx_all_gather_tasks: gathering counts (n_tasks={n_tasks})") | ||
| all_counts = cast( |
There was a problem hiding this comment.
mx_all_gather_tasks() is invoked as part of task/cancellation agreement; logging at INFO on every call can be very noisy in multi-rank runs. Consider switching this to DEBUG (or gating behind a verbose/debug env var) to avoid log spam.
| mx.array([n_tasks]), group=group, stream=cpu_stream | ||
| ).tolist(), | ||
| ) | ||
| logger.info(f"mx_all_gather_tasks: counts gathered: {all_counts}") |
There was a problem hiding this comment.
This second INFO log in mx_all_gather_tasks() will be emitted on every agreement round and includes per-rank counts, which can be high-volume. Recommend using DEBUG level or a debug flag so production logs aren’t flooded.
| logger.info(f"mx_all_gather_tasks: counts gathered: {all_counts}") | |
| logger.debug(f"mx_all_gather_tasks: counts gathered: {all_counts}") |
| assert self._pending_values is not None | ||
| self._pending_keys = mx.concatenate( | ||
| [self._pending_keys, keys.astype(mx.float16)], | ||
| axis=2, | ||
| ) | ||
| self._pending_values = mx.concatenate( | ||
| [self._pending_values, values.astype(mx.float16)], | ||
| axis=2, | ||
| ) |
There was a problem hiding this comment.
_append_pending() repeatedly concatenates growing fp16 buffers. For long prompts split into many prefill chunks, this becomes O(n²) in both time and allocations. Consider switching to a growable preallocated buffer (similar to _expand_storage) or accumulating chunks in a list and concatenating once during _flush_deferred.
| rotated by a fixed unit quaternion via the Hamilton product | ||
| ``q_L * v`` (treating ``v`` as a pure quaternion). The inverse uses | ||
| ``conj(q_L) * v`` because the rotation quaternion is unit. |
There was a problem hiding this comment.
The module docstring says the rotation uses q_L * v "treating v as a pure quaternion", but the implementation treats each 4D group as a general quaternion (v0,v1,v2,v3) (no zeroing of the scalar part) and applies left-multiplication. Please adjust the wording to avoid implying the standard 3D pure-quaternion rotation formula.
| rotated by a fixed unit quaternion via the Hamilton product | |
| ``q_L * v`` (treating ``v`` as a pure quaternion). The inverse uses | |
| ``conj(q_L) * v`` because the rotation quaternion is unit. | |
| transformed by left Hamilton multiplication with a fixed unit quaternion, | |
| ``q_L * v``, where ``v`` is the full 4D group interpreted as a general | |
| quaternion ``(v0, v1, v2, v3)``. The inverse applies ``conj(q_L) * v``, | |
| which undoes the left-multiplication because ``q_L`` is unit. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3100207151
ℹ️ 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".
| local_backend = os.environ.get("SKULK_KV_CACHE_BACKEND") | ||
| if local_backend: |
There was a problem hiding this comment.
The sync logic only preserves the local KV backend when local_backend is truthy, so an explicit empty SKULK_KV_CACHE_BACKEND is treated as if no override exists. In this codebase, blank SKULK_* values are intentionally meaningful for overriding legacy fallback behavior, so this path can overwrite inference.kv_cache_backend in skulk.yaml with the cluster value despite the launch-time override flag being set. That makes the override non-persistent across restarts and can unexpectedly switch the backend on next boot.
Useful? React with 👍 / 👎.
The distributed_prompt_progress_callback (agree_on_tasks/cancellations) does all_gather collectives. When called from stream_generate's progress callback, these collectives deadlock on 3-node pipelines because: 1. mlx_lm's generate_step prefetches the next decode token via _step(y) with mx.async_eval BEFORE firing the callback. That prefetch dispatches pipeline communication (sends/recvs or all_sum) to generation_stream. JACCL cannot run an all_gather while other collectives are in-flight. 2. Warmup also uses stream_generate and abandons the generator after one token, potentially leaving generation_stream in an inconsistent state for subsequent collectives. 3. For short prompts (< effective_prefill_step_size) the entire prefill takes milliseconds — no meaningful window for task/cancellation checks. Fix: pass only the local progress_callback (no distributed collectives) to stream_generate. pipeline_parallel_prefill keeps its distributed callback since it handles inter-chunk synchronization explicitly. Also keeps the earlier CPU stream fix for mx_all_gather_tasks to protect any remaining callers that run inside non-default stream contexts. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
JACCL's all_gather deadlocks intermittently on 3-node pipelines — the first all_gather in a sequence may succeed but subsequent ones hang indefinitely. all_sum is proven reliable across the entire codebase (mx_any, mx_barrier, PipelineLastLayer decode broadcast). Replace all_gather everywhere with all_sum using the zero-contribution slot pattern: each rank fills its own slot in a [world_size, ...] tensor and contributes zeros elsewhere. all_sum merges non-overlapping slots, giving the same result as all_gather without the JACCL deadlock. Changed: - mx_all_gather_tasks: both count and task-ID exchanges now use all_sum - warmup check_for_cancel_every: slotted all_sum instead of all_gather - KVCacheMemoryPressureTracker: slotted all_sum for pressure exchange - Tests: update mocks from fake_all_gather to fake_all_sum Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 26 out of 28 changed files in this pull request and generated 2 comments.
Files not reviewed (1)
- dashboard-react/package-lock.json: Language not supported
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def _cache_dep_anchor(cache: object) -> mx.array | None: | ||
| """Return an array from ``cache`` suitable for ``mx.depends`` ordering. | ||
|
|
||
| Standard ``KVCache`` exposes ``.keys``; quantized caches like | ||
| ``RotorQuantKVCache`` do not. Fall back to the first leaf of | ||
| ``.state`` so we can order cache evaluation relative to distributed | ||
| sends on *any* cache backend. | ||
| """ | ||
| # Fast path: mlx-lm KVCache and friends | ||
| keys = getattr(cache, "keys", None) | ||
| if isinstance(keys, mx.array): | ||
| return keys | ||
|
|
||
| # Generic path: first array in cache.state | ||
| state: object = getattr(cache, "state", None) | ||
| if isinstance(state, (list, tuple)): | ||
| for leaf in list(cast(list[object], state)): | ||
| if isinstance(leaf, mx.array): | ||
| return leaf | ||
| return None |
There was a problem hiding this comment.
_cache_dep_anchor falls back to using the first leaf of cache.state for dependency ordering. For RotorQuantKVCache, state intentionally returns slices up to offset (not the full preallocated backing arrays). When _set_cache_dep_anchor writes that depended slice back into _key_indices, it discards the cache’s extra capacity, so the next token append will trigger _expand_storage and reallocate/concatenate again (effectively once per token in pipeline mode). Prefer anchoring on a stable backing array (e.g., check for _key_indices/_pending_keys before state, or add an explicit public anchor on caches) so mx.depends ordering doesn’t collapse preallocation capacity.
| # --- Phase 1: exchange counts via all_sum --- | ||
| # Each rank contributes its count in its own slot of a [world_size] | ||
| # vector; other slots are zero. all_sum merges them. | ||
| n_tasks = len(tasks) | ||
| logger.info(f"mx_all_gather_tasks: gathering counts (n_tasks={n_tasks})") | ||
| counts_vec = [0] * world_size | ||
| counts_vec[rank] = n_tasks | ||
| all_counts = cast( | ||
| list[int], | ||
| mx.distributed.all_gather(mx.array([n_tasks]), group=group).tolist(), | ||
| mx.distributed.all_sum( | ||
| mx.array(counts_vec, dtype=mx.int32), | ||
| group=group, | ||
| stream=cpu_stream, | ||
| ).tolist(), | ||
| ) | ||
| logger.info(f"mx_all_gather_tasks: counts gathered: {all_counts}") | ||
| max_tasks = max(all_counts) |
There was a problem hiding this comment.
mx_all_gather_tasks adds multiple unconditional INFO logs (gathering counts / counts gathered). Since this is called from agree_on_tasks/agree_on_cancellations in the hot path, it can flood logs and add overhead. Consider lowering these to DEBUG or gating them behind a troubleshooting flag so normal inference logs stay clean.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 13a5f47c03
ℹ️ 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".
| if len(seq) == 0: | ||
| return |
There was a problem hiding this comment.
The state setter documents that a 0-length sequence represents an uninitialized cache, but this branch immediately returns without clearing offset or any pending/live storage fields. If a previously used cache instance is later restored with state=[] (for example, when loading or resetting an empty cache snapshot), stale KV contents and token counts are retained instead of producing a clean cache, which can leak old state into subsequent generations.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 27 out of 29 changed files in this pull request and generated 1 comment.
Files not reviewed (1)
- dashboard-react/package-lock.json: Language not supported
Comments suppressed due to low confidence (1)
src/exo/worker/engines/mlx/utils_mlx.py:1218
- encode_task_id() uses int.from_bytes(utf8_task_id[i:i+1]) without providing the required byteorder argument, which will raise TypeError at runtime the first time mx_all_gather_tasks() runs. For single-byte extraction you can just use utf8_task_id[i] (bytes indexing already yields an int), or call int.from_bytes(..., byteorder="little").
def encode_task_id(task_id: TaskId) -> list[int]:
utf8_task_id = task_id.encode()
return [
int.from_bytes(utf8_task_id[i : i + 1]) for i in range(len(utf8_task_id))
]
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def decode_task_id(encoded_task_id: list[int]) -> TaskId: | ||
| return TaskId( | ||
| bytes.decode(b"".join((x).to_bytes(length=1) for x in encoded_task_id)) | ||
| ) |
There was a problem hiding this comment.
decode_task_id() builds bytes via (x).to_bytes(length=1) but int.to_bytes also requires a byteorder argument; this will raise TypeError at runtime. Use x.to_bytes(1, byteorder="little") (or "big") consistently with encode_task_id, or simplify by constructing bytes(encoded_task_id) if values are already 0-255.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 04893db6d8
ℹ️ 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".
| if len(seq) == 2: | ||
| self._pending_keys = seq[0] | ||
| self._pending_values = seq[1] | ||
| self._pending_active = True |
There was a problem hiding this comment.
The len(seq) == 2 restore path only assigns _pending_keys/_pending_values and flips _pending_active, but it does not clear any existing quantized _key_*/_value_* tensors from prior use. If a reused cache object is restored from a deferred snapshot, _flush_deferred() may treat stale quantized storage as valid capacity (since _expand_storage checks capacity, not shape compatibility) and then write pending tensors into mismatched old shapes on the first decode-shaped update, causing runtime shape errors. Reset the non-deferred storage when entering deferred mode.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6c4ba6f3dd
ℹ️ 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".
| kv_group_size=KV_GROUP_SIZE, | ||
| kv_bits=KV_BITS, | ||
| prompt_progress_callback=combined_progress_callback, | ||
| prompt_progress_callback=progress_callback, |
There was a problem hiding this comment.
In the non-pipeline prefill path, stream_generate now receives prompt_progress_callback=progress_callback instead of the combined callback, so distributed_prompt_progress_callback is never invoked during prefill. This regresses distributed tensor-parallel behavior (where group is set but _has_pipeline_communication_layer(model) is false): cancellation and cross-rank task agreement checks are deferred until decode starts, which can leave long prefill requests effectively un-cancellable and delay queue synchronization by the full prompt prefill time.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 32 out of 34 changed files in this pull request and generated 1 comment.
Files not reviewed (1)
- dashboard-react/package-lock.json: Language not supported
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| uuid_byte_length = 36 | ||
| cpu_stream = mx.default_stream(mx.Device(mx.cpu)) | ||
| world_size: int = 1 if group is None else group.size() | ||
| rank: int = 0 if group is None else group.rank() | ||
|
|
||
| # --- Phase 1: exchange counts via all_sum --- | ||
| # Each rank contributes its count in its own slot of a [world_size] | ||
| # vector; other slots are zero. all_sum merges them. | ||
| n_tasks = len(tasks) | ||
| logger.debug(f"mx_all_gather_tasks[{label}]: gathering counts (n_tasks={n_tasks})") | ||
| counts_vec = [0] * world_size | ||
| counts_vec[rank] = n_tasks | ||
| all_counts = cast( | ||
| list[int], | ||
| mx.distributed.all_gather(mx.array([n_tasks]), group=group).tolist(), | ||
| mx.distributed.all_sum( | ||
| mx.array(counts_vec, dtype=mx.int32), | ||
| group=group, | ||
| stream=cpu_stream, | ||
| ).tolist(), | ||
| ) |
There was a problem hiding this comment.
mx_all_gather_tasks() is called from agree_on_tasks/agree_on_cancellations even when group is None (single-node / non-distributed runner). In that case this implementation still calls mx.distributed.all_sum(..., group=group), which will likely fail when MLX distributed hasn’t been initialized. Consider adding an early fast-path like if group is None: return tasks, [] (or otherwise avoid collectives) so single-node runs don’t attempt distributed ops.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d612d14775
ℹ️ 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".
| "SKULK_KV_CACHE_BACKEND", | ||
| os.environ.get("EXO_KV_CACHE_BACKEND", ""), | ||
| ) | ||
| if launch_backend: |
There was a problem hiding this comment.
_user_set_kv_backend is true whenever SKULK_KV_CACHE_BACKEND exists, including when it is intentionally set to an empty string, but this branch only runs the sync when launch_backend is truthy. That means an explicit blank override (used in this codebase to force default behavior) is silently skipped, so skulk.yaml keeps the old backend and can be rebroadcast/reloaded later, undoing the operator’s launch-time override on restart. Treat key presence as the signal and normalize blank to default before syncing.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a50a53785e
ℹ️ 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".
| self._key_indices = seq[0] | ||
| self._key_norms = seq[1] | ||
| self._value_indices = seq[2] | ||
| self._value_norms = seq[3] | ||
| self.offset = int(self._key_indices.shape[2]) | ||
| self._pending_active = False |
There was a problem hiding this comment.
When state is restored from a 4-leaf quantized snapshot, this branch switches _pending_active to False but never clears _pending_keys/_pending_values. If the cache object was previously in deferred mode, those fp16 pending tensors stay resident and continue to contribute to nbytes, which can double memory footprint and trigger avoidable cache eviction/OOM under long-context loads. The setter should drop deferred buffers when entering the quantized-state path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 33 out of 35 changed files in this pull request and generated 4 comments.
Files not reviewed (1)
- dashboard-react/package-lock.json: Language not supported
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| uuid_byte_length = 36 | ||
| cpu_stream = mx.default_stream(mx.Device(mx.cpu)) | ||
| world_size: int = 1 if group is None else group.size() | ||
| rank: int = 0 if group is None else group.rank() | ||
|
|
||
| # --- Phase 1: exchange counts via all_sum --- | ||
| # Each rank contributes its count in its own slot of a [world_size] | ||
| # vector; other slots are zero. all_sum merges them. | ||
| n_tasks = len(tasks) | ||
| logger.debug(f"mx_all_gather_tasks[{label}]: gathering counts (n_tasks={n_tasks})") | ||
| counts_vec = [0] * world_size | ||
| counts_vec[rank] = n_tasks | ||
| all_counts = cast( | ||
| list[int], | ||
| mx.distributed.all_gather(mx.array([n_tasks]), group=group).tolist(), | ||
| mx.distributed.all_sum( | ||
| mx.array(counts_vec, dtype=mx.int32), | ||
| group=group, | ||
| stream=cpu_stream, | ||
| ).tolist(), | ||
| ) | ||
| logger.debug(f"mx_all_gather_tasks[{label}]: counts gathered: {all_counts}") | ||
| max_tasks = max(all_counts) | ||
| world_size: int = 1 if group is None else group.size() | ||
|
|
||
| if max_tasks == 0: | ||
| return [], [] | ||
|
|
||
| padded = [encode_task_id(task.task_id) for task in tasks] + [ | ||
| # --- Phase 2: exchange task IDs via all_sum --- | ||
| # Build a [world_size, max_tasks, uuid_byte_length] tensor. This rank | ||
| # fills its own slice; all other slices are zero. all_sum merges them | ||
| # without overlap (each rank owns a unique slice). | ||
| padded_ids = [encode_task_id(task.task_id) for task in tasks] + [ | ||
| [0] * uuid_byte_length | ||
| ] * (max_tasks - n_tasks) | ||
| assert all(len(eid) == uuid_byte_length for eid in padded_ids) | ||
|
|
||
| assert all(len(encoded_task_id) == uuid_byte_length for encoded_task_id in padded) | ||
| full = [[[0] * uuid_byte_length] * max_tasks for _ in range(world_size)] | ||
| full[rank] = padded_ids |
There was a problem hiding this comment.
mx_all_gather_tasks hard-codes uuid_byte_length = 36 and asserts every encoded task_id matches that length. TaskId is not guaranteed to be a 36-char UUID (e.g. there are many non-UUID TaskIds in the repo like "shutdown" / "warmup" in src/exo/worker/tests/constants.py), so this will raise (or corrupt IDs) as soon as such a task participates in the collective. Consider gathering the maximum UTF-8 byte length across ranks (via the same slotted all_sum pattern) and padding/trimming per-rank IDs accordingly (and stripping trailing 0s when decoding), or explicitly enforcing UUID TaskIds at the type boundary if that’s the intended invariant.
| logger.info( | ||
| f"agree_on_tasks: entering with {len(self._maybe_queue)} pending tasks" | ||
| ) | ||
| agreed, different = mx_all_gather_tasks( | ||
| self._maybe_queue, self.group, label="task-queue" | ||
| ) | ||
| logger.info(f"agree_on_tasks: {len(agreed)} agreed, {len(different)} different") |
There was a problem hiding this comment.
agree_on_tasks now logs at INFO level every time it enters / completes a task agreement. This method is invoked from distributed_prompt_progress_callback and on_generation_token, so under load it can emit a large volume of logs and drown out more actionable messages. Consider downgrading these to DEBUG (or sampling/rate-limiting) to avoid log spam in normal operation.
| logger.info( | |
| f"agree_on_tasks: entering with {len(self._maybe_queue)} pending tasks" | |
| ) | |
| agreed, different = mx_all_gather_tasks( | |
| self._maybe_queue, self.group, label="task-queue" | |
| ) | |
| logger.info(f"agree_on_tasks: {len(agreed)} agreed, {len(different)} different") | |
| logger.debug( | |
| f"agree_on_tasks: entering with {len(self._maybe_queue)} pending tasks" | |
| ) | |
| agreed, different = mx_all_gather_tasks( | |
| self._maybe_queue, self.group, label="task-queue" | |
| ) | |
| logger.debug( | |
| f"agree_on_tasks: {len(agreed)} agreed, {len(different)} different" | |
| ) |
| logger.info( | ||
| f"agree_on_tasks: entering with {len(self._maybe_queue)} pending tasks" | ||
| ) | ||
| agreed, different = mx_all_gather_tasks( | ||
| self._maybe_queue, self.group, label="task-queue" | ||
| ) | ||
| logger.info(f"agree_on_tasks: {len(agreed)} agreed, {len(different)} different") |
There was a problem hiding this comment.
Same as above for BatchGenerator: INFO logs inside agree_on_tasks can be very noisy because task agreement is called from prefill/generation progress callbacks. Consider using DEBUG (and keep INFO for exceptional cases) to avoid high-volume logging in production.
| logger.info( | |
| f"agree_on_tasks: entering with {len(self._maybe_queue)} pending tasks" | |
| ) | |
| agreed, different = mx_all_gather_tasks( | |
| self._maybe_queue, self.group, label="task-queue" | |
| ) | |
| logger.info(f"agree_on_tasks: {len(agreed)} agreed, {len(different)} different") | |
| logger.debug( | |
| f"agree_on_tasks: entering with {len(self._maybe_queue)} pending tasks" | |
| ) | |
| agreed, different = mx_all_gather_tasks( | |
| self._maybe_queue, self.group, label="task-queue" | |
| ) | |
| if different: | |
| logger.info( | |
| f"agree_on_tasks: {len(agreed)} agreed, {len(different)} different" | |
| ) | |
| else: | |
| logger.debug( | |
| f"agree_on_tasks: {len(agreed)} agreed, {len(different)} different" | |
| ) |
| _kv_cache_backend_value = preferred_env_value( | ||
| "SKULK_KV_CACHE_BACKEND", | ||
| "EXO_KV_CACHE_BACKEND", | ||
| DEFAULT_KV_CACHE_BACKEND, | ||
| ) | ||
| KV_CACHE_BACKEND: KVCacheBackend = cast( | ||
| KVCacheBackend, | ||
| _kv_cache_backend_value if _kv_cache_backend_value else DEFAULT_KV_CACHE_BACKEND, | ||
| ) | ||
| KV_CACHE_BACKEND: KVCacheBackend = resolve_kv_cache_backend(_kv_cache_backend_value) | ||
| TURBOQUANT_K_BITS: int | None = ( |
There was a problem hiding this comment.
KV_CACHE_BACKEND is still computed once at import time from env. Elsewhere in this PR, get_kv_cache_backend() was changed to read env at call time so dashboard/config-sync updates take effect in-process; any code that continues to use this constant (e.g. in generator/generate.py) can observe a stale backend and make inconsistent decisions. Consider deprecating/removing KV_CACHE_BACKEND in favor of calling resolve_kv_cache_backend(preferred_env_value(...)) at use sites, or making KV_CACHE_BACKEND a thin function/property instead of a frozen module constant.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 856b30f7a8
ℹ️ 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".
| # pipeline mode, the MLX BatchGenerator path currently produces | ||
| # degenerate token repetition after a valid prefill, while the | ||
| # sequential path is stable with the same prompt and default KV cache. | ||
| force_sequential_for_gemma4 = _is_gemma4_model(self.model_id, self.model_card) |
There was a problem hiding this comment.
The new Gemma 4 guard is applied unconditionally, so any Gemma 4 model is forced onto SequentialGenerator even on single-node/non-distributed deployments. The surrounding comment and warning message both scope this workaround to distributed pipeline behavior, so this broad condition introduces an avoidable throughput regression for local Gemma 4 inference that previously could use BatchGenerator.
Useful? React with 👍 / 👎.
Addresses 9 sev-4 review threads on PR #133. Each flagged a fact that drifted from the actual code: 1. **DOWNLOAD_COMMANDS missing from topic inventory.** Added to the Router topic list in architecture.md and to the Pubsub topics table in architecture-reference.md. The topic carries SyncConfig and model-store coordination commands; it's separate from COMMANDS for payload-size and retry-semantics reasons. 2. **TaskFinished mis-classified as event.** Per src/exo/shared/types/ commands.py:66 it's a Command, not an Event. Replaced in event-list examples in architecture.md and the events table in architecture- reference.md. Natural completion is represented by TaskStatusUpdated(Complete) which is driven by the TaskFinished command. 3. **CancelTask mis-classified as command.** Per src/exo/shared/types/ tasks.py:69 it's a Task. Removed from the Commands table and added a "Tasks (not commands)" note explaining the distinction; tasks are work units delivered to runners over mp.Queue, commands are imperative requests to the master. 4. **SyncConfig in wrong topic table.** It's a DownloadCommand carried on DOWNLOAD_COMMANDS, not a Command on COMMANDS. Split the "Commands" section in architecture-reference.md into COMMANDS-topic and DOWNLOAD_COMMANDS-topic subsections so the routing surface is accurate. 5. **Pubsub topic payload types were inner types, not network types.** The reference doc is supposed to be file:line accurate; corrected the payload column to show the wire types (GlobalForwarderEvent, LocalForwarderEvent, ForwarderCommand, ForwarderDownloadCommand, StateSyncMessage) with the inner payload broken out separately. 6. **/place_instance documented as alias for /instance.** They're different routes with different param shapes (PlaceInstanceParams vs CreateInstanceParams) bound to different handlers (place_instance vs create_instance). Replaced the alias claim with a note clarifying the distinction. 7. **/admin/restart endpoint shape wrong.** Documented as /admin/restart/node/{node_id}; actual route is /admin/restart with optional node_id query param. Corrected in both architecture.md note and architecture-reference.md table. 8. **rotorquant listed as a shipped KV cache backend.** RotorQuant is research code in PR #103 (unmerged); not in src/exo/worker/engines/mlx/constants.py. Removed from the backend table and added a parenthetical pointing to the PR for current status. Sev 1 thread (heading typo "Tools / store / store / admin") and a sev 1 thread claiming the markdown tables render with extra empty columns (false positive — tables use standard `| col | col |` syntax) are ignored per the strict review-loop rule. Validation: - npm run build (Docusaurus) — both pages render - All file:line refs in commit message resolve to current HEAD - Branch verified docs/architecture-foundation before commit Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Summary
rotorquantandrotorquant_adaptiveKV cache backends — pure-MLX port of IsoQuant 3-bit (block-diagonal quaternion rotations + Lloyd-Max centroids) fromscrya-com/rotorquantandjohndpope/llama-cpp-turboquant, both MIT#ifdef GGML_USE_CUDAinllama-context.cpp:1690); this is the first MLX/Metal implementationn_kv_heads != n_headsEverything Different About Skulkliving section inREADME.mdandwebsite/docs/everything-different.mdcapturing every divergence from upstream exo in one canonical placeTracks #102.
Why
OptiQ is currently the only KV cache backend in Skulk that realizes the inference-time benefit of rotation-based quantization (centroid-space storage + rotated-space SDPA). The two native TurboQuant backends quantize and immediately dequantize, paying the rotation cost on every step and gaining only storage savings. RotorQuant is the next generation of this family with two distinct improvements:
O(d)instead ofO(d²)per token, with 44× fewer rotation parametersThe upstream published numbers (5.3× prefill, 1.28× decode, PPL 6.91 vs TurboQuant) come from the llama.cpp fork on RTX 5090. Most of the prefill speedup comes from deferred prefill, which Metal has not had until now.
What's in
New module:
src/exo/worker/engines/mlx/rotorquant/tables.py— 32 hardcoded unit quaternions and the 8-entry 3-bit Lloyd-Max centroid table, lifted verbatim fromggml-iso-quant.cso the math agrees with the C referencerotation.py— Hamilton-product forward (q_L * v) and inverse (conj(q_L) * v) in pure MLX, reshape-to-quads + per-block broadcastquantizer.py—IsoQuantizerwith the norm-correction trick fromggml-planar-quant.c:111-114(stored norm =||x|| / sqrt(Σ centroid²)). Critical and undocumented in the upstream README — it absorbs the L2 shrinkage induced by centroid quantization so SDPA gets unbiased magnitudescache.py—RotorQuantKVCachewith deferred-prefill state machine + factory functions for the standard and adaptive variantstests/test_iso_quantizer.py— rotation orthogonality, R⁻¹R round-trip, index-range, 3-bit reconstruction MSE, norm-correction unbiasedness, shape contracts (7 tests)tests/test_cache.py— deferred prefill keeps storage unallocated during prefill, flushes on first decode, returns lossless fp16 during prefill, non-deferred path quantizes immediately, decode-after-flush appends correctly, trim works in both phases (7 tests)Wired in
constants.py—KVCacheBackendliteral extended;ROTORQUANT_FP16_LAYERSandROTORQUANT_DEFER_PREFILLenv varscache.py:make_kv_cache— two new branches before the final fallback, mirroring the optiq wiring withmake_cachetemplate handling and SSM/RotatingKVCache passthroughrunner.py— added to theforce_sequentiallist (same constraint as the other quantized backends)store/config.py— extendedInferenceConfig.kv_cache_backendliteral so the cluster-wide config endpoint accepts the new valuesSettingsPanel.tsx— two new dropdown entries (Adaptive marked recommended), tooltip rewrittenDocs
website/docs/kv-cache-backends.md— RotorQuant section with deferred-prefill explanation, env vars, expectations tablewebsite/docs/everything-different.md— new living page bulleting every Skulk-vs-exo divergence, organized into 9 categories. Includes a footer with rules for keeping it accuratewebsite/sidebars.ts+intro.md— wired the new page into the docs navREADME.md— new "Everything Different About Skulk" section mirroring the docs page (single source of truth, same categories), plus updates to "What Skulk Is Good At", "Core Features", and the env var tableWhat's not in (deferred to follow-ups)
These are explicitly out of scope per the design plan and tracked for v2:
block_iso3_0storage (50 bytes per 128 elements). Current uint8 storage gives 2× compression vs fp16; packed would give 5×. Purely additive, no algorithmic riskupdate_and_fetchand uses standard SDPA, which matches the llama.cpp FA kernel pattern and is portablemx.fast.metal_kernel.mx.fast.scaled_dot_product_attentionalready saturates memory bandwidth and the per-step rotation math is tiny relative to the SDPA call — a kernel would be premature optimization. Revisit only if profiling shows a hot spotTest plan
uv run pytest src/exo/worker/engines/mlx/rotorquant/tests/— 14/14 passuv run pytest src/exo/worker/tests/unittests/test_mlx/— 47/47 pass (no regressions in turboquant or kv prefix cache)uv run pytest— 462/463 pass (1 pre-existing rust binding failure on main, unrelated)uv run basedpyright src/exo/worker/engines/mlx/rotorquant src/exo/worker/engines/mlx/cache.py src/exo/worker/engines/mlx/constants.py src/exo/store/config.py src/exo/worker/runner/llm_inference/runner.py— 0 new errors (2 pre-existing incache.pyuntouched)uv run ruff checkon edited files — cleancd dashboard-react && npm run build— cleancd website && npm run build— clean (the new "Everything Different About Skulk" page renders)SKULK_KV_CACHE_BACKEND=rotorquant_adaptive uv run skulk, load a small Llama 3.2 1B or Qwen 1.5B (deliberately a GQA model to exercise the GQA-native path), confirm chat output is coherent and logs showUsing rotorquant adaptive KV cache. Pending review.