Skip to content

safetensors: accept the VMLX_ spelling of the mmap variables - #4

Open
rcfa wants to merge 1 commit into
osaurus-ai:qwen4-expfrom
rcfa:fix/vmlx-env-prefix-safetensors
Open

safetensors: accept the VMLX_ spelling of the mmap variables#4
rcfa wants to merge 1 commit into
osaurus-ai:qwen4-expfrom
rcfa:fix/vmlx-env-prefix-safetensors

Conversation

@rcfa

@rcfa rcfa commented Aug 28, 2026

Copy link
Copy Markdown

mmap_tensor_buffers_enabled here already reads both VMLINUX_ and VMLX_,
and VMLX_QMV_MR in the Metal backend uses the current prefix. The other five
mmap variables were left on the legacy spelling alone, so which name works
depends on which knob you reach for.

All three spellings are accepted after this; nothing that works today stops
working.

Why it isn't only consistency

The Swift side has already moved. vmlx-swift's ModelFactory does:

setenv("VMLX_MMAP_SAFETENSORS", "1", 1)

to force mmap for a load — and nothing here reads that name:

return env_truthy("MLX_SAFETENSORS_MMAP") ||
    env_truthy("VMLINUX_MMAP_SAFETENSORS");

So that write reaches nothing. mmap keeps working only because
MLX_SAFETENSORS_MMAP is set beside it. For this one variable the new name
is the weaker one — a user who sets VMLX_MMAP_SAFETENSORS gets less than one
who sets the legacy spelling.

Booleans and values are not the same problem

For the ||-chained flags, precedence is unobservable, so the new spelling is
simply added — the shape mmap_tensor_buffers_enabled already uses.

mmap_start_cold_pct and mmap_cold_advice return a value, so order is
observable. There the current spelling is consulted first; appending it after
the legacy one would mean that setting both silently yields the old value.

MLX_SAFETENSORS_MMAP_COLD_ADVICEVMLX_…  →  VMLINUX_…  →  default

Companion PR

osaurus-ai/vmlx-swift#322 finishes the same migration on the Swift side. Until
this lands, that PR writes both spellings at each setenv site so mmap keeps
being forced regardless of which mlx revision is pinned. Once this merges, the
legacy write can be dropped there.

A note on the base branch

This targets qwen4-exp, not main — the mmap safetensors code does not
exist on main, and qwen4-exp is the branch vmlx-swift pins as its Cmlx
submodule (currently 0c7cf413). Happy to retarget if this work is headed for
main by another route.

Verified by compiling through the vmlx-swift build, which is what consumes
this file.

`mmap_tensor_buffers_enabled` already reads both `VMLINUX_` and `VMLX_`, and
`VMLX_QMV_MR` in the Metal backend uses the current prefix. The other five mmap
variables were left on the legacy spelling alone, so which name works depends on
which knob you reach for.

It matters beyond consistency, because the Swift side has already moved:
vmlx-swift's `ModelFactory` sets `VMLX_MMAP_SAFETENSORS` to force mmap for a
load, and nothing here reads it. That write reaches nothing today; mmap keeps
working only because `MLX_SAFETENSORS_MMAP` is set beside it. For this one
variable the NEW name is the weaker one.

All three spellings are accepted; nothing that works today stops working.

The two value-returning knobs differ from the booleans and are handled
accordingly: for `||`-chained flags precedence is unobservable, but
`mmap_start_cold_pct` and `mmap_cold_advice` RETURN a value, so the current
spelling is consulted first — otherwise setting both would silently yield the
legacy one.

Companion to osaurus-ai/vmlx-swift#322, which finishes the same migration on the
Swift side. Once this lands, the legacy write can be dropped there.
@rcfa
rcfa marked this pull request as ready for review August 28, 2026 19:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant