fix(vision): align Gemma 3n VLM per-layer-inputs to padded prefill by inureyes · Pull Request #36 · lablup/mlxcel

inureyes · 2026-05-19T23:14:22Z

Summary

Fixes #35 — Gemma 3n VLM single-row prefill aborted with a [broadcast_shapes] mismatch (e.g. (1,288,256) vs (1,273,256)) whenever the prefill path padded the token stream and merged embeddings up to a tile-aligned length. The projected per-layer-inputs tensor is captured in get_input_embeddings before that generic padding step, so the per-layer blend in forward_with_embeddings tried to broadcast the unpadded per-layer tensor against the padded embeddings and threw.

Fix

Add Gemma3nVLModel::align_per_layer_inputs_to_embeddings: pads the per-layer tensor with zeros when shorter than the embeddings sequence, slices when longer, and returns None (no clone) when they already match.
The VLM prefill path runs the cached per-layer tensor through this helper before forward_with_inputs, so a tile-padded embedding stream no longer mismatches the per-layer blend.

Tooling (`mlxcel-bench-decode`)

Port the mlxcel-bench-decode same-process benchmark harness: load the model once, run warmup, reset model/cache state, then run the measured pass in the same process — mirroring Python stream_generate timing far better than two cold mlxcel generate invocations.
Rewrite scripts/bench_decode.sh to drive this binary (single load, --no-chat-template passthrough, combined warmup+measured timeout).
Add scripts/bench_mlxlm.py as the matching mlx-lm / mlx-vlm baseline sweep.

Docs

Refresh the README parity tables and M5 Max benchmark docs: the three Gemma 3n VLM rows move from FAIL (broadcast abort) to passing decode throughput, and a Gemma 3n E2B row joins the representative decode table.

Test plan

cargo fmt --check
cargo clippy --features metal,accelerate --all-targets -- -D warnings
cargo build --release --features metal,accelerate --bin mlxcel-bench-decode (binary links)
align_per_layer_inputs_* unit tests (pad / slice / no-op) — 3 passed
gemma3n_vlm_single_row_bench_prefill_accepts_m5_tile_padding — ignored, weight-gated; run locally with cargo test --release --test vlm_concurrency -- --ignored against a Gemma 3n VLM checkpoint

Gemma 3n VLM single-row prefill aborted with a `[broadcast_shapes]` mismatch (e.g. `(1,288,256)` vs `(1,273,256)`) whenever the prefill path padded the token stream and merged embeddings up to a tile-aligned length. The projected per-layer-inputs tensor is captured in `get_input_embeddings` before that generic padding step runs, so the per-layer blend in `forward_with_embeddings` then tried to broadcast the unpadded per-layer tensor against the padded embeddings and threw. Hidden dim 256 is Gemma 3n's per-layer-inputs hidden size, confirming the broken broadcast is between the per-layer tensor and the main token stream. Add `Gemma3nVLModel::align_per_layer_inputs_to_embeddings`, which compares the per-layer tensor's sequence length against the embeddings sequence length and pads with zeros when shorter, slices when longer, or returns `None` (no clone) when they already match. The VLM prefill path aligns the cached per-layer tensor through this helper before handing it to `forward_with_inputs`, so a tile-padded embedding stream no longer mismatches the per-layer blend. Three unit tests cover the pad, slice, and no-op branches. Port the `mlxcel-bench-decode` same-process benchmark harness from the internal tree so the VLM prefill path is reproducible and testable here. The binary loads a model once, runs the warmup pass, resets model/cache state, then runs the measured pass in the same process, mirroring the Python `stream_generate` timing far better than two cold `mlxcel generate` invocations. `scripts/bench_decode.sh` is rewritten to drive this binary (single load, `--no-chat-template` passthrough, combined warmup+measured timeout), and `scripts/bench_mlxlm.py` is added as the matching `mlx-lm` / `mlx-vlm` baseline sweep. Add an ignored, weight-gated regression test (`gemma3n_vlm_single_row_bench_prefill_accepts_m5_tile_padding`) that drives `mlxcel-bench-decode` against a real Gemma 3n VLM checkpoint so future shape drift on the tile-padded path is caught. Refresh the README parity tables and the M5 Max benchmark docs: the three Gemma 3n VLM rows move from FAIL (broadcast abort) to passing decode throughput, and a Gemma 3n E2B row joins the representative decode table.

Bump version to 0.0.29. Highlights: model-owned VLM fallback state reset (#34), Gemma 3n VLM padded prefill alignment plus the new mlxcel-bench-decode same-process benchmark harness (#36), and benchmark report reorganization with refreshed Apple Silicon and M5 Max VLM tables (#32, #33).

… (carry-across items banked where surgery reads) Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

…als, synth + detach round-trip THE FIELD DECISION (Violet deep-read bar item 2, the audit): kvarn_v_bits on KVCache (8 default, 4 = K8V4), mode STAYS KVarN8. A new enum variant would silently fall OUT of every matches!(mode, KVarN8) gate — is_trimmable, the scheduler tripwire's can_trim_prompt_cache, paged bypass, detach/snapshot refusals, D1 empty-downgrade — the day's own silent-wrong class. The field keeps every mode gate correct by construction; the inversion cost (sites needing V4-different behavior must OPT IN) is paid with loud reader guards. MODE-SITE AUDIT (complete, per site): - update dispatch (1070): bits consulted inside update_kvarn8 [HANDLED] - supports_block_fetch (1378): true for v4 -> selection proceeds -> the blocks-fetch guard panics; either dispatch route ends at a loud refusal [CORRECT-WITH-GUARD] - kvarn_qmm_state: PANIC not None — v_s_row is None on V4, so the '?' chain would return None and C would silently fall through wearing the wrong-core label [GUARDED] - fetch_kvarn8 / fetch_kvarn8_blocks: loud refusals until read paths land behind the construction key (design §5) [GUARDED] - is_trimmable / tripwire / trim live_start arm / sliding-window refusal / paged bypass: kvarn-generic, correct for v4 as-is [OK] - D1 downgrade: empty-only, no V data exists [OK] - FromStr/Display: construction+display strings are the CLI commit's surface; v4 unconstructable until then [DEFERRED WITH CLI] - gemma4 snapshot: refuses all non-Fp16 upstream (gemma4.rs:901) — the mode-tag mapping is dead code for kvarn [OK] - turbo boundary.rs: KVarN8 -> nominal, width-agnostic [OK] - DetachedKVCache nbytes: sums actual array sizes, auto-correct [OK] DETACH FINDING (design §3.3 'verified absence' was STALE): kvarn state round-trips detach/adopt (prompt-cache donation) field-by-field, and mode alone cannot distinguish k8v8 from k8v4 — a v4 donation re-adopted into a fresh slot would resurrect as v_bits=8, silently mislabeled, PASSING the reader guards. kvarn_v_bits now travels: DetachedKVCache field + clone_handle + install_detached + pool_backed_handle_clone. Detached trim_to verified width-agnostic (slices take s[3] from the existing shape). BIRTH CONSTRAINT (lablup#36 seam): full_len in update_kvarn8 is THE finalize-boundary variable, comment names it and cites DESIGN_finalize_cap_at_true_length — the cap lands as min(natural, cap) there and nowhere else. V4 branch routes through the same variable. SYNTH SAME-COMMIT (design §3.3 critical addition): synth_kvarn_state gains v_bits; synth_kvarn8_state delegates at 8 (callers unchanged); layout authority pinned to the engine writer by test. TESTS (cache:: 466/466, kvarn 38/38 green): layout pin, THE wiring pin (stored fields bitwise == kvarn_quantize_v4 on the update path's own tiles), K-side bitwise-equal across widths, k8v8 regression pin (bitwise == kvarn_quantize direct — guards the branch restructure), three loud-refusal should_panics, multi-update append consistency, synth-vs-engine layout equality, detach/adopt v_bits round-trip incl. V4-shape-preserving trim. MUTATIONS OBSERVED RED: (2) install-side v_bits restore dropped -> roundtrip test red at the named assertion (left 8, right 4). Mutation (1) scale/zp destination swap -> wiring test red — observed pre-replay; re-proven against THIS committed base in the follow-up run recorded in the PR notes. HONEST RECORD: mid-surgery I destroyed the uncommitted cache.rs edits with a reflexive 'git checkout' meant to revert a 10-line test mutation — the same broad-revert-for-narrow-undo class as the cycle-57 rm. All edits replayed verbatim from context; this commit exists BEFORE further mutation work because mutations belong on committed bases. FULL-SUITE CAVEAT: mlxcel-core full suite cannot complete on this host right now — the resident 295K mlxcel-server holds unified memory and the speculative/pipeline tests die variably (SIGABRT/SIGTRAP/SIGSEGV at different tests across runs; stash-verified PRE-EXISTING on the clean branch). 901+ tests pass before memory death; the surgery's blast radius (cache::, kvarn) is deterministically green. Full-suite verification owed at the next no-resident-server window or CI. Also: Gate B results doc range-shorthand fix (25-29 wrongly included layer 27 at 0.0187; enumerated set 25,26,28,29 — Violet QE catch). Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

… state The doc claimed 'detach/snapshot refusal sits at mode level upstream of the fields' as a verified absence. Stale-wrong for the in-memory donation path (DetachedKVCache carries the full kvarn block through clone_handle/install_detached); true only for gemma4's persisted snapshots (non-Fp16 refused upstream). The doc had conflated two paths. Corrected in place with the surgery commit cited; codified-wrong is worse than not-codified. Status block gains the landed build state (34df406, Xander APPROVE, Violet deep read pending, lablup#36 at junction). Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

THE FIX (DESIGN_finalize_cap_at_true_length, Violet's design, built at the junction): pending_finalize_cap: Option<i32> on KVCache — absolute exclusive position, set by the scheduler pre-forward at padded sites, CONSUMED (take at entry, one-shot) by update_kvarn8. The cap lands at the finalize-boundary variable and nowhere else (the K8V4 birth constraint paying out): a tile finalizes only if it lies entirely below the cap, so prefill padding stays in the fp16 end-state. Both padding wrongs die at the root — no tile/offset desync (wrong 1), no Sinkhorn pollution (wrong 2, asserted bitwise). DESIGN AMENDMENT (found at source, Violet-confirmed + blessed): the doc's premise 'the existing dense trim already handles the tail' was wrong for kvarn — trim() touched zero kvarn fields, so capped padding would have desynced IN THE TAIL (wrong-1 relocated). trim() now dispatches KVarN8 to trim_kvarn_fp16_end: tail first, sink only while no tiles exist (short padded prefills park padding in the SINK — the cap deliberately does not stop the fp16 sink fill), REFUSE loudly with zero mutation if quantized tiles would be reached. is_trimmable stays false (general rewind still unsupported; spec-decode keeps gating). m3_idx_offset rolls back in lockstep — the idx accumulation received the padded rows too, and selection reading idx past offset is the cycle-79 desync class. [FILED, not fixed here: the DENSE trim has the same m3_idx gap on fp16 M3 caches under padded concurrency — behavior change on fp16 production, wants its own sizing like lablup#37.] PREDICATE SHARPENED: padding_trim_would_corrupt now asks per-cache can_trim_padding(excess) — true exactly when the arm can act (fp16 end-state suffices). The tripwire passes exactly when trimming is genuinely safe and stays armed for uncapped leaks (defense-in-depth). Existing empty-cache contract holds (no fp16 end-state → corrupt). TESTS (cache:: 475/475): six design edges — mid-tile cap, identity, beyond-end clamp, consume-and-clear, THE padded-roundtrip golden (stored fields AND fetched window AND offset-reconciliation bitwise vs the unpadded reference, THROUGH the arm), Sinkhorn purity field-level — plus refuse-zero-mutation, tail-bounded success with m3 lockstep, sink-case at the hist-empty boundary, and predicate-mirrors-arm. Named mutations (consume-clear drop; refuse-check drop) run against THIS commit next — mutations only ever on committed bases. Scheduler site placements (4) follow as the wiring commit. Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

…m landed, wiring to the fresh window Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

…ed for the fresh window Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

…-prefill sites The mechanism (99d7ff6, Xander mechanism APPROVE) consumes a one-shot cap in update_kvarn8; this commit sets it. One helper carries the load-bearing form, four sites carry their gates and lengths: - mlxcel_core::cache::set_prefill_finalize_caps(caches, true_new_rows): per-cache cap = c.offset + true_new_rows, sibling to padding_trim_would_corrupt (the pre-forward arm beside the post-forward tripwire). The form equals the design doc's actual_len[i] when offset==0 (fresh batched prefill) and stays correct where the doc's fresh-prefill shorthand would not: adopted prompt-cache prefixes and chunk continuations, where the cache already holds rows. - execute_batched_prefill: per sequence per layer cache, unconditional (identity for max-len sequences — golden-pinned bitwise at the edges). - execute_full_prefill / start_chunked_prefill / continue_chunked_prefill: gated on pad_mask_opt.is_some() — exactly the condition under which each site's own post-forward trim runs. The tripwire at all four sites STAYS — defense-in-depth for any uncapped path (DESIGN_finalize_cap_at_true_length §tripwire-relationship). Wiring pins live at the consumer (scheduler prefill paths need a live model; the scheduler test harness is policy-level by design): - sites_helper_general_form_golden_over_prior_rows: prior rows + helper cap + padded suffix + arm trim == unpadded twin, bitwise (kvarn_state_snapshot). Named mutation: drop 'c.offset +' from the helper — update_kvarn8's cap_rel >= 0 scheduler-bug assert fires red. - sites_helper_mode_blind_across_mixed_slice: the live M3 slice shape (Fp16 D1-downgraded dense-prefix + KVarN8 MSA): kvarn consumes at entry; fp16 carries the field inertly, bit-identical to a never-capped twin. Named mutation: mode-gate the helper — the fp16-armed assert fires red. cache:: 477/477 (475 mechanism + 2 wiring pins). cargo check clean on mlxcel-core + mlxcel. Toolchain note: local rustc 1.96.0 (Homebrew) vs CI's pinned 1.93.1 for cargo fmt --all --check — local fmt wanted to reformat 10 files this branch never touched, so no fmt sweep rides this commit (diff is 166 pure additions); any fmt pass belongs to PR-prep under the pinned toolchain. Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

…ws pending Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

…57850e cherry-picked, base topology recorded Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

…(nbytes kvarn-blind, confirmed), finding-3 refuted at source, lablup#36 to Clement Co-Authored-By: violet-14057653 <violet@sjstargetedsolutions.co.nz> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…tion witnessed), Clement crossing, wiring = far-side pickup, lablup#38 filed Co-Authored-By: violet-14057653 <violet@sjstargetedsolutions.co.nz> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…orm at all four sites, cross-commit defense-in-depth proven, cap ARMED Co-Authored-By: violet-14057653 <violet@sjstargetedsolutions.co.nz> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

… approve 18:09); f57850e cherry-picked to branch, board stays base-side Co-Authored-By: violet-14057653 <violet@sjstargetedsolutions.co.nz> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ablup#36) Gemma 3n VLM single-row prefill aborted with a `[broadcast_shapes]` mismatch (e.g. `(1,288,256)` vs `(1,273,256)`) whenever the prefill path padded the token stream and merged embeddings up to a tile-aligned length. The projected per-layer-inputs tensor is captured in `get_input_embeddings` before that generic padding step runs, so the per-layer blend in `forward_with_embeddings` then tried to broadcast the unpadded per-layer tensor against the padded embeddings and threw. Hidden dim 256 is Gemma 3n's per-layer-inputs hidden size, confirming the broken broadcast is between the per-layer tensor and the main token stream. Add `Gemma3nVLModel::align_per_layer_inputs_to_embeddings`, which compares the per-layer tensor's sequence length against the embeddings sequence length and pads with zeros when shorter, slices when longer, or returns `None` (no clone) when they already match. The VLM prefill path aligns the cached per-layer tensor through this helper before handing it to `forward_with_inputs`, so a tile-padded embedding stream no longer mismatches the per-layer blend. Three unit tests cover the pad, slice, and no-op branches. Port the `mlxcel-bench-decode` same-process benchmark harness from the internal tree so the VLM prefill path is reproducible and testable here. The binary loads a model once, runs the warmup pass, resets model/cache state, then runs the measured pass in the same process, mirroring the Python `stream_generate` timing far better than two cold `mlxcel generate` invocations. `scripts/bench_decode.sh` is rewritten to drive this binary (single load, `--no-chat-template` passthrough, combined warmup+measured timeout), and `scripts/bench_mlxlm.py` is added as the matching `mlx-lm` / `mlx-vlm` baseline sweep. Add an ignored, weight-gated regression test (`gemma3n_vlm_single_row_bench_prefill_accepts_m5_tile_padding`) that drives `mlxcel-bench-decode` against a real Gemma 3n VLM checkpoint so future shape drift on the tile-padded path is caught. Refresh the README parity tables and the M5 Max benchmark docs: the three Gemma 3n VLM rows move from FAIL (broadcast abort) to passing decode throughput, and a Gemma 3n E2B row joins the representative decode table.

Bump version to 0.0.29. Highlights: model-owned VLM fallback state reset (lablup#34), Gemma 3n VLM padded prefill alignment plus the new mlxcel-bench-decode same-process benchmark harness (lablup#36), and benchmark report reorganization with refreshed Apple Silicon and M5 Max VLM tables (lablup#32, lablup#33).

inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:high High priority area:models Model architectures, weights, loading, metadata labels May 19, 2026

inureyes merged commit e75cd70 into main May 19, 2026
4 checks passed

inureyes deleted the fix/gemma3n-vlm-padded-prefill branch May 19, 2026 23:15

inureyes self-assigned this May 19, 2026

inureyes added this to the 0.1 milestone Jun 21, 2026

sjswerdloff added a commit to sjswerdloff/mlxcel that referenced this pull request Jul 11, 2026

handoff: Xander mechanism APPROVE recorded — lablup#36 wiring unblock…

436c49d

…ed for the fresh window Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

sjswerdloff added a commit to sjswerdloff/mlxcel that referenced this pull request Jul 11, 2026

handoff: lablup#36 wiring landed @ 52a2154 — evidence recorded, revie…

0960b23

…ws pending Co-Authored-By: clement-7074f29f <clement-7074f29f@sjstargetedsolutions.co.nz>

inureyes added the status:done Completed label Sep 6, 2026

Navigation Menu

Repository navigation

fix(vision): align Gemma 3n VLM per-layer-inputs to padded prefill - #36

fix(vision): align Gemma 3n VLM per-layer-inputs to padded prefill#36
inureyes merged 1 commit into
mainfrom
fix/gemma3n-vlm-padded-prefill

inureyes commented May 19, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

inureyes commented May 19, 2026

Summary

Fix

Tooling (mlxcel-bench-decode)

Docs

Test plan

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Tooling (`mlxcel-bench-decode`)