Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ Epic #1431 pinned llama.cpp release b10621 (commit `c1d0e7a0`, published 2026-08

### Fixed

- **A malformed `MLXCEL_PAGED_SLAB_BLOCKS` now falls back to the derived slab size, as its warning always claimed** (#1137). `resolve_paged_slab_blocks` returned early on a parse error, which is the `0` pin: the pool kept the historical 32-block slab while the log said the derived size was in use, so a typo in the variable quietly made the fused paged decode path unreachable past one slab. The parse-error arm now falls through to the derivation, `0` still pins the pool default, and a unit test keeps the three arms distinct.
- **Speculative-burst responses report `prompt_ms` and `predicted_ms` the way classic decode does** (#1592). The B=1 and batched bursts run every verify round before streaming, and the first-token stamp was taken when the finished token vector was replayed, so `prompt_ms` absorbed drafter load, prefill and all verify rounds while `predicted_ms` reported the sub-millisecond replay (`predicted_per_second` read 95000 on a 96-token DFlash request). The burst now carries the instant its target prefill finished and stamps the sequence with it, so `predicted_ms` covers the verify rounds; the drafter's once-per-process disk load is excluded from every `timings` field and stays on the `Drafter loaded` log line. Classic decode is unchanged.
- **Youtu-VL restores raster order after the windowed vision tower** (#1600). `reverse_window_indices` sorted `(value, index)` pairs and wrote the rank back by index, which for a permutation of `0..N` reproduces `window_index` itself rather than its inverse, so the merged vision tokens were un-reordered by applying the window permutation a second time. That only restores raster order when the permutation is an involution, and with Youtu-VL's 8x8 merged-token window even a 512x512 image (a 16x16 merged grid) is not one: 192 of its 256 merged tokens reached the language model out of place. The helper now builds `argsort(window_index)` directly, bounds-checked like the Qwen2.5-VL fix in #1601, with checkpoint-free tests pinning the non-involution, single-window, and multi-image cases. This corrects the token order at that stage only, and does not on its own make the family describe images correctly: the processor emits patches in raster order where the tower's rotary positions and the patch merger both expect merge-block-major, which is filed separately as #1610.
- **`mlxcel_core::is_gpu_available()` is renamed `default_device_is_gpu()`, and `gpu_backend_available()` answers whether a GPU backend exists** (#1421). The old name read as a hardware query but reported whether the MLX default device was currently the GPU, so it turned `false` the moment anything called `set_default_device(false)` on a machine with a GPU. `gpu_backend_available()` is the unclamped `device_count(Device::gpu) > 0`: `true` on Metal and on a CUDA build with a device, `false` on a CPU-only build and on a CUDA build without a driver, and unchanged by the default device. `is_gpu_available()` remains as a deprecated shim for one release. `initialize_runtime()` now resolves `RuntimeSetup.device` from backend availability and reports `MLXCEL_DEVICE=cpu` separately in the new `RuntimeSetup.cpu_override` field, which the server and CLI startup lines print, so a CPU device is distinguishable as requested or as the only option. Every CPU resolution is applied to MLX's default device, not only the operator's override: the CUDA backend answers `gpu::is_available()` unconditionally while `device_count(Device::gpu)` reports what the driver found, so a CUDA build on a driverless host would otherwise report `Cpu` while MLX kept dispatching to the unusable GPU. Test modules that moved the default device to the CPU inside a `Once` and never moved it back, which made every real-checkpoint gate sorting after them measure the CPU backend under `--test-threads=1`, now hold the new `mlxcel_core::streams::DefaultDeviceGuard` per test, and the unconditional GPU pin PR #1420 added to `mlx_test_guard` is reduced to an assertion so a future leak fails at the first gate after it instead of being silently repaired. Restoring the device is not the same as excluding other movers, so there is also one process-wide lock for it, `mlxcel_core::streams::lock_default_device`, held by every test that moves the default device and by `mlx_test_guard` for the duration of the tests it serializes: a lock private to each module leaves the modules interleaving, which under the parallel `cargo test --lib` in `scripts/run_quality_gate.sh` reported another module's live guard as a leak in 4 of 300 runs and would have let a real-checkpoint gate score on the CPU backend without saying so.
Expand Down
2 changes: 1 addition & 1 deletion docs/environment-variables.md
Original file line number Diff line number Diff line change
Expand Up @@ -397,7 +397,7 @@ declines to serve, see

| Variable | Values | Default | Notes |
|----------|--------|---------|-------|
| `MLXCEL_PAGED_SLAB_BLOCKS` | non-negative integer blocks | derived (`ceil(per_slot_ctx / block_size) * batch`, floored at the 32-block pool default and capped at the per-layer share of the KV block budget) | Paged pool slab size in blocks. The fused kernels read one contiguous buffer per side, so a layer whose rows outgrow its first slab silently returns to the gather path; the server derives the slab from `--ctx-size` and `--parallel` and logs the resolved value at startup (`Paged KV slab size: N blocks per layer`). Raise it when serving contexts longer than `--ctx-size` implies. `0` pins the historical 32-block pool default, which keeps pre-#899 allocation behaviour and, as a side effect, keeps the fused path unreachable for anything past one slab. A value that is not a non-negative integer warns and leaves that same 32-block default in place. |
| `MLXCEL_PAGED_SLAB_BLOCKS` | positive integer blocks, used verbatim (neither floored nor capped); `0` pins the pool default | derived (`ceil(per_slot_ctx / block_size) * batch`, floored at the 32-block pool default and capped at the per-layer share of the KV block budget) | Paged pool slab size in blocks. The fused kernels read one contiguous buffer per side, so a layer whose rows outgrow its first slab silently returns to the gather path; the server derives the slab from `--ctx-size` and `--parallel` and logs the resolved value at startup (`Paged KV slab size: N blocks per layer`). Raise it when serving contexts longer than `--ctx-size` implies. `0` pins the historical 32-block pool default, which keeps pre-#899 allocation behaviour and, as a side effect, keeps the fused path unreachable for anything past one slab. A value that is not a non-negative integer warns and is ignored, so the derived size applies exactly as if the variable were unset (#1137). |
| `MLXCEL_PAGED_V2_MIN_KV_TOKENS` | non-negative integer tokens | `4096` | Visible KV tokens a **lone-request** launch must reach before v2 is dispatched. 4096 is the weakest measured single-request win (1.08x); batch 1 at 1024 tokens is the only measured loss (0.91x), which is what the floor excludes. An unparseable value falls back to the default, so a typo cannot silently install an unmeasured policy. |
| `MLXCEL_PAGED_V2_MIN_KV_TOKENS_PER_REQUEST` | non-negative integer tokens | `512` | Visible KV tokens **per request** a multi-request launch must reach. The smallest measured multi-request win is batch 4 at 1024 tokens per request (1.41x); the floor sits below that with margin so a prompt landing just under a nominal 1K is not declined for no measured reason. Same fallback-on-typo behaviour as the row above. To benchmark the declined regime, force the path on with `MLXCEL_PAGED_ATTENTION_NATIVE=1` rather than driving a floor to zero. |

Expand Down
62 changes: 50 additions & 12 deletions src/execution/memory_estimate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1295,7 +1295,9 @@ pub fn paged_v2_workspace_reserve_bytes(model_dir: &Path, num_layers: usize, bat
/// `0` pins the pool default ([`mlxcel_core::cache::POOL_SLAB_BLOCKS`]), which
/// keeps the pre-#899 allocation behaviour and, as a side effect, keeps the
/// fused decode path unreachable for anything past one slab. Any other positive
/// integer is used verbatim, which is how a benchmark sweeps the setting.
/// integer is used verbatim, neither floored nor budget-capped, which is how
/// a benchmark sweeps the setting. Anything else warns and is ignored, so the
/// slab is derived exactly as if the variable were unset (#1137).
pub const PAGED_SLAB_BLOCKS_ENV: &str = "MLXCEL_PAGED_SLAB_BLOCKS";

/// Resolve the paged pool's slab size in blocks (issue #899).
Expand Down Expand Up @@ -1353,17 +1355,16 @@ pub fn resolve_paged_slab_blocks(
return None;
}
if let Ok(raw) = std::env::var(PAGED_SLAB_BLOCKS_ENV) {
return match raw.trim().parse::<usize>() {
Ok(0) => None,
Ok(n) => Some(n),
Err(_) => {
tracing::warn!(
"{PAGED_SLAB_BLOCKS_ENV}={raw:?} is not a non-negative integer; \
using the derived slab size"
);
None
}
};
match raw.trim().parse::<usize>() {
Ok(0) => return None,
Ok(n) => return Some(n),
Err(_) => tracing::warn!(
env_var = PAGED_SLAB_BLOCKS_ENV,
value = raw,
"{PAGED_SLAB_BLOCKS_ENV} must be a non-negative integer; ignoring it and \
deriving the slab as if unset (0 pins the pool default)",
),
}
}
// The geometry probe doubles as the "is this model pool-eligible at all"
// check: without it the byte clamp below would be meaningless.
Expand Down Expand Up @@ -1516,6 +1517,13 @@ mod tests {
unsafe { std::env::set_var(key, value) };
Self { key, previous }
}

fn unset(key: &'static str) -> Self {
let previous = std::env::var(key).ok();
// SAFETY: as for `set`.
unsafe { std::env::remove_var(key) };
Self { key, previous }
}
}

impl Drop for EnvRestore {
Expand Down Expand Up @@ -1922,6 +1930,36 @@ mod tests {
assert_eq!(est.available_bytes, 512 * 1024 * 1024);
}

/// Issue #1137: the warning for a malformed `MLXCEL_PAGED_SLAB_BLOCKS`
/// says the derived slab size is used, and the code used to return the
/// `0` pin instead. The three arms must stay distinct: unset and
/// malformed derive, `0` pins the pool default (`None`), a number wins
/// verbatim.
#[test]
fn paged_slab_blocks_malformed_env_derives_like_unset() {
let _env = crate::test_support::env_lock::env_lock();
let tmp = tempfile::tempdir().unwrap();
write_minimal_config(tmp.path());
let resolve = || resolve_paged_slab_blocks(tmp.path(), 32, 32, 4, 8192, false, None);

let derived = {
let _unset = EnvRestore::unset(PAGED_SLAB_BLOCKS_ENV);
resolve()
};
assert_eq!(derived, Some(8192 / 32 * 4));

let _malformed = EnvRestore::set(PAGED_SLAB_BLOCKS_ENV, "lots");
assert_eq!(resolve(), derived);

let _zero = EnvRestore::set(PAGED_SLAB_BLOCKS_ENV, "0");
assert_eq!(resolve(), None);

// Trimmed and verbatim: not floored at the 32-block pool default the
// derived path applies.
let _explicit = EnvRestore::set(PAGED_SLAB_BLOCKS_ENV, " 8 ");
assert_eq!(resolve(), Some(8));
}

#[test]
fn parse_optional_memory_size_rejects_non_positive_and_non_finite() {
assert_eq!(parse_optional_memory_size_bytes("0"), None);
Expand Down