mlx-train-perf 0.1.0__tar.gz

mlx_train_perf-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,48 @@
+name: ci
+on: [push, pull_request]
+jobs:
+  checks:
+    runs-on: macos-15
+    steps:
+      - uses: actions/checkout@v4
+        with: { fetch-depth: 0 }
+      - uses: astral-sh/setup-uv@v5
+      - run: uv sync --all-extras
+      - run: uv run ruff check src/ tests/
+      - run: uv run mypy src/
+      # coverage REPORTED but not gated here: the metal marker gates CORE code (kernel
+      # launch/loss paths), so the default lane can't reach fail_under=85 by construction.
+      # The gate lives in the metal-enabled lane (added in Task 7 once the probe verdict is
+      # known) and in the Task 18 local pre-release check.
+      - run: uv run pytest tests/ --cov=mlx_train_perf --cov-fail-under=0
+  metal-probe:
+    runs-on: macos-15
+    continue-on-error: true
+    steps:
+      - uses: actions/checkout@v4
+        with: { fetch-depth: 0 }
+      - uses: astral-sh/setup-uv@v5
+      - run: uv sync
+      - name: Can this runner JIT + run mx.fast.metal_kernel?
+        run: |
+          uv run python - <<'EOF'
+          import mlx.core as mx
+          k = mx.fast.metal_kernel(name="probe", input_names=["inp"], output_names=["out"],
+                                   source="uint e = thread_position_in_grid.x; out[e] = inp[e] * 2.0f;")
+          x = mx.arange(16, dtype=mx.float32)
+          (y,) = k(inputs=[x], grid=(16, 1, 1), threadgroup=(16, 1, 1),
+                   output_shapes=[(16,)], output_dtypes=[mx.float32])
+          assert (y == x * 2).all().item(), "wrong result"
+          print("METAL PROBE: OK")
+          EOF
+  # Task-2 probe verdict: METAL PROBE: OK on macos-15 — this lane runs the Metal-gated
+  # kernel parity suite and carries the fail_under=85 coverage gate (the default `checks`
+  # job above can't reach it: the metal marker gates the kernel launch/loss paths).
+  metal-tests:
+    runs-on: macos-15
+    steps:
+      - uses: actions/checkout@v4
+        with: { fetch-depth: 0 }
+      - uses: astral-sh/setup-uv@v5
+      - run: uv sync --all-extras
+      - run: uv run pytest tests/ --run-metal --cov=mlx_train_perf

mlx_train_perf-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,34 @@
+name: publish
+# Tag-triggered release: build the sdist + wheel and publish to PyPI via Trusted
+# Publishing (OIDC, no stored token). The package is pure Python (the Metal kernel is
+# JIT-compiled at runtime through mx.fast.metal_kernel), so the build is platform-neutral.
+on:
+  push:
+    tags: ["v*"]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0   # hatch-vcs derives the version from the git tag/history
+      - uses: astral-sh/setup-uv@v5
+      - run: uv build
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write   # required for PyPI Trusted Publishing
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

mlx_train_perf-0.1.0/.gitignore

@@ -0,0 +1,8 @@
+/CLAUDE.md
+docs/superpowers/
+.venv/
+__pycache__/
+*.egg-info/
+.coverage
+dist/
+_artifacts/

mlx_train_perf-0.1.0/CHANGELOG.md

@@ -0,0 +1,41 @@
+# Changelog
+
+All notable changes to this project are documented here. The format follows
+[Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and versions follow
+[Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [0.1.0] - 2026-07-08
+
+First release. A fused, logit-free linear-cross-entropy loss for MLX training on Apple
+Silicon, with an mlx-lm adapter, a RAM-fit planner, and a benchmark harness.
+
+### Added
+- Fused Metal cross-entropy kernel that never materializes the `(N, V)` logits: about
+  3900x less loss-layer memory than materialized logits in isolation, at a 1.64x forward
+  cost (n=8192, V=151936, D=4096, bf16; `scripts/bench_loss_layer.py`). Exact value and
+  gradient parity against a materialized reference.
+- Three implementations behind one `impl` argument: `kernel` (the fused Metal path),
+  `chunked` (a pure-MLX fallback bounded by a fixed vocabulary tile), and `naive` (the
+  materialized correctness oracle). `auto` selects the kernel when the mlx version is
+  verified and the head and dtype are supported, and otherwise raises a typed error
+  naming the reason and the alternatives. It never silently downgrades.
+- Quantized-head support: 4-bit group-size-64 heads (the mlx-community QLoRA default) run
+  through the kernel.
+- mlx-lm training adapter (`make_loss_fn`) that plugs the loss into `mlx_lm`'s compiled
+  trainer, with per-step loss curves matching the stock trainer to bf16 tolerance
+  (about 2e-3). End-to-end throughput is roughly 8-12% slower per step at bf16
+  (`scripts/bench_train_step.py`).
+- RAM-fit planner that estimates peak training memory and accounts for the O(N^2)
+  attention backward that dominates at long context. Fit to measured Qwen3-8B train-step
+  peaks and cross-model validated on Llama-3.2-3B to within about 9%.
+- Benchmark harness and committed scripts for every published number.
+
+### Known limits
+- The fused loss frees memory at a given context but does not, on its own, extend the
+  maximum trainable context on MLX. `mx.fast.scaled_dot_product_attention` has an O(N^2)
+  backward that materializes the `(N, N)` attention matrix and becomes the memory
+  bottleneck at long context. A memory-efficient attention backward is the next step (see
+  `ROADMAP.md`).
+- Architectures: Llama and Qwen3 only. Training: LoRA / QLoRA. Apple Silicon only.
+
+[0.1.0]: https://github.com/IonDen/mlx-train-perf/releases/tag/v0.1.0

mlx_train_perf-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Denis Ineshin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mlx_train_perf-0.1.0/PKG-INFO

@@ -0,0 +1,123 @@
+Metadata-Version: 2.4
+Name: mlx-train-perf
+Version: 0.1.0
+Summary: Fused, logit-free linear-cross-entropy loss, RAM-fit planner, and benchmark harness for MLX fine-tuning on Apple Silicon
+Author-email: Denis Ineshin <denis.ineshin@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 4 - Beta
+Classifier: Operating System :: MacOS
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: mlx>=0.31.2
+Provides-Extra: mlx-lm
+Requires-Dist: mlx-lm<0.32,>=0.31.3; extra == 'mlx-lm'
+Requires-Dist: transformers<5.13,>=5.0; extra == 'mlx-lm'
+Provides-Extra: probe
+Requires-Dist: pyobjc-framework-metal; extra == 'probe'
+Description-Content-Type: text/markdown
+
+# mlx-train-perf
+
+A fused, logit-free linear-cross-entropy loss for training on Apple Silicon with [MLX](https://github.com/ml-explore/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](https://arxiv.org/abs/2411.09009) and [Liger-Kernel](https://github.com/linkedin/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 were measured on an M1 Max (32 GB), mlx 0.31.2.
+
+## The problem it solves
+
+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.
+
+## What this does and does not buy you
+
+I want to be precise here, because the honest end-to-end story is narrower than the loss-layer number suggests.
+
+In a real fine-tune step the fused loss frees the memory the logit tensor would have taken. That headroom goes to a larger batch or a longer sequence at the *same* peak. The 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).
+
+What it does *not* buy you on MLX today is a longer maximum context before you hit an out-of-memory error. We measured this directly (`scripts/northstar_context_sweep.py`): with gradient checkpointing on, ours and the stock trainer hit essentially the same context ceiling for an 8B QLoRA on 32 GB — ours about 8450 tokens, stock about 8700, a gap of one 256-token probe step. The reason is that `mx.fast.scaled_dot_product_attention` has a memory-efficient forward but an O(N²) backward. It materializes the `(N,N)` attention matrix one layer at a time during training, and that term dominates the peak at long context. Once the logits are gone, attention is the bottleneck, not the loss. A memory-efficient attention backward is the piece that would move the context ceiling on MLX, and it is the next thing on the [roadmap](ROADMAP.md).
+
+So this frees real memory at a given context, and it is the right building block, but the flagship "train much longer sequences" win needs the attention backward too.
+
+## Install
+
+```bash
+pip install mlx-train-perf            # the loss kernel + planner
+pip install "mlx-train-perf[mlx-lm]"  # plus the mlx-lm training adapter
+```
+
+Apple Silicon only. Requires mlx (>=0.31.2 recommended; the kernel's JIT contract is verified against it).
+
+## Use it in an mlx-lm fine-tune
+
+The adapter builds a loss callable with the same signature `mlx_lm`'s trainer expects, so you pass it straight to `train(...)`:
+
+```python
+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
+
+model, tokenizer = load("mlx-community/Qwen3-8B-4bit")
+model.set_dtype(mx.bfloat16)  # 4-bit checkpoints compute in fp16; the kernel needs bf16/fp32
+# ... freeze the base model and apply linear_to_lora_layers as in a normal mlx-lm LoRA run ...
+
+loss_fn = make_loss_fn(model, impl="auto")
+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.
+
+## Implementations
+
+`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 than `kernel` but far less than `naive`. 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.
+
+## RAM-fit planner
+
+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:
+
+```bash
+mlx-train-perf plan --config path/to/config.json --batch 1 --seq-len 4096 --lora-rank 8
+```
+
+The 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.
+
+## Supported models
+
+- Architectures: Llama and Qwen3. 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.
+
+## Reproducing the numbers
+
+Each claim above has one script. They run on the GPU, take real wall-clock time, and print the artifacts they measured:
+
+```bash
+python scripts/bench_loss_layer.py        # the ~3900x loss-layer memory number
+python scripts/bench_train_step.py --model mlx-community/Qwen3-8B-4bit --seq-len 1024 2048 \
+    --impl kernel --compute-dtype bfloat16 --grad-checkpoint   # end-to-end tok/s vs stock
+python scripts/northstar_context_sweep.py # the max-context sweep (1-2 h; heavy)
+```
+
+## License
+
+MIT. See [LICENSE](LICENSE).

mlx_train_perf-0.1.0/README.md

@@ -0,0 +1,101 @@
+# mlx-train-perf
+
+A fused, logit-free linear-cross-entropy loss for training on Apple Silicon with [MLX](https://github.com/ml-explore/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](https://arxiv.org/abs/2411.09009) and [Liger-Kernel](https://github.com/linkedin/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 were measured on an M1 Max (32 GB), mlx 0.31.2.
+
+## The problem it solves
+
+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.
+
+## What this does and does not buy you
+
+I want to be precise here, because the honest end-to-end story is narrower than the loss-layer number suggests.
+
+In a real fine-tune step the fused loss frees the memory the logit tensor would have taken. That headroom goes to a larger batch or a longer sequence at the *same* peak. The 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).
+
+What it does *not* buy you on MLX today is a longer maximum context before you hit an out-of-memory error. We measured this directly (`scripts/northstar_context_sweep.py`): with gradient checkpointing on, ours and the stock trainer hit essentially the same context ceiling for an 8B QLoRA on 32 GB — ours about 8450 tokens, stock about 8700, a gap of one 256-token probe step. The reason is that `mx.fast.scaled_dot_product_attention` has a memory-efficient forward but an O(N²) backward. It materializes the `(N,N)` attention matrix one layer at a time during training, and that term dominates the peak at long context. Once the logits are gone, attention is the bottleneck, not the loss. A memory-efficient attention backward is the piece that would move the context ceiling on MLX, and it is the next thing on the [roadmap](ROADMAP.md).
+
+So this frees real memory at a given context, and it is the right building block, but the flagship "train much longer sequences" win needs the attention backward too.
+
+## Install
+
+```bash
+pip install mlx-train-perf            # the loss kernel + planner
+pip install "mlx-train-perf[mlx-lm]"  # plus the mlx-lm training adapter
+```
+
+Apple Silicon only. Requires mlx (>=0.31.2 recommended; the kernel's JIT contract is verified against it).
+
+## Use it in an mlx-lm fine-tune
+
+The adapter builds a loss callable with the same signature `mlx_lm`'s trainer expects, so you pass it straight to `train(...)`:
+
+```python
+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
+
+model, tokenizer = load("mlx-community/Qwen3-8B-4bit")
+model.set_dtype(mx.bfloat16)  # 4-bit checkpoints compute in fp16; the kernel needs bf16/fp32
+# ... freeze the base model and apply linear_to_lora_layers as in a normal mlx-lm LoRA run ...
+
+loss_fn = make_loss_fn(model, impl="auto")
+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.
+
+## Implementations
+
+`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 than `kernel` but far less than `naive`. 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.
+
+## RAM-fit planner
+
+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:
+
+```bash
+mlx-train-perf plan --config path/to/config.json --batch 1 --seq-len 4096 --lora-rank 8
+```
+
+The 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.
+
+## Supported models
+
+- Architectures: Llama and Qwen3. 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.
+
+## Reproducing the numbers
+
+Each claim above has one script. They run on the GPU, take real wall-clock time, and print the artifacts they measured:
+
+```bash
+python scripts/bench_loss_layer.py        # the ~3900x loss-layer memory number
+python scripts/bench_train_step.py --model mlx-community/Qwen3-8B-4bit --seq-len 1024 2048 \
+    --impl kernel --compute-dtype bfloat16 --grad-checkpoint   # end-to-end tok/s vs stock
+python scripts/northstar_context_sweep.py # the max-context sweep (1-2 h; heavy)
+```
+
+## License
+
+MIT. See [LICENSE](LICENSE).

mlx_train_perf-0.1.0/ROADMAP.md

@@ -0,0 +1,35 @@
+# Roadmap
+
+## Released
+
+### 0.1.0
+- Fused, logit-free linear-cross-entropy Metal kernel (forward), with a chunked pure-MLX
+  fallback and a materialized correctness oracle.
+- Quantized (4-bit group-size-64) head support.
+- mlx-lm training adapter, RAM-fit planner, and benchmark harness.
+
+## Planned
+
+### Memory-efficient attention backward
+This is the highest-leverage item. The measured bottleneck for long-context training on
+MLX is the attention backward, which materializes the `(N, N)` score matrix one layer at
+a time. A tiled, recompute-based attention backward (the FlashAttention backward, written
+as a Metal kernel) is what would actually raise the maximum trainable context on Apple
+Silicon. The fused loss frees the logit memory; this frees the attention memory, and
+together they are what "train much longer sequences" needs.
+
+### Fused backward kernel for the loss
+The loss ships with a fused forward and a proven chunked backward. A fully fused backward
+kernel is a further memory reduction for the loss layer itself.
+
+### Sequence packing
+Pack variable-length examples into fixed blocks to cut padding waste in SFT training.
+
+### Planner inverse queries
+Ask the planner for the largest batch or sequence length that fits a given memory budget,
+instead of checking one config at a time.
+
+## Not doing
+- Trainer UX or a training-loop framework. That is the lane of tools like mlx-lm-lora and
+  mlx-tune; this project is a layer they can import.
+- Non-Apple-Silicon backends.

mlx_train_perf-0.1.0/pyproject.toml

@@ -0,0 +1,112 @@
+[build-system]
+requires = ["hatchling", "hatch-vcs"]
+build-backend = "hatchling.build"
+
+[project]
+name = "mlx-train-perf"
+dynamic = ["version"]
+description = "Fused, logit-free linear-cross-entropy loss, RAM-fit planner, and benchmark harness for MLX fine-tuning on Apple Silicon"
+readme = "README.md"
+requires-python = ">=3.11"
+license = "MIT"
+authors = [{ name = "Denis Ineshin", email = "denis.ineshin@gmail.com" }]
+dependencies = ["mlx>=0.31.2"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Operating System :: MacOS",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+[project.optional-dependencies]
+# transformers>=5.13 breaks `import mlx_lm` outright: mlx_lm/tokenizer_utils.py:505 calls
+# AutoTokenizer.register("NewlineTokenizer", ...) with a string class name, and
+# transformers 5.13's auto_factory.py:680 requires an actual class object there
+# (AttributeError: 'str' object has no attribute '__module__'). Verified clean on
+# 5.0.0/5.4.0/5.8.0/5.11.0/5.12.0; verified broken on 5.13.0.
+mlx-lm = ["mlx-lm>=0.31.3,<0.32", "transformers>=5.0,<5.13"]
+probe = ["pyobjc-framework-Metal"]
+
+[project.scripts]
+mlx-train-perf = "mlx_train_perf.cli:main"
+
+[tool.hatch.version]
+source = "vcs"
+
+[tool.hatch.build.targets.sdist]
+# Explicit allowlist: the repo carries gitignored working-state (.superpowers/ SDD
+# ledger, .codegraph/ index, docs/, _artifacts/, CLAUDE.md) that hatchling would
+# otherwise sweep into the sdist. Ship only the source, tests, bench scripts, CI, and
+# the user-facing docs. (The wheel already ships just the `mlx_train_perf` package.)
+include = [
+    "/src",
+    "/tests",
+    "/scripts",
+    "/.github",
+    "/README.md",
+    "/CHANGELOG.md",
+    "/ROADMAP.md",
+    "/LICENSE",
+]
+
+[dependency-groups]
+dev = [
+    "pytest>=8", "pytest-cov", "hypothesis", "ruff", "mypy", "basedpyright",
+    # transformers>=5.13 breaks `import mlx_lm` (see the mlx-lm extra above for the
+    # verified-broken/verified-clean version detail); pinned here too since dev installs
+    # both groups.
+    "mlx-lm>=0.31.3,<0.32", "transformers>=5.0,<5.13", "pyobjc-framework-Metal",
+]
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I", "N", "UP", "B", "A", "C4", "PT", "RET", "SIM", "ARG", "PTH", "PL", "RUF"]
+ignore = ["PLR0913", "PLR2004"]
+
+[tool.ruff.lint.per-file-ignores]
+# Verbatim-ported MSL kernel body (mlx-train-perf-spike/kernel_v2e.py) — line length is
+# inherent to the source text; reformatting would break the byte-for-byte port fidelity
+# the parity cross-check (task-7-report.md) depends on.
+"src/mlx_train_perf/core/kernel/source.py" = ["E501"]
+
+[tool.mypy]
+strict = true
+python_version = "3.12"
+
+[[tool.mypy.overrides]]
+module = ["mlx_lm.*", "Metal"]
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = ["numpy", "numpy.*"]
+follow_imports = "skip"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+# `mlx/nn/__init__.py` re-exports its public API via `from .layers import *` with no
+# `__all__` -- under `--strict`'s `no_implicit_reexport`, mypy doesn't treat that as a
+# legitimate re-export, so `nn.Module` / `nn.Linear` / etc. fail as "not defined" / "has
+# no attribute" at every call site (verified: `reveal_type(mlx.nn)` -> a bare
+# `types.ModuleType`, not an error in mlx.nn's own source). `implicit_reexport = true`
+# fixes exactly that resolution gap while still fully attribute-checking `nn.*` (e.g. a
+# typo'd `nn.QuantizedLinaer` is still caught) -- unlike `follow_imports = "skip"`,
+# which would also treat `nn.*` as `Any` and silently swallow that kind of typo.
+module = ["mlx.nn", "mlx.nn.*"]
+implicit_reexport = true
+
+[tool.pytest.ini_options]
+addopts = "-ra"
+markers = [
+    "metal: needs a Metal GPU + JIT (collection-gated: --run-metal)",
+    "smoke: loads a real model (collection-gated: --run-smoke)",
+    "benchmark: timing runs (collection-gated: --run-benchmark)",
+    "network: touches the network (collection-gated: --run-network)",
+]
+
+[tool.coverage.report]
+fail_under = 85