Feat/mps gemm 4bit bf16

[eaglstun]

Retargeted from an accidental upstream PR (bitsandbytes-foundation#2070) onto this fork. Same branch, same 23 commits.

Summary

Adds a native Metal (MPS) backend for Apple Silicon: hand-written Metal kernels for the
blockwise quantize/dequantize ops and both 4-bit matmuls, replacing the pure-PyTorch fallbacks
that MPS currently resolves to. Everything is gated behind -DCOMPUTE_BACKEND=mps and is inert
on every other platform.

Important

This branch has picked up scope beyond its title, and I'd rather flag that than have a
reviewer discover it. See Scope at the bottom — there is an unrelated (and I think real)
Lion weight-decay bug fix in here that deserves its own PR, plus some local agent-tooling docs
that almost certainly should not land upstream. Happy to split before review.

What runs natively now

op	before (on mps)	after
quantize_blockwise / dequantize_blockwise	pure-torch fallback	hand-written Metal
quantize_4bit / dequantize_4bit (+.out)	pure-torch fallback	hand-written Metal
gemv_4bit (M == 1)	dequant + F.linear	fused dequant+dot Metal kernel
gemm_4bit (general M)	dequant + F.linear	chunked dequant → scratch → GEMM
gemm_4bit_backward (new op)	composed inline in Python	same, fused

Each native path keeps its fallback and routes back to it whenever a guard fails (unsupported
blocksize, K % 32 != 0, missing symbol in a stale dylib, and so on), so a partial or older build
degrades to today's behaviour rather than breaking. BNB_MPS_REQUIRE_NATIVE=1 turns those silent
fallbacks into hard failures, which is what the test suite runs under.

bf16 needed a different GEMM. MPSMatrixMultiplication hard-asserts on anything but
fp32/fp16/int8/int16, so bf16 originally kept the fallback — awkward, since bf16 is the dtype most
QLoRA fine-tuning actually uses. MPSGraph does have a bf16 matmul, and it shares a command
buffer with our own compute encoders (MPSCommandBuffer conforms to MTLCommandBuffer), so bf16
now gets the same structure through a graph instead.

gemm_4bit_backward is a new op. MatMul4Bit.backward previously computed
grad_output @ dequantize_4bit(B) inline, which on MPS means a native dequant on one queue and a
matmul on torch's — two cross-queue round trips per Linear4bit per step, on roughly half of a
training step. The op's default kernel is exactly that composition, so every other backend is
unaffected and the fused MPS kernel has an oracle to be checked against.

Correctness

tests/test_mps_parity.py (new, ~1180 lines) checks every op against a CPU oracle with seeded
inputs, and asserts via spy that the native path is the one that actually ran — a parity test
that silently exercised the fallback would prove nothing.

• quantize_4bit / dequantize_4bit / dequantize_blockwise: bit-exact across dtypes and
  blocksizes.
• 4-bit matmuls: within documented per-dtype tolerances, and where a bit-exact claim is available
  it is asserted as one — the bf16 forward reproduces the fallback bit-exactly without bias, and
  the fused backward does so unconditionally (there is no bias epilogue in the backward to
  double-round).
• Cross-queue sync discipline has a dedicated race stress test, verified to have teeth: with the
  pre-dispatch torch.mps.synchronize() no-op'd it fails 30/30 iterations.

Known failures: 26, all quantize_blockwise (int8), and they are not MPS bugs. The MPS kernel
matches an exact float64 reference on all 1,048,576 values; the CPU kernel misses 1365 because
it snaps the normalized value to a 65536-point LUT before the codebook lookup. A Python emulation
of that LUT reproduces the CPU kernel bit-for-bit. The test therefore asserts bit-exactness between
an exact kernel and an approximate oracle. I have deliberately not "fixed" this by loosening the
test — whether the CPU kernel or the test should change is a maintainer call, and I'd like
guidance. These do not affect non-macOS CI.

Performance

M4 Max, macOS 26.5. Every table was taken on an idle machine with a fixed-size clone() control
read before and after; contention on this box has not merely added noise but reversed the winner
of an A/B, so anything measured under load was discarded rather than reported.

• gemv_4bit (M == 1, the decode case): 3.4–6.2x over dequant + F.linear across
  fp16/bf16/fp32.
• gemm_4bit: the win is the single sync, so it scales with how small the op is —
  bf16 2.09x at M=8, 1.94x at M=64, 1.09x at M=512, ~1.0x at M=2048, where the GEMM itself
  dominates and MPS's GEMM ≈ F.linear's.
• gemm_4bit_backward: 1.14–1.31x at realistic training shapes, slightly better than the
  forward at the same shapes — no bias epilogue to encode and no transpose for the GEMM to absorb.

End to end on a real workload (CogView4-6B QLoRA, 512×512, batch 1), toggling the native paths in
place against one binary via env switches, n=6 per arm with half the reps in reversed arm
order:

	forward	backward	step
fallback	4.784 ± 0.363 s	4.340 ± 0.375 s	9.394 ± 0.716 s
native	4.055 ± 0.264 s	3.638 ± 0.179 s	7.960 ± 0.431 s
	−15.2%	−16.2%	−15.3%

The measurement carries its own error bar: arms that should not move a given stage move it by
1.8–3.9%, so that is the noise floor, and the effects above sit well clear of it.

Design decisions recorded, not just made

docs/apple_silicon/MPS_STATUS.md is the op-by-op reference and also records what was tried and
rejected, with the evidence — because the code cannot show that. Most relevant to review:

• Why the native ops run on a private MTLCommandQueue with two syncs per call, costing a
  fixed ~0.15 ms. torch exposes no queue or stream handle; reaching MPSStream's internals means
  dlsym-ing mangled C++ inlines at header-derived offsets. Rejected as an ABI trap. The two
  sanctioned alternatives (a libtorch-linked extension, or torch.mps.compile_shader) are both
  re-architectures and are written up rather than half-done.
• Why a view's data_ptr() is cloned rather than bound at an offset. A view's data_ptr() is
  raw pointer arithmetic, not an id<MTLBuffer>; objc-probing one SIGSEGVs uncatchably. A
  load-time guard (bnb_mps_check_buffer_contract) pins the undocumented torch contract that an
  MPS tensor's data_ptr() is its MTLBuffer, so a future torch that changes it fails loudly
  instead of casting garbage into a kernel.

Scope — please read before reviewing

Three separable things ended up on this branch:

1. The MPS backend — csrc/mps_*, bitsandbytes/backends/mps/, cextension.py, _ops.py,
   autograd/_functions.py, backends/default/ops.py, tests/test_mps_parity.py,
   docs/apple_silicon/, benchmarks_wip/, and the build/packaging changes.
2. A Lion weight-decay fix, unrelated to MPS and worth its own PR. The Triton 8-bit blockwise
   kernel gates decoupled weight decay on OPTIMIZER_ID == 2, which is Adagrad — so Lion got
   coupled decay (corrupting its sign update) and Adagrad got decoupled decay instead of the L2
   fold it expects. bitsandbytes/backends/triton/kernels_optim.py + regression tests in
   tests/test_optim.py. This affects CUDA/Triton users and has nothing to do with this branch.
3. Local agent/dispatch tooling docs — .agents/**, agents/**, _typos.toml.