save_safetensors: write files the mmap loader can actually map

[rcfa]

tensor_from_mmap shares the mapping only when a tensor's absolute byte offset is a multiple of its
dtype width, and copies into a freshly allocated aligned buffer otherwise. save_safetensors
guarantees the copy: the header is written unpadded, so the data section begins at
8 + strlen(json) — an arbitrary byte, landing on a multiple of 8 about one time in eight — and
tensors are then packed back to back at arbitrary relative offsets.

So this fork reads its own output on the slow path.

What it costs

Measured on a 95 GB bundle, 128 GB M4 Max:

	as written	after aligning the same shards
tensors on the zero-copy path	1,224 / 2,999	2,999 / 2,999
extra anonymous memory at load	59.3 GB	0
memory compressor	60.1 GB	0.8 GB
generation	~3 tokens/minute	normal

The 59.3 GB is a second copy of most of the weights on top of the 95 GB mapping. Below the RAM
ceiling nobody notices; above it the machine swaps and throughput collapses.

Output is unaffected — the same bytes are copied, so greedy decoding is bit-identical before and
after. That is why this can sit undetected: it is invisible in accuracy and in any test that reads
the tensors, and shows up only as memory.

Scanning 128 locally held MLX bundles, 49 were affected, across 11 publisher namespaces — it
tracks the writer, not the packager.

The change

Pad the header to 8 bytes, and lay each tensor on an 8-byte boundary. Both are needed: an aligned
relative offset is only aligned absolutely if the data section itself starts aligned. Eight
covers every dtype safetensors defines, so one constant serves them all, at a cost of at most 7
bytes per tensor — about 21 KB across a 95 GB model.

Header padding is sanctioned by the format rather than merely tolerated. The reference
implementation carries a test for it:

/// Test that the JSON header may be leading-padded with JSON whitespace characters.
/// This is intentional: writers may pad the header to align the data section to a
/// page boundary, so readers must tolerate leading whitespace.

so a padded file stays readable by every existing consumer, including older MLX and non-MLX loaders.

Relationship to a828cb4

This complements "Fix mmap loading for unaligned safetensors tensors",
which made unaligned files load correctly. This stops producing them, so the fast path becomes
reachable rather than merely survivable.

Verification

I did not build the full framework for this. Instead the layout arithmetic was extracted verbatim
into a standalone program and run over tensors with deliberately hostile sizes — 3, 10, 1 and 6
bytes, which misalign everything after them under back-to-back packing — with the output parsed by
an independent reader:

• patched: data section at (8 + header_len) % 8 == 0, 6 of 6 tensors aligned, every
  tensor's bytes intact and distinct
• control, same test on the pre-change arithmetic: data section at 3 mod 8, 2 of 6
  misaligned — so the test discriminates rather than passing vacuously

A build-and-run check against the real save_safetensors / load_safetensors round trip would be
worth doing before merge; I flag that rather than imply it was done.

Context

Filed upstream as ml-explore/mlx#4439 as a feature
request — upstream never mmaps safetensors, so alignment costs it nothing there. Here the reader and
writer are in the same file and disagree, which is why the fix belongs on this branch.

Existing published files can be repaired in place without altering a single tensor byte — only
positions change. Detection needs just the first few KB of each shard: misaligned bundles start
their data section at (8 + header_len) % 8 != 0.