Skip to content

chore(execution): parse memory sizes with one shared grammar - #1573

Merged
inureyes merged 2 commits into
mainfrom
refactor/issue-1317-one-memory-size-grammar
Sep 2, 2026
Merged

inureyes merged 2 commits into
mainfrom
refactor/issue-1317-one-memory-size-grammar

Conversation

@inureyes

@inureyes inureyes commented Sep 2, 2026 •

Copy link
Copy Markdown
Member

Summary

MLXCEL_MEMORY_LIMIT was read by two parsers with two grammars, so MLXCEL_MEMORY_LIMIT=4G capped the MLX allocator and was silently ignored by the memory-estimation preflight, which then reported availability from the machine's total memory. This keeps one parser in src/execution/runtime.rs, exposes it to src/execution/memory_estimate.rs, and pins the grammar with tests on both sides.

What changed

  • src/execution/runtime.rs: parse_memory_size becomes pub(crate) fn parse_memory_size(&str) -> Option<u64> and is the single grammar for MLXCEL_MEMORY_LIMIT, MLXCEL_WIRED_LIMIT and MLXCEL_CACHE_LIMIT. It adds K/KB beside the existing M/MB and G/GB, rejects a negative, NaN or infinite numeric part instead of casting it to zero, floors a fractional value, and states the overflow saturation explicitly rather than leaning on the implicit saturating float-to-int cast that both old parsers relied on. Every spelling that parsed before parses to the same value.
  • src/execution/runtime.rs: the three resolvers keep their own 0 / none / empty (and max) handling and narrow to usize at the MLX setter boundary through a new clamp_to_usize, so RuntimeSetup and the mlxcel_core setters see the types they saw before.
  • src/execution/memory_estimate.rs: parse_scaled_memory_size is deleted and parse_optional_memory_size_bytes keeps only its unset check before delegating to crate::execution::runtime::parse_memory_size(...).filter(|b| *b > 0). The preflight and the allocator cap now resolve one string to one number.
  • src/execution/runtime_tests.rs: the five existing parse_memory_size_* tests stay as regression pins for the old spellings, and parse_memory_size_accepts_every_suffix_spelling, parse_memory_size_fractional_is_exact_floor, parse_memory_size_rejects_garbage and parse_memory_size_saturates_instead_of_wrapping are added beside them.
  • src/execution/memory_estimate.rs tests: available_memory_honors_short_suffix_env_limit drives MLXCEL_MEMORY_LIMIT=512M through estimate_total_memory and asserts the same 512 MiB its 512MB sibling asserts, and parse_optional_memory_size_accepts_the_runtime_grammar pins 4G == 4GB == 4gb.
  • docs/environment-variables.md: the three size-valued rows list every accepted suffix, and a paragraph under the table states the shared grammar once (binary units, fraction allowed on a suffixed value, floored, integer-only for a bare byte count).
  • src/main.rs: the --help environment block now reads GB/G, MB/M, KB/K, or bytes for MLXCEL_WIRED_LIMIT and MLXCEL_MEMORY_LIMIT.

Blast radius

parse_memory_size also governs MLXCEL_WIRED_LIMIT and MLXCEL_CACHE_LIMIT, so the grammar and return type here move allocator caps for every run. The old spellings are pinned by the pre-existing tests, which are unchanged apart from the usize to u64 literal in parse_memory_size_fractional_gb. One outcome changes, and only for garbage that used to cast to zero: a negative or non-finite suffixed value (-1GB, NaNGB, infGB) now returns None rather than Some(0). For MLXCEL_MEMORY_LIMIT and MLXCEL_CACHE_LIMIT that is the same end state (both mapped a zero to unset already); for MLXCEL_WIRED_LIMIT it means -1GB falls back to gpu_max_memory_size() instead of silently disabling the wired limit, which is what MLXCEL_WIRED_LIMIT=abc always did. Neither parser ever wrapped: Rust's float-to-int as cast has saturated since 1.45, so 1e30GB resolved to the maximum before this change too. resolve_paged_slab_blocks in memory_estimate.rs is untouched (issue #1137 owns it).

Test plan

  • cargo test --profile test-fast --features metal,accelerate --lib execution:: passes: 126 passed, 0 failed, including all 13 parser and preflight tests named above.
  • cargo clippy --profile test-fast --lib --tests --features metal,accelerate -- -D warnings passes.
  • cargo fmt --all -- --check passes.
  • python3 scripts/ci/check_cross_repo_refs.py passes.
  • Real-binary acceptance run, below, is left to the orchestrator because it needs a release rebuild.

Acceptance commands for the release binary

Note that mlxcel inspect takes the model through -m/--model, not positionally, and the checkpoint present locally is models/mlx/qwen3-0.6b-4bit; the issue body's mlxcel inspect models/qwen3-0.6b-4bit is wrong on both counts.

cargo build --release --features metal,accelerate
MLXCEL_MEMORY_LIMIT=4G  ./target/release/mlxcel inspect -m models/mlx/qwen3-0.6b-4bit | grep -E 'Available:|Total estimate'
MLXCEL_MEMORY_LIMIT=4GB ./target/release/mlxcel inspect -m models/mlx/qwen3-0.6b-4bit | grep -E 'Available:|Total estimate'
MLXCEL_MEMORY_LIMIT=4096M ./target/release/mlxcel inspect -m models/mlx/qwen3-0.6b-4bit | grep -E 'Available:'
./target/release/mlxcel inspect -m models/mlx/qwen3-0.6b-4bit | grep -E 'Available:'

Expected: the first three print an identical Available: line reading 4.00 GB, and the fourth prints the machine figure (the host's unified memory, so far larger). On main the 4G and 4096M invocations print the machine figure instead, which is the defect.

The allocator side, to confirm both readers agree on one string:

MLXCEL_MEMORY_LIMIT=4G ./target/release/mlxcel generate -m models/mlx/qwen3-0.6b-4bit --no-chat-template -p "Hello" -n 8

Expected: a MLX allocator memory limit: 4.0 GB (MLXCEL_MEMORY_LIMIT) line at startup and normal generation. On main that line already reads 4.0 GB for 4G, which is exactly the half that the preflight disagreed with.

Closes #1317

@inureyes inureyes added status:review Under review type:chore Maintenance tasks (build, CI, etc.) priority:low Low priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:cli Command-line interface / CLI flags labels Sep 2, 2026
`MLXCEL_MEMORY_LIMIT` was read by two parsers with two grammars. `parse_memory_size` in `src/execution/runtime.rs` served the allocator cap and accepted `4G`, `4GB`, `512M`, `512MB` and plain bytes; the memory-estimation preflight in `src/execution/memory_estimate.rs` carried its own `parse_optional_memory_size_bytes` plus `parse_scaled_memory_size`, which took `GB` and `MB` only. So `MLXCEL_MEMORY_LIMIT=4G` capped the allocator and was silently dropped by `mlxcel inspect` and `--estimate-memory`, which then reported availability from the machine's total memory instead. `inspect` runs before runtime bring-up, so the MLX-limit fallback is zero there too and nothing caught the divergence.

`parse_memory_size` is now the one grammar and is `pub(crate)`, so the preflight resolves a string to the same number the allocator cap will apply. It returns `u64`, adds `K`/`KB` alongside the existing `M`/`MB` and `G`/`GB`, rejects a negative, `NaN` or infinite numeric part instead of casting it to zero, and states the overflow saturation explicitly rather than leaning on the implicit saturating float-to-int cast that both old parsers relied on. Every spelling that parsed before still parses to the same value; the three resolvers keep their own `0`/`none`/empty (and `max`) handling and narrow to `usize` at the MLX setter boundary through `clamp_to_usize`. One outcome does change: a garbage suffixed value such as `MLXCEL_WIRED_LIMIT=-1GB` used to cast to zero and silently disable the wired limit, and now falls back to `gpu_max_memory_size()` the way `MLXCEL_WIRED_LIMIT=abc` always did. The preflight's `parse_optional_memory_size_bytes` keeps only its unset check and delegates the rest, and `parse_scaled_memory_size` is gone.

The blast radius is the reason the old spellings are pinned first: `parse_memory_size` also governs `MLXCEL_WIRED_LIMIT` and `MLXCEL_CACHE_LIMIT`, so a grammar or return-type change moves allocator caps for every run. `parse_memory_size_gb`, `_mb`, `_bytes`, `_fractional_gb` and `_invalid` stay as they were, and `parse_memory_size_accepts_every_suffix_spelling`, `_fractional_is_exact_floor`, `_rejects_garbage` and `_saturates_instead_of_wrapping` are added beside them. On the preflight side, `available_memory_honors_short_suffix_env_limit` drives `MLXCEL_MEMORY_LIMIT=512M` through `estimate_total_memory` and asserts the same 512 MiB the `512MB` sibling asserts, and `parse_optional_memory_size_accepts_the_runtime_grammar` pins `4G` == `4GB` == `4gb`.

Documentation follows the code: the three size-valued rows in `docs/environment-variables.md` list every accepted suffix, a new paragraph under the table states the shared grammar once (binary units, fractional allowed on a suffixed value, floored, integer-only for a bare byte count), and the `mlxcel --help` environment block now reads `GB/G, MB/M, KB/K, or bytes`.

Validated with `cargo test --profile test-fast --features metal,accelerate --lib execution::` (126 passed, 0 failed), `cargo clippy --profile test-fast --lib --tests --features metal,accelerate -- -D warnings`, and `cargo fmt --all -- --check`.

Refs #1317
@inureyes
inureyes force-pushed the refactor/issue-1317-one-memory-size-grammar branch from 3a396d0 to 078cf8a Compare September 2, 2026 00:35
@inureyes inureyes added status:done Completed and removed status:review Under review labels Sep 2, 2026
@inureyes
inureyes merged commit 4d71eba into main Sep 2, 2026
13 checks passed
@inureyes
inureyes deleted the refactor/issue-1317-one-memory-size-grammar branch September 2, 2026 00:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:cli Command-line interface / CLI flags area:core mlxcel-core: MLX FFI, primitives, KV cache, layers priority:low Low priority status:done Completed type:chore Maintenance tasks (build, CI, etc.)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

chore(execution): parse MLXCEL_MEMORY_LIMIT with one size grammar for the allocator cap and the preflight

1 participant