Repository navigation
llama: published MLX affine quants (8-bit and 4-bit) read in place, beside the bf16 default (builds on #1415) - #1416
BobClawblaw wants to merge 4 commits into
Conversation
|
Nothing converts at load, which is right. But when Also: |
A new `llama` family for the native Metal engine: dense Llama-architecture checkpoints (heads of 128, default RoPE, no biases, an untied head) in bf16, read in place from the shards with nothing converted. It keeps HF modeling_llama's rounding points (a bf16 residual, LlamaRMSNorm, RoPE with bf16 cos and sin, SiLU), runs prompt chunks on gemm_nax's tensor-unit GEMM and decode windows on a simdgroup-a-row matvec, decode attention on sdpa_vector lanes at absolute positions, and sampled draws on tf_sample_full. tf-llama-run checks greedy tokens, every position's logits, and that windows and one-row steps agree bit for bit. Measured on an M5 Ultra against HF's bf16 run of the release: 1,500 positions teacher-forced at mean KL 0.0015 and top-1 99.1%; decode 16.4 tokens/s a stream and 28-32 with two; cold prompts 1.3k tokens/s at 8k tokens and 563 at 131k. docs/recipes/nanbeige-4.3-33b-metal.md has the rest.
The family's Metal kernels through gates 1-5 (zig/tests/contract/metal_llama*.zig, the roll call and sources in contracts.zig): the glue (embedding, the norms at HF's rounding points, the residual add, RoPE with bf16 cos and sin, the SiLU gate, the decode attention's partials and their merge, the argmax draw by its stated NaN rule), the six row-projection forms, the four GEMM forms and the prompt attention at both served geometries, each with its float64 reference, served shapes and tiles, launched as forward.zig launches it; gate 5 reads the pipelines on this GPU. tf-kernel-contract --only llama on the M5 Ultra: 21 kernels, 21 pass gates 1-4, 0 fail, 0 known findings, 0 limit problems.
…out tensor units by name; the teacher-forced misses and the licence in the doc tf-llama-run takes exact's window width as an argument and loses its environment switches (LLAMA_LAYERS, LLAMA_WINDOW).
…eside the bf16 default
The family reads an MLX affine checkpoint as mlx_lm writes it: each projection
its own {u32 words, bf16 scales, bf16 biases} triple in groups of 64 at 4 or 8
bits, the embedding and head quantized, the norms bf16, and config.json's
uniform quantization table. Nothing is converted, by us or at load. Decode
windows run the core's row kernels, prompt chunks its tensor-unit affine
matmul; the embedding's rows come from a new family kernel, ll_embed_affine,
MLX's dequantize bf16(bf16(scale * q) + bias), registered with the contract.
affine_mm's 8-bit groups-of-64 form is registered too. capabilities lists
bf16, mlx-q4g64 and mlx-q8g64. The sums scratch is allocated only for
quantized checkpoints, so the bf16 path is as before.
Nanbeige4.3-33B on an M5 Ultra, two published quants pinned
(pratikgx/Nanbeige4.3-33B-MLX-4bit, fcmeyer/Nanbeige4.3-33B-MLX-8bit),
against HF's bf16 run of the release over 1,500 positions: 8-bit mean KL
0.00234, top-1 98.7% (MLX on the same files: 0.00272); 4-bit 0.0849, 91.1%
(MLX: 0.0851). Windows and one-row steps agree bit for bit on both; the
release's logits are bit-identical to before. The recipe doc has the tables;
the quants are lossy and the release stays the default.
73e9b7f to
0500646
Compare
|
Thanks, agreed. Pushed 0500646 (one commit on #1415's head f92916d, which is rebased onto main 15a0b2a): MLX's own layout, read in place. Checked on two published community quants, pinned:
The engine is as close to the release as MLX is on each quant. 16- and 7-row windows and one-row steps agree bit for bit on both. Every harness serving check passes on both: structured 18/18, tools 5/5, sampling 9/9, turns, soak 60/60, concurrent 12/12, streaming, multi-turn, stop, thinking, max_tokens. Contract. Capabilities. bf16 A/B. The The PR body is rewritten for the new scope (published quants, no conversion). On these commits over main a1e6d34: |
BobClawblaw
changed the title
|
Thanks @BobClawblaw. The published split MLX layout is now read in place ( Two asks:
#1415's second point (the exactness mode always exits 0) applies to this branch too. |
What
The
llamafamily serves published MLX affine quants (8-bit and 4-bit, groups of 64) of a Llama release in place, asmlx_lmwrites them, beside its bf16 default (builds on #1415). Nothing is converted, by us or at load. Checked on two community quants of Nanbeige4.3-33B, pinned. They are lossy; their cost against the release is measured below, and the bf16 release stays the recipe's default.How
weights.zig:Linear = union{bf16, affine}. Each of q, k, v, o, gate, up and down is read as its own{weight (u32 words), scales, biases}triple, and the embedding and head asmlx_lmquantizes them. Shapes, dtypes and byte extents are checked at load.config.zigadmits a uniformquantizationtable (affine, groups of 64, 4 or 8 bits) and refuses anything else by name.row_projection.zig), and prompt chunks the core's tensor-unit affine matmul (affine_mm.zig), as the other affine families do.ll_embed_affine(zig/kernels/metal/llama/glue.metal), which computes MLX's dequantizebf16(bf16(scale * q) + bias).ll_embed_affinehas gates 1-4 at 4 and 8 bits (zig/tests/contract/metal_llama.zig) and a limits entry.affine_mmat 8 bits in groups of 64, a shape the published 8-bit quants use, is registered inmetal_main.zigand the registry.tf-kernel-contract --only llama:SUMMARY 23 kernels: 23 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems;affine_mmq8 g64 passes gates 1-4.capabilities --jsonlistsbf16,mlx-q4g64andmlx-q8g64forllama.sumsscratch is allocated only for quantized checkpoints, so the bf16 path is as in llama: Nanbeige4.3-33B on Metal, the bf16 release served as published, every kernel in the contract #1415.Both state Apache-2.0, like the release.
Cost against the release (M5 Ultra; the HF reference and harness plans of #1415)
The reference is HF's bf16 run of Nanbeige/Nanbeige4.3-33B @ c1093c501cc6460427c0ac634db7223d0450c0ee over 1,500 teacher-forced positions. MLX is MLX 0.32.4 with
mlx_lm0.32.2 running the same quant files.The engine is as close to the release as MLX is on each quant. 16- and 7-row windows and one-row steps agree bit for bit on both. Every harness serving check passes on both: structured 18/18 (drafted == plain 12/12), tools 5/5, sampling 9/9, turns resumed 5,493 of 5,649 tokens, soak 60/60, concurrent 12/12, streaming, multi-turn, stop, thinking, max_tokens.
Receipt
--context 131072 --parallel 2 --prompt-cache-gib 8andLLAMA_LOAD_LIMIT_GB=150, with an idle server of another model loaded beside it.tools/bench_concurrent.py --levels 1,2 --tokens 128 --temperatures 0,0.7 --alone --serialon each quant: every reply equals its request alone and with"draft": false(alone 12/12, serial 12/12, failed 0, on each).tools/bench_openai.py --tokens 256 --reps 3, bf16 → 8-bit → 4-bit:bench_concurrent, greedy: 31.3-34.4 → 49.8-50.4 → 76.4-76.8tools/bench_cache.py --require-cold, a 9,000-token prompt cold: 7.18 s (8-bit) and 6.64 s (4-bit). Repeated raw-token prompts inbench_cache.pyreportcached_tokens0: the family's prompt cache resumes chat turns, and the harness's turns check resumes 5,493 of 5,649 tokens in 0.24-0.28 s against 4.15-4.49 s cold. The cold first request is what the receipt asks for.zig build test,test-kernel-contract,test-golden,fuzz-cortheon,dev-tools,test-server-security,test-pull-securityall exit 0;zig fmt --checkis clean; the lean check finds 0 new problems (428 on main and on the head);tf-kernel-contract --only llamaon the M5 Ultra:SUMMARY 23 kernels: 23 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems. Platforms: Metal on the M5 only; M1-M4 Macs, CUDA and HIP not run (no machine here; the family is Metal-only).tf-kernel-contract --only llama,test-kernel-contract,zig fmt --checkand the bf16, quant and drafted bit-identity checks, and all pass.