From 328e3fd20c75fba6b5f2d16ead88b59282b4f327 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 21 Sep 2026 20:21:37 +0900 Subject: [PATCH] perf(metal): raise the command-buffer input budget for decode only On M1 through M4 mlxcel raises `MLX_MAX_OPS_PER_BUFFER` to 1000 (#353) but leaves MLX's second commit trigger, `MLX_MAX_MB_PER_BUFFER`, at 40-50. That budget counts input elements (`data_size() >> 20`), not bytes, and decode reads the whole weight set every token, so a 4-bit 7B model commits about 23 command buffers per token and the GPU idles about 40 us at each boundary (xctrace: 89.5% busy). A budget of 1000 lifts decode on M1 Ultra by up to 21% (command-r7b +7%, Qwen3-30B-A3B +20%, Mixtral +21%, Llama 3.1 8B bf16 +17%, Gemma 3 4B flat), but applied globally it also keeps prefill activations alive per buffer and roughly doubles peak memory on a 2048-token prompt (Qwen3-30B-A3B 19.8 to 36.2 GB). It also hurts a synchronous decode step, which encodes one large buffer before the GPU can start. `Device` latches the budget once, so a `device.cpp` overlay adds a runtime override that `needs_commit` reads, and `DecodeCommandBufferBudget` raises it only around pipelined decode: the generate loops and the server's lookahead tick and prime. `MLXCEL_DECODE_MB_PER_BUFFER` tunes or disables it; an explicit `MLX_MAX_MB_PER_BUFFER` still pins both phases. command-r7b 4-bit on M1 Ultra, same binary with the switch on and off (`MLXCEL_DECODE_MB_PER_BUFFER=0`): CLI harness decode 104.2-104.4 to 110.1-111.2 tok/s, server lookahead decode median 110.0 to 118.5 (ten requests, two starts), synchronous server decode unchanged, prefill and pp2048 peak memory unchanged. Speculative loops are not wired yet. Measurements are in docs/benchmark_results/metal-mb-per-buffer-m1ultra-2026-09-21.md. --- .../metal-mb-per-buffer-m1ultra-2026-09-21.md | 120 ++ docs/environment-variables.md | 3 +- src/lib/mlx-cpp/CMakeLists.txt | 6 + .../patches/mlx/backend/metal/device.cpp | 1015 +++++++++++++++++ src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp | 23 + src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 8 + .../mlxcel-core/src/command_buffer_budget.rs | 88 ++ .../src/command_buffer_budget_tests.rs | 66 ++ src/lib/mlxcel-core/src/generate.rs | 14 + src/lib/mlxcel-core/src/hardware.rs | 170 +++ src/lib/mlxcel-core/src/lib.rs | 16 + src/main.rs | 5 + src/server/batch/scheduler/decode_tick.rs | 9 + 13 files changed, 1542 insertions(+), 1 deletion(-) create mode 100644 docs/benchmark_results/metal-mb-per-buffer-m1ultra-2026-09-21.md create mode 100644 src/lib/mlx-cpp/patches/mlx/backend/metal/device.cpp create mode 100644 src/lib/mlxcel-core/src/command_buffer_budget.rs create mode 100644 src/lib/mlxcel-core/src/command_buffer_budget_tests.rs diff --git a/docs/benchmark_results/metal-mb-per-buffer-m1ultra-2026-09-21.md b/docs/benchmark_results/metal-mb-per-buffer-m1ultra-2026-09-21.md new file mode 100644 index 000000000..407702442 --- /dev/null +++ b/docs/benchmark_results/metal-mb-per-buffer-m1ultra-2026-09-21.md @@ -0,0 +1,120 @@ +# Metal command-buffer input budget during decode, M1 Ultra, 2026-09-21 + +Starting point: command-r7b 4-bit decode on this host read 102-104 tok/s, and a Python MLX replica of the same decode graph (same primitives, same kernels, same MLX era) read 108.5. The 0.37 ms per token between them was independent of context length (16, 500 and 2048 tokens), so it was not attention. + +Headline: **mlxcel's decode on M1 through M4 has been splitting every token into about 23 Metal command buffers since `MLX_MAX_OPS_PER_BUFFER` was raised to 1000 (#353), because MLX's second commit trigger, the input budget `MLX_MAX_MB_PER_BUFFER`, was left at its default of 50.** Raising that budget to 1000 for decode moved decode by +3% to +21% on eight of the nine checkpoints measured (Gemma 3 4B was flat) and removed most of the gap to the replica. The same budget during prefill roughly doubles peak memory on long prompts, so it is applied around decode steps only, through a runtime override in the `mlx/backend/metal/device.cpp` overlay. + +## Environment + +| Field | Value | +|---|---| +| Host | Mac Studio, Apple M1 Ultra, 128 GB unified memory, macOS 27.0 (26A428) | +| Apple GPU generation | 13 (`d`, Ultra): MLX defaults are 50 ops and 50 "MB" per command buffer; mlxcel raises ops to 1000 | +| Base | `main` at `0bbfa95d`, MLX pin `81ba1c6a` | +| Build | `cargo build --release --features metal,accelerate` | +| Harness | `mlxcel-bench-decode` with the `scripts/bench_decode.sh` shape: `--prompt-tokens 500 -n 128 --ignore-eos --warmup-tokens 20` unless noted; cells interleaved, three runs per cell unless noted | +| Background | Photos library analysis running throughout (`mediaanalysisd` 60-140% CPU for two days, `com.apple.photos.ImageConversionService` 3.2-3.5% GPU time), load average about 4, Time Machine not running | + +The background load depresses both arms of every comparison below equally. By the final comparison it had subsided, and suspending the indexers outright did not move `main`'s absolute number (see the last section). + +## Mechanism + +MLX commits a Metal command buffer when either counter passes its cap (`CommandEncoder::needs_commit`, `mlx/backend/metal/device.cpp`): + +```cpp +return (buffer_ops_ > max_ops) || ((buffer_sizes_ >> 20) > max_mb); +``` + +`buffer_sizes_` sums `array::data_size()` over the distinct input arrays of the buffer. That is an element count, not bytes. A 4-bit 7B model stores about 1.1G packed `u32` elements, and decode reads all of them every token, so with the op cap at 1000 the input cap of 50 (52M elements) commits a buffer every one to two layers: about 23 per token on command-r7b. A bf16 checkpoint stores one element per weight rather than eight weights per packed word, so it commits more often still. + +A Metal System Trace attached to a running decode (`xctrace record --attach`, 2 s) showed the GPU busy 89.5% of the time, with an idle gap between consecutive command buffers of 42.8 us median (p90 52 us). The Python replica, traced the same way, showed the same gap structure (88.1% busy, 41.8 us median), which is why the replica also gains from a larger budget, only less (+2.5%, see below). + +## Budget sweep, command-r7b 4-bit + +`MLX_MAX_MB_PER_BUFFER` set in the environment, op cap at mlxcel's 1000. + +| Budget | Prefill tok/s | Decode tok/s | +|---|---|---| +| 50 (default) | 665-668 | 101.4-103.3 | +| 100 | 661-676 | 107.4-109.2 | +| 200 | 657-678 | 106.4-107.5 | +| 400 | 655-671 | 107.3-109.6 | +| 1000 | 656-675 | 109.4-110.3 | +| 4000 | 638-654 | 109.2-109.9 | +| 100000 | 650-654 | 108.7-109.6 | + +Decode saturates by 1000; prefill starts losing at 4000. The Python replica under the MLX wheel's own defaults (50 ops, 50 budget): 107.9-108.4 at 50, 110.7 at 1000. + +## Across families, default versus 1000 + +| Checkpoint | Decode, default | Decode, 1000 | Prefill, default | Prefill, 1000 | +|---|---|---|---|---| +| command-r7b 4-bit | 101.4-103.3 | 109.4-110.3 (+7%) | 665-668 | 656-675 | +| Llama 3.1 8B Instruct 4-bit | 99.0-99.2 | 103.4-105.1 (+5.6%) | 713-716 | 696-700 | +| Qwen2.5 7B Instruct 4-bit | 98.2-100.1 | 106.4-108.1 (+8%) | 754-762 | 722-745 | +| Gemma 3n E4B 4-bit | 59.7-60.6 | 62.1-62.4 (+3%) | 759-769 | 735-762 | +| Gemma 3 4B 4-bit | 95.6-97.3 | 96.3-96.9 (flat) | 917-935 | 913-953 | +| Granite 4.0 H Tiny 4-bit | 103.3-104.0 | 113.4-113.8 (+10%) | 1605-1612 | 1607-1620 | +| Qwen3-30B-A3B 4-bit | 72.9-76.0 | 90.0-90.4 (+20%) | 837-845 | 780-851 | +| Mixtral 8x7B Instruct 4-bit | 51.0-52.1 | 61.9-62.8 (+21%) | 323-326 | 320-327 | +| Llama 3.1 8B Instruct bf16 | 34.2-34.7 | 40.1-40.4 (+17%) | 769-777 | 751-755 | + +A second, interleaved pass on the two dense models that showed a prefill loss, at two prompt lengths: + +| Checkpoint | Prompt | Budget | Prefill | Decode | +|---|---|---|---|---| +| Llama 3.1 8B 4-bit | 500 | 50 / 400 / 1000 | 717-718 / 691-714 / 715-721 | 96.3-99.6 / 103.6-104.9 / 103.6-105.6 | +| Llama 3.1 8B 4-bit | 2048 | 50 / 400 / 1000 | 751-759 / 737-742 / 742-755 | 91.3-93.0 / 92.5-93.0 / 92.5-93.9 | +| Qwen2.5 7B 4-bit | 500 | 50 / 400 / 1000 | 747-762 / 749-760 / 732-747 | 98.6-100.6 / 106.6-108.0 / 106.6-108.3 | +| Qwen2.5 7B 4-bit | 2048 | 50 / 400 / 1000 | 799-808 / 790-797 / 790-794 | 96.7-96.9 / 99.0-100.5 / 100.3-101.5 | + +No decode cell regressed. Prefill at 1000 is flat to -1.4% on the 4-bit models and -2.7% on the bf16 one. + +## Why the budget is scoped to decode + +The input budget is also what bounds how long prefill activations stay alive: a buffer's intermediates are released when it completes. MLX peak memory at a 2048-token prompt, 32 generated tokens: + +| Checkpoint | 50 | 100 | 200 | 400 | 1000 | +|---|---|---|---|---|---| +| Qwen2.5 7B 4-bit | 6.01 GB | 6.90 | 8.93 | 11.09 | 12.77 | +| Qwen3-30B-A3B 4-bit | 19.75 GB | 20.74 | 23.36 | 28.78 | 36.19 | +| command-r7b 4-bit | 7.11 GB | | | | 13.94 | +| Llama 3.1 8B bf16 | 17.93 GB | | | | 24.24 | + +With one generated token the delta is the same (Qwen2.5: 6.01 to 12.59 GB at 2048, 5.07 to 6.56 GB at 512, +0.43 GB at a 64-token prompt), so it is prefill, and it grows with prompt length. No constant serves both phases, and `Device` latches the budget from the environment once. The overlay adds a process-wide override that `needs_commit` consults when it is non-zero; `DecodeCommandBufferBudget` raises it after a prompt has been encoded and restores it when the decode loop or decode step ends. + +## Result with the decode-only switch + +Peak memory at a 2048-token prompt with the switch: Qwen2.5 7B 4-bit 6.07 GB, Qwen3-30B-A3B 4-bit 19.75 GB, the same as the device default (6.01 and 19.75), against 12.77 and 36.19 with the budget applied to both phases. + +command-r7b 4-bit, same binary, `MLXCEL_DECODE_MB_PER_BUFFER` unset versus `0`, three interleaved runs: decode 110.1-111.2 versus 104.2-104.4, prefill 671-680 versus 674-676. + +Final interleaved comparison against `main` at `0bbfa95d`, eight ABBA pairs per arm, `--prompt-tokens 500 -n 128 --ignore-eos`. The measured branch also carried a cohere2 residual-add fusion that is submitted separately; on its own it is worth +1.1% (measured with the budget pinned on both arms), so about 6% of the gain below is this switch. The same-binary on/off comparison above isolates it. + +| Condition | `main` decode, median (range) | Branch decode, median (range) | Prefill, `main` / branch | +|---|---|---|---| +| As found (Photos analysis idle by then, load average 3.3) | 102.69 (100.33-103.02) | 109.88 (108.94-110.37) | 670.5 / 671.5 | +| Indexers suspended (`scripts/with_indexers_paused.sh`, plus `suggestd`) | 102.77 (102.23-103.18) | 109.50 (108.73-110.06) | 670.4 / 670.6 | + +The op cap was re-checked with the switch active: 50 gives 106.3, 100 gives 108.1, 200 and 400 give 108.4-109.8, and 1000 (mlxcel's default since #353) gives 109.2-110.1, so it stays. + +## Server path + +`mlxcel-server` on command-r7b, 128-token `/completion` requests at temperature 0, ten consecutive requests per server start, two starts per arm. Requests the lookahead gate admits (the common case) decode pipelined; requests it rejects (`ignore_eos` and other token bias, penalties, per-token logprobs, grammar) decode synchronously, one step encoded and then waited on. + +| Path | Switch off | Switch on | +|---|---|---| +| Lookahead (no `ignore_eos`), `mlxcel_batch_decode_lookahead_steps_total` +125 per request | 106.8-110.1 (median 110.0) | 115.8-119.1 (median 118.5), +7.7% | +| Synchronous (`ignore_eos`), budget raised around the step | 97.8-101.8 | 67.6-102.1, erratic | +| Synchronous, budget left on the device default (shipped) | 100.6-101.1 | 100.5-102.1 | + +A synchronous step with one large buffer cannot start on the GPU until the whole step is encoded, so the raised budget removes the overlap between CPU encoding and GPU execution inside the step, and the step time picks up the variance of CPU encoding. The switch is therefore applied only around pipelined work: the server's steady lookahead tick and its prime, and the generate loops (except under `MLXCEL_FORCE_SYNC`). + +Two earlier readings in this investigation were artifacts of measuring the server with `ignore_eos`, which takes the synchronous path: "server decode runs 9% below the CLI harness" (it does not; pipelined server decode is faster than the harness, which itself decodes with the `ignore_eos` bias), and "the first request after a start is slower with the switch on" (it was the synchronous-path instability, seen on whichever request it hit). The CLI generate loop has no first-request penalty either: single-shot `mlxcel generate -n 128` with no warmup reads 104.4-104.8 on `main` and 113.9-114.2 on the branch. + +## Not measured + +- M2, M3 and M4 base/Pro/Max. The switch follows the `MLX_MAX_OPS_PER_BUFFER` gate (M1 through M4), which was also set from Ultra measurements; the mechanism is the same, the magnitude is not known. +- M5 and later: MLX's defaults are kept in both phases, as for the op cap. +- Speculative decoding loops (DFlash round loop, MTP verify): not wired to the switch. +- Decode beyond 2048 tokens of context. A 600-token generation, which crosses the 256-token cache clear twice, holds the gain (104.5-104.7 on `main`, 110.8-111.2 on the branch, three runs each). diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 2a344bfaf..d2707408d 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -114,7 +114,8 @@ Both server entry points implement llama-server b10621's Vertex AI custom-contai | `MLX_CUDA_GRAPH_CACHE_SIZE` (MLX-native, CUDA only) | unsigned integer capacity | `2000` on CUDA builds (MLX's own default is `400`) | LRU capacity for MLX's captured CUDA-graph cache. MLX keys the cache by graph shape and its `lru_cache.h` also counts a lifetime miss counter that never resets (not on hits, not on trim), throwing a fatal `Cache thrashing` `runtime_error` once lifetime misses pass `2 * capacity` (800 at MLX's default 400). Any long-lived, shape-diverse CUDA server crosses that threshold over its lifetime, and speculative or batched decode reaches it fastest because draft/verify phases times varying batch sizes and sequence-length buckets multiply the number of distinct graph shapes; the throw is a whole-process abort, not a request-level error, so it drops every in-flight request (issue #818). mlxcel raises the default to `2000`, validated sufficient (13/13 requests across bursts on GB10). This is an LRU cap, not a preallocation, so it only costs memory as distinct graph shapes accumulate. An explicit `MLX_CUDA_GRAPH_CACHE_SIZE` always overrides. Read only by MLX's CUDA backend, so it is a harmless no-op on Metal/CPU. | | `MLX_CUDA_SDPA_CACHE_SIZE` (MLX-native, CUDA only) | unsigned integer capacity | `2000` on CUDA builds (MLX's own default is `256`) | LRU capacity for MLX's cuDNN SDPA execution-plan cache, keyed by the exact query, key, value and mask shapes and strides. It is built on the same `lru_cache.h` as the graph cache above, so it has the same lifetime miss counter and the same fatal `Cache thrashing` abort once lifetime misses pass `2 * capacity` (512 at MLX's default). Every distinct prompt length that prefills through cuDNN is one miss per attention layer class, so a long-lived server crosses 512 on prompt diversity alone, and before `MLXCEL_SDPA_FALLBACK_MAX_QUERIES` a multi-row speculative verify crossed it in a few hundred rounds (a 400-token Laguna DFlash run at block 2 aborted on it, issue #1799). mlxcel raises the default to `2000` with the same caveats as the graph cache: an LRU cap, not a preallocation, and a larger lifetime budget rather than a fix. The cache is `thread_local`, so the cap is per MLX eval thread rather than per process, and each entry is a built cuDNN graph with a selected execution plan; unlike the graph-cache raise there is no validated-sufficient datapoint for this cache, and the memory of a full cache has not been measured. An explicit value always overrides. No-op off CUDA. | | `MLX_MAX_OPS_PER_BUFFER` (MLX-native, Metal and CUDA) | positive integer op count | Metal: `1000` on pre-M5 Apple Silicon (#353), untouched elsewhere. CUDA: `100` on compute capability 12.1 (GB10) for the family whose every measured workload gains from it (`model_type` `laguna`, #1798), untouched for every other checkpoint and on every other capability (MLX's own table: 20 on A100 and 12.1, 100 on H100, B200 and consumer Blackwell, 20 for anything unlisted) | Number of ops MLX captures into one command buffer (Metal) or one CUDA graph before committing it. On CUDA this binds together with `MLX_MAX_MB_PER_BUFFER` below: MLX commits when either budget is exceeded (`needs_commit`, `mlx/backend/cuda/device.cpp`), so on GB10's default 25 "MB" the op budget is rarely reached and raising it alone reads as inert; raising both is what the GB10 MoE default does. The gate is a family allowlist, not the device alone and not a shape rule, because raising both budgets measured +17% on Laguna in every workload, -7% on gpt_oss (with +7 GB of peak memory), flat on gemma4 MoE, +4% on Qwen 3.5 4B dense and -9% on Llama 3.1 8B, ranges disjoint; no config-level shape rule separated the winners from the losers. The two Qwen MoE families are the reason the bar is every workload rather than decode alone: both gained on single-stream decode (+21% and +22%) and then failed on the serving path at n = 3, qwen3_moe flat at concurrency 4 and 8 and qwen3_5_moe -10.7% at concurrency 8 with disjoint ranges, so neither is listed. The raised budgets also hold more temporaries per graph: about +0.5 GB of peak memory at short prompts and +7 GB per 2048-token prefill chunk on Laguna. An explicit value always overrides, per variable. Measurement: `docs/benchmark_results/cuda-graph-budget-gb10-2026-09-12.md`. | -| `MLX_MAX_MB_PER_BUFFER` (MLX-native, CUDA only) | positive integer, nominally MB | `1000` on compute capability 12.1 (GB10) for the family whose every measured workload gains from it (`model_type` `laguna`, #1798), untouched for every other checkpoint and elsewhere (MLX's own table: 400 on A100, 1000 on H100, B200 and consumer Blackwell, 25 on 12.1, 100 for anything unlisted) | Input-size budget of one captured CUDA graph. The counter is not bytes: MLX sums `array::data_size()` over the graph's input arrays, an element count, and compares `count >> 20` against this value, so on the 25 default any op reading an array over 26.2M elements commits its own graph. That is every 256-expert NVFP4 expert stack on Laguna (about 120 `gather_qmm` per token) and every 4-bit `lm_head` or tied embedding over 26.2M packed words (Qwen 3.5 4B, Llama 3.1 8B, Gemma 3 4B). On a model where that happens per layer, capture at 25 costs more than graphs off; whether raising the budget helps is a per-family measurement (see `MLX_MAX_OPS_PER_BUFFER` above). An explicit value always overrides. No-op off CUDA. | +| `MLX_MAX_MB_PER_BUFFER` (MLX-native, Metal and CUDA) | positive integer, nominally MB | Metal: unset, so MLX's own table applies (40 on base/Pro/phone, 50 on Max and Ultra); mlxcel raises the budget for decode steps only, see `MLXCEL_DECODE_MB_PER_BUFFER`. CUDA: `1000` on compute capability 12.1 (GB10) for the family whose every measured workload gains from it (`model_type` `laguna`, #1798), untouched for every other checkpoint and elsewhere (MLX's own table: 400 on A100, 1000 on H100, B200 and consumer Blackwell, 25 on 12.1, 100 for anything unlisted) | Input-size budget of one captured CUDA graph. The counter is not bytes: MLX sums `array::data_size()` over the graph's input arrays, an element count, and compares `count >> 20` against this value, so on the 25 default any op reading an array over 26.2M elements commits its own graph. That is every 256-expert NVFP4 expert stack on Laguna (about 120 `gather_qmm` per token) and every 4-bit `lm_head` or tied embedding over 26.2M packed words (Qwen 3.5 4B, Llama 3.1 8B, Gemma 3 4B). On a model where that happens per layer, capture at 25 costs more than graphs off; whether raising the budget helps is a per-family measurement (see `MLX_MAX_OPS_PER_BUFFER` above). On Metal the same element-count budget decides when a command buffer is committed. Setting it here pins it for prefill and decode alike and turns off the decode-only switch below. An explicit value always overrides. | +| `MLXCEL_DECODE_MB_PER_BUFFER` | `0`/`off`/`false`/`no` (disable), positive integer | `1000` on pre-M5 Apple Silicon (the same gate as `MLX_MAX_OPS_PER_BUFFER`), off elsewhere | Metal command-buffer input budget applied during decode steps only. MLX commits a command buffer once the element count of its distinct inputs, shifted right by 20, passes `MLX_MAX_MB_PER_BUFFER` (40-50 by default). With `MLX_MAX_OPS_PER_BUFFER` raised to 1000 that budget is the only cap that binds, and in decode, where every token reads the whole weight set, the default commits a buffer every one to two layers (about 23 per token on command-r7b 4-bit) and idles the GPU at each boundary. mlxcel raises the budget around pipelined decode only (the generate loops and the server's lookahead decode) and leaves prefill and synchronous decode steps on the device default, through a runtime override in the `mlx/backend/metal/device.cpp` overlay. A synchronous step encodes and then waits, so one large buffer there only removes the overlap between CPU encoding and GPU execution. Measured on M1 Ultra (500-token prompt, 128 generated tokens, three interleaved runs per cell), decode at 1000 versus the default: command-r7b 4-bit +7%, Llama 3.1 8B 4-bit +5.6%, Qwen2.5 7B 4-bit +8%, Gemma 3n E4B +3%, Granite 4.0 H Tiny +10%, Qwen3-30B-A3B +20%, Mixtral 8x7B +21%, Llama 3.1 8B bf16 +17%, Gemma 3 4B flat. Applying the same value to prefill as well would roughly double peak memory on long prompts (2048 tokens: Qwen2.5 7B 4-bit 6.0 to 12.6 GB, Qwen3-30B-A3B 19.8 to 36.2 GB) and cost up to 2.7% prefill throughput, which is why it is scoped to decode. Speculative draft/verify loops are not covered yet. Ignored when `MLX_MAX_MB_PER_BUFFER` is set. Measurement: `docs/benchmark_results/metal-mb-per-buffer-m1ultra-2026-09-21.md`. | | `MLXCEL_HEADROOM_FACTOR` | positive `f64` | `1.20` | Runtime/activation headroom multiplier used by the unified memory estimator (`mlxcel inspect`, `--estimate-memory`, `--recommend-quant`). Positive values `<= 1.0` disable the headroom term; invalid or non-positive values warn and fall back to the default. Override only for calibration runs — see the in-code recipe in `src/execution/memory_estimate.rs`. | | `MLXCEL_CACHE_DIR` | directory path | `$HOME/.cache/mlxcel` | Root for mlxcel's on-disk caches. The tokenizer language-analysis disk cache (language-bias features) lives under `tokenizer-scripts/`, and the location-independent global model store lives under `models//` when `MLXCEL_MODELS_DIR` and the store-root flag (`--model-store-root` on the servers, `--models-dir` on the subcommands) are both unset. | | `MLXCEL_MODELS_DIR` | directory path | unset (falls back to `${MLXCEL_CACHE_DIR:-$HOME/.cache/mlxcel}/models`) | Dedicated model-store root. Snapshots live directly at `$MLXCEL_MODELS_DIR//` with no `models/` subdir, so the whole store can sit on a separate volume without dragging the tokenizer-script cache along. Read by `mlxcel download`, the `-m/--model` resolver (`generate` / `serve` / `inspect` / `run`), the `mlxcel-server -m/--model` resolver, and `list` / `rm`. Resolution precedence for the models root: the CLI flag (`--model-store-root ` on `mlxcel-server` / `mlxcel serve` since #1438 reserved `--models-dir` for b10621 router mode; still `--models-dir ` on the `download` / `list` / `rm` / `generate` subcommands), then `MLXCEL_MODELS_DIR`, then `${MLXCEL_CACHE_DIR:-$HOME/.cache/mlxcel}/models`. (`download --local-dir ` is separate: it writes the snapshot verbatim at that exact path.) | diff --git a/src/lib/mlx-cpp/CMakeLists.txt b/src/lib/mlx-cpp/CMakeLists.txt index 133b53368..b5163a7f0 100644 --- a/src/lib/mlx-cpp/CMakeLists.txt +++ b/src/lib/mlx-cpp/CMakeLists.txt @@ -33,8 +33,14 @@ function(mlx_apply_source_overlays mlx_source_dir) # `use_qmv_wide` (issue #1187). Upstream has no equivalent knob and the # predicate is unchanged as of pin 81ba1c6a, so it has to be overlaid here. # Refresh the file wholesale on a bump and re-apply that one hunk. + # device.cpp carries a runtime override of the per-command-buffer input + # budget (`mlxcel_set_mb_per_buffer_override`), which mlxcel raises around + # decode steps and leaves at the device default for prefill. Upstream latches + # the budget once from MLX_MAX_MB_PER_BUFFER and has no setter. Same refresh + # rule: take the new upstream file and re-apply the three marked hunks. set(_metal_patch_files "mlx/backend/metal/compiled.cpp" + "mlx/backend/metal/device.cpp" "mlx/backend/metal/kernels/utils.h" "mlx/backend/metal/quantized.cpp") foreach(_patch_file ${_metal_patch_files}) diff --git a/src/lib/mlx-cpp/patches/mlx/backend/metal/device.cpp b/src/lib/mlx-cpp/patches/mlx/backend/metal/device.cpp new file mode 100644 index 000000000..824aa4ce8 --- /dev/null +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/device.cpp @@ -0,0 +1,1015 @@ +// Copyright © 2023-2024 Apple Inc. +// Patched by mlxcel: a runtime override for the per-command-buffer input +// budget. Synced to upstream 81ba1c6a. The delta is three hunks: the +// include below, the override read in `CommandEncoder::needs_commit`, and the +// two bridge entry points `mlxcel_set_mb_per_buffer_override()` / +// `mlxcel_mb_per_buffer_override()` at the end of the file. Everything else is +// upstream verbatim, so a bump refreshes this file and re-applies those hunks. +// +// Why it exists. MLX commits a command buffer once the element count of its +// distinct inputs passes `max_mb_per_buffer_ << 20`, and `Device` latches that +// cap from `MLX_MAX_MB_PER_BUFFER` once, at construction. The right cap differs +// by phase: decode streams the whole weight set every token, so the default +// (40-50) commits a buffer every one to two layers and leaves the GPU idle at +// each boundary (+3% to +21% decode on M1 Ultra at 1000), while prefill holds +// its activations until the buffer completes, so the same 1000 roughly doubles +// peak memory on a 2048-token prompt. A process-wide constant cannot serve +// both; mlxcel raises the cap around decode steps only and leaves prefill on +// the device default. Upstream exposes no setter, so the switch lives here. +// +// Zero means "no override": `needs_commit` then behaves exactly as upstream. +// Relaxed ordering is enough because the value only decides where a command +// buffer boundary falls, never what is computed, and the caller sets it on the +// thread that encodes the work it means to affect. + +#include +#include +#include + +#include + +#define NS_PRIVATE_IMPLEMENTATION +#define CA_PRIVATE_IMPLEMENTATION +#define MTL_PRIVATE_IMPLEMENTATION + +#include "mlx/backend/common/utils.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/event.h" +#include "mlx/backend/metal/metal.h" +#include "mlx/backend/metal/utils.h" +#include "mlx/utils.h" + +namespace std { + +// Required for putting the pointer in unordered_set. +template +struct hash> { + size_t operator()(const NS::SharedPtr& p) const { + return std::hash{}(p.get()); + } +}; + +} // namespace std + +namespace mlx::core::metal { + +namespace { + +constexpr const char* default_mtllib_path = METAL_PATH; + +void set_compile_options( + MTL::CompileOptions* mtl_options, + const CompileOptions& compile_options) { + if (__builtin_available(macOS 15, iOS 18, tvOS 18, visionOS 2, *)) { + switch (compile_options.math_mode) { + case MathMode::Safe: + mtl_options->setMathMode(MTL::MathModeSafe); + break; + case MathMode::Relaxed: + mtl_options->setMathMode(MTL::MathModeRelaxed); + break; + case MathMode::Fast: + mtl_options->setMathMode(MTL::MathModeFast); + break; + default: + throw std::invalid_argument("[metal::Device] Invalid math mode."); + } + } else { + if (compile_options.math_mode == MathMode::Relaxed) { + throw std::runtime_error( + "[metal::Device] Metal math mode `relaxed` requires macOS 15, " + "iOS 18, tvOS 18, or visionOS 2."); + } + mtl_options->setFastMathEnabled( + compile_options.math_mode == MathMode::Fast); + } +} + +auto get_metal_version() { + auto get_metal_version_ = []() { + if (__builtin_available(macOS 27, iOS 27, tvOS 27, visionOS 27, *)) { + // TODO: Use MTL::LanguageVersion4_1 after metal-cpp_27 is released. + return static_cast((4 << 16) + 1); + } else if (__builtin_available(macOS 26, iOS 26, tvOS 26, visionOS 26, *)) { + return MTL::LanguageVersion4_0; + } else if (__builtin_available(macOS 15, iOS 18, tvOS 18, visionOS 2, *)) { + return MTL::LanguageVersion3_2; + } else { + return MTL::LanguageVersion3_1; + } + }; + static auto metal_version_ = get_metal_version_(); + return metal_version_; +} + +NS::SharedPtr load_device() { + auto pool = new_scoped_memory_pool(); + auto devices = NS::TransferPtr(MTL::CopyAllDevices()); + // In headless, sandboxed, or virtualized macOS sessions CopyAllDevices() + // returns an empty NSArray. Indexing object(0) on an empty array raises + // an unrecoverable NSRangeException, so guard the access and fall back to + // CreateSystemDefaultDevice (which can also return null). + MTL::Device* first = nullptr; + if (devices && devices->count() > 0) { + first = static_cast(devices->object(0)); + } + auto device = first ? NS::RetainPtr(first) + : NS::TransferPtr(MTL::CreateSystemDefaultDevice()); + if (!device) { + throw std::runtime_error( + "[metal::load_device] No Metal device available. This typically " + "occurs in headless, sandboxed, or virtualized macOS sessions " + "where the GPU is not accessible."); + } + return device; +} + +std::pair load_library_from_path( + MTL::Device* device, + const char* path) { + auto library = NS::String::string(path, NS::UTF8StringEncoding); + NS::Error* error; + auto lib = device->newLibrary(library, &error); + + return std::make_pair(lib, error); +} + +#ifdef SWIFTPM_BUNDLE +MTL::Library* try_load_bundle( + MTL::Device* device, + NS::URL* url, + const std::string& lib_name) { + std::string bundle_path = std::string(url->fileSystemRepresentation()) + "/" + + SWIFTPM_BUNDLE + ".bundle"; + auto bundle = NS::Bundle::alloc()->init( + NS::String::string(bundle_path.c_str(), NS::UTF8StringEncoding)); + if (bundle != nullptr) { + std::string resource_path = + std::string(bundle->resourceURL()->fileSystemRepresentation()) + "/" + + lib_name + ".metallib"; + auto [lib, error] = load_library_from_path(device, resource_path.c_str()); + if (lib) { + return lib; + } + } + return nullptr; +} + +MTL::Library* try_load_framework( + MTL::Device* device, + NS::URL* url, + const std::string& lib_name) { + std::string resource_path = std::string(url->fileSystemRepresentation()) + + "/" + lib_name + ".metallib"; + auto [lib, error] = load_library_from_path(device, resource_path.c_str()); + if (lib) { + return lib; + } + return nullptr; +} +#endif + +// Firstly, search for the metallib in the same path as this binary +std::pair load_colocated_library( + MTL::Device* device, + const std::string& relative_path) { + auto path = current_binary_dir() / relative_path; + if (!path.has_extension()) { + path.replace_extension(".metallib"); + } + + return load_library_from_path(device, path.c_str()); +} + +std::pair load_swiftpm_library( + MTL::Device* device, + const std::string& lib_name) { +#ifdef SWIFTPM_BUNDLE + MTL::Library* library = + try_load_bundle(device, NS::Bundle::mainBundle()->bundleURL(), lib_name); + if (library != nullptr) { + return {library, nullptr}; + } + auto bundles = NS::Bundle::allBundles(); + for (int i = 0, c = (int)bundles->count(); i < c; i++) { + auto bundle = reinterpret_cast(bundles->object(i)); + library = try_load_bundle(device, bundle->resourceURL(), lib_name); + if (library != nullptr) { + return {library, nullptr}; + } + } + // if SWIFTPM_BUNDLE is a framework identifier, try loading from that + auto frameworks = NS::Bundle::allFrameworks(); + for (int i = 0, c = (int)frameworks->count(); i < c; i++) { + const auto bundle = reinterpret_cast(frameworks->object(i)); + const auto identifier = bundle->bundleIdentifier(); + if (identifier != nullptr && + !strcmp(identifier->utf8String(), SWIFTPM_BUNDLE)) { + library = try_load_framework(device, bundle->resourceURL(), lib_name); + if (library != nullptr) { + return {library, nullptr}; + } + } + } +#endif + return {nullptr, nullptr}; +} + +MTL::Library* load_default_library(MTL::Device* device) { + // Check override path before automatic lookup + if (!get_metallib_path().empty()) { + auto [lib, error] = + load_library_from_path(device, get_metallib_path().c_str()); + if (!lib) { + throw std::runtime_error( + fmt::format( + "Can not load metallib from specified location \"{}\": {}.", + get_metallib_path(), + error->localizedDescription()->utf8String())); + } + return lib; + } + + NS::Error* error[5]; + MTL::Library* lib; + // First try the colocated mlx.metallib + std::tie(lib, error[0]) = load_colocated_library(device, "mlx"); + if (lib) { + return lib; + } + + std::tie(lib, error[1]) = load_colocated_library(device, "Resources/mlx"); + if (lib) { + return lib; + } + + // Then try default.metallib in a SwiftPM bundle if we have one + std::tie(lib, error[2]) = load_swiftpm_library(device, "default"); + if (lib) { + return lib; + } + + // Try lo load resources from Framework resources if SwiftPM wrapped as a + // dynamic framework. + std::tie(lib, error[3]) = load_colocated_library(device, "Resources/default"); + if (lib) { + return lib; + } + + // Finally try default_mtllib_path + std::tie(lib, error[4]) = load_library_from_path(device, default_mtllib_path); + if (!lib) { + std::ostringstream msg; + msg << "Failed to load the default metallib. "; + for (int i = 0; i < 5; i++) { + if (error[i] != nullptr) { + msg << error[i]->localizedDescription()->utf8String() << " "; + } + } + throw std::runtime_error(msg.str()); + } + return lib; +} + +MTL::Library* load_library( + MTL::Device* device, + const std::string& lib_name, + const std::string& lib_path) { + // We have been given a path that ends in metallib so try to load it + if (lib_path.size() > 9 && + std::equal(lib_path.end() - 9, lib_path.end(), ".metallib")) { + auto [lib, error] = load_library_from_path(device, lib_path.c_str()); + if (!lib) { + std::ostringstream msg; + msg << "Failed to load the metallib from <" << lib_path << "> with error " + << error->localizedDescription()->utf8String(); + throw std::runtime_error(msg.str()); + } + return lib; + } + + // We have been given a path so try to load from lib_path / lib_name.metallib + if (lib_path.size() > 0) { + std::string full_path = lib_path + "/" + lib_name + ".metallib"; + auto [lib, error] = load_library_from_path(device, full_path.c_str()); + if (!lib) { + std::ostringstream msg; + msg << "Failed to load the metallib from <" << full_path + << "> with error " << error->localizedDescription()->utf8String(); + throw std::runtime_error(msg.str()); + } + return lib; + } + + // Try to load the colocated library + { + auto [lib, error] = load_colocated_library(device, lib_name); + if (lib) { + return lib; + } + } + + // Try to load the library from swiftpm + { + auto [lib, error] = load_swiftpm_library(device, lib_name); + if (lib) { + return lib; + } + } + + std::ostringstream msg; + msg << "Failed to load the metallib " << lib_name << ".metallib. " + << "We attempted to load it from <" << current_binary_dir() << "/" + << lib_name << ".metallib>"; +#ifdef SWIFTPM_BUNDLE + msg << " and from the Swift PM bundle."; +#endif + throw std::runtime_error(msg.str()); +} + +} // namespace + +CommandEncoder::CommandEncoder( + Device& d, + int index, + ResidencySets& residency_sets) + : device_(d), residency_sets_(residency_sets) { + auto pool = new_scoped_memory_pool(); + queue_ = NS::TransferPtr(device_.mtl_device()->newCommandQueue()); + if (!queue_) { + throw std::runtime_error( + "[metal::CommandEncoder] Failed to make new command queue."); + } + // Sets created later are attached in commit(). + residency_sets_.attach_new_sets(queue_.get(), sets_attached_); + debug_set_stream_queue_label(queue_.get(), index); + buffer_ = NS::RetainPtr(queue_->commandBufferWithUnretainedReferences()); +} + +CommandEncoder::~CommandEncoder() { + exiting_ = true; + synchronize(); + auto pool = new_scoped_memory_pool(); + buffer_.reset(); + queue_.reset(); +} + +void CommandEncoder::set_buffer( + const MTL::Buffer* buf, + int idx, + int64_t offset /* = 0 */) { + // Record as both input and output to ensure synchronization between command + // buffers + all_inputs_.insert((void*)buf); + all_outputs_.insert((void*)buf); + get_command_encoder()->setBuffer(buf, offset, idx); +} + +void CommandEncoder::set_input_array( + const array& a, + int idx, + int64_t offset /* = 0 */) { + if (all_inputs_.insert(a.buffer().ptr()).second) { + buffer_sizes_ += a.data_size(); + } + auto r_buf = static_cast(const_cast(a.buffer().ptr())); + next_inputs_.insert(r_buf); + needs_barrier_ = + needs_barrier_ | (prev_outputs_.find(r_buf) != prev_outputs_.end()); + auto a_buf = static_cast(a.buffer().ptr()); + get_command_encoder()->setBuffer(a_buf, a.offset() + offset, idx); +} + +void CommandEncoder::set_output_array( + array& a, + int idx, + int64_t offset /* = 0 */) { + // Add barriers before adding the output to the output set + set_input_array(a, idx, offset); + register_output_array(a); +} + +void CommandEncoder::register_output_array(const array& a) { + all_outputs_.insert(a.buffer().ptr()); + + auto buf = static_cast(const_cast(a.buffer().ptr())); + if (concurrent_) { + concurrent_outputs_.insert(buf); + } else { + next_outputs_.insert(buf); + needs_barrier_ = + needs_barrier_ | (prev_inputs_.find(buf) != prev_inputs_.end()); + } +} + +void CommandEncoder::add_temporary(array arr) { + temporaries_.push_back(std::move(arr)); +} + +void CommandEncoder::add_temporaries(std::vector arrays) { + temporaries_.insert( + temporaries_.end(), + std::make_move_iterator(arrays.begin()), + std::make_move_iterator(arrays.end())); +} + +void CommandEncoder::maybeInsertBarrier() { + if (needs_barrier_) { + get_command_encoder()->memoryBarrier(MTL::BarrierScopeBuffers); + needs_barrier_ = false; + // Preserve the hash tables' buckets for reuse across barrier epochs. + prev_inputs_.swap(next_inputs_); + prev_outputs_.swap(next_outputs_); + } else { + prev_inputs_.insert(next_inputs_.begin(), next_inputs_.end()); + prev_outputs_.insert(next_outputs_.begin(), next_outputs_.end()); + } + next_inputs_.clear(); + next_outputs_.clear(); +} + +void CommandEncoder::dispatch_threadgroups( + MTL::Size grid_dims, + MTL::Size group_dims) { + maybeInsertBarrier(); + buffer_ops_++; + get_command_encoder()->dispatchThreadgroups(grid_dims, group_dims); +} + +void CommandEncoder::dispatch_threads( + MTL::Size grid_dims, + MTL::Size group_dims) { + maybeInsertBarrier(); + buffer_ops_++; + get_command_encoder()->dispatchThreads(grid_dims, group_dims); +} + +void CommandEncoder::barrier() { + get_command_encoder()->memoryBarrier(MTL::BarrierScopeBuffers); +} + +void CommandEncoder::end_encoding() { + // Each command encoder has a unique fence. We also store a map of + // all previous outputs of command encoders to their corresponding fence. + // - The command encoder records its inputs and outputs. + // - Wait on a fence if any inputs in the encoder are outputs of a previous + // encoder. + // - Update the map of outputs to include this command encoder's outputs. + // - Always signal this command encoders fence. + // - Add a completion handler for this command encoder that removes outputs + // from the map to limit the growth of the map and avoid unnecessary waits + // - Temporaries are a special case as they do not cross command encoder + // boundaries. These can be removed early from the encoders inputs and + // outputs since they don't need synchronization. + if (!encoder_) { + return; + } + + // Remove temporaries from inputs and outputs. + for (auto& t : temporaries_) { + all_outputs_.erase(t.buffer().ptr()); + all_inputs_.erase(t.buffer().ptr()); + } + + // Keep references to the fences we waited on and put them in the completion + // handler so they are not prematurely released. + std::unordered_set> waiting_on; + { + std::lock_guard lk(outputs_mtx_); + for (auto& in : all_inputs_) { + if (auto it = prev_ce_outputs_.find(in); it != prev_ce_outputs_.end()) { + // If we've already waited on a fence, don't wait on it again. + if (waiting_on.find(it->second) == waiting_on.end()) { + encoder_->waitForFence(it->second.get()); + waiting_on.insert(it->second); + } + } + } + for (auto& out : all_outputs_) { + prev_ce_outputs_[out] = fence_; + } + } + + encoder_->updateFence(fence_.get()); + buffer_->addCompletedHandler([this, + fence = std::move(fence_), + temporaries = std::move(temporaries_), + all_outputs = std::move(all_outputs_), + waiting_on = std::move(waiting_on)]( + MTL::CommandBuffer*) mutable { + std::lock_guard lk(outputs_mtx_); + for (auto& o : all_outputs) { + if (auto it = prev_ce_outputs_.find(o); it != prev_ce_outputs_.end()) { + if (it->second == fence) { + prev_ce_outputs_.erase(it); + } + } + } + }); + + encoder_->endEncoding(); + encoder_.reset(); + needs_barrier_ = false; + concurrent_ = false; + prev_inputs_.clear(); + next_inputs_.clear(); + prev_outputs_.clear(); + next_outputs_.clear(); + concurrent_outputs_.clear(); + all_inputs_.clear(); +} + +void CommandEncoder::signal_event(Event event, uint64_t value) { + end_encoding(); + buffer_->encodeSignalEvent(event.cast().mtl_event(), value); + signal_events_.push_back({std::move(event), value}); +} + +void CommandEncoder::wait_event(Event event, uint64_t value) { + end_encoding(); + buffer_->encodeWait(event.cast().mtl_event(), value); + wait_events_.push_back(std::move(event)); +} + +// mlxcel: runtime override of the input budget (see the header comment). +static std::atomic& mlxcel_mb_per_buffer_override_slot() { + static std::atomic slot{0}; + return slot; +} + +bool CommandEncoder::needs_commit() const { + auto [max_ops, max_mb] = device_.get_max_ops_mb_per_buffer(); + if (int override_mb = + mlxcel_mb_per_buffer_override_slot().load(std::memory_order_relaxed); + override_mb > 0) { + max_mb = override_mb; + } + return (buffer_ops_ > max_ops) || ((buffer_sizes_ >> 20) > max_mb); +} + +void CommandEncoder::commit(std::function completion) { + // Metal locks a command buffer's residency at commit time, so attach any + // sets created since the last commit first. + residency_sets_.attach_new_sets(queue_.get(), sets_attached_); + buffer_->addCompletedHandler( + [&error_ = error_, + wait_events = std::move(wait_events_), + signal_events = std::move(signal_events_), + completion = std::move(completion)](MTL::CommandBuffer* cbuf) mutable { + if (completion) { + completion(); + } + // If any of the waited event has error in it, poison the encoder. + for (auto& event : wait_events) { + if (error_.store_if_valid(event.load_error())) { + break; + } + } + // Set error only when no error happended before, to preserve the + // earliest error. + bool has_error = error_.valid(); + if (!has_error && cbuf->status() == MTL::CommandBufferStatusError) { + error_.set_message( + std::make_shared(fmt::format( + "[METAL] Command buffer execution failed: {}.", + cbuf->error()->localizedDescription()->utf8String()))); + has_error = true; + } + // Poison all the signaled events when error happened. + if (has_error) { + for (auto& [event, value] : signal_events) { + event.set_error(error_); + } + } + // Metal won't signal the events for us on error, manually signal them + // to avoid infinite waiting. + if (cbuf->status() == MTL::CommandBufferStatusError) { + for (auto& [event, value] : signal_events) { + event.cast().signal(value); + } + } + }); + buffer_->commit(); + buffer_ = NS::RetainPtr(queue_->commandBufferWithUnretainedReferences()); + buffer_ops_ = 0; + buffer_sizes_ = 0; +} + +void CommandEncoder::synchronize() { + auto pool = new_scoped_memory_pool(); + auto cbuf = buffer_; // retained + end_encoding(); + commit(); + cbuf->waitUntilCompleted(); + + if (!exiting_) { + error_.check(); + } +} + +MTL::ComputeCommandEncoder* CommandEncoder::get_command_encoder() { + if (!encoder_) { + error_.check(); + encoder_ = NS::RetainPtr( + buffer_->computeCommandEncoder(MTL::DispatchTypeConcurrent)); + fence_ = NS::TransferPtr(device_.mtl_device()->newFence()); + } + return encoder_.get(); +} + +Device::Device() : device_(load_device()), residency_sets_(device_.get()) { + auto pool = new_scoped_memory_pool(); + default_library_ = NS::TransferPtr(load_default_library(device_.get())); + arch_ = env::metal_gpu_arch(); + if (arch_.empty()) { + arch_ = std::string(device_->architecture()->name()->utf8String()); + } + int ag_tens = 0; + int ag_ones = 0; + if (arch_.size() >= 3) { + ag_tens = arch_[arch_.size() - 3] - '0'; + ag_ones = arch_[arch_.size() - 2] - '0'; + ag_tens = (ag_tens < 10 && ag_tens >= 0) ? ag_tens : 0; + ag_ones = (ag_ones < 10 && ag_ones >= 0) ? ag_ones : 0; + } + arch_gen_ = ag_tens * 10 + ag_ones; + auto arch = arch_.back(); + switch (arch) { + case 'p': // phone + max_ops_per_buffer_ = 20; + max_mb_per_buffer_ = 40; + break; + case 'g': // base, pro + max_ops_per_buffer_ = 40; + max_mb_per_buffer_ = 40; + break; + case 's': // max + max_ops_per_buffer_ = 50; + max_mb_per_buffer_ = 50; + break; + case 'd': // ultra + max_ops_per_buffer_ = 50; + max_mb_per_buffer_ = 50; + break; + default: // default to medium + max_ops_per_buffer_ = 40; + max_mb_per_buffer_ = 40; + break; + } + max_ops_per_buffer_ = env::max_ops_per_buffer(max_ops_per_buffer_); + max_mb_per_buffer_ = env::max_mb_per_buffer(max_mb_per_buffer_); +} + +Device::~Device() = default; + +MTL::Library* Device::get_library( + const std::string& name, + const std::string& path /* = "" */) { + { + std::shared_lock rlock(library_mtx_); + if (auto it = library_map_.find(name); it != library_map_.end()) { + return it->second.get(); + } + } + + std::unique_lock wlock(library_mtx_); + if (auto it = library_map_.find(name); it != library_map_.end()) { + return it->second.get(); + } + + auto new_lib = load_library(device_.get(), name, path.c_str()); + library_map_.insert({name, NS::TransferPtr(new_lib)}); + return new_lib; +} + +NS::SharedPtr Device::build_library_( + const std::string& source_string, + const CompileOptions& compile_options) { + auto pool = new_scoped_memory_pool(); + + auto ns_code = + NS::String::string(source_string.c_str(), NS::ASCIIStringEncoding); + + NS::Error* error = nullptr; + auto options = MTL::CompileOptions::alloc()->init()->autorelease(); + set_compile_options(options, compile_options); + options->setLanguageVersion(get_metal_version()); +#ifndef NDEBUG + if (options->languageVersion() >= MTL::LanguageVersion3_2) { + options->setEnableLogging(true); + } +#endif + auto mtl_lib = NS::TransferPtr(device_->newLibrary(ns_code, options, &error)); + + // Throw error if unable to compile library + if (!mtl_lib) { + std::ostringstream msg; + msg << "[metal::Device] Unable to build metal library from source\n"; + if (error) { + msg << error->localizedDescription()->utf8String() << "\n"; + } + throw std::runtime_error(msg.str()); + } + + return mtl_lib; +} + +NS::SharedPtr Device::get_function_( + const std::string& name, + MTL::Library* mtl_lib) { + auto pool = new_scoped_memory_pool(); + // Pull kernel from library + auto ns_name = NS::String::string(name.c_str(), NS::ASCIIStringEncoding); + return NS::TransferPtr(mtl_lib->newFunction(ns_name)); +} + +NS::SharedPtr Device::get_function_( + const std::string& name, + const std::string& specialized_name, + const MTLFCList& func_consts, + MTL::Library* mtl_lib) { + if (func_consts.empty() && (specialized_name == name)) { + return get_function_(name, mtl_lib); + } + + auto pool = new_scoped_memory_pool(); + + // Prepare function constants + auto mtl_func_consts = + MTL::FunctionConstantValues::alloc()->init()->autorelease(); + + for (auto [value, type, index] : func_consts) { + mtl_func_consts->setConstantValue(value, type, index); + } + + // Prepare function desc + auto desc = MTL::FunctionDescriptor::functionDescriptor(); + desc->setName(NS::String::string(name.c_str(), NS::ASCIIStringEncoding)); + desc->setSpecializedName( + NS::String::string(specialized_name.c_str(), NS::ASCIIStringEncoding)); + desc->setConstantValues(mtl_func_consts); + + // Pull kernel from library + NS::Error* error = nullptr; + auto mtl_function = NS::TransferPtr(mtl_lib->newFunction(desc, &error)); + + // Throw error if unable to build metal function + if (!mtl_function) { + std::ostringstream msg; + msg << "[metal::Device] Unable to load function " << name << "\n"; + if (error) { + msg << error->localizedDescription()->utf8String() << "\n"; + } + throw std::runtime_error(msg.str()); + } + + return mtl_function; +} + +NS::SharedPtr Device::get_kernel_( + const std::string& name, + const MTL::Function* mtl_function) { + // Compile kernel to compute pipeline + NS::Error* error = nullptr; + NS::SharedPtr kernel; + + if (mtl_function) { + kernel = + NS::TransferPtr(device_->newComputePipelineState(mtl_function, &error)); + } + + // Throw error if unable to compile metal function + if (!mtl_function || !kernel) { + std::ostringstream msg; + msg << "[metal::Device] Unable to load kernel " << name << "\n"; + if (error) { + msg << error->localizedDescription()->utf8String() << "\n"; + } + throw std::runtime_error(msg.str()); + } + + return kernel; +} + +NS::SharedPtr Device::get_kernel_( + const std::string& name, + const MTL::Function* mtl_function, + const MTL::LinkedFunctions* linked_functions) { + // Check inputs + if (!linked_functions) { + return get_kernel_(name, mtl_function); + } + + if (!mtl_function) { + std::ostringstream msg; + msg << "[metal::Device] Unable to load kernel " << name << "\n"; + throw std::runtime_error(msg.str()); + } + + auto pool = new_scoped_memory_pool(); + + // Prepare compute pipeline state descriptor + auto desc = MTL::ComputePipelineDescriptor::alloc()->init()->autorelease(); + desc->setComputeFunction(mtl_function); + desc->setLinkedFunctions(linked_functions); + + // Compile kernel to compute pipeline + NS::Error* error = nullptr; + auto kernel = NS::TransferPtr(device_->newComputePipelineState( + desc, MTL::PipelineOptionNone, nullptr, &error)); + + // Throw error if unable to compile metal function + if (!kernel) { + std::ostringstream msg; + msg << "[metal::Device] Unable to load kernel " << name << "\n"; + if (error) { + msg << error->localizedDescription()->utf8String() << "\n"; + } + throw std::runtime_error(msg.str()); + } + + return kernel; +} + +MTL::Library* Device::get_library( + const std::string& name, + const CompileOptions& compile_options, + const std::function& builder) { + { + std::shared_lock rlock(library_mtx_); + if (auto it = library_map_.find(name); it != library_map_.end()) { + return it->second.get(); + } + } + + std::unique_lock wlock(library_mtx_); + if (auto it = library_map_.find(name); it != library_map_.end()) { + return it->second.get(); + } + + auto mtl_lib = build_library_(builder(), compile_options); + library_map_.insert({name, mtl_lib}); + return mtl_lib.get(); +} + +void Device::clear_library(const std::string& name) { + std::unique_lock wlock(library_mtx_); + if (auto it = library_map_.find(name); it != library_map_.end()) { + library_kernels_.erase(it->second.get()); + library_map_.erase(it); + } +} + +NS::SharedPtr Device::get_linked_functions_( + const std::vector& funcs) { + if (funcs.empty()) { + return nullptr; + } + + auto pool = new_scoped_memory_pool(); + auto lfuncs = NS::TransferPtr(MTL::LinkedFunctions::linkedFunctions()); + NS::Array* funcs_arr = NS::Array::array( + reinterpret_cast(funcs.data()), funcs.size()); + lfuncs->setPrivateFunctions(funcs_arr); + return lfuncs; +} + +MTL::ComputePipelineState* Device::get_kernel_( + const std::string& base_name, + MTL::Library* mtl_lib, + const std::string& hash_name, + const MTLFCList& func_consts /* = {} */, + const std::vector& linked_functions /* = {} */) { + // Single writer allowed + std::unique_lock wlock(kernel_mtx_); + + // Try loading again to avoid loading twice + auto& kernel_map_ = library_kernels_[mtl_lib]; + if (auto it = kernel_map_.find(hash_name); it != kernel_map_.end()) { + return it->second.get(); + } + + auto pool = new_scoped_memory_pool(); + + // Pull kernel from library + auto mtl_function = get_function_(base_name, hash_name, func_consts, mtl_lib); + + // Compile kernel to compute pipeline + auto mtl_linked_funcs = get_linked_functions_(linked_functions); + auto kernel = + get_kernel_(hash_name, mtl_function.get(), mtl_linked_funcs.get()); + + // Add kernel to cache + kernel_map_.insert({hash_name, kernel}); + + return kernel.get(); +} + +MTL::ComputePipelineState* Device::get_kernel( + const std::string& base_name, + MTL::Library* mtl_lib, + const std::string& hash_name /* = "" */, + const MTLFCList& func_consts /* = {} */, + const std::vector& linked_functions /* = {} */) { + const auto& kname = hash_name.empty() ? base_name : hash_name; + { + // Multiple readers allowed + std::shared_lock lock(kernel_mtx_); + + // Look for cached kernel + auto library_it = library_kernels_.find(mtl_lib); + if (library_it != library_kernels_.end()) { + auto kernel_it = library_it->second.find(kname); + if (kernel_it != library_it->second.end()) { + return kernel_it->second.get(); + } + } + } + return get_kernel_(base_name, mtl_lib, kname, func_consts, linked_functions); +} + +MTL::ComputePipelineState* Device::get_kernel( + const std::string& base_name, + const std::string& hash_name /* = "" */, + const MTLFCList& func_consts /* = {} */, + const std::vector& linked_functions /* = {} */) { + return get_kernel( + base_name, + default_library_.get(), + hash_name, + func_consts, + linked_functions); +} + +Device& device(mlx::core::Device) { + // Leak singleton device intentionally, to avoid cases where a compute kernel + // returns and tries to access the object after it has been freed by the main + // thread teardown. + static Device* metal_device = new Device; + return *metal_device; +} + +CommandEncoder& get_command_encoder(Stream s) { + auto& encoders = get_command_encoders(); + auto it = encoders.find(s.index); + if (it == encoders.end()) { + auto& global_encoders = get_global_command_encoders(); + it = global_encoders.find(s.index); + if (it == global_encoders.end()) { + throw std::runtime_error( + fmt::format( + "There is no Stream(gpu, {}) in current thread.", s.index)); + } + } + return it->second; +} + +std::unordered_map& get_command_encoders() { + static thread_local std::unordered_map encoders; + return encoders; +} + +std::unordered_map& get_global_command_encoders() { + static std::unordered_map encoders; + return encoders; +} + +NS::SharedPtr new_scoped_memory_pool() { + return NS::TransferPtr(NS::AutoreleasePool::alloc()->init()); +} + +bool is_nax_available() { +#ifdef MLX_METAL_NO_NAX + return false; +#else + auto _check_nax = []() { + bool can_use_nax = false; + if (__builtin_available( + macOS 26.2, iOS 26.2, tvOS 26.2, visionOS 26.2, *)) { + can_use_nax = true; + } + auto& d = metal::device(mlx::core::Device::gpu); + auto arch = d.get_architecture().back(); + auto gen = d.get_architecture_gen(); + can_use_nax &= gen >= (arch == 'p' ? 18 : 17); + return can_use_nax; + }; + static bool is_nax_available_ = _check_nax(); + return is_nax_available_; +#endif +} + +// mlxcel: bridge entry points for the input-budget override, called through +// the cxx bridge (mlxcel_core::set_metal_mb_per_buffer_override). Declared here +// rather than in a header because upstream owns every header in this tree and +// an overlay that adds one would drift on the next pin bump. +void mlxcel_set_mb_per_buffer_override(int mb) { + mlxcel_mb_per_buffer_override_slot().store( + mb > 0 ? mb : 0, std::memory_order_relaxed); +} + +int mlxcel_mb_per_buffer_override() { + return mlxcel_mb_per_buffer_override_slot().load(std::memory_order_relaxed); +} + +} // namespace mlx::core::metal diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index da77371b1..da0754a28 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -33,6 +33,12 @@ namespace mlx::core { void mlxcel_set_qmv_wide(bool enabled); bool mlxcel_qmv_wide(void); } // namespace mlx::core +// Defined by the mlx/backend/metal/device.cpp overlay: the per-command-buffer +// input budget override that mlxcel raises around decode steps. +namespace mlx::core::metal { +void mlxcel_set_mb_per_buffer_override(int mb); +int mlxcel_mb_per_buffer_override(); +} // namespace mlx::core::metal #endif namespace mlx_cxx { @@ -1344,6 +1350,23 @@ bool qmv_wide_enabled() { } #endif +#ifdef MLXCEL_BRIDGE_METAL_BACKEND +void set_metal_mb_per_buffer_override(int32_t mb) { + ::mlx::core::metal::mlxcel_set_mb_per_buffer_override(mb); +} + +int32_t metal_mb_per_buffer_override() { + return ::mlx::core::metal::mlxcel_mb_per_buffer_override(); +} +#else +// No Metal command buffers to size: the override is inert and reads as unset. +void set_metal_mb_per_buffer_override(int32_t) {} + +int32_t metal_mb_per_buffer_override() { + return 0; +} +#endif + std::unique_ptr random_categorical(const MlxArray& logits, int32_t axis) { return std::make_unique(mlx::core::random::categorical(logits.inner, axis)); } diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 79e410679..156c967cc 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -424,6 +424,14 @@ void random_seed(uint64_t seed); void set_qmv_wide(bool enabled); bool qmv_wide_enabled(); +// Override MLX's per-command-buffer input budget (MLX_MAX_MB_PER_BUFFER) for +// work encoded from now on; 0 restores the device default. Lives in the +// mlx/backend/metal/device.cpp overlay. mlxcel raises it around decode steps +// only, because the same budget during prefill multiplies peak memory. Inert, +// and reads back 0, on a build without the Metal backend. +void set_metal_mb_per_buffer_override(int32_t mb); +int32_t metal_mb_per_buffer_override(); + // Random categorical sampling std::unique_ptr random_categorical(const MlxArray& logits, int32_t axis); diff --git a/src/lib/mlxcel-core/src/command_buffer_budget.rs b/src/lib/mlxcel-core/src/command_buffer_budget.rs new file mode 100644 index 000000000..3bb1d45c1 --- /dev/null +++ b/src/lib/mlxcel-core/src/command_buffer_budget.rs @@ -0,0 +1,88 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Decode-only Metal command-buffer input budget. +//! +//! MLX commits a Metal command buffer once the element count of its distinct +//! inputs passes `MLX_MAX_MB_PER_BUFFER << 20`. During decode that cap, at +//! MLX's default, splits every token into one command buffer per one or two +//! layers and idles the GPU at each boundary; during prefill the same cap is +//! what keeps a long prompt's activations from piling up until the buffer +//! completes. [`crate::hardware::decode_mb_per_buffer`] says what the decode +//! budget should be, and [`DecodeCommandBufferBudget`] applies it for the +//! lifetime of a decode loop or decode step and restores the previous value +//! when dropped, so prefill always encodes against the device default. +//! +//! Only pipelined decode benefits, where step n+1 is encoded while the GPU +//! still runs step n. A synchronous step that encodes and then waits should +//! stay on the device default: with one large buffer the GPU cannot start +//! until the whole step is encoded, which on M1 Ultra made server decode both +//! slower and erratic. +//! +//! The override is read by MLX while it encodes, which happens on the thread +//! that calls `eval` / `async_eval`. Enter the guard on that thread, after the +//! prefill work has been encoded (an `async_eval` returns only once its graph +//! is encoded and committed), and keep it alive until the last decode eval. +//! The override is process-wide: work that another thread encodes while a +//! guard is alive also sees the decode budget. That can only move where a +//! command buffer boundary falls, never what is computed. + +use crate::ffi; + +/// RAII guard that raises MLX's per-command-buffer input budget to the decode +/// value while it is alive. A no-op when no decode budget applies (M5+, +/// non-Apple, `MLXCEL_DECODE_MB_PER_BUFFER=0`, or an operator-set +/// `MLX_MAX_MB_PER_BUFFER`). Guards nest: each restores the value it found. +#[must_use = "the budget is restored as soon as the guard is dropped"] +pub struct DecodeCommandBufferBudget { + previous: Option, +} + +impl DecodeCommandBufferBudget { + /// Apply the process's decode budget + /// ([`crate::hardware::decode_mb_per_buffer`]). + pub fn enter() -> Self { + Self::with_budget(crate::hardware::decode_mb_per_buffer()) + } + + /// Apply an explicit budget; `None` leaves the current value untouched. + pub fn with_budget(budget: Option) -> Self { + let Some(mb) = budget.and_then(|mb| i32::try_from(mb).ok()) else { + return Self { previous: None }; + }; + let previous = ffi::metal_mb_per_buffer_override(); + ffi::set_metal_mb_per_buffer_override(mb); + Self { + previous: Some(previous), + } + } + + /// Whether this guard changed the budget (and will restore it on drop). + #[must_use] + pub fn is_active(&self) -> bool { + self.previous.is_some() + } +} + +impl Drop for DecodeCommandBufferBudget { + fn drop(&mut self) { + if let Some(previous) = self.previous { + ffi::set_metal_mb_per_buffer_override(previous); + } + } +} + +#[cfg(test)] +#[path = "command_buffer_budget_tests.rs"] +mod tests; diff --git a/src/lib/mlxcel-core/src/command_buffer_budget_tests.rs b/src/lib/mlxcel-core/src/command_buffer_budget_tests.rs new file mode 100644 index 000000000..611c26cb7 --- /dev/null +++ b/src/lib/mlxcel-core/src/command_buffer_budget_tests.rs @@ -0,0 +1,66 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::*; +use std::sync::Mutex; + +// The override is one process-wide slot, so the tests that write it must not +// interleave with each other. +static SLOT: Mutex<()> = Mutex::new(()); + +#[cfg(feature = "metal")] +#[test] +fn override_round_trips_through_the_device_overlay() { + let _lock = SLOT.lock().unwrap_or_else(|e| e.into_inner()); + ffi::set_metal_mb_per_buffer_override(777); + assert_eq!(ffi::metal_mb_per_buffer_override(), 777); + // Non-positive values mean "no override", not a zero budget: a zero cap + // would commit a buffer after every op. + ffi::set_metal_mb_per_buffer_override(-5); + assert_eq!(ffi::metal_mb_per_buffer_override(), 0); + ffi::set_metal_mb_per_buffer_override(0); + assert_eq!(ffi::metal_mb_per_buffer_override(), 0); +} + +#[cfg(feature = "metal")] +#[test] +fn guard_applies_the_budget_and_restores_what_it_found() { + let _lock = SLOT.lock().unwrap_or_else(|e| e.into_inner()); + ffi::set_metal_mb_per_buffer_override(0); + { + let outer = DecodeCommandBufferBudget::with_budget(Some(1000)); + assert!(outer.is_active()); + assert_eq!(ffi::metal_mb_per_buffer_override(), 1000); + { + let inner = DecodeCommandBufferBudget::with_budget(Some(400)); + assert!(inner.is_active()); + assert_eq!(ffi::metal_mb_per_buffer_override(), 400); + } + assert_eq!(ffi::metal_mb_per_buffer_override(), 1000); + } + assert_eq!(ffi::metal_mb_per_buffer_override(), 0); +} + +#[test] +fn guard_without_a_budget_leaves_the_slot_alone() { + let _lock = SLOT.lock().unwrap_or_else(|e| e.into_inner()); + ffi::set_metal_mb_per_buffer_override(0); + let before = ffi::metal_mb_per_buffer_override(); + { + let guard = DecodeCommandBufferBudget::with_budget(None); + assert!(!guard.is_active()); + assert_eq!(ffi::metal_mb_per_buffer_override(), before); + } + assert_eq!(ffi::metal_mb_per_buffer_override(), before); +} diff --git a/src/lib/mlxcel-core/src/generate.rs b/src/lib/mlxcel-core/src/generate.rs index d35d208e2..c1af83fba 100644 --- a/src/lib/mlxcel-core/src/generate.rs +++ b/src/lib/mlxcel-core/src/generate.rs @@ -1795,6 +1795,11 @@ impl CxxGenerator { }; ffi::async_eval(&y); self.prepare_turbo4_delegated_before_decode(max_tokens); + // Prefill is encoded by now; raise the command-buffer input budget for + // the decode loop only (see `DecodeCommandBufferBudget`). Not under + // MLXCEL_FORCE_SYNC: a synchronous eval gains nothing from larger + // buffers and loses the CPU-encode / GPU-execute overlap inside a step. + let _decode_budget = (!force_sync).then(crate::DecodeCommandBufferBudget::enter); // Main generation loop - matches Python exactly: // 1. Start next step computation @@ -2116,6 +2121,9 @@ impl CxxGenerator { }; ffi::async_eval(&y); self.prepare_turbo4_delegated_before_decode(max_tokens); + // Prefill is encoded by now; raise the command-buffer input budget for + // the decode loop only (see `DecodeCommandBufferBudget`). + let _decode_budget = crate::DecodeCommandBufferBudget::enter(); // Decode loop — identical to standard generation (no embeddings needed) let mut n = 0; @@ -2294,6 +2302,9 @@ impl CxxGenerator { let ttft_eval_ns = ttft_eval_start.map_or(0, |t| t.elapsed().as_nanos()); let ttft_post_start = profile_ttft.then(Instant::now); self.prepare_turbo4_delegated_before_decode(max_tokens); + // Prefill is encoded by now; raise the command-buffer input budget for + // the decode loop only (see `DecodeCommandBufferBudget`). + let _decode_budget = crate::DecodeCommandBufferBudget::enter(); let ttft_post_ns = ttft_post_start.map_or(0, |t| t.elapsed().as_nanos()); let prefill_time = prefill_start.elapsed(); if profile_ttft { @@ -2566,6 +2577,9 @@ impl CxxGenerator { let ttft_eval_ns = ttft_eval_start.map_or(0, |t| t.elapsed().as_nanos()); let ttft_post_start = profile_ttft.then(Instant::now); self.prepare_turbo4_delegated_before_decode(max_tokens); + // Prefill is encoded by now; raise the command-buffer input budget for + // the decode loop only (see `DecodeCommandBufferBudget`). + let _decode_budget = crate::DecodeCommandBufferBudget::enter(); let ttft_post_ns = ttft_post_start.map_or(0, |t| t.elapsed().as_nanos()); let prefill_time = prefill_start.elapsed(); if profile_ttft { diff --git a/src/lib/mlxcel-core/src/hardware.rs b/src/lib/mlxcel-core/src/hardware.rs index 930e708d0..1d43cabf1 100644 --- a/src/lib/mlxcel-core/src/hardware.rs +++ b/src/lib/mlxcel-core/src/hardware.rs @@ -386,6 +386,104 @@ pub fn apply_metal_ops_per_buffer_default() { } } +/// The per-command-buffer input budget to use during decode steps on an Apple +/// Silicon class, or `None` to leave decode on MLX's device default. +/// +/// MLX commits a Metal command buffer when either its op count passes +/// `MLX_MAX_OPS_PER_BUFFER` or its input budget passes `MLX_MAX_MB_PER_BUFFER` +/// (`command_buffer_needs_commit` in `mlx/backend/metal/device.cpp`). The budget +/// is not bytes: MLX sums `array::data_size()`, an element count, over the +/// distinct input arrays and compares `count >> 20` against the cap, which it +/// defaults to 40-50 per device class. Once [`metal_ops_per_buffer_default`] +/// lifts the op cap to 1000, the input budget is the one that binds during +/// decode, because every token reads the whole weight set: a 4-bit 7B model +/// packs its weights into about 1.1G `u32` elements, so the default 50 commits a +/// buffer every one to two layers (about 23 per token on command-r7b), and each +/// commit leaves the GPU idle for tens of microseconds before the next buffer +/// starts. A bf16 checkpoint counts eight times as many elements per layer and +/// commits more often still. +/// +/// Measured on M1 Ultra, 500-token prompt, 128 generated tokens, three +/// interleaved runs per cell, decode at MLX's default versus 1000: command-r7b +/// 4-bit +7%, Llama 3.1 8B 4-bit +5.6%, Qwen2.5 7B 4-bit +8%, Gemma 3n E4B +3%, +/// Granite 4.0 H Tiny +10%, Qwen3-30B-A3B +20%, Mixtral 8x7B +21%, Llama 3.1 +/// 8B bf16 +17%, Gemma 3 4B flat. See +/// docs/benchmark_results/metal-mb-per-buffer-m1ultra-2026-09-21.md. +/// +/// The value applies to decode only. The same budget during prefill keeps a +/// whole prompt's activations alive until the buffer completes: at 1000 the +/// peak for a 2048-token prompt went from 6.0 to 12.6 GB on Qwen2.5 7B and from +/// 19.8 to 36.2 GB on Qwen3-30B-A3B, and prefill throughput lost up to 2.7%. +/// [`crate::DecodeCommandBufferBudget`] applies it around pipelined decode only +/// (the generate loops and the server's lookahead decode) and leaves prefill +/// and synchronous decode steps on the device default: a step that encodes +/// and then waits loses its CPU-encode / GPU-execute overlap in one large +/// buffer. +/// +/// Gated like [`metal_ops_per_buffer_default`]: the input budget only binds +/// once the op cap is raised, which happens only on M1 through M4. M5+ was not +/// measured and keeps MLX's default in both phases. +#[must_use] +pub fn metal_decode_mb_per_buffer_default( + r#gen: AppleSiliconGen, + has_neural_accelerator: bool, +) -> Option { + metal_ops_per_buffer_default(r#gen, has_neural_accelerator).map(|_| 1000) +} + +/// Environment variable that sets the decode-step input budget explicitly. +/// `0`, `off`, `false` or `no` disables the decode switch; a positive integer +/// replaces the hardware default. Ignored when `MLX_MAX_MB_PER_BUFFER` is set. +pub const DECODE_MB_PER_BUFFER_ENV: &str = "MLXCEL_DECODE_MB_PER_BUFFER"; + +/// Resolve the decode-step input budget from the two environment variables +/// and the hardware default. Pure, so the precedence is unit-testable. +/// +/// An operator-set `MLX_MAX_MB_PER_BUFFER` wins over everything: it pins the +/// budget for prefill and decode alike, and a decode-only switch on top of it +/// would silently override the value the operator chose. Otherwise +/// `MLXCEL_DECODE_MB_PER_BUFFER` decides, and an unparseable value falls back +/// to the hardware default. +#[must_use] +pub fn resolve_decode_mb_per_buffer( + mlx_max_mb_per_buffer: Option<&str>, + decode_override: Option<&str>, + hardware_default: Option, +) -> Option { + if mlx_max_mb_per_buffer.is_some() { + return None; + } + let Some(value) = decode_override.map(str::trim).filter(|v| !v.is_empty()) else { + return hardware_default; + }; + if matches!( + value.to_ascii_lowercase().as_str(), + "0" | "off" | "false" | "no" + ) { + return None; + } + value + .parse::() + .ok() + .filter(|n| *n > 0) + .or(hardware_default) +} + +/// The decode-step input budget for this process, resolved once. See +/// [`metal_decode_mb_per_buffer_default`] and [`resolve_decode_mb_per_buffer`]. +#[must_use] +pub fn decode_mb_per_buffer() -> Option { + static CACHED: OnceLock> = OnceLock::new(); + *CACHED.get_or_init(|| { + let hw = get_hardware(); + resolve_decode_mb_per_buffer( + std::env::var("MLX_MAX_MB_PER_BUFFER").ok().as_deref(), + std::env::var(DECODE_MB_PER_BUFFER_ENV).ok().as_deref(), + metal_decode_mb_per_buffer_default(hw.silicon_gen, hw.has_neural_accelerator), + ) + }) +} + /// The `MLX_CUDA_GRAPH_CACHE_SIZE` default to apply on a CUDA build, or `None` /// off CUDA. /// @@ -1154,6 +1252,78 @@ mod tests { ); } + #[test] + fn decode_mb_per_buffer_default_follows_the_ops_default_gate() { + // The input budget only binds once the op cap is raised, so the two + // defaults must apply on exactly the same hardware. Decoupling them + // would either leave M1-M4 decode committing a buffer every layer or + // raise the budget where nothing was measured (M5+, non-Apple). + for (r#gen, na) in [ + (AppleSiliconGen::M1, false), + (AppleSiliconGen::M2, false), + (AppleSiliconGen::M3, false), + (AppleSiliconGen::M4, false), + (AppleSiliconGen::M5, true), + (AppleSiliconGen::Unknown, false), + ] { + let ops = metal_ops_per_buffer_default(r#gen, na); + let mb = metal_decode_mb_per_buffer_default(r#gen, na); + assert_eq!(ops.is_some(), mb.is_some(), "{gen:?} na={na}"); + } + assert_eq!( + metal_decode_mb_per_buffer_default(AppleSiliconGen::M1, false), + Some(1000) + ); + } + + #[test] + fn operator_mlx_max_mb_per_buffer_disables_the_decode_switch() { + // An explicit MLX_MAX_MB_PER_BUFFER pins both phases; the decode switch + // must not override it, whatever MLXCEL_DECODE_MB_PER_BUFFER says. + assert_eq!( + resolve_decode_mb_per_buffer(Some("50"), None, Some(1000)), + None + ); + assert_eq!( + resolve_decode_mb_per_buffer(Some("50"), Some("2000"), Some(1000)), + None + ); + } + + #[test] + fn decode_mb_per_buffer_env_overrides_and_disables() { + assert_eq!( + resolve_decode_mb_per_buffer(None, None, Some(1000)), + Some(1000) + ); + assert_eq!(resolve_decode_mb_per_buffer(None, None, None), None); + assert_eq!( + resolve_decode_mb_per_buffer(None, Some("400"), Some(1000)), + Some(400) + ); + // An explicit value applies even where the hardware default is off. + assert_eq!( + resolve_decode_mb_per_buffer(None, Some("400"), None), + Some(400) + ); + for off in ["0", "off", "OFF", "false", "no", " 0 "] { + assert_eq!( + resolve_decode_mb_per_buffer(None, Some(off), Some(1000)), + None, + "{off}" + ); + } + // Unparseable or empty falls back to the hardware default. + assert_eq!( + resolve_decode_mb_per_buffer(None, Some("lots"), Some(1000)), + Some(1000) + ); + assert_eq!( + resolve_decode_mb_per_buffer(None, Some(""), Some(1000)), + Some(1000) + ); + } + #[test] fn cuda_sdpa_cache_default_matches_build_feature() { #[cfg(feature = "cuda")] diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index e4b881e04..a2edee5c6 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -565,6 +565,17 @@ mod ffi { /// without the Metal backend. fn qmv_wide_enabled() -> bool; + /// Override MLX's per-command-buffer input budget + /// (`MLX_MAX_MB_PER_BUFFER`, an element count >> 20) for work encoded + /// from now on; `0` restores the device default. Process-wide and + /// immediate. Use [`crate::DecodeCommandBufferBudget`] rather than + /// calling this directly. Inert on a build without the Metal backend. + fn set_metal_mb_per_buffer_override(mb: i32); + + /// Current override set by [`set_metal_mb_per_buffer_override`], `0` + /// when none is active (always `0` without the Metal backend). + fn metal_mb_per_buffer_override() -> i32; + /// Random categorical sampling fn random_categorical(logits: &MlxArray, axis: i32) -> UniquePtr; @@ -3547,6 +3558,11 @@ pub mod cuda_graph_budget; // Public so that mlxcel (the main crate) can log hardware info at startup. pub mod hardware; +// Decode-only raise of MLX's per-command-buffer input budget; prefill keeps +// the device default so long prompts do not hold their activations longer. +pub mod command_buffer_budget; +pub use command_buffer_budget::DecodeCommandBufferBudget; + // Typed wrappers around MLX's runtime memory accounting APIs (issue #55). // Public so that the CLI generate path can surface post-load resident // memory and the preflight (#56) can call `set_memory_limit` to fail fast. diff --git a/src/main.rs b/src/main.rs index 2adea6c8a..400dcb6fb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -60,6 +60,11 @@ Environment Variables: unset, auto-defaults to 1000 on pre-M5 Apple Silicon (M1-M4), MLX default on M5+ and non-Apple (hardware-gated, #353) explicit value always wins (manual override / sweeps) + MLXCEL_DECODE_MB_PER_BUFFER Metal command-buffer input budget, decode steps only + unset, 1000 on pre-M5 Apple Silicon (M1-M4), off elsewhere; + prefill keeps MLX's default to bound peak memory + \"0\"/\"off\", disable; a positive value replaces the default + ignored when MLX_MAX_MB_PER_BUFFER is set (pins both phases) MLX_CUDA_GRAPH_CACHE_SIZE CUDA graph-cache LRU capacity (CUDA only) unset, auto-defaults to 2000 on CUDA builds (MLX default 400 aborts long-lived shape-diverse decode, #818) diff --git a/src/server/batch/scheduler/decode_tick.rs b/src/server/batch/scheduler/decode_tick.rs index cfd74524b..c66759dca 100644 --- a/src/server/batch/scheduler/decode_tick.rs +++ b/src/server/batch/scheduler/decode_tick.rs @@ -288,8 +288,15 @@ impl BatchScheduler { pub(super) fn run_decode_tick(&mut self, seq_ids: &[SequenceId]) { let params = self.lookahead_params(seq_ids); + // The raised command-buffer input budget (`DecodeCommandBufferBudget`) + // is applied only around pipelined work, where step n+1 is encoded + // while the GPU still runs step n. A synchronous step encodes and then + // waits, so one large buffer there removes the CPU-encode / GPU-execute + // overlap inside the step: on M1 Ultra, command-r7b sync decode went + // from a steady 98-102 tok/s to 68-100 with the budget raised. match self.decode_lookahead.take() { Some(la) if la.ids == seq_ids && params.is_some() && self.lookahead_safe() => { + let _decode_budget = mlxcel_core::DecodeCommandBufferBudget::enter(); self.pipelined_steady_decode(la, seq_ids, ¶ms.unwrap()); } Some(la) => { @@ -299,10 +306,12 @@ impl BatchScheduler { self.apply_lookahead_trim(&la.ids, lookahead_teardown_positions(false)); drop(la); self.dispatch_sync_decode(seq_ids); + let _decode_budget = mlxcel_core::DecodeCommandBufferBudget::enter(); self.maybe_prime_lookahead(seq_ids); } None => { self.dispatch_sync_decode(seq_ids); + let _decode_budget = mlxcel_core::DecodeCommandBufferBudget::enter(); self.maybe_prime_lookahead(seq_ids); } }