Memory Estimation¶

mlxcel can estimate a model load before it touches the MLX runtime. Use this to size a model, context length, batch shape, and KV-cache mode before running generation or starting a server.

The estimator is intentionally pre-load and read-only. It reads config.json, safetensors metadata, environment limits, and Linux memory information when available; it does not allocate tensors, initialize MLX, or load model weights.

What Is Counted¶

The estimate has these components:

Component	Source
Weights	Exact safetensors footprint when possible. Sharded indexes use validated `metadata.total_size`; stale indexes fall back to shard/header inspection. Single-file safetensors are read from the binary header without loading tensor data.
KV cache	`config.json` architecture fields: layers, KV heads, `head_dim`, element width, context length, and batch, with the architecture taken into account (sliding-window, MLA, hybrid and pure-SSM layouts are sized by their own cache shape, not the flat per-layer formula). Context is rounded up to the next 256-token allocation block.
Allocator overhead	`MLXCEL_HEADROOM_FACTOR` (default `1.20`) applied to `weights + KV cache`, covering MLX allocator and graph working set. The factor scales only this term.
Activation reserve	`MLXCEL_ACTIVATION_MULT` (default `2.0`) times `batch x min(context, prefill chunk) x (hidden + intermediate) x 2 bytes`, plus the last-token logit buffer `batch x vocab x 2 bytes`. The prefill chunk is `MLXCEL_PREFILL_CHUNK`, or 2048 when unset.
Backend in-flight	ROCm builds only: the `MLX_ROCM_MAX_INFLIGHT_MB` budget (default 1024 MiB) that committed but unfinished command batches may hold. Every other build reserves `0`.
Host shadows	Windows ROCm builds only: pinned host copies the allocator keeps for GPU buffers the CPU touches, estimated per sequence as one f32 logits row (`vocab x 4`) plus the int32 token ids of the context (`context x 4`). They draw from host RAM, so they are compared against host RAM separately and are not added to the total; `fits` requires both comparisons to pass. Every other build reserves `0`.
Available memory	Apple unified memory, Linux `/proc/meminfo`, Windows `GlobalMemoryStatusEx`, or an explicit `MLXCEL_MEMORY_LIMIT` cap when set.

The runtime headroom reported by inspect is the sum of the allocator overhead, the activation reserve, and the backend in-flight reserve; the total is weights + KV cache + runtime headroom. The server uses the same estimator to size the paged KV --kv-cache-budget auto budget and the model-aware prompt-cache capacity defaults.

The KV estimate honors INT8 KV cache settings from --kv-cache-mode int8 or split --cache-type-k int8 --cache-type-v int8. TurboQuant modes are reported through the same flag surface, but the current estimator only distinguishes the INT8 byte-width case from the default FP16-equivalent KV cache.

Inspect a Model¶

mlxcel inspect prints the same byte breakdown used by the preflight checks, then exits without loading the model.

mlxcel inspect --model ./models/Llama-3.1-8B-Instruct-4bit

Common sizing overrides:

# Estimate a longer context window.
mlxcel inspect -m ./models/qwen3-8b --max-tokens 32768

# Estimate concurrent server slots.
mlxcel inspect -m ./models/qwen3-8b --max-tokens 8192 --batch 4

# Estimate INT8 KV cache.
mlxcel inspect -m ./models/qwen3-8b --kv-cache-mode int8

--quant fp16|int8|int4 labels the intended weight mode in the output. It does not rewrite the weight-byte total because mlxcel takes the safetensors header at face value and quantized checkpoints already encode their real on-disk tensor dtype and shape.

JSON Report¶

Use --json when a recipe builder, scheduler, or CI check needs machine-readable fit data:

mlxcel inspect --json -m mlx-community/Qwen3-4B-4bit

The command writes exactly one JSON object to stdout on success. Resolution, download, and estimator errors still use the existing stderr path and exit nonzero, so a pipeline can treat stdout as parseable JSON only after the command succeeds.

The report fields are raw bytes unless noted otherwise:

Field	Meaning
`mlxcel_version`, `model`, `model_type`, `family`	Binary version, resolved model path, raw top-level `config.json.model_type`, and the registry family slug used by architecture metadata.
`inputs`	Effective `max_tokens`, `batch`, `kv_cache_mode`, and `quant` labels used for the estimate.
`weights_bytes`, `weights_source`	Weight footprint and whether it came from a safetensors header, analytical config estimate, or fallback.
`kv_bytes_per_token`	FP16 and INT8 per-token KV rates when model geometry is available; `turbo4` is `null` until the estimator models TurboQuant storage widths directly.
`kv_bytes_total`, `kv_detail`	Total KV allocation for the requested context and batch, plus the architecture detail string shown in the text banner.
`per_slot_overhead_bytes`	Paged-decode workspace reserve for default paged families; `null` for dense-default or unavailable geometry.
`activation_bytes`, `headroom_bytes`, `headroom_factor`	Activation reserve, total runtime headroom (allocator overhead plus activation reserve plus backend in-flight reserve), and the multiplier used for allocator overhead.
`backend_inflight_bytes`	Backend in-flight reserve included in `headroom_bytes`. Non-zero only on ROCm builds, where it equals `MLX_ROCM_MAX_INFLIGHT_MB`; the text banner prints it as a `Backend in-flight` line.
`host_shadow_bytes`, `host_ram_available_bytes`, `host_shadow_fits`	Windows ROCm host-shadow reserve, the host RAM it is compared against, and whether it fits. `0`, `0` (or the RAM figure) and `true` on other builds.
`budget_bytes`, `total_bytes`, `fits`	Available-memory budget, total estimate, and the same fit boolean used by preflight checks. `fits` is false when either the device budget or the host-shadow comparison fails.

Preflight Generation¶

mlxcel generate --estimate-memory runs after prompt rendering/tokenization and before model load. It sizes the KV cache from:

rendered prompt tokens + --max-tokens

If the estimated total exceeds available memory, generation aborts before loading the model. Override only after checking the breakdown:

mlxcel generate \
  --model ./models/qwen3-8b \
  --prompt "Summarize this deployment plan." \
  --max-tokens 1024 \
  --estimate-memory

mlxcel generate ... --estimate-memory --force

--no-memory-check is an alias for --force.

Preflight Server Startup¶

mlxcel serve --estimate-memory and mlxcel-server --estimate-memory run before worker startup. They use the same server context semantics as runtime:

effective_context = floor(--ctx-size / active_slots)
active_slots      = --max-batch-size when set, otherwise --n-parallel

With --no-batch, active_slots is 1. If --ctx-size is 0, the estimator uses its default 8192-token context length because the model-default context is not known until load time.

mlxcel serve \
  --model ./models/qwen3-8b \
  --ctx-size 32768 \
  --n-parallel 4 \
  --estimate-memory

For this example, the server preflight sizes each active slot with an 8192-token context and a batch of 4. /health.context_size and /slots[].context_size report the same effective per-slot context window once the server is running.

Environment Variables¶

Variable	Default	Effect
`MLXCEL_MEMORY_LIMIT`	unset	Soft allocator memory cap. Nonzero values become the authoritative available-memory budget for `inspect`, `--estimate-memory`, and MLX runtime setup. Supports values such as `48GB`, `512MB`, or raw bytes.
`MLXCEL_HEADROOM_FACTOR`	`1.20`	Allocator-overhead multiplier on `weights + KV cache`. It does not scale the activation reserve or the backend in-flight reserve. Positive values `<= 1.0` disable the allocator-overhead term; invalid or non-positive values warn and fall back to the default.
`MLXCEL_ACTIVATION_MULT`	`2.0`	Multiplier in the activation reserve `batch x min(context, prefill chunk) x (hidden + intermediate) x 2 bytes`. Invalid or non-positive values warn and fall back to the default.
`MLX_ROCM_MAX_INFLIGHT_MB`	`1024`	ROCm builds only. Backend in-flight budget in MiB, reserved by the estimator as `backend_inflight_bytes`. `0` removes the backend bound, and the estimator then reserves nothing.

Use MLXCEL_MEMORY_LIMIT for container, CI, or shared-host deployments where the OS-reported memory is not the real budget available to the process.

Relationship to Quantization Advice¶

--recommend-quant uses the same weight and KV-cache estimator as inspect and memory preflight. That keeps the recommendation, printed memory table, and startup refusal condition aligned for the same model directory and KV-cache mode.