A fused, logit-free linear-cross-entropy loss for training on Apple Silicon with MLX, plus a RAM-fit planner and an honest benchmark harness. It drops into an mlx-lm LoRA/QLoRA fine-tune as the loss function.
The idea is the same one behind Cut Cross-Entropy and Liger-Kernel on the CUDA side, ported to a Metal kernel: compute the cross-entropy loss and its gradient without ever building the full (N, V) logits tensor. For a large vocabulary that tensor is the single biggest allocation in the training step, and it is pure waste. You only need the per-token loss and a gradient back into the hidden states.
Released on PyPI as mlx-train-perf. Every number below has a committed script under scripts/ that reproduces it, all measured on one M1 Max (32 GB, macOS 26.5). The loss-layer figures were taken on mlx 0.31.2 and reproduce on the pinned 0.32.0; the flash-attention memory figures were taken on 0.32.0 in 0.2.0, and the 0.3.0 context-ceiling figures on 0.32.0.
Standard cross-entropy in a trainer materializes logits of shape (batch·seq, vocab). At Qwen3-8B's vocabulary (151,936) and a 2048-token sequence, that is a 0.6 GB tensor in bf16, plus another for the softmax gradient in the backward pass. The fused kernel never allocates it: the forward regenerates logits in registers tile-by-tile over the vocabulary and returns three N-length arrays (the per-token NLL, the log-sum-exp, and the target logit); the backward recomputes the needed tiles instead of reading a stored matrix.
Measured in isolation, at n=8192, V=151936, D=4096, bf16 (scripts/bench_loss_layer.py):
| loss layer | peak memory | forward wall |
|---|---|---|
| naive (materialized logits) | 2.318 GB | 1.0× |
| kernel (this project) | 0.0006 GB | 1.64× |
About 3900× less memory for the loss layer, at a 1.64× cost on the forward pass.
The fused loss is exact to bf16 tolerance against the stock trainer (per-step loss curves match to about 2e-3), and the throughput cost is small: roughly 8–12% slower per step at bf16 (scripts/bench_train_step.py, Qwen3-8B-4bit, LoRA r=8, gradient checkpointing on).
New in 0.2.0 and opt-in. Removing the logit tensor frees real memory, but on its own it barely moves the training peak at long context. The reason is attention. mx.fast.scaled_dot_product_attention has a memory-light forward and an O(N²) backward that rebuilds the (N, N) score matrix one layer at a time. Once the logits are gone, that backward is what sets the peak.
0.2.0 adds a flash-attention path with a Metal forward and a Metal backward, neither of which materializes the score matrix. It keeps O(N) saved state — the attention output and the log-sum-exp — and recomputes the tiles it needs. You switch it on per model with enable_flash_attention and train exactly as before.
On Qwen3-8B-4bit (LoRA rank 8, batch 1, gradient checkpointing on, bf16) the two attention paths are close at a 2048-token sequence. At 8192 they are not: the flash path halves the whole step's peak memory.
| seq 8192, Qwen3-8B-4bit | total peak | marginal peak |
|---|---|---|
| stock attention | 25.68 GiB | 21.31 GiB |
| flash attention | 12.75 GiB | 8.37 GiB |
(scripts/bench_train_step.py; M1 Max 32 GB, macOS 26.5, mlx 0.32.0. These memory and throughput figures are the 0.2.0 measurements, carried into 0.3.0 unchanged: 0.3.0 changed how the backward splits its kernel launches, not what it allocates, and the 0.3.0 context sweep below — measured fresh — confirms the flash path's memory still scales linearly in sequence length.)
The 32 GB machine that peaked near its ceiling with stock attention now runs the same step at half the memory. That is real headroom: a longer sequence, or a second job on the GPU.
The attention op itself, timed alone at the flagship shape (batch 1, 32 query / 8 KV heads, 8192 tokens, head_dim 128), is 0.186 s on the forward and 0.576 s on the full backward, 3.1× the forward (scripts/bench_attention_op.py). As the sequence doubles, the flash op's peak grows about 2.00× (2048→4096) and about 3.06× (4096→8192). The second step is above 2× because the chained backward split adds a small, budget-bounded constant of at most ~0.3 GB, not because the O(N) growth law changed. Stock attention over the same doublings grows 3.76× then 3.05×, and that last figure is flattered by paging: at 8192 the stock op allocates about 32.4 GB on a 32 GB machine and its wall time degrades roughly 41× as it pages, which puts the single flash op about 45× below stock there. Those stock figures are a pre-net measurement of the exact hazard the safety net now prevents. They were taken before the memory watchdog shipped, and on a 32 GB machine scripts/bench_attention_op.py now aborts that condition by design (aborted_memory_ceiling) rather than paging into it. The flash-side numbers all reproduce. Flash is not universally cheaper, though — below about 2100 tokens the stock op's simpler bookkeeping wins, and the two curves cross there. The win is at real training context, not tiny shapes. (At 16384 the flash op now runs rather than refusing, since the launch guard no longer caps the chain; stock attention still cannot reach that context on 32 GB.)
Turning flash attention on is not free. On the stock-loss path at 8192 it costs 5.3% of tokens/sec (74.0 vs 78.1); at 2048 the cost is 5.5% on the fused-loss path (86.4 vs 91.5) and 5.9% on the stock-loss path (92.1 vs 97.8). The fused-loss comparison at 8192 has no stock-attention number to pair with: on this 32 GB machine that baseline condition crosses the memory safety net's ceiling and records an abort instead of a number. That baseline running out of room is the problem flash attention exists to remove. Under flash attention the fused cross-entropy and mlx-lm's stock cross-entropy stay close: 0.94× at 2048 (86.4 vs 92.1 tok/s) and 0.99× at 8192 on Qwen3-8B, 0.92× and 0.97× on Llama-3.2-3B. The loss values match to bf16 tolerance throughout — the worst per-step difference across every measured pair is 2.4e-3. The worst attention-arm throughput ratio measured is 0.94× stock; 0.85× is the maintainer's release acceptance bar for that path, provisional and pinned at the PR.
0.2.0 shipped this path with a caveat — it halved the memory but did not extend the longest sequence you could train, because a launch-safety budget capped it before memory did. That cap turned out to be guarding the wrong thing. Re-reading how mlx schedules Metal work, and re-running the crash that motivated the budget, showed the GPU watchdog acts on a single command buffer, not on a chain of them, and at training shapes each backward dispatch already runs in its own buffer. 0.3.0 replaces the per-chain budget with a per-buffer one that models what the scheduler actually commits. No safety margin moved; what changed is what the margin is measured against. (scripts/probe_command_buffer_packing.py reproduces the evidence.)
With that cap gone, the flash path is bound by memory, the same thing that bounds stock attention — and it needs far less of it. Measured the same day with the same search (scripts/northstar_context_sweep.py, Qwen3-8B-4bit QLoRA, gradient checkpointing, bf16):
| max trainable context, 32 GB | tokens | peak at the ceiling |
|---|---|---|
| stock attention | 7,936 | 24.5 GiB |
| flash attention | 23,040 | 24.5 GiB |
Both arms stop at the same ~24.5 GiB, the effective memory ceiling on this machine at run time. Under that one budget the flash path trains 2.9× the context, because it holds O(N) saved state where stock holds the O(N²) score matrix. The ratio is the part that travels: raise the available memory and both ceilings rise together (a freshly booted or larger machine lets both climb toward the 28 GiB static ceiling), but the flash path keeps its roughly threefold reach. On the same machine in 0.2.0 this path was launch-capped near 10k tokens; removing that cap is what moved it. (These figures are this release's measurement; the two arms are comparable to each other, taken together, not to 0.2.0's numbers.)
enable_flash_attention is causal-only and training-only, and it refuses anything outside that up front rather than failing mid-run:
| Condition | When | Error |
|---|---|---|
| Model family other than Llama, Qwen2, or Qwen3 | at enable | UnsupportedAttentionError |
Sliding-window or mixed attention (layer_types not all full_attention) |
at enable | UnsupportedAttentionError |
head_dim outside {64, 96, 128} |
at enable | UnsupportedAttentionError |
| Non-zero attention dropout | at enable | UnsupportedAttentionError |
| An array attention mask (sliding-window or additive) | first attention call | AttentionInputError |
| A KV cache present (inference) | first attention call | AttentionInputError |
from mlx_train_perf.attention import enable_flash_attention
enable_flash_attention(model, seq_len=8192, batch_size=1)Call it in place on a loaded model, after you set the compute dtype and before you build the loss and call train. mlx-lm's train wraps the step in mx.compile, and the kernel calibrates itself with a one-time host sync that a compiled step cannot run inside its traced region. Passing seq_len (and batch_size) does that calibration up front, at your training shape, so the compiled step traces with warm caches. Match them to the shape you actually train: batch_size defaults to 1 and must equal your training batch. The hints are load-bearing for compiled training: if a compiled train traces at a shape the caches were not warmed for, the calibration host-sync fires inside the compiled region and raises during the trace, a loud failure rather than a silent slowdown. Omit the hints and the call still succeeds — eager and mx.grad callers calibrate lazily on the first attention call — but a compiled train run should always pass them.
pip install mlx-train-perf # the loss kernel + planner
pip install "mlx-train-perf[mlx-lm]" # plus the mlx-lm training adapterApple Silicon only. Requires mlx >=0.32.0,<0.33 — the version the kernels' JIT contract is verified against. The mlx-lm adapter and the flash-attention wrapper need the optional mlx-lm extra.
The adapter builds a loss callable with the same signature mlx_lm's trainer expects, so you pass it straight to train(...). enable_flash_attention is the second, independent lever — turn on either, both, or neither:
import mlx.core as mx
from mlx_lm import load
from mlx_lm.tuner.trainer import train
from mlx_train_perf.adapters.mlx_lm import make_loss_fn
from mlx_train_perf.attention import enable_flash_attention
model, tokenizer = load("mlx-community/Qwen3-8B-4bit")
model.set_dtype(mx.bfloat16) # 4-bit checkpoints compute in fp16; the kernels need bf16/fp32
# ... freeze the base model and apply linear_to_lora_layers as in a normal mlx-lm LoRA run ...
enable_flash_attention(model, seq_len=8192, batch_size=1) # O(N) attention backward
loss_fn = make_loss_fn(model, impl="auto") # logit-free cross-entropy
train(model=model, optimizer=opt, train_dataset=ds, args=args, loss=loss_fn)make_loss_fn splits the model into its trunk and its output head and routes the loss through the fused kernel; enable_flash_attention swaps each layer's attention for the flash path in place.
impl picks how the loss is computed. "auto" is the default and the one to use.
kernel— the fused Metal kernel."auto"resolves here when the mlx version is verified and the head/dtype are supported (dense or tied fp32/bf16 head; 4-bit group-size-64 quantized head; hidden states in fp32 or bf16). It never materializes(N, V).chunked— a pure-MLX fallback that processes the vocabulary in fixed tiles. No Metal kernel, works anywhere MLX does, uses more memory thankernelbut far less thannaive. This is also the backward path the kernel forward pairs with today.naive— materializes the full logits. It is the correctness oracle the other two are tested against, not something to train with.
"auto" never silently downgrades. If it cannot use the kernel (unverified mlx, an unsupported head, fp16 hidden states) it raises a typed error naming the reason and the alternatives, so you always know which path ran.
Before a run, the planner estimates the peak training memory for a config and tells you whether it fits, or suggests a smaller batch or sequence length that would:
mlx-train-perf plan --config path/to/config.json --batch 1 --seq-len 4096 --lora-rank 8The memory model is fit to measured Qwen3-8B train-step peaks and cross-model validated on Llama-3.2-3B to within about 9%. It accounts for the O(N²) attention backward described above, so it does not under-predict at long context the way a linear model would. It is an estimate, and it errs toward over-predicting, which is the safe direction for a tool whose job is to keep you off the OOM cliff.
Pass --attention flash to price the flash-attention path instead of the stock backward:
mlx-train-perf plan --config path/to/config.json --batch 1 --seq-len 8192 --lora-rank 8 --attention flashThe flash model is an analytic saved-state term plus one measured linear coefficient. Across the four measured anchors up to 8192 tokens it over-predicts the cushioned total by 8% to 20% — again on the safe side.
- Architectures: Llama, Qwen2 (the Qwen2.5 family), and Qwen3, for both the loss adapter and the flash-attention wrapper. The adapter's model splitter handles these; others raise a typed error.
- Quantization: 4-bit group-size-64 (the mlx-community QLoRA default), or a dense fp32/bf16 head.
- Training: LoRA / QLoRA. Full fine-tuning is estimated by the planner but is not the case this is tuned for.
- Hardware: Apple Silicon.
Each claim above has one script. They run on the GPU, take real wall-clock time, and print the artifacts they measured:
python scripts/bench_loss_layer.py # the ~3900x loss-layer memory number
python scripts/bench_attention_op.py # the single-op flash vs stock memory + timing
# the 12.75 vs 25.68 GiB training table: run each attention arm into its own --out dir
python scripts/bench_train_step.py --model mlx-community/Qwen3-8B-4bit --seq-len 8192 \
--attention flash --impl kernel --compute-dtype bfloat16 --grad-checkpoint --out _artifacts/flash
python scripts/bench_train_step.py --model mlx-community/Qwen3-8B-4bit --seq-len 8192 \
--attention stock --impl kernel --compute-dtype bfloat16 --grad-checkpoint --out _artifacts/stock
python scripts/northstar_context_sweep.py # the max-context sweep (1-2 h; heavy)Every benchmark and contribution run is fenced by a device-relative memory guard. A GPU over-allocation on Apple Silicon does not always fail cleanly: mx.set_memory_limit is advisory, so an allocation past the soft cap pages instead of raising, and a hard enough paging storm can panic the machine rather than kill the process.
The guard sets an active-memory ceiling from the machine's own RAM. It is anchored at 28 GiB on a 32 GB Mac — above the largest legitimate peak measured here (25.68 GiB) and below physical RAM — and scales from that anchor across the range from 16 GB up to a 1 TB machine. At start it takes the smaller of that static ceiling and what the machine actually has free right now, minus a 2 GiB cushion. A daemon thread samples active memory throughout the run and aborts the moment it reaches the ceiling, writing an honest aborted-status artifact instead of letting the storm build. If the machine is already too loaded to start safely — less than a quarter of RAM effectively available — the run refuses up front with a typed error. Between those two points it proceeds but prints a warning naming how much memory it expected free for the machine's class against how much it measured, so a crowded machine is visible rather than silent.
The guard is rank-local: every input it reads is this node's own RAM, availability, and process memory. On a multi-node mx.distributed job each rank sizes its own ceiling and flags its own crowding, and a breach hard-exits that rank — so run distributed training under a launcher (mpirun or mlx.launch) that propagates a rank failure to the whole job.
Every number above is from an M1 Max (32 GB), the machine this is developed on. Whether the kernel and the flash-attention path scale the way the memory model expects on larger machines is a question only other people's hardware can answer, so there is a one-command way to measure it and send the numbers back:
mlx-train-perf contribute --tier quick # ~10-15 min; --tier full loads a model, ~1-2 hIt detects your machine, picks shapes for your RAM, prints a time estimate, runs the committed benches with the same memory guardrails the project uses, and writes one provenance-complete file plus a ready-to-paste PR. The three-step submission flow is in community-benchmarks/README.md.
Submitted results are folded into the table below (python scripts/aggregate_community.py).
Each row is measured on that contributor's own hardware and reported as-is — nothing is
extrapolated to machines no one has run, and a row does not imply a trainable-context
ceiling beyond what that machine measured. The stock-attention baseline comparison is run
on reference hardware by the maintainer, not asked of contributors.
| Chip | RAM (GB) | mlx | Tier | Loss kernel peak (GB) | Attn flash 2x ratio | Train tok/s (flash) | PR |
|---|---|---|---|---|---|---|---|
| Apple M1 Max | 32 | 0.32.0 | quick | 0.0006 | 3.06 | — | — |
The "Attn flash 2x ratio" column is how the flash forward+backward peak grows per sequence doubling. The O(N) target is about 2×. On the reference machine the largest measured pair (4096→8192) reads 3.06× rather than 2×, because the chained backward split adds a small, budget-bounded constant of at most ~0.3 GB; the growth law is still linear, so a reading a little above 2× on a given machine is expected, not a regression.
MIT. See LICENSE.