perf(lfm2): quantized flat-default (~1.84×) + compiled flat/paged decode (~2.34×) + prefill-tps telemetry fix - #67
Conversation
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
… fix paged prefill-tps telemetry Quantized lfm2_moe checkpoints carry `.scales` tensors, so they can never register with the compiled-PAGED C++ path; with `use_block_paged_cache` defaulting ON they fell onto the eager-PAGED loop (~12 `synchronize_mlx()`/token, blocking `y.eval()`, no async double-buffering) — ~1.84× slower than the eager-FLAT path (in-graph KVCache + `async_eval_arrays`) on the measured mxfp8 LFM2.5-8B-A1B (74 → 131 tok/s on M5 Max, closing most of the gap vs oMLX's 187). Fix #1: resolve the `use_block_paged_cache` default in `Lfm2Inner::load_from_dir`, keyed on the authoritative `.scales` tensor signal (the SAME signal that gates compiled registration, so the two can't diverge for a checkpoint whose quantization block lacks top-level bits/mode) — quantized + unset → flat; bf16 → paged (PR #66 compiled-PAGED unaffected); an explicit config.json value always wins. The server auto-routes via `has_block_paged_cache()`, so batched / cross-request reuse is unaffected (bf16 stays paged; quantized moves to the warm-slot path the server already supports). Single-stream only. Fix #4: the paged prefill-throughput telemetry divided ttft by the attention SUFFIX (`tokens.len() - cached_prefix_len`), but paged prefill reprocesses the FULL prompt through the conv layers, so ttft is full-prompt scale — this under-reported prefill tok/s by the cache-hit ratio (~37 vs thousands on warm content-addressed prefix reuse). Use the full prompt count at both paged call sites. Qwen3.5's identical expression is correct (no conv reprocessing) and is left untouched. Deferred (plan only): quantized compiled-PAGED decode (#2, for batched serving) and a cross-family paged-loop wired_limit (#3a, no-op for the flat default on a 128GB host). Tests: pure-helper policy unit test; a parse_config-does-not-resolve contract test; a compute_performance_metrics divider test; and two #[ignore] real-weights regressions (quantized default load → flat via has_block_paged_cache(); paged warm-reuse prefill-tps is full-prompt scale). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ager-paged)
Enable the compiled C++ decode path for quantized lfm2 / lfm2_moe
(LFM2.5-8B-A1B class), previously bf16-only. A per-projection quant-info
registry (`mlx_store_quant_info`, keyed on each `.scales` prefix) makes the
C++ (mode,bits,group_size) dispatch authoritative instead of relying on the
companion-tensor heuristic (which mislabels mxfp4/nvfp4 as mxfp8); the
heuristic is retained only as a fallback.
Gating:
- model.rs: `activations_bf16 = is_quantized ? (quant_compiled_enabled() &&
non_quant_floats_bf16) : all_float_weights_bf16` guards the paged C++ pre-path.
- persistence.rs: `should_register_compiled(...)` — quantized registers when
eligible (flat, or paged with non_quant_floats_bf16 + block-size-ok); a packed
embedding (`embed_tokens.scales`) bars compiled (the C++ does a dense `take`).
- `MLX_LFM2_DISABLE_QUANT_COMPILED` env escape hatch.
Correctness: byte-identical to pure-Rust eager across {mxfp8, 4-bit affine} x
{flat, paged}, proven via the model-id eviction oracle (the evicted eager-ref
runs the independent QuantizedLinear/QuantizedSwitchLinear modules, so a C++
mislabel diverges early). Perf: compiled-paged ~2.34x over eager-paged; flat
compiled ~= flat eager (no regression).
Review: codex adversarial-review runtime hung twice mid quant-dispatch
cross-reference; a Claude-subagent adversarial review cleared it SHIP with no
blocking bug (the truncated codex "No-ship" reconstructed as a false alarm on
packed-embedding / unsupported-mode / flat-bf16-invariant). Deferred follow-ups:
a synthetic non-gated quantized parity test (parity is currently operator-verified
via LFM2_COMPILED_E2E=1), and a pre-existing 2D-quant-weight transpose waste.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
4676007 to
ee0a61f
Compare
Repo-wide comment-only cleanup across 56 files. Removes development-arc scaffolding references (Phase P1/P2, Phase 4b/B3.1, W6.x work-item tags, E53, PR-#/session/Codex-round narration) and compresses over-verbose block comments, while preserving genuine technical content: ABI/locking contracts, dtype/f32-promotion gotchas, safety invariants, kernel-parity reasoning, and load-bearing vLLM/mlx-lm cross-references. Comments only — no behavior change. Verified that every file is byte-identical after comment stripping; no string/log/error literals or code identifiers were touched. index.d.cts regenerated from the cleaned NAPI doc-comments (JSDoc text only; all signatures unchanged). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Force flat clones to disable paged cache
For bf16 LFM2 checkpoints where config.json omits use_block_paged_cache, the loader defaults to paged (None becomes true in Lfm2Inner::new). Since clone_model_dir(..., false) leaves the copied config untouched here, the flat_dir used by the parity tests can load the paged path too, making the flat-vs-paged comparisons pass trivially and miss regressions in the flat path. Set use_block_paged_cache to the requested boolean in both branches, as the compiled E2E helper does.
ℹ️ 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".
Addresses Codex review on PR #67. `clone_model_dir(..., false)` only mutated config.json inside the `if use_block_paged` branch, so the flat clone of a bf16 checkpoint that omits `use_block_paged_cache` left the key absent -> `Lfm2Inner::new`'s `unwrap_or(true)` silently loaded the PAGED path. The flat-vs-paged parity tests then compared paged-vs-paged and could not catch flat-path regressions. Write `use_block_paged_cache` unconditionally in both branches (mirroring `lfm2_compiled_e2e.rs`); keep the paged-only pool sizing inside the `if`. The flat clone now persists `Some(false)` and loads the genuine flat path. Test-only change; compiles clean (cargo build --tests -p mlx-core). Tests are #[ignore]'d and env-gated on MLX_TEST_MODEL_PATH, so this restores their fidelity for local runs without affecting CI. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Addressed in a4a760d — confirmed real (deep control-flow trace below) and fixed. Why the finding holds
So FixWrite |
Summary
Closes most of the LFM2.5-8B-A1B quantized single-stream decode gap to oMLX by fixing the default decode path and extending the compiled C++ path to quantized weights. Two commits:
fb234575— quantized → flat default (~1.84×) + paged prefill-tps telemetry fix46760077— quantized compiled flat+paged decode (~2.34× paged over eager-paged)#1 — quantized single-stream defaults to FLAT decode (~1.84×)
Quantized
lfm2/lfm2_moewas silently defaulting to the eager-PAGED loop (~12synchronize_mlx()/token + blockingy.eval()+ no async double-buffering), ~1.84× slower than FLAT on the measured mxfp8 8B-A1B (74 → 131 tok/s, M5 Max). The default is now keyed on the authoritative.scalestensor signal: quantized → FLAT, bf16 → PAGED (unchanged). Explicituse_block_paged_cacheinconfig.jsonalways wins.#4 — paged prefill-tps telemetry fix
The paged path reported a bogus ~37
prefillTokensPerSecond(it divided full-prompt ttft by the attention suffix count on warm prefix-cache hits). Now uses the full-prompt count as the numerator; guarded bylfm2_paged_prefill_tps_is_full_prompt_scale_on_warm_reuse.#2 — quantized compiled flat+paged decode (~2.34× over eager-paged)
Extends the compiled C++ decode path (previously bf16-only) to quantized
lfm2/lfm2_moe. A per-projection quant-info registry (mlx_store_quant_info, keyed on each.scalesprefix) makes the C++(mode, bits, group_size)dispatch authoritative instead of the companion-tensor heuristic (which mislabels mxfp4/nvfp4 as mxfp8); the heuristic is retained only as a fallback. Compiled-PAGED is ~2.34× over eager-PAGED, rescuing the pinned-paged quant path (e.g. server/batched). A packed embedding (embed_tokens.scales) bars the compiled path (C++ does a densetake). Env escape hatch:MLX_LFM2_DISABLE_QUANT_COMPILED.Correctness
Byte-identical to the pure-Rust eager path across {mxfp8, 4-bit affine} × {flat, paged}, proven via the model-id eviction oracle in
lfm2_compiled_e2e.rs(quant_compiled_vs_eager_parity): loading the compiled model evicts the eager-ref's process-global weights, so the eager-ref runs the independentQuantizedLinear/QuantizedSwitchLinearmodules — a C++ dispatch mislabel would diverge early. This is stronger than a same-graphMLX_NO_COMPILEreference.Perf context
This is the quantized path — the relevant one for oMLX's 8-bit headline. Separately verified this session: for bf16, our decode (~110 tok/s) is at exact op-for-op parity with mlx-lm and is memory-bandwidth-bound (MoE gather already saturates ~404 GB/s at the k=4 decode shape, ~80% of the M5 Max ceiling); the residual bf16 gap to oMLX is host/measurement, not software. The real lever for absolute decode speed is reducing bytes-per-token (quantization) — which is exactly what these changes make fast.
Test plan
cargo clippy --all-targets -- -D warnings— cleancargo fmt --check— cleancargo test -p mlx-core, incl. the compiled-registration gate tests)LFM2_COMPILED_E2E=1+LFM2_QUANT_MODEL_PATH)yarn build:nativeclean; noindex.d.ctsdriftReview status
The mandated
codex:adversarial-reviewruntime hung twice mid quant-dispatch cross-reference (a codex-runtime issue, not a code signal). A thorough Claude-subagent adversarial review cleared it SHIP / no blocking bug — verifying dispatch parity for every projection class (MoE experts, router gate, dense-MLP, attention q/k/v/out, conv, untied lm_head) and ruling out the truncated codex concern on all three plausible completions (packed-embedding guard, registry-authoritative quant modes, pre-existing flat bf16 invariant).Deferred follow-ups (non-blocking):
LFM2_COMPILED_E2E=1; the synthetic harness only generates bf16 weights, and the completenessdebug_assert_eq!is compiled out in release).mlx_store_weighttransposes packed 2D quant.weightintog_weight_transposesthat's never read (pre-existing waste, surfaced not introduced).🤖 Generated with Claude Code
Note
Medium Risk
Changes default decode routing and global compiled weight registration for quantized LFM2, where incorrect quant dispatch or gating would affect correctness and performance; mitigated by expanded unit tests and documented escape hatches.
Overview
LFM2 load and decode routing now treat quantized checkpoints differently: when
use_block_paged_cacheis unset, presence of.scalestensors defaults to flat decode (instead of paged), with resolution moved fromparse_configtoload_from_dirso it matches the registration gate. Explicitconfig.jsonvalues still win.Quantized models can use the compiled C++ path (flat and paged): registration publishes per-projection quant info via
mlx_store_quant_info,should_register_compiledandpaged_compiled_decode_setupusenon_quant_floats_bf16plusMLX_LFM2_DISABLE_QUANT_COMPILED, and packedembed_tokensblocks compiled registration because the C++ path does a dense embedding lookup.Paged chat performance metrics use the full prompt token count for prefill throughput (conv layers re-run the full prompt), fixing inflated TTFT/prefill-tps on warm prefix-cache hits.
Most other diff hunks are comment and docstring cleanup (phase/W6/PR ticket references removed); behavior in convert, MTP, Qwen3, and banded-attention modules is unchanged aside from wording.
Reviewed by Cursor Bugbot for commit a4a760d. Bugbot is set up for automated code reviews on this repo. Configure here.