Skip to content

lfm3: add native Metal vision support above the text decoder - #1502

Draft
BlivionIaG wants to merge 131 commits into
ashhart:mainfrom
BlivionIaG:feature/lfm3-vl-support
Draft

BlivionIaG wants to merge 131 commits into
ashhart:mainfrom
BlivionIaG:feature/lfm3-vl-support

Conversation

@BlivionIaG

@BlivionIaG BlivionIaG commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor

What changes

Add native Metal image preparation, vision encoder and projector for LFM3-VL-2B from the official MLX 4/6/8-bit checkpoints. Projected image rows feed the LFM3 decoder and existing lane engine, including image order/content handling, prompt caching, verification, snapshots and concurrent serving. Python is used only for development references and clients.

Refs #1491. Depends on admission #1513 and text #1514. The PRs target upstream main until their foundations land; review the vision-only difference.

flowchart LR
    I[Image pixels] --> V[Vision encoder and projector]
    V --> E[Decoder embeddings]
    T[Text tokens] --> E
    E --> D[LFM3 decoder and existing lane engine]
Loading

Code pointers

Validation and numerical evidence

Pixel normalization correction: fused float32 normalization now matches the cited transformers/torchvision processor patch hash. The new regression fails before the change and passes after; host build 12/12 and image tests 26/26. Corrected-pixel vision checks reproduce all 143,360 projected words for VL4 and VL6 using the named rounded graph. Actual vision tensor, parameter, input and output identity supports VL6-to-VL8 vision-only transfer. New ideal FP64 metrics remain recorded. Earlier image receipts retain their original pixel scope; Corrected-pixel serving checks pass 183/183 at VL4/6/8, including 27 cancellation checks and synthetic ordered-image requests. Corrected-input decoder references reproduce all 34,560,000 native BF16 words under the prescribed rounded graph across three independent 90-row runs. Ideal FP64 all-row top differences/disjoint sets are 2/1, 4/4 and 7/2 at 4/6/8 bits; suffix counts are 0/0, 1/1 and 1/0. Every difference and both margins remain listed; no ideal-FP64 acceptance is asserted.

Recipe and checkpoint revisions · Receipts and reproduction tools

  • The combined range-correction/norm-fusion candidate passed locked local M4 qualification across all six checkpoints: 408/408 state checks, 303/303 serving checks, including sampled/cancellation/image cases, and 394,752,000 logit bytes identical to the previous public runtime on the retained checkpoint fixtures. The range receipt separates candidate execution from the matching applied-source build: 78/78 steps, formatting passed, and zero new lean findings over 428 inherited.
  • The attention correction preserves finite contributions lost by FP32 underflow/overflow. Registered original-fail/new-pass cases use unchanged FP64 bounds; disabling repair or clearing risk makes the failures return. Guarded M5 testing separately passed the current 70-kernel LFM3 contract set and normal limits. These targeted results are not a complete-backend pass.
  • Paired projections, shared convolution and fused post-add/next-norm dispatch retain the original rounding boundaries. Their individual receipts identify tested source/binary hashes, exact-output checks and controls. Live CI status is shown by the PR checks.

Independent prompt evidence retains all 5,581 teacher-forced rows, changed top tokens and both margins. Named production-rounding replay resolves disjoint maximum sets; 19 lowest-index differences remain overlapping maxima. All six genuine one-row decode captures after 63 identical history tokens have zero top differences. No new tie policy or tolerance is asserted.

The frozen real-cat whole-pipeline comparison matches all 16 suffix tops at 4/6 bits but has two disjoint maxima against ideal FP64 at 8 bits. An independent CPU arithmetic replay, explicitly delegating named scalar functions to Metal, reproduces all native vision/decoder stages and logits byte-for-byte. This establishes rounding causality, not independent-FP64 acceptance. The FP32 supplement retains both differences in pure FP32; with model BF16 boundaries it has zero disjoint suffix maxima, while differences remain in the full 90-row scope. All original comparisons remain available for review.

The earlier frozen full M5 backend run at 8955faee0 reports 310 kernels: 309 pass, one inherited Qwen4 HC failure, two known findings; all 317 controls red; zero limit problems. The Qwen4 failure reproduces on pinned base b30d5b9d. It is not a global pass and does not qualify later kernels by implication.

Measurements and qualification

Current M5 long-context checks: all 26 requests completed: VL6/8 text prompts at 32k/64k, text4/6/8 at 131072/262128 prompt tokens, and VL4/6/8 at 127984 prompt tokens, each repeated in a fresh process. All return 16 tokens with zero cached input; all 13 repeated output-token hashes match. All 16 servers exit zero, all three packet leases are released, and the controller exits zero. These are native text-context checks, including VL checkpoints; they do not establish image-context, FP64 or comparative speed acceptance.

Extracted text artifact: 15 focused state checks, 120 serving checks and 18 cancellation checks passed across text 4/6/8-bit. Six prompt/decode comparisons match 49,920,000 bytes against the combined artifact. Public native/forward/exact build: 65/65 steps.

Final 6/8-bit checks: 272 state checks, 202 serving checks and 30 cancellation checks passed. Current outputs match 145,981,440 retained native bytes. Six distinct VL6 numerical contexts reproduce native logits exactly under the prescribed rounded graph; ideal FP64 differences remain listed. Long-context checks are recorded below.

Current M5 backend: 315/316 kernels passed, one Qwen4 HC failure matching the recorded baseline case, two known affine nonfinite classes, 321/321 controls red, zero limit problems. Actual contract exit: 1.

M5 four-bit code decode, solo / four streams: text 315.75 / 861.85 tok/s, standard 344.20 / 995.0; VL 298.9 / 835.65, standard 325.8 / 947.75. Requests produce 256 tokens. Cold requests produce 16 tokens with zero cached input and a 70k request budget. Cache allocation differs between engines; standard reply hashes are unavailable.

Four-bit M5 measurements, 6/8-bit M5 measurements, M4 measurements. The 6/8-bit packet records 192 repetitions, 480 primary replies and 32 cold replies. Native cold 2k / 8k TTFT spans 0.85–0.96 / 4.51–4.99 s; standard spans 0.15–0.17 / 0.38–0.39 s.

Final local artifact: build 72/72; tests 7,232 passed, 45 skipped; LFM3 contracts 70/70, two controls red, zero limit problems; four-bit serving 101/101, including 15 cancellation checks. Formatting passes; lean has zero new findings over 428 inherited. Full-backend execution on M4 stops at a tensor-unit requirement. Final 6/8-bit short-model qualification is recorded in the linked receipt; long-context checks are recorded above.

Four-bit numerical ledger: all four distinct observed disjoint contexts are reproduced by the independent prescribed-rounding graph. This includes all 70 projected image rows and both disputed image-placeholder heads. Ideal FP64 token differences and margins remain recorded. 32k/64k checks cover text4 and VL4 text prompts, with matching fresh-process token hashes.

Remaining checks: shared prompt/decode path requirements, merged-chain integration, and maintainer review of recorded FP64 differences and performance comparisons. No optimization experiments are scheduled.

Weights: LFM Open License v1.0. Checked machines: M4 Pro 48 GB and M5 Ultra 256 GB; these measurements do not establish smaller-memory qualification. ATC issue 1491 remains held.

Prepare text and VL 4/6/8-bit checkpoints with revision-pinned header inspection and native metadata validation. Inference remains unregistered pending operator and model qualification.

Refs ashhart#1491
Serve LFM3-1.7B and LFM3-VL-2B MLX 4/6/8-bit checkpoints through shared lane verification, prompt snapshots and native image rows. Preserve checkpoint quantization and BF16 activation boundaries with FP32 reductions.

Add independent MLX references, real-checkpoint exactness and served cancellation checks, and Metal kernel contracts. Refs ashhart#1491.
@ashhart

ashhart commented Oct 27, 2026

Copy link
Copy Markdown
Owner

Thanks @BlivionIaG for a careful port and an honest receipt. The same-engine checks are thorough (231 served, with cancellation and recovery), and the rotary fix comes with a test that fails on the old formula.

The merge with main is clean. Here is our answer on the gates.

Fidelity:

  • Compare teacher-forced against an FP64 forward of the dequantized checkpoint. MLX is a second opinion, not the reference.
  • Compare prompt rows and one-row decode rows separately, on identical input rows, at 4, 6 and 8 bits.
  • List every top-1 difference with both tokens and both margins, and report the worst rows, not only an average.
  • A difference passes only as a near-tie that a reference rounding where the engine rounds reproduces, as lfm3: serve LFM3-1.7B on CUDA: NVFP4 and FP8 weight-only, MLX 4-bit (stacked on #1513, #1492 and #1503) #1508 does for the CUDA side. The 18 rows with disjoint top sets need tracing first.
  • For VL: the projected rows of real images against FP64, then teacher-forced text after them.

Performance, on the same Mac and checkpoint against mlx-lm (mlx-vlm for VL), at each width:

  • one-stream decode and the 4-stream aggregate, with outputs of equal length;
  • cold prefill at 2k, 8k, 32k and 64k. Decode should be at least the standard server's. Cold prefill should take the prompt tiles, as below; a remaining prefill gap is fine to track as its own issue, as we do for the 27B.

Before this leaves draft:

  1. Prompt rows go through the row kernels and the decode attention (zig/src/families/lfm3/forward.zig:39-49, :196-216). That is why cold prefill is 7 to 18 times slower than mlx-lm's. LFM2's Metal family takes tile_qmm's prompt tiles with a decode twin of the same bits (zig/src/families/lfm2_moe/forward.zig:30), and the tensor-unit prompt attention on an M5 (:244). Please build on those, and on LFM2: shared-round decode #1377's shared rounds for the 4-stream aggregate.
  2. Please split the PR: the admission and the text decoder first, with VL stacked on them. lfm3: serve LFM3-1.7B on CUDA: NVFP4 and FP8 weight-only, MLX 4-bit (stacked on #1513, #1492 and #1503) #1508 already builds on the admission.
  3. Please state the weights' licence, and the Macs it was checked on, in the recipe.

Once the gates hold, we'll run the full Metal contract and the served checks on our M5 Ultra.

Add a direct-rounding CPU FP64 oracle and image-conditioned tracing. Exercise production-shared FP32 reductions with real reduced-precision controls. Bound sliding attention dispatch without changing partial slots or merge order, and reject malformed diagnostic inputs.
@BlivionIaG BlivionIaG changed the title lfm3: add native Metal text and vision support for MLX 4/6/8-bit lfm3: add native Metal vision support above the text decoder Oct 28, 2026
@jschmied

Copy link
Copy Markdown
Contributor

@BlivionIaG, one small pixel-parity note on families/lfm3/image.zig's patchify, from the LFM3-VL processor work on CUDA.

transformers' Lfm3VlImageProcessor (5.19, torchvision backend) folds the rescale into the normalization. It computes (x - mean / rescale) / (std / rescale) in float32, which for LFM3 is (x - 127.5) / 127.5. The lookup table here computes x * f32(1/255) first and then (v - 0.5) / 0.5. The two disagree on 207 of the 256 entries, by at most 1.2e-7 (one ulp near ±1). So the patches can't be bit-equal to the processor's pixel_values.

The resize (.torch) and the geometry already match. Fixing it is one line in the table:

value.* = (@as(f32, @floatFromInt(i)) - 127.5) / 127.5;

With that, the processor's output should be reproduced bit for bit.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants